diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index 1b02c4ae25d3..8fed531abbee 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/parquet/src/arrow/arrow_reader/mod.rs @@ -1368,12 +1368,15 @@ impl ParquetRecordBatchReader { if batch_size == 0 { return Ok(None); } + let page_boundaries = self.read_plan.page_boundaries(); match self.read_plan.row_selection_cursor_mut() { RowSelectionCursor::Mask(mask_cursor) => { // Stream the record batch reader using contiguous segments of the selection // mask, avoiding the need to materialize intermediate `RowSelector` ranges. while !mask_cursor.is_empty() { - let Some(mask_chunk) = mask_cursor.next_mask_chunk(batch_size) else { + let Some(mask_chunk) = + mask_cursor.next_mask_chunk(batch_size, page_boundaries.as_deref()) + else { return Ok(None); }; diff --git a/parquet/src/arrow/arrow_reader/read_plan.rs b/parquet/src/arrow/arrow_reader/read_plan.rs index 7c9eb36befe3..38f9d5b96585 100644 --- a/parquet/src/arrow/arrow_reader/read_plan.rs +++ b/parquet/src/arrow/arrow_reader/read_plan.rs @@ -28,6 +28,7 @@ use crate::errors::{ParquetError, Result}; use arrow_array::Array; use arrow_select::filter::prep_null_mask_filter; use std::collections::VecDeque; +use std::sync::Arc; /// A builder for [`ReadPlan`] #[derive(Clone, Debug)] @@ -37,6 +38,8 @@ pub struct ReadPlanBuilder { selection: Option, /// Policy to use when materializing the row selection row_selection_policy: RowSelectionPolicy, + /// Precomputed page boundary row indices for mask chunking + page_boundaries: Option>, } impl ReadPlanBuilder { @@ -46,6 +49,7 @@ impl ReadPlanBuilder { batch_size, selection: None, row_selection_policy: RowSelectionPolicy::default(), + page_boundaries: None, } } @@ -175,6 +179,12 @@ impl ReadPlanBuilder { Ok(self) } + /// Set page boundary rows directly for mask chunking + pub(crate) fn with_page_boundaries(mut self, boundaries: Option>) -> Self { + self.page_boundaries = boundaries; + self + } + /// Create a final `ReadPlan` the read plan for the scan pub fn build(mut self) -> ReadPlan { // If selection is empty, truncate @@ -189,6 +199,7 @@ impl ReadPlanBuilder { batch_size, selection, row_selection_policy: _, + page_boundaries: _, } = self; let selection = selection.map(|s| s.trim()); @@ -209,6 +220,7 @@ impl ReadPlanBuilder { ReadPlan { batch_size, row_selection_cursor, + page_boundaries: self.page_boundaries, } } } @@ -307,6 +319,8 @@ pub struct ReadPlan { batch_size: usize, /// Row ranges to be selected from the data source row_selection_cursor: RowSelectionCursor, + /// Precomputed page boundary row indices for mask chunking + page_boundaries: Option>, } impl ReadPlan { @@ -330,6 +344,11 @@ impl ReadPlan { pub fn batch_size(&self) -> usize { self.batch_size } + + /// Return the page boundary row indices used for mask chunking + pub fn page_boundaries(&self) -> Option> { + self.page_boundaries.clone() + } } #[cfg(test)] diff --git a/parquet/src/arrow/arrow_reader/selection.rs b/parquet/src/arrow/arrow_reader/selection.rs index 2ddf812f9c39..455ef849edaf 100644 --- a/parquet/src/arrow/arrow_reader/selection.rs +++ b/parquet/src/arrow/arrow_reader/selection.rs @@ -271,8 +271,64 @@ impl RowSelection { }) } - /// Returns true if selectors should be forced, preventing mask materialisation - pub(crate) fn should_force_selectors( + /// Returns row offsets for the starts of skipped pages across projected columns + fn skipped_page_row_offsets( + &self, + projection: &ProjectionMask, + columns: &[OffsetIndexMetaData], + ) -> Vec { + let mut skipped_page_rows: Vec = Vec::new(); + + for (leaf_idx, column) in columns.iter().enumerate() { + if !projection.leaf_included(leaf_idx) { + continue; + } + + let locations = column.page_locations(); + if locations.is_empty() { + continue; + } + + let selected_ranges = self.scan_ranges(locations); + let mut selected_idx = 0usize; + for page in locations { + let page_start = page.offset as u64; + + while selected_idx < selected_ranges.len() + && selected_ranges[selected_idx].start < page_start + { + selected_idx += 1; + } + + let selected = selected_idx < selected_ranges.len() + && selected_ranges[selected_idx].start == page_start; + + if selected { + selected_idx += 1; + } else { + skipped_page_rows.push(page.first_row_index as usize); + } + } + } + + skipped_page_rows.sort_unstable(); + skipped_page_rows.dedup(); + skipped_page_rows + } + + /// Returns row offsets for skipped page starts when page-aware mask chunking is needed + pub(crate) fn page_aware_mask_boundaries( + &self, + projection: &ProjectionMask, + offset_index: Option<&[OffsetIndexMetaData]>, + ) -> Option> { + offset_index + .map(|columns| self.skipped_page_row_offsets(projection, columns)) + .filter(|offsets| !offsets.is_empty()) + } + + /// Returns true if bitmasks should be page aware + pub(crate) fn requires_page_aware_mask( &self, projection: &ProjectionMask, offset_index: Option<&[OffsetIndexMetaData]>, @@ -770,6 +826,9 @@ pub struct MaskCursor { mask: BooleanBuffer, /// Current absolute offset into the selection position: usize, + /// Index of the next page boundary candidate. This advances monotonically + /// as `position` advances. + next_boundary_idx: usize, } impl MaskCursor { @@ -778,8 +837,13 @@ impl MaskCursor { self.position >= self.mask.len() } - /// Advance through the mask representation, producing the next chunk summary - pub fn next_mask_chunk(&mut self, batch_size: usize) -> Option { + /// Advance through the mask representation, producing the next chunk summary. + /// Optionally clips chunk boundaries to the next page boundary. + pub fn next_mask_chunk( + &mut self, + batch_size: usize, + page_boundaries: Option<&[usize]>, + ) -> Option { let (initial_skip, chunk_rows, selected_rows, mask_start, end_position) = { let mask = &self.mask; @@ -791,6 +855,7 @@ impl MaskCursor { let mut cursor = start_position; let mut initial_skip = 0; + // Skip unselected rows while cursor < mask.len() && !mask.value(cursor) { initial_skip += 1; cursor += 1; @@ -800,10 +865,23 @@ impl MaskCursor { let mut chunk_rows = 0; let mut selected_rows = 0; - // Advance until enough rows have been selected to satisfy the batch size, - // or until the mask is exhausted. This mirrors the behaviour of the legacy - // `RowSelector` queue-based iteration. - while cursor < mask.len() && selected_rows < batch_size { + let max_chunk_rows = page_boundaries + .and_then(|boundaries| { + while self.next_boundary_idx < boundaries.len() + && boundaries[self.next_boundary_idx] <= mask_start + { + self.next_boundary_idx += 1; + } + boundaries + .get(self.next_boundary_idx) + .and_then(|&start| (start > mask_start).then_some(start - mask_start)) + }) + .unwrap_or(usize::MAX); + + // Advance until enough rows have been selected to satisfy batch_size, + // or until the mask is exhausted or until a page boundary. + while cursor < mask.len() && selected_rows < batch_size && chunk_rows < max_chunk_rows { + // Increment counters chunk_rows += 1; if mask.value(cursor) { selected_rows += 1; @@ -906,6 +984,7 @@ impl RowSelectionCursor { Self::Mask(MaskCursor { mask: boolean_mask_from_selectors(&selectors), position: 0, + next_boundary_idx: 0, }) } @@ -1537,6 +1616,60 @@ mod tests { assert_eq!(ranges, vec![10..20, 20..30, 30..40]); } + #[test] + fn test_page_aware_mask_boundaries_skipped_pages_only() { + let selection = RowSelection::from(vec![ + RowSelector::skip(10), + RowSelector::select(10), + RowSelector::skip(10), + RowSelector::select(20), + RowSelector::skip(20), + ]); + + let page_locations = vec![ + PageLocation { + offset: 0, + compressed_page_size: 10, + first_row_index: 0, + }, + PageLocation { + offset: 10, + compressed_page_size: 10, + first_row_index: 10, + }, + PageLocation { + offset: 20, + compressed_page_size: 10, + first_row_index: 20, + }, + PageLocation { + offset: 30, + compressed_page_size: 10, + first_row_index: 30, + }, + PageLocation { + offset: 40, + compressed_page_size: 10, + first_row_index: 40, + }, + PageLocation { + offset: 50, + compressed_page_size: 10, + first_row_index: 50, + }, + ]; + + let offsets = selection.page_aware_mask_boundaries( + &ProjectionMask::all(), + Some(&[OffsetIndexMetaData { + page_locations, + unencoded_byte_array_data_bytes: None, + }]), + ); + + assert_eq!(offsets, Some(vec![0, 20, 50])); + } + #[test] fn test_from_ranges() { let ranges = [1..3, 4..6, 6..6, 8..8, 9..10]; diff --git a/parquet/src/arrow/push_decoder/reader_builder/mod.rs b/parquet/src/arrow/push_decoder/reader_builder/mod.rs index d3d78ca7c263..e8522d3b6763 100644 --- a/parquet/src/arrow/push_decoder/reader_builder/mod.rs +++ b/parquet/src/arrow/push_decoder/reader_builder/mod.rs @@ -22,9 +22,9 @@ use crate::DecodeResult; use crate::arrow::ProjectionMask; use crate::arrow::array_reader::{ArrayReaderBuilder, CacheOptions, RowGroupCache}; use crate::arrow::arrow_reader::metrics::ArrowReaderMetrics; -use crate::arrow::arrow_reader::selection::RowSelectionStrategy; use crate::arrow::arrow_reader::{ ParquetRecordBatchReader, ReadPlanBuilder, RowFilter, RowSelection, RowSelectionPolicy, + selection::RowSelectionStrategy, }; use crate::arrow::in_memory_row_group::ColumnChunkData; use crate::arrow::push_decoder::reader_builder::data::DataRequestBuilder; @@ -444,16 +444,13 @@ impl RowGroupReaderBuilder { // pages for subsequent predicates. plan_builder = plan_builder.with_row_selection_policy(self.row_selection_policy); - // Prepare to evaluate the filter. - // Note: first update the selection strategy to properly handle any pages - // pruned during fetch - plan_builder = override_selector_strategy_if_needed( - plan_builder, + let page_boundaries = self.compute_page_aware_boundaries( + row_group_idx, + plan_builder.selection(), predicate.projection(), - self.row_group_offset_index(row_group_idx), + plan_builder.resolve_selection_strategy() == RowSelectionStrategy::Mask, ); - // `with_predicate` actually evaluates the filter - + plan_builder = plan_builder.with_page_boundaries(page_boundaries); plan_builder = plan_builder.with_predicate(array_reader, filter_info.current_mut())?; @@ -553,12 +550,6 @@ impl RowGroupReaderBuilder { plan_builder = plan_builder.with_row_selection_policy(self.row_selection_policy); - plan_builder = override_selector_strategy_if_needed( - plan_builder, - &self.projection, - self.row_group_offset_index(row_group_idx), - ); - let row_group_info = RowGroupInfo { row_group_idx, row_count, @@ -605,6 +596,17 @@ impl RowGroupReaderBuilder { &mut self.buffers, )?; + // For mask-based selection, chunk only at starts of skipped pages to + // coalesce contiguous selected and skipped page runs. + let page_boundaries = self.compute_page_aware_boundaries( + row_group_idx, + plan_builder.selection(), + &self.projection, + plan_builder.resolve_selection_strategy() == RowSelectionStrategy::Mask, + ); + + let plan_builder = plan_builder.with_page_boundaries(page_boundaries); + let plan = plan_builder.build(); // if we have any cached results, connect them up @@ -661,7 +663,34 @@ impl RowGroupReaderBuilder { mask.without_nested_types(self.metadata.file_metadata().schema_descr()) } - /// Get the offset index for the specified row group, if any + /// Compute page-aware mask chunk boundaries for the given projection, if needed. + fn compute_page_aware_boundaries( + &self, + row_group_idx: usize, + selection: Option<&RowSelection>, + projection: &ProjectionMask, + mask_strategy: bool, + ) -> Option> { + if !mask_strategy + || selection.is_none() + || self.row_group_offset_index(row_group_idx).is_none() + { + return None; + } + + selection + .and_then(|selection| { + let offset_index = self.row_group_offset_index(row_group_idx); + if selection.requires_page_aware_mask(projection, offset_index) { + selection.page_aware_mask_boundaries(projection, offset_index) + } else { + None + } + }) + .map(Arc::<[usize]>::from) + } + + /// Get the column offset indexes for the specified row group, if any fn row_group_offset_index(&self, row_group_idx: usize) -> Option<&[OffsetIndexMetaData]> { self.metadata .offset_index() @@ -671,57 +700,6 @@ impl RowGroupReaderBuilder { } } -/// Override the selection strategy if needed. -/// -/// Some pages can be skipped during row-group construction if they are not read -/// by the selections. This means that the data pages for those rows are never -/// loaded and definition/repetition levels are never read. When using -/// `RowSelections` selection works because `skip_records()` handles this -/// case and skips the page accordingly. -/// -/// However, with the current mask design, all values must be read and decoded -/// and then a mask filter is applied. Thus if any pages are skipped during -/// row-group construction, the data pages are missing and cannot be decoded. -/// -/// A simple example: -/// * the page size is 2, the mask is 100001, row selection should be read(1) skip(4) read(1) -/// * the `ColumnChunkData` would be page1(10), page2(skipped), page3(01) -/// -/// Using the row selection to skip(4), page2 won't be read at all, so in this -/// case we can't decode all the rows and apply a mask. To correctly apply the -/// bit mask, we need all 6 values be read, but page2 is not in memory. -fn override_selector_strategy_if_needed( - plan_builder: ReadPlanBuilder, - projection_mask: &ProjectionMask, - offset_index: Option<&[OffsetIndexMetaData]>, -) -> ReadPlanBuilder { - // override only applies to Auto policy, If the policy is already Mask or Selectors, respect that - let RowSelectionPolicy::Auto { .. } = plan_builder.row_selection_policy() else { - return plan_builder; - }; - - let preferred_strategy = plan_builder.resolve_selection_strategy(); - - let force_selectors = matches!(preferred_strategy, RowSelectionStrategy::Mask) - && plan_builder.selection().is_some_and(|selection| { - selection.should_force_selectors(projection_mask, offset_index) - }); - - let resolved_strategy = if force_selectors { - RowSelectionStrategy::Selectors - } else { - preferred_strategy - }; - - // override the plan builder strategy with the resolved one - let new_policy = match resolved_strategy { - RowSelectionStrategy::Mask => RowSelectionPolicy::Mask, - RowSelectionStrategy::Selectors => RowSelectionPolicy::Selectors, - }; - - plan_builder.with_row_selection_policy(new_policy) -} - #[cfg(test)] mod tests { use super::*; @@ -729,6 +707,6 @@ mod tests { #[test] // Verify that the size of RowGroupDecoderState does not grow too large fn test_structure_size() { - assert_eq!(std::mem::size_of::(), 200); + assert_eq!(std::mem::size_of::(), 216); } } diff --git a/parquet/tests/arrow_reader/row_filter/async.rs b/parquet/tests/arrow_reader/row_filter/async.rs index 66840bb8147b..746bd94a850b 100644 --- a/parquet/tests/arrow_reader/row_filter/async.rs +++ b/parquet/tests/arrow_reader/row_filter/async.rs @@ -42,6 +42,7 @@ use parquet::{ metadata::{PageIndexPolicy, ParquetMetaDataReader}, properties::WriterProperties, }, + schema::types::ColumnPath, }; #[tokio::test] @@ -634,3 +635,267 @@ async fn test_multi_predicate_mask_policy_carryover() { // Plus even-indexed rows in [200,250) with value<250 → rows 200,202,...,248 (25 rows) assert_eq!(batch.num_rows(), 75); } + +/// Regression test for adaptive selection switching to mask materialization while +/// evaluating a predicate that references multiple columns with different page boundaries. +#[tokio::test] +async fn test_complex_predicate_pushdown_with_skipped_pages() { + const NUM_ROWS: usize = 240; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("int_flag", DataType::Int8, false), + Field::new("text_flag", DataType::Utf8, false), + ])); + + let huge_true = "T".repeat(128); + let huge_false = "F".repeat(128); + + let ids = Int32Array::from_iter_values(0..NUM_ROWS as i32); + let int_flag = Int8Array::from_iter_values((0..NUM_ROWS).map(|i| (i % 2 == 0) as i8)); + let text_flag = StringArray::from_iter_values((0..NUM_ROWS).map(|i| { + if i % 3 == 0 { + huge_true.as_str() + } else { + huge_false.as_str() + } + })); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(ids), Arc::new(int_flag), Arc::new(text_flag)], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_write_batch_size(1) + .set_data_page_size_limit(512) + .set_data_page_row_count_limit(128) + .set_column_data_page_size_limit(ColumnPath::from("text_flag"), 256) + .set_column_dictionary_enabled(ColumnPath::from("text_flag"), false) + .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 builder = ParquetRecordBatchStreamBuilder::new_with_options( + TestReader::new(data), + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), + ) + .await + .unwrap(); + + let schema_descr = builder.parquet_schema().clone(); + + // Start with a sparse RLE row selection, then apply a predicate that produces + // a dense, fragmented mask. + let selection = RowSelection::from(vec![ + RowSelector::select(60), + RowSelector::skip(120), + RowSelector::select(60), + ]); + + let huge_true_scalar = StringArray::from_iter_values([huge_true.as_str()]); + let predicate = + ArrowPredicateFn::new(ProjectionMask::roots(&schema_descr, [1, 2]), move |batch| { + let int_true = eq(batch.column(0), &Int8Array::new_scalar(1))?; + let text_true = eq(batch.column(1), &Scalar::new(&huge_true_scalar))?; + or(&int_true, &text_true) + }); + + let stream = builder + .with_row_selection(selection) + .with_row_filter(RowFilter::new(vec![Box::new(predicate)])) + .with_projection(ProjectionMask::roots(&schema_descr, [0])) + .with_batch_size(NUM_ROWS) + .with_row_selection_policy(RowSelectionPolicy::Mask) + .build() + .unwrap(); + + let batches: Vec = stream.try_collect().await.unwrap(); + let result = concat_batches(&batches[0].schema(), &batches).unwrap(); + + let expected_ids = + Int32Array::from_iter_values((0..60).chain(180..240).filter(|i| i % 2 == 0 || i % 3 == 0)); + + assert_eq!(result.num_columns(), 1); + assert_eq!( + result + .column(0) + .as_primitive::(), + &expected_ids + ); +} + +/// Regression test for mask materialization when a manual selection skips full pages +/// and projected columns have different page boundaries. +#[tokio::test] +async fn test_mask_selection_projection_with_skipped_pages() { + const NUM_ROWS: usize = 240; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("int_flag", DataType::Int8, false), + Field::new("text_flag", DataType::Utf8, false), + ])); + + let huge_true = "T".repeat(128); + let huge_false = "F".repeat(128); + + let ids = Int32Array::from_iter_values(0..NUM_ROWS as i32); + let int_flag = Int8Array::from_iter_values((0..NUM_ROWS).map(|i| (i % 2 == 0) as i8)); + let text_flag = StringArray::from_iter_values((0..NUM_ROWS).map(|i| { + if i % 3 == 0 { + huge_true.as_str() + } else { + huge_false.as_str() + } + })); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(ids), Arc::new(int_flag), Arc::new(text_flag)], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_write_batch_size(1) + .set_data_page_size_limit(512) + .set_data_page_row_count_limit(128) + .set_column_data_page_size_limit(ColumnPath::from("text_flag"), 256) + .set_column_dictionary_enabled(ColumnPath::from("text_flag"), false) + .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 builder = ParquetRecordBatchStreamBuilder::new_with_options( + TestReader::new(data), + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), + ) + .await + .unwrap(); + + let schema_descr = builder.parquet_schema().clone(); + let selection = RowSelection::from(vec![ + RowSelector::select(60), + RowSelector::skip(120), + RowSelector::select(60), + ]); + + let stream = builder + .with_row_selection(selection) + .with_row_selection_policy(RowSelectionPolicy::Mask) + .with_projection(ProjectionMask::roots(&schema_descr, [0, 2])) + .with_batch_size(NUM_ROWS) + .build() + .unwrap(); + + let batches: Vec = stream.try_collect().await.unwrap(); + let result = concat_batches(&batches[0].schema(), &batches).unwrap(); + + let expected_ids = Int32Array::from_iter_values((0..60).chain(180..240)); + + assert_eq!(result.num_columns(), 2); + assert_eq!( + result + .column(0) + .as_primitive::(), + &expected_ids + ); + assert_eq!(result.num_rows(), expected_ids.len()); +} + +/// Regression test for mask materialization during filtering when predicates span +/// columns with different page boundaries. +#[tokio::test] +async fn test_mask_selection_multi_col_predicate_with_skipped_pages() { + const NUM_ROWS: usize = 240; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("int_flag", DataType::Int8, false), + Field::new("text_flag", DataType::Utf8, false), + ])); + + let huge_true = "T".repeat(128); + let huge_false = "F".repeat(128); + + let ids = Int32Array::from_iter_values(0..NUM_ROWS as i32); + let int_flag = Int8Array::from_iter_values((0..NUM_ROWS).map(|i| (i % 2 == 0) as i8)); + let text_flag = StringArray::from_iter_values((0..NUM_ROWS).map(|i| { + if i % 3 == 0 { + huge_true.as_str() + } else { + huge_false.as_str() + } + })); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(ids), Arc::new(int_flag), Arc::new(text_flag)], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_write_batch_size(1) + .set_data_page_size_limit(512) + .set_data_page_row_count_limit(128) + .set_column_data_page_size_limit(ColumnPath::from("text_flag"), 256) + .set_column_dictionary_enabled(ColumnPath::from("text_flag"), false) + .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 builder = ParquetRecordBatchStreamBuilder::new_with_options( + TestReader::new(data), + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), + ) + .await + .unwrap(); + + let schema_descr = builder.parquet_schema().clone(); + let selection = RowSelection::from(vec![ + RowSelector::select(60), + RowSelector::skip(120), + RowSelector::select(60), + ]); + + let huge_true_scalar = StringArray::from_iter_values([huge_true.as_str()]); + let predicate = + ArrowPredicateFn::new(ProjectionMask::roots(&schema_descr, [1, 2]), move |batch| { + let int_true = eq(batch.column(0), &Int8Array::new_scalar(1))?; + let text_true = eq(batch.column(1), &Scalar::new(&huge_true_scalar))?; + or(&int_true, &text_true) + }); + + let stream = builder + .with_row_selection(selection) + .with_row_filter(RowFilter::new(vec![Box::new(predicate)])) + .with_projection(ProjectionMask::roots(&schema_descr, [0])) + .with_row_selection_policy(RowSelectionPolicy::Mask) + .with_batch_size(NUM_ROWS) + .build() + .unwrap(); + + let batches: Vec = stream.try_collect().await.unwrap(); + let result = concat_batches(&batches[0].schema(), &batches).unwrap(); + + let expected_ids = + Int32Array::from_iter_values((0..60).chain(180..240).filter(|i| i % 2 == 0 || i % 3 == 0)); + + assert_eq!(result.num_columns(), 1); + assert_eq!( + result + .column(0) + .as_primitive::(), + &expected_ids + ); +}