hhhizzz commented on code in PR #12019:
URL: https://github.com/apache/arrow-rs/pull/12019#discussion_r4533431349
##########
parquet/src/arrow/array_reader/builder.rs:
##########
@@ -476,7 +550,25 @@ impl<'a> ArrayReaderBuilder<'a> {
));
let page_iterator = self.row_groups.column_chunks(col_idx)?;
- let arrow_type = Some(field.arrow_type.clone());
+ let arrow_type = Some(match (&field.arrow_type, physical_type) {
+ (DataType::Utf8, PhysicalType::BYTE_ARRAY)
Review Comment:
Consolidated this into `stage_byte_array`. The size-statistics check is only
reached for top-level `BYTE_ARRAY` columns requested as `Utf8` or `Binary`; the
match then selects the staged or native type. Addressed in
[06918f486](https://github.com/hhhizzz/arrow-rs/commit/06918f486ec407f92139abb00657714f9a10803f).
##########
parquet/src/arrow/array_reader/builder.rs:
##########
@@ -112,6 +114,8 @@ struct ReaderArgs<'a> {
field: &'a ParquetField,
/// Which leaf columns are being read.
mask: &'a ProjectionMask,
+ /// Logical field depth: the root struct is zero, its direct fields are
one.
+ depth: usize,
Review Comment:
Added `ReaderArgs::is_top_level()` to name the direct-child-of-root check.
`with_field` still advances the depth when descending into a child, and the
staging guard uses the named helper. Addressed in
[06918f486](https://github.com/hhhizzz/arrow-rs/commit/06918f486ec407f92139abb00657714f9a10803f).
##########
parquet/src/arrow/array_reader/builder.rs:
##########
@@ -197,6 +218,59 @@ impl<'a> ArrayReaderBuilder<'a> {
Ok(reader)
}
+ /// Build once and collect the logical fields as projected root children
+ /// are accepted. Schema metadata is omitted, matching
+ /// `ParquetRecordBatchReader::new`.
+ pub(crate) fn build_array_reader_with_logical_schema(
+ &self,
+ field: Option<&ParquetField>,
+ mask: &ProjectionMask,
+ ) -> Result<(Box<dyn ArrayReader>, SchemaRef)> {
+ let mut logical = SchemaBuilder::new();
+ let reader = match field {
+ Some(field) => {
+ if !matches!(field.arrow_type, DataType::Struct(_)) {
+ return Err(general_err!("array reader root must be
Struct"));
+ }
+ self.build_struct_reader(
+ ReaderArgs {
+ field,
+ mask,
+ depth: 0,
+ padding_threshold: None,
+ },
+ Some(&mut logical),
+ )?
+ }
+ None => None,
+ }
+ .unwrap_or_else(|| make_empty_array_reader(self.num_rows()));
+ Ok((reader, Arc::new(logical.finish())))
+ }
+
+ /// A bound for the complete scan, including repeated row groups. Page
+ /// pruning and predicate-cache roles must not change a column's type.
+ fn byte_array_fits(&self, column: usize, capacity: usize) -> bool {
+ let mut has_metadata = false;
+ let total = self
+ .row_groups
+ .row_groups()
+ .try_fold(0_i64, |total, group| {
+ has_metadata = true;
+ let bytes = group
+ .columns()
+ .get(column)?
+ .unencoded_byte_array_data_bytes()?;
+ if bytes < 0 {
+ return None;
+ }
+ total
+ .checked_add(bytes)
+ .filter(|&total| total <= capacity as i64)
+ });
+ has_metadata && total.is_some()
Review Comment:
Renamed the flag to `has_row_groups` and documented that an empty collection
keeps the conservative staging default, even though there are no rows to
decode. The existing behavior is unchanged. Addressed in
[06918f486](https://github.com/hhhizzz/arrow-rs/commit/06918f486ec407f92139abb00657714f9a10803f).
##########
parquet/src/arrow/arrow_reader/capacity_staging.rs:
##########
@@ -0,0 +1,247 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Restore top-level 32-bit byte arrays after decoding with 64-bit offsets.
+//!
+//! A fitting window shares its payload buffers. Oversized windows are split
+//! at complete rows. Each staged column shares slices using more than half its
+//! visible window payload, provided the reported backing capacity is at most
+//! twice that payload. Other slices are copied. Sharing avoids large copies
+//! but can keep the entire allocation alive. Other columns and null
+//! buffers may also share storage; this is not a reader memory bound.
+//! Nested columns retain their requested representation.
+
+use crate::errors::{ParquetError, Result};
+use arrow_array::cast::AsArray;
+use arrow_array::types::{ByteArrayType, GenericBinaryType, GenericStringType};
+use arrow_array::{Array, ArrayRef, GenericByteArray, RecordBatch,
RecordBatchOptions};
+use arrow_buffer::{Buffer, OffsetBuffer, ScalarBuffer};
+use arrow_schema::{DataType, SchemaRef};
+use std::sync::Arc;
+
+/// Restores staged top-level columns and splits their payloads at `max_bytes`.
+pub(crate) struct CapacitySplitter {
+ schema: SchemaRef,
+ max_bytes: i64,
+ pending: Option<PendingBatch>,
+}
+
+struct PendingBatch {
+ batch: RecordBatch,
+ offset: usize,
+ split: bool,
+}
+
+impl CapacitySplitter {
+ pub(crate) fn new(schema: SchemaRef, max_bytes: usize) -> Self {
+ Self {
+ schema,
+ max_bytes: max_bytes.min(i32::MAX as usize) as i64,
+ pending: None,
+ }
+ }
+
+ /// Accept a decoded window after the previous window has been drained.
+ pub(crate) fn push(&mut self, batch: RecordBatch) -> Result<()> {
+ if self.pending.is_some() {
+ return Err(general_err!(
+ "Internal Error: capacity splitter must be drained before push"
+ ));
+ }
+ if batch.num_columns() != self.schema.fields().len() {
+ return Err(general_err!(
+ "Internal Error: staged batch has {} columns, expected {}",
+ batch.num_columns(),
+ self.schema.fields().len()
+ ));
+ }
+ let mut split = false;
+ for ((array, from), to) in batch
+ .columns()
+ .iter()
+ .zip(batch.schema().fields())
+ .zip(self.schema.fields())
+ {
+ let offsets = staged_offsets(array.as_ref(), to.data_type());
+ if from.name() != to.name()
+ || (array.data_type() != to.data_type() && offsets.is_none())
+ {
+ return Err(general_err!(
+ "Internal Error: cannot restore staged field {from} as
{to}"
+ ));
+ }
+ if let Some(offsets) = offsets {
+ split |= offsets[offsets.len() - 1] - offsets[0] >
self.max_bytes;
+ }
+ }
+ self.pending = Some(PendingBatch {
+ batch,
+ offset: 0,
+ split,
+ });
+ Ok(())
+ }
+
+ /// Emit the next batch. An error releases the pending window.
+ pub(crate) fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
+ let Some(mut pending) = self.pending.take() else {
+ return Ok(None);
+ };
+ let mut end = pending.batch.num_rows();
+ if pending.split {
+ for (array, field) in
pending.batch.columns().iter().zip(self.schema.fields()) {
+ let Some(offsets) = staged_offsets(array.as_ref(),
field.data_type()) else {
+ continue;
+ };
+ let start = offsets[pending.offset];
+ let limit = start.saturating_add(self.max_bytes);
+ let count = offsets[pending.offset + 1..].partition_point(|&o|
o <= limit);
+ if count == 0 {
+ return Err(general_err!(
+ "a single row needs {} bytes of {} data for column {},
which exceeds \
+ the {} byte limit of one {} array; read this column
as a Large or \
+ View type instead",
+ offsets[pending.offset + 1] - start,
+ field.data_type(),
+ field.name(),
+ self.max_bytes,
+ field.data_type()
+ ));
+ }
+ end = end.min(pending.offset + count);
+ }
+ }
+ let sliced;
+ let slice = if pending.offset == 0 && end == pending.batch.num_rows() {
+ &pending.batch
+ } else {
+ sliced = pending.batch.slice(pending.offset, end - pending.offset);
+ &sliced
+ };
+ let columns = slice
+ .columns()
+ .iter()
+ .zip(pending.batch.columns())
+ .zip(self.schema.fields())
+ .map(
+ |((array, window), field)| match (array.data_type(),
field.data_type()) {
+ (DataType::LargeUtf8, DataType::Utf8) => {
+ narrow::<GenericStringType<i64>,
GenericStringType<i32>>(
+ array,
+ pending.split.then_some(window),
+ )
+ }
+ (DataType::LargeBinary, DataType::Binary) => {
+ narrow::<GenericBinaryType<i64>,
GenericBinaryType<i32>>(
+ array,
+ pending.split.then_some(window),
+ )
+ }
+ _ => Ok(array.clone()),
+ },
+ )
+ .collect::<Result<Vec<_>>>()?;
+ let output = RecordBatch::try_new_with_options(
+ self.schema.clone(),
+ columns,
+ &RecordBatchOptions::new().with_row_count(Some(slice.num_rows())),
+ )?;
+ pending.offset = end;
+ if end < pending.batch.num_rows() {
+ self.pending = Some(pending);
+ }
+ Ok(Some(output))
+ }
+}
+
+/// Whether a reader root differs only by optional top-level byte-array
staging.
+pub(crate) fn can_restore_type(decoded: &DataType, logical: &DataType) -> bool
{
+ let (DataType::Struct(decoded), DataType::Struct(logical)) = (decoded,
logical) else {
+ return decoded == logical;
+ };
+ decoded.len() == logical.len()
+ && decoded.iter().zip(logical).all(|(from, to)| {
+ from.name() == to.name()
+ && from.is_nullable() == to.is_nullable()
+ && from.metadata() == to.metadata()
+ && (from.data_type() == to.data_type()
+ || matches!(
+ (from.data_type(), to.data_type()),
+ (DataType::LargeUtf8, DataType::Utf8)
+ | (DataType::LargeBinary, DataType::Binary)
+ ))
+ })
+}
+
+fn staged_offsets<'a>(array: &'a dyn Array, target: &DataType) -> Option<&'a
[i64]> {
+ match (array.data_type(), target) {
+ (DataType::LargeUtf8, DataType::Utf8) =>
Some(array.as_string::<i64>().value_offsets()),
+ (DataType::LargeBinary, DataType::Binary) =>
Some(array.as_binary::<i64>().value_offsets()),
+ _ => None,
+ }
+}
+
+/// Rebase and narrow only the visible slice's offsets, preserving its
validity.
+fn narrow<F, T>(array: &ArrayRef, split_window: Option<&ArrayRef>) ->
Result<ArrayRef>
+where
+ F: ByteArrayType<Offset = i64>,
+ T: ByteArrayType<Offset = i32, Native = F::Native>,
+{
+ let source = array.as_bytes::<F>();
+ let offsets = source.value_offsets();
+ let first = offsets[0];
+ let span = offsets[offsets.len() - 1] - first;
+ if span > i32::MAX as i64 {
+ return Err(general_err!(
+ "Internal Error: staged slice of {span} bytes exceeds 32-bit
offsets"
+ ));
+ }
+ let narrowed = offsets
+ .iter()
+ .map(|o| (*o - first) as i32)
+ .collect::<Vec<_>>();
+ let (first, span) = (first as usize, span as usize);
+ // Use the original column window's visible span to identify a dominant
+ // slice, allowing up to 2x capacity for decoder Vec growth. Small slices
+ // and windows backed by an oversized allocation still get compact copies.
+ // For split windows, sharing retains less than 4x the slice's payload in
+ // reported capacity; external owners may retain unaccounted storage.
+ let copy_values = split_window.is_some_and(|window| {
+ let offsets = window.as_bytes::<F>().value_offsets();
+ let window_span = (offsets[offsets.len() - 1] - offsets[0]) as usize;
+ let capacity = source.values().capacity();
+ capacity == 0 || span <= window_span / 2 || capacity >
window_span.saturating_mul(2)
Review Comment:
Extracted `should_copy_payload` and documented the three cases: zero
reported backing capacity, a slice containing at most half the visible window
payload, and backing capacity above twice the window. The thresholds are
unchanged. Addressed in
[06918f486](https://github.com/hhhizzz/arrow-rs/commit/06918f486ec407f92139abb00657714f9a10803f).
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -259,6 +259,11 @@ pub struct ArrowReaderBuilder<T> {
pub(crate) metrics: ArrowReaderMetrics,
pub(crate) max_predicate_cache_size: usize,
+
+ /// Byte limit of each `Utf8`/`Binary` array in an output batch, see
+ /// [`capacity_staging`]. Always `i32::MAX` outside of tests; `None`
+ /// disables staging so tests can compare against the unstaged readers.
+ pub(crate) byte_array_capacity: Option<usize>,
Review Comment:
Added the public `with_byte_array_capacity_staging(bool)` option to
`ArrowReaderBuilder`, enabled by default. The setting carries through sync,
async, push and predicate/cache reads. The only clamp now lives in the numeric
test override; the internal entry points assert the bound. Tests cover
disabling staging and retaining the setting when rebuilding a push decoder.
Addressed in
[06918f486](https://github.com/hhhizzz/arrow-rs/commit/06918f486ec407f92139abb00657714f9a10803f).
--
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]
