raulcd commented on code in PR #54926:
URL: https://github.com/apache/arrow/pull/54926#discussion_r4441902970
##########
cpp/src/arrow/adapters/orc/adapter_test.cc:
##########
@@ -18,6 +18,8 @@
#include <gmock/gmock-matchers.h>
#include <gtest/gtest.h>
+#include <iostream>
Review Comment:
to be removed if we remove `std::cout`
##########
python/pyarrow/orc.py:
##########
@@ -165,6 +166,61 @@ 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. Names and integer indices cannot be mixed.
+ 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)
Review Comment:
this code seems to be overly defensive AI writing code. From the
`operator.index` to the validation of `OverflowErrors`, plus validation of the
columns `Integral`. I am not saying this is wrong but makes reading the code
slightly harder than it is strictly required as both the language and the
underlying layers (cython/C++) already test that. As an example on parquet
`iter_batches` where we do a similar work with batch_size and row_groups that
code is much more readable and I'd argue that those checks are done too but
either on C++ or Cython:
https://github.com/apache/arrow/blob/bf0c7f06cc453196e99bffa2fccce2c24aae78bc/python/pyarrow/parquet/core.py#L595-L606
@rok @AlenkaF what do you think?
##########
cpp/src/arrow/adapters/orc/adapter_test.cc:
##########
@@ -456,6 +461,196 @@ 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,
std::nullopt));
+ 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,
std::vector<std::string>{"missing"}));
+ ASSERT_OK_AND_ASSIGN(auto projected,
+ reader->GetStripeReader(
+ 1, 10000, std::vector<std::string>{"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());
+ for (int stripe : {2, 0}) {
+ tracked->ResetStats();
+ 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()) {
+ for (int other = 0; other < reader->NumberOfStripes(); ++other) {
+ if (other == stripe) continue;
+ auto unselected = reader->GetStripeInformation(other);
+ ASSERT_TRUE(range.offset + range.length <= unselected.offset ||
+ range.offset >= unselected.offset + unselected.length);
+ }
+ }
+ }
+ }
+}
+
+TEST(TestAdapterRead, StripeBatchMemoryBounded) {
+ constexpr int64_t batch_size = 256;
+ int64_t small_peak = 0;
+ for (int64_t rows : {262144, 1048576}) {
+ auto buffer = MakeStripeBatchFile(rows, 1);
+ int64_t full_peak;
+ {
+ ProxyMemoryPool pool(default_memory_pool());
+ auto file = std::make_shared<io::BufferReader>(buffer);
+ ASSERT_OK_AND_ASSIGN(auto reader,
adapters::orc::ORCFileReader::Open(file, &pool));
+ ASSERT_EQ(1, reader->NumberOfStripes());
+ ASSERT_EQ(rows, reader->GetStripeInformation(0).num_rows);
+ ASSERT_OK_AND_ASSIGN(auto batch, reader->ReadStripe(0));
+ ASSERT_EQ(rows, batch->num_rows());
+ full_peak = pool.max_memory();
+ }
+ ProxyMemoryPool pool(default_memory_pool());
+ auto file = std::make_shared<io::BufferReader>(buffer);
+ ASSERT_OK_AND_ASSIGN(auto reader, adapters::orc::ORCFileReader::Open(file,
&pool));
+ ASSERT_EQ(0, pool.bytes_allocated());
+ ASSERT_OK_AND_ASSIGN(auto batches, reader->GetStripeReader(0, batch_size));
+ ASSERT_EQ(0, pool.bytes_allocated()); // No whole-stripe Arrow allocation.
+ int64_t count = 0;
+ int64_t first_peak = 0;
+ while (true) {
+ std::shared_ptr<RecordBatch> batch;
+ ASSERT_OK(batches->ReadNext(&batch));
+ if (!batch) break;
+ ASSERT_LE(batch->num_rows(), batch_size);
+ if (count == 0) first_peak = pool.max_memory();
+ ASSERT_EQ(StripeBatchValue(count, 0),
+ checked_cast<const Int64Array&>(*batch->column(0)).Value(0));
+ count += batch->num_rows();
+ }
+ ASSERT_EQ(rows, count);
+ ASSERT_EQ(0, pool.bytes_allocated());
+ ASSERT_LT(first_peak, full_peak / 4);
+ ASSERT_LT(pool.max_memory(), full_peak / 4);
+ if (small_peak == 0) small_peak = pool.max_memory();
+ ASSERT_LE(pool.max_memory(), 2 * small_peak);
+ std::cout << "ORC_BATCH_MEMORY rows=" << rows << " batch_size=" <<
batch_size
Review Comment:
In general we don't `std::cout` on tests, is this relevant enough to keep it?
##########
python/pyarrow/orc.py:
##########
@@ -165,6 +166,61 @@ 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. Names and integer indices cannot be mixed.
+ 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 = list(columns)
+ if any(isinstance(col, Integral) for col in columns) and any(
+ not isinstance(col, Integral) for col in columns):
+ raise TypeError("columns cannot mix names and integer indices")
Review Comment:
Initial review: on `read_stripe` and `read` we don't do this distinction and
checks. We probably want to match existing APIs, is there a reason to allow a
different type of `columns` here?
After further revision this seems to come from a previous comment but the
bug seems to come from `self._select_names` which will also affect `read` and
`read_stripe`. I'd drop it from there and fix it there even though we can argue
this is an existing issue and we can fix it on a follow up PR.
##########
python/pyarrow/_orc.pyx:
##########
@@ -364,6 +366,54 @@ 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
+ optional[std_vector[c_string]] selected_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]
+ selected_names = optional[std_vector[c_string]](names)
+
+ # Keep self (and its source) alive for the generator's entire lifetime.
Review Comment:
I am not sure I understand this comment is this referring to
`stripe_indices`? I'd argue if we want to add a comment could be about the
`None` sentinel value indicating reading the full file.
##########
python/pyarrow/_orc.pyx:
##########
@@ -364,6 +366,54 @@ 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
+ optional[std_vector[c_string]] selected_names
+ int64_t stripe
+
+ if batch_size <= 0:
+ raise ValueError("batch_size must be greater than zero")
Review Comment:
we already check this on the Python side, feels slightly redundant.
--
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]
