hhhizzz commented on code in PR #12019:
URL: https://github.com/apache/arrow-rs/pull/12019#discussion_r4524686236
##########
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:
The `Utf8` and `Binary` arms repeat the same guard. Computing it once reads
more easily, e.g. `let stage = args.depth == 1 && physical_type ==
PhysicalType::BYTE_ARRAY && self.byte_array_capacity.is_some_and(|c|
!self.byte_array_fits(col_idx, c));` followed by a plain type mapping.
##########
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:
`depth == 1` as the meaning of "top-level column" depends on every recursion
going through `with_field`. A named helper such as `is_top_level()` would make
the intent explicit and is harder to break when a new recursive path is added.
##########
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:
`has_metadata` makes an empty row-group list count as "does not fit",
although zero bytes trivially fit. The tests pin this behaviour but nothing
says why. A one-line comment on the intent would help.
##########
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:
This field exists for tests but is threaded through `ArrowReaderBuilder`,
`RowGroupReaderBuilder`, its `Parts`, `PredicateOptions` and
`ArrayReaderBuilder`, and the clamp to `i32::MAX` is repeated in both setters
and in `CapacitySplitter::new`. If it became a real option (see the review
summary) the plumbing would be justified; otherwise clamping in one place would
be enough.
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -1833,20 +1907,28 @@ 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 {
+ let mut reader = Self::new_with_capacity_staging(
array_reader,
- schema: Arc::new(Schema::new(levels.fields.clone())),
read_plan,
- })
+ schema,
+ Some(i32::MAX as usize),
+ );
+ reader.schema = Arc::new(Schema::new(levels.fields.clone()));
Review Comment:
Overwriting `reader.schema` after construction is awkward to read. Passing
the schema to keep into the constructor, or building the struct directly here,
would avoid the two-step setup.
##########
parquet/tests/arrow_reader/large_string_overflow.rs:
##########
@@ -15,120 +15,142 @@
// specific language governing permissions and limitations
// under the License.
+//! Reader regressions for a logical byte-array window above i32::MAX.
+//! Build legal input batches, stream each output, and never concatenate them.
+//! Run these memory-intensive tests explicitly with one test thread.
+
use std::sync::Arc;
-use arrow_array::builder::BinaryBuilder;
-use arrow_array::{ArrayRef, RecordBatch};
+use arrow_array::{Array, ArrayRef, BinaryArray, Int64Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
-
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use parquet::arrow::arrow_writer::ArrowWriter;
use parquet::basic::Encoding;
use parquet::file::properties::WriterProperties;
-
+use parquet::schema::types::ColumnPath;
use tempfile::tempfile;
-/// Number of rows written
-const ROWS: usize = 1024;
-
-/// Size of each binary value (~3MB)
-/// 1024 * 3MB ≈ 3GB total → guaranteed offset overflow for i32
+const ROWS: usize = 700;
+const WRITE_ROWS: usize = 10;
const VALUE_SIZE: usize = 3 * 1024 * 1024;
-fn make_large_binary_array() -> ArrayRef {
- let mut builder = BinaryBuilder::new();
-
- for _ in 0..ROWS {
- let data = vec![b'a'; VALUE_SIZE];
- builder.append_value(&data);
- }
-
- Arc::new(builder.finish()) as ArrayRef
-}
-
-fn write_parquet_with_encoding(array: ArrayRef, encoding: Encoding) ->
std::fs::File {
- let schema = Arc::new(Schema::new(vec![Field::new(
- "col",
- DataType::Binary,
- false,
- )]));
-
- let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
-
- let file = tempfile().unwrap();
-
- let builder = WriterProperties::builder();
- let builder = match encoding {
- Encoding::RLE_DICTIONARY => builder.set_dictionary_enabled(true),
- _ => builder.set_dictionary_enabled(false).set_encoding(encoding),
+fn roundtrip_large_binary(encoding: Encoding) {
+ assert!(ROWS * VALUE_SIZE > i32::MAX as usize);
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("col", DataType::Binary, false),
+ Field::new("row_id", DataType::Int64, false),
+ ]));
+ let value = vec![b'a'; VALUE_SIZE];
+ // Reuse a legal 30 MiB array: an oversized input builder would panic
+ // before reaching the Parquet reader and would not test this regression.
+ let values: ArrayRef =
Arc::new(BinaryArray::from_iter_values(std::iter::repeat_n(
+ value.as_slice(),
+ WRITE_ROWS,
+ )));
+ let column = ColumnPath::from("col");
+ let properties = WriterProperties::builder()
+ .set_dictionary_enabled(false)
+ .set_dictionary_page_size_limit(8 * 1024 * 1024)
+ .set_max_row_group_row_count(Some(ROWS))
+ .set_write_batch_size(WRITE_ROWS)
+ .set_data_page_row_count_limit(WRITE_ROWS);
+ let properties = if encoding == Encoding::RLE_DICTIONARY {
+ properties.set_column_dictionary_enabled(column, true)
+ } else {
+ properties.set_column_encoding(column, encoding)
};
-
- let props = builder.build();
-
- let mut writer = ArrowWriter::try_new(file.try_clone().unwrap(), schema,
Some(props)).unwrap();
-
- writer.write(&batch).unwrap();
+ // Compression keeps the fixture small; decoding must still materialize
+ // the complete >2 GiB logical window.
+ #[cfg(feature = "zstd")]
+ let properties =
+
properties.set_compression(parquet::basic::Compression::ZSTD(Default::default()));
+ let file = tempfile().unwrap();
+ let mut writer = ArrowWriter::try_new(
+ file.try_clone().unwrap(),
+ schema.clone(),
+ Some(properties.build()),
+ )
+ .unwrap();
+ for start in (0..ROWS).step_by(WRITE_ROWS) {
+ let ids: ArrayRef = Arc::new(Int64Array::from_iter_values(
+ start as i64..(start + WRITE_ROWS) as i64,
+ ));
+ writer
+ .write(&RecordBatch::try_new(schema.clone(), vec![values.clone(),
ids]).unwrap())
+ .unwrap();
+ }
writer.close().unwrap();
-
- file
+ drop(values);
+
+ let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+ assert_eq!(builder.metadata().num_row_groups(), 1);
+ assert!(
+ builder
+ .metadata()
+ .row_group(0)
+ .column(0)
+ .encodings()
+ .any(|actual| actual == encoding),
+ "fixture must exercise {encoding:?}"
+ );
+ let mut rows = 0;
+ let mut output_rows = vec![];
+ for batch in builder.with_batch_size(ROWS).build().unwrap() {
+ let batch = batch.expect("the reader must split a wide logical
window");
+ assert_eq!(batch.schema(), schema);
+ let values = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<BinaryArray>()
+ .unwrap();
+ let ids = batch
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ assert!(i32::try_from(values.values().len()).is_ok());
+ for index in 0..batch.num_rows() {
+ assert_eq!(ids.value(index), (rows + index) as i64);
+ assert!(
+ values.value(index) == value.as_slice(),
+ "{encoding:?}: payload differs at row {}",
+ rows + index
+ );
+ }
+ output_rows.push(batch.num_rows());
+ rows += batch.num_rows();
+ // Drop this output before requesting the next. The pending window
+ // and a near-2 GiB output still coexist; this remains a heavy test.
+ }
+ assert_eq!(rows, ROWS);
+ let first_rows = i32::MAX as usize / VALUE_SIZE;
+ assert_eq!(output_rows, vec![first_rows, ROWS - first_rows]);
+ println!(
+ "{encoding:?}: streamed {rows} rows / {} bytes as {output_rows:?}",
+ ROWS * VALUE_SIZE
+ );
}
#[test]
-// Panics until https://github.com/apache/arrow-rs/issues/7973 is fixed
-#[should_panic(expected = "byte array offset overflow")]
+#[ignore = "requires several GiB of peak memory; run --ignored
--test-threads=1"]
Review Comment:
With these four tests ignored, CI no longer has an end-to-end read above 2
GiB; coverage relies on the small-capacity tests. That trade-off seems right
given the memory cost, but it is worth stating explicitly in the description.
##########
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:
This is the rule in the PR that needs the most explanation, and it is three
conditions on one line inside a closure. Consider a named function (e.g.
`should_copy_payload(span, window_span, capacity)`) with one comment per
condition: unknown capacity, minority slice, over-reserved backing buffer.
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -1569,13 +1603,22 @@ impl<T: ChunkReader + 'static> PageIterator for
ReaderPageIterator<T> {}
///
/// Typically, either reads from a file or an in memory buffer [`Bytes`]
///
+/// When byte-array capacity staging is needed, this reader ends after
returning
+/// an error, releasing its pending window instead of continuing past rows that
+/// could not be returned. This matches the push decoder's terminal error
behavior.
+///
/// Created by [`ParquetRecordBatchReaderBuilder`]
///
/// [`Bytes`]: bytes::Bytes
pub struct ParquetRecordBatchReader {
array_reader: Box<dyn ArrayReader>,
schema: SchemaRef,
read_plan: ReadPlan,
+ /// Restores staged top-level `Utf8`/`Binary` columns, splitting oversized
windows.
+ /// `None` when the projection has no staged columns.
+ splitter: Option<capacity_staging::CapacitySplitter>,
+ /// Set after an error on the staged path, which ends the stream.
+ splitter_failed: bool,
Review Comment:
`splitter: Option<_>` plus `splitter_failed: bool` encode three states (not
staged / active / failed). A small enum would state that directly and remove
the two `unwrap()` calls in `next_restored`.
--
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]
