This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 66abc6cebc Move Row Filter tests to a separate module (#9335)
66abc6cebc is described below
commit 66abc6cebc9e338748fa0baf24a87d3dab640822
Author: Kosta Tarasov <[email protected]>
AuthorDate: Tue Feb 3 14:18:58 2026 -0500
Move Row Filter tests to a separate module (#9335)
# Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax.
-->
- Part of #9269.
# Rationale for this change
Check issue
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
# What changes are included in this PR?
Moved tests created in #8733 to a separate module ->
`parquet/tests/arrow_reader/row_filter.rs`
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
# Are these changes tested?
No need, code movement.
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
# Are there any user-facing changes?
No, code movement.
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please call them out.
-->
---
parquet/src/arrow/arrow_reader/mod.rs | 163 +------------------------
parquet/tests/arrow_reader/mod.rs | 1 +
parquet/tests/arrow_reader/row_filter.rs | 196 +++++++++++++++++++++++++++++++
3 files changed, 198 insertions(+), 162 deletions(-)
diff --git a/parquet/src/arrow/arrow_reader/mod.rs
b/parquet/src/arrow/arrow_reader/mod.rs
index 2c3e0a6e38..3a31c69ff3 100644
--- a/parquet/src/arrow/arrow_reader/mod.rs
+++ b/parquet/src/arrow/arrow_reader/mod.rs
@@ -1582,7 +1582,7 @@ pub(crate) mod tests {
use crate::arrow::arrow_reader::{
ArrowPredicateFn, ArrowReaderBuilder, ArrowReaderMetadata,
ArrowReaderOptions,
ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder, RowFilter,
RowSelection,
- RowSelectionPolicy, RowSelector,
+ RowSelector,
};
use crate::arrow::schema::{
add_encoded_arrow_schema_to_metadata,
@@ -1602,8 +1602,6 @@ pub(crate) mod tests {
use crate::schema::parser::parse_message_type;
use crate::schema::types::{Type, TypePtr};
use crate::util::test_common::rand_gen::RandGen;
- use arrow::compute::kernels::cmp::eq;
- use arrow::compute::or;
use arrow_array::builder::*;
use arrow_array::cast::AsArray;
use arrow_array::types::{
@@ -5125,93 +5123,6 @@ pub(crate) mod tests {
assert_eq!(out, batch.slice(2, 1));
}
- #[test]
- fn test_row_selection_interleaved_skip() -> Result<()> {
- let schema = Arc::new(Schema::new(vec![Field::new(
- "v",
- ArrowDataType::Int32,
- false,
- )]));
-
- let values = Int32Array::from(vec![0, 1, 2, 3, 4]);
- let batch = RecordBatch::try_from_iter([("v", Arc::new(values) as
ArrayRef)]).unwrap();
-
- let mut buffer = Vec::with_capacity(1024);
- let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(),
None).unwrap();
- writer.write(&batch)?;
- writer.close()?;
-
- let selection = RowSelection::from(vec![
- RowSelector::select(1),
- RowSelector::skip(2),
- RowSelector::select(2),
- ]);
-
- let mut reader =
ParquetRecordBatchReaderBuilder::try_new(Bytes::from(buffer))?
- .with_batch_size(4)
- .with_row_selection(selection)
- .build()?;
-
- let out = reader.next().unwrap()?;
- assert_eq!(out.num_rows(), 3);
- let values = out
- .column(0)
- .as_primitive::<arrow_array::types::Int32Type>()
- .values();
- assert_eq!(values, &[0, 3, 4]);
- assert!(reader.next().is_none());
- Ok(())
- }
-
- #[test]
- fn test_row_selection_mask_sparse_rows() -> Result<()> {
- let schema = Arc::new(Schema::new(vec![Field::new(
- "v",
- ArrowDataType::Int32,
- false,
- )]));
-
- let values = Int32Array::from((0..30).collect::<Vec<i32>>());
- let batch = RecordBatch::try_from_iter([("v", Arc::new(values) as
ArrayRef)])?;
-
- let mut buffer = Vec::with_capacity(1024);
- let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(),
None)?;
- writer.write(&batch)?;
- writer.close()?;
-
- let total_rows = batch.num_rows();
- let ranges = (1..total_rows)
- .step_by(2)
- .map(|i| i..i + 1)
- .collect::<Vec<_>>();
- let selection =
RowSelection::from_consecutive_ranges(ranges.into_iter(), total_rows);
-
- let selectors: Vec<RowSelector> = selection.clone().into();
- assert!(total_rows < selectors.len() * 8);
-
- let bytes = Bytes::from(buffer);
-
- let reader = ParquetRecordBatchReaderBuilder::try_new(bytes.clone())?
- .with_batch_size(7)
- .with_row_selection(selection)
- .build()?;
-
- let mut collected = Vec::new();
- for batch in reader {
- let batch = batch?;
- collected.extend_from_slice(
- batch
- .column(0)
- .as_primitive::<arrow_array::types::Int32Type>()
- .values(),
- );
- }
-
- let expected: Vec<i32> = (1..total_rows).step_by(2).map(|i| i as
i32).collect();
- assert_eq!(collected, expected);
- Ok(())
- }
-
fn test_decimal32_roundtrip() {
let d = |values: Vec<i32>, p: u8| {
let iter = values.into_iter();
@@ -5594,78 +5505,6 @@ pub(crate) mod tests {
c0.iter().zip(c1.iter()).for_each(|(l, r)| assert_eq!(l, r));
}
- #[test]
- fn test_row_filter_full_page_skip_is_handled() {
- let first_value: i64 = 1111;
- let last_value: i64 = 9999;
- let num_rows: usize = 12;
-
- // build data with row selection average length 4
- // The result would be (1111 XXXX) ... (4 page in the middle)... (XXXX
9999)
- // The Row Selection would be [1111, (skip 10), 9999]
- let schema = Arc::new(Schema::new(vec![
- Field::new("key", arrow_schema::DataType::Int64, false),
- Field::new("value", arrow_schema::DataType::Int64, false),
- ]));
-
- let mut int_values: Vec<i64> = (0..num_rows as i64).collect();
- int_values[0] = first_value;
- int_values[num_rows - 1] = last_value;
- let keys = Int64Array::from(int_values.clone());
- let values = Int64Array::from(int_values.clone());
- let batch = RecordBatch::try_new(
- Arc::clone(&schema),
- vec![Arc::new(keys) as ArrayRef, Arc::new(values) as ArrayRef],
- )
- .unwrap();
-
- let props = WriterProperties::builder()
- .set_write_batch_size(2)
- .set_data_page_row_count_limit(2)
- .build();
-
- let mut buffer = Vec::new();
- let mut writer = ArrowWriter::try_new(&mut buffer, schema,
Some(props)).unwrap();
- writer.write(&batch).unwrap();
- writer.close().unwrap();
- let data = Bytes::from(buffer);
-
- let options =
ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required);
- let builder =
-
ParquetRecordBatchReaderBuilder::try_new_with_options(data.clone(),
options).unwrap();
- let schema = builder.parquet_schema().clone();
- let filter_mask = ProjectionMask::leaves(&schema, [0]);
-
- let make_predicate = |mask: ProjectionMask| {
- ArrowPredicateFn::new(mask, move |batch: RecordBatch| {
- let column = batch.column(0);
- let match_first = eq(column,
&Int64Array::new_scalar(first_value))?;
- let match_second = eq(column,
&Int64Array::new_scalar(last_value))?;
- or(&match_first, &match_second)
- })
- };
-
- let options =
ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required);
- let predicate = make_predicate(filter_mask.clone());
-
- // The batch size is set to 12 to read all rows in one go after
filtering
- // If the Reader chooses mask to handle filter, it might cause panic
because the mid 4 pages may not be decoded.
- let reader =
ParquetRecordBatchReaderBuilder::try_new_with_options(data.clone(), options)
- .unwrap()
- .with_row_filter(RowFilter::new(vec![Box::new(predicate)]))
- .with_row_selection_policy(RowSelectionPolicy::Auto { threshold:
32 })
- .with_batch_size(12)
- .build()
- .unwrap();
-
- // Predicate pruning used to panic once mask-backed plans removed
whole pages.
- // Collecting into batches validates the plan now downgrades to
selectors instead.
- let schema = reader.schema().clone();
- let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
- let result = concat_batches(&schema, &batches).unwrap();
- assert_eq!(result.num_rows(), 2);
- }
-
#[test]
fn test_get_row_group_column_bloom_filter_with_length() {
// convert to new parquet file with bloom_filter_length
diff --git a/parquet/tests/arrow_reader/mod.rs
b/parquet/tests/arrow_reader/mod.rs
index 3d566306a9..9acfebda48 100644
--- a/parquet/tests/arrow_reader/mod.rs
+++ b/parquet/tests/arrow_reader/mod.rs
@@ -45,6 +45,7 @@ mod int96_stats_roundtrip;
mod io;
#[cfg(feature = "async")]
mod predicate_cache;
+mod row_filter;
mod statistics;
// returns a struct array with columns "int32_col", "float32_col" and
"float64_col" with the specified values
diff --git a/parquet/tests/arrow_reader/row_filter.rs
b/parquet/tests/arrow_reader/row_filter.rs
new file mode 100644
index 0000000000..78ba7569a4
--- /dev/null
+++ b/parquet/tests/arrow_reader/row_filter.rs
@@ -0,0 +1,196 @@
+// 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 std::sync::Arc;
+
+use arrow::{
+ array::AsArray,
+ compute::{concat_batches, kernels::cmp::eq, or},
+};
+use arrow_array::{ArrayRef, Int32Array, Int64Array, RecordBatch,
RecordBatchReader};
+use arrow_schema::{DataType as ArrowDataType, Field, Schema};
+use bytes::Bytes;
+use parquet::{
+ arrow::{
+ ArrowWriter, ProjectionMask,
+ arrow_reader::{
+ ArrowPredicateFn, ArrowReaderOptions,
ParquetRecordBatchReaderBuilder, RowFilter,
+ RowSelection, RowSelectionPolicy, RowSelector,
+ },
+ },
+ errors::Result,
+ file::{metadata::PageIndexPolicy, properties::WriterProperties},
+};
+
+#[test]
+fn test_row_selection_interleaved_skip() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "v",
+ ArrowDataType::Int32,
+ false,
+ )]));
+
+ let values = Int32Array::from(vec![0, 1, 2, 3, 4]);
+ let batch = RecordBatch::try_from_iter([("v", Arc::new(values) as
ArrayRef)]).unwrap();
+
+ let mut buffer = Vec::with_capacity(1024);
+ let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(),
None).unwrap();
+ writer.write(&batch)?;
+ writer.close()?;
+
+ let selection = RowSelection::from(vec![
+ RowSelector::select(1),
+ RowSelector::skip(2),
+ RowSelector::select(2),
+ ]);
+
+ let mut reader =
ParquetRecordBatchReaderBuilder::try_new(Bytes::from(buffer))?
+ .with_batch_size(4)
+ .with_row_selection(selection)
+ .build()?;
+
+ let out = reader.next().unwrap()?;
+ assert_eq!(out.num_rows(), 3);
+ let values = out
+ .column(0)
+ .as_primitive::<arrow_array::types::Int32Type>()
+ .values();
+ assert_eq!(values, &[0, 3, 4]);
+ assert!(reader.next().is_none());
+ Ok(())
+}
+
+#[test]
+fn test_row_selection_mask_sparse_rows() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "v",
+ ArrowDataType::Int32,
+ false,
+ )]));
+
+ let values = Int32Array::from((0..30).collect::<Vec<i32>>());
+ let batch = RecordBatch::try_from_iter([("v", Arc::new(values) as
ArrayRef)])?;
+
+ let mut buffer = Vec::with_capacity(1024);
+ let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(), None)?;
+ writer.write(&batch)?;
+ writer.close()?;
+
+ let total_rows = batch.num_rows();
+ let ranges = (1..total_rows)
+ .step_by(2)
+ .map(|i| i..i + 1)
+ .collect::<Vec<_>>();
+ let selection = RowSelection::from_consecutive_ranges(ranges.into_iter(),
total_rows);
+
+ let selectors: Vec<RowSelector> = selection.clone().into();
+ assert!(total_rows < selectors.len() * 8);
+
+ let bytes = Bytes::from(buffer);
+
+ let reader = ParquetRecordBatchReaderBuilder::try_new(bytes.clone())?
+ .with_batch_size(7)
+ .with_row_selection(selection)
+ .build()?;
+
+ let mut collected = Vec::new();
+ for batch in reader {
+ let batch = batch?;
+ collected.extend_from_slice(
+ batch
+ .column(0)
+ .as_primitive::<arrow_array::types::Int32Type>()
+ .values(),
+ );
+ }
+
+ let expected: Vec<i32> = (1..total_rows).step_by(2).map(|i| i as
i32).collect();
+ assert_eq!(collected, expected);
+ Ok(())
+}
+
+#[test]
+fn test_row_filter_full_page_skip_is_handled() {
+ let first_value: i64 = 1111;
+ let last_value: i64 = 9999;
+ let num_rows: usize = 12;
+
+ // build data with row selection average length 4
+ // The result would be (1111 XXXX) ... (4 page in the middle)... (XXXX
9999)
+ // The Row Selection would be [1111, (skip 10), 9999]
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("key", arrow_schema::DataType::Int64, false),
+ Field::new("value", arrow_schema::DataType::Int64, false),
+ ]));
+
+ let mut int_values: Vec<i64> = (0..num_rows as i64).collect();
+ int_values[0] = first_value;
+ int_values[num_rows - 1] = last_value;
+ let keys = Int64Array::from(int_values.clone());
+ let values = Int64Array::from(int_values.clone());
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(keys) as ArrayRef, Arc::new(values) as ArrayRef],
+ )
+ .unwrap();
+
+ let props = WriterProperties::builder()
+ .set_write_batch_size(2)
+ .set_data_page_row_count_limit(2)
+ .build();
+
+ let mut buffer = Vec::new();
+ let mut writer = ArrowWriter::try_new(&mut buffer, schema,
Some(props)).unwrap();
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+ let data = Bytes::from(buffer);
+
+ let options =
ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required);
+ let builder =
+ ParquetRecordBatchReaderBuilder::try_new_with_options(data.clone(),
options).unwrap();
+ let schema = builder.parquet_schema().clone();
+ let filter_mask = ProjectionMask::leaves(&schema, [0]);
+
+ let make_predicate = |mask: ProjectionMask| {
+ ArrowPredicateFn::new(mask, move |batch: RecordBatch| {
+ let column = batch.column(0);
+ let match_first = eq(column,
&Int64Array::new_scalar(first_value))?;
+ let match_second = eq(column,
&Int64Array::new_scalar(last_value))?;
+ or(&match_first, &match_second)
+ })
+ };
+
+ let options =
ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required);
+ let predicate = make_predicate(filter_mask.clone());
+
+ // The batch size is set to 12 to read all rows in one go after filtering
+ // If the Reader chooses mask to handle filter, it might cause panic
because the mid 4 pages may not be decoded.
+ let reader =
ParquetRecordBatchReaderBuilder::try_new_with_options(data.clone(), options)
+ .unwrap()
+ .with_row_filter(RowFilter::new(vec![Box::new(predicate)]))
+ .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 32 })
+ .with_batch_size(12)
+ .build()
+ .unwrap();
+
+ // Predicate pruning used to panic once mask-backed plans removed whole
pages.
+ // Collecting into batches validates the plan now downgrades to selectors
instead.
+ let schema = reader.schema().clone();
+ let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
+ let result = concat_batches(&schema, &batches).unwrap();
+ assert_eq!(result.num_rows(), 2);
+}
