From 11ed9b12d6d001aa1b70754846da7dd6d8336598 Mon Sep 17 00:00:00 2001
From: Sutou Kouhei <kou@clear-code.com>
Date: Sat, 25 May 2024 12:33:25 +0900
Subject: [PATCH] xxx

---
 cpp/src/arrow/util/iterator.h       | 12 +++++++-
 cpp/src/arrow/util/iterator_test.cc | 43 +++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/cpp/src/arrow/util/iterator.h b/cpp/src/arrow/util/iterator.h
index 5e716d0fd113d..95ae018775658 100644
--- a/cpp/src/arrow/util/iterator.h
+++ b/cpp/src/arrow/util/iterator.h
@@ -107,7 +107,17 @@ class Iterator : public util::EqualityComparable<Iterator<T>> {
   /// \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()); }
+  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
index ba21ddcced209..a247ba13aef73 100644
--- 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