hhhizzz commented on code in PR #12019:
URL: https://github.com/apache/arrow-rs/pull/12019#discussion_r4533649642
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -1833,20 +1935,26 @@ impl ParquetRecordBatchReader {
) -> Result<Self> {
// note metrics are not supported in this API
let metrics = ArrowReaderMetrics::disabled();
- let array_reader = ArrayReaderBuilder::new(row_groups, &metrics)
+ let (array_reader, schema) = ArrayReaderBuilder::new(row_groups,
&metrics)
.with_batch_size(batch_size)
.with_parquet_metadata(row_groups.metadata())
- .build_array_reader(levels.levels.as_ref(),
&ProjectionMask::all())?;
+ .with_byte_array_capacity(Some(i32::MAX as usize))
Review Comment:
The docs for both entry points now explain the scope.
`try_new_with_row_groups` always enables byte-array batch splitting, and the
builder option configures sync, async, and push readers created through the
builder.
[904457618](https://github.com/hhhizzz/arrow-rs/commit/904457618fd4976570cc2e4f7b735373c15a2662)
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -1833,20 +1935,26 @@ impl ParquetRecordBatchReader {
) -> Result<Self> {
// note metrics are not supported in this API
let metrics = ArrowReaderMetrics::disabled();
- let array_reader = ArrayReaderBuilder::new(row_groups, &metrics)
+ let (array_reader, schema) = ArrayReaderBuilder::new(row_groups,
&metrics)
.with_batch_size(batch_size)
.with_parquet_metadata(row_groups.metadata())
- .build_array_reader(levels.levels.as_ref(),
&ProjectionMask::all())?;
+ .with_byte_array_capacity(Some(i32::MAX as usize))
+ .build_array_reader_with_logical_schema(
+ levels.levels.as_ref(),
+ &ProjectionMask::all(),
+ )?;
+ debug_assert_eq!(schema.fields(), &levels.fields);
let read_plan = ReadPlanBuilder::new(batch_size)
.with_selection(selection)
.build();
- Ok(Self {
+ Ok(Self::new_with_capacity_staging(
array_reader,
- schema: Arc::new(Schema::new(levels.fields.clone())),
read_plan,
- })
+ Arc::new(Schema::new(levels.fields.clone())),
Review Comment:
The low-level constructor now uses the logical schema collected while
building its column readers. A regression test supplies a mismatched internal
`FieldLevels.fields` list and checks that the reader reports the collected
schema and reads the original rows.
[904457618](https://github.com/hhhizzz/arrow-rs/commit/904457618fd4976570cc2e4f7b735373c15a2662)
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -323,12 +333,42 @@ impl<T> ArrowReaderBuilder<T> {
/// guarantee exact internal buffer capacities.
///
/// If `batch_size` is more than the file row count, use the file row
count.
+ ///
+ /// When capacity staging is enabled, a batch can have fewer rows when a
+ /// top-level `Utf8`/`Binary` column would contain more than `i32::MAX`
data bytes.
+ /// The remaining rows are returned in later batches, each containing
+ /// complete top-level rows. Requesting `LargeUtf8`/`LargeBinary` or
+ /// `Utf8View`/`BinaryView` for the affected column avoids this byte limit.
+ /// Nested byte arrays retain their existing capacity limits.
pub fn with_batch_size(self, batch_size: usize) -> Self {
// Try to avoid allocate large buffer
let batch_size =
batch_size.min(self.metadata.file_metadata().num_rows() as usize);
Self { batch_size, ..self }
}
+ /// Enable or disable capacity staging for top-level `Utf8`/`Binary`
columns.
+ ///
+ /// Enabled by default, staging decodes eligible columns with 64-bit
offsets
+ /// and splits output batches when needed to fit 32-bit offsets. It is
+ /// skipped when size statistics prove the scanned data fits. Disabling it
+ /// can avoid staging overhead, but a decoded batch that exceeds `i32::MAX`
+ /// bytes for a column can return an offset overflow error, including when
+ /// unselected rows are decoded before filtering. The setting also applies
+ /// to predicate evaluation, the predicate cache, and the async and push
+ /// readers. It does not impose a peak-memory limit.
+ pub fn with_byte_array_capacity_staging(mut self, enabled: bool) -> Self {
Review Comment:
Renamed the new option to `with_byte_array_batch_splitting` to describe the
output behavior. It is on `ArrowReaderBuilder` alongside `with_batch_size`,
where callers configure reading, and the docs and tests use the new name.
[904457618](https://github.com/hhhizzz/arrow-rs/commit/904457618fd4976570cc2e4f7b735373c15a2662)
--
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]
