hhhizzz commented on code in PR #12019:
URL: https://github.com/apache/arrow-rs/pull/12019#discussion_r4533494665
##########
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:
`try_new_with_row_groups` still hard-codes `Some(i32::MAX as usize)`, so
callers of this public entry point cannot disable staging. The docs for
`with_byte_array_capacity_staging` describe the setting as applying across the
read paths; this is an exception and should at least be documented here.
##########
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 schema passed to `new_with_capacity_staging` is now built from
`levels.fields` rather than the one collected while building the reader. If the
two ever differ for a reason other than staging, the reader becomes `Active`
and fails on the first batch with "cannot restore staged field". I have not
confirmed that such a mismatch is reachable, but the only guard is the
`debug_assert_eq!` above, which does nothing in release builds. Passing the
collected `schema` here, or returning an error on mismatch, would make this
robust.
##########
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:
Naming: `with_byte_array_capacity_staging` exposes the internal term
"capacity staging" in the public API. Reviewers may prefer a name that
describes the user-visible effect (splitting oversized `Utf8`/`Binary`
batches), or placing the option on `ArrowReaderOptions`. Style only.
--
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]
