diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index bb4d350e5a..585723d416 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -79,6 +79,7 @@ pub fn iceberg::arrow::ArrowReaderBuilder::with_range_coalesce_bytes(self, range pub fn iceberg::arrow::ArrowReaderBuilder::with_range_fetch_concurrency(self, range_fetch_concurrency: usize) -> Self pub fn iceberg::arrow::ArrowReaderBuilder::with_row_group_filtering_enabled(self, row_group_filtering_enabled: bool) -> Self pub fn iceberg::arrow::ArrowReaderBuilder::with_row_selection_enabled(self, row_selection_enabled: bool) -> Self +pub fn iceberg::arrow::ArrowReaderBuilder::with_runtime_predicate_provider(self, runtime_predicate_provider: alloc::sync::Arc) -> Self pub struct iceberg::arrow::PartitionValueCalculator impl iceberg::arrow::PartitionValueCalculator pub fn iceberg::arrow::PartitionValueCalculator::calculate(&self, batch: &arrow_array::record_batch::RecordBatch) -> iceberg::Result @@ -93,6 +94,13 @@ pub fn iceberg::arrow::RecordBatchPartitionSplitter::split(&self, batch: &arrow_ pub fn iceberg::arrow::RecordBatchPartitionSplitter::try_new(iceberg_schema: iceberg::spec::SchemaRef, partition_spec: iceberg::spec::PartitionSpecRef, calculator: core::option::Option) -> iceberg::Result pub fn iceberg::arrow::RecordBatchPartitionSplitter::try_new_with_computed_values(iceberg_schema: iceberg::spec::SchemaRef, partition_spec: iceberg::spec::PartitionSpecRef) -> iceberg::Result pub fn iceberg::arrow::RecordBatchPartitionSplitter::try_new_with_precomputed_values(iceberg_schema: iceberg::spec::SchemaRef, partition_spec: iceberg::spec::PartitionSpecRef) -> iceberg::Result +pub struct iceberg::arrow::RuntimePredicateSnapshot +impl iceberg::arrow::RuntimePredicateSnapshot +pub fn iceberg::arrow::RuntimePredicateSnapshot::generation(&self) -> u64 +pub fn iceberg::arrow::RuntimePredicateSnapshot::into_predicate(self) -> core::option::Option +pub fn iceberg::arrow::RuntimePredicateSnapshot::new(predicate: core::option::Option, generation: u64) -> Self +impl core::fmt::Debug for iceberg::arrow::RuntimePredicateSnapshot +pub fn iceberg::arrow::RuntimePredicateSnapshot::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub struct iceberg::arrow::ScanMetrics impl iceberg::scan::ScanMetrics pub fn iceberg::scan::ScanMetrics::bytes_read(&self) -> u64 @@ -124,6 +132,9 @@ pub fn iceberg::arrow::ArrowSchemaVisitor::primitive(&mut self, p: &arrow_schema pub fn iceberg::arrow::ArrowSchemaVisitor::schema(&mut self, schema: &arrow_schema::schema::Schema, values: alloc::vec::Vec) -> iceberg::Result pub fn iceberg::arrow::ArrowSchemaVisitor::struct(&mut self, fields: &arrow_schema::fields::Fields, results: alloc::vec::Vec) -> iceberg::Result pub fn iceberg::arrow::ArrowSchemaVisitor::variant(&mut self, field: &arrow_schema::field::FieldRef) -> iceberg::Result where Self: core::marker::Sized +pub trait iceberg::arrow::RuntimePredicateProvider: core::marker::Send + core::marker::Sync +pub fn iceberg::arrow::RuntimePredicateProvider::generation(&self) -> u64 +pub fn iceberg::arrow::RuntimePredicateProvider::snapshot(&self) -> iceberg::Result pub fn iceberg::arrow::arrow_primitive_to_literal(primitive_array: &arrow_array::array::ArrayRef, ty: &iceberg::spec::Type) -> iceberg::Result>> pub fn iceberg::arrow::arrow_schema_to_schema(schema: &arrow_schema::schema::Schema) -> iceberg::Result pub fn iceberg::arrow::arrow_schema_to_schema_auto_assign_ids(schema: &arrow_schema::schema::Schema) -> iceberg::Result diff --git a/crates/iceberg/src/arrow/reader/mod.rs b/crates/iceberg/src/arrow/reader/mod.rs index 9d9c8be48e..2afc00261a 100644 --- a/crates/iceberg/src/arrow/reader/mod.rs +++ b/crates/iceberg/src/arrow/reader/mod.rs @@ -17,6 +17,8 @@ //! Parquet file data reader +use std::sync::Arc; + use crate::arrow::caching_delete_file_loader::CachingDeleteFileLoader; use crate::io::FileIO; use crate::runtime::Runtime; @@ -42,6 +44,9 @@ mod predicate_visitor; mod projection; mod row_filter; mod row_lineage; +mod runtime_predicate; +#[cfg(test)] +mod runtime_predicate_tests; pub use file_reader::ArrowFileReader; pub(crate) use options::ParquetReadOptions; use predicate_visitor::{CollectFieldIdVisitor, PredicateConverter}; @@ -49,6 +54,7 @@ use projection::{ add_fallback_field_ids_to_arrow_schema, apply_name_mapping_to_arrow_schema, find_leaf_by_field_id, }; +pub use runtime_predicate::{RuntimePredicateProvider, RuntimePredicateSnapshot}; /// Builder to create ArrowReader pub struct ArrowReaderBuilder { @@ -58,6 +64,7 @@ pub struct ArrowReaderBuilder { row_group_filtering_enabled: bool, row_selection_enabled: bool, bloom_filter_enabled: bool, + runtime_predicate_provider: Option>, parquet_read_options: ParquetReadOptions, runtime: Runtime, } @@ -74,6 +81,7 @@ impl ArrowReaderBuilder { row_group_filtering_enabled: true, row_selection_enabled: false, bloom_filter_enabled: false, + runtime_predicate_provider: None, parquet_read_options: ParquetReadOptions::builder().build(), runtime, } @@ -119,6 +127,21 @@ impl ArrowReaderBuilder { self } + /// Supplies execution-time predicates, sampled when each data-file task + /// starts and combined with the task's filters using AND. See + /// [`RuntimePredicateProvider`]. + /// + /// Only the reader accepts a provider. `TableScan::to_arrow` cannot pass + /// one, so wiring it through `TableScan` and execution engines is a + /// follow-up. + pub fn with_runtime_predicate_provider( + mut self, + runtime_predicate_provider: Arc, + ) -> Self { + self.runtime_predicate_provider = Some(runtime_predicate_provider); + self + } + /// Provide a hint as to the number of bytes to prefetch for parsing the Parquet metadata /// /// This hint can help reduce the number of fetch requests. For more details see the @@ -159,6 +182,7 @@ impl ArrowReaderBuilder { row_group_filtering_enabled: self.row_group_filtering_enabled, row_selection_enabled: self.row_selection_enabled, bloom_filter_enabled: self.bloom_filter_enabled, + runtime_predicate_provider: self.runtime_predicate_provider, parquet_read_options: self.parquet_read_options, } } @@ -177,5 +201,6 @@ pub struct ArrowReader { row_group_filtering_enabled: bool, row_selection_enabled: bool, bloom_filter_enabled: bool, + runtime_predicate_provider: Option>, parquet_read_options: ParquetReadOptions, } diff --git a/crates/iceberg/src/arrow/reader/pipeline.rs b/crates/iceberg/src/arrow/reader/pipeline.rs index 11fa200ee9..70fd4de42d 100644 --- a/crates/iceberg/src/arrow/reader/pipeline.rs +++ b/crates/iceberg/src/arrow/reader/pipeline.rs @@ -20,19 +20,22 @@ //! predicates, row-group / row selection, and delete handling into a stream //! of transformed Arrow `RecordBatch`es. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::sync::atomic::AtomicU64; use arrow_schema::{DataType, Field}; use futures::{StreamExt, TryStreamExt}; -use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; +use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions, RowFilter}; use parquet::arrow::{ PARQUET_FIELD_ID_META_KEY, ParquetRecordBatchStreamBuilder, ProjectionMask, RowNumber, }; use parquet::encryption::decrypt::FileDecryptionProperties; use super::row_lineage::synthesize_row_id_column; +use super::runtime_predicate::{ + RuntimePredicates, check_runtime_predicate_columns, intersect_page_selection, intersect_sorted, +}; use super::{ ArrowFileReader, ArrowReader, ParquetReadOptions, add_fallback_field_ids_to_arrow_schema, apply_name_mapping_to_arrow_schema, find_leaf_by_field_id, @@ -76,6 +79,9 @@ impl ArrowReader { row_selection_enabled: self.row_selection_enabled, bloom_filter_enabled: self.bloom_filter_enabled, parquet_read_options: self.parquet_read_options, + runtime_predicates: self + .runtime_predicate_provider + .map(|provider| Arc::new(RuntimePredicates::new(provider))), scan_metrics: scan_metrics.clone(), }; @@ -130,12 +136,32 @@ struct FileScanTaskReader { row_selection_enabled: bool, bloom_filter_enabled: bool, parquet_read_options: ParquetReadOptions, + /// Shared across tasks so each publication is bound once. + runtime_predicates: Option>, scan_metrics: ScanMetrics, } +/// A predicate resolved against one file, with its statistics-based row-group selection. +struct PlannedPredicate { + predicate: Arc, + field_ids: HashSet, + field_id_map: HashMap, + row_groups: Option>, + /// Runtime predicates are advisory: a planning or page-pruning failure skips them. + advisory: bool, +} + impl FileScanTaskReader { async fn process(self, task: FileScanTask) -> Result { - let should_load_page_index = (self.row_selection_enabled && task.predicate().is_some()) + let runtime_predicate = self.runtime_predicates.as_ref().and_then(|predicates| { + predicates.current( + &task.schema_ref(), + task.case_sensitive(), + task.data_file_path(), + ) + }); + let should_load_page_index = (self.row_selection_enabled + && (task.predicate().is_some() || runtime_predicate.is_some())) || !task.deletes().is_empty(); let mut parquet_read_options = self.parquet_read_options; parquet_read_options.preload_page_index = should_load_page_index; @@ -470,7 +496,7 @@ impl FileScanTaskReader { // we also have an optional predicate resulting from equality delete files. // If both are present, we logical-AND them together to form a single filter // predicate that we can pass to the `RecordBatchStreamBuilder`. - let final_predicate = match (task.predicate(), delete_predicate) { + let planned_predicate = match (task.predicate(), delete_predicate) { (None, None) => None, (Some(predicate), None) => Some(predicate.clone()), (None, Some(ref predicate)) => Some(predicate.clone()), @@ -508,76 +534,138 @@ impl FileScanTaskReader { selected_row_group_indices = Some(byte_range_filtered_row_groups); } - if let Some(predicate) = final_predicate { - let (iceberg_field_ids, field_id_map) = ArrowReader::build_field_id_set_and_map( - record_batch_stream_builder.parquet_schema(), - record_batch_stream_builder.schema(), - &predicate, + let mut plans = Vec::with_capacity(2); + if let Some(predicate) = planned_predicate { + plans.push(self.plan_predicate( + Arc::new(predicate), + false, + &record_batch_stream_builder, + &task, use_position_fallback, - )?; - - let row_filter = ArrowReader::get_row_filter( + )?); + } + // Skip the runtime predicate for this file if it does not store its + // columns exactly as the table types them, or if it cannot be planned. + if let Some(predicate) = runtime_predicate { + let planned = check_runtime_predicate_columns( &predicate, record_batch_stream_builder.parquet_schema(), - &iceberg_field_ids, - &field_id_map, - )?; - record_batch_stream_builder = record_batch_stream_builder.with_row_filter(row_filter); - - if self.row_group_filtering_enabled { - let predicate_filtered_row_groups = ArrowReader::get_selected_row_group_indices( - &predicate, - record_batch_stream_builder.metadata(), - &field_id_map, - task.schema(), - )?; + record_batch_stream_builder.schema(), + task.schema(), + use_position_fallback, + ) + .and_then(|()| { + self.plan_predicate( + predicate, + true, + &record_batch_stream_builder, + &task, + use_position_fallback, + ) + }); + match planned { + Ok(plan) => plans.push(plan), + Err(error) => tracing::debug!( + "Skipping runtime predicate for {}: {error}", + task.data_file_path() + ), + } + } + // The planned and runtime predicates form one Arrow predicate, so + // columns they share are decoded once and neither runs on the other's + // survivors. Separate Arrow predicates decoded shared columns twice and + // measured up to 2.4x slower than planned-only when the runtime + // predicate removed few rows. + let arrow_predicate = { + // Scoped so the builder borrow does not live across later awaits. + let row_filter_predicate = |plans: &[PlannedPredicate]| { + let predicates: Vec<_> = plans + .iter() + .map(|plan| (plan.predicate.as_ref(), &plan.field_ids, &plan.field_id_map)) + .collect(); + ArrowReader::get_arrow_predicate( + &predicates, + record_batch_stream_builder.parquet_schema(), + ) + }; + if plans.is_empty() { + None + } else { + match row_filter_predicate(&plans) { + Ok(arrow_predicate) => Some(arrow_predicate), + Err(error) if plans.last().is_some_and(|plan| plan.advisory) => { + tracing::debug!( + "Skipping runtime predicate for {}: {error}", + task.data_file_path() + ); + plans.pop(); + if plans.is_empty() { + None + } else { + Some(row_filter_predicate(&plans)?) + } + } + Err(error) => return Err(error), + } + } + }; + if let Some(arrow_predicate) = arrow_predicate { + record_batch_stream_builder = + record_batch_stream_builder.with_row_filter(RowFilter::new(vec![arrow_predicate])); + } + for plan in &plans { + if let Some(predicate_filtered_row_groups) = &plan.row_groups { // Merge predicate-based filtering with byte range filtering (if present) // by taking the intersection of both filters selected_row_group_indices = match selected_row_group_indices { - Some(byte_range_filtered) => { - // Keep only row groups that are in both filters - let intersection: Vec = byte_range_filtered - .into_iter() - .filter(|idx| predicate_filtered_row_groups.contains(idx)) - .collect(); - Some(intersection) - } - None => Some(predicate_filtered_row_groups), + // Keep only row groups that are in both filters + Some(byte_range_filtered) => Some(intersect_sorted( + &byte_range_filtered, + predicate_filtered_row_groups, + )), + None => Some(predicate_filtered_row_groups.clone()), }; } + } - if self.bloom_filter_enabled { - let all_rgs; - let candidate_rgs = match &selected_row_group_indices { - Some(indices) => indices.as_slice(), - None => { - all_rgs = (0..record_batch_stream_builder.metadata().num_row_groups()) - .collect::>(); - &all_rgs - } - }; + if self.bloom_filter_enabled && !plans.is_empty() { + let all_rgs; + let candidate_rgs = match &selected_row_group_indices { + Some(indices) => indices.as_slice(), + None => { + all_rgs = (0..record_batch_stream_builder.metadata().num_row_groups()) + .collect::>(); + &all_rgs + } + }; - let bloom_filtered = Self::filter_row_groups_by_bloom_filter( - &predicate, - &mut record_batch_stream_builder, - candidate_rgs, - &field_id_map, - ) - .await?; + let bloom_filtered = Self::filter_row_groups_by_bloom_filter( + &plans, + &mut record_batch_stream_builder, + candidate_rgs, + ) + .await?; - if bloom_filtered.len() < candidate_rgs.len() { - selected_row_group_indices = Some(bloom_filtered); - } + if bloom_filtered.len() < candidate_rgs.len() { + selected_row_group_indices = Some(bloom_filtered); } + } - if self.row_selection_enabled { - row_selection = ArrowReader::get_row_selection_for_filter_predicate( - &predicate, + if self.row_selection_enabled { + for plan in &plans { + let selection = ArrowReader::get_row_selection_for_filter_predicate( + &plan.predicate, record_batch_stream_builder.metadata(), &selected_row_group_indices, - &field_id_map, + &plan.field_id_map, task.schema(), + ); + row_selection = intersect_page_selection( + row_selection, + selection, + plan.advisory, + task.data_file_path(), )?; } } @@ -697,22 +785,64 @@ impl FileScanTaskReader { }) } - /// Reads bloom filters for relevant columns and evaluates the predicate - /// against them to filter out row groups that definitely don't match. + /// Resolves `predicate` against the file and plans its row-group selection. + fn plan_predicate( + &self, + predicate: Arc, + advisory: bool, + builder: &ParquetRecordBatchStreamBuilder, + task: &FileScanTask, + use_position_fallback: bool, + ) -> Result { + let (iceberg_field_ids, field_id_map) = ArrowReader::build_field_id_set_and_map( + builder.parquet_schema(), + builder.schema(), + &predicate, + use_position_fallback, + )?; + let row_groups = if self.row_group_filtering_enabled { + Some(ArrowReader::get_selected_row_group_indices( + &predicate, + builder.metadata(), + &field_id_map, + task.schema(), + )?) + } else { + None + }; + if self.bloom_filter_enabled { + collect_bloom_filter_field_ids(&predicate)?; + } + Ok(PlannedPredicate { + predicate, + field_ids: iceberg_field_ids, + field_id_map, + row_groups, + advisory, + }) + } + + /// Reads bloom filters for relevant columns and evaluates each predicate + /// against them to filter out row groups that definitely don't match. A + /// column's filter is read at most once per row group, even when several + /// predicates use it. async fn filter_row_groups_by_bloom_filter( - predicate: &BoundPredicate, + plans: &[PlannedPredicate], builder: &mut ParquetRecordBatchStreamBuilder, candidate_row_groups: &[usize], - field_id_map: &HashMap, ) -> Result> { // Only collect field IDs from eq/in predicates — the only types // bloom filters can help with. Skip columns not in the parquet schema. - let bloom_filter_field_ids: Vec = collect_bloom_filter_field_ids(predicate)? - .into_iter() - .filter(|id| field_id_map.contains_key(id)) - .collect(); + let mut bloom_filter_field_ids = Vec::with_capacity(plans.len()); + for plan in plans { + let field_ids: Vec = collect_bloom_filter_field_ids(&plan.predicate)? + .into_iter() + .filter(|id| plan.field_id_map.contains_key(id)) + .collect(); + bloom_filter_field_ids.push(field_ids); + } - if bloom_filter_field_ids.is_empty() { + if bloom_filter_field_ids.iter().all(Vec::is_empty) { return Ok(candidate_row_groups.to_vec()); } @@ -720,49 +850,66 @@ impl FileScanTaskReader { for &rg_idx in candidate_row_groups { let mut bloom_filters: HashMap = HashMap::new(); + let mut attempted: HashSet = HashSet::new(); + let mut might_match = true; - for &field_id in &bloom_filter_field_ids { - let col_idx = field_id_map[&field_id]; - let col_meta = builder.metadata().row_group(rg_idx).column(col_idx); - - // Only attempt to load if this column chunk actually has a bloom filter - if col_meta.bloom_filter_offset().is_none() { + for (plan, field_ids) in plans.iter().zip(&bloom_filter_field_ids) { + if field_ids.is_empty() { continue; } + for &field_id in field_ids { + if !attempted.insert(field_id) { + continue; + } + let col_idx = plan.field_id_map[&field_id]; + let col_meta = builder.metadata().row_group(rg_idx).column(col_idx); + + // Only attempt to load if this column chunk actually has a bloom filter + if col_meta.bloom_filter_offset().is_none() { + continue; + } + + let physical_type = col_meta.column_type(); + let type_length = col_meta.column_descr().type_length(); + + match builder + .get_row_group_column_bloom_filter(rg_idx, col_idx) + .await + { + Ok(Some(sbbf)) => { + bloom_filters.insert( + field_id, + ColumnBloomFilter::new(sbbf, physical_type, type_length), + ); + } + Ok(None) => {} + Err(e) => { + // Left absent from the map, so the evaluator treats the column + // as might-match and the row group survives. + tracing::debug!( + "Bloom filter for field {field_id} in row group {rg_idx} could not be read: {e}" + ); + } + } + } - let physical_type = col_meta.column_type(); - let type_length = col_meta.column_descr().type_length(); - - match builder - .get_row_group_column_bloom_filter(rg_idx, col_idx) - .await - { - Ok(Some(sbbf)) => { - bloom_filters.insert( - field_id, - ColumnBloomFilter::new(sbbf, physical_type, type_length), - ); + match BloomFilterEvaluator::eval(&plan.predicate, &bloom_filters) { + Ok(true) => {} + Ok(false) => { + // Row group pruned by bloom filter + might_match = false; + break; } - Ok(None) => {} Err(e) => { - // Left absent from the map, so the evaluator treats the column - // as might-match and the row group survives. tracing::debug!( - "Bloom filter for field {field_id} in row group {rg_idx} could not be read: {e}" + "Bloom filter evaluation failed for row group {rg_idx}, including it: {e}" ); } } } - match BloomFilterEvaluator::eval(predicate, &bloom_filters) { - Ok(true) => result.push(rg_idx), - Ok(false) => { /* Row group pruned by bloom filter */ } - Err(e) => { - tracing::debug!( - "Bloom filter evaluation failed for row group {rg_idx}, including it: {e}" - ); - result.push(rg_idx); - } + if might_match { + result.push(rg_idx); } } diff --git a/crates/iceberg/src/arrow/reader/row_filter.rs b/crates/iceberg/src/arrow/reader/row_filter.rs index 34b7391134..2982a378bf 100644 --- a/crates/iceberg/src/arrow/reader/row_filter.rs +++ b/crates/iceberg/src/arrow/reader/row_filter.rs @@ -23,13 +23,14 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; +use arrow_arith::boolean::and_kleene; use parquet::arrow::ProjectionMask; -use parquet::arrow::arrow_reader::{ArrowPredicateFn, RowFilter, RowSelection}; +use parquet::arrow::arrow_reader::{ArrowPredicate, ArrowPredicateFn, RowSelection}; use parquet::file::metadata::ParquetMetaData; use parquet::schema::types::SchemaDescriptor; use super::{ArrowReader, PredicateConverter}; -use crate::error::Result; +use crate::error::{Result, invalid_data}; use crate::expr::BoundPredicate; use crate::expr::visitors::bound_predicate_visitor::visit; use crate::expr::visitors::page_index_evaluator::PageIndexEvaluator; @@ -37,33 +38,54 @@ use crate::expr::visitors::row_group_metrics_evaluator::RowGroupMetricsEvaluator use crate::spec::Schema; impl ArrowReader { - pub(super) fn get_row_filter( - predicates: &BoundPredicate, + /// Converts predicates into one Arrow predicate of a `RowFilter` that ANDs + /// them. Each is resolved against its own field-id map, and the union of + /// their columns is decoded once. + pub(super) fn get_arrow_predicate( + predicates: &[(&BoundPredicate, &HashSet, &HashMap)], parquet_schema: &SchemaDescriptor, - iceberg_field_ids: &HashSet, - field_id_map: &HashMap, - ) -> Result { + ) -> Result> { // Collect Parquet column indices from field ids. // If the field id is not found in Parquet schema, it will be ignored due to schema evolution. - let mut column_indices = iceberg_field_ids + let mut column_indices = predicates .iter() - .filter_map(|field_id| field_id_map.get(field_id).cloned()) + .flat_map(|(_, field_ids, field_id_map)| { + field_ids + .iter() + .filter_map(move |field_id| field_id_map.get(field_id).cloned()) + }) .collect::>(); column_indices.sort(); - - // The converter that converts `BoundPredicates` to `ArrowPredicates` - let mut converter = PredicateConverter { - parquet_schema, - column_map: field_id_map, - column_indices: &column_indices, - }; - - // After collecting required leaf column indices used in the predicate, - // creates the projection mask for the Arrow predicates. - let projection_mask = ProjectionMask::leaves(parquet_schema, column_indices.clone()); - let predicate_func = visit(&mut converter, predicates)?; - let arrow_predicate = ArrowPredicateFn::new(projection_mask, predicate_func); - Ok(RowFilter::new(vec![Box::new(arrow_predicate)])) + column_indices.dedup(); + + let mut predicate_funcs = Vec::with_capacity(predicates.len()); + for (predicate, _, field_id_map) in predicates { + // The converter that converts `BoundPredicates` to `ArrowPredicates` + let mut converter = PredicateConverter { + parquet_schema, + column_map: field_id_map, + column_indices: &column_indices, + }; + predicate_funcs.push(visit(&mut converter, predicate)?); + } + let predicate_func = predicate_funcs + .into_iter() + .reduce(|mut lhs, mut rhs| { + Box::new(move |batch| { + let left = lhs(batch.clone())?; + let right = rhs(batch)?; + and_kleene(&left, &right) + }) + }) + .ok_or_else(|| invalid_data!("Row filter requires at least one predicate"))?; + + // After collecting required leaf column indices used in the predicates, + // creates the projection mask for the Arrow predicate. + let projection_mask = ProjectionMask::leaves(parquet_schema, column_indices); + Ok(Box::new(ArrowPredicateFn::new( + projection_mask, + predicate_func, + ))) } pub(super) fn get_selected_row_group_indices( diff --git a/crates/iceberg/src/arrow/reader/runtime_predicate.rs b/crates/iceberg/src/arrow/reader/runtime_predicate.rs new file mode 100644 index 0000000000..1746911483 --- /dev/null +++ b/crates/iceberg/src/arrow/reader/runtime_predicate.rs @@ -0,0 +1,595 @@ +// 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. + +//! Execution-time predicates published to the Arrow reader. + +use std::sync::{Arc, Mutex}; + +use arrow_schema::SchemaRef as ArrowSchemaRef; +use parquet::arrow::arrow_reader::RowSelection; +use parquet::schema::types::SchemaDescriptor; + +use super::ArrowReader; +use crate::arrow::record_batch_transformer::column_needs_type_promotion; +use crate::arrow::type_to_arrow_type; +use crate::expr::{Bind, BoundPredicate, Predicate}; +use crate::spec::{Schema, SchemaRef}; +use crate::{Error, ErrorKind, Result}; + +/// An immutable view of a runtime predicate at one point in execution. +#[derive(Debug)] +pub struct RuntimePredicateSnapshot { + predicate: Option, + generation: u64, +} + +impl RuntimePredicateSnapshot { + /// Creates a runtime predicate snapshot. A `None` predicate adds no + /// restriction to tasks that start with it. It does not undo pruning + /// already applied. + pub fn new(predicate: Option, generation: u64) -> Self { + Self { + predicate, + generation, + } + } + + /// Returns the publication generation this predicate belongs to. + pub fn generation(&self) -> u64 { + self.generation + } + + /// Consumes the snapshot, returning its predicate without cloning it. + pub fn into_predicate(self) -> Option { + self.predicate + } +} + +/// Supplies execution-time predicates, such as a completed join build's key +/// range or a TopK/MIN/MAX bound, to the Arrow reader. +/// +/// Contract: +/// * Data skipped under a published predicate must never be required by a +/// later publication. Once published, a predicate stays safe to combine with +/// the planned predicate for every row read after it, including by tasks still +/// using that generation. A later publication may therefore tighten it, but +/// not loosen it. The one exception is `None`, which stops runtime pruning +/// for tasks that start afterwards. It does not revoke predicates already +/// captured by earlier tasks or restore data they skipped. +/// * Generations increase monotonically on every change, are never reused, +/// and become visible only after their predicate. +/// * Provider, snapshot, binding, planning and page-pruning failures are +/// advisory: the reader skips the runtime predicate and does not retry them. +/// Errors while decoding its columns or evaluating it on rows are not +/// skipped. +/// +/// Pass a provider to `ArrowReaderBuilder::with_runtime_predicate_provider`. +/// This is the reader layer only: `TableScan::to_arrow` does not accept a +/// provider yet, so engines that plan through `TableScan` need that wiring as +/// a follow-up. +/// +/// # Example +/// +/// A provider that publishes an upper bound on `id`. The predicate and its +/// generation change together under a lock, and the cheap generation is +/// stored last so it is never visible before its predicate. +/// +/// ``` +/// use std::sync::Mutex; +/// use std::sync::atomic::{AtomicU64, Ordering}; +/// +/// use iceberg::Result; +/// use iceberg::arrow::{RuntimePredicateProvider, RuntimePredicateSnapshot}; +/// use iceberg::expr::{Predicate, Reference}; +/// use iceberg::spec::Datum; +/// +/// struct UpperBound { +/// published: Mutex<(u64, Option)>, +/// generation: AtomicU64, +/// } +/// +/// impl UpperBound { +/// /// Publishes a bound that must be no looser than the previous one. +/// fn publish(&self, bound: i64) { +/// let mut published = self.published.lock().unwrap(); +/// let generation = published.0 + 1; +/// let predicate = Reference::new("id").less_than_or_equal_to(Datum::long(bound)); +/// *published = (generation, Some(predicate)); +/// self.generation.store(generation, Ordering::Release); +/// } +/// } +/// +/// impl RuntimePredicateProvider for UpperBound { +/// fn generation(&self) -> u64 { +/// self.generation.load(Ordering::Acquire) +/// } +/// +/// fn snapshot(&self) -> Result { +/// let (generation, predicate) = self.published.lock().unwrap().clone(); +/// Ok(RuntimePredicateSnapshot::new(predicate, generation)) +/// } +/// } +/// ``` +pub trait RuntimePredicateProvider: Send + Sync { + /// Returns the current publication generation. Called for every task, so + /// keep it cheap, for example an atomic load with acquire ordering. + fn generation(&self) -> u64; + + /// Returns the current runtime predicate snapshot. The predicate and + /// generation must be a self-consistent pair, and the generation must not + /// be older than one `generation()` has already returned; the reader + /// ignores an older snapshot. + /// + /// Implementations should use in-memory, non-blocking work because refresh + /// serializes data-file tasks while taking the snapshot and binding it. + fn snapshot(&self) -> Result; +} + +/// Retains one bound predicate across a scan's tasks, keyed by generation, +/// schema and case policy. A different schema or case policy replaces the +/// binding; bindings for multiple schemas are not retained simultaneously. +/// Failures are cached as `None` for that key. +pub(super) struct RuntimePredicates { + provider: Arc, + cached: Mutex>, +} + +struct CachedPredicate { + generation: u64, + schema: SchemaRef, + case_sensitive: bool, + predicate: Option>, +} + +impl CachedPredicate { + fn matches(&self, generation: u64, schema: &SchemaRef, case_sensitive: bool) -> bool { + // A snapshot can observe a publication newer than the cheap check. + self.generation >= generation + && self.case_sensitive == case_sensitive + && (Arc::ptr_eq(&self.schema, schema) || *self.schema == **schema) + } +} + +impl RuntimePredicates { + pub(super) fn new(provider: Arc) -> Self { + Self { + provider, + cached: Mutex::new(None), + } + } + + /// Returns the current predicate bound to `schema`, with NOT pushed to the + /// leaves because statistics evaluators cannot negate a "might match". + pub(super) fn current( + &self, + schema: &SchemaRef, + case_sensitive: bool, + data_file_path: &str, + ) -> Option> { + // Hold the lock for the whole refresh so concurrent misses wait for it + // instead of repeating the snapshot and binding, and a slow refresh + // cannot overwrite a newer one. Read the generation after locking so + // waiters reuse a completed refresh, including a failed one. + let mut cache = self.cached.lock().unwrap(); + let generation = self.provider.generation(); + if let Some(cached) = cache.as_ref() + && cached.matches(generation, schema, case_sensitive) + { + return cached.predicate.clone(); + } + + // Never replace a newer cached generation, even if the binding context + // changes or the provider returns an outdated snapshot. Such a snapshot + // is advisory and is cached as a failure at the observed generation. + let generation = cache + .as_ref() + .map_or(generation, |cached| generation.max(cached.generation)); + let (generation, predicate) = match self.provider.snapshot() { + Ok(snapshot) if snapshot.generation() < generation => { + tracing::debug!("Skipping stale runtime predicate for {data_file_path}"); + (generation, None) + } + Ok(snapshot) => { + let generation = snapshot.generation(); + let bound = snapshot + .into_predicate() + .map(|predicate| predicate.rewrite_not().bind(schema.clone(), case_sensitive)) + .transpose(); + match bound { + Ok(bound) => (generation, bound.map(Arc::new)), + Err(error) => { + tracing::debug!("Skipping runtime predicate for {data_file_path}: {error}"); + (generation, None) + } + } + } + Err(error) => { + tracing::debug!("Skipping runtime predicate for {data_file_path}: {error}"); + (generation, None) + } + }; + let result = predicate.clone(); + *cache = Some(CachedPredicate { + generation, + schema: schema.clone(), + case_sensitive, + predicate, + }); + result + } +} + +/// Fails if `predicate` references a column the file lacks (its default is +/// invisible to physical filters) or stores with a promoted type (a literal +/// cast before promotion can round or overflow and reject valid rows). +pub(super) fn check_runtime_predicate_columns( + predicate: &BoundPredicate, + parquet_schema: &SchemaDescriptor, + arrow_schema: &ArrowSchemaRef, + table_schema: &Schema, + use_position_fallback: bool, +) -> Result<()> { + let (field_ids, field_id_map) = ArrowReader::build_field_id_set_and_map( + parquet_schema, + arrow_schema, + predicate, + use_position_fallback, + )?; + // Arrow and Parquet use the same leaf-column order, including for nested + // fields. Resolve by column index so name mapping and position fallback + // follow the same field resolution as predicate planning. + let mut source_fields = Vec::new(); + arrow_schema.fields().filter_leaves(|_, field| { + source_fields.push(field.clone()); + false + }); + for field_id in field_ids { + let unusable = |reason: &str| { + Error::new( + ErrorKind::FeatureUnsupported, + format!("Runtime predicate field {field_id} {reason}"), + ) + }; + let column = *field_id_map + .get(&field_id) + .ok_or_else(|| unusable("is not stored in this file"))?; + let field = table_schema + .field_by_id(field_id) + .ok_or_else(|| unusable("is not in the table schema"))?; + field + .field_type + .as_primitive_type() + .ok_or_else(|| unusable("is not a primitive column"))?; + let source_field = source_fields + .get(column) + .ok_or_else(|| unusable("has no Arrow leaf column"))?; + let target_type = type_to_arrow_type(&field.field_type)?; + if column_needs_type_promotion(source_field.data_type(), &target_type) { + return Err(unusable("is stored with a promoted physical type")); + } + } + Ok(()) +} + +/// Intersects two ascending row-group index lists in linear time. +pub(super) fn intersect_sorted(left: &[usize], right: &[usize]) -> Vec { + let (mut left, mut right) = (left.iter().peekable(), right.iter().peekable()); + let mut intersection = Vec::new(); + while let (Some(&&l), Some(&&r)) = (left.peek(), right.peek()) { + match l.cmp(&r) { + std::cmp::Ordering::Less => { + left.next(); + } + std::cmp::Ordering::Greater => { + right.next(); + } + std::cmp::Ordering::Equal => { + intersection.push(l); + left.next(); + right.next(); + } + } + } + intersection +} + +/// Intersects a predicate's page-index selection with `current`, where `None` +/// selects every row. An advisory predicate's failure leaves `current` as is. +pub(super) fn intersect_page_selection( + current: Option, + selection: Result>, + advisory: bool, + data_file_path: &str, +) -> Result> { + let selection = match selection { + Ok(selection) => selection, + Err(error) if advisory => { + tracing::debug!("Skipping runtime page pruning for {data_file_path}: {error}"); + None + } + Err(error) => return Err(error), + }; + Ok(match (current, selection) { + (Some(current), Some(selection)) => Some(current.intersection(&selection)), + (current, selection) => current.or(selection), + }) +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicU64, Ordering}; + use std::sync::mpsc::{self, Receiver, Sender}; + use std::sync::{Barrier, OnceLock, Weak}; + use std::time::Duration; + + use super::*; + use crate::expr::Reference; + use crate::spec::{Datum, NestedField, PrimitiveType, Type}; + + /// Holds the first snapshot open so concurrent callers overlap a refresh. + struct GatedProvider { + generation: AtomicU64, + snapshots: AtomicU64, + started: Sender<()>, + release: Mutex>, + fail_first_generation: bool, + predicates: OnceLock>, + } + + impl RuntimePredicateProvider for GatedProvider { + fn generation(&self) -> u64 { + self.generation.load(Ordering::Acquire) + } + + fn snapshot(&self) -> Result { + let generation = self.generation(); + let first = self.snapshots.fetch_add(1, Ordering::Relaxed) == 0; + self.started.send(()).unwrap(); + if first { + // The cache stays locked for the whole refresh, so concurrent + // misses wait for this snapshot instead of starting their own. + let predicates = self.predicates.get().unwrap().upgrade().unwrap(); + assert!(predicates.cached.try_lock().is_err()); + self.release + .lock() + .unwrap() + .recv_timeout(Duration::from_secs(10)) + .unwrap(); + } + if self.fail_first_generation && generation == 1 { + return Err(Error::new(ErrorKind::Unexpected, "snapshot failed")); + } + Ok(RuntimePredicateSnapshot::new( + Some( + Reference::new("id") + .greater_than_or_equal_to(Datum::long((generation * 100) as i64)), + ), + generation, + )) + } + } + + fn concurrent_refresh(fail_first_generation: bool, publish_during_snapshot: bool) { + const CALLERS: usize = 8; + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), + ]) + .build() + .unwrap(), + ); + let (started_tx, started_rx) = mpsc::channel(); + let (release_tx, release_rx) = mpsc::channel(); + let provider = Arc::new(GatedProvider { + generation: AtomicU64::new(1), + snapshots: AtomicU64::new(0), + started: started_tx, + release: Mutex::new(release_rx), + fail_first_generation, + predicates: OnceLock::new(), + }); + let predicates = Arc::new(RuntimePredicates::new(provider.clone())); + provider + .predicates + .set(Arc::downgrade(&predicates)) + .unwrap(); + let barrier = Arc::new(Barrier::new(CALLERS + 1)); + let callers: Vec<_> = (0..CALLERS) + .map(|_| { + let predicates = predicates.clone(); + let schema = schema.clone(); + let barrier = barrier.clone(); + std::thread::spawn(move || { + barrier.wait(); + predicates.current(&schema, false, "concurrent.parquet") + }) + }) + .collect(); + barrier.wait(); + started_rx.recv_timeout(Duration::from_secs(10)).unwrap(); + if publish_during_snapshot { + provider.generation.store(2, Ordering::Release); + } + release_tx.send(()).unwrap(); + let results: Vec<_> = callers.into_iter().map(|t| t.join().unwrap()).collect(); + + let current = predicates.current(&schema, false, "later.parquet"); + let expected_generation = if publish_during_snapshot { 2 } else { 1 }; + assert_eq!( + predicates + .cached + .lock() + .unwrap() + .as_ref() + .unwrap() + .generation, + expected_generation + ); + assert_eq!( + provider.snapshots.load(Ordering::Relaxed), + expected_generation, + "one snapshot and binding attempt per generation" + ); + if fail_first_generation { + assert!(results.iter().all(Option::is_none)); + assert!(current.is_none()); + // A failed generation is retained, but does not suppress a later + // successful publication. + provider.generation.store(2, Ordering::Release); + let recovered = predicates + .current(&schema, false, "recovered.parquet") + .unwrap(); + assert!(Arc::ptr_eq( + &recovered, + &predicates + .current(&schema, false, "reused.parquet") + .unwrap() + )); + assert_eq!(provider.snapshots.load(Ordering::Relaxed), 2); + } else { + let current = current.unwrap(); + let reused = results + .iter() + .filter(|result| Arc::ptr_eq(result.as_ref().unwrap(), ¤t)) + .count(); + assert_eq!(reused, CALLERS - usize::from(publish_during_snapshot)); + // Arc identity verifies that waiters reuse the bound predicate, + // rather than just deduplicating snapshots and binding separately. + assert!(Arc::ptr_eq( + ¤t, + &predicates + .current(&schema, false, "reused.parquet") + .unwrap() + )); + } + } + + #[test] + fn intersect_sorted_keeps_common_indices_in_order() { + assert_eq!( + intersect_sorted(&[0, 2, 3, 7, 9], &[1, 2, 3, 8, 9, 10]), + vec![2, 3, 9] + ); + assert_eq!(intersect_sorted(&[1, 2], &[]), Vec::::new()); + assert_eq!(intersect_sorted(&[], &[1, 2]), Vec::::new()); + assert_eq!(intersect_sorted(&[4, 5], &[4, 5]), vec![4, 5]); + assert_eq!(intersect_sorted(&[0, 1], &[2, 3]), Vec::::new()); + } + + #[test] + fn concurrent_runtime_predicate_refresh_snapshots_and_binds_once() { + concurrent_refresh(false, false); + } + + #[test] + fn concurrent_runtime_predicate_publication_does_not_regress_cache() { + concurrent_refresh(false, true); + } + + #[test] + fn concurrent_runtime_predicate_failure_is_cached_until_next_generation() { + concurrent_refresh(true, false); + } + + #[test] + fn runtime_predicate_stale_snapshot_is_cached_without_regression() { + struct Provider { + generation: AtomicU64, + snapshot_generation: AtomicU64, + snapshots: AtomicU64, + } + impl RuntimePredicateProvider for Provider { + fn generation(&self) -> u64 { + self.generation.load(Ordering::Acquire) + } + + fn snapshot(&self) -> Result { + self.snapshots.fetch_add(1, Ordering::Relaxed); + Ok(RuntimePredicateSnapshot::new( + Some(Reference::new("id").greater_than(Datum::long(0))), + self.snapshot_generation.load(Ordering::Acquire), + )) + } + } + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), + ]) + .build() + .unwrap(), + ); + let provider = Arc::new(Provider { + generation: AtomicU64::new(1), + snapshot_generation: AtomicU64::new(2), + snapshots: AtomicU64::new(0), + }); + let predicates = RuntimePredicates::new(provider.clone()); + let bound = predicates.current(&schema, false, "ahead.parquet").unwrap(); + // Snapshot publication can precede the cheap generation publication. + assert!(Arc::ptr_eq( + &bound, + &predicates + .current(&schema, false, "reused.parquet") + .unwrap() + )); + assert_eq!(provider.snapshots.load(Ordering::Relaxed), 1); + + // Rebinding under another case policy must not lower the high-water + // generation even if a provider supplies an outdated snapshot. + provider.snapshot_generation.store(1, Ordering::Release); + assert!(predicates.current(&schema, true, "stale.parquet").is_none()); + assert_eq!( + predicates + .cached + .lock() + .unwrap() + .as_ref() + .unwrap() + .generation, + 2 + ); + assert!( + predicates + .current(&schema, true, "cached-failure.parquet") + .is_none() + ); + assert_eq!(provider.snapshots.load(Ordering::Relaxed), 2); + + // A stale snapshot following a newer cheap check is also a cached + // failure, rather than work to retry for every data-file task. + provider.generation.store(3, Ordering::Release); + assert!(predicates.current(&schema, true, "newer.parquet").is_none()); + assert!( + predicates + .current(&schema, true, "no-storm.parquet") + .is_none() + ); + assert_eq!(provider.snapshots.load(Ordering::Relaxed), 3); + assert_eq!( + predicates + .cached + .lock() + .unwrap() + .as_ref() + .unwrap() + .generation, + 3 + ); + } +} diff --git a/crates/iceberg/src/arrow/reader/runtime_predicate_tests.rs b/crates/iceberg/src/arrow/reader/runtime_predicate_tests.rs new file mode 100644 index 0000000000..0b9c20f00d --- /dev/null +++ b/crates/iceberg/src/arrow/reader/runtime_predicate_tests.rs @@ -0,0 +1,1327 @@ +// 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::collections::HashMap; +use std::fs::File; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use arrow_array::{ + ArrayRef, Decimal128Array, Float32Array, Float64Array, Int32Array, Int64Array, RecordBatch, + StringArray, +}; +use arrow_schema::{DataType, Field, Schema as ArrowSchema}; +use futures::{StreamExt, TryStreamExt}; +use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY}; +use parquet::basic::Compression; +use parquet::file::properties::WriterProperties; +use tempfile::TempDir; + +use super::{ArrowReaderBuilder, RuntimePredicateProvider, RuntimePredicateSnapshot}; +use crate::arrow::ScanMetrics; +use crate::expr::{Bind, Predicate, Reference}; +use crate::io::FileIO; +use crate::scan::{ + ArrowRecordBatchStream, FileScanTask, FileScanTaskDeleteFile, FileScanTaskStream, +}; +use crate::spec::{ + DataContentType, DataFileFormat, Datum, NestedField, PrimitiveType, Schema, SchemaRef, Type, +}; +use crate::{Result, Runtime}; + +#[derive(Debug)] +struct FixedRuntimePredicate { + predicate: Predicate, + generation: u64, + snapshots: AtomicU64, +} + +impl FixedRuntimePredicate { + fn new(predicate: Predicate) -> Self { + Self { + predicate, + generation: 1, + snapshots: AtomicU64::new(0), + } + } + + fn snapshots(&self) -> u64 { + self.snapshots.load(Ordering::Relaxed) + } +} + +impl RuntimePredicateProvider for FixedRuntimePredicate { + fn generation(&self) -> u64 { + self.generation + } + + fn snapshot(&self) -> Result { + self.snapshots.fetch_add(1, Ordering::Relaxed); + Ok(RuntimePredicateSnapshot::new( + Some(self.predicate.clone()), + self.generation, + )) + } +} + +fn iceberg_schema() -> SchemaRef { + Arc::new( + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::required(2, "payload", Type::Primitive(PrimitiveType::String)).into(), + ]) + .build() + .unwrap(), + ) +} + +fn write_three_row_group_file(dir: &str, name: &str) -> String { + write_row_group_file(dir, name, &[0, 100, 200]) +} + +fn write_row_group_file(dir: &str, name: &str, bases: &[i32]) -> String { + write_row_group_file_with_page_size(dir, name, bases, 1024) +} + +fn write_row_group_file_with_page_size( + dir: &str, + name: &str, + bases: &[i32], + page_rows: usize, +) -> String { + let id_field = Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])); + let payload_field = Field::new("payload", DataType::Utf8, false).with_metadata(HashMap::from( + [(PARQUET_FIELD_ID_META_KEY.to_string(), "2".to_string())], + )); + let arrow_schema = Arc::new(ArrowSchema::new(vec![id_field, payload_field])); + + let file_path = format!("{dir}/{name}"); + let file = File::create(&file_path).unwrap(); + let props = WriterProperties::builder() + .set_compression(Compression::UNCOMPRESSED) + .set_max_row_group_row_count(Some(4)) + .set_data_page_row_count_limit(page_rows) + .set_write_batch_size(page_rows) + .set_dictionary_enabled(false) + .build(); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&arrow_schema), Some(props)).unwrap(); + + for &base in bases { + let ids: Vec = (base..base + 4).collect(); + let payloads: Vec = (0..4) + .map(|row| { + (0..16_384) + .map(|offset| { + let value = ((base as usize) + row + offset * 17) % 94; + char::from_u32(33 + value as u32).unwrap() + }) + .collect() + }) + .collect(); + let columns: Vec = vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(payloads)), + ]; + let batch = RecordBatch::try_new(Arc::clone(&arrow_schema), columns).unwrap(); + writer.write(&batch).unwrap(); + writer.flush().unwrap(); + } + writer.close().unwrap(); + file_path +} + +fn scan_task( + file_path: String, + schema: SchemaRef, + predicate: Option, +) -> FileScanTask { + scan_task_with_deletes(file_path, schema, predicate, vec![]) +} + +fn scan_task_with_deletes( + file_path: String, + schema: SchemaRef, + predicate: Option, + deletes: Vec, +) -> FileScanTask { + scan_task_with_deletes_and_projection(file_path, schema, predicate, deletes, vec![1, 2]) +} + +fn scan_task_with_deletes_and_projection( + file_path: String, + schema: SchemaRef, + predicate: Option, + deletes: Vec, + project_field_ids: Vec, +) -> FileScanTask { + FileScanTask::builder() + .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len()) + .with_start(0) + .with_length(0) + .with_data_file_path(file_path) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(schema) + .with_project_field_ids(project_field_ids) + .with_predicate(predicate) + .with_deletes(deletes) + .with_case_sensitive(false) + .build() + .unwrap() +} + +fn ids(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + .collect::>() + }) + .collect() +} + +fn write_delete(path: &str, fields: Vec, arrays: Vec) { + let schema = Arc::new(ArrowSchema::new(fields)); + let batch = RecordBatch::try_new(Arc::clone(&schema), arrays).unwrap(); + let mut writer = ArrowWriter::try_new(File::create(path).unwrap(), schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); +} + +fn field(name: &str, data_type: DataType, id: i32) -> Field { + Field::new(name, data_type, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + id.to_string(), + )])) +} + +/// Writes one four-row group per base with small payloads. +fn write_groups(path: &str, bases: &[i32], field_ids: bool, key: Option<&[u8]>, bloom: bool) { + let metadata = |id: &str| { + if field_ids { + HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), id.to_string())]) + } else { + HashMap::new() + } + }; + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(metadata("1")), + Field::new("payload", DataType::Utf8, false).with_metadata(metadata("2")), + ])); + let mut properties = WriterProperties::builder() + .set_compression(Compression::UNCOMPRESSED) + .set_max_row_group_row_count(Some(4)) + .set_bloom_filter_enabled(bloom); + if let Some(key) = key { + properties = properties.with_file_encryption_properties( + parquet::encryption::encrypt::FileEncryptionProperties::builder(key.to_vec()) + .build() + .unwrap(), + ); + } + let mut writer = ArrowWriter::try_new( + File::create(path).unwrap(), + Arc::clone(&schema), + Some(properties.build()), + ) + .unwrap(); + for &base in bases { + let ids: Vec = (base..base + 4).collect(); + let payloads: Vec = ids.iter().map(|id| format!("payload-{id}")).collect(); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(payloads)), + ]) + .unwrap(); + writer.write(&batch).unwrap(); + writer.flush().unwrap(); + } + writer.close().unwrap(); +} + +#[derive(Default)] +struct FailedProvider { + snapshots: AtomicU64, +} + +impl RuntimePredicateProvider for FailedProvider { + fn generation(&self) -> u64 { + 0 + } + + fn snapshot(&self) -> Result { + self.snapshots.fetch_add(1, Ordering::Relaxed); + Err(crate::Error::new( + crate::ErrorKind::Unexpected, + "publication failed", + )) + } +} + +async fn execute_tasks( + tasks: Vec, + provider: Option>, + bloom_filter: bool, +) -> (Vec, ScanMetrics) { + execute_tasks_with_concurrency(tasks, provider, bloom_filter, 1).await +} + +async fn execute_tasks_with_concurrency( + tasks: Vec, + provider: Option>, + bloom_filter: bool, + concurrency: usize, +) -> (Vec, ScanMetrics) { + let mut builder = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()) + .with_data_file_concurrency_limit(concurrency) + .with_row_selection_enabled(true) + .with_bloom_filter_enabled(bloom_filter); + if let Some(provider) = provider { + builder = builder.with_runtime_predicate_provider(provider); + } + let tasks = Box::pin(futures::stream::iter(tasks.into_iter().map(Ok))) as FileScanTaskStream; + let scan = builder.build().read(tasks).unwrap(); + let metrics = scan.metrics().clone(); + let batches = scan.stream().try_collect().await.unwrap(); + (batches, metrics) +} + +async fn execute( + task: FileScanTask, + provider: Option>, +) -> (Vec, ScanMetrics) { + execute_tasks(vec![task], provider, false).await +} + +fn all_ids() -> Vec { + [0, 100, 200] + .into_iter() + .flat_map(|base| base..base + 4) + .collect() +} + +#[tokio::test] +async fn runtime_predicate_prunes_row_groups_and_bytes() { + let temp = TempDir::new().unwrap(); + let path = write_three_row_group_file(temp.path().to_str().unwrap(), "prune.parquet"); + let task = scan_task(path, iceberg_schema(), None); + let (_, baseline) = execute(task.clone(), None).await; + + let provider = Arc::new(FixedRuntimePredicate::new( + Reference::new("id") + .greater_than_or_equal_to(Datum::int(100)) + .and(Reference::new("id").less_than_or_equal_to(Datum::int(103))), + )); + let (batches, metrics) = execute(task, Some(provider.clone())).await; + assert_eq!(ids(&batches), vec![100, 101, 102, 103]); + assert_eq!(provider.snapshots(), 1); + assert!( + metrics.bytes_read() < baseline.bytes_read(), + "runtime={} baseline={}", + metrics.bytes_read(), + baseline.bytes_read() + ); +} + +#[tokio::test] +async fn runtime_predicate_is_anded_with_task_predicate() { + let temp = TempDir::new().unwrap(); + let path = write_three_row_group_file(temp.path().to_str().unwrap(), "and.parquet"); + let planned = Reference::new("id") + .greater_than_or_equal_to(Datum::int(100)) + .bind(iceberg_schema(), false) + .unwrap(); + let provider = Arc::new(FixedRuntimePredicate::new( + Reference::new("id").less_than_or_equal_to(Datum::int(103)), + )); + let (batches, _) = execute( + scan_task(path, iceberg_schema(), Some(planned)), + Some(provider), + ) + .await; + assert_eq!(ids(&batches), vec![100, 101, 102, 103]); +} + +#[tokio::test] +async fn runtime_predicate_failures_keep_the_planned_filter() { + let temp = TempDir::new().unwrap(); + let path = write_three_row_group_file(temp.path().to_str().unwrap(), "fail.parquet"); + let planned = || { + Some( + Reference::new("id") + .equal_to(Datum::int(101)) + .bind(iceberg_schema(), false) + .unwrap(), + ) + }; + // A failing provider, and a predicate that cannot be bound. + for provider in [ + Arc::new(FailedProvider::default()) as Arc, + Arc::new(FixedRuntimePredicate::new( + Reference::new("missing").equal_to(Datum::int(1)), + )), + ] { + let (batches, _) = execute( + scan_task(path.clone(), iceberg_schema(), planned()), + Some(provider), + ) + .await; + assert_eq!(ids(&batches), vec![101]); + let (batches, _) = execute(scan_task(path.clone(), iceberg_schema(), None), None).await; + assert_eq!(ids(&batches), all_ids()); + } +} + +#[tokio::test] +async fn runtime_not_predicates_keep_matching_rows() { + let temp = TempDir::new().unwrap(); + let path = write_three_row_group_file(temp.path().to_str().unwrap(), "not.parquet"); + // RG1 holds 100..=103, so NOT(id < 102) must keep 102 and 103. + let provider = Arc::new(FixedRuntimePredicate::new( + !Reference::new("id").less_than(Datum::int(102)), + )); + let (batches, _) = execute(scan_task(path, iceberg_schema(), None), Some(provider)).await; + assert_eq!(ids(&batches), vec![102, 103, 200, 201, 202, 203]); +} + +#[tokio::test] +async fn runtime_predicate_on_column_missing_from_file_is_ignored() { + use crate::spec::Literal; + + let temp = TempDir::new().unwrap(); + let path = write_three_row_group_file(temp.path().to_str().unwrap(), "evolved.parquet"); + // `added` was added after this file was written. Its value in this file is + // the initial default (7) or null, which physical filters never see. + for default in [Some(Literal::int(7)), None] { + let mut added = NestedField::optional(3, "added", Type::Primitive(PrimitiveType::Int)); + if let Some(default) = default.clone() { + added = added.with_initial_default(default); + } + let schema: SchemaRef = Arc::new( + Schema::builder() + .with_schema_id(2) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::required(2, "payload", Type::Primitive(PrimitiveType::String)) + .into(), + added.into(), + ]) + .build() + .unwrap(), + ); + for predicate in [ + Reference::new("added").is_null(), + Reference::new("added").greater_than_or_equal_to(Datum::int(5)), + Reference::new("added").less_than(Datum::int(5)), + ] { + let provider = Arc::new(FixedRuntimePredicate::new(predicate.clone())); + let (batches, _) = execute( + scan_task(path.clone(), Arc::clone(&schema), None), + Some(provider), + ) + .await; + assert_eq!( + ids(&batches), + all_ids(), + "{predicate} with default {default:?}" + ); + } + } +} + +#[tokio::test] +async fn runtime_predicate_on_promoted_column_is_ignored() { + let temp = TempDir::new().unwrap(); + // The file stores `id` as INT; the table has since promoted it to BIGINT. + let path = write_three_row_group_file(temp.path().to_str().unwrap(), "promoted.parquet"); + let schema: SchemaRef = Arc::new( + Schema::builder() + .with_schema_id(2) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), + NestedField::required(2, "payload", Type::Primitive(PrimitiveType::String)).into(), + ]) + .build() + .unwrap(), + ); + // Above i32::MAX: a literal cast down to the physical INT type would overflow. + let provider = Arc::new(FixedRuntimePredicate::new( + Reference::new("id").less_than(Datum::long(i64::from(i32::MAX) + 1)), + )); + let (batches, _) = execute(scan_task(path, schema, None), Some(provider)).await; + let values: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert_eq!( + values, + all_ids().into_iter().map(i64::from).collect::>() + ); +} + +#[tokio::test] +async fn runtime_predicate_preserves_position_and_equality_deletes() { + const FIELD_ID_POSITIONAL_DELETE_FILE_PATH: i32 = 2_147_483_546; + const FIELD_ID_POSITIONAL_DELETE_POS: i32 = 2_147_483_545; + + let temp = TempDir::new().unwrap(); + let dir = temp.path().to_str().unwrap(); + let data_path = write_three_row_group_file(dir, "data.parquet"); + let position_path = format!("{dir}/positions.parquet"); + let equality_path = format!("{dir}/equalities.parquet"); + // Position 5 is id=101 in the second row group. If positions were + // renumbered after the first group is pruned, the wrong row would go. + write_delete( + &position_path, + vec![ + field( + "file_path", + DataType::Utf8, + FIELD_ID_POSITIONAL_DELETE_FILE_PATH, + ), + field("pos", DataType::Int64, FIELD_ID_POSITIONAL_DELETE_POS), + ], + vec![ + Arc::new(StringArray::from(vec![data_path.as_str()])), + Arc::new(Int64Array::from(vec![5])), + ], + ); + write_delete(&equality_path, vec![field("id", DataType::Int32, 1)], vec![ + Arc::new(Int32Array::from(vec![102])), + ]); + let delete = |path: String, file_type, equality_ids| { + FileScanTaskDeleteFile::builder() + .with_file_size_in_bytes(std::fs::metadata(&path).unwrap().len()) + .with_file_path(path) + .with_file_type(file_type) + .with_file_format(DataFileFormat::Parquet) + .with_partition_spec_id(0) + .with_equality_ids(equality_ids) + .build() + }; + let position = delete(position_path, DataContentType::PositionDeletes, None); + let equality = delete( + equality_path, + DataContentType::EqualityDeletes, + Some(vec![1]), + ); + for (deletes, expected) in [ + (vec![position.clone()], vec![100, 102, 103]), + (vec![equality.clone()], vec![100, 101, 103]), + (vec![position, equality], vec![100, 103]), + ] { + let task = scan_task_with_deletes(data_path.clone(), iceberg_schema(), None, deletes); + let (baseline, baseline_metrics) = execute(task.clone(), None).await; + let provider = Arc::new(FixedRuntimePredicate::new( + Reference::new("id") + .greater_than_or_equal_to(Datum::int(100)) + .and(Reference::new("id").less_than_or_equal_to(Datum::int(103))), + )); + let (runtime, metrics) = execute(task, Some(provider)).await; + assert_eq!(ids(&runtime), expected); + assert_eq!( + ids(&baseline) + .into_iter() + .filter(|id| (100..=103).contains(id)) + .collect::>(), + expected + ); + assert!(metrics.bytes_read() < baseline_metrics.bytes_read()); + } +} + +#[tokio::test] +async fn runtime_predicate_respects_task_byte_range() { + use parquet::file::reader::{FileReader, SerializedFileReader}; + + let temp = TempDir::new().unwrap(); + let path = write_three_row_group_file(temp.path().to_str().unwrap(), "split.parquet"); + let parquet = SerializedFileReader::new(File::open(&path).unwrap()).unwrap(); + let start = 4 + parquet.metadata().row_group(0).compressed_size() as u64; + let file_size = std::fs::metadata(&path).unwrap().len(); + let task = FileScanTask::builder() + .with_file_size_in_bytes(file_size) + .with_start(start) + .with_length(file_size - start) + .with_data_file_path(path) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(iceberg_schema()) + .with_project_field_ids(vec![1, 2]) + .with_case_sensitive(false) + .build() + .unwrap(); + // RG0 belongs to another split, so even a predicate matching it reads nothing there. + let provider = Arc::new(FixedRuntimePredicate::new( + Reference::new("id") + .less_than(Datum::int(4)) + .or(Reference::new("id").greater_than_or_equal_to(Datum::int(200))), + )); + let (batches, _) = execute(task, Some(provider)).await; + assert_eq!(ids(&batches), vec![200, 201, 202, 203]); +} + +#[tokio::test] +async fn runtime_predicate_with_bloom_filters_keeps_planned_equality() { + let temp = TempDir::new().unwrap(); + let path = format!("{}/bloom.parquet", temp.path().to_str().unwrap()); + write_groups(&path, &[0, 100, 200], true, None, true); + let planned = Reference::new("id") + .equal_to(Datum::int(201)) + .bind(iceberg_schema(), false) + .unwrap(); + let provider = Arc::new(FixedRuntimePredicate::new( + Reference::new("id").greater_than_or_equal_to(Datum::int(100)), + )); + let (batches, _) = execute_tasks( + vec![scan_task(path, iceberg_schema(), Some(planned))], + Some(provider), + true, + ) + .await; + assert_eq!(ids(&batches), vec![201]); +} + +/// A provider whose predicate changes when [`Self::publish`] is called. +struct ChangingRuntimePredicate { + generation: AtomicU64, + predicate: std::sync::Mutex>, + snapshots: AtomicU64, +} + +impl ChangingRuntimePredicate { + fn new(predicate: Option, generation: u64) -> Self { + Self { + generation: AtomicU64::new(generation), + predicate: std::sync::Mutex::new(predicate), + snapshots: AtomicU64::new(0), + } + } + + fn publish(&self, predicate: Option, generation: u64) { + *self.predicate.lock().unwrap() = predicate; + self.generation.store(generation, Ordering::Release); + } + + fn snapshots(&self) -> u64 { + self.snapshots.load(Ordering::Relaxed) + } +} + +impl RuntimePredicateProvider for ChangingRuntimePredicate { + fn generation(&self) -> u64 { + self.generation.load(Ordering::Acquire) + } + + fn snapshot(&self) -> Result { + self.snapshots.fetch_add(1, Ordering::Relaxed); + let predicate = self.predicate.lock().unwrap(); + Ok(RuntimePredicateSnapshot::new( + predicate.clone(), + self.generation.load(Ordering::Acquire), + )) + } +} + +#[tokio::test] +async fn runtime_predicate_is_bound_once_per_generation_across_tasks() { + let temp = TempDir::new().unwrap(); + let dir = temp.path().to_str().unwrap(); + let tasks: Vec<_> = ["a.parquet", "b.parquet", "c.parquet"] + .into_iter() + .map(|name| { + scan_task( + write_three_row_group_file(dir, name), + iceberg_schema(), + None, + ) + }) + .collect(); + let provider = Arc::new(ChangingRuntimePredicate::new( + Some(Reference::new("id").greater_than_or_equal_to(Datum::int(100))), + 1, + )); + let scan = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()) + .with_data_file_concurrency_limit(1) + .with_runtime_predicate_provider(provider.clone()) + .build() + .read(Box::pin(futures::stream::iter(tasks.into_iter().map(Ok))) as FileScanTaskStream) + .unwrap(); + let mut stream = scan.stream(); + // The first task reads RG1 and RG2 under generation 1. + let mut batches = vec![ + stream.try_next().await.unwrap().unwrap(), + stream.try_next().await.unwrap().unwrap(), + ]; + assert_eq!(ids(&batches), vec![100, 101, 102, 103, 200, 201, 202, 203]); + assert_eq!(provider.snapshots(), 1); + // A tighter generation is picked up by the tasks that start afterwards and + // bound once for both of them. + provider.publish( + Some(Reference::new("id").greater_than_or_equal_to(Datum::int(200))), + 2, + ); + batches.extend(stream.try_collect::>().await.unwrap()); + let mut expected = vec![100, 101, 102, 103, 200, 201, 202, 203]; + for _ in 0..2 { + expected.extend([200, 201, 202, 203]); + } + assert_eq!(ids(&batches), expected); + assert_eq!(provider.snapshots(), 2); +} + +#[tokio::test] +async fn runtime_predicate_on_a_column_that_is_not_projected() { + let temp = TempDir::new().unwrap(); + let path = write_three_row_group_file(temp.path().to_str().unwrap(), "filter-only.parquet"); + // Only `payload` is projected; the runtime predicate filters on `id`. + let task = scan_task_with_deletes_and_projection(path, iceberg_schema(), None, vec![], vec![2]); + let (baseline, baseline_metrics) = execute(task.clone(), None).await; + let provider = Arc::new(FixedRuntimePredicate::new( + Reference::new("id").greater_than_or_equal_to(Datum::int(200)), + )); + let (batches, metrics) = execute(task, Some(provider)).await; + let rows = |batches: &[RecordBatch]| batches.iter().map(RecordBatch::num_rows).sum::(); + assert_eq!(rows(&baseline), 12); + assert_eq!(rows(&batches), 4); + assert!(batches.iter().all(|batch| batch.num_columns() == 1)); + assert!(metrics.bytes_read() < baseline_metrics.bytes_read()); +} + +#[tokio::test] +async fn runtime_predicate_resolves_files_without_field_ids() { + use crate::spec::{MappedField, NameMapping}; + + let temp = TempDir::new().unwrap(); + let path = format!("{}/no-ids.parquet", temp.path().to_str().unwrap()); + write_groups(&path, &[0, 100, 200], false, None, false); + let mapping = Arc::new(NameMapping::new(vec![ + MappedField::new(Some(1), vec!["id".to_string()], vec![]), + MappedField::new(Some(2), vec!["payload".to_string()], vec![]), + ])); + // A name mapping assigns ids by name; without one, ids follow column positions. + for name_mapping in [Some(mapping), None] { + let task = FileScanTask::builder() + .with_file_size_in_bytes(std::fs::metadata(&path).unwrap().len()) + .with_start(0) + .with_length(0) + .with_data_file_path(path.clone()) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(iceberg_schema()) + .with_project_field_ids(vec![1, 2]) + .with_name_mapping(name_mapping) + .with_case_sensitive(false) + .build() + .unwrap(); + let provider = Arc::new(FixedRuntimePredicate::new( + Reference::new("id").greater_than_or_equal_to(Datum::int(200)), + )); + let (batches, _) = execute(task, Some(provider)).await; + assert_eq!(ids(&batches), vec![200, 201, 202, 203]); + } +} + +#[tokio::test] +async fn runtime_predicate_prunes_pages_within_a_row_group() { + use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; + + const PAGE_ROWS: i32 = 1024; + let temp = TempDir::new().unwrap(); + let path = format!("{}/pages.parquet", temp.path().to_str().unwrap()); + // One row group of four pages. Pages hold many rows so the reader skips + // unselected pages instead of decoding them under a row mask. + let schema = Arc::new(ArrowSchema::new(vec![field("id", DataType::Int32, 1)])); + let props = WriterProperties::builder() + .set_compression(Compression::UNCOMPRESSED) + .set_dictionary_enabled(false) + .set_data_page_row_count_limit(PAGE_ROWS as usize) + .set_write_batch_size(PAGE_ROWS as usize) + .build(); + let mut writer = ArrowWriter::try_new( + File::create(&path).unwrap(), + Arc::clone(&schema), + Some(props), + ) + .unwrap(); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new( + Int32Array::from_iter_values(0..4 * PAGE_ROWS), + )]) + .unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .parse_and_finish(&File::open(&path).unwrap()) + .unwrap(); + assert_eq!(metadata.num_row_groups(), 1); + assert_eq!( + metadata.offset_index().unwrap()[0][0] + .page_locations() + .len(), + 4 + ); + + let read = |row_selection: bool| { + let task = scan_task_with_deletes_and_projection( + path.clone(), + iceberg_schema(), + None, + vec![], + vec![1], + ); + async move { + let provider = Arc::new(FixedRuntimePredicate::new( + Reference::new("id") + .greater_than_or_equal_to(Datum::int(2 * PAGE_ROWS)) + .and(Reference::new("id").less_than(Datum::int(3 * PAGE_ROWS))), + )); + let scan = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()) + .with_row_group_filtering_enabled(false) + .with_row_selection_enabled(row_selection) + // Keep the skipped page ranges separate; the default 1 MiB + // coalescing threshold would read them along with selected pages. + .with_range_coalesce_bytes(0) + .with_runtime_predicate_provider(provider) + .build() + .read(Box::pin(futures::stream::iter([Ok(task)])) as FileScanTaskStream) + .unwrap(); + let metrics = scan.metrics().clone(); + let batches: Vec = scan.stream().try_collect().await.unwrap(); + (ids(&batches), metrics.bytes_read()) + } + }; + let expected: Vec = (2 * PAGE_ROWS..3 * PAGE_ROWS).collect(); + // Without page selection the row filter decodes every page. + let (unpruned_ids, unpruned) = read(false).await; + let (pruned_ids, pruned) = read(true).await; + assert_eq!(unpruned_ids, expected); + assert_eq!(pruned_ids, expected); + assert!(pruned < unpruned, "pruned={pruned} unpruned={unpruned}"); +} + +#[tokio::test] +async fn runtime_predicate_that_cannot_be_planned_keeps_the_planned_filter() { + use arrow_array::StructArray; + + let temp = TempDir::new().unwrap(); + let path = format!("{}/nested.parquet", temp.path().to_str().unwrap()); + let x = field("x", DataType::Int32, 4); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + field("id", DataType::Int32, 1), + field("s", DataType::Struct(vec![x.clone()].into()), 3), + ])); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(4)) + .build(); + let mut writer = ArrowWriter::try_new( + File::create(&path).unwrap(), + Arc::clone(&arrow_schema), + Some(props), + ) + .unwrap(); + for base in [0, 100, 200] { + let s = StructArray::from(vec![( + Arc::new(x.clone()), + Arc::new(Int32Array::from(vec![-1; 4])) as ArrayRef, + )]); + let batch = RecordBatch::try_new(Arc::clone(&arrow_schema), vec![ + Arc::new(Int32Array::from((base..base + 4).collect::>())), + Arc::new(s), + ]) + .unwrap(); + writer.write(&batch).unwrap(); + writer.flush().unwrap(); + } + writer.close().unwrap(); + + let schema: SchemaRef = Arc::new( + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::required( + 3, + "s", + Type::Struct(crate::spec::StructType::new(vec![ + NestedField::required(4, "x", Type::Primitive(PrimitiveType::Int)).into(), + ])), + ) + .into(), + ]) + .build() + .unwrap(), + ); + let task = |predicate: Option| { + scan_task_with_deletes_and_projection( + path.clone(), + Arc::clone(&schema), + predicate, + vec![], + vec![1], + ) + }; + let planned = Reference::new("id") + .greater_than_or_equal_to(Datum::int(100)) + .bind(Arc::clone(&schema), false) + .unwrap(); + // The reader cannot build row filters on nested columns, so this passes the + // column checks but fails planning; applied, it would reject every row. + let provider = Arc::new(FixedRuntimePredicate::new( + Reference::new("s.x").greater_than_or_equal_to(Datum::int(0)), + )); + let (_, baseline) = execute(task(None), None).await; + let (batches, metrics) = execute(task(Some(planned)), Some(provider)).await; + assert_eq!(ids(&batches), vec![100, 101, 102, 103, 200, 201, 202, 203]); + // The planned predicate still prunes RG0. + assert!(metrics.bytes_read() < baseline.bytes_read()); +} + +#[test] +fn runtime_page_selection_failure_keeps_the_planned_selection() { + use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; + use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; + + use super::ArrowReader; + use super::runtime_predicate::intersect_page_selection; + + let temp = TempDir::new().unwrap(); + let path = write_row_group_file_with_page_size( + temp.path().to_str().unwrap(), + "pages.parquet", + &[0], + 1, + ); + let metadata = Arc::new( + ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .parse_and_finish(&File::open(&path).unwrap()) + .unwrap(), + ); + // The page-index evaluator rejects NOT, which the reader rewrites away + // before planning, so this yields a real evaluation error. + let predicate = (!Reference::new("id").less_than(Datum::int(2))) + .bind(iceberg_schema(), false) + .unwrap(); + let field_id_map = HashMap::from([(1, 0), (2, 1)]); + let selection = || { + ArrowReader::get_row_selection_for_filter_predicate( + &predicate, + &metadata, + &None, + &field_id_map, + &iceberg_schema(), + ) + }; + assert!(selection().is_err()); + + let planned = RowSelection::from(vec![RowSelector::skip(1), RowSelector::select(3)]); + let combined = + intersect_page_selection(Some(planned.clone()), selection(), true, &path).unwrap(); + assert_eq!(combined, Some(planned)); + assert!(intersect_page_selection(None, selection(), false, &path).is_err()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn runtime_predicate_concurrent_tasks_reuse_success_and_failure() { + let temp = TempDir::new().unwrap(); + let path = write_three_row_group_file(temp.path().to_str().unwrap(), "concurrent.parquet"); + let tasks: Vec<_> = (0..12) + .map(|_| scan_task(path.clone(), iceberg_schema(), None)) + .collect(); + let success = Arc::new(FixedRuntimePredicate::new( + Reference::new("id").greater_than_or_equal_to(Datum::int(200)), + )); + let absent = Arc::new(ChangingRuntimePredicate::new(None, 1)); + let failure = Arc::new(FailedProvider::default()); + let binding_failure = Arc::new(FixedRuntimePredicate::new( + Reference::new("missing").equal_to(Datum::int(0)), + )); + for (provider, expected) in [ + (success.clone() as Arc, vec![ + 200, 201, 202, 203, + ]), + ( + absent.clone() as Arc, + all_ids(), + ), + ( + failure.clone() as Arc, + all_ids(), + ), + ( + binding_failure.clone() as Arc, + all_ids(), + ), + ] { + let (batches, _) = + execute_tasks_with_concurrency(tasks.clone(), Some(provider), false, 4).await; + let mut actual = ids(&batches); + actual.sort_unstable(); + let mut expected = expected.repeat(tasks.len()); + expected.sort_unstable(); + assert_eq!(actual, expected); + } + assert_eq!(success.snapshots(), 1); + assert_eq!(absent.snapshots(), 1); + assert_eq!(failure.snapshots.load(Ordering::Relaxed), 1); + assert_eq!(binding_failure.snapshots(), 1); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn runtime_predicate_concurrent_tasks_pick_up_new_generation_without_regression() { + let temp = TempDir::new().unwrap(); + let path = write_three_row_group_file(temp.path().to_str().unwrap(), "generations.parquet"); + let task = scan_task(path, iceberg_schema(), None); + let provider = Arc::new(ChangingRuntimePredicate::new( + Some(Reference::new("id").greater_than_or_equal_to(Datum::int(100))), + 1, + )); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let next = task.clone(); + // Withhold the second wave until the four first-wave tasks have completed. + // This tests a publication boundary without depending on file-read timing. + let tasks = futures::stream::iter(vec![task.clone(); 4].into_iter().map(Ok)) + .chain(futures::stream::once(async move { + release_rx.await.unwrap(); + Ok(next) + })) + .chain(futures::stream::iter(vec![task; 3].into_iter().map(Ok))); + let scan = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()) + .with_data_file_concurrency_limit(4) + .with_runtime_predicate_provider(provider.clone()) + .build() + .read(Box::pin(tasks) as FileScanTaskStream) + .unwrap(); + let mut stream = scan.stream(); + let mut first_wave = vec![]; + // Each first-wave file produces two four-row batches (RG1 and RG2). + for _ in 0..8 { + first_wave.push(stream.try_next().await.unwrap().unwrap()); + } + let mut actual = ids(&first_wave); + actual.sort_unstable(); + let mut expected = [100, 101, 102, 103, 200, 201, 202, 203].repeat(4); + expected.sort_unstable(); + assert_eq!(actual, expected); + assert_eq!(provider.snapshots(), 1); + + provider.publish( + Some(Reference::new("id").greater_than_or_equal_to(Datum::int(200))), + 2, + ); + release_tx.send(()).unwrap(); + let second_wave = stream.try_collect::>().await.unwrap(); + let mut actual = ids(&second_wave); + actual.sort_unstable(); + let mut expected = [200, 201, 202, 203].repeat(4); + expected.sort_unstable(); + assert_eq!(actual, expected); + assert_eq!(provider.snapshots(), 2); +} + +async fn check_promoted_runtime_column( + path: &str, + file_type: DataType, + table_type: PrimitiveType, + values: ArrayRef, + predicate: Predicate, +) -> Vec { + write_delete( + path, + vec![ + field("id", DataType::Int32, 1), + field("value", file_type, 3), + ], + vec![Arc::new(Int32Array::from(vec![0, 1, 2, 3])), values], + ); + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::required(3, "value", Type::Primitive(table_type)).into(), + ]) + .build() + .unwrap(), + ); + let planned = Reference::new("id") + .greater_than_or_equal_to(Datum::int(1)) + .bind(schema.clone(), false) + .unwrap(); + let task = scan_task_with_deletes_and_projection( + path.to_string(), + schema, + Some(planned), + vec![], + vec![1, 3], + ); + let (baseline, _) = execute(task.clone(), None).await; + let provider = Arc::new(FixedRuntimePredicate::new(predicate)); + let (batches, _) = execute(task, Some(provider.clone())).await; + assert_eq!(ids(&batches), vec![1, 2, 3]); + assert_eq!( + batches, baseline, + "promotion must fail open and preserve the planned filter" + ); + assert_eq!(provider.snapshots(), 1); + batches +} + +#[tokio::test] +async fn runtime_predicate_on_float_promoted_to_double_is_ignored() { + let temp = TempDir::new().unwrap(); + let path = temp.path().join("float.parquet"); + // The DOUBLE boundary is between adjacent FLOAT values. Narrowing it + // changes the comparison even though it lies within the FLOAT range. + let boundary = 1.0 + f64::from(f32::EPSILON) / 2.0; + let batches = check_promoted_runtime_column( + path.to_str().unwrap(), + DataType::Float32, + PrimitiveType::Double, + Arc::new(Float32Array::from(vec![ + 1.0, + 1.0, + 1.0 + f32::EPSILON, + 1.0 + f32::EPSILON, + ])), + Reference::new("value").less_than(Datum::double(boundary)), + ) + .await; + assert_eq!(batches[0].column(1).data_type(), &DataType::Float64); + let values = batches[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.values().as_ref(), &[ + 1.0, + f64::from(1.0 + f32::EPSILON), + f64::from(1.0 + f32::EPSILON) + ]); +} + +#[tokio::test] +async fn runtime_predicate_on_widened_decimal_precision_is_ignored() { + let temp = TempDir::new().unwrap(); + // 8 -> 9 can share the same physical width; precision itself must be + // checked, rather than only the Parquet physical storage type. + for precision in [9, 12] { + let path = temp.path().join(format!("decimal-{precision}.parquet")); + let batches = check_promoted_runtime_column( + path.to_str().unwrap(), + DataType::Decimal128(8, 2), + PrimitiveType::Decimal { + precision, + scale: 2, + }, + Arc::new( + Decimal128Array::from(vec![123, 456, 789, 1000]) + .with_precision_and_scale(8, 2) + .unwrap(), + ), + // Fits the table precision but exceeds the file's precision. + Reference::new("value").greater_than(Datum::decimal_from_str("1000000.00").unwrap()), + ) + .await; + assert_eq!( + batches[0].column(1).data_type(), + &DataType::Decimal128(precision as u8, 2) + ); + let values = batches[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.values().as_ref(), &[456, 789, 1000]); + } +} + +#[tokio::test] +async fn runtime_predicate_row_group_pruning_preserves_projected_positions() { + use crate::metadata_columns::RESERVED_FIELD_ID_POS; + + let temp = TempDir::new().unwrap(); + let path = write_three_row_group_file(temp.path().to_str().unwrap(), "positions.parquet"); + let task = scan_task_with_deletes_and_projection(path, iceberg_schema(), None, vec![], vec![ + 1, + 2, + RESERVED_FIELD_ID_POS, + ]); + let positions = |batches: &[RecordBatch]| { + batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("_pos") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>() + }; + let (baseline, baseline_metrics) = execute(task.clone(), None).await; + assert_eq!(positions(&baseline), (0..12).collect::>()); + let provider = Arc::new(FixedRuntimePredicate::new( + Reference::new("id") + .greater_than_or_equal_to(Datum::int(100)) + .and(Reference::new("id").less_than_or_equal_to(Datum::int(103))), + )); + let (batches, metrics) = execute(task, Some(provider)).await; + assert_eq!(ids(&batches), vec![100, 101, 102, 103]); + // RG0 was pruned, but RowNumber must still use RG1's original file ordinals. + assert_eq!(positions(&batches), vec![4, 5, 6, 7]); + assert!(metrics.bytes_read() < baseline_metrics.bytes_read()); +} + +#[tokio::test] +async fn runtime_predicate_shares_bloom_filter_reads_with_planned_predicate() { + let temp = TempDir::new().unwrap(); + let path = format!("{}/shared-bloom.parquet", temp.path().to_str().unwrap()); + // Each row group spans the probed values, so statistics prune nothing and + // only the bloom filters can prune row groups 0 and 2. + let schema = Arc::new(ArrowSchema::new(vec![ + field("id", DataType::Int32, 1), + field("payload", DataType::Utf8, 2), + ])); + let props = WriterProperties::builder() + .set_compression(Compression::UNCOMPRESSED) + .set_max_row_group_row_count(Some(4)) + .set_bloom_filter_enabled(true) + .build(); + let mut writer = ArrowWriter::try_new( + File::create(&path).unwrap(), + Arc::clone(&schema), + Some(props), + ) + .unwrap(); + for ids in [[0, 500, 1, 501], [2, 502, 3, 503], [4, 504, 5, 505]] { + let payloads: Vec = ids.iter().map(|id| format!("payload-{id}")).collect(); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![ + Arc::new(Int32Array::from(ids.to_vec())), + Arc::new(StringArray::from(payloads)), + ]) + .unwrap(); + writer.write(&batch).unwrap(); + writer.flush().unwrap(); + } + writer.close().unwrap(); + + let bind = |predicate: Predicate| predicate.bind(iceberg_schema(), false).unwrap(); + let planned = Reference::new("id").equal_to(Datum::int(3)); + let runtime = Reference::new("id").is_in([Datum::int(3), Datum::int(77)]); + // The same restriction as one planned predicate, which reads each bloom + // filter once. + let (expected, expected_metrics) = execute_tasks( + vec![scan_task( + path.clone(), + iceberg_schema(), + Some(bind(planned.clone().and(runtime.clone()))), + )], + None, + true, + ) + .await; + let (batches, metrics) = execute_tasks( + vec![scan_task(path, iceberg_schema(), Some(bind(planned)))], + Some(Arc::new(FixedRuntimePredicate::new(runtime))), + true, + ) + .await; + assert_eq!(ids(&expected), vec![3]); + assert_eq!(ids(&batches), vec![3]); + // A bloom filter is far larger than the rows read, so a repeated read of + // the overlapping column would show up here. + assert_eq!(metrics.bytes_read(), expected_metrics.bytes_read()); +} + +async fn next_ids(stream: &mut ArrowRecordBatchStream, batches: usize) -> Vec { + let mut read = vec![]; + for _ in 0..batches { + read.push(stream.try_next().await.unwrap().unwrap()); + } + ids(&read) +} + +#[tokio::test] +async fn runtime_predicate_transitions_to_none_and_then_to_a_tighter_predicate() { + let temp = TempDir::new().unwrap(); + let path = write_three_row_group_file(temp.path().to_str().unwrap(), "transitions.parquet"); + let task = scan_task(path, iceberg_schema(), None); + let provider = Arc::new(ChangingRuntimePredicate::new( + Some(Reference::new("id").greater_than_or_equal_to(Datum::int(100))), + 1, + )); + // Tasks are released one at a time, so each is planned after the + // preceding publication. + let (tasks, released) = futures::channel::mpsc::unbounded::(); + let scan = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()) + .with_data_file_concurrency_limit(1) + .with_runtime_predicate_provider(provider.clone()) + .build() + .read(Box::pin(released.map(Ok)) as FileScanTaskStream) + .unwrap(); + let mut stream = scan.stream(); + + tasks.unbounded_send(task.clone()).unwrap(); + assert_eq!(next_ids(&mut stream, 2).await, vec![ + 100, 101, 102, 103, 200, 201, 202, 203 + ]); + assert_eq!(provider.snapshots(), 1); + + // `None` stops pruning for later tasks; the cached predicate must not be reused. + provider.publish(None, 2); + tasks.unbounded_send(task.clone()).unwrap(); + assert_eq!(next_ids(&mut stream, 3).await, all_ids()); + assert_eq!(provider.snapshots(), 2); + + // A later predicate tighter than the first one applies again. + provider.publish( + Some(Reference::new("id").greater_than_or_equal_to(Datum::int(200))), + 3, + ); + tasks.unbounded_send(task.clone()).unwrap(); + assert_eq!(next_ids(&mut stream, 1).await, vec![200, 201, 202, 203]); + assert_eq!(provider.snapshots(), 3); + + // The unchanged generation is reused, not snapshotted again. + tasks.unbounded_send(task).unwrap(); + assert_eq!(next_ids(&mut stream, 1).await, vec![200, 201, 202, 203]); + assert_eq!(provider.snapshots(), 3); + + drop(tasks); + assert!(stream.try_next().await.unwrap().is_none()); +} diff --git a/crates/iceberg/src/arrow/record_batch_transformer.rs b/crates/iceberg/src/arrow/record_batch_transformer.rs index 19976a81c5..4d1bc38468 100644 --- a/crates/iceberg/src/arrow/record_batch_transformer.rs +++ b/crates/iceberg/src/arrow/record_batch_transformer.rs @@ -133,6 +133,12 @@ fn constants_map( Ok(constants) } +/// Whether a file column needs casting to the table's Arrow type. Shared with +/// runtime-filter planning, which must not filter values before this cast. +pub(crate) fn column_needs_type_promotion(source_type: &DataType, target_type: &DataType) -> bool { + !source_type.equals_datatype(target_type) +} + /// Indicates how a particular column in a processed RecordBatch should /// be sourced. #[derive(Debug)] @@ -831,7 +837,7 @@ impl RecordBatchTransformer { // No conflict detection needed - schema resolution happened in reader.rs. let field_by_id = field_id_to_source_schema_map.get(field_id).map( |(source_field, source_index)| { - if source_field.data_type().equals_datatype(target_type) { + if !column_needs_type_promotion(source_field.data_type(), target_type) { ColumnSource::PassThrough { source_index: *source_index, }