jiangxt2 commented on code in PR #54926:
URL: https://github.com/apache/arrow/pull/54926#discussion_r4434229516
##########
python/pyarrow/_orc.pyx:
##########
@@ -364,6 +366,57 @@ cdef class ORCReader(_Weakrefable):
return pyarrow_wrap_batch(sp_record_batch)
+ def iter_batches(self, int64_t batch_size, stripes=None, columns=None):
+ cdef:
+ shared_ptr[CRecordBatchReader] batch_reader
+ shared_ptr[CRecordBatch] batch
+ std_vector[c_string] names
+ int64_t stripe
+
+ if batch_size <= 0:
+ raise ValueError("batch_size must be greater than zero")
+ if columns is not None:
+ names = [tobytes(name) for name in columns]
+
+ # Keep self (and its source) alive for the generator's entire lifetime.
+ if stripes is None and (columns is None or len(columns) != 0):
+ stripe_indices = [None]
+ elif stripes is None:
+ stripe_indices = range(self.nstripes())
+ else:
+ stripe_indices = stripes
+
+ try:
+ for index in stripe_indices:
+ if index is None:
+ with nogil:
+ batch_reader = GetResultValue(
+ self.reader.get().GetRecordBatchReader(batch_size,
names))
+ else:
+ stripe = index
+ if columns is None:
+ with nogil:
+ batch_reader = GetResultValue(
+ self.reader.get().GetStripeReader(stripe,
batch_size))
+ else:
+ with nogil:
+ batch_reader = GetResultValue(
+ self.reader.get().GetStripeReader(
+ stripe, batch_size, names))
+ while True:
+ if self.source_handle.get().closed():
Review Comment:
Thanks for flagging this. The existing ORC open path already calls
`get_reader()` → `get_random_access_file()` → `_assert_readable()` →
`_assert_open()`, which checks `self.closed`
([source](https://github.com/apache/arrow/blob/949bf7baadf04da7842065c5e7cb99f9da124a29/python/pyarrow/io.pxi#L209-L227)).
For Python-backed inputs, this reaches `PythonFile::closed()` before
`ORCFileReader::Open`, so a missing `.closed` already causes opening to fail
([source](https://github.com/apache/arrow/blob/949bf7baadf04da7842065c5e7cb99f9da124a29/python/pyarrow/src/arrow/python/io.cc#L73-L88)).
I have retained the per-batch check to detect a caller closing the source
while ORC still has buffered rows. The existing
`test_iter_batches_closed_source_with_buffered_rows` covers this behavior.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
