wgtmac commented on code in PR #54926:
URL: https://github.com/apache/arrow/pull/54926#discussion_r4421095053
##########
cpp/src/arrow/adapters/orc/adapter.h:
##########
@@ -135,6 +135,25 @@ class ARROW_EXPORT ORCFileReader {
/// \param[in] row_number the rows number to seek
Status Seek(int64_t row_number);
+ /// \brief Read a designated stripe incrementally without changing the Seek
position.
+ ///
+ /// \param[in] stripe the stripe index
+ /// \param[in] batch_size the maximum number of rows per batch, which must
be positive
+ /// \return an independent reader for all fields in the stripe
+ Result<std::shared_ptr<RecordBatchReader>> GetStripeReader(int64_t stripe,
+ int64_t
batch_size);
+
+ /// \brief Read selected fields of a designated stripe incrementally.
+ ///
+ /// This does not change the Seek position. Field order follows the file
schema.
+ /// An empty selection produces zero-column batches with the stripe's row
count.
+ /// \param[in] stripe the stripe index
+ /// \param[in] batch_size the maximum number of rows per batch, which must
be positive
+ /// \param[in] include_names the selected field names, including nested
field paths
+ /// \return an independent reader for the selected fields in the stripe
+ Result<std::shared_ptr<RecordBatchReader>> GetStripeReader(
+ int64_t stripe, int64_t batch_size, const std::vector<std::string>&
include_names);
Review Comment:
Now `include_names` treats an empty name list as zero-column projection,
while `GetRecordBatchReader()` and `NextStripeReader()` treat an empty
selection as “all fields”. Unfortunately this inconsistency already exists
prior to this change.
Perhaps define it as `std::optional<std::vector<std::string>>`,
`std::optional<std::reference_wrapper<const std::vector<std::string>>>` or
`std::optional<std::span<const std::string>>` to explicitly provide `none`,
`all` and `partial` column selection in a single function?
##########
cpp/src/arrow/adapters/orc/adapter_test.cc:
##########
@@ -456,6 +460,193 @@ TEST(TestAdapterRead,
ReadIntAndStringFileMultipleStripes) {
EXPECT_EQ(num_rows / reader_batch_size, batches);
}
+int64_t StripeBatchValue(int64_t row, int column) {
+ return static_cast<int64_t>(
+ (static_cast<uint64_t>(row + 1) * (0x123456789abcdefULL + column)) &
+ 0x7fffffffffffffffULL);
+}
+
+std::shared_ptr<Buffer> MakeStripeBatchFile(
+ int64_t rows_per_stripe, int stripes,
+ liborc::CompressionKind compression = liborc::CompressionKind_NONE) {
+ MemoryOutputStream stream(64 * 1024 * 1024);
+ auto type = liborc::Type::buildTypeFromString(
+ "struct<col0:bigint,col1:bigint,col2:bigint,col3:bigint>");
+ liborc::WriterOptions options;
+ options.setCompression(compression);
+ options.setRowIndexStride(0);
+ options.setStripeSize(stripes == 1 ? 256 * 1024 * 1024 : 1);
+ auto writer = liborc::createWriter(*type, &stream, options);
+ auto batch = writer->createRowBatch(16384);
+ auto& root = checked_cast<liborc::StructVectorBatch&>(*batch);
+ const int64_t total = rows_per_stripe * stripes;
+ for (int64_t offset = 0; offset < total; offset += 16384) {
+ root.numElements = std::min<int64_t>(16384, total - offset);
+ for (int col = 0; col < 4; ++col) {
+ auto& values = checked_cast<liborc::LongVectorBatch&>(*root.fields[col]);
+ values.numElements = root.numElements;
+ for (uint64_t row = 0; row < root.numElements; ++row) {
+ values.data[row] = StripeBatchValue(offset + row, col);
+ }
+ }
+ writer->add(root);
+ }
+ writer->close();
+ return Buffer::FromString(std::string(stream.getData(), stream.getLength()));
+}
+
+TEST(TestAdapterRead, StripeBatchSelectionAndState) {
+ auto file = std::make_shared<io::BufferReader>(MakeStripeBatchFile(16384,
4));
+ ASSERT_OK_AND_ASSIGN(auto reader,
+ adapters::orc::ORCFileReader::Open(file,
default_memory_pool()));
+ ASSERT_EQ(4, reader->NumberOfStripes());
+ ASSERT_OK(reader->Seek(16384 + 7));
+ ASSERT_OK_AND_ASSIGN(auto first, reader->GetStripeReader(2, 1000));
+ ASSERT_OK_AND_ASSIGN(auto second, reader->GetStripeReader(0, 1));
+ std::shared_ptr<RecordBatch> batch;
+ ASSERT_OK(second->ReadNext(&batch));
+ ASSERT_EQ(StripeBatchValue(0, 0),
+ checked_cast<const Int64Array&>(*batch->column(0)).Value(0));
+ ASSERT_OK_AND_ASSIGN(auto ignored, reader->ReadStripe(3));
+ int64_t rows = 0;
+ while (true) {
+ ASSERT_OK(first->ReadNext(&batch));
+ if (!batch) break;
+ ASSERT_OK(batch->ValidateFull());
+ ASSERT_LE(batch->num_rows(), 1000);
+ for (int64_t row = 0; row < batch->num_rows(); ++row) {
+ ASSERT_EQ(StripeBatchValue(2 * 16384 + rows + row, 0),
+ checked_cast<const Int64Array&>(*batch->column(0)).Value(row));
+ }
+ rows += batch->num_rows();
+ }
+ ASSERT_EQ(16384, rows);
+ ASSERT_OK_AND_ASSIGN(auto cursor_reader, reader->NextStripeReader(3));
+ ASSERT_OK(cursor_reader->ReadNext(&batch));
+ ASSERT_EQ(StripeBatchValue(16384 + 7, 0),
+ checked_cast<const Int64Array&>(*batch->column(0)).Value(0));
+ reader.reset();
+ ASSERT_OK(second->ReadNext(&batch));
+ ASSERT_EQ(StripeBatchValue(1, 0),
+ checked_cast<const Int64Array&>(*batch->column(0)).Value(0));
+}
+
+TEST(TestAdapterRead, StripeBatchProjectionAndArguments) {
+ auto file = std::make_shared<io::BufferReader>(MakeStripeBatchFile(16384,
3));
+ ASSERT_OK_AND_ASSIGN(auto reader,
+ adapters::orc::ORCFileReader::Open(file,
default_memory_pool()));
+ ASSERT_RAISES(Invalid, reader->GetStripeReader(-1, 10));
+ ASSERT_RAISES(Invalid, reader->GetStripeReader(3, 10));
+ ASSERT_RAISES(Invalid, reader->GetStripeReader(0, 0));
+ ASSERT_RAISES(Invalid, reader->GetStripeReader(0, -1));
+ ASSERT_RAISES(IOError, reader->GetStripeReader(0, 10, {"missing"}));
+ ASSERT_OK_AND_ASSIGN(auto projected,
+ reader->GetStripeReader(1, 10000, {"col2", "col0",
"col0"}));
+ ASSERT_EQ(projected->schema()->field_names(),
+ (std::vector<std::string>{"col0", "col2"}));
+ std::shared_ptr<RecordBatch> batch;
+ ASSERT_OK(projected->ReadNext(&batch));
+ ASSERT_EQ(10000, batch->num_rows());
+ ASSERT_EQ(StripeBatchValue(16384, 2),
+ checked_cast<const Int64Array&>(*batch->column(1)).Value(0));
+ ASSERT_OK_AND_ASSIGN(auto zero,
+ reader->GetStripeReader(2, 1000,
std::vector<std::string>{}));
+ int64_t rows = 0;
+ while (true) {
+ ASSERT_OK(zero->ReadNext(&batch));
+ if (!batch) break;
+ ASSERT_OK(batch->ValidateFull());
+ ASSERT_EQ(0, batch->num_columns());
+ ASSERT_GT(batch->num_rows(), 0);
+ ASSERT_LE(batch->num_rows(), 1000);
+ rows += batch->num_rows();
+ }
+ ASSERT_EQ(16384, rows);
+ ProxyMemoryPool zero_pool(default_memory_pool());
+ ASSERT_OK_AND_ASSIGN(auto zero_reader,
+ adapters::orc::ORCFileReader::Open(file, &zero_pool));
+ ASSERT_OK_AND_ASSIGN(auto zero_batches,
+ zero_reader->GetStripeReader(0, 1000,
std::vector<std::string>{}));
+ ASSERT_OK(zero_batches->ReadNext(&batch));
+ ASSERT_EQ(1000, batch->num_rows());
+ ASSERT_EQ(0, zero_pool.max_memory()); // No full-column decode followed by
projection.
+}
+
+TEST(TestAdapterRead, StripeBatchReadsOnlySelectedRange) {
+ for (auto codec : {liborc::CompressionKind_NONE,
liborc::CompressionKind_ZLIB}) {
+ auto file = std::make_shared<io::BufferReader>(MakeStripeBatchFile(16384,
4, codec));
+ std::shared_ptr<io::TrackedRandomAccessFile> tracked =
+ io::TrackedRandomAccessFile::Make(file.get());
+ ASSERT_OK_AND_ASSIGN(
+ auto reader, adapters::orc::ORCFileReader::Open(tracked,
default_memory_pool()));
+ ASSERT_EQ(4, reader->NumberOfStripes());
+ tracked->ResetStats();
+ for (int stripe : {2, 0}) {
+ ASSERT_OK_AND_ASSIGN(auto batches, reader->GetStripeReader(stripe, 17));
+ std::shared_ptr<RecordBatch> batch;
+ ASSERT_OK(batches->ReadNext(&batch));
+ ASSERT_EQ(17, batch->num_rows());
+ ASSERT_GT(tracked->bytes_read(), 0);
+ auto reads = tracked->num_reads();
+ batches.reset(); // Early stop must not read the rest or another stripe.
+ ASSERT_EQ(reads, tracked->num_reads());
+ }
+ for (const auto& range : tracked->get_read_ranges()) {
Review Comment:
ResetStats() is outside the loop and the final check only excludes stripes 1
and 3. As a result, reading stripe 0 while selecting stripe 2, or vice versa,
would still pass. Could we reset inside each iteration and compare every read
range against all stripes except the one currently selected?
##########
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:
**Disclaimer: suggested by AI**
`get_native_file()` accepts any object with `read()`, but this per-batch
check requires a `closed` attribute; `PythonFile::closed()` treats a missing
attribute as closed. Could we avoid making `.closed` a new prerequisite, or
document and test it explicitly?
##########
python/pyarrow/orc.py:
##########
@@ -165,6 +166,56 @@ def read_stripe(self, n, columns=None):
columns = self._select_names(columns)
return self.reader.read_stripe(n, columns=columns)
+ def iter_batches(self, batch_size=65536, stripes=None, columns=None):
+ """Read streaming batches from the file or selected stripes.
+
+ Parameters
+ ----------
+ batch_size : int, default 65536
+ Maximum number of rows per batch. Must be positive.
+ stripes : sequence of int, optional
+ Zero-based stripe indices, read in the supplied order. Repeated
+ indices repeat the corresponding rows. None reads the entire file;
+ an empty sequence produces no batches.
+ columns : list, optional
+ Column names, nested field paths, or top-level integer indices,
+ as in read_stripe. Output follows the file's column order. None
+ reads all columns; an empty list reads no columns but preserves
+ the number of rows.
+
+ Yields
+ ------
+ pyarrow.RecordBatch
+ A batch containing at most batch_size rows. Stripe boundaries may
+ produce smaller batches.
+
+ Notes
+ -----
+ Reading starts when the iterator is consumed. Call its close() method
+ to release the reader when stopping early. This does not close a source
+ supplied by the caller. Returned batches remain valid after closing.
+
+ Batches are decoded incrementally, without first constructing an Arrow
+ batch for the entire stripe. batch_size limits rows, not bytes or total
+ process memory: ORC dictionaries and decompression buffers may still
+ depend on the stripe size.
+ """
+ batch_size = operator.index(batch_size)
+ if batch_size > 2**63 - 1 or batch_size < -2**63:
+ raise OverflowError("batch_size does not fit in int64")
+ if batch_size <= 0:
+ raise ValueError("batch_size must be greater than zero")
+ if stripes is not None:
+ stripes = tuple(operator.index(stripe) for stripe in stripes)
+ for stripe in stripes:
+ if stripe > 2**63 - 1 or stripe < -2**63:
+ raise OverflowError("stripe index does not fit in int64")
+ if not 0 <= stripe < self.nstripes:
+ raise ValueError(f"Out of bounds stripe: {stripe}")
+ if columns is not None:
+ columns = self._select_names(list(columns))
Review Comment:
**Disclaimer: suggested by AI**
Mixed column selectors do not work: `_select_names()` returns the original
list as soon as it encounters a name, so `columns=[0, "other"]` reaches
`tobytes(0)` and fails. Please either normalize each entry independently or
explicitly reject mixed selectors, and add a mixed projection test.
--
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]
