diff --git a/datafusion/core/tests/parquet/mod.rs b/datafusion/core/tests/parquet/mod.rs index 70806377632ec..308d7176c455e 100644 --- a/datafusion/core/tests/parquet/mod.rs +++ b/datafusion/core/tests/parquet/mod.rs @@ -113,7 +113,7 @@ enum Unit { struct ContextWithParquet { /// temp file parquet data is written to. The file is cleaned up /// when dropped - _file: NamedTempFile, + file: NamedTempFile, provider: Arc, ctx: SessionContext, } @@ -361,7 +361,7 @@ impl ContextWithParquet { ctx.register_table("t", provider.clone()).unwrap(); Self { - _file: file, + file, provider, ctx, } diff --git a/datafusion/core/tests/parquet/page_pruning.rs b/datafusion/core/tests/parquet/page_pruning.rs index 372a7a601d492..70052ec6b7d60 100644 --- a/datafusion/core/tests/parquet/page_pruning.rs +++ b/datafusion/core/tests/parquet/page_pruning.rs @@ -17,7 +17,7 @@ use std::sync::Arc; -use crate::parquet::Unit::Page; +use crate::parquet::Unit::{Page, RowGroupAndPage}; use crate::parquet::{ContextWithParquet, Scenario}; use arrow::array::{Int32Array, RecordBatch}; @@ -31,7 +31,7 @@ use datafusion::datasource::source::DataSourceExec; use datafusion::execution::context::SessionState; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::metrics::MetricValue; -use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; use datafusion_common::{ScalarValue, ToDFSchema}; use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::{Expr, col, lit}; @@ -1051,6 +1051,237 @@ async fn test_pages_with_null_values() { .await; } +async fn page_limit_context(values: Vec>) -> ContextWithParquet { + let row_count = values.len(); + page_limit_context_with_row_group(values, row_count).await +} + +async fn page_limit_context_with_row_group( + values: Vec>, + row_group_rows: usize, +) -> ContextWithParquet { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(values))], + ) + .unwrap(); + + ContextWithParquet::with_custom_data( + Scenario::Int, + RowGroupAndPage(row_group_rows, 3), + schema, + vec![batch], + ) + .await +} + +#[tokio::test] +async fn page_level_limit_pruning() { + // The middle page is fully matched. The surrounding pages each contain a + // match, but their statistics cannot prove that every row matches. + let mut context = page_limit_context(vec![ + Some(0), + Some(10), + Some(0), + Some(5), + Some(6), + Some(7), + Some(0), + Some(8), + Some(0), + ]) + .await; + + let output = context.query("SELECT a FROM t WHERE a >= 5 LIMIT 3").await; + + assert_eq!(output.result_rows, 3); + for value in [5, 6, 7] { + assert!(output.pretty_results().contains(&format!("| {value} "))); + } + assert!(!output.pretty_results().contains("| 10 ")); + assert_eq!(output.metric_value("limit_pruned_rows"), Some(6)); +} + +#[tokio::test] +async fn page_level_limit_pruning_preserves_eliminated_sort() { + // File order is by id, not a. Skipping the first partially matched page + // would lose id=1 even though a later page can satisfy the entire LIMIT. + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("a", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from_iter_values(0..9)), + Arc::new(Int32Array::from(vec![0, 10, 0, 5, 6, 7, 0, 8, 0])), + ], + ) + .unwrap(); + let mut context = ContextWithParquet::with_custom_data( + Scenario::Int, + RowGroupAndPage(9, 3), + schema, + vec![batch], + ) + .await; + context.ctx.deregister_table("t").unwrap(); + context + .ctx + .register_parquet( + "t", + context.file.path().to_str().unwrap(), + ParquetReadOptions::default() + .file_sort_order(vec![vec![col("id").sort(true, false)]]), + ) + .await + .unwrap(); + + let sql = "SELECT id FROM t WHERE a >= 5 ORDER BY id LIMIT 3"; + let plan = context + .ctx + .sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let plan = datafusion::physical_plan::displayable(plan.as_ref()) + .indent(true) + .to_string(); + assert!(!plan.contains("SortExec"), "{plan}"); + assert!(plan.contains("limit=3"), "{plan}"); + + let output = context.query(sql).await; + assert_eq!(output.result_rows, 3); + assert_eq!( + output.pretty_results(), + "+----+\n| id |\n+----+\n| 1 |\n| 3 |\n| 4 |\n+----+" + ); + assert_eq!(output.metric_value("limit_pruned_rows"), None); +} + +#[tokio::test] +async fn page_level_limit_pruning_combines_row_groups_and_pages() { + // RG 0 is partial but has one fully matched page. RG 1 is fully matched, + // but is not large enough to satisfy LIMIT 8 on its own. + let mut context = page_limit_context_with_row_group( + vec![ + Some(0), + Some(10), + Some(0), + Some(5), + Some(6), + Some(7), + Some(8), + Some(9), + Some(10), + Some(11), + Some(12), + Some(13), + ], + 6, + ) + .await; + + let output = context.query("SELECT a FROM t WHERE a >= 5 LIMIT 8").await; + + assert_eq!(output.result_rows, 8); + for value in 5..=12 { + assert!(output.pretty_results().contains(&format!("| {value} "))); + } + assert_eq!(output.metric_value("limit_pruned_rows"), Some(3)); +} + +#[tokio::test] +async fn page_level_limit_pruning_is_null_safe() { + // Page 0 has min=5/max=6, but its NULL row does not pass `a >= 5`. + // Page 1 is the only page with three guaranteed matches. + let mut context = page_limit_context(vec![ + None, + Some(5), + Some(6), + Some(7), + Some(8), + Some(9), + Some(0), + Some(10), + Some(0), + ]) + .await; + + let output = context.query("SELECT a FROM t WHERE a >= 5 LIMIT 3").await; + + assert_eq!(output.result_rows, 3); + for value in [7, 8, 9] { + assert!(output.pretty_results().contains(&format!("| {value} "))); + } + assert_eq!(output.metric_value("limit_pruned_rows"), Some(6)); +} + +#[tokio::test] +async fn page_level_limit_pruning_requires_every_conjunct() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![5, 0, 0, 5, 6, 7, 8, 9, 0])), + Arc::new(Int32Array::from(vec![1, 0, 0, 0, 0, 0, 1, 1, 0])), + Arc::new(Int32Array::from(vec![1; 9])), + ], + ) + .unwrap(); + let mut context = ContextWithParquet::with_custom_data( + Scenario::Int, + RowGroupAndPage(9, 3), + schema, + vec![batch], + ) + .await; + + // `b = c` needs two columns and cannot be proven from one column's page + // statistics. The fully matched `a` page must therefore not be selected. + let output = context + .query("SELECT a FROM t WHERE a >= 5 AND b = c LIMIT 3") + .await; + + assert_eq!(output.result_rows, 3); + for value in [5, 8, 9] { + assert!(output.pretty_results().contains(&format!("| {value} "))); + } + assert_eq!(output.metric_value("limit_pruned_rows"), None); +} + +#[tokio::test] +async fn page_level_limit_pruning_needs_enough_guaranteed_rows() { + // Only the final two-row page is fully matched, which is insufficient for + // LIMIT 4. The ordinary page-pruned plan must be retained. + let mut context = page_limit_context(vec![ + Some(5), + Some(0), + Some(0), + Some(6), + Some(0), + Some(0), + Some(7), + Some(8), + ]) + .await; + + let output = context.query("SELECT a FROM t WHERE a >= 5 LIMIT 4").await; + + assert_eq!(output.result_rows, 4); + for value in [5, 6, 7, 8] { + assert!(output.pretty_results().contains(&format!("| {value} "))); + } + assert_eq!(output.metric_value("limit_pruned_rows"), None); +} + fn cast_count_metric(metric: MetricValue) -> Option { match metric { MetricValue::Count { count, .. } => Some(count.value()), @@ -1100,7 +1331,10 @@ async fn test_parquet_opener_without_page_index() { // Query the table // If the bug exists, this might fail because Opener tries to load PageIndex forcefully - let df = ctx.sql("SELECT * FROM t").await.unwrap(); + let df = ctx + .sql("SELECT * FROM t WHERE a >= 2 LIMIT 1") + .await + .unwrap(); let batches = df .collect() .await @@ -1108,5 +1342,5 @@ async fn test_parquet_opener_without_page_index() { // We expect this to succeed, but currently it might fail assert_eq!(batches.len(), 1); - assert_eq!(batches[0].num_rows(), 3); + assert_eq!(batches[0].num_rows(), 1); } diff --git a/datafusion/datasource-parquet/src/metrics.rs b/datafusion/datasource-parquet/src/metrics.rs index c9a908a989924..f9db118e7ec0b 100644 --- a/datafusion/datasource-parquet/src/metrics.rs +++ b/datafusion/datasource-parquet/src/metrics.rs @@ -377,6 +377,28 @@ impl ParquetFileMetrics { count.add(n); } + /// Record rows skipped by page-level limit pruning. + /// + /// This metric is registered lazily to avoid adding per-file metric setup + /// overhead when the optimization does not apply. + pub(crate) fn add_limit_pruned_rows( + metrics: &ExecutionPlanMetricsSet, + partition: usize, + filename: &str, + n: usize, + ) { + if n == 0 { + return; + } + + let count = MetricBuilder::new(metrics) + .with_new_label("filename", filename.to_string()) + .with_type(MetricType::Summary) + .with_category(MetricCategory::Rows) + .counter("limit_pruned_rows", partition); + count.add(n); + } + /// Record that page index I/O was skipped because row-group statistics /// already proved page index could not prune further. pub(crate) fn add_page_index_load_skipped( diff --git a/datafusion/datasource-parquet/src/mod.rs b/datafusion/datasource-parquet/src/mod.rs index 23348f94aae8f..80d2205a1fab0 100644 --- a/datafusion/datasource-parquet/src/mod.rs +++ b/datafusion/datasource-parquet/src/mod.rs @@ -34,6 +34,7 @@ mod nested_schema_pruning; mod opener; mod page_filter; mod projection_read_plan; +mod pruning; mod push_decoder; mod reader; mod row_filter; diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 25c3bc9a77851..ec24462db0564 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -1401,6 +1401,11 @@ impl RowGroupsPrunedParquetOpen { reader_metadata.parquet_schema(), file_metadata.as_ref(), &prepared.file_metrics, + if prepared.preserve_order { + None + } else { + prepared.limit + }, ); access_plan = page_pruning_result.access_plan; ParquetFileMetrics::add_page_index_pages_skipped_by_fully_matched( @@ -1409,6 +1414,12 @@ impl RowGroupsPrunedParquetOpen { &prepared.file_name, page_pruning_result.pages_skipped_by_fully_matched, ); + ParquetFileMetrics::add_limit_pruned_rows( + &prepared.metrics, + prepared.partition_index, + &prepared.file_name, + page_pruning_result.limit_pruned_rows, + ); } // Prepare access plans, then apply row-group ordering tweaks per @@ -1909,7 +1920,7 @@ mod test { use object_store::{ObjectStore, ObjectStoreExt, memory::InMemory, path::Path}; use parquet::arrow::{ArrowSchemaConverter, ArrowWriter}; use parquet::file::metadata::{ColumnChunkMetaData, FileMetaData, ParquetMetaData}; - use parquet::file::properties::WriterProperties; + use parquet::file::properties::{EnabledStatistics, WriterProperties}; use parquet::schema::types::SchemaDescPtr; use std::collections::VecDeque; use std::sync::Arc; @@ -2248,6 +2259,12 @@ mod test { self } + /// Set whether the scan must preserve file order. + fn with_preserve_order(mut self, enable: bool) -> Self { + self.preserve_order = enable; + self + } + /// Build the ParquetMorselizer instance, unwrapping validation errors. /// /// # Panics @@ -3833,6 +3850,46 @@ mod test { assert_eq!(values, vec![3, 4, 5, 6]); } + #[tokio::test] + async fn test_page_limit_pruning_preserves_order() { + let store = Arc::new(InMemory::new()) as Arc; + let batch = + record_batch!(("a", Int32, vec![0, 10, 0, 5, 6, 7, 0, 8, 0])).unwrap(); + let schema = batch.schema(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(9)) + .set_data_page_row_count_limit(3) + .set_write_batch_size(3) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let data_len = write_parquet_batches( + Arc::clone(&store), + "test.parquet", + vec![batch], + Some(props), + ) + .await; + let file = PartitionedFile::new( + "test.parquet".to_string(), + u64::try_from(data_len).unwrap(), + ); + let predicate = logical2physical(&col("a").gt_eq(lit(5)), &schema); + + let opener = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(schema) + .with_predicate(predicate) + .with_pushdown_filters(true) + .with_row_group_stats_pruning(true) + .with_enable_page_index(true) + .with_preserve_order(true) + .with_limit(3) + .build(); + + let values = collect_int32_values(open_file(&opener, file).await.unwrap()).await; + assert_eq!(values, vec![10, 5, 6]); + } + #[tokio::test] async fn test_fully_matched_runs_preserve_reverse_order() { let store = Arc::new(InMemory::new()) as Arc; diff --git a/datafusion/datasource-parquet/src/page_filter.rs b/datafusion/datasource-parquet/src/page_filter.rs index cd28eb24fdc1a..f6a267a58208f 100644 --- a/datafusion/datasource-parquet/src/page_filter.rs +++ b/datafusion/datasource-parquet/src/page_filter.rs @@ -21,8 +21,9 @@ use std::collections::HashSet; use std::sync::Arc; use super::metrics::ParquetFileMetrics; -use crate::ParquetAccessPlan; use crate::metadata::has_untrusted_min_max_order; +use crate::pruning::build_inverted_predicate; +use crate::{ParquetAccessPlan, RowGroupAccess}; use arrow::array::BooleanArray; use arrow::{ @@ -114,6 +115,9 @@ pub struct PagePruningAccessPlanFilter { /// single column predicates (e.g. (`col = 5`) extracted from the overall /// predicate. Must all be true for a row to be included in the result. predicates: Vec, + /// Whether every original conjunct can be evaluated from one column's page + /// statistics. Otherwise no range can be classified as fully matched. + all_predicates_supported: bool, } /// Result of applying page-index pruning to a [`ParquetAccessPlan`]. @@ -122,16 +126,20 @@ pub(crate) struct PagePruningResult { /// Pages skipped because the containing row group was fully matched by /// row-group statistics. pub(crate) pages_skipped_by_fully_matched: usize, + /// Rows skipped by page-level limit pruning. + pub(crate) limit_pruned_rows: usize, } impl PagePruningResult { fn new( access_plan: ParquetAccessPlan, pages_skipped_by_fully_matched: usize, + limit_pruned_rows: usize, ) -> Self { Self { access_plan, pages_skipped_by_fully_matched, + limit_pruned_rows, } } } @@ -150,36 +158,44 @@ impl PagePruningAccessPlanFilter { schema: &SchemaRef, max_in_list_size: usize, ) -> Self { - // extract any single column predicates - let predicates = split_conjunction(expr) - .into_iter() - .filter_map(|predicate| { - let pp = match PruningPredicateBuilder::new() - .with_file_schema(Arc::clone(schema)) - .with_max_in_list_size(max_in_list_size) - .try_build(Arc::clone(predicate)) - { - Ok(pp) => pp, - Err(e) => { - debug!("Ignoring error creating page pruning predicate: {e}"); - return None; - } - }; + let mut predicates = vec![]; + let mut all_predicates_supported = true; - if pp.always_true() { - debug!("Ignoring always true page pruning predicate: {predicate}"); - return None; + // Ordinary pruning can use any supported conjunct. Fully matched + // ranges require every original conjunct to be supported. + for predicate in split_conjunction(expr) { + let pp = match PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(schema)) + .with_max_in_list_size(max_in_list_size) + .try_build(Arc::clone(predicate)) + { + Ok(pp) => pp, + Err(e) => { + debug!("Ignoring error creating page pruning predicate: {e}"); + all_predicates_supported = false; + continue; } + }; - if pp.required_columns().single_column().is_none() { - debug!("Ignoring multi-column page pruning predicate: {predicate}"); - return None; - } + if pp.always_true() { + debug!("Ignoring always true page pruning predicate: {predicate}"); + all_predicates_supported = false; + continue; + } - Some(pp) - }) - .collect::>(); - Self { predicates } + if pp.required_columns().single_column().is_none() { + debug!("Ignoring multi-column page pruning predicate: {predicate}"); + all_predicates_supported = false; + continue; + } + + predicates.push(pp); + } + + Self { + predicates, + all_predicates_supported, + } } /// Returns an updated [`ParquetAccessPlan`] by applying predicates to the @@ -198,6 +214,7 @@ impl PagePruningAccessPlanFilter { parquet_schema, parquet_metadata, file_metrics, + None, ) .access_plan } @@ -211,18 +228,19 @@ impl PagePruningAccessPlanFilter { parquet_schema: &SchemaDescriptor, parquet_metadata: &ParquetMetaData, file_metrics: &ParquetFileMetrics, + limit: Option, ) -> PagePruningResult { // scoped timer updates on drop let _timer_guard = file_metrics.page_index_eval_time.timer(); if self.predicates.is_empty() { - return PagePruningResult::new(access_plan, 0); + return PagePruningResult::new(access_plan, 0, 0); } let page_index_predicates = &self.predicates; let groups = parquet_metadata.row_groups(); if groups.is_empty() { - return PagePruningResult::new(access_plan, 0); + return PagePruningResult::new(access_plan, 0, 0); } if parquet_metadata.offset_index().is_none() @@ -233,9 +251,28 @@ impl PagePruningAccessPlanFilter { parquet_metadata.offset_index().is_some(), parquet_metadata.column_index().is_some() ); - return PagePruningResult::new(access_plan, 0); + return PagePruningResult::new(access_plan, 0, 0); } + // Build inverse predicates only for unordered LIMIT queries. Ordinary + // page pruning does not need them. + let inverted_predicates: Option> = + if limit.is_some() && self.all_predicates_supported { + self.predicates + .iter() + .map(|predicate| { + let inverted = build_inverted_predicate(predicate, arrow_schema)?; + inverted + .required_columns() + .single_column() + .is_some() + .then_some(inverted) + }) + .collect() + } else { + None + }; + // track the total number of rows that should be skipped let mut total_skip = 0; // track the total number of rows that should not be skipped @@ -247,6 +284,10 @@ impl PagePruningAccessPlanFilter { // track pages for which page-index pruning was skipped because the // containing row group was already proven fully matched by statistics let mut total_pages_skipped_by_fully_matched = 0; + let mut fully_matched_page_selections = inverted_predicates + .as_ref() + .map(|_| vec![None; groups.len()]); + let mut guaranteed_rows = 0; // for each row group specified in the access plan let row_group_indexes = access_plan.row_group_indexes(); @@ -257,6 +298,10 @@ impl PagePruningAccessPlanFilter { let page_count = fully_matched_page_count(row_group_index, parquet_metadata); total_pages_skipped_by_fully_matched += page_count; + guaranteed_rows += selected_rows_in_access( + &access_plan.inner()[row_group_index], + &groups[row_group_index], + ); continue; } @@ -274,35 +319,11 @@ impl PagePruningAccessPlanFilter { HashSet::from_iter(0..total_pages_in_group); for predicate in page_index_predicates { - let Some(column) = predicate.required_columns().single_column() else { - debug!( - "Ignoring multi-column page pruning predicate: {:?}", - predicate.predicate_expr() - ); - continue; - }; - - let converter = StatisticsConverter::try_new( - column.name(), - arrow_schema, - parquet_schema, - ); - - let converter = match converter { - Ok(converter) => converter, - Err(e) => { - debug!( - "Could not create statistics converter for column {}: {e}", - column.name() - ); - continue; - } - }; - - let selection = prune_pages_in_one_row_group( + let selection = prune_pages_for_predicate( row_group_index, predicate, - converter, + arrow_schema, + parquet_schema, parquet_metadata, file_metrics, ); @@ -365,6 +386,55 @@ impl PagePruningAccessPlanFilter { parquet_metadata.row_group(row_group_index).num_rows() as usize; } + if guaranteed_rows < limit.unwrap_or(0) + && access_plan.should_scan(row_group_index) + && let Some(inverted_predicates) = &inverted_predicates + { + let mut fully_matched_selection = None; + let mut complete = true; + + for predicate in inverted_predicates { + let Some((selection, _)) = prune_pages_for_predicate( + row_group_index, + predicate, + arrow_schema, + parquet_schema, + parquet_metadata, + file_metrics, + ) else { + complete = false; + break; + }; + + fully_matched_selection = update_selection( + fully_matched_selection, + complement_selection(selection), + ); + + // An empty intersection cannot gain rows from later conjuncts. + if !fully_matched_selection + .as_ref() + .is_some_and(|selection| selection.selects_any()) + { + break; + } + } + + if complete + && let Some(selections) = fully_matched_page_selections.as_mut() + && let Some(selection) = fully_matched_selection + { + guaranteed_rows += match &access_plan.inner()[row_group_index] { + RowGroupAccess::Skip => 0, + RowGroupAccess::Scan => selection.row_count(), + RowGroupAccess::Selection(existing) => { + existing.intersection(&selection).row_count() + } + }; + selections[row_group_index] = Some(selection); + } + } + let pages_matched = matched_pages_in_group.len(); total_pages_select += pages_matched; total_pages_skip += total_pages_in_group - pages_matched; @@ -380,7 +450,30 @@ impl PagePruningAccessPlanFilter { file_metrics .page_index_pages_pruned .add_matched(total_pages_select); - PagePruningResult::new(access_plan, total_pages_skipped_by_fully_matched) + + let mut limit_pruned_rows = 0; + if let Some(limit) = limit + && let Some(fully_matched_page_selections) = + fully_matched_page_selections.as_deref() + && let Some((limit_plan, pruned_rows, pruned_row_groups)) = limit_pruned_plan( + &access_plan, + fully_matched_page_selections, + limit, + groups, + ) + { + access_plan = limit_plan; + limit_pruned_rows = pruned_rows; + file_metrics + .limit_pruned_row_groups + .add_pruned(pruned_row_groups); + } + + PagePruningResult::new( + access_plan, + total_pages_skipped_by_fully_matched, + limit_pruned_rows, + ) } /// Returns the number of filters in the [`PagePruningAccessPlanFilter`] @@ -409,6 +502,96 @@ fn update_selection( } } +fn complement_selection(selection: RowSelection) -> RowSelection { + let selectors: Vec = selection.into(); + selectors + .into_iter() + .map(|selector| { + if selector.skip { + RowSelector::select(selector.row_count) + } else { + RowSelector::skip(selector.row_count) + } + }) + .collect() +} + +/// Build a plan containing only ranges known to satisfy the predicate, if +/// those ranges contain enough rows to satisfy `limit`. +fn limit_pruned_plan( + access_plan: &ParquetAccessPlan, + fully_matched_page_selections: &[Option], + limit: usize, + groups: &[RowGroupMetaData], +) -> Option<(ParquetAccessPlan, usize, usize)> { + let original_rows = selected_row_count(access_plan, groups); + let original_row_groups = access_plan.row_group_indexes().len(); + let mut candidate_plan = ParquetAccessPlan::new_none(access_plan.len()); + if limit == 0 { + return Some((candidate_plan, original_rows, original_row_groups)); + } + + let mut candidate_rows = 0; + + for row_group_index in access_plan.row_group_indexes() { + let access = if access_plan.is_fully_matched(row_group_index) { + access_plan.inner()[row_group_index].clone() + } else { + let Some(selection) = fully_matched_page_selections[row_group_index].as_ref() + else { + continue; + }; + match &access_plan.inner()[row_group_index] { + RowGroupAccess::Skip => continue, + RowGroupAccess::Scan => RowGroupAccess::Selection(selection.clone()), + RowGroupAccess::Selection(existing) => { + RowGroupAccess::Selection(existing.intersection(selection)) + } + } + }; + + let rows = selected_rows_in_access(&access, &groups[row_group_index]); + if rows == 0 { + continue; + } + + candidate_plan.set(row_group_index, access); + if access_plan.is_fully_matched(row_group_index) { + candidate_plan.mark_fully_matched(row_group_index); + } + candidate_rows += rows; + + if candidate_rows >= limit { + let pruned_rows = original_rows.saturating_sub(candidate_rows); + let pruned_row_groups = original_row_groups + .saturating_sub(candidate_plan.row_group_indexes().len()); + return Some((candidate_plan, pruned_rows, pruned_row_groups)); + } + } + + None +} + +fn selected_row_count( + access_plan: &ParquetAccessPlan, + groups: &[RowGroupMetaData], +) -> usize { + access_plan + .inner() + .iter() + .zip(groups) + .map(|(access, group)| selected_rows_in_access(access, group)) + .sum() +} + +fn selected_rows_in_access(access: &RowGroupAccess, group: &RowGroupMetaData) -> usize { + match access { + RowGroupAccess::Skip => 0, + RowGroupAccess::Scan => group.num_rows() as usize, + RowGroupAccess::Selection(selection) => selection.row_count(), + } +} + /// Returns the number of pages for which page-index pruning is skipped because /// the containing row group is fully matched by row-group statistics. fn fully_matched_page_count( @@ -422,6 +605,39 @@ fn fully_matched_page_count( }) } +fn prune_pages_for_predicate( + row_group_index: usize, + predicate: &PruningPredicate, + arrow_schema: &Schema, + parquet_schema: &SchemaDescriptor, + parquet_metadata: &ParquetMetaData, + metrics: &ParquetFileMetrics, +) -> Option<(RowSelection, Vec)> { + #[cfg(test)] + tests::PAGE_PREDICATE_CALLS.set(tests::PAGE_PREDICATE_CALLS.get() + 1); + + let column = predicate.required_columns().single_column()?; + let converter = + match StatisticsConverter::try_new(column.name(), arrow_schema, parquet_schema) { + Ok(converter) => converter, + Err(e) => { + debug!( + "Could not create statistics converter for column {}: {e}", + column.name() + ); + return None; + } + }; + + prune_pages_in_one_row_group( + row_group_index, + predicate, + converter, + parquet_metadata, + metrics, + ) +} + /// Returns a [`RowSelection`] for the rows in this row group to scan, in addition to a vec of /// booleans that state if each page was matched (true) or not (false). /// @@ -655,3 +871,202 @@ impl PruningStatistics for PagesPruningStatistics<'_> { None } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int32Array; + use arrow::datatypes::{DataType, Field}; + use arrow::record_batch::RecordBatch; + use bytes::Bytes; + use datafusion_expr::Operator; + use datafusion_physical_expr::expressions::{BinaryExpr, Column, lit}; + use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::{ + ArrowReaderOptions, ParquetRecordBatchReaderBuilder, + }; + use parquet::basic::Type as PhysicalType; + use parquet::file::metadata::{ColumnChunkMetaData, PageIndexPolicy}; + use parquet::file::properties::WriterProperties; + use parquet::schema::types::{SchemaDescriptor, Type as SchemaType}; + + thread_local! { + pub(super) static PAGE_PREDICATE_CALLS: std::cell::Cell = const { std::cell::Cell::new(0) }; + } + + #[test] + fn inverse_pruning_short_circuits() { + // Each row group has three two-row pages. Only the middle page can + // prove a > 5; mixed pages must still be scanned by ordinary pruning. + for (middle, limit, selected, fully_matched, expected_calls, expected_rows) in [ + ([0, 10], Some(2), 6, false, 9, 18), // Empty inverse: skip the second conjunct. + ([10, 10], Some(2), 6, false, 8, 2), // First group covers LIMIT. + ([10, 10], Some(3), 6, false, 10, 4), // Two groups cover LIMIT. + ([10, 10], Some(7), 6, false, 12, 18), // Insufficient proof: retain the plan. + ([10, 10], None, 6, false, 6, 18), // Ordinary pruning only. + ([10, 10], Some(0), 6, false, 6, 0), + ([10, 10], Some(2), 3, false, 10, 3), // Only one guaranteed row in RG0. + ([10, 10], Some(2), 1, true, 6, 3), // Fully matched RG0 contributes one row. + ] { + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let props = WriterProperties::builder() + .set_data_page_row_count_limit(2) + .set_write_batch_size(2) + .build(); + let mut writer = + ArrowWriter::try_new(Vec::new(), Arc::clone(&schema), Some(props)) + .unwrap(); + for row_group in 0..3 { + let values = if fully_matched && row_group == 0 { + vec![10; 6] + } else { + vec![0, 10, middle[0], middle[1], 0, 10] + }; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(values))], + ) + .unwrap(); + writer.write(&batch).unwrap(); + writer.flush().unwrap(); + } + let options = ArrowReaderOptions::new() + .with_page_index_policy(PageIndexPolicy::Required); + let reader = ParquetRecordBatchReaderBuilder::try_new_with_options( + Bytes::from(writer.into_inner().unwrap()), + options, + ) + .unwrap(); + let column: Arc = Arc::new(Column::new("a", 0)); + let expr: Arc = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::clone(&column), + Operator::Gt, + lit(5i32), + )), + Operator::And, + Arc::new(BinaryExpr::new(column, Operator::Lt, lit(20i32))), + )); + let filter = PagePruningAccessPlanFilter::new(&expr, Arc::clone(&schema)); + let mut plan = ParquetAccessPlan::new_all(3); + plan.scan_selection( + 0, + vec![ + RowSelector::select(selected), + RowSelector::skip(6 - selected), + ] + .into(), + ); + if fully_matched { + plan.mark_fully_matched(0); + } + let metrics = + ParquetFileMetrics::new(0, "test", &ExecutionPlanMetricsSet::new()); + PAGE_PREDICATE_CALLS.set(0); + let result = filter.prune_plan_with_page_index_and_metrics( + plan, + &schema, + reader.parquet_schema(), + reader.metadata(), + &metrics, + limit, + ); + assert_eq!( + PAGE_PREDICATE_CALLS.get(), + expected_calls, + "middle={middle:?}, limit={limit:?}, selected={selected}, fully_matched={fully_matched}" + ); + assert_eq!( + selected_row_count(&result.access_plan, reader.metadata().row_groups()), + expected_rows + ); + } + } + + #[test] + fn limit_plan_preserves_existing_selections() { + let groups = vec![row_group_metadata(6), row_group_metadata(6)]; + let mut access_plan = ParquetAccessPlan::new(vec![ + RowGroupAccess::Selection( + vec![ + RowSelector::skip(2), + RowSelector::select(1), + RowSelector::skip(3), + ] + .into(), + ), + RowGroupAccess::Selection( + vec![ + RowSelector::select(2), + RowSelector::skip(1), + RowSelector::select(3), + ] + .into(), + ), + ]); + access_plan.mark_fully_matched(0); + + let fully_matched_pages = vec![ + None, + Some( + vec![ + RowSelector::skip(1), + RowSelector::select(4), + RowSelector::skip(1), + ] + .into(), + ), + ]; + + // One selected row from the fully matched row group plus three rows + // from the fully matched page satisfy the limit. + let (plan, pruned_rows, _) = + limit_pruned_plan(&access_plan, &fully_matched_pages, 4, &groups).unwrap(); + assert_eq!(selected_row_count(&plan, &groups), 4); + assert_eq!(pruned_rows, 2); + assert!(plan.is_fully_matched(0)); + assert_eq!(plan.inner()[0], access_plan.inner()[0]); + assert_eq!( + plan.inner()[1], + RowGroupAccess::Selection( + vec![ + RowSelector::skip(1), + RowSelector::select(1), + RowSelector::skip(1), + RowSelector::select(2), + RowSelector::skip(1), + ] + .into() + ) + ); + + // The same candidates are insufficient for a larger limit. + assert!( + limit_pruned_plan(&access_plan, &fully_matched_pages, 5, &groups).is_none() + ); + } + + fn row_group_metadata(num_rows: i64) -> RowGroupMetaData { + let field = SchemaType::primitive_type_builder("a", PhysicalType::INT32) + .build() + .unwrap(); + let schema = Arc::new(SchemaDescriptor::new(Arc::new( + SchemaType::group_type_builder("schema") + .with_fields(vec![Arc::new(field)]) + .build() + .unwrap(), + ))); + let column = ColumnChunkMetaData::builder(schema.column(0)) + .set_num_values(num_rows) + .build() + .unwrap(); + + RowGroupMetaData::builder(schema) + .set_num_rows(num_rows) + .set_column_metadata(vec![column]) + .build() + .unwrap() + } +} diff --git a/datafusion/datasource-parquet/src/pruning.rs b/datafusion/datasource-parquet/src/pruning.rs new file mode 100644 index 0000000000000..9ed2b12518dcb --- /dev/null +++ b/datafusion/datasource-parquet/src/pruning.rs @@ -0,0 +1,71 @@ +// 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. + +//! Shared predicate construction for row-group and page-index pruning. + +use std::sync::Arc; + +use arrow::datatypes::Schema; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::{BinaryExpr, IsNullExpr, NotExpr}; +use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr::{PhysicalExpr, PhysicalExprSimplifier}; +use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder}; + +/// Build a null-safe inverse used to prove every row matches `predicate`. +/// +/// Rows where a filter evaluates to NULL do not pass it, so nullable referenced +/// columns are included in the inverse. If the inverse can be pruned, every row +/// is guaranteed to satisfy the original predicate. +pub(crate) fn build_inverted_predicate( + predicate: &PruningPredicate, + arrow_schema: &Schema, +) -> Option { + // Some pruning rewrites preserve only whether rows can be TRUE, while + // full-match inference must distinguish FALSE from UNKNOWN. + if !predicate.can_be_inverted_for_full_match() { + return None; + } + let mut inverted_expr: Arc = + Arc::new(NotExpr::new(Arc::clone(predicate.orig_expr()))); + + let mut columns = collect_columns(predicate.orig_expr()) + .into_iter() + .filter(|column| arrow_schema.field(column.index()).is_nullable()) + .collect::>(); + columns.sort_by(|a, b| { + a.index() + .cmp(&b.index()) + .then_with(|| a.name().cmp(b.name())) + }); + + for column in columns { + inverted_expr = Arc::new(BinaryExpr::new( + inverted_expr, + Operator::Or, + Arc::new(IsNullExpr::new(Arc::new(column))), + )); + } + + let simplifier = PhysicalExprSimplifier::new(arrow_schema); + let inverted_expr = simplifier.simplify(inverted_expr).ok()?; + PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(predicate.schema())) + .with_max_in_list_size(predicate.max_in_list_size()) + .try_build(inverted_expr) + .ok() +} diff --git a/datafusion/datasource-parquet/src/row_group_filter.rs b/datafusion/datasource-parquet/src/row_group_filter.rs index 5ca99f2498e75..7f733058e6faa 100644 --- a/datafusion/datasource-parquet/src/row_group_filter.rs +++ b/datafusion/datasource-parquet/src/row_group_filter.rs @@ -21,17 +21,14 @@ use std::sync::Arc; use super::{ParquetAccessPlan, ParquetFileMetrics, RowGroupAccess}; use crate::bloom_filter::BloomFilterStatistics; use crate::metadata::{has_untrusted_byte_array_stats, has_untrusted_min_max_order}; +use crate::pruning::build_inverted_predicate; use arrow::array::{ArrayRef, BooleanArray, UInt64Array}; use arrow::compute::nullif; use arrow::datatypes::Schema; use datafusion_common::pruning::PruningStatistics; use datafusion_common::{Column, Result, ScalarValue}; use datafusion_datasource::FileRange; -use datafusion_expr::Operator; -use datafusion_physical_expr::expressions::{BinaryExpr, IsNullExpr, NotExpr}; -use datafusion_physical_expr::utils::collect_columns; -use datafusion_physical_expr::{PhysicalExpr, PhysicalExprSimplifier}; -use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder}; +use datafusion_pruning::PruningPredicate; use parquet::arrow::arrow_reader::statistics::StatisticsConverter; use parquet::basic::ColumnOrder; use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData}; @@ -395,50 +392,9 @@ impl RowGroupAccessPlanFilter { if candidate_row_group_indices.is_empty() { return; } - // Some pruning rewrites preserve only whether rows can be TRUE, while - // full-match inference must distinguish FALSE from UNKNOWN. - if !predicate.can_be_inverted_for_full_match() { - return; - } let arrow_schema = pruning_stats.arrow_schema; - let mut inverted_expr: Arc = - Arc::new(NotExpr::new(Arc::clone(predicate.orig_expr()))); - - // Rows where the predicate evaluates to NULL do not pass the filter. - // Include NULL checks in the inverted expression so a row group is only - // considered fully matched when every referenced column is known non-null. - // This is conservative for null-accepting predicates, but fully matched - // row groups must not have false positives. - let mut columns = collect_columns(predicate.orig_expr()) - .into_iter() - .filter(|column| arrow_schema.field(column.index()).is_nullable()) - .collect::>(); - columns.sort_by(|a, b| { - a.index() - .cmp(&b.index()) - .then_with(|| a.name().cmp(b.name())) - }); - - for column in columns { - inverted_expr = Arc::new(BinaryExpr::new( - inverted_expr, - Operator::Or, - Arc::new(IsNullExpr::new(Arc::new(column))), - )); - } - - // Simplify the inverted expression (e.g., NOT(c1 = 0) -> c1 != 0) - // before building the pruning predicate - let simplifier = PhysicalExprSimplifier::new(arrow_schema); - let Ok(inverted_expr) = simplifier.simplify(inverted_expr) else { - return; - }; - - let Ok(inverted_predicate) = PruningPredicateBuilder::new() - .with_file_schema(Arc::clone(predicate.schema())) - .with_max_in_list_size(predicate.max_in_list_size()) - .try_build(inverted_expr) + let Some(inverted_predicate) = build_inverted_predicate(predicate, arrow_schema) else { return; }; @@ -644,8 +600,10 @@ mod tests { use arrow::datatypes::DataType::Decimal128; use arrow::datatypes::{DataType, Field}; use datafusion_expr::{cast, col, lit}; + use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::planner::logical2physical; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; + use datafusion_pruning::PruningPredicateBuilder; use parquet::arrow::ArrowSchemaConverter; use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; use parquet::basic::LogicalType; diff --git a/datafusion/physical-expr-common/src/metrics/value.rs b/datafusion/physical-expr-common/src/metrics/value.rs index 37ab5194b2cc4..64416b010442d 100644 --- a/datafusion/physical-expr-common/src/metrics/value.rs +++ b/datafusion/physical-expr-common/src/metrics/value.rs @@ -1050,11 +1050,9 @@ impl MetricValue { Self::SpilledRows(_) => 12, Self::CurrentMemoryUsage(_) => 13, Self::Count { name, .. } => match name.as_ref() { - // This Parquet page-index metric is a plain Count because it - // records pages that skipped page-index evaluation, not a - // pruned/matched pair. Keep it grouped with the other + // Keep lazy Parquet pruning counters grouped with the other // page-index pruning metrics in EXPLAIN output. - "page_index_pages_skipped_by_fully_matched" => 8, + "page_index_pages_skipped_by_fully_matched" | "limit_pruned_rows" => 8, _ => 14, }, Self::PeakMemoryUsage { .. } => 13, diff --git a/datafusion/sqllogictest/src/engines/conversion.rs b/datafusion/sqllogictest/src/engines/conversion.rs index d22b518234803..13178e1db2c1e 100644 --- a/datafusion/sqllogictest/src/engines/conversion.rs +++ b/datafusion/sqllogictest/src/engines/conversion.rs @@ -96,7 +96,7 @@ pub(crate) fn arrow_decimal_to_str( } #[cfg(feature = "postgres")] -pub(crate) fn decimal_to_str(value: BigDecimal) -> String { +pub(crate) fn decimal_to_str(value: &BigDecimal) -> String { value.to_plain_string() } diff --git a/datafusion/sqllogictest/src/engines/postgres_engine/mod.rs b/datafusion/sqllogictest/src/engines/postgres_engine/mod.rs index 7ab5f1977d0e4..daa5acf165891 100644 --- a/datafusion/sqllogictest/src/engines/postgres_engine/mod.rs +++ b/datafusion/sqllogictest/src/engines/postgres_engine/mod.rs @@ -377,7 +377,7 @@ fn cell_to_string(row: &SimpleQueryRow, column_type: &Type, idx: usize) -> Strin (&Type::INT4, Some(value)) => value.parse::().unwrap().to_string(), (&Type::INT8, Some(value)) => value.parse::().unwrap().to_string(), (&Type::NUMERIC, Some(value)) => { - decimal_to_str(BigDecimal::from_str(value).unwrap()) + decimal_to_str(&BigDecimal::from_str(value).unwrap()) } // Parse date/time strings explicitly to avoid locale-specific formatting. (&Type::DATE, Some(value)) => NaiveDate::parse_from_str(value, "%Y-%m-%d") diff --git a/docs/source/user-guide/explain-usage.md b/docs/source/user-guide/explain-usage.md index 43ea6caac3a7a..b073313c6a0a4 100644 --- a/docs/source/user-guide/explain-usage.md +++ b/docs/source/user-guide/explain-usage.md @@ -233,6 +233,7 @@ When predicate pushdown is enabled, `DataSourceExec` with `ParquetSource` gains - `row_groups_pruned_bloom_filter`: number of row groups evaluated by Bloom Filters, reporting both total checked groups and groups that matched. - `row_groups_pruned_statistics`: number of row groups evaluated by row-group statistics (min/max), reporting both total checked groups and groups that matched. - `limit_pruned_row_groups`: number of row groups pruned by the limit. +- `limit_pruned_rows`: number of rows skipped by limit pruning when fully matched page ranges contain enough rows to satisfy the limit. - `pushdown_rows_matched`: rows that were tested by any of the above filters, and passed all of them. - `pushdown_rows_pruned`: rows that were tested by any of the above filters, and did not pass at least one of them. - `predicate_evaluation_errors`: number of times evaluating the filter expression failed (expected to be zero in normal operation)