hhhizzz commented on code in PR #12019:
URL: https://github.com/apache/arrow-rs/pull/12019#discussion_r4533432831
##########
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:
Replaced the two fields with `CapacityStaging::{Disabled, Active, Failed}`
and pass the splitter directly to `next_restored`, removing the unwraps. Staged
errors release the pending window and end the reader; the sync, async and push
terminal-error tests pass. Addressed in
[06918f486](https://github.com/hhhizzz/arrow-rs/commit/06918f486ec407f92139abb00657714f9a10803f).
##########
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:
`try_new_with_row_groups` now passes the logical schema directly to the
constructor so the reader and splitter both use it from construction. The
unstaged branch also uses the supplied schema directly, avoiding a temporary
schema allocation. Addressed in
[06918f486](https://github.com/hhhizzz/arrow-rs/commit/06918f486ec407f92139abb00657714f9a10803f).
##########
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:
The PR description now states that ordinary CI uses the small-limit tests
and skips the four real >2 GiB cases. All four were run separately on this
reader implementation, verified the payload and schema, and returned batches of
`[682, 18]` rows.
--
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]
