diff --git a/datafusion-examples/examples/custom_data_source/custom_file_format.rs b/datafusion-examples/examples/custom_data_source/custom_file_format.rs index 0cfbe11877e4d..bbd3223da1b16 100644 --- a/datafusion-examples/examples/custom_data_source/custom_file_format.rs +++ b/datafusion-examples/examples/custom_data_source/custom_file_format.rs @@ -38,7 +38,7 @@ use datafusion::{ error::Result, execution::session_state::SessionStateBuilder, physical_expr_common::sort_expr::LexRequirement, - physical_plan::ExecutionPlan, + physical_plan::{ExecutionPlan, PhysicalExpr}, prelude::SessionContext, }; @@ -147,8 +147,11 @@ impl FileFormat for TSVFileFormat { &self, state: &dyn Session, conf: FileScanConfig, + filters: &[Arc], ) -> Result> { - self.csv_file_format.create_physical_plan(state, conf).await + self.csv_file_format + .create_physical_plan(state, conf, filters) + .await } async fn create_writer_physical_plan( diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index 3ff86662f6b0b..f0359ca733fd2 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -769,11 +769,39 @@ impl ListingTable { .with_expr_adapter(self.expr_adapter_factory.clone()) .build(); + // Pass the (non partition) filters to the file format, which may use + // them e.g. to prune files while planning. Filters are still evaluated + // above the scan (see `supports_filters_pushdown`), so filters that + // cannot be converted are simply not passed and never fail the scan. + let physical_filters = if filters.is_empty() { + vec![] + } else { + match DFSchema::try_from(Arc::clone(&self.table_schema)) { + Ok(df_schema) => filters + .iter() + .filter_map(|filter| { + state + .create_physical_expr(filter.clone(), &df_schema) + .inspect_err(|e| { + log::debug!( + "Not passing filter {filter} to the file format: {e}" + ) + }) + .ok() + }) + .collect::>(), + Err(e) => { + log::debug!("Not passing filters to the file format: {e}"); + vec![] + } + } + }; + // create the execution plan let plan = self .options .format - .create_physical_plan(state, scan_config) + .create_physical_plan(state, scan_config, &physical_filters) .await?; Ok(ScanResult::new(plan)) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 4246dcc451d45..10a37237768e5 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -585,6 +585,75 @@ impl Display for SpillCompression { } } +/// How much pruning work the Parquet scan performs during planning to refine +/// the scan's statistics. +/// +/// Pruning normally happens only when a `DataSourceExec` is executed, so the +/// optimizer sees row counts for whole files even when a filter can skip most +/// of them. When eager pruning is enabled, pruning is performed while the scan +/// is planned and its results are used to tighten the (inexact) row count and +/// byte size estimates of the scan, e.g. for join ordering. +/// +/// The levels form a ladder: each level includes the work of all lower levels. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum EagerParquetPruning { + /// No pruning during planning + #[default] + Disabled, + /// Prune row groups using row group statistics (min/max) + RowGroups, + /// Additionally prune data pages using the page index + PageIndex, + /// Additionally prune row groups using bloom filters + BloomFilters, +} + +impl EagerParquetPruning { + /// Returns true if this level performs the work of `level` + pub fn includes(self, level: EagerParquetPruning) -> bool { + self >= level + } +} + +impl FromStr for EagerParquetPruning { + type Err = DataFusionError; + + fn from_str(s: &str) -> Result { + match s.to_ascii_lowercase().as_str() { + "disabled" | "" => Ok(Self::Disabled), + "row_groups" => Ok(Self::RowGroups), + "page_index" => Ok(Self::PageIndex), + "bloom_filters" => Ok(Self::BloomFilters), + other => Err(DataFusionError::Configuration(format!( + "Invalid eager parquet pruning level: {other}. Expected one of: disabled, row_groups, page_index, bloom_filters" + ))), + } + } +} + +impl ConfigField for EagerParquetPruning { + fn visit(&self, v: &mut V, key: &str, description: &'static str) { + v.some(key, self, description) + } + + fn set(&mut self, _: &str, value: &str) -> Result<()> { + *self = EagerParquetPruning::from_str(value)?; + Ok(()) + } +} + +impl Display for EagerParquetPruning { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let str = match self { + Self::Disabled => "disabled", + Self::RowGroups => "row_groups", + Self::PageIndex => "page_index", + Self::BloomFilters => "bloom_filters", + }; + write!(f, "{str}") + } +} + /// A `usize` configuration value that rejects zero when set from strings. /// /// Use this for options where zero is never a meaningful runtime value. @@ -1423,6 +1492,23 @@ config_namespace! { /// Defaults to 20. pub max_in_list_size: usize, default = 20 + /// (reading) How much pruning the parquet scan performs while the query + /// is planned, to tighten the scan's row count and byte size estimates. + /// When the scan executes, it starts from the pruned files and row + /// groups, but evaluates its own predicate again. + /// + /// Valid values are `disabled`, `row_groups` (row group statistics), + /// `page_index` (additionally the page index) and `bloom_filters` + /// (additionally bloom filters). Each level includes the work of the + /// levels before it and may add I/O and latency to planning. Statistics + /// derived from eager pruning are always inexact. + pub eager_pruning: EagerParquetPruning, default = EagerParquetPruning::Disabled + + /// (reading) Eager pruning is skipped for scans with more files than + /// this limit, bounding the planning time spent by + /// `eager_pruning`. A limit of 0 skips eager pruning for every scan. + pub eager_pruning_file_limit: usize, default = 256 + // The following options affect writing to parquet files // and map to parquet::file::properties::WriterProperties diff --git a/datafusion/common/src/file_options/parquet_writer.rs b/datafusion/common/src/file_options/parquet_writer.rs index af4554c4eef84..864ab5edcd801 100644 --- a/datafusion/common/src/file_options/parquet_writer.rs +++ b/datafusion/common/src/file_options/parquet_writer.rs @@ -249,6 +249,8 @@ impl ParquetOptions { skip_arrow_metadata: _, max_predicate_cache_size: _, max_in_list_size: _, + eager_pruning: _, // reads not used for writer props + eager_pruning_file_limit: _, // reads not used for writer props } = self; let mut builder = WriterProperties::builder() @@ -409,6 +411,8 @@ mod tests { enable_page_index: defaults.enable_page_index, pruning: defaults.pruning, max_in_list_size: defaults.max_in_list_size, + eager_pruning: defaults.eager_pruning, + eager_pruning_file_limit: defaults.eager_pruning_file_limit, skip_metadata: defaults.skip_metadata, metadata_size_hint: defaults.metadata_size_hint, pushdown_filters: defaults.pushdown_filters, @@ -534,6 +538,9 @@ mod tests { enable_page_index: global_options_defaults.enable_page_index, pruning: global_options_defaults.pruning, max_in_list_size: global_options_defaults.max_in_list_size, + eager_pruning: global_options_defaults.eager_pruning, + eager_pruning_file_limit: global_options_defaults + .eager_pruning_file_limit, skip_metadata: global_options_defaults.skip_metadata, metadata_size_hint: global_options_defaults.metadata_size_hint, pushdown_filters: global_options_defaults.pushdown_filters, diff --git a/datafusion/core/src/datasource/file_format/mod.rs b/datafusion/core/src/datasource/file_format/mod.rs index c46b472bd6404..b4072f1d54d58 100644 --- a/datafusion/core/src/datasource/file_format/mod.rs +++ b/datafusion/core/src/datasource/file_format/mod.rs @@ -87,6 +87,7 @@ pub(crate) mod test_util { .with_projection_indices(projection)? .with_limit(limit) .build(), + &[], ) .await?; Ok(exec) diff --git a/datafusion/core/tests/parquet/eager_pruning.rs b/datafusion/core/tests/parquet/eager_pruning.rs new file mode 100644 index 0000000000000..7f497eaf63a0b --- /dev/null +++ b/datafusion/core/tests/parquet/eager_pruning.rs @@ -0,0 +1,594 @@ +// 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. + +//! Tests for eager parquet pruning (`datafusion.execution.parquet.eager_pruning`) + +use std::fs::File; +use std::path::Path; +use std::sync::Arc; + +use arrow::array::{Int64Array, RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::util::pretty::pretty_format_batches; +use datafusion::datasource::physical_plan::{FileScanConfig, ParquetSource}; +use datafusion::datasource::source::DataSourceExec; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use datafusion_common::ScalarValue; +use datafusion_common::stats::Precision; +use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_datasource::file_groups::FileGroup; +use datafusion_datasource::file_scan_config::FileScanConfigBuilder; +use datafusion_datasource::{FileExtensions, PartitionedFile}; +use datafusion_datasource_parquet::{ + EagerPruningStats, EagerPruningSummary, ParquetAccessPlan, +}; +use datafusion_physical_plan::{ExecutionPlan, collect, displayable}; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::WriterProperties; +use parquet::schema::types::ColumnPath; +use tempfile::TempDir; + +/// Rows in each test file +const ROWS_PER_FILE: i64 = 1000; + +/// Writes a parquet file with 10 row groups of 100 rows, each with data pages +/// of 10 rows and a bloom filter on column `s`. +/// +/// Column `a` is sorted and holds `offset..offset + ROWS_PER_FILE`. Column `s` +/// holds values `s000` to `s999` scattered across all row groups, so that its +/// row group statistics cannot prune, but its bloom filters can. Column `c` +/// holds the index of the row group, so it has a single value in each row +/// group but not in the file. +fn write_file(path: &Path, offset: i64) { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("s", DataType::Utf8, false), + Field::new("c", DataType::Int64, false), + ])); + let a = Int64Array::from_iter_values(offset..offset + ROWS_PER_FILE); + let s = StringArray::from_iter_values( + (0..ROWS_PER_FILE).map(|i| format!("s{:03}", (i * 37) % 1000)), + ); + let c = Int64Array::from_iter_values((0..ROWS_PER_FILE).map(|i| i / 100)); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(a), Arc::new(s), Arc::new(c)], + ) + .unwrap(); + + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(100)) + .set_data_page_row_count_limit(10) + .set_write_batch_size(10) + .set_column_bloom_filter_enabled(ColumnPath::from("s"), true) + .set_column_bloom_filter_fpp(ColumnPath::from("s"), 0.0001) + .build(); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + let mut writer = + ArrowWriter::try_new(File::create(path).unwrap(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); +} + +/// A session with a table `t` of `num_files` test files, whose column `a` +/// holds `0..num_files * ROWS_PER_FILE` +struct TestTable { + ctx: SessionContext, + dir: TempDir, +} + +impl TestTable { + async fn new(config: &[(&str, &str)], num_files: usize) -> Self { + let mut session_config = SessionConfig::new().with_target_partitions(2); + for (key, value) in config { + session_config = session_config.set_str(key, value); + } + let ctx = SessionContext::new_with_config(session_config); + let dir = TempDir::new().unwrap(); + for i in 0..num_files { + write_file( + &dir.path().join(format!("{i}.parquet")), + i as i64 * ROWS_PER_FILE, + ); + } + ctx.register_parquet( + "t", + dir.path().to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await + .unwrap(); + Self { ctx, dir } + } + + /// A table with eager pruning at `level` and otherwise default settings + async fn with_level(level: &str, num_files: usize) -> Self { + Self::new( + &[("datafusion.execution.parquet.eager_pruning", level)], + num_files, + ) + .await + } + + async fn scan(&self, sql: &str) -> Scan { + Scan::plan(&self.ctx, sql).await + } + + async fn results(&self, sql: &str) -> String { + let batches = self.ctx.sql(sql).await.unwrap().collect().await.unwrap(); + pretty_format_batches(&batches).unwrap().to_string() + } +} + +/// The parquet scan of a planned query +struct Scan { + conf: FileScanConfig, + source: ParquetSource, + plan: String, +} + +impl Scan { + async fn plan(ctx: &SessionContext, sql: &str) -> Self { + let plan = ctx + .sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let (conf, source) = find_parquet_scan(plan.as_ref()) + .unwrap_or_else(|| panic!("no parquet scan in plan for {sql}")); + Self { + conf, + source, + plan: displayable(plan.as_ref()).indent(true).to_string(), + } + } + + fn num_rows(&self) -> Precision { + self.conf.statistics().num_rows + } + + fn num_files(&self) -> usize { + self.conf.file_groups.iter().map(|group| group.len()).sum() + } + + fn num_files_with_access_plan(&self) -> usize { + self.conf + .file_groups + .iter() + .flat_map(|group| group.iter()) + .filter(|file| file.extensions.contains::()) + .count() + } + + fn summary(&self) -> Option<&EagerPruningSummary> { + self.source.eager_pruning_summary() + } + + fn pruned_stats(&self) -> &EagerPruningStats { + match self.summary() { + Some(EagerPruningSummary::Pruned(stats)) => stats, + other => panic!("expected eager pruning, got {other:?}\n{}", self.plan), + } + } +} + +fn find_parquet_scan( + plan: &dyn ExecutionPlan, +) -> Option<(FileScanConfig, ParquetSource)> { + if let Some(exec) = plan.downcast_ref::() + && let Some((conf, source)) = exec.downcast_to_file_source::() + { + return Some((conf.clone(), source.clone())); + } + plan.children() + .into_iter() + .find_map(|child| find_parquet_scan(child.as_ref())) +} + +#[tokio::test] +async fn disabled_by_default() { + let table = TestTable::new(&[], 1).await; + let scan = table.scan("SELECT * FROM t WHERE a < 150").await; + + assert_eq!(scan.summary(), None); + assert_eq!(scan.num_rows(), Precision::Inexact(1000)); + assert_eq!(scan.num_files_with_access_plan(), 0); + assert!(!scan.plan.contains("eager_pruning"), "{}", scan.plan); +} + +#[tokio::test] +async fn row_groups() { + let table = TestTable::with_level("row_groups", 1).await; + let scan = table.scan("SELECT * FROM t WHERE a < 150").await; + + assert_eq!( + scan.pruned_stats(), + &EagerPruningStats { + level: datafusion_common::config::EagerParquetPruning::RowGroups, + files: 1, + files_not_evaluated: 0, + files_pruned: 0, + row_groups: 10, + row_groups_pruned: 8, + rows: 1000, + rows_pruned: 800, + } + ); + assert_eq!(scan.num_rows(), Precision::Inexact(200)); + assert_eq!(scan.num_files_with_access_plan(), 1); + // Column statistics describe the rows that are read, so that estimates + // derived from them do not account for the pruning twice + let statistics = scan.conf.statistics(); + assert_eq!( + statistics.column_statistics[0].max_value, + Precision::Inexact(ScalarValue::Int64(Some(199))) + ); + assert!( + scan.plan.contains( + "eager_pruning=[level=row_groups, files_pruned=0/1, \ + row_groups_pruned=8/10, rows_pruned=800/1000]" + ), + "{}", + scan.plan + ); +} + +#[tokio::test] +async fn page_index() { + let sql = "SELECT * FROM t WHERE a = 123"; + + let table = TestTable::with_level("row_groups", 1).await; + assert_eq!(table.scan(sql).await.num_rows(), Precision::Inexact(100)); + + let table = TestTable::with_level("page_index", 1).await; + let scan = table.scan(sql).await; + assert_eq!(scan.num_rows(), Precision::Inexact(10)); + assert_eq!(scan.pruned_stats().rows_pruned, 990); +} + +#[tokio::test] +async fn page_index_respects_enable_page_index() { + let table = TestTable::new( + &[ + ("datafusion.execution.parquet.eager_pruning", "page_index"), + ("datafusion.execution.parquet.enable_page_index", "false"), + ], + 1, + ) + .await; + let scan = table.scan("SELECT * FROM t WHERE a = 123").await; + assert_eq!(scan.num_rows(), Precision::Inexact(100)); +} + +#[tokio::test] +async fn bloom_filters() { + // 's5005' is within the min/max range of every row group, but absent + let sql = "SELECT * FROM t WHERE s = 's5005'"; + + let table = TestTable::with_level("row_groups", 1).await; + let scan = table.scan(sql).await; + assert_eq!(scan.pruned_stats().row_groups_pruned, 0); + assert_eq!(scan.num_files(), 1); + // Nothing was pruned, so the file is left unchanged + assert_eq!(scan.num_files_with_access_plan(), 0); + + let table = TestTable::with_level("bloom_filters", 1).await; + let scan = table.scan(sql).await; + let stats = scan.pruned_stats(); + assert_eq!(stats.files_pruned, 1); + assert_eq!(stats.row_groups_pruned, 10); + assert_eq!(scan.num_files(), 0); + assert_eq!(scan.num_rows(), Precision::Inexact(0)); +} + +#[tokio::test] +async fn prunes_files() { + let table = TestTable::with_level("row_groups", 3).await; + let scan = table.scan("SELECT * FROM t WHERE a < 150").await; + + let stats = scan.pruned_stats(); + assert_eq!(stats.files, 3); + assert_eq!(stats.files_pruned, 2); + assert_eq!(stats.row_groups_pruned, 28); + assert_eq!(scan.num_files(), 1); + assert_eq!(scan.num_rows(), Precision::Inexact(200)); +} + +#[tokio::test] +async fn file_limit() { + let table = TestTable::new( + &[ + ("datafusion.execution.parquet.eager_pruning", "row_groups"), + ("datafusion.execution.parquet.eager_pruning_file_limit", "2"), + ], + 3, + ) + .await; + let scan = table.scan("SELECT * FROM t WHERE a < 150").await; + + assert_eq!( + scan.summary(), + Some(&EagerPruningSummary::Skipped { + level: datafusion_common::config::EagerParquetPruning::RowGroups, + reason: "3 files exceed eager_pruning_file_limit 2".to_string(), + }) + ); + assert_eq!(scan.num_files(), 3); + assert_eq!(scan.num_rows(), Precision::Inexact(3000)); +} + +#[tokio::test] +async fn no_filters() { + let table = TestTable::with_level("bloom_filters", 1).await; + let scan = table.scan("SELECT * FROM t").await; + + assert_eq!(scan.summary(), None); + assert_eq!(scan.num_rows(), Precision::Exact(1000)); +} + +#[tokio::test] +async fn nothing_to_prune_keeps_exact_file_statistics() { + let table = TestTable::with_level("row_groups", 1).await; + let scan = table.scan("SELECT * FROM t WHERE a >= 0").await; + + assert_eq!(scan.pruned_stats().rows_pruned, 0); + assert_eq!(scan.num_files_with_access_plan(), 0); + let file = &scan.conf.file_groups[0].files()[0]; + assert_eq!( + file.statistics.as_ref().unwrap().num_rows, + Precision::Exact(1000) + ); +} + +#[tokio::test] +async fn without_collected_statistics() { + let table = TestTable::new( + &[ + ("datafusion.execution.parquet.eager_pruning", "row_groups"), + ("datafusion.execution.collect_statistics", "false"), + ], + 1, + ) + .await; + let scan = table.scan("SELECT * FROM t WHERE a < 150").await; + + assert_eq!(scan.num_rows(), Precision::Inexact(200)); +} + +#[tokio::test] +async fn without_collected_statistics_and_unpruned_files() { + let table = TestTable::new( + &[ + ("datafusion.execution.parquet.eager_pruning", "row_groups"), + ("datafusion.execution.collect_statistics", "false"), + ], + 2, + ) + .await; + // The second file is read entirely, but its row count is known from the + // metadata eager pruning reads, so the row count of the scan is known too + let scan = table + .scan("SELECT * FROM t WHERE a < 150 OR a >= 1000") + .await; + + assert_eq!(scan.pruned_stats().rows_pruned, 800); + assert_eq!(scan.num_rows(), Precision::Inexact(1200)); +} + +#[tokio::test] +async fn page_index_without_pruned_rows() { + // Row group 0 is not fully matched by `a <> 5`, but no page can be pruned + let table = TestTable::with_level("page_index", 1).await; + let scan = table.scan("SELECT * FROM t WHERE a <> 5").await; + + assert_eq!(scan.pruned_stats().rows_pruned, 0); + // Nothing was pruned, so the file is left unchanged + assert_eq!(scan.num_files_with_access_plan(), 0); + let file = &scan.conf.file_groups[0].files()[0]; + assert_eq!( + file.statistics.as_ref().unwrap().num_rows, + Precision::Exact(1000) + ); +} + +/// Removes the access plans attached to the files of parquet scans, like +/// serializing a plan does (file statistics are serialized, extensions are not) +fn without_access_plans(plan: Arc) -> Arc { + plan.transform_up(|node| { + let Some(exec) = node.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + let Some((conf, _)) = exec.downcast_to_file_source::() else { + return Ok(Transformed::no(node)); + }; + let file_groups = conf + .file_groups + .iter() + .map(|group| { + FileGroup::new( + group + .iter() + .map(|file| PartitionedFile { + extensions: FileExtensions::new(), + ..file.clone() + }) + .collect(), + ) + }) + .collect(); + let conf = FileScanConfigBuilder::from(conf.clone()) + .with_file_groups(file_groups) + .build(); + Ok(Transformed::yes( + DataSourceExec::from_data_source(conf) as Arc + )) + }) + .unwrap() + .data +} + +#[tokio::test] +async fn correct_without_access_plans() { + // Eager pruning reads only row group 0, in which `c` has a single value. + // The statistics of the file must still be correct for the entire file + // when the attached access plan is lost, otherwise the scan could replace + // `c` with that value in the row groups it now reads. + let sql = "SELECT c, count(*) FROM t WHERE c = 0 GROUP BY c"; + for level in ["row_groups", "page_index"] { + for pushdown_filters in ["false", "true"] { + let table = TestTable::new( + &[ + ("datafusion.execution.parquet.eager_pruning", level), + ( + "datafusion.execution.parquet.pushdown_filters", + pushdown_filters, + ), + ], + 1, + ) + .await; + let plan = table + .ctx + .sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let plan = without_access_plans(plan); + let batches = collect(plan, table.ctx.task_ctx()).await.unwrap(); + assert_eq!( + pretty_format_batches(&batches).unwrap().to_string(), + "+---+----------+\n\ + | c | count(*) |\n\ + +---+----------+\n\ + | 0 | 100 |\n\ + +---+----------+", + "level={level}, pushdown_filters={pushdown_filters}" + ); + } + } +} + +#[tokio::test] +async fn table_option() { + let table = TestTable::new(&[], 1).await; + let location = table.dir.path().to_str().unwrap(); + table + .ctx + .sql(&format!( + "CREATE EXTERNAL TABLE t_eager STORED AS PARQUET LOCATION '{location}' \ + OPTIONS ('format.eager_pruning' 'row_groups')" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + + let scan = table.scan("SELECT * FROM t_eager WHERE a < 150").await; + assert_eq!(scan.num_rows(), Precision::Inexact(200)); + + // The session default is unchanged + let scan = table.scan("SELECT * FROM t WHERE a < 150").await; + assert_eq!(scan.summary(), None); +} + +#[tokio::test] +async fn partition_columns() { + let ctx = SessionContext::new_with_config( + SessionConfig::new() + .set_str("datafusion.execution.parquet.eager_pruning", "row_groups"), + ); + let dir = TempDir::new().unwrap(); + write_file(&dir.path().join("part=1").join("0.parquet"), 0); + write_file(&dir.path().join("part=2").join("0.parquet"), 0); + ctx.register_parquet( + "t", + dir.path().to_str().unwrap(), + ParquetReadOptions::default() + .table_partition_cols(vec![("part".to_string(), DataType::Int32)]), + ) + .await + .unwrap(); + + // The filter cannot be evaluated using partition values alone, but after + // substituting them it prunes all rows with part = 1 and a >= 150 + let sql = "SELECT count(*) FROM t WHERE a < 150 OR part = 2"; + let scan = Scan::plan(&ctx, sql).await; + assert_eq!(scan.num_rows(), Precision::Inexact(1200)); + + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + assert_eq!( + pretty_format_batches(&batches).unwrap().to_string(), + "+----------+\n\ + | count(*) |\n\ + +----------+\n\ + | 1150 |\n\ + +----------+" + ); +} + +#[tokio::test] +async fn results_match_disabled() { + let queries = [ + "SELECT count(*) FROM t WHERE a < 150", + "SELECT a, s FROM t WHERE a = 123", + "SELECT count(*) FROM t WHERE s = 's5005'", + "SELECT count(*) FROM t WHERE s = 's500'", + "SELECT a FROM t WHERE a >= 95 AND a < 105 ORDER BY a", + "SELECT count(*), min(a), max(a) FROM t WHERE a < 450 AND s > 's500'", + "SELECT a FROM t WHERE a < 2500 ORDER BY a DESC LIMIT 3", + "SELECT t.a FROM t JOIN (VALUES (5), (250), (1500)) AS v(x) ON t.a = v.x \ + WHERE t.a < 500 ORDER BY t.a", + "SELECT count(*) FROM t WHERE a > 5000", + ]; + + let expected = { + let table = TestTable::new(&[], 3).await; + let mut expected = Vec::with_capacity(queries.len()); + for sql in queries { + expected.push(table.results(sql).await); + } + expected + }; + + for level in ["row_groups", "page_index", "bloom_filters"] { + for pushdown_filters in ["false", "true"] { + let table = TestTable::new( + &[ + ("datafusion.execution.parquet.eager_pruning", level), + ( + "datafusion.execution.parquet.pushdown_filters", + pushdown_filters, + ), + ], + 3, + ) + .await; + for (sql, expected) in queries.iter().zip(&expected) { + assert_eq!( + &table.results(sql).await, + expected, + "level={level}, pushdown_filters={pushdown_filters}: {sql}" + ); + } + } + } +} diff --git a/datafusion/core/tests/parquet/mod.rs b/datafusion/core/tests/parquet/mod.rs index 308d7176c455e..d6174874278c6 100644 --- a/datafusion/core/tests/parquet/mod.rs +++ b/datafusion/core/tests/parquet/mod.rs @@ -47,6 +47,7 @@ use tempfile::NamedTempFile; mod content_defined_chunking; mod custom_reader; mod dynamic_row_group_pruning; +mod eager_pruning; #[cfg(feature = "parquet_encryption")] mod encryption; mod expr_adapter; diff --git a/datafusion/datasource-arrow/src/file_format.rs b/datafusion/datasource-arrow/src/file_format.rs index b874466e2f7d7..1062640d95b91 100644 --- a/datafusion/datasource-arrow/src/file_format.rs +++ b/datafusion/datasource-arrow/src/file_format.rs @@ -57,7 +57,9 @@ use datafusion_datasource::file_format::{FileFormat, FileFormatFactory}; use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig}; use datafusion_datasource::source::DataSourceExec; use datafusion_datasource::write::demux::DemuxedStreamReceiver; -use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, +}; use datafusion_session::Session; use futures::StreamExt; use futures::stream::BoxStream; @@ -185,6 +187,7 @@ impl FileFormat for ArrowFormat { &self, state: &dyn Session, conf: FileScanConfig, + _filters: &[Arc], ) -> Result> { let object_store = state.runtime_env().object_store(&conf.object_store_url)?; let object_location = &conf diff --git a/datafusion/datasource-avro/src/file_format.rs b/datafusion/datasource-avro/src/file_format.rs index 0c081b0ee9b64..ec668ccb34503 100644 --- a/datafusion/datasource-avro/src/file_format.rs +++ b/datafusion/datasource-avro/src/file_format.rs @@ -51,7 +51,9 @@ use datafusion_datasource::write::{ use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr_common::sort_expr::LexRequirement; -use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, +}; use datafusion_session::Session; use async_trait::async_trait; @@ -167,6 +169,7 @@ impl FileFormat for AvroFormat { &self, _state: &dyn Session, conf: FileScanConfig, + _filters: &[Arc], ) -> Result> { Ok(DataSourceExec::from_data_source(conf)) } diff --git a/datafusion/datasource-csv/src/file_format.rs b/datafusion/datasource-csv/src/file_format.rs index 15e24df4bf7c7..9bc127df7c303 100644 --- a/datafusion/datasource-csv/src/file_format.rs +++ b/datafusion/datasource-csv/src/file_format.rs @@ -52,7 +52,9 @@ use datafusion_datasource::write::orchestration::spawn_writer_tasks_and_join; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr_common::sort_expr::LexRequirement; -use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, +}; use datafusion_session::Session; use async_trait::async_trait; @@ -434,6 +436,7 @@ impl FileFormat for CsvFormat { &self, state: &dyn Session, conf: FileScanConfig, + _filters: &[Arc], ) -> Result> { // Consult configuration options for default values let has_header = self diff --git a/datafusion/datasource-json/src/file_format.rs b/datafusion/datasource-json/src/file_format.rs index aef89560a3b5f..cadd641b9d4c3 100644 --- a/datafusion/datasource-json/src/file_format.rs +++ b/datafusion/datasource-json/src/file_format.rs @@ -55,7 +55,9 @@ use datafusion_datasource::write::orchestration::spawn_writer_tasks_and_join; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr_common::sort_expr::LexRequirement; -use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, +}; use datafusion_session::Session; use crate::utils::JsonArrayToNdjsonReader; @@ -345,6 +347,7 @@ impl FileFormat for JsonFormat { &self, _state: &dyn Session, conf: FileScanConfig, + _filters: &[Arc], ) -> Result> { let conf = FileScanConfigBuilder::from(conf) .with_file_compression_type(FileCompressionType::from( diff --git a/datafusion/datasource-parquet/src/access_plan.rs b/datafusion/datasource-parquet/src/access_plan.rs index b72612134d7db..f5374d6e9ec4f 100644 --- a/datafusion/datasource-parquet/src/access_plan.rs +++ b/datafusion/datasource-parquet/src/access_plan.rs @@ -380,6 +380,42 @@ impl ParquetAccessPlan { &self.fully_matched } + /// Clears all fully matched flags. + /// + /// Fully matched flags are only valid for the predicate they were computed + /// with. They must be cleared before the plan is reused with a predicate + /// that may differ, such as a plan computed while planning a scan that is + /// later executed with additional (e.g. dynamic) filters. + pub(crate) fn clear_fully_matched(&mut self) { + self.fully_matched.fill(false); + } + + /// Returns the number of rows this plan selects, given the metadata of + /// the file's row groups. + /// + /// Returns an error if `row_group_meta_data` does not describe the same + /// number of row groups as this plan. + pub fn selected_row_count( + &self, + row_group_meta_data: &[RowGroupMetaData], + ) -> Result { + assert_eq_or_internal_err!( + row_group_meta_data.len(), + self.len(), + "ParquetAccessPlan and row group metadata describe a different number of row groups" + ); + Ok(self + .row_groups + .iter() + .zip(row_group_meta_data) + .map(|(access, row_group)| match access { + RowGroupAccess::Skip => 0, + RowGroupAccess::Scan => row_group.num_rows() as usize, + RowGroupAccess::Selection(selection) => selection.row_count(), + }) + .sum()) + } + /// Set to scan only the [`RowSelection`] in the specified row group. /// /// Behavior is different depending on the existing access diff --git a/datafusion/datasource-parquet/src/bloom_filter.rs b/datafusion/datasource-parquet/src/bloom_filter.rs index 2384f95df9885..1393eaed092f5 100644 --- a/datafusion/datasource-parquet/src/bloom_filter.rs +++ b/datafusion/datasource-parquet/src/bloom_filter.rs @@ -22,12 +22,80 @@ use std::collections::{HashMap, HashSet}; use arrow::array::{ArrayRef, BooleanArray}; +use arrow::datatypes::Schema; use datafusion_common::pruning::PruningStatistics; use datafusion_common::{Column, ScalarValue}; +use datafusion_physical_plan::metrics::Count; +use datafusion_pruning::PruningPredicate; +use log::debug; +use parquet::arrow::ParquetRecordBatchStreamBuilder; +use parquet::arrow::async_reader::AsyncFileReader; +use parquet::arrow::parquet_column; use parquet::basic::Type; use parquet::bloom_filter::Sbbf; use parquet::data_type::Decimal; +/// Reads the bloom filters needed to evaluate `predicate` on the row groups in +/// `row_group_indexes`. +/// +/// Returns one [`BloomFilterStatistics`] per row group in the file. Row groups +/// not listed in `row_group_indexes` and columns without a bloom filter get no +/// bloom filters. Errors reading an individual bloom filter are ignored (and +/// counted in `predicate_evaluation_errors`): a missing bloom filter only means +/// less pruning. +pub(crate) async fn load_row_group_bloom_filters( + builder: &mut ParquetRecordBatchStreamBuilder, + predicate: &PruningPredicate, + physical_file_schema: &Schema, + row_group_indexes: &[usize], + predicate_evaluation_errors: &Count, +) -> Vec +where + T: AsyncFileReader + Send + 'static, +{ + let mut row_group_bloom_filters = + vec![BloomFilterStatistics::new(); builder.metadata().num_row_groups()]; + + let parquet_columns: Vec<(String, usize, Type, i32)> = predicate + .literal_columns() + .into_iter() + .filter_map(|column_name| { + let parquet_schema = builder.parquet_schema(); + let (column_idx, _) = + parquet_column(parquet_schema, physical_file_schema, &column_name)?; + Some(( + column_name, + column_idx, + parquet_schema.column(column_idx).physical_type(), + parquet_schema.column(column_idx).type_length(), + )) + }) + .collect(); + + for &idx in row_group_indexes { + let mut row_group_filters = + BloomFilterStatistics::with_capacity(parquet_columns.len()); + for (column_name, column_idx, physical_type, type_length) in &parquet_columns { + let bf: Sbbf = match builder + .get_row_group_column_bloom_filter(idx, *column_idx) + .await + { + Ok(Some(bf)) => bf, + Ok(None) => continue, + Err(e) => { + debug!("Ignoring error reading bloom filter: {e}"); + predicate_evaluation_errors.add(1); + continue; + } + }; + row_group_filters.insert(column_name, bf, *physical_type, *type_length); + } + row_group_bloom_filters[idx] = row_group_filters; + } + + row_group_bloom_filters +} + /// In memory Parquet Split Block Bloom Filters (SBBF). /// /// This structure implements [`PruningStatistics`] and is used to prune @@ -663,55 +731,16 @@ mod tests { let access_plan = ParquetAccessPlan::new_all(builder.metadata().num_row_groups()); let mut pruned_row_groups = RowGroupAccessPlanFilter::new(access_plan); - let literal_columns = pruning_predicate.literal_columns(); - let parquet_columns: Vec<_> = literal_columns - .into_iter() - .filter_map(|column_name| { - let (column_idx, _) = parquet::arrow::parquet_column( - builder.parquet_schema(), - pruning_predicate.schema(), - &column_name, - )?; - Some(( - column_name.to_string(), - column_idx, - builder.parquet_schema().column(column_idx).physical_type(), - builder.parquet_schema().column(column_idx).type_length(), - )) - }) - .collect::>(); - let mut row_group_bloom_filters = - Vec::with_capacity(builder.metadata().num_row_groups()); - row_group_bloom_filters.resize_with( - builder.metadata().num_row_groups(), - BloomFilterStatistics::new, - ); - for idx in pruned_row_groups.row_group_indexes() { - let mut bloom_filters = - BloomFilterStatistics::with_capacity(parquet_columns.len()); - for (column_name, column_idx, physical_type, type_length) in &parquet_columns - { - let bf = match builder - .get_row_group_column_bloom_filter(idx, *column_idx) - .await - { - Ok(Some(bf)) => bf, - Ok(None) => continue, - Err(e) => { - log::debug!("Ignoring error reading bloom filter: {e}"); - file_metrics.predicate_evaluation_errors.add(1); - continue; - } - }; - bloom_filters.insert( - column_name.clone(), - bf, - *physical_type, - *type_length, - ); - } - row_group_bloom_filters[idx] = bloom_filters; - } + let row_group_indexes: Vec = + pruned_row_groups.row_group_indexes().collect(); + let row_group_bloom_filters = load_row_group_bloom_filters( + &mut builder, + pruning_predicate, + pruning_predicate.schema(), + &row_group_indexes, + &file_metrics.predicate_evaluation_errors, + ) + .await; pruned_row_groups.prune_by_bloom_filters( pruning_predicate, &file_metrics, diff --git a/datafusion/datasource-parquet/src/eager_pruning.rs b/datafusion/datasource-parquet/src/eager_pruning.rs new file mode 100644 index 0000000000000..8834babeab166 --- /dev/null +++ b/datafusion/datasource-parquet/src/eager_pruning.rs @@ -0,0 +1,1245 @@ +// 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. + +//! Eager pruning: pruning parquet files while a scan is planned. +//! +//! Pruning normally happens only when a parquet scan executes, so during +//! planning the scan reports statistics for entire files, even when a filter +//! can skip most of their data. Eager pruning (see [`EagerParquetPruning`]) +//! runs the same row group, page index and bloom filter pruning while the scan +//! is planned, and uses the results to: +//! +//! * drop files in which no row can match the filters, +//! * attach the resulting [`ParquetAccessPlan`] to each file, so the scan +//! starts from the already pruned plan, and +//! * tighten the (always inexact) statistics of the files and of the scan. +//! +//! The access plan attached to a file never marks row groups as fully +//! matched: the scan may execute with a different (e.g. dynamic) predicate, +//! so it re-evaluates its own predicate starting from the attached plan. + +use std::collections::HashMap; +use std::fmt; +use std::sync::Arc; + +use arrow::datatypes::{Schema, TimeUnit}; +use datafusion_common::config::{EagerParquetPruning, TableParquetOptions}; +use datafusion_common::stats::{NdvFallback, Precision}; +use datafusion_common::{ColumnStatistics, Result, ScalarValue, Statistics}; +use datafusion_datasource::file_groups::FileGroup; +use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; +use datafusion_datasource::{PartitionedFile, TableSchema}; +use datafusion_physical_expr::conjunction; +use datafusion_physical_expr::simplifier::PhysicalExprSimplifier; +use datafusion_physical_expr_adapter::{ + DefaultPhysicalExprAdapterFactory, PhysicalExprAdapterFactory, + replace_columns_with_literals, +}; +use datafusion_physical_plan::PhysicalExpr; +use datafusion_physical_plan::metrics::{Count, ExecutionPlanMetricsSet}; +use futures::StreamExt; +use log::debug; +use parquet::arrow::ParquetRecordBatchStreamBuilder; +use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; +use parquet::arrow::async_reader::AsyncFileReader; +use parquet::file::metadata::{PageIndexPolicy, RowGroupMetaData}; + +use crate::access_plan::RowGroupAccess; +use crate::bloom_filter::load_row_group_bloom_filters; +use crate::metadata::DFParquetMetadata; +use crate::opener::{ + build_page_pruning_predicate, build_pruning_predicates, constant_columns_from_stats, + create_initial_plan, load_page_index, +}; +use crate::row_group_filter::RowGroupAccessPlanFilter; +use crate::schema_coercion::coerce_physical_file_schema; +use crate::source::{parse_coerce_int96_string, parse_coerce_int96_tz_string}; +use crate::{ + ParquetAccessPlan, ParquetFileMetrics, ParquetFileReaderFactory, ParquetRowSelection, +}; + +/// The outcome of eager pruning for a parquet scan. +/// +/// See [`EagerParquetPruning`]. Shown in the `EXPLAIN` output of the scan. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum EagerPruningSummary { + /// Eager pruning was enabled but not performed + Skipped { + /// The configured eager pruning level + level: EagerParquetPruning, + /// Why eager pruning was not performed + reason: String, + }, + /// Eager pruning was performed + Pruned(EagerPruningStats), +} + +/// Counts describing the work done by eager pruning. +/// +/// Row group and row counts only include files that were evaluated. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct EagerPruningStats { + /// The configured eager pruning level + pub level: EagerParquetPruning, + /// Number of files in the scan + pub files: usize, + /// Number of files that were not evaluated and left unchanged, for + /// example because their metadata could not be read + pub files_not_evaluated: usize, + /// Number of files removed from the scan because no row can match + pub files_pruned: usize, + /// Number of row groups in the evaluated files + pub row_groups: usize, + /// Number of row groups that will not be read + pub row_groups_pruned: usize, + /// Number of rows in the evaluated files + pub rows: usize, + /// Number of rows that will not be read + pub rows_pruned: usize, +} + +impl fmt::Display for EagerPruningSummary { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Skipped { level, reason } => { + write!(f, "[level={level}, skipped={reason}]") + } + Self::Pruned(stats) => { + let EagerPruningStats { + level, + files, + files_not_evaluated, + files_pruned, + row_groups, + row_groups_pruned, + rows, + rows_pruned, + } = stats; + write!( + f, + "[level={level}, files_pruned={files_pruned}/{files}, \ + row_groups_pruned={row_groups_pruned}/{row_groups}, \ + rows_pruned={rows_pruned}/{rows}" + )?; + if *files_not_evaluated > 0 { + write!(f, ", files_not_evaluated={files_not_evaluated}")?; + } + write!(f, "]") + } + } + } +} + +/// Performs eager pruning for a parquet scan. +pub(crate) struct EagerPruner { + level: EagerParquetPruning, + /// The conjunction of the scan's filters, against the table schema + predicate: Arc, + options: TableParquetOptions, + reader_factory: Arc, + concurrency: usize, +} + +/// The result of eager pruning a single file +enum FileOutcome { + /// The file was not evaluated and is kept unchanged + NotEvaluated(PartitionedFile), + /// The file was evaluated + Evaluated { + /// The file to scan, or `None` if no row in the file can match + file: Option, + /// Whether the file changed, either what is read from it or its + /// statistics + changed: bool, + row_groups: usize, + row_groups_pruned: usize, + rows: usize, + rows_pruned: usize, + }, +} + +/// Per scan state needed to prune individual files +struct ScanContext<'a> { + table_schema: &'a TableSchema, + expr_adapter_factory: Arc, + coerce_int96: Option, + coerce_int96_tz: Option>, +} + +impl EagerPruner { + /// Create a pruner for a scan with the given `filters`. + /// + /// Returns `None` if eager pruning is disabled or there is nothing to prune + /// with. + pub(crate) fn try_new( + options: &TableParquetOptions, + filters: &[Arc], + reader_factory: Arc, + concurrency: usize, + ) -> Option { + let level = options.global.eager_pruning; + if level == EagerParquetPruning::Disabled || filters.is_empty() { + return None; + } + Some(Self { + level, + predicate: conjunction(filters.iter().cloned()), + options: options.clone(), + reader_factory, + concurrency: concurrency.max(1), + }) + } + + /// Eagerly prune the files of `conf`. + /// + /// Returns the updated configuration and a summary of the work done. + /// Errors pruning individual files are not returned: such files are left + /// unchanged. + pub(crate) async fn prune( + &self, + table_schema: &TableSchema, + conf: FileScanConfig, + ) -> Result<(FileScanConfig, EagerPruningSummary)> { + let num_files: usize = conf.file_groups.iter().map(FileGroup::len).sum(); + if let Some(reason) = self.skip_reason(table_schema, num_files) { + debug!("Skipping eager parquet pruning: {reason}"); + let summary = EagerPruningSummary::Skipped { + level: self.level, + reason, + }; + return Ok((conf, summary)); + } + + let global = &self.options.global; + let context = ScanContext { + table_schema, + expr_adapter_factory: conf + .expr_adapter_factory + .clone() + .unwrap_or_else(|| Arc::new(DefaultPhysicalExprAdapterFactory) as _), + coerce_int96: global + .coerce_int96 + .as_deref() + .map(parse_coerce_int96_string) + .transpose()?, + coerce_int96_tz: global + .coerce_int96_tz + .as_deref() + .map(parse_coerce_int96_tz_string) + .transpose()?, + }; + + let group_lens: Vec = + conf.file_groups.iter().map(FileGroup::len).collect(); + let files = conf + .file_groups + .iter() + .flat_map(|group| group.iter().cloned()) + .collect::>(); + let outcomes: Vec = futures::stream::iter(files) + .map(|file| self.prune_file_or_keep(&context, file)) + .buffered(self.concurrency) + .collect() + .await; + + let mut stats = EagerPruningStats { + level: self.level, + files: num_files, + ..Default::default() + }; + let mut changed = false; + let mut outcomes = outcomes.into_iter(); + let mut file_groups = Vec::with_capacity(group_lens.len()); + for len in group_lens { + let mut files = Vec::with_capacity(len); + for outcome in outcomes.by_ref().take(len) { + match outcome { + FileOutcome::NotEvaluated(file) => { + stats.files_not_evaluated += 1; + files.push(file); + } + FileOutcome::Evaluated { + file, + changed: file_changed, + row_groups, + row_groups_pruned, + rows, + rows_pruned, + } => { + changed |= file_changed; + stats.row_groups += row_groups; + stats.row_groups_pruned += row_groups_pruned; + stats.rows += rows; + stats.rows_pruned += rows_pruned; + match file { + Some(file) => files.push(file), + None => stats.files_pruned += 1, + } + } + } + } + file_groups.push(files); + } + + let summary = EagerPruningSummary::Pruned(stats); + if !changed { + // Nothing is pruned: keep the original files and (possibly exact) + // statistics + return Ok((conf, summary)); + } + + let table_schema = table_schema.table_schema(); + let all_files_have_statistics = file_groups + .iter() + .flatten() + .all(|file| file.statistics.is_some()); + let file_groups = file_groups + .into_iter() + .zip(&conf.file_groups) + .map(|(files, original_group)| { + let statistics = if all_files_have_statistics { + Some(group_statistics(&files, table_schema)?) + } else { + // Files may have been removed from the group, so its + // original statistics are an upper bound + original_group + .file_statistics(None) + .map(|statistics| statistics.clone().to_inexact()) + }; + let group = FileGroup::new(files); + Ok(match statistics { + Some(statistics) => group.with_statistics(Arc::new(statistics)), + None => group, + }) + }) + .collect::>>()?; + + let statistics = if all_files_have_statistics { + let mut non_empty_groups = file_groups + .iter() + .filter(|group| !group.is_empty()) + .filter_map(|group| group.file_statistics(None)) + .peekable(); + if non_empty_groups.peek().is_none() { + empty_statistics(table_schema) + } else { + Statistics::try_merge_iter_with_ndv_fallback( + non_empty_groups, + table_schema, + NdvFallback::Max, + )? + .to_inexact() + } + } else { + // The number of rows is unknown for some files: keep the statistics + // of the scan, which are now an upper bound + conf.statistics().to_inexact() + }; + + let conf = FileScanConfigBuilder::from(conf) + .with_file_groups(file_groups) + .with_statistics(statistics) + .build(); + Ok((conf, summary)) + } + + /// Returns why eager pruning cannot be performed for a scan, if it cannot + fn skip_reason( + &self, + table_schema: &TableSchema, + num_files: usize, + ) -> Option { + let file_limit = self.options.global.eager_pruning_file_limit; + if num_files > file_limit { + return Some(format!( + "{num_files} files exceed eager_pruning_file_limit {file_limit}" + )); + } + let crypto = &self.options.crypto; + if crypto.file_decryption.is_some() || crypto.factory_id.is_some() { + return Some("encrypted files are not supported".to_string()); + } + if !table_schema.virtual_columns().is_empty() { + return Some("virtual columns are not supported".to_string()); + } + None + } + + /// Eagerly prune `file`, keeping the file unchanged if that fails. + async fn prune_file_or_keep( + &self, + context: &ScanContext<'_>, + file: PartitionedFile, + ) -> FileOutcome { + match self.prune_file(context, &file).await { + Ok(outcome) => outcome, + Err(e) => { + debug!( + "Eager parquet pruning of {} failed, leaving it unchanged: {e}", + file.object_meta.location + ); + FileOutcome::NotEvaluated(file) + } + } + } + + /// Eagerly prune a single file. + /// + /// This mirrors the pruning the parquet opener performs when the file is + /// scanned. + async fn prune_file( + &self, + context: &ScanContext<'_>, + file: &PartitionedFile, + ) -> Result { + if file.extensions.contains::() { + // A ParquetAccessPlan cannot be attached next to a row selection + return Ok(FileOutcome::NotEvaluated(file.clone())); + } + + let global = &self.options.global; + let table_schema = context.table_schema; + let logical_file_schema = table_schema.file_schema(); + let file_name = file.object_meta.location.to_string(); + + // Replace partition columns and columns that are constant in the file + // with their values, as the opener does + let mut literal_columns: HashMap = table_schema + .table_partition_cols() + .iter() + .zip(file.partition_values.iter()) + .map(|(field, value)| (field.name().clone(), value.clone())) + .collect(); + literal_columns.extend(constant_columns_from_stats( + file.statistics.as_deref(), + logical_file_schema, + )); + let mut predicate = Arc::clone(&self.predicate); + if !literal_columns.is_empty() { + predicate = replace_columns_with_literals(predicate, &literal_columns)?; + } + + // Metrics are required by the pruning code but are not reported: the + // scan reports the pruning it performs itself when it executes + let metrics = ExecutionPlanMetricsSet::new(); + let file_metrics = ParquetFileMetrics::new(0, &file_name, &metrics); + let mut reader: Box = self.reader_factory.create_reader( + 0, + file.clone(), + file.metadata_size_hint.or(global.metadata_size_hint), + &metrics, + )?; + + // Like the opener, defer loading the page index until it is needed + let mut options = + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Skip); + if let Some(schema) = file.arrow_schema.as_ref() { + options = options.with_schema(Arc::clone(schema)); + } + let mut reader_metadata = + ArrowReaderMetadata::load_async(&mut reader, options.clone()).await?; + + let physical_file_schema = coerce_physical_file_schema( + logical_file_schema, + reader_metadata.schema(), + reader_metadata.parquet_schema(), + context.coerce_int96.as_ref(), + context.coerce_int96_tz.clone(), + ) + .unwrap_or_else(|| Arc::clone(reader_metadata.schema())); + + // Adapt the predicate to the physical file schema + let rewriter = context.expr_adapter_factory.create( + Arc::clone(logical_file_schema), + Arc::clone(&physical_file_schema), + )?; + let predicate = PhysicalExprSimplifier::new(&physical_file_schema) + .simplify(rewriter.rewrite(predicate)?)?; + + let parquet_metadata = Arc::clone(reader_metadata.metadata()); + let rg_metadata = parquet_metadata.row_groups(); + let mut row_groups = RowGroupAccessPlanFilter::new(create_initial_plan( + &file_name, + &file.extensions, + rg_metadata, + )?); + if let Some(range) = file.range.as_ref() { + row_groups.prune_by_range(rg_metadata, range); + } + let initial_plan = row_groups.access_plan().clone(); + let initial_row_groups = read_row_group_count(&initial_plan, rg_metadata); + let initial_rows = initial_plan.selected_row_count(rg_metadata)?; + + let predicate_creation_errors = Count::new(); + let pruning_predicate = build_pruning_predicates( + Some(&predicate), + &physical_file_schema, + &predicate_creation_errors, + global.max_in_list_size, + ); + let page_pruning_predicate = + (self.level.includes(EagerParquetPruning::PageIndex) + && global.enable_page_index) + .then(|| { + build_page_pruning_predicate( + &predicate, + &physical_file_schema, + global.max_in_list_size, + ) + }) + .filter(|p| p.filter_number() > 0); + + // Row group statistics + if let Some(pruning_predicate) = pruning_predicate.as_deref() + && global.pruning + && !row_groups.is_empty() + { + row_groups.prune_by_statistics_with_metadata( + &physical_file_schema, + &parquet_metadata, + pruning_predicate, + &file_metrics, + ); + } + + // Load the page index if it can prune any remaining row group. This + // must happen before reading bloom filters, which consumes the reader. + if page_pruning_predicate.is_some() && !row_groups.is_empty() { + let fully_matched = row_groups.is_fully_matched(); + if !row_groups.row_group_indexes().all(|idx| fully_matched[idx]) { + reader_metadata = load_page_index( + reader_metadata, + &mut reader, + options.with_page_index_policy(PageIndexPolicy::Optional), + ) + .await?; + } + } + + // Bloom filters + if let Some(pruning_predicate) = pruning_predicate.as_deref() + && self.level.includes(EagerParquetPruning::BloomFilters) + && global.bloom_filter_on_read + && !row_groups.is_empty() + { + let mut builder = ParquetRecordBatchStreamBuilder::new_with_metadata( + reader, + reader_metadata.clone(), + ); + // Statistics have already proved that every row in a fully matched + // group matches the predicate, so its Bloom filters cannot prune it. + let row_group_indexes: Vec = row_groups + .row_group_indexes() + .filter(|idx| !row_groups.access_plan().is_fully_matched(*idx)) + .collect(); + let bloom_filters = load_row_group_bloom_filters( + &mut builder, + pruning_predicate, + &physical_file_schema, + &row_group_indexes, + &file_metrics.predicate_evaluation_errors, + ) + .await; + row_groups.prune_by_bloom_filters( + pruning_predicate, + &file_metrics, + &bloom_filters, + ); + } + + let mut access_plan = row_groups.build(); + + // Page index + if let Some(page_pruning_predicate) = page_pruning_predicate + && access_plan.row_group_index_iter().next().is_some() + { + access_plan = page_pruning_predicate.prune_plan_with_page_index( + access_plan, + &physical_file_schema, + reader_metadata.parquet_schema(), + reader_metadata.metadata(), + &file_metrics, + ); + } + + // Fully matched flags only hold for this predicate, see module docs + access_plan.clear_fully_matched(); + + let rows = access_plan.selected_row_count(rg_metadata)?; + let row_groups_read = read_row_group_count(&access_plan, rg_metadata); + let outcome = |file, changed| FileOutcome::Evaluated { + file, + changed, + row_groups: initial_row_groups, + row_groups_pruned: initial_row_groups.saturating_sub(row_groups_read), + rows: initial_rows, + rows_pruned: initial_rows.saturating_sub(rows), + }; + + if rows == 0 { + return Ok(outcome(None, true)); + } + if rows == initial_rows { + // Nothing is pruned, so the file is kept unchanged. The access plan + // can still differ from the initial plan, as page index pruning + // selects all rows of a row group in which no page is pruned. + if file + .statistics + .as_ref() + .is_some_and(|statistics| statistics.num_rows.get_value().is_some()) + { + return Ok(outcome(Some(file.clone()), false)); + } + // The row count of the file is unknown, e.g. because statistics are + // not collected. Use the metadata that was read to fill it in, so + // the row count of the scan does not become unknown. + let file_statistics = DFParquetMetadata::statistics_from_parquet_metadata( + &parquet_metadata, + logical_file_schema, + )?; + let file = file.clone().with_statistics(Arc::new(file_statistics)); + return Ok(outcome(Some(file), true)); + } + + // Summarize only the row groups that are read, so that column + // statistics such as min/max values describe the rows that are read + // rather than the entire file. Otherwise estimates derived from them, + // such as filter selectivity, would account for the pruning twice. + let read_row_groups: Vec = access_plan + .row_group_index_iter() + .map(|idx| rg_metadata[idx].clone()) + .collect(); + let read_row_groups_rows = read_row_groups + .iter() + .map(|row_group| row_group.num_rows() as usize) + .sum(); + let row_group_statistics = DFParquetMetadata::statistics_from_row_groups( + &parquet_metadata, + &read_row_groups, + logical_file_schema, + )?; + // `with_statistics` also adds statistics for partition columns + let mut file = file.clone().with_statistics(Arc::new(row_group_statistics)); + if let Some(statistics) = file.statistics.as_deref() { + file.statistics = Some(Arc::new(pruned_statistics( + statistics, + read_row_groups_rows, + rows, + ))); + } + file.extensions.insert(access_plan); + + Ok(outcome(Some(file), true)) + } +} + +/// Returns the number of row groups from which `access_plan` reads any rows +fn read_row_group_count( + access_plan: &ParquetAccessPlan, + rg_metadata: &[RowGroupMetaData], +) -> usize { + access_plan + .inner() + .iter() + .zip(rg_metadata) + .filter(|(access, row_group)| match access { + RowGroupAccess::Skip => false, + RowGroupAccess::Scan => row_group.num_rows() > 0, + RowGroupAccess::Selection(selection) => selection.selects_any(), + }) + .count() +} + +/// Returns statistics describing the `selected_rows` rows read from row +/// groups with `row_group_rows` rows, given the `statistics` of these row +/// groups. +/// +/// The number of rows is the (upper bound) estimate from pruning, and byte +/// sizes are scaled accordingly. All values are inexact, except a null count +/// of zero. +/// +/// The statistics must stay correct for the entire file if the attached +/// access plan is not applied, for example after the plan is serialized (file +/// statistics are serialized, file extensions are not). Min and max values +/// only bound the rows that are read, so they must not be exact: the scan +/// replaces a column whose exact min and max are equal with that value, which +/// is wrong for row groups the statistics do not cover. A null count of zero +/// stays exact because every row that is not read fails the filters, which are +/// always applied to the output of the scan. +fn pruned_statistics( + statistics: &Statistics, + row_group_rows: usize, + selected_rows: usize, +) -> Statistics { + let scale = |value: &Precision| match value.get_value() { + Some(value) if row_group_rows > 0 => { + let scaled = + (*value as u128 * selected_rows as u128).div_ceil(row_group_rows as u128); + Precision::Inexact(usize::try_from(scaled).unwrap_or(usize::MAX)) + } + _ => value.to_inexact(), + }; + + let column_statistics = statistics + .column_statistics + .iter() + .map(|column| { + ColumnStatistics { + null_count: if column.null_count == Precision::Exact(0) { + Precision::Exact(0) + } else { + column.null_count.to_inexact() + }, + min_value: column.min_value.clone().to_inexact(), + max_value: column.max_value.clone().to_inexact(), + // The sum of the entire file says little about the sum of the + // selected rows + sum_value: Precision::Absent, + distinct_count: column.distinct_count.to_inexact(), + byte_size: scale(&column.byte_size), + } + }) + .collect(); + + Statistics { + num_rows: Precision::Inexact(selected_rows), + total_byte_size: scale(&statistics.total_byte_size), + column_statistics, + } +} + +/// Merges the statistics of the files in a file group +fn group_statistics( + files: &[PartitionedFile], + table_schema: &Schema, +) -> Result { + if files.is_empty() { + return Ok(empty_statistics(table_schema)); + } + Statistics::try_merge_iter_with_ndv_fallback( + files.iter().filter_map(|file| file.statistics.as_deref()), + table_schema, + NdvFallback::Max, + ) +} + +/// Statistics for a scan, or file group, from which no rows are read +fn empty_statistics(table_schema: &Schema) -> Statistics { + Statistics::new_unknown(table_schema) + .with_num_rows(Precision::Inexact(0)) + .with_total_byte_size(Precision::Inexact(0)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::DefaultParquetFileReaderFactory; + use crate::source::ParquetSource; + + use arrow::array::{Int64Array, RecordBatch}; + use arrow::datatypes::{DataType, Field, SchemaRef}; + use bytes::{BufMut, BytesMut}; + use datafusion_datasource::{FileRange, TableSchemaBuilder}; + use datafusion_execution::object_store::ObjectStoreUrl; + use datafusion_expr::{Expr, col, lit}; + use datafusion_physical_expr::planner::logical2physical; + use object_store::memory::InMemory; + use object_store::path::Path; + use object_store::{ObjectStore, ObjectStoreExt}; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; + use parquet::file::properties::WriterProperties; + + /// A parquet file whose column `a` holds `0..400` in 4 row groups of 100 + /// rows, in an in memory object store + struct TestFile { + store: Arc, + schema: SchemaRef, + size: u64, + } + + impl TestFile { + async fn new() -> Self { + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from_iter_values(0..400))], + ) + .unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(100)) + .build(); + let mut out = BytesMut::new().writer(); + { + let mut writer = + ArrowWriter::try_new(&mut out, Arc::clone(&schema), Some(props)) + .unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + } + let data = out.into_inner().freeze(); + let size = data.len() as u64; + let store: Arc = Arc::new(InMemory::new()); + store + .put(&Path::from("test.parquet"), data.into()) + .await + .unwrap(); + Self { + store, + schema, + size, + } + } + + fn file(&self) -> PartitionedFile { + PartitionedFile::new("test.parquet", self.size) + } + + fn table_schema(&self) -> TableSchema { + TableSchema::from(&self.schema) + } + + /// Prunes `files` as a single file group + async fn prune_files( + &self, + options: &TableParquetOptions, + table_schema: &TableSchema, + predicate: Expr, + files: Vec, + ) -> (FileScanConfig, EagerPruningSummary) { + let conf = FileScanConfigBuilder::new( + ObjectStoreUrl::local_filesystem(), + Arc::new(ParquetSource::new(table_schema.clone())), + ) + .with_file_group(FileGroup::new(files)) + .build(); + + let filters = [logical2physical(&predicate, &self.schema)]; + let pruner = EagerPruner::try_new( + options, + &filters, + Arc::new(DefaultParquetFileReaderFactory::new(Arc::clone( + &self.store, + ))), + 4, + ) + .unwrap(); + pruner.prune(table_schema, conf).await.unwrap() + } + } + + fn eager_options(level: EagerParquetPruning) -> TableParquetOptions { + let mut options = TableParquetOptions::default(); + options.global.eager_pruning = level; + options + } + + /// Prunes the test file using `predicate` + async fn prune( + level: EagerParquetPruning, + file_limit: usize, + predicate: Expr, + ) -> (FileScanConfig, EagerPruningSummary) { + let test_file = TestFile::new().await; + let mut options = eager_options(level); + options.global.eager_pruning_file_limit = file_limit; + let files = vec![test_file.file()]; + test_file + .prune_files(&options, &test_file.table_schema(), predicate, files) + .await + } + + #[tokio::test] + async fn prunes_row_groups_without_fully_matched_row_groups() { + // Row groups 0 and 1 fully match `a < 250`, row group 2 partially matches + let (conf, summary) = + prune(EagerParquetPruning::RowGroups, 1, col("a").lt(lit(250i64))).await; + + assert_eq!( + summary, + EagerPruningSummary::Pruned(EagerPruningStats { + level: EagerParquetPruning::RowGroups, + files: 1, + files_not_evaluated: 0, + files_pruned: 0, + row_groups: 4, + row_groups_pruned: 1, + rows: 400, + rows_pruned: 100, + }) + ); + + let file = &conf.file_groups[0].files()[0]; + let access_plan = file.extensions.get::().unwrap(); + assert_eq!(access_plan.row_group_indexes(), vec![0, 1, 2]); + // The scan may execute with a different predicate, so no row group + // may be marked as fully matched + assert!( + access_plan.fully_matched().iter().all(|matched| !matched), + "{access_plan:?}" + ); + + // Statistics describe the row groups that are read + let file_statistics = file.statistics.as_ref().unwrap(); + assert_eq!(file_statistics.num_rows, Precision::Inexact(300)); + let column_statistics = &file_statistics.column_statistics[0]; + assert_eq!( + column_statistics.min_value, + Precision::Inexact(ScalarValue::Int64(Some(0))) + ); + assert_eq!( + column_statistics.max_value, + Precision::Inexact(ScalarValue::Int64(Some(299))) + ); + assert_eq!(column_statistics.null_count, Precision::Exact(0)); + assert_eq!(conf.statistics().num_rows, Precision::Inexact(300)); + } + + #[tokio::test] + async fn prunes_entire_file() { + let (conf, summary) = + prune(EagerParquetPruning::RowGroups, 1, col("a").gt(lit(1000i64))).await; + + let EagerPruningSummary::Pruned(stats) = summary else { + panic!("expected eager pruning"); + }; + assert_eq!(stats.files_pruned, 1); + assert_eq!(conf.file_groups.len(), 1); + assert!(conf.file_groups[0].is_empty()); + assert_eq!(conf.statistics().num_rows, Precision::Inexact(0)); + assert_eq!( + conf.file_groups[0].file_statistics(None).unwrap().num_rows, + Precision::Inexact(0) + ); + } + + #[tokio::test] + async fn skips_scans_above_file_limit() { + let (conf, summary) = + prune(EagerParquetPruning::RowGroups, 0, col("a").lt(lit(250i64))).await; + + assert_eq!( + summary, + EagerPruningSummary::Skipped { + level: EagerParquetPruning::RowGroups, + reason: "1 files exceed eager_pruning_file_limit 0".to_string(), + } + ); + let file = &conf.file_groups[0].files()[0]; + assert!(!file.extensions.contains::()); + } + + #[tokio::test] + async fn skips_encrypted_files() { + let test_file = TestFile::new().await; + let mut options = eager_options(EagerParquetPruning::RowGroups); + options.crypto.factory_id = Some("test_factory".to_string()); + let files = vec![test_file.file()]; + + let (conf, summary) = test_file + .prune_files( + &options, + &test_file.table_schema(), + col("a").lt(lit(250i64)), + files, + ) + .await; + + assert_eq!( + summary, + EagerPruningSummary::Skipped { + level: EagerParquetPruning::RowGroups, + reason: "encrypted files are not supported".to_string(), + } + ); + assert!( + !conf.file_groups[0].files()[0] + .extensions + .contains::() + ); + } + + #[tokio::test] + async fn skips_scans_with_virtual_columns() { + let test_file = TestFile::new().await; + let table_schema = TableSchemaBuilder::new(Arc::clone(&test_file.schema)) + .with_virtual_columns(vec![Arc::new(Field::new( + "row_number", + DataType::Int64, + true, + ))]) + .build(); + let files = vec![test_file.file()]; + + let (_, summary) = test_file + .prune_files( + &eager_options(EagerParquetPruning::RowGroups), + &table_schema, + col("a").lt(lit(250i64)), + files, + ) + .await; + + assert_eq!( + summary, + EagerPruningSummary::Skipped { + level: EagerParquetPruning::RowGroups, + reason: "virtual columns are not supported".to_string(), + } + ); + } + + #[tokio::test] + async fn keeps_files_that_cannot_be_read() { + let test_file = TestFile::new().await; + // The first file does not exist, so its metadata cannot be read + let files = vec![ + PartitionedFile::new("missing.parquet", 1024), + test_file.file(), + ]; + + let (conf, summary) = test_file + .prune_files( + &eager_options(EagerParquetPruning::RowGroups), + &test_file.table_schema(), + col("a").lt(lit(250i64)), + files, + ) + .await; + + let EagerPruningSummary::Pruned(stats) = summary else { + panic!("expected eager pruning"); + }; + assert_eq!(stats.files, 2); + assert_eq!(stats.files_not_evaluated, 1); + assert_eq!(stats.rows_pruned, 100); + + // Both files are kept, the one that could not be read unchanged + let files = conf.file_groups[0].files(); + assert_eq!(files.len(), 2); + assert!(!files[0].extensions.contains::()); + assert!(files[0].statistics.is_none()); + assert!(files[1].extensions.contains::()); + // The row count of one file is unknown, so it is unknown for the scan + assert_eq!(conf.statistics().num_rows, Precision::Absent); + } + + #[tokio::test] + async fn skips_files_with_a_row_selection() { + let test_file = TestFile::new().await; + let selection = ParquetRowSelection::new(RowSelection::from(vec![ + RowSelector::select(200), + RowSelector::skip(200), + ])); + let files = vec![test_file.file().with_extension(selection)]; + + let (conf, summary) = test_file + .prune_files( + &eager_options(EagerParquetPruning::RowGroups), + &test_file.table_schema(), + col("a").lt(lit(250i64)), + files, + ) + .await; + + let EagerPruningSummary::Pruned(stats) = summary else { + panic!("expected eager pruning"); + }; + assert_eq!(stats.files_not_evaluated, 1); + assert_eq!(stats.rows_pruned, 0); + let file = &conf.file_groups[0].files()[0]; + assert!(!file.extensions.contains::()); + assert!(file.extensions.contains::()); + } + + #[tokio::test] + async fn prunes_row_groups_within_a_file_range() { + let test_file = TestFile::new().await; + // A range that covers the first part of the file only + let mut file = test_file.file(); + file.range = Some(FileRange { + start: 0, + end: (test_file.size / 2) as i64, + }); + let files = vec![file]; + + let (conf, summary) = test_file + .prune_files( + &eager_options(EagerParquetPruning::RowGroups), + &test_file.table_schema(), + col("a").lt(lit(150i64)), + files, + ) + .await; + + let EagerPruningSummary::Pruned(stats) = summary else { + panic!("expected eager pruning"); + }; + // Row groups outside the range are not part of the evaluated rows + assert!(stats.rows < 400, "{stats:?}"); + assert_eq!(stats.rows_pruned, stats.rows - 200); + let access_plan = conf.file_groups[0].files()[0] + .extensions + .get::() + .unwrap(); + assert_eq!(access_plan.row_group_indexes(), vec![0, 1]); + } + + #[tokio::test] + async fn with_int96_coercion_and_file_arrow_schema() { + let test_file = TestFile::new().await; + let mut options = eager_options(EagerParquetPruning::RowGroups); + options.global.coerce_int96 = Some("ms".to_string()); + options.global.coerce_int96_tz = Some("UTC".to_string()); + let mut file = test_file.file(); + file.arrow_schema = Some(Arc::clone(&test_file.schema)); + let files = vec![file]; + + let (conf, summary) = test_file + .prune_files( + &options, + &test_file.table_schema(), + col("a").lt(lit(250i64)), + files, + ) + .await; + + let EagerPruningSummary::Pruned(stats) = summary else { + panic!("expected eager pruning"); + }; + assert_eq!(stats.rows_pruned, 100); + assert_eq!(conf.statistics().num_rows, Precision::Inexact(300)); + } + + #[test] + fn disabled_or_without_filters() { + let reader_factory: Arc = Arc::new( + DefaultParquetFileReaderFactory::new(Arc::new(InMemory::new())), + ); + let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]); + let filters = [logical2physical(&col("a").lt(lit(1i64)), &schema)]; + + let mut options = TableParquetOptions::default(); + assert!( + EagerPruner::try_new(&options, &filters, Arc::clone(&reader_factory), 1) + .is_none() + ); + + options.global.eager_pruning = EagerParquetPruning::BloomFilters; + assert!( + EagerPruner::try_new(&options, &[], Arc::clone(&reader_factory), 1).is_none() + ); + assert!(EagerPruner::try_new(&options, &filters, reader_factory, 1).is_some()); + } + + #[test] + fn statistics_of_pruned_file() { + let column = |min: i64, max: i64, null_count: usize| ColumnStatistics { + null_count: Precision::Exact(null_count), + min_value: Precision::Exact(ScalarValue::Int64(Some(min))), + max_value: Precision::Exact(ScalarValue::Int64(Some(max))), + sum_value: Precision::Exact(ScalarValue::Int64(Some(1000))), + distinct_count: Precision::Exact(10), + byte_size: Precision::Exact(800), + }; + let statistics = Statistics { + num_rows: Precision::Exact(100), + total_byte_size: Precision::Exact(1000), + column_statistics: vec![column(5, 5, 0), column(1, 10, 0), column(1, 10, 3)], + }; + + let pruned = pruned_statistics(&statistics, 100, 25); + + assert_eq!(pruned.num_rows, Precision::Inexact(25)); + assert_eq!(pruned.total_byte_size, Precision::Inexact(250)); + let [single_value, no_nulls, nulls] = pruned.column_statistics.as_slice() else { + panic!("expected 3 columns"); + }; + + // Even a single value is inexact: it only describes the rows that are + // read, and an exact single value would let the scan replace the + // column with it in rows it reads without the attached access plan + assert_eq!(single_value.null_count, Precision::Exact(0)); + assert_eq!( + single_value.min_value, + Precision::Inexact(ScalarValue::Int64(Some(5))) + ); + assert_eq!( + single_value.max_value, + Precision::Inexact(ScalarValue::Int64(Some(5))) + ); + + assert_eq!(no_nulls.null_count, Precision::Exact(0)); + assert_eq!( + no_nulls.min_value, + Precision::Inexact(ScalarValue::Int64(Some(1))) + ); + assert_eq!( + no_nulls.max_value, + Precision::Inexact(ScalarValue::Int64(Some(10))) + ); + + assert_eq!(nulls.null_count, Precision::Inexact(3)); + assert_eq!( + nulls.min_value, + Precision::Inexact(ScalarValue::Int64(Some(1))) + ); + + for column in &pruned.column_statistics { + assert_eq!(column.sum_value, Precision::Absent); + assert_eq!(column.distinct_count, Precision::Inexact(10)); + assert_eq!(column.byte_size, Precision::Inexact(200)); + } + } + + #[test] + fn display() { + let stats = EagerPruningStats { + level: EagerParquetPruning::PageIndex, + files: 4, + files_not_evaluated: 0, + files_pruned: 1, + row_groups: 120, + row_groups_pruned: 117, + rows: 1000, + rows_pruned: 990, + }; + assert_eq!( + EagerPruningSummary::Pruned(stats.clone()).to_string(), + "[level=page_index, files_pruned=1/4, row_groups_pruned=117/120, \ + rows_pruned=990/1000]" + ); + assert_eq!( + EagerPruningSummary::Pruned(EagerPruningStats { + files_not_evaluated: 2, + ..stats + }) + .to_string(), + "[level=page_index, files_pruned=1/4, row_groups_pruned=117/120, \ + rows_pruned=990/1000, files_not_evaluated=2]" + ); + assert_eq!( + EagerPruningSummary::Skipped { + level: EagerParquetPruning::RowGroups, + reason: "some reason".to_string(), + } + .to_string(), + "[level=row_groups, skipped=some reason]" + ); + } +} diff --git a/datafusion/datasource-parquet/src/file_format.rs b/datafusion/datasource-parquet/src/file_format.rs index dc0f9a7be434e..ec8fc82950c87 100644 --- a/datafusion/datasource-parquet/src/file_format.rs +++ b/datafusion/datasource-parquet/src/file_format.rs @@ -56,9 +56,11 @@ use datafusion_datasource::sink::DataSinkExec; use datafusion_datasource::write::get_writer_schema; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr_common::sort_expr::{LexOrdering, LexRequirement}; -use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; use datafusion_session::Session; +use crate::ParquetFileReaderFactory; +use crate::eager_pruning::EagerPruner; use crate::metadata::{DFParquetMetadata, lex_ordering_to_sorting_columns}; use crate::reader::CachedParquetFileReaderFactory; use crate::source::{ @@ -496,6 +498,7 @@ impl FileFormat for ParquetFormat { &self, state: &dyn Session, conf: FileScanConfig, + filters: &[Arc], ) -> Result> { let mut metadata_size_hint = None; @@ -515,9 +518,10 @@ impl FileFormat for ParquetFormat { let store = state .runtime_env() .object_store(conf.object_store_url.clone())?; - let cached_parquet_read_factory = + let cached_parquet_read_factory: Arc = Arc::new(CachedParquetFileReaderFactory::new(store, metadata_cache)); - source = source.with_parquet_file_reader_factory(cached_parquet_read_factory); + source = source + .with_parquet_file_reader_factory(Arc::clone(&cached_parquet_read_factory)); if let Some(metadata_size_hint) = metadata_size_hint { source = source.with_metadata_size_hint(metadata_size_hint) @@ -525,6 +529,25 @@ impl FileFormat for ParquetFormat { source = self.set_source_encryption_factory(source, state)?; + // Eagerly prune files using the filters, to refine the statistics + // of the scan (see `EagerParquetPruning`) + let mut conf = conf; + if let Some(pruner) = EagerPruner::try_new( + &self.options, + filters, + cached_parquet_read_factory, + state + .config_options() + .execution + .meta_fetch_concurrency + .get(), + ) { + let (pruned_conf, summary) = + pruner.prune(source.table_schema(), conf).await?; + conf = pruned_conf; + source = source.with_eager_pruning_summary(summary); + } + let conf = FileScanConfigBuilder::from(conf) .with_source(Arc::new(source)) .build(); @@ -681,6 +704,10 @@ impl From<&ParquetFormatFactory> for protobuf::TableParquetOptions { }), max_row_group_size: global_options.global.max_row_group_size as u64, max_in_list_size: global_options.global.max_in_list_size as u64, + eager_pruning: global_options.global.eager_pruning.to_string(), + eager_pruning_file_limit_opt: Some(parquet_options::EagerPruningFileLimitOpt::EagerPruningFileLimit( + global_options.global.eager_pruning_file_limit as u64, + )), created_by: global_options.global.created_by.clone(), column_index_truncate_length_opt: global_options.global.column_index_truncate_length.map(|length| { parquet_options::ColumnIndexTruncateLengthOpt::ColumnIndexTruncateLength(length as u64) diff --git a/datafusion/datasource-parquet/src/metadata.rs b/datafusion/datasource-parquet/src/metadata.rs index 51dc98e4f2f02..2fd9da1dbb56e 100644 --- a/datafusion/datasource-parquet/src/metadata.rs +++ b/datafusion/datasource-parquet/src/metadata.rs @@ -508,8 +508,21 @@ impl<'a> DFParquetMetadata<'a> { metadata: &ParquetMetaData, logical_file_schema: &SchemaRef, ) -> Result { - let row_groups_metadata = metadata.row_groups(); + Self::statistics_from_row_groups( + metadata, + metadata.row_groups(), + logical_file_schema, + ) + } + /// Like [`Self::statistics_from_parquet_metadata`], but only summarizes + /// `row_groups_metadata`, a subset of the row groups of the file described + /// by `metadata`. + pub(crate) fn statistics_from_row_groups( + metadata: &ParquetMetaData, + row_groups_metadata: &[RowGroupMetaData], + logical_file_schema: &SchemaRef, + ) -> Result { // Use Statistics::default() as opposed to Statistics::new_unknown() // because we are going to replace the column statistics below // and we don't want to initialize them twice. diff --git a/datafusion/datasource-parquet/src/mod.rs b/datafusion/datasource-parquet/src/mod.rs index 80d2205a1fab0..40c05aece73b1 100644 --- a/datafusion/datasource-parquet/src/mod.rs +++ b/datafusion/datasource-parquet/src/mod.rs @@ -27,6 +27,7 @@ pub mod access_plan; mod bloom_filter; mod decoder_projection; +mod eager_pruning; pub mod file_format; pub mod metadata; mod metrics; @@ -53,6 +54,7 @@ mod writer; pub use access_plan::{ParquetAccessPlan, ParquetRowSelection, RowGroupAccess}; pub use bloom_filter::BloomFilterStatistics; +pub use eager_pruning::{EagerPruningStats, EagerPruningSummary}; pub use file_format::*; pub use metrics::ParquetFileMetrics; pub use page_filter::PagePruningAccessPlanFilter; diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index eb329fb3037c0..c240cc6a69e3a 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -24,6 +24,7 @@ use self::early_stop::EarlyStoppingStream; #[cfg(feature = "parquet_encryption")] use self::encryption::EncryptionContext; use crate::access_plan::PreparedAccessPlan; +use crate::bloom_filter::load_row_group_bloom_filters; use crate::decoder_projection::DecoderProjection; use crate::metrics::{ByteProgress, RowFilterSkippedFullyMatchedMetric}; use crate::page_filter::PagePruningAccessPlanFilter; @@ -32,10 +33,10 @@ use crate::push_decoder::{ RowFilterContext, RowGroupPruner, }; use crate::row_group_filter::{RowGroupAccessPlanFilter, row_group_in_range}; +use crate::schema_coercion::coerce_physical_file_schema; use crate::{ - BloomFilterStatistics, Int96Coercer, ParquetAccessPlan, ParquetFileMetrics, + BloomFilterStatistics, ParquetAccessPlan, ParquetFileMetrics, ParquetFileReaderFactory, ParquetRowSelection, ParquetVirtualColumn, RowGroupAccess, - apply_file_schema_type_coercions, }; use arrow::array::RecordBatch; use arrow::datatypes::DataType; @@ -74,7 +75,6 @@ use datafusion_common::config::EncryptionFactoryOptions; #[cfg(feature = "parquet_encryption")] use datafusion_execution::parquet_encryption::EncryptionFactory; use futures::{FutureExt, StreamExt, future::BoxFuture, stream::BoxStream}; -use log::debug; use parquet::arrow::ParquetRecordBatchStreamBuilder; use parquet::arrow::arrow_reader::metrics::ArrowReaderMetrics; use parquet::arrow::arrow_reader::{ @@ -82,8 +82,6 @@ use parquet::arrow::arrow_reader::{ }; use parquet::arrow::async_reader::AsyncFileReader; use parquet::arrow::parquet_column; -use parquet::basic::Type; -use parquet::bloom_filter::Sbbf; use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader, RowGroupMetaData}; /// Morselizer-level state for virtual columns, precomputed once per scan @@ -1109,25 +1107,14 @@ impl MetadataLoadedParquetOpen { // desired schema (for example if we want to instruct the parquet // reader to read strings using Utf8View instead). Update if necessary let mut metadata_dirty = false; - if let Some(merged) = apply_file_schema_type_coercions( + if let Some(coerced) = coerce_physical_file_schema( &prepared.logical_file_schema, &physical_file_schema, + reader_metadata.parquet_schema(), + prepared.coerce_int96.as_ref(), + prepared.coerce_int96_tz.clone(), ) { - physical_file_schema = Arc::new(merged); - options = options.with_schema(Arc::clone(&physical_file_schema)); - metadata_dirty = true; - } - - if let Some(ref coerce) = prepared.coerce_int96 - && let Some(merged) = Int96Coercer::new( - reader_metadata.parquet_schema(), - &physical_file_schema, - coerce, - ) - .with_timezone(prepared.coerce_int96_tz.clone()) - .coerce() - { - physical_file_schema = Arc::new(merged); + physical_file_schema = coerced; options = options.with_schema(Arc::clone(&physical_file_schema)); metadata_dirty = true; } @@ -1505,57 +1492,21 @@ impl RowGroupsPrunedParquetOpen { mem::replace(&mut prepared.async_file_reader, replacement_reader), reader_metadata, ); - let parquet_columns: Vec<(String, usize, Type, i32)> = predicate - .literal_columns() - .into_iter() - .filter_map(|column_name| { - let parquet_schema = builder.parquet_schema(); - let (column_idx, _) = parquet_column( - parquet_schema, - &prepared.physical_file_schema, - &column_name, - )?; - Some(( - column_name, - column_idx, - parquet_schema.column(column_idx).physical_type(), - parquet_schema.column(column_idx).type_length(), - )) - }) + // Statistics have already proved that every row in a fully matched + // group matches the predicate, so its Bloom filters cannot prune it. + let row_group_indexes: Vec = self + .row_groups + .row_group_indexes() + .filter(|idx| !self.row_groups.access_plan().is_fully_matched(*idx)) .collect(); - - for idx in self.row_groups.row_group_indexes() { - // Statistics have already proved that every row in this group - // matches the predicate, so its Bloom filters cannot prune it. - if self.row_groups.access_plan().is_fully_matched(idx) { - continue; - } - let mut row_group_filters = - BloomFilterStatistics::with_capacity(parquet_columns.len()); - for (column_name, column_idx, physical_type, type_length) in - &parquet_columns - { - let bf: Sbbf = match builder - .get_row_group_column_bloom_filter(idx, *column_idx) - .await - { - Ok(Some(bf)) => bf, - Ok(None) => continue, - Err(e) => { - debug!("Ignoring error reading bloom filter: {e}"); - prepared.file_metrics.predicate_evaluation_errors.add(1); - continue; - } - }; - row_group_filters.insert( - column_name, - bf, - *physical_type, - *type_length, - ); - } - row_group_bloom_filters[idx] = row_group_filters; - } + row_group_bloom_filters = load_row_group_bloom_filters( + &mut builder, + predicate, + &prepared.physical_file_schema, + &row_group_indexes, + &prepared.file_metrics.predicate_evaluation_errors, + ) + .await; } Ok(BloomFiltersLoadedParquetOpen { @@ -1923,7 +1874,7 @@ fn row_group_bytes(rg_meta: &RowGroupMetaData) -> u64 { u64::try_from(rg_meta.compressed_size()).unwrap_or(0) } -type ConstantColumns = HashMap; +pub(crate) type ConstantColumns = HashMap; /// True when substituting `stats_constants` — the columns this file's /// statistics prove constant — is on its own enough to collapse `predicate` @@ -1958,7 +1909,7 @@ fn stats_alone_unsatisfiable( } /// Extract constant column values from statistics, keyed by column name in the logical file schema. -fn constant_columns_from_stats( +pub(crate) fn constant_columns_from_stats( statistics: Option<&Statistics>, file_schema: &SchemaRef, ) -> ConstantColumns { @@ -2025,7 +1976,7 @@ fn constant_value_from_stats( /// Returns an error if an invalid parquet access extension is provided. /// /// Note: file_name is only used for error messages -fn create_initial_plan( +pub(crate) fn create_initial_plan( file_name: &str, extensions: &datafusion_datasource::FileExtensions, rg_metadata: &[RowGroupMetaData], @@ -2090,7 +2041,7 @@ pub(crate) fn build_pruning_predicates( /// Returns a `ArrowReaderMetadata` with the page index loaded, loading /// it from the underlying `AsyncFileReader` if necessary. -async fn load_page_index( +pub(crate) async fn load_page_index( reader_metadata: ArrowReaderMetadata, input: &mut T, options: ArrowReaderOptions, diff --git a/datafusion/datasource-parquet/src/schema_coercion.rs b/datafusion/datasource-parquet/src/schema_coercion.rs index 828070659b707..7548f5d2c3b00 100644 --- a/datafusion/datasource-parquet/src/schema_coercion.rs +++ b/datafusion/datasource-parquet/src/schema_coercion.rs @@ -42,6 +42,41 @@ use arrow::datatypes::{DataType, Field, FieldRef, Fields, Schema, TimeUnit}; use parquet::basic::Type; use parquet::schema::types::SchemaDescriptor; +/// Applies all type coercions DataFusion performs when reading a parquet file +/// to `physical_file_schema`, the arrow schema read from the file's metadata: +/// the coercions of [`apply_file_schema_type_coercions`] followed by the +/// optional INT96 timestamp coercion (see [`Int96Coercer`]). +/// +/// Returns `None` if no coercion is needed. +pub(crate) fn coerce_physical_file_schema( + logical_file_schema: &Schema, + physical_file_schema: &Arc, + parquet_schema: &SchemaDescriptor, + coerce_int96: Option<&TimeUnit>, + coerce_int96_tz: Option>, +) -> Option> { + let mut coerced = None; + if let Some(merged) = + apply_file_schema_type_coercions(logical_file_schema, physical_file_schema) + { + coerced = Some(Arc::new(merged)); + } + + if let Some(coerce) = coerce_int96 + && let Some(merged) = Int96Coercer::new( + parquet_schema, + coerced.as_deref().unwrap_or(physical_file_schema), + coerce, + ) + .with_timezone(coerce_int96_tz) + .coerce() + { + coerced = Some(Arc::new(merged)); + } + + coerced +} + /// Apply necessary schema type coercions to make file schema match table schema. /// /// This function performs two main types of transformations in a single pass: diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index a5c165abb1294..878388f09cd98 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -21,6 +21,7 @@ use std::fmt::Formatter; use std::sync::Arc; use crate::DefaultParquetFileReaderFactory; +use crate::EagerPruningSummary; use crate::ParquetFileReaderFactory; use crate::ParquetFileSchemaProvider; use crate::opener::ParquetMorselizer; @@ -324,6 +325,9 @@ pub struct ParquetSource { /// Sort order driving `PreparedAccessPlan::reorder_by_statistics` /// in the opener. sort_order_for_reorder: Option, + /// The outcome of eager pruning while the scan was planned, if eager + /// pruning was attempted. Only used for display. + eager_pruning_summary: Option, } impl ParquetSource { @@ -351,9 +355,24 @@ impl ParquetSource { encryption_factory: None, reverse_row_groups: false, sort_order_for_reorder: None, + eager_pruning_summary: None, } } + /// Record the outcome of eager pruning, shown in `EXPLAIN` output. + /// + /// See [`EagerParquetPruning`](datafusion_common::config::EagerParquetPruning). + pub fn with_eager_pruning_summary(mut self, summary: EagerPruningSummary) -> Self { + self.eager_pruning_summary = Some(summary); + self + } + + /// The outcome of eager pruning, if eager pruning was attempted while the + /// scan was planned. + pub fn eager_pruning_summary(&self) -> Option<&EagerPruningSummary> { + self.eager_pruning_summary.as_ref() + } + /// Set the `TableParquetOptions` for this ParquetSource. pub fn with_table_parquet_options( mut self, @@ -844,6 +863,10 @@ impl FileSource for ParquetSource { )?; } } + + if let Some(summary) = &self.eager_pruning_summary { + write!(f, ", eager_pruning={summary}")?; + } Ok(()) } DisplayFormatType::TreeRender => { @@ -1144,6 +1167,8 @@ impl FileSource for ParquetSource { encryption_factory: _, reverse_row_groups, sort_order_for_reorder, + // Only used to display the outcome of eager pruning in `EXPLAIN`. + eager_pruning_summary: _, } = self; if schema_provider.is_some() { diff --git a/datafusion/datasource/src/file_format.rs b/datafusion/datasource/src/file_format.rs index 2c2a94e8cc4b4..52688aa6d3f2a 100644 --- a/datafusion/datasource/src/file_format.rs +++ b/datafusion/datasource/src/file_format.rs @@ -35,7 +35,7 @@ use datafusion_common::{ }; use datafusion_physical_expr::LexRequirement; use datafusion_physical_expr_common::sort_expr::LexOrdering; -use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; use datafusion_session::Session; use async_trait::async_trait; @@ -187,10 +187,22 @@ pub trait FileFormat: Any + Send + Sync + fmt::Debug { /// Take a list of files and convert it to the appropriate executor /// according to this file format. + /// + /// `filters` are the filters of the query that apply to this scan, + /// expressed against the table schema (file columns plus partition + /// columns) and combined with an implicit `AND`. An empty slice means there + /// are no filters. + /// + /// The filters are hints: callers still evaluate them on the output of the + /// returned plan, so a format may use them to skip data that provably + /// cannot match (e.g. to refine statistics during planning), but must not + /// rely on them for correctness. They are distinct from filters pushed into + /// the [`FileSource`] by the physical optimizer. async fn create_physical_plan( &self, state: &dyn Session, conf: FileScanConfig, + filters: &[Arc], ) -> Result>; /// Take a list of files and the configuration to convert it to the diff --git a/datafusion/proto-common/proto/datafusion_common.proto b/datafusion/proto-common/proto/datafusion_common.proto index bdabb0777ef81..0f9e6ac0cc039 100644 --- a/datafusion/proto-common/proto/datafusion_common.proto +++ b/datafusion/proto-common/proto/datafusion_common.proto @@ -626,6 +626,14 @@ message ParquetOptions { uint64 max_in_list_size = 38; + string eager_pruning = 39; // default = "disabled" + + // Optional so that messages without it decode to the default (256) rather + // than 0, which would skip eager pruning for every scan + oneof eager_pruning_file_limit_opt { + uint64 eager_pruning_file_limit = 40; + } + string created_by = 16; oneof coerce_int96_opt { diff --git a/datafusion/proto-common/src/from_proto/mod.rs b/datafusion/proto-common/src/from_proto/mod.rs index 0bde696c7c46b..32164e7bb3811 100644 --- a/datafusion/proto-common/src/from_proto/mod.rs +++ b/datafusion/proto-common/src/from_proto/mod.rs @@ -1219,6 +1219,16 @@ impl TryFrom<&protobuf::ParquetOptions> for ParquetOptions { .transpose()?, max_row_group_size: to_usize(value.max_row_group_size, "max_row_group_size")?, max_in_list_size: to_usize(value.max_in_list_size, "max_in_list_size")?, + eager_pruning: value.eager_pruning.parse()?, + eager_pruning_file_limit: value + .eager_pruning_file_limit_opt + .map(|opt| match opt { + protobuf::parquet_options::EagerPruningFileLimitOpt::EagerPruningFileLimit(v) => { + to_usize(v, "eager_pruning_file_limit") + } + }) + .transpose()? + .unwrap_or_else(|| ParquetOptions::default().eager_pruning_file_limit), created_by: value.created_by.clone(), column_index_truncate_length: value .column_index_truncate_length_opt.as_ref() @@ -1436,7 +1446,8 @@ where #[cfg(test)] mod tests { use datafusion_common::config::{ - MaxRowGroupBytes, ParquetCdcOptions, ParquetOptions, TableParquetOptions, + EagerParquetPruning, MaxRowGroupBytes, ParquetCdcOptions, ParquetOptions, + TableParquetOptions, }; use datafusion_common::parquet_config::{DFParquetCompression, DFParquetStatistics}; @@ -1625,6 +1636,36 @@ mod tests { ); } + #[test] + fn test_parquet_options_eager_pruning_round_trip() { + let opts = ParquetOptions { + eager_pruning: EagerParquetPruning::PageIndex, + eager_pruning_file_limit: 7, + ..ParquetOptions::default() + }; + let recovered = parquet_options_proto_round_trip(opts); + assert_eq!(recovered.eager_pruning, EagerParquetPruning::PageIndex); + assert_eq!(recovered.eager_pruning_file_limit, 7); + } + + #[test] + fn test_parquet_options_eager_pruning_defaults_when_unset() { + // Messages without the eager pruning fields (e.g. written before they + // existed) decode to the config defaults, not to 0 / empty values + let mut proto: crate::protobuf_common::ParquetOptions = + (&ParquetOptions::default()).try_into().expect("to_proto"); + proto.eager_pruning = String::new(); + proto.eager_pruning_file_limit_opt = None; + + let recovered = ParquetOptions::try_from(&proto).expect("from_proto"); + let defaults = ParquetOptions::default(); + assert_eq!(recovered.eager_pruning, defaults.eager_pruning); + assert_eq!( + recovered.eager_pruning_file_limit, + defaults.eager_pruning_file_limit + ); + } + #[test] fn test_table_parquet_options_coerce_int96_tz_round_trip() { let mut opts = TableParquetOptions::default(); diff --git a/datafusion/proto-common/src/generated/pbjson.rs b/datafusion/proto-common/src/generated/pbjson.rs index a1d86c5fe3187..483ab10e300d6 100644 --- a/datafusion/proto-common/src/generated/pbjson.rs +++ b/datafusion/proto-common/src/generated/pbjson.rs @@ -6491,6 +6491,9 @@ impl serde::Serialize for ParquetOptions { if self.max_in_list_size != 0 { len += 1; } + if !self.eager_pruning.is_empty() { + len += 1; + } if !self.created_by.is_empty() { len += 1; } @@ -6524,6 +6527,9 @@ impl serde::Serialize for ParquetOptions { if self.bloom_filter_ndv_opt.is_some() { len += 1; } + if self.eager_pruning_file_limit_opt.is_some() { + len += 1; + } if self.coerce_int96_opt.is_some() { len += 1; } @@ -6616,6 +6622,9 @@ impl serde::Serialize for ParquetOptions { #[allow(clippy::needless_borrows_for_generic_args)] struct_ser.serialize_field("maxInListSize", ToString::to_string(&self.max_in_list_size).as_str())?; } + if !self.eager_pruning.is_empty() { + struct_ser.serialize_field("eagerPruning", &self.eager_pruning)?; + } if !self.created_by.is_empty() { struct_ser.serialize_field("createdBy", &self.created_by)?; } @@ -6693,6 +6702,15 @@ impl serde::Serialize for ParquetOptions { } } } + if let Some(v) = self.eager_pruning_file_limit_opt.as_ref() { + match v { + parquet_options::EagerPruningFileLimitOpt::EagerPruningFileLimit(v) => { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("eagerPruningFileLimit", ToString::to_string(&v).as_str())?; + } + } + } if let Some(v) = self.coerce_int96_opt.as_ref() { match v { parquet_options::CoerceInt96Opt::CoerceInt96(v) => { @@ -6776,6 +6794,8 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "maxRowGroupSize", "max_in_list_size", "maxInListSize", + "eager_pruning", + "eagerPruning", "created_by", "createdBy", "content_defined_chunking", @@ -6796,6 +6816,8 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "bloomFilterFpp", "bloom_filter_ndv", "bloomFilterNdv", + "eager_pruning_file_limit", + "eagerPruningFileLimit", "coerce_int96", "coerceInt96", "max_predicate_cache_size", @@ -6829,6 +6851,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { DataPageRowCountLimit, MaxRowGroupSize, MaxInListSize, + EagerPruning, CreatedBy, ContentDefinedChunking, MetadataSizeHint, @@ -6840,6 +6863,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { Encoding, BloomFilterFpp, BloomFilterNdv, + EagerPruningFileLimit, CoerceInt96, MaxPredicateCacheSize, MaxRowGroupBytes, @@ -6886,6 +6910,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "dataPageRowCountLimit" | "data_page_row_count_limit" => Ok(GeneratedField::DataPageRowCountLimit), "maxRowGroupSize" | "max_row_group_size" => Ok(GeneratedField::MaxRowGroupSize), "maxInListSize" | "max_in_list_size" => Ok(GeneratedField::MaxInListSize), + "eagerPruning" | "eager_pruning" => Ok(GeneratedField::EagerPruning), "createdBy" | "created_by" => Ok(GeneratedField::CreatedBy), "contentDefinedChunking" | "content_defined_chunking" => Ok(GeneratedField::ContentDefinedChunking), "metadataSizeHint" | "metadata_size_hint" => Ok(GeneratedField::MetadataSizeHint), @@ -6897,6 +6922,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "encoding" => Ok(GeneratedField::Encoding), "bloomFilterFpp" | "bloom_filter_fpp" => Ok(GeneratedField::BloomFilterFpp), "bloomFilterNdv" | "bloom_filter_ndv" => Ok(GeneratedField::BloomFilterNdv), + "eagerPruningFileLimit" | "eager_pruning_file_limit" => Ok(GeneratedField::EagerPruningFileLimit), "coerceInt96" | "coerce_int96" => Ok(GeneratedField::CoerceInt96), "maxPredicateCacheSize" | "max_predicate_cache_size" => Ok(GeneratedField::MaxPredicateCacheSize), "maxRowGroupBytes" | "max_row_group_bytes" => Ok(GeneratedField::MaxRowGroupBytes), @@ -6941,6 +6967,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { let mut data_page_row_count_limit__ = None; let mut max_row_group_size__ = None; let mut max_in_list_size__ = None; + let mut eager_pruning__ = None; let mut created_by__ = None; let mut content_defined_chunking__ = None; let mut metadata_size_hint_opt__ = None; @@ -6952,6 +6979,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { let mut encoding_opt__ = None; let mut bloom_filter_fpp_opt__ = None; let mut bloom_filter_ndv_opt__ = None; + let mut eager_pruning_file_limit_opt__ = None; let mut coerce_int96_opt__ = None; let mut max_predicate_cache_size_opt__ = None; let mut max_row_group_bytes_opt__ = None; @@ -7100,6 +7128,12 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) ; } + GeneratedField::EagerPruning => { + if eager_pruning__.is_some() { + return Err(serde::de::Error::duplicate_field("eagerPruning")); + } + eager_pruning__ = Some(map_.next_value()?); + } GeneratedField::CreatedBy => { if created_by__.is_some() { return Err(serde::de::Error::duplicate_field("createdBy")); @@ -7166,6 +7200,12 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { } bloom_filter_ndv_opt__ = map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| parquet_options::BloomFilterNdvOpt::BloomFilterNdv(x.0)); } + GeneratedField::EagerPruningFileLimit => { + if eager_pruning_file_limit_opt__.is_some() { + return Err(serde::de::Error::duplicate_field("eagerPruningFileLimit")); + } + eager_pruning_file_limit_opt__ = map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| parquet_options::EagerPruningFileLimitOpt::EagerPruningFileLimit(x.0)); + } GeneratedField::CoerceInt96 => { if coerce_int96_opt__.is_some() { return Err(serde::de::Error::duplicate_field("coerceInt96")); @@ -7214,6 +7254,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { data_page_row_count_limit: data_page_row_count_limit__.unwrap_or_default(), max_row_group_size: max_row_group_size__.unwrap_or_default(), max_in_list_size: max_in_list_size__.unwrap_or_default(), + eager_pruning: eager_pruning__.unwrap_or_default(), created_by: created_by__.unwrap_or_default(), content_defined_chunking: content_defined_chunking__, metadata_size_hint_opt: metadata_size_hint_opt__, @@ -7225,6 +7266,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { encoding_opt: encoding_opt__, bloom_filter_fpp_opt: bloom_filter_fpp_opt__, bloom_filter_ndv_opt: bloom_filter_ndv_opt__, + eager_pruning_file_limit_opt: eager_pruning_file_limit_opt__, coerce_int96_opt: coerce_int96_opt__, max_predicate_cache_size_opt: max_predicate_cache_size_opt__, max_row_group_bytes_opt: max_row_group_bytes_opt__, diff --git a/datafusion/proto-common/src/generated/prost.rs b/datafusion/proto-common/src/generated/prost.rs index f856c1aa3959c..877f9df490473 100644 --- a/datafusion/proto-common/src/generated/prost.rs +++ b/datafusion/proto-common/src/generated/prost.rs @@ -875,6 +875,9 @@ pub struct ParquetOptions { pub max_row_group_size: u64, #[prost(uint64, tag = "38")] pub max_in_list_size: u64, + /// default = "disabled" + #[prost(string, tag = "39")] + pub eager_pruning: ::prost::alloc::string::String, #[prost(string, tag = "16")] pub created_by: ::prost::alloc::string::String, #[prost(message, optional, tag = "35")] @@ -907,6 +910,12 @@ pub struct ParquetOptions { pub bloom_filter_fpp_opt: ::core::option::Option, #[prost(oneof = "parquet_options::BloomFilterNdvOpt", tags = "22")] pub bloom_filter_ndv_opt: ::core::option::Option, + /// Optional so that messages without it decode to the default (256) rather + /// than 0, which would skip eager pruning for every scan + #[prost(oneof = "parquet_options::EagerPruningFileLimitOpt", tags = "40")] + pub eager_pruning_file_limit_opt: ::core::option::Option< + parquet_options::EagerPruningFileLimitOpt, + >, #[prost(oneof = "parquet_options::CoerceInt96Opt", tags = "32")] pub coerce_int96_opt: ::core::option::Option, #[prost(oneof = "parquet_options::MaxPredicateCacheSizeOpt", tags = "33")] @@ -971,6 +980,13 @@ pub mod parquet_options { #[prost(uint64, tag = "22")] BloomFilterNdv(u64), } + /// Optional so that messages without it decode to the default (256) rather + /// than 0, which would skip eager pruning for every scan + #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Oneof)] + pub enum EagerPruningFileLimitOpt { + #[prost(uint64, tag = "40")] + EagerPruningFileLimit(u64), + } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Oneof)] pub enum CoerceInt96Opt { #[prost(string, tag = "32")] diff --git a/datafusion/proto-common/src/to_proto/mod.rs b/datafusion/proto-common/src/to_proto/mod.rs index b8eb46f2c8da8..2c72e441bb488 100644 --- a/datafusion/proto-common/src/to_proto/mod.rs +++ b/datafusion/proto-common/src/to_proto/mod.rs @@ -948,6 +948,8 @@ impl TryFrom<&ParquetOptions> for protobuf::ParquetOptions { statistics_enabled_opt: value.statistics_enabled.map(|v| protobuf::parquet_options::StatisticsEnabledOpt::StatisticsEnabled(v.to_string())), max_row_group_size: value.max_row_group_size as u64, max_in_list_size: value.max_in_list_size as u64, + eager_pruning: value.eager_pruning.to_string(), + eager_pruning_file_limit_opt: Some(protobuf::parquet_options::EagerPruningFileLimitOpt::EagerPruningFileLimit(value.eager_pruning_file_limit as u64)), created_by: value.created_by.clone(), column_index_truncate_length_opt: value.column_index_truncate_length.map(|v| protobuf::parquet_options::ColumnIndexTruncateLengthOpt::ColumnIndexTruncateLength(v as u64)), statistics_truncate_length_opt: value.statistics_truncate_length.map(|v| protobuf::parquet_options::StatisticsTruncateLengthOpt::StatisticsTruncateLength(v as u64)), diff --git a/datafusion/proto-models/src/from_proto.rs b/datafusion/proto-models/src/from_proto.rs index fdf32396475b7..25c980471194e 100644 --- a/datafusion/proto-models/src/from_proto.rs +++ b/datafusion/proto-models/src/from_proto.rs @@ -411,6 +411,17 @@ impl TryFrom<&ParquetOptionsProto> for ParquetOptions { .transpose()?, max_row_group_size: to_usize(proto.max_row_group_size, "max_row_group_size")?, max_in_list_size: to_usize(proto.max_in_list_size, "max_in_list_size")?, + eager_pruning: proto.eager_pruning.parse()?, + eager_pruning_file_limit: proto + .eager_pruning_file_limit_opt + .as_ref() + .map(|opt| match opt { + parquet_options::EagerPruningFileLimitOpt::EagerPruningFileLimit( + limit, + ) => to_usize(*limit, "eager_pruning_file_limit"), + }) + .transpose()? + .unwrap_or_else(|| ParquetOptions::default().eager_pruning_file_limit), created_by: proto.created_by.clone(), column_index_truncate_length: proto .column_index_truncate_length_opt @@ -571,6 +582,7 @@ impl TryFrom<&TableParquetOptionsProto> for TableParquetOptions { #[cfg(test)] mod tests { use super::*; + use datafusion_common::config::EagerParquetPruning; #[test] fn rejects_invalid_parquet_statistics() { @@ -589,4 +601,38 @@ mod tests { .contains("Invalid parquet statistics setting: invalid") ); } + + #[test] + fn eager_pruning_defaults_when_unset() { + // Messages without the eager pruning fields (e.g. written before they + // existed) decode to the config defaults, not to 0 / empty values + let proto = ParquetOptionsProto { + writer_version: "1.0".to_string(), + ..Default::default() + }; + + let options = ParquetOptions::try_from(&proto).unwrap(); + let defaults = ParquetOptions::default(); + assert_eq!(options.eager_pruning, defaults.eager_pruning); + assert_eq!( + options.eager_pruning_file_limit, + defaults.eager_pruning_file_limit + ); + } + + #[test] + fn eager_pruning_from_proto() { + let proto = ParquetOptionsProto { + writer_version: "1.0".to_string(), + eager_pruning: "page_index".to_string(), + eager_pruning_file_limit_opt: Some( + parquet_options::EagerPruningFileLimitOpt::EagerPruningFileLimit(7), + ), + ..Default::default() + }; + + let options = ParquetOptions::try_from(&proto).unwrap(); + assert_eq!(options.eager_pruning, EagerParquetPruning::PageIndex); + assert_eq!(options.eager_pruning_file_limit, 7); + } } diff --git a/datafusion/proto-models/src/generated/datafusion_proto_common.rs b/datafusion/proto-models/src/generated/datafusion_proto_common.rs index f856c1aa3959c..877f9df490473 100644 --- a/datafusion/proto-models/src/generated/datafusion_proto_common.rs +++ b/datafusion/proto-models/src/generated/datafusion_proto_common.rs @@ -875,6 +875,9 @@ pub struct ParquetOptions { pub max_row_group_size: u64, #[prost(uint64, tag = "38")] pub max_in_list_size: u64, + /// default = "disabled" + #[prost(string, tag = "39")] + pub eager_pruning: ::prost::alloc::string::String, #[prost(string, tag = "16")] pub created_by: ::prost::alloc::string::String, #[prost(message, optional, tag = "35")] @@ -907,6 +910,12 @@ pub struct ParquetOptions { pub bloom_filter_fpp_opt: ::core::option::Option, #[prost(oneof = "parquet_options::BloomFilterNdvOpt", tags = "22")] pub bloom_filter_ndv_opt: ::core::option::Option, + /// Optional so that messages without it decode to the default (256) rather + /// than 0, which would skip eager pruning for every scan + #[prost(oneof = "parquet_options::EagerPruningFileLimitOpt", tags = "40")] + pub eager_pruning_file_limit_opt: ::core::option::Option< + parquet_options::EagerPruningFileLimitOpt, + >, #[prost(oneof = "parquet_options::CoerceInt96Opt", tags = "32")] pub coerce_int96_opt: ::core::option::Option, #[prost(oneof = "parquet_options::MaxPredicateCacheSizeOpt", tags = "33")] @@ -971,6 +980,13 @@ pub mod parquet_options { #[prost(uint64, tag = "22")] BloomFilterNdv(u64), } + /// Optional so that messages without it decode to the default (256) rather + /// than 0, which would skip eager pruning for every scan + #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Oneof)] + pub enum EagerPruningFileLimitOpt { + #[prost(uint64, tag = "40")] + EagerPruningFileLimit(u64), + } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Oneof)] pub enum CoerceInt96Opt { #[prost(string, tag = "32")] diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index f6633a165df18..666ea6e3833d7 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -250,6 +250,8 @@ datafusion.execution.parquet.data_page_row_count_limit 20000 datafusion.execution.parquet.data_pagesize_limit 1048576 datafusion.execution.parquet.dictionary_enabled true datafusion.execution.parquet.dictionary_page_size_limit 1048576 +datafusion.execution.parquet.eager_pruning disabled +datafusion.execution.parquet.eager_pruning_file_limit 256 datafusion.execution.parquet.enable_page_index true datafusion.execution.parquet.encoding NULL datafusion.execution.parquet.force_filter_selections false @@ -412,6 +414,8 @@ datafusion.execution.parquet.data_page_row_count_limit 20000 (writing) Sets best datafusion.execution.parquet.data_pagesize_limit 1048576 (writing) Sets best effort maximum size of data page in bytes datafusion.execution.parquet.dictionary_enabled true (writing) Sets if dictionary encoding is enabled. If NULL, uses default parquet writer setting datafusion.execution.parquet.dictionary_page_size_limit 1048576 (writing) Sets best effort maximum dictionary page size, in bytes +datafusion.execution.parquet.eager_pruning disabled (reading) How much pruning the parquet scan performs while the query is planned, to tighten the scan's row count and byte size estimates. When the scan executes, it starts from the pruned files and row groups, but evaluates its own predicate again. Valid values are `disabled`, `row_groups` (row group statistics), `page_index` (additionally the page index) and `bloom_filters` (additionally bloom filters). Each level includes the work of the levels before it and may add I/O and latency to planning. Statistics derived from eager pruning are always inexact. +datafusion.execution.parquet.eager_pruning_file_limit 256 (reading) Eager pruning is skipped for scans with more files than this limit, bounding the planning time spent by `eager_pruning`. A limit of 0 skips eager pruning for every scan. datafusion.execution.parquet.enable_page_index true (reading) If true, reads the Parquet data page level metadata (the Page Index), if present, to reduce the I/O and number of rows decoded. datafusion.execution.parquet.encoding NULL (writing) Sets default encoding for any column. Valid values are: plain, plain_dictionary, rle, bit_packed, delta_binary_packed, delta_length_byte_array, delta_byte_array, rle_dictionary, and byte_stream_split. These values are not case sensitive. If NULL, uses default parquet writer setting datafusion.execution.parquet.force_filter_selections false (reading) Force the use of RowSelections for filter results, when pushdown_filters is enabled. If false, the reader will automatically choose between a RowSelection and a Bitmap based on the number and pattern of selected rows. diff --git a/datafusion/sqllogictest/test_files/parquet_eager_pruning.slt b/datafusion/sqllogictest/test_files/parquet_eager_pruning.slt new file mode 100644 index 0000000000000..ecec2bf5d2c9b --- /dev/null +++ b/datafusion/sqllogictest/test_files/parquet_eager_pruning.slt @@ -0,0 +1,239 @@ +# 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. + +# Tests for eager parquet pruning (datafusion.execution.parquet.eager_pruning) +# +# The test file has 10 row groups of 100 rows, with data pages of 10 rows and a +# bloom filter on column `s`. Column `a` is sorted and holds 0 to 999. Column +# `s` holds 's000' to 's999' scattered across all row groups. +# +# NOTE that parquet options are read when a table is created, so the table +# must be recreated to see the effects of a setting + +statement ok +set datafusion.execution.target_partitions = 1; + +query I +COPY ( + SELECT value AS a, concat('s', lpad(CAST((value * 37) % 1000 AS VARCHAR), 3, '0')) AS s + FROM generate_series(0, 999) + ORDER BY a +) +TO 'test_files/scratch/parquet_eager_pruning/t/0.parquet' +STORED AS PARQUET +OPTIONS ( + 'format.max_row_group_size' '100', + 'format.data_page_row_count_limit' '10', + 'format.write_batch_size' '10', + 'format.bloom_filter_enabled::s' true, + 'format.bloom_filter_fpp::s' 0.0001 +); +---- +1000 + +statement ok +set datafusion.explain.physical_plan_only = true; + +statement ok +set datafusion.explain.show_statistics = true; + +###### +# Eager pruning is disabled by default +###### + +statement ok +CREATE EXTERNAL TABLE t +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_eager_pruning/t/'; + +query TT +EXPLAIN SELECT a FROM t WHERE a < 150; +---- +physical_plan +01)FilterExec: a@0 < 150, statistics=[Rows=Inexact(150), Bytes=Inexact(1200), [(Col[0]: Min=Inexact(Int64(0)) Max=Inexact(Int64(149)) Null=Exact(0) ScanBytes=Inexact(1200))]] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_eager_pruning/t/0.parquet]]}, projection=[a], output_ordering=[a@0 ASC NULLS LAST], file_type=parquet, predicate=a@0 < 150, pruning_predicate=a_null_count@1 != row_count@2 AND a_min@0 < 150, required_guarantees=[], statistics=[Rows=Inexact(1000), Bytes=Inexact(8000), [(Col[0]: Min=Inexact(Int64(0)) Max=Inexact(Int64(999)) Null=Inexact(0) ScanBytes=Inexact(8000))]] + +statement ok +DROP TABLE t; + +###### +# Row group statistics +###### + +statement ok +set datafusion.execution.parquet.eager_pruning = 'row_groups'; + +statement ok +CREATE EXTERNAL TABLE t +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_eager_pruning/t/'; + +query TT +EXPLAIN SELECT a FROM t WHERE a < 150; +---- +physical_plan +01)FilterExec: a@0 < 150, statistics=[Rows=Inexact(150), Bytes=Inexact(1200), [(Col[0]: Min=Inexact(Int64(0)) Max=Inexact(Int64(149)) Null=Exact(0) ScanBytes=Inexact(1200))]] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_eager_pruning/t/0.parquet]]}, projection=[a], output_ordering=[a@0 ASC NULLS LAST], file_type=parquet, predicate=a@0 < 150, pruning_predicate=a_null_count@1 != row_count@2 AND a_min@0 < 150, required_guarantees=[], eager_pruning=[level=row_groups, files_pruned=0/1, row_groups_pruned=8/10, rows_pruned=800/1000], statistics=[Rows=Inexact(200), Bytes=Inexact(1600), [(Col[0]: Min=Inexact(Int64(0)) Max=Inexact(Int64(199)) Null=Inexact(0) ScanBytes=Inexact(1600))]] + +query I +SELECT count(*) FROM t WHERE a < 150; +---- +150 + +# Row group statistics cannot prune `a = 123` below a row group +query TT +EXPLAIN SELECT a FROM t WHERE a = 123; +---- +physical_plan +01)FilterExec: a@0 = 123, statistics=[Rows=Inexact(1), Bytes=Inexact(8), [(Col[0]: Min=Exact(Int64(123)) Max=Exact(Int64(123)) Null=Exact(0) Distinct=Inexact(1) ScanBytes=Inexact(8))]] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_eager_pruning/t/0.parquet]]}, projection=[a], output_ordering=[a@0 ASC NULLS LAST], file_type=parquet, predicate=a@0 = 123, pruning_predicate=a_null_count@2 != row_count@3 AND a_min@0 <= 123 AND 123 <= a_max@1, required_guarantees=[a in (123)], eager_pruning=[level=row_groups, files_pruned=0/1, row_groups_pruned=9/10, rows_pruned=900/1000], statistics=[Rows=Inexact(100), Bytes=Inexact(800), [(Col[0]: Min=Inexact(Int64(100)) Max=Inexact(Int64(199)) Null=Inexact(0) ScanBytes=Inexact(800))]] + +# No rows can be pruned +query TT +EXPLAIN SELECT a FROM t WHERE a >= 0; +---- +physical_plan +01)FilterExec: a@0 >= 0, statistics=[Rows=Inexact(1000), Bytes=Inexact(8000), [(Col[0]: Min=Inexact(Int64(0)) Max=Inexact(Int64(999)) Null=Exact(0) ScanBytes=Inexact(8000))]] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_eager_pruning/t/0.parquet]]}, projection=[a], output_ordering=[a@0 ASC NULLS LAST], file_type=parquet, predicate=a@0 >= 0, pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 >= 0, required_guarantees=[], eager_pruning=[level=row_groups, files_pruned=0/1, row_groups_pruned=0/10, rows_pruned=0/1000], statistics=[Rows=Inexact(1000), Bytes=Inexact(8000), [(Col[0]: Min=Inexact(Int64(0)) Max=Inexact(Int64(999)) Null=Inexact(0) ScanBytes=Inexact(8000))]] + +statement ok +DROP TABLE t; + +###### +# Page index +###### + +statement ok +set datafusion.execution.parquet.eager_pruning = 'page_index'; + +statement ok +CREATE EXTERNAL TABLE t +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_eager_pruning/t/'; + +query TT +EXPLAIN SELECT a FROM t WHERE a = 123; +---- +physical_plan +01)FilterExec: a@0 = 123, statistics=[Rows=Inexact(1), Bytes=Inexact(1), [(Col[0]: Min=Exact(Int64(123)) Max=Exact(Int64(123)) Null=Exact(0) Distinct=Inexact(1) ScanBytes=Inexact(1))]] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_eager_pruning/t/0.parquet]]}, projection=[a], output_ordering=[a@0 ASC NULLS LAST], file_type=parquet, predicate=a@0 = 123, pruning_predicate=a_null_count@2 != row_count@3 AND a_min@0 <= 123 AND 123 <= a_max@1, required_guarantees=[a in (123)], eager_pruning=[level=page_index, files_pruned=0/1, row_groups_pruned=9/10, rows_pruned=990/1000], statistics=[Rows=Inexact(10), Bytes=Inexact(80), [(Col[0]: Min=Inexact(Int64(100)) Max=Inexact(Int64(199)) Null=Inexact(0) ScanBytes=Inexact(80))]] + +query IT +SELECT a, s FROM t WHERE a = 123; +---- +123 s551 + +statement ok +DROP TABLE t; + +###### +# Bloom filters +###### + +statement ok +set datafusion.execution.parquet.eager_pruning = 'bloom_filters'; + +statement ok +CREATE EXTERNAL TABLE t +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_eager_pruning/t/'; + +# 's5005' is within the min/max range of every row group, but absent +query TT +EXPLAIN SELECT a FROM t WHERE s = 's5005'; +---- +physical_plan +01)FilterExec: s@1 = s5005, projection=[a@0], statistics=[Rows=Inexact(0), Bytes=Inexact(0), [(Col[0]:)]] +02)--DataSourceExec: file_groups={1 group: [[]]}, projection=[a, s], output_ordering=[a@0 ASC NULLS LAST], file_type=parquet, predicate=s@1 = s5005, pruning_predicate=s_null_count@2 != row_count@3 AND s_min@0 <= s5005 AND s5005 <= s_max@1, required_guarantees=[s in (s5005)], eager_pruning=[level=bloom_filters, files_pruned=1/1, row_groups_pruned=10/10, rows_pruned=1000/1000], statistics=[Rows=Inexact(0), Bytes=Inexact(0), [(Col[0]:),(Col[1]:)]] + +query I +SELECT count(*) FROM t WHERE s = 's5005'; +---- +0 + +query I +SELECT count(*) FROM t WHERE s = 's500'; +---- +1 + +statement ok +DROP TABLE t; + +###### +# Scans with more files than eager_pruning_file_limit are not pruned +###### + +statement ok +set datafusion.execution.parquet.eager_pruning_file_limit = 0; + +statement ok +CREATE EXTERNAL TABLE t +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_eager_pruning/t/'; + +query TT +EXPLAIN SELECT a FROM t WHERE a < 150; +---- +physical_plan +01)FilterExec: a@0 < 150, statistics=[Rows=Inexact(150), Bytes=Inexact(1200), [(Col[0]: Min=Inexact(Int64(0)) Max=Inexact(Int64(149)) Null=Exact(0) ScanBytes=Inexact(1200))]] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_eager_pruning/t/0.parquet]]}, projection=[a], output_ordering=[a@0 ASC NULLS LAST], file_type=parquet, predicate=a@0 < 150, pruning_predicate=a_null_count@1 != row_count@2 AND a_min@0 < 150, required_guarantees=[], eager_pruning=[level=bloom_filters, skipped=1 files exceed eager_pruning_file_limit 0], statistics=[Rows=Inexact(1000), Bytes=Inexact(8000), [(Col[0]: Min=Inexact(Int64(0)) Max=Inexact(Int64(999)) Null=Inexact(0) ScanBytes=Inexact(8000))]] + +statement ok +DROP TABLE t; + +statement ok +reset datafusion.execution.parquet.eager_pruning_file_limit; + +statement ok +reset datafusion.execution.parquet.eager_pruning; + +###### +# Eager pruning can be enabled for a single table +###### + +statement ok +CREATE EXTERNAL TABLE t +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_eager_pruning/t/' +OPTIONS ('format.eager_pruning' 'row_groups'); + +query TT +EXPLAIN SELECT a FROM t WHERE a < 150; +---- +physical_plan +01)FilterExec: a@0 < 150, statistics=[Rows=Inexact(150), Bytes=Inexact(1200), [(Col[0]: Min=Inexact(Int64(0)) Max=Inexact(Int64(149)) Null=Exact(0) ScanBytes=Inexact(1200))]] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_eager_pruning/t/0.parquet]]}, projection=[a], output_ordering=[a@0 ASC NULLS LAST], file_type=parquet, predicate=a@0 < 150, pruning_predicate=a_null_count@1 != row_count@2 AND a_min@0 < 150, required_guarantees=[], eager_pruning=[level=row_groups, files_pruned=0/1, row_groups_pruned=8/10, rows_pruned=800/1000], statistics=[Rows=Inexact(200), Bytes=Inexact(1600), [(Col[0]: Min=Inexact(Int64(0)) Max=Inexact(Int64(199)) Null=Inexact(0) ScanBytes=Inexact(1600))]] + +statement ok +DROP TABLE t; + +###### +# Invalid values +###### + +statement error Invalid eager parquet pruning level: full\. Expected one of: disabled, row_groups, page_index, bloom_filters +set datafusion.execution.parquet.eager_pruning = 'full'; + +statement ok +set datafusion.explain.show_statistics = false; + +statement ok +set datafusion.explain.physical_plan_only = false; + +# The SLT runner sets `target_partitions` to 4 +statement ok +set datafusion.execution.target_partitions = 4; diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index be71743f1afa9..dfa55446f5ba6 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -122,6 +122,84 @@ Rewriting affected files with a current DataFusion version, for example using restores optimizations that depend on them. Preserve the required data ordering and ordering metadata when rewriting sorted files. +### `FileFormat::create_physical_plan` receives the scan's filters + +`FileFormat::create_physical_plan` has a new `filters: &[Arc]` +argument. `ListingTable` passes the (non partition) filters of the query that +apply to the scan, expressed against the table schema and combined with an +implicit `AND`. The filters are hints: they are still evaluated above the scan, +so a format may use them to skip data that cannot match, for example to refine +statistics during planning, but must not rely on them for correctness. +`ParquetFormat` uses them for eager pruning, see +`datafusion.execution.parquet.eager_pruning`. + +**Who is affected:** + +- Implementations of the `FileFormat` trait. +- Code that calls `FileFormat::create_physical_plan` directly. + +**Migration guide:** + +Add the argument to implementations. Formats that do not use the filters can +ignore it: + +```rust,ignore +// Before +async fn create_physical_plan( + &self, + state: &dyn Session, + conf: FileScanConfig, +) -> Result> { + // ... +} + +// After +async fn create_physical_plan( + &self, + state: &dyn Session, + conf: FileScanConfig, + _filters: &[Arc], +) -> Result> { + // ... +} +``` + +Callers without filters pass an empty slice: + +```rust,ignore +format.create_physical_plan(state, conf, &[]).await?; +``` + +Formats that wrap another format should forward the filters. + +### `ParquetOptions` has new `eager_pruning` and `eager_pruning_file_limit` fields + +`ParquetOptions` gained the public fields `eager_pruning: EagerParquetPruning` +(default `EagerParquetPruning::Disabled`) and `eager_pruning_file_limit: usize` +(default `256`). With eager pruning enabled, parquet scans prune row groups (and +optionally data pages and bloom filters) while the query is planned, to refine +the statistics of the scan. The generated protobuf `ParquetOptions` message +gained the corresponding `eager_pruning` field and the optional +`eager_pruning_file_limit_opt` field. + +**Who is affected:** + +- Users constructing `ParquetOptions` or the generated protobuf + `ParquetOptions` with an exhaustive struct literal. Reading or `SET`-ing + configuration is unaffected. + +**Migration guide:** + +Set the new fields, or fill them from `Default`: + +```rust,ignore +ParquetOptions { + pruning: true, + // ... other fields ... + ..Default::default() +} +``` + ### `Timestamp(u, Some(tz)) - Timestamp(u, None)` reads the naive operand in `tz` Subtracting a timezone-naive timestamp from a timezone-aware one _at the same diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index ba8ea52bb0850..24aef2ffcf183 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -94,6 +94,8 @@ The following configuration settings are available: | datafusion.execution.parquet.bloom_filter_on_read | true | (reading) Use any available bloom filters when reading parquet files | | datafusion.execution.parquet.max_predicate_cache_size | NULL | (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. | | datafusion.execution.parquet.max_in_list_size | 20 | Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger literal lists use a compact representation when the column type is string, variable-length binary, integer, decimal, date, time, timestamp, or duration. This applies to both `IN` and `NOT IN`, including lists with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot match any rows. Compact lists containing NULL do not use the fully-matched-row-group optimization. Floating-point and other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. | +| datafusion.execution.parquet.eager_pruning | disabled | (reading) How much pruning the parquet scan performs while the query is planned, to tighten the scan's row count and byte size estimates. When the scan executes, it starts from the pruned files and row groups, but evaluates its own predicate again. Valid values are `disabled`, `row_groups` (row group statistics), `page_index` (additionally the page index) and `bloom_filters` (additionally bloom filters). Each level includes the work of the levels before it and may add I/O and latency to planning. Statistics derived from eager pruning are always inexact. | +| datafusion.execution.parquet.eager_pruning_file_limit | 256 | (reading) Eager pruning is skipped for scans with more files than this limit, bounding the planning time spent by `eager_pruning`. A limit of 0 skips eager pruning for every scan. | | datafusion.execution.parquet.data_pagesize_limit | 1048576 | (writing) Sets best effort maximum size of data page in bytes | | datafusion.execution.parquet.write_batch_size | 1024 | (writing) Sets write_batch_size in rows | | datafusion.execution.parquet.writer_version | 1.0 | (writing) Sets parquet writer version valid values are "1.0" and "2.0" |