apacheGH-41771: [C++] Iterator releases its resource immediately when…

… it reads all values (apache#41824) ### Rationale for this change `Iterator` keeps its resource (`ptr_`) until it's deleted but we can release its resource immediately when it reads all values. If `Iterator` keeps its resource until it's deleted, it may block closing a file. See apacheGH-41771 for this case. ### What changes are included in this PR? Releases `ptr_` when `Next()` returns the end. ### Are these changes tested? Yes. ### Are there any user-facing changes? Yes. * GitHub Issue: apache#41771 Authored-by: Sutou Kouhei <[email protected]> Signed-off-by: Benjamin Kietzman <[email protected]>
mathworks · May 28, 2024 · e6e00e7 · e6e00e7
1 parent fe2d926
commit e6e00e7
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 3 deletions.
diff --git a/cpp/src/arrow/util/iterator.h b/cpp/src/arrow/util/iterator.h
@@ -105,9 +105,18 @@ class Iterator : public util::EqualityComparable<Iterator<T>> {
   Iterator() : ptr_(NULLPTR, [](void*) {}) {}
 
   /// \brief Return the next element of the sequence, IterationTraits<T>::End() when the
-  /// iteration is completed. Calling this on a default constructed Iterator
-  /// will result in undefined behavior.
-  Result<T> Next() { return next_(ptr_.get()); }
+  /// iteration is completed.
+  Result<T> Next() {
+    if (ptr_) {
+      auto next_result = next_(ptr_.get());
+      if (next_result.ok() && IsIterationEnd(next_result.ValueUnsafe())) {
+        ptr_.reset(NULLPTR);
+      }
+      return next_result;
+    } else {
+      return IterationTraits<T>::End();
+    }
+  }
 
   /// Pass each element of the sequence to a visitor. Will return any error status
   /// returned by the visitor, terminating iteration.

diff --git a/cpp/src/arrow/util/iterator_test.cc b/cpp/src/arrow/util/iterator_test.cc
@@ -146,6 +146,49 @@ void AssertIteratorNext(T expected, Iterator<T>& it) {
   ASSERT_EQ(expected, actual);
 }
 
+template <typename T>
+class DeleteDetectableIterator {
+ public:
+  explicit DeleteDetectableIterator(std::vector<T> values, bool* deleted)
+      : values_(std::move(values)), i_(0), deleted_(deleted) {}
+
+  DeleteDetectableIterator(DeleteDetectableIterator&& source)
+      : values_(std::move(source.values_)), i_(source.i_), deleted_(source.deleted_) {
+    source.deleted_ = nullptr;
+  }
+
+  ~DeleteDetectableIterator() {
+    if (deleted_) {
+      *deleted_ = true;
+    }
+  }
+
+  Result<T> Next() {
+    if (i_ == values_.size()) {
+      return IterationTraits<T>::End();
+    }
+    return std::move(values_[i_++]);
+  }
+
+ private:
+  std::vector<T> values_;
+  size_t i_;
+  bool* deleted_;
+};
+
+// Generic iterator tests
+
+TEST(TestIterator, DeleteOnEnd) {
+  bool deleted = false;
+  Iterator<TestInt> it(DeleteDetectableIterator<TestInt>({1}, &deleted));
+  ASSERT_FALSE(deleted);
+  AssertIteratorNext({1}, it);
+  ASSERT_FALSE(deleted);
+  ASSERT_OK_AND_ASSIGN(auto value, it.Next());
+  ASSERT_TRUE(IsIterationEnd(value));
+  ASSERT_TRUE(deleted);
+}
+
 // --------------------------------------------------------------------
 // Synchronous iterator tests