Repository navigation
Conversation
…stripes Expose streaming batches from complete ORC files or selected stripes. Reuse native decoding through a stateless stripe reader factory and preserve row counts for zero-column projections. Add validation, lifecycle, range-reading and memory regression coverage. Generated-by: OpenAI Codex Signed-off-by: jiangxt2 <[email protected]>
|
|
github-actions
Bot
added
Component: C++
Component: Python
Component: Documentation
labels
Generate shared fixture data with Python integers and isolate NumPy integer arguments in a marked test. Cover whole-file and selected-stripe iterator lifetimes across all five codecs. Generated-by: OpenAI Codex Signed-off-by: jiangxt2 <[email protected]>
Merge upstream main at d24791f. Retain batch iteration and the removal of legacy ORC writer argument-order compatibility. Generated-by: OpenAI Codex Signed-off-by: jiangxt2 <[email protected]>
Remove a redundant column-list copy and reduce duplicate test combinations while preserving argument snapshots, codec coverage, and both iterator lifecycle paths. Generated-by: OpenAI Codex Signed-off-by: jiangxt2 <[email protected]>
| /// \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); |
There was a problem hiding this 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?
There was a problem hiding this comment.
Updated GetStripeReader() to a single API using std::optional<std::vector<std::string>>: std::nullopt selects all fields, an empty vector selects no fields while preserving row counts, and a nonempty vector selects the requested fields. The Cython binding preserves this distinction, and the C++ tests cover all three states.
| 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()) { |
There was a problem hiding this 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?
There was a problem hiding this comment.
ResetStats() now runs inside each iteration, and every read range is checked against all stripes except the selected one. This catches unintended reads between stripes 0 and 2 as well. The updated test passes for NONE and ZLIB, including the early-close assertion.
| 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)) |
There was a problem hiding this 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.
There was a problem hiding this comment.
Mixed column names and integer indices now raise an explicit TypeError at the iter_batches() call, before iteration starts. Added tests for both [0, "other"] and ["value", 1], and documented the restriction. Both cases pass with NumPy and in the isolated environment without NumPy.
| self.reader.get().GetStripeReader( | ||
| stripe, batch_size, names)) | ||
| while True: | ||
| if self.source_handle.get().closed(): |
There was a problem hiding this 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?
There was a problem hiding this 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). For Python-backed inputs, this reaches PythonFile::closed() before ORCFileReader::Open, so a missing .closed already causes opening to fail (source).
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.
Use optional column selection, strengthen per-stripe read-range assertions, and explicitly reject mixed Python column selectors. Generated-by: OpenAI Codex Signed-off-by: jiangxt2 <[email protected]>
Rationale for this change
ORCFile.read_stripe()materializes an entire stripe as one record batch. Dataset already scans ORC in batches; this adds native batching with explicit stripe selection toORCFile, supporting incremental consumption and early termination.Closes #54921.
What changes are included in this PR?
ORCFile.iter_batches(batch_size=65536, stripes=None, columns=None)with column projection.GetStripeReaderAPI with explicit optional column selection that decodes directly into batches without changing the reader's Seek position.Are these changes tested?
Validated locally on macOS arm64 with matching Release C++/Python builds, Python 3.13.7 and ORC 2.2.1.
44 passed; one existing timezone-database test restricted to ORC < 2.0 was skipped.
From
python/:140 passed with NumPy 2.5.4. In an isolated environment without NumPy, the
iter_batchestests andtest_without_numpy.pyhad 122 passes; only the marked NumPy integer-input test was skipped. Coverage includes five codecs, stripe order and duplicates, batch limits, nested/zero-column projections, invalid arguments, interleaved iterators and read failures. Both environments cover owner lifetime, early close and resource release for whole-file and selected-stripe iteration.The C++ memory regression uses four int64 columns and batches of 256 rows. For single stripes of 262144 and 1048576 rows, whole-stripe Arrow allocation peaks were 16.125 and 64.5 MiB; streaming peaked at 16.25 KiB for both sizes. NONE/ZLIB range tests verify that consuming the first batch and closing early does not read unselected stripe data.
Changed-file format/lint checks and
git diff --checkpassed.Are there any user-facing changes?
Yes. The new API yields
RecordBatchobjects.stripes=Nonereads all stripes; explicit sequences preserve order and duplicates, and empty sequences yield no batches.columns=[]preserves row counts. Closing the iterator releases its reader without closing the input; returned batches remain valid.batch_sizelimits output rows, not total process memory. ORC buffers and dictionaries require additional memory, and small batches can reduce throughput. Reading the same stripes does not inherently reduce I/O.Was AI used for this PR?
PR code and description written by:
Reviewed before submission by:
The human contributor directed the design and reviewed the proposal and PR text. OpenAI Codex assisted with code, tests, documentation and static review.