This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11977-0b347943c541f4d51f61b182c4b3d6621d6e8dcb in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit c1ce2b9cdccf50ccd826108027b59892bccb426e Author: don <[email protected]> AuthorDate: Thu Oct 29 16:22:52 2026 +0000 test(parquet): move offset-index async reader tests into their own module (#11977) > [!IMPORTANT] > > CI is enabled on my fork. # Which issue does this PR close? - Part of #9348 # Rationale for this change #11063 moved the public-API async reader tests into one integration file. #9348 was reopened so those tests can be split into topical modules, following `parquet/tests/arrow_reader/row_filter/`. This PR establishes `parquet/tests/arrow_reader/async_reader/` and moves the offset-index tests. The remaining groups stay in `async_reader/mod.rs` for follow-up PRs. # What changes are included in this PR? - Rename `parquet/tests/arrow_reader/async_reader.rs` to `parquet/tests/arrow_reader/async_reader/mod.rs`. - Move these tests into `async_reader/offset_index.rs`: - `non_empty_offset_index_doesnt_panic_in_read_row_group` - `empty_offset_index_doesnt_panic_in_column_chunks` - `test_cached_array_reader_sparse_offset_error` No production code, assertions, or public APIs change. Tests that use private helpers stay in `parquet/src/arrow/async_reader/mod.rs`. # Are these changes tested? `cargo test -p parquet --features async --test arrow_reader -- offset_index` # Are there any user-facing changes? No. # AI usage disclosure This pull request was prepared with an AI assistant in Cursor. --- .../{async_reader.rs => async_reader/mod.rs} | 123 +----------------- .../arrow_reader/async_reader/offset_index.rs | 143 +++++++++++++++++++++ 2 files changed, 145 insertions(+), 121 deletions(-) diff --git a/parquet/tests/arrow_reader/async_reader.rs b/parquet/tests/arrow_reader/async_reader/mod.rs similarity index 85% rename from parquet/tests/arrow_reader/async_reader.rs rename to parquet/tests/arrow_reader/async_reader/mod.rs index f65c1732f9..fc476de6c2 100644 --- a/parquet/tests/arrow_reader/async_reader.rs +++ b/parquet/tests/arrow_reader/async_reader/mod.rs @@ -37,12 +37,13 @@ use parquet::arrow::{ ArrowWriter, AsyncArrowWriter, ParquetRecordBatchStreamBuilder, ProjectionMask, }; use parquet::errors::Result; -use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader}; +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; use parquet::file::properties::WriterProperties; use rand::{RngExt, rng}; use std::collections::HashMap; use std::sync::Arc; use tempfile::tempfile; +mod offset_index; #[tokio::test] async fn test_async_reader() { @@ -752,126 +753,6 @@ async fn test_nested_skip() { } } -#[tokio::test] -async fn non_empty_offset_index_doesnt_panic_in_read_row_group() { - use tokio::fs::File; - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/alltypes_tiny_pages.parquet"); - let mut file = File::open(&path).await.unwrap(); - let file_size = file.metadata().await.unwrap().len(); - let metadata = ParquetMetaDataReader::new() - .with_page_index_policy(PageIndexPolicy::Required) - .load_and_finish(&mut file, file_size) - .await - .unwrap(); - - let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); - let arrow_reader_metadata = ArrowReaderMetadata::try_new(metadata.into(), options).unwrap(); - let reader = ParquetRecordBatchStreamBuilder::new_with_metadata(file, arrow_reader_metadata) - .build() - .unwrap(); - - let result = reader.try_collect::<Vec<_>>().await.unwrap(); - assert_eq!(result.len(), 8); -} - -#[tokio::test] -async fn empty_offset_index_doesnt_panic_in_column_chunks() { - use tempfile::TempDir; - use tokio::fs::File; - fn write_metadata_to_local_file(metadata: ParquetMetaData, file: impl AsRef<std::path::Path>) { - use parquet::file::metadata::ParquetMetaDataWriter; - use std::fs::File; - let file = File::create(file).unwrap(); - ParquetMetaDataWriter::new(file, &metadata) - .finish() - .unwrap() - } - - fn read_metadata_from_local_file(file: impl AsRef<std::path::Path>) -> ParquetMetaData { - use std::fs::File; - let file = File::open(file).unwrap(); - ParquetMetaDataReader::new() - .with_page_index_policy(PageIndexPolicy::Required) - .parse_and_finish(&file) - .unwrap() - } - - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/alltypes_plain.parquet"); - let mut file = File::open(&path).await.unwrap(); - let file_size = file.metadata().await.unwrap().len(); - let metadata = ParquetMetaDataReader::new() - .with_page_index_policy(PageIndexPolicy::Required) - .load_and_finish(&mut file, file_size) - .await - .unwrap(); - - let tempdir = TempDir::new().unwrap(); - let metadata_path = tempdir.path().join("thrift_metadata.dat"); - write_metadata_to_local_file(metadata, &metadata_path); - let metadata = read_metadata_from_local_file(&metadata_path); - - let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); - let arrow_reader_metadata = ArrowReaderMetadata::try_new(metadata.into(), options).unwrap(); - let reader = ParquetRecordBatchStreamBuilder::new_with_metadata(file, arrow_reader_metadata) - .build() - .unwrap(); - - // Panics here - let result = reader.try_collect::<Vec<_>>().await.unwrap(); - assert_eq!(result.len(), 1); -} - -#[tokio::test] -async fn test_cached_array_reader_sparse_offset_error() { - use futures::TryStreamExt; - - use arrow_array::{BooleanArray, RecordBatch}; - use parquet::arrow::arrow_reader::{ArrowPredicateFn, RowFilter, RowSelection, RowSelector}; - - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet"); - let data = Bytes::from(std::fs::read(path).unwrap()); - - let async_reader = TestReader::new(data); - - // Enable page index so the fetch logic loads only required pages - let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); - let builder = ParquetRecordBatchStreamBuilder::new_with_options(async_reader, options) - .await - .unwrap(); - - // Skip the first 22 rows (entire first Parquet page) and then select the - // next 3 rows (22, 23, 24). This means the fetch step will not include - // the first page starting at file offset 0. - let selection = RowSelection::from(vec![RowSelector::skip(22), RowSelector::select(3)]); - - // Trivial predicate on column 0 that always returns `true`. Using the - // same column in both predicate and projection activates the caching - // layer (Producer/Consumer pattern). - let parquet_schema = builder.parquet_schema(); - let proj = ProjectionMask::leaves(parquet_schema, vec![0]); - let always_true = ArrowPredicateFn::new(proj.clone(), |batch: RecordBatch| { - Ok(BooleanArray::from(vec![true; batch.num_rows()])) - }); - let filter = RowFilter::new(vec![Box::new(always_true)]); - - // Build the stream with batch size 8 so the cache reads whole batches - // that straddle the requested row range (rows 0-7, 8-15, 16-23, …). - let stream = builder - .with_batch_size(8) - .with_projection(proj) - .with_row_selection(selection) - .with_row_filter(filter) - .build() - .unwrap(); - - // Collecting the stream should fail with the sparse column chunk offset - // error we want to reproduce. - let _result: Vec<_> = stream.try_collect().await.unwrap(); -} - #[tokio::test] async fn test_predicate_cache_disabled() { let k = Int32Array::from_iter_values(0..10); diff --git a/parquet/tests/arrow_reader/async_reader/offset_index.rs b/parquet/tests/arrow_reader/async_reader/offset_index.rs new file mode 100644 index 0000000000..043de47bea --- /dev/null +++ b/parquet/tests/arrow_reader/async_reader/offset_index.rs @@ -0,0 +1,143 @@ +// 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. + +use crate::io::TestReader; +use bytes::Bytes; +use futures::TryStreamExt; +use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; +use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask}; +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader}; + +#[tokio::test] +async fn non_empty_offset_index_doesnt_panic_in_read_row_group() { + use tokio::fs::File; + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_tiny_pages.parquet"); + let mut file = File::open(&path).await.unwrap(); + let file_size = file.metadata().await.unwrap().len(); + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .load_and_finish(&mut file, file_size) + .await + .unwrap(); + + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let arrow_reader_metadata = ArrowReaderMetadata::try_new(metadata.into(), options).unwrap(); + let reader = ParquetRecordBatchStreamBuilder::new_with_metadata(file, arrow_reader_metadata) + .build() + .unwrap(); + + let result = reader.try_collect::<Vec<_>>().await.unwrap(); + assert_eq!(result.len(), 8); +} + +#[tokio::test] +async fn empty_offset_index_doesnt_panic_in_column_chunks() { + use tempfile::TempDir; + use tokio::fs::File; + fn write_metadata_to_local_file(metadata: ParquetMetaData, file: impl AsRef<std::path::Path>) { + use parquet::file::metadata::ParquetMetaDataWriter; + use std::fs::File; + let file = File::create(file).unwrap(); + ParquetMetaDataWriter::new(file, &metadata) + .finish() + .unwrap() + } + + fn read_metadata_from_local_file(file: impl AsRef<std::path::Path>) -> ParquetMetaData { + use std::fs::File; + let file = File::open(file).unwrap(); + ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .parse_and_finish(&file) + .unwrap() + } + + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_plain.parquet"); + let mut file = File::open(&path).await.unwrap(); + let file_size = file.metadata().await.unwrap().len(); + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .load_and_finish(&mut file, file_size) + .await + .unwrap(); + + let tempdir = TempDir::new().unwrap(); + let metadata_path = tempdir.path().join("thrift_metadata.dat"); + write_metadata_to_local_file(metadata, &metadata_path); + let metadata = read_metadata_from_local_file(&metadata_path); + + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let arrow_reader_metadata = ArrowReaderMetadata::try_new(metadata.into(), options).unwrap(); + let reader = ParquetRecordBatchStreamBuilder::new_with_metadata(file, arrow_reader_metadata) + .build() + .unwrap(); + + // Panics here + let result = reader.try_collect::<Vec<_>>().await.unwrap(); + assert_eq!(result.len(), 1); +} + +#[tokio::test] +async fn test_cached_array_reader_sparse_offset_error() { + use futures::TryStreamExt; + + use arrow_array::{BooleanArray, RecordBatch}; + use parquet::arrow::arrow_reader::{ArrowPredicateFn, RowFilter, RowSelection, RowSelector}; + + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet"); + let data = Bytes::from(std::fs::read(path).unwrap()); + + let async_reader = TestReader::new(data); + + // Enable page index so the fetch logic loads only required pages + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let builder = ParquetRecordBatchStreamBuilder::new_with_options(async_reader, options) + .await + .unwrap(); + + // Skip the first 22 rows (entire first Parquet page) and then select the + // next 3 rows (22, 23, 24). This means the fetch step will not include + // the first page starting at file offset 0. + let selection = RowSelection::from(vec![RowSelector::skip(22), RowSelector::select(3)]); + + // Trivial predicate on column 0 that always returns `true`. Using the + // same column in both predicate and projection activates the caching + // layer (Producer/Consumer pattern). + let parquet_schema = builder.parquet_schema(); + let proj = ProjectionMask::leaves(parquet_schema, vec![0]); + let always_true = ArrowPredicateFn::new(proj.clone(), |batch: RecordBatch| { + Ok(BooleanArray::from(vec![true; batch.num_rows()])) + }); + let filter = RowFilter::new(vec![Box::new(always_true)]); + + // Build the stream with batch size 8 so the cache reads whole batches + // that straddle the requested row range (rows 0-7, 8-15, 16-23, …). + let stream = builder + .with_batch_size(8) + .with_projection(proj) + .with_row_selection(selection) + .with_row_filter(filter) + .build() + .unwrap(); + + // Collecting the stream should fail with the sparse column chunk offset + // error we want to reproduce. + let _result: Vec<_> = stream.try_collect().await.unwrap(); +}
