From 0b0506a9acab9d5892ecf7e89243c3b34664bcc6 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 8 Sep 2026 21:13:44 +0800 Subject: [PATCH] feat: prune Parquet row groups for struct field predicates --- datafusion/common/src/pruning.rs | 79 ++++++ .../core/tests/parquet/filter_pushdown.rs | 194 ++++++++++++++ .../src/row_group_filter.rs | 237 ++++++++++++++++- datafusion/pruning/src/pruning_predicate.rs | 251 ++++++++++++++++-- .../test_files/projection_pushdown.slt | 14 +- 5 files changed, 746 insertions(+), 29 deletions(-) diff --git a/datafusion/common/src/pruning.rs b/datafusion/common/src/pruning.rs index 9f6a95c0978d4..76a82eaa582ba 100644 --- a/datafusion/common/src/pruning.rs +++ b/datafusion/common/src/pruning.rs @@ -88,6 +88,55 @@ pub trait PruningStatistics { /// Note: the returned array must contain [`Self::num_containers`] rows fn max_values(&self, column: &Column) -> Option; + /// Return minimum values for a field beneath `column`. + /// + /// `field_path` contains literal struct field names, without splitting dots. + /// An empty path selects the column itself. A non-empty path includes nulls + /// inherited from every struct ancestor. The bounds and array-length + /// requirements of [`Self::min_values`] apply. Providers that do not support + /// nested statistics return `None`. + fn min_values_for_path( + &self, + column: &Column, + field_path: &[String], + ) -> Option { + if field_path.is_empty() { + self.min_values(column) + } else { + None + } + } + + /// Return maximum values for a struct field, following the path and bound + /// conventions of [`Self::min_values_for_path`] and [`Self::max_values`]. + fn max_values_for_path( + &self, + column: &Column, + field_path: &[String], + ) -> Option { + if field_path.is_empty() { + self.max_values(column) + } else { + None + } + } + + /// Return null counts for a struct field, including nulls inherited from + /// its ancestors. Paths follow [`Self::min_values_for_path`]; unknown counts + /// follow [`Self::null_counts`]. Repeated list or map elements are not rows + /// and must not be reported as struct-field statistics. + fn null_counts_for_path( + &self, + column: &Column, + field_path: &[String], + ) -> Option { + if field_path.is_empty() { + self.null_counts(column) + } else { + None + } + } + /// Return the number of containers (e.g. Row Groups) being pruned with /// these statistics. /// @@ -480,6 +529,36 @@ impl CompositePruningStatistics { #[expect(deprecated)] impl PruningStatistics for CompositePruningStatistics { + fn min_values_for_path( + &self, + column: &Column, + field_path: &[String], + ) -> Option { + self.statistics + .iter() + .find_map(|stats| stats.min_values_for_path(column, field_path)) + } + + fn max_values_for_path( + &self, + column: &Column, + field_path: &[String], + ) -> Option { + self.statistics + .iter() + .find_map(|stats| stats.max_values_for_path(column, field_path)) + } + + fn null_counts_for_path( + &self, + column: &Column, + field_path: &[String], + ) -> Option { + self.statistics + .iter() + .find_map(|stats| stats.null_counts_for_path(column, field_path)) + } + fn min_values(&self, column: &Column) -> Option { for stats in &self.statistics { if let Some(array) = stats.min_values(column) { diff --git a/datafusion/core/tests/parquet/filter_pushdown.rs b/datafusion/core/tests/parquet/filter_pushdown.rs index da9f008d7ce21..df4b6349ec130 100644 --- a/datafusion/core/tests/parquet/filter_pushdown.rs +++ b/datafusion/core/tests/parquet/filter_pushdown.rs @@ -1041,3 +1041,197 @@ async fn custom_struct_accessor_pushdown_and_schema_adaptation() { } } } + +#[tokio::test] +async fn struct_field_row_group_statistics() { + use arrow::array::{Array, StructArray}; + use arrow::buffer::NullBuffer; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_expr::{ScalarUDF, Signature, Volatility}; + use parquet::file::properties::EnabledStatistics; + + for nullable_parent in [false, true] { + let values = Arc::new(Int32Array::from(vec![ + Some(1), + Some(2), + Some(3), + Some(10), + None, + Some(12), + nullable_parent.then_some(100), + nullable_parent.then_some(100), + nullable_parent.then_some(100), + Some(20), + Some(21), + Some(22), + ])); + let deep = StructArray::from(vec![( + Arc::new(Field::new("value", DataType::Int32, true)), + values.clone() as ArrayRef, + )]); + let a = StructArray::from(vec![( + Arc::new(Field::new("b", DataType::Int32, false)), + Arc::new(Int32Array::from(vec![-100; 12])) as ArrayRef, + )]); + let s = StructArray::new( + vec![ + Arc::new(Field::new("value", DataType::Int32, true)), + Arc::new(Field::new("deep", deep.data_type().clone(), false)), + Arc::new(Field::new("a.b", DataType::Int32, true)), + Arc::new(Field::new("a", a.data_type().clone(), false)), + ] + .into(), + vec![values.clone(), Arc::new(deep), values, Arc::new(a)], + nullable_parent.then(|| { + NullBuffer::from(vec![ + true, true, true, true, true, true, false, false, false, true, true, + true, + ]) + }), + ); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("s", s.data_type().clone(), nullable_parent), + Field::new("id", DataType::Int32, false), + Field::new("s.a.b", DataType::Int32, false), + // Must not collide with the first synthetic pruning column. + Field::new("__datafusion_struct_field_4", DataType::Int32, false), + ])), + vec![ + Arc::new(s), + Arc::new(Int32Array::from_iter_values(0..12)), + Arc::new(Int32Array::from(vec![-100; 12])), + Arc::new(Int32Array::from(vec![123; 12])), + ], + ) + .unwrap(); + for statistics in [EnabledStatistics::None, EnabledStatistics::Chunk] { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("struct-statistics.parquet"); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .set_statistics_enabled(statistics) + .build(); + let mut writer = ArrowWriter::try_new( + File::create(&path).unwrap(), + batch.schema(), + Some(props), + ) + .unwrap(); + writer.write(&batch).unwrap(); + assert_eq!(writer.close().unwrap().row_groups().len(), 4); + for pruning in [false, true] { + let mut config = SessionConfig::new() + .with_target_partitions(1) + .with_parquet_pruning(pruning) + .with_parquet_bloom_filter_pruning(false) + .with_parquet_page_index_pruning(false); + config.options_mut().execution.parquet.pushdown_filters = false; + let ctx = SessionContext::new_with_config(config); + for (name, declare_access) in + [("alias_field_at", true), ("opaque_field_at", false)] + { + ctx.register_udf( + ScalarUDF::from(FieldAt { + declare_access, + signature: Signature::any(2, Volatility::Immutable), + }) + .with_aliases([name]), + ); + } + ctx.register_parquet( + "t", + path.to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await + .unwrap(); + for (predicate, expected, pruned) in [ + ("s['value'] > 5", vec![3, 5, 9, 10, 11], 2), + ("s['deep']['value'] > 5", vec![3, 5, 9, 10, 11], 2), + ("s['a.b'] > 5", vec![3, 5, 9, 10, 11], 2), + ("alias_field_at('value', s) > 5", vec![3, 5, 9, 10, 11], 2), + ( + "alias_field_at('value', alias_field_at('deep', s)) > 5", + vec![3, 5, 9, 10, 11], + 2, + ), + ("opaque_field_at('value', s) > 5", vec![3, 5, 9, 10, 11], 0), + ("s['value'] = 2", vec![1], 3), + ("s['value'] IS NULL", vec![4, 6, 7, 8], 2), + ("s['value'] IS NOT NULL", vec![0, 1, 2, 3, 5, 9, 10, 11], 1), + ("s['value'] IN (2, 12)", vec![1, 5], 2), + ("s['value'] NOT IN (1, 2, 3)", vec![3, 5, 9, 10, 11], 1), + ("s['value'] BETWEEN 10 AND 12", vec![3, 5], 3), + ( + "s['value'] <= 3 OR s['value'] > 20", + vec![0, 1, 2, 10, 11], + 2, + ), + ("CAST(s['value'] AS BIGINT) > 5", vec![3, 5, 9, 10, 11], 2), + ( + "s['value'] > 5 AND __datafusion_struct_field_4 = 123", + vec![3, 5, 9, 10, 11], + 2, + ), + ("s['a']['b'] > 5", vec![], 4), + ("\"s.a.b\" > 5", vec![], 4), + ] { + let plan = ctx + .sql(&format!("SELECT id FROM t WHERE {predicate} ORDER BY id")) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let batches = collect(plan.clone(), ctx.task_ctx()).await.unwrap(); + let actual = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + assert_eq!( + actual, expected, + "{predicate}, nullable_parent={nullable_parent}" + ); + let metrics = TestParquetFile::parquet_metrics(&plan).unwrap(); + let expected_pruned = + if pruning && statistics == EnabledStatistics::Chunk { + pruned + } else { + 0 + }; + assert_eq!( + get_value(&metrics, "row_groups_pruned_statistics"), + expected_pruned, + "{predicate}, pruning={pruning}, statistics={statistics:?}, nullable_parent={nullable_parent}\n{}", + displayable(plan.as_ref()).indent(false) + ); + assert_eq!(get_value(&metrics, "predicate_evaluation_errors"), 0); + assert_eq!(get_value(&metrics, "pushdown_rows_pruned"), 0); + if predicate == "s['value'] BETWEEN 10 AND 12" { + let Some(MetricValue::PruningMetrics { + pruning_metrics, .. + }) = metrics.sum_by_name("row_groups_pruned_statistics") + else { + panic!("missing statistics metric"); + }; + assert_eq!( + pruning_metrics.fully_matched(), + 0, + "the matching row group contains a null leaf" + ); + } + } + } + } + } +} diff --git a/datafusion/datasource-parquet/src/row_group_filter.rs b/datafusion/datasource-parquet/src/row_group_filter.rs index 5ca99f2498e75..3877441021ec9 100644 --- a/datafusion/datasource-parquet/src/row_group_filter.rs +++ b/datafusion/datasource-parquet/src/row_group_filter.rs @@ -23,14 +23,17 @@ use crate::bloom_filter::BloomFilterStatistics; use crate::metadata::{has_untrusted_byte_array_stats, has_untrusted_min_max_order}; use arrow::array::{ArrayRef, BooleanArray, UInt64Array}; use arrow::compute::nullif; -use arrow::datatypes::Schema; +use arrow::datatypes::{DataType, Schema}; use datafusion_common::pruning::PruningStatistics; +use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_common::{Column, Result, ScalarValue}; use datafusion_datasource::FileRange; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, IsNullExpr, NotExpr}; use datafusion_physical_expr::utils::collect_columns; -use datafusion_physical_expr::{PhysicalExpr, PhysicalExprSimplifier}; +use datafusion_physical_expr::{ + PhysicalExpr, PhysicalExprSimplifier, ScalarFunctionExpr, +}; use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder}; use parquet::arrow::arrow_reader::statistics::StatisticsConverter; use parquet::basic::ColumnOrder; @@ -428,6 +431,28 @@ impl RowGroupAccessPlanFilter { )); } + // A non-null struct may still contain a null leaf. Guard field accesses + // as well as root columns before proving that every row matches. + if predicate + .orig_expr() + .apply(|expr| { + if let Some(function) = expr.downcast_ref::() + && function.struct_field_access().is_some() + && !expr.data_type(arrow_schema)?.is_nested() + { + inverted_expr = Arc::new(BinaryExpr::new( + Arc::clone(&inverted_expr), + Operator::Or, + Arc::new(IsNullExpr::new(Arc::clone(expr))), + )); + } + Ok(TreeNodeRecursion::Continue) + }) + .is_err() + { + return; + } + // Simplify the inverted expression (e.g., NOT(c1 = 0) -> c1 != 0) // before building the pruning predicate let simplifier = PhysicalExprSimplifier::new(arrow_schema); @@ -557,11 +582,63 @@ impl<'a> RowGroupPruningStatistics<'a> { .with_missing_null_counts_as_zero(self.missing_null_counts_as_zero)) } + fn statistics_converter_for_path( + &self, + column: &Column, + field_path: &[String], + ) -> Option> { + if field_path.is_empty() { + return self.statistics_converter(column).ok(); + } + let mut roots = self + .arrow_schema + .fields() + .iter() + .filter(|f| f.name() == &column.name); + let mut field = roots.next()?.as_ref(); + if roots.next().is_some() { + return None; + } + for name in field_path { + let DataType::Struct(fields) = field.data_type() else { + return None; + }; + let mut matches = fields.iter().filter(|f| f.name() == name); + field = matches.next()?.as_ref(); + if matches.next().is_some() { + return None; + } + } + if field.data_type().is_nested() { + return None; + } + let mut matches = + self.parquet_schema + .columns() + .iter() + .enumerate() + .filter(|(_, leaf)| { + let parts = leaf.path().parts(); + leaf.max_rep_level() == 0 + && parts.first() == Some(&column.name) + && parts[1..] == *field_path + }); + let (index, _) = matches.next()?; + if matches.next().is_some() { + return None; + } + StatisticsConverter::from_column_index(index, field, self.parquet_schema) + .ok() + // Missing counts must not prove that a nested field has no nulls. + .map(|converter| converter.with_missing_null_counts_as_zero(false)) + } + fn min_max_statistics_converter( &self, column: &Column, + field_path: &[String], ) -> Option> { - let converter = self.statistics_converter(column).ok()?; + let converter = self.statistics_converter_for_path(column, field_path)?; let parquet_index = converter.parquet_column_index(); if parquet_index.is_some_and(|index| { has_untrusted_min_max_order(self.parquet_schema, self.column_orders, index) @@ -594,13 +671,29 @@ impl<'a> RowGroupPruningStatistics<'a> { impl PruningStatistics for RowGroupPruningStatistics<'_> { fn min_values(&self, column: &Column) -> Option { - let converter = self.min_max_statistics_converter(column)?; + self.min_values_for_path(column, &[]) + } + + fn min_values_for_path( + &self, + column: &Column, + field_path: &[String], + ) -> Option { + let converter = self.min_max_statistics_converter(column, field_path)?; let values = converter.row_group_mins(self.metadata_iter()).ok()?; self.mask_untrusted_byte_array_stats(converter.parquet_column_index(), values) } fn max_values(&self, column: &Column) -> Option { - let converter = self.min_max_statistics_converter(column)?; + self.max_values_for_path(column, &[]) + } + + fn max_values_for_path( + &self, + column: &Column, + field_path: &[String], + ) -> Option { + let converter = self.min_max_statistics_converter(column, field_path)?; let values = converter.row_group_maxes(self.metadata_iter()).ok()?; self.mask_untrusted_byte_array_stats(converter.parquet_column_index(), values) } @@ -610,8 +703,16 @@ impl PruningStatistics for RowGroupPruningStatistics<'_> { } fn null_counts(&self, column: &Column) -> Option { - self.statistics_converter(column) - .and_then(|c| Ok(c.row_group_null_counts(self.metadata_iter())?)) + self.null_counts_for_path(column, &[]) + } + + fn null_counts_for_path( + &self, + column: &Column, + field_path: &[String], + ) -> Option { + self.statistics_converter_for_path(column, field_path)? + .row_group_null_counts(self.metadata_iter()) .ok() .map(|counts| Arc::new(counts) as ArrayRef) } @@ -1522,6 +1623,128 @@ mod tests { assert_eq!(metrics.row_groups_pruned_bloom_filter.matched(), 1); } + #[test] + fn nested_statistics_resolve_literal_paths_and_unknown_null_counts() { + use arrow::array::{Array, Int32Array}; + + let schema = Schema::new(vec![ + Field::new( + "s", + DataType::Struct( + vec![ + Field::new("text", DataType::Utf8, true), + Field::new("a.b", DataType::Int32, true), + Field::new( + "a", + DataType::Struct( + vec![Field::new("b", DataType::Int32, true)].into(), + ), + true, + ), + Field::new( + "items", + DataType::List(Arc::new(Field::new( + "item", + DataType::Int32, + true, + ))), + true, + ), + ] + .into(), + ), + true, + ), + Field::new("s.a.b", DataType::Int32, true), + ]); + let parquet_schema = + Arc::new(ArrowSchemaConverter::new().convert(&schema).unwrap()); + let group = get_row_group_meta_data( + &parquet_schema, + vec![ + ParquetStatistics::byte_array( + Some(ByteArray::from("a")), + Some(ByteArray::from("z")), + None, + Some(0), + false, + ), + ParquetStatistics::int32(Some(11), Some(19), None, None, false), + ParquetStatistics::int32(Some(22), Some(29), None, Some(3), false), + ParquetStatistics::int32(Some(33), Some(39), None, Some(0), false), + ParquetStatistics::int32(Some(44), Some(49), None, Some(0), false), + ], + ); + let stats = RowGroupPruningStatistics { + parquet_schema: &parquet_schema, + column_orders: None, + row_group_metadatas: vec![&group], + arrow_schema: &schema, + missing_null_counts_as_zero: true, + }; + for (root, path, expected) in [ + ("s", vec!["a.b".into()], 11), + ("s", vec!["a".into(), "b".into()], 22), + ("s.a.b", vec![], 44), + ] { + let values = stats + .min_values_for_path(&Column::from_name(root), &path) + .unwrap(); + assert_eq!( + values + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + expected + ); + } + let root = Column::from_name("s"); + assert!( + stats.min_values_for_path(&root, &["text".into()]).is_none(), + "nested byte-array bounds need a trusted column order too" + ); + let orders = + vec![ + ColumnOrder::TYPE_DEFINED_ORDER(parquet::basic::SortOrder::UNSIGNED); + parquet_schema.num_columns() + ]; + let trusted_stats = RowGroupPruningStatistics { + column_orders: Some(&orders), + row_group_metadatas: vec![&group], + parquet_schema: &parquet_schema, + arrow_schema: &schema, + missing_null_counts_as_zero: true, + }; + assert!( + trusted_stats + .min_values_for_path(&root, &["text".into()]) + .is_some() + ); + let nulls = stats.null_counts_for_path(&root, &["a.b".into()]).unwrap(); + assert!(nulls.is_null(0), "absent nested counts must remain unknown"); + let nulls = stats + .null_counts_for_path(&root, &["a".into(), "b".into()]) + .unwrap(); + assert_eq!( + nulls + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 3 + ); + for path in [ + vec!["missing".into()], + vec!["items".into()], + vec!["items".into(), "item".into()], + ] { + assert!(stats.min_values_for_path(&root, &path).is_none()); + assert!(stats.max_values_for_path(&root, &path).is_none()); + assert!(stats.null_counts_for_path(&root, &path).is_none()); + } + } + fn get_row_group_meta_data( schema_descr: &SchemaDescPtr, column_statistics: Vec, diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index c3362c63299e1..4736a55017150 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -19,7 +19,7 @@ //! based on statistics (e.g. Parquet Row Groups) //! //! [`Expr`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/enum.Expr.html -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use crate::string_in_list::{SetMembership, StringInListPruningExpr}; @@ -27,7 +27,7 @@ use crate::string_in_list::{SetMembership, StringInListPruningExpr}; use arrow::array::AsArray; use arrow::{ array::{ArrayRef, BooleanArray, new_null_array}, - datatypes::{DataType, Field, Schema, SchemaRef}, + datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}, record_batch::{RecordBatch, RecordBatchOptions}, }; // pub use for backwards compatibility @@ -48,7 +48,9 @@ use datafusion_common::{ use datafusion_expr_common::casts::try_cast_literal_to_type; use datafusion_expr_common::operator::Operator; use datafusion_physical_expr::utils::{Guarantee, LiteralGuarantee}; -use datafusion_physical_expr::{PhysicalExprRef, expressions as phys_expr}; +use datafusion_physical_expr::{ + PhysicalExprRef, ScalarFunctionExpr, expressions as phys_expr, +}; use datafusion_physical_expr_common::physical_expr::snapshot_physical_expr_opt; use datafusion_physical_plan::{ColumnarValue, PhysicalExpr}; @@ -539,10 +541,12 @@ impl<'a> PruningPredicateBuilder<'a> { // build predicate expression once let mut required_columns = RequiredColumns::new(); + let (prunable_predicate, pruning_schema) = + required_columns.rewrite_struct_fields(&predicate, &file_schema)?; let mut properties = PruningExpressionProperties::default(); let predicate_expr = build_predicate_expression( - &predicate, - &file_schema, + &prunable_predicate, + &pruning_schema, &mut required_columns, &unhandled_hook, self.max_in_list_size, @@ -864,6 +868,9 @@ pub struct RequiredColumns { /// * The field the statistics value should be placed in for /// pruning predicate evaluation (e.g. `min_value` or `max_value`) columns: Vec<(phys_expr::Column, StatisticsType, Field)>, + /// Synthetic columns used only while building the pruning expression. + /// The root and path stay separate so dotted names cannot alias a field. + nested_columns: HashMap)>, } impl RequiredColumns { @@ -871,15 +878,75 @@ impl RequiredColumns { Self::default() } + /// Replace exact struct-field access with scalar columns for the existing + /// pruning rules. Neither the original expression nor its schema is changed. + fn rewrite_struct_fields( + &mut self, + predicate: &PhysicalExprRef, + schema: &SchemaRef, + ) -> Result<(PhysicalExprRef, SchemaRef)> { + let mut fields = None; + let rewritten = Arc::clone(predicate) + .transform_down(|expr| { + if expr.downcast_ref::().is_none() { + return Ok(Transformed::no(expr)); + } + let Some((root, path, field)) = struct_field_column(&expr, schema) else { + return Ok(Transformed::no(expr)); + }; + if field.data_type().is_nested() { + return Ok(Transformed::no(expr)); + } + let column = if let Some((column, _)) = + self.nested_columns + .iter() + .find(|(_, (column, field_path))| { + *column == root && *field_path == path + }) { + column.clone() + } else { + let fields = fields.get_or_insert_with(|| schema.fields().to_vec()); + let index = fields.len(); + let mut name = format!("__datafusion_struct_field_{index}"); + while fields.iter().any(|field| field.name() == &name) { + name.push('_'); + } + fields.push(Arc::new( + field.as_ref().clone().with_name(&name).with_nullable(true), + )); + let column = phys_expr::Column::new(&name, index); + self.nested_columns.insert(column.clone(), (root, path)); + column + }; + Ok(Transformed::yes(Arc::new(column) as _)) + })? + .data; + let schema = fields.map_or_else( + || Arc::clone(schema), + |fields| { + Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())) + }, + ); + Ok((rewritten, schema)) + } + /// Returns Some(column) if this is a single column predicate. /// - /// Returns None if this is a multi-column predicate. + /// Returns None for multi-column predicates or nested field statistics. /// /// Examples: /// * `a > 5 OR a < 10` returns `Some(a)` /// * `a > 5 OR b < 10` returns `None` /// * `true` returns None pub fn single_column(&self) -> Option<&phys_expr::Column> { + // Callers of this API expect a top-level schema column (e.g. page pruning). + if self + .columns + .iter() + .any(|(column, _, _)| self.nested_columns.contains_key(column)) + { + return None; + } if self.columns.windows(2).all(|w| { // check if all columns are the same (ignoring statistics and field) let c1 = &w[0].0; @@ -1020,9 +1087,53 @@ impl RequiredColumns { } } +/// Resolve only exact accessor chains rooted at a column. Casts of whole +/// structs and repeated fields are deliberately left to the normal fallback. +fn struct_field_column( + expr: &PhysicalExprRef, + schema: &Schema, +) -> Option<(Column, Vec, FieldRef)> { + if let Some(column) = expr.downcast_ref::() { + let field = schema.fields().get(column.index())?; + if schema + .fields() + .iter() + .filter(|f| f.name() == field.name()) + .count() + != 1 + { + return None; + } + return Some((Column::from_name(field.name()), vec![], Arc::clone(field))); + } + let function = expr.downcast_ref::()?; + let access = function.struct_field_access()?; + let (column, mut path, mut field) = + struct_field_column(&function.args()[access.source_arg], schema)?; + for name in access.field_path { + let DataType::Struct(children) = field.data_type() else { + return None; + }; + let mut matches = children.iter().filter(|child| child.name() == &name); + let child = Arc::clone(matches.next()?); + if matches.next().is_some() { + return None; + } + field = child; + path.push(name); + } + if expr.data_type(schema).ok()? != *field.data_type() { + return None; + } + Some((column, path, field)) +} + impl From> for RequiredColumns { fn from(columns: Vec<(phys_expr::Column, StatisticsType, Field)>) -> Self { - Self { columns } + Self { + columns, + ..Default::default() + } } } @@ -1058,15 +1169,20 @@ fn build_statistics_record_batch( let mut arrays = Vec::::new(); // For each needed statistics column: for (column, statistics_type, stat_field) in required_columns.iter() { - let column = Column::from_name(column.name()); + let root = Column::from_name(column.name()); + let (column, path) = required_columns + .nested_columns + .get(column) + .map(|(column, path)| (column, path.as_slice())) + .unwrap_or((&root, &[])); let data_type = stat_field.data_type(); let num_containers = statistics.num_containers(); let array = match statistics_type { - StatisticsType::Min => statistics.min_values(&column), - StatisticsType::Max => statistics.max_values(&column), - StatisticsType::NullCount => statistics.null_counts(&column), + StatisticsType::Min => statistics.min_values_for_path(column, path), + StatisticsType::Max => statistics.max_values_for_path(column, path), + StatisticsType::NullCount => statistics.null_counts_for_path(column, path), StatisticsType::RowCount => statistics.row_counts(), }; let array = array.unwrap_or_else(|| new_null_array(data_type, num_containers)); @@ -1266,10 +1382,6 @@ fn rewrite_expr_to_prunable( Arc::clone(cast.target_field()), None, )); - // PruningPredicate does not support pruning on nested fields yet. - // End-to-end nested-field pruning also requires Parquet statistics - // extraction to agree with PruningPredicate on a stats representation - // for nested field expressions. Ok((left, op, right)) } else if let Some(try_cast) = column_expr.downcast_ref::() { // `try_cast(col) op lit()` @@ -2717,6 +2829,115 @@ mod tests { } } + #[test] + fn struct_field_pruning_paths_and_unsupported_statistics() { + use datafusion_expr::{ + ScalarUDF, ScalarUDFImpl, Signature, StructFieldAccess, Volatility, + }; + + #[derive(Debug, PartialEq, Eq, Hash)] + struct FieldProbe { + path: Vec, + signature: Signature, + } + impl ScalarUDFImpl for FieldProbe { + fn name(&self) -> &str { + "field_probe" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Int32) + } + fn invoke_with_args( + &self, + _: datafusion_expr::ScalarFunctionArgs, + ) -> Result { + unreachable!( + "pruning must evaluate statistics, not the original accessor" + ) + } + fn struct_field_access( + &self, + _: &[Option], + ) -> Option { + Some(StructFieldAccess { + source_arg: 0, + field_path: self.path.clone(), + }) + } + } + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Struct( + vec![ + Field::new("a.b", DataType::Int32, true), + Field::new( + "a", + DataType::Struct( + vec![Field::new("b", DataType::Int32, true)].into(), + ), + true, + ), + Field::new( + "items", + DataType::List(Arc::new(Field::new( + "item", + DataType::Int32, + true, + ))), + true, + ), + ] + .into(), + ), + true, + )])); + // This provider knows only top-level columns. Its bounds must never + // accidentally be used for a nested field. + let statistics = OneContainerStats { + min_values: Some(Arc::new(Int32Array::from(vec![0]))), + max_values: Some(Arc::new(Int32Array::from(vec![0]))), + num_containers: 1, + }; + for (path, supported) in [ + (vec!["a.b"], true), + (vec!["a", "b"], true), + (vec!["missing"], false), + (vec!["items", "item"], false), + (vec!["a"], false), + ] { + let path = path.into_iter().map(String::from).collect::>(); + let udf = ScalarUDF::from(FieldProbe { + path: path.clone(), + signature: Signature::any(1, Volatility::Immutable), + }) + .with_aliases(["alias_field_probe"]); + let expr = logical2physical(&udf.call(vec![col("s")]).gt(lit(5)), &schema); + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(Arc::clone(&expr)) + .unwrap(); + assert_eq!(predicate.always_true(), !supported, "{path:?}"); + assert_eq!(predicate.schema(), &schema); + assert_eq!(predicate.orig_expr().as_ref(), expr.as_ref()); + assert_eq!(predicate.prune(&statistics).unwrap(), vec![true]); + assert!(predicate.required_columns().single_column().is_none()); + if supported { + assert_eq!(predicate.required_columns.nested_columns.len(), 1); + let (root, actual_path) = predicate + .required_columns + .nested_columns + .values() + .next() + .unwrap(); + assert_eq!(root, &Column::from_name("s")); + assert_eq!(actual_path, &path); + } + } + } + /// Row count should only be referenced once in the pruning expression, even if we need the row count /// for multiple columns. #[test] diff --git a/datafusion/sqllogictest/test_files/projection_pushdown.slt b/datafusion/sqllogictest/test_files/projection_pushdown.slt index 1f9176f7137bd..5842ba3e8e22b 100644 --- a/datafusion/sqllogictest/test_files/projection_pushdown.slt +++ b/datafusion/sqllogictest/test_files/projection_pushdown.slt @@ -292,7 +292,7 @@ logical_plan physical_plan 01)ProjectionExec: expr=[id@0 as id, __datafusion_extracted_2@1 as simple_struct.s[label]] 02)--FilterExec: __datafusion_extracted_1@0 > 150, projection=[id@1, __datafusion_extracted_2@2] -03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id, get_field(s@1, label) as __datafusion_extracted_2], file_type=parquet, predicate=get_field(s@1, value) > 150 +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id, get_field(s@1, label) as __datafusion_extracted_2], file_type=parquet, predicate=get_field(s@1, value) > 150, pruning_predicate=__datafusion_struct_field_2_null_count@1 != row_count@2 AND __datafusion_struct_field_2_max@0 > 150, required_guarantees=[] # Verify correctness query IT @@ -852,7 +852,7 @@ logical_plan physical_plan 01)ProjectionExec: expr=[id@0 as id, __datafusion_extracted_2@1 as nullable_struct.s[label]] 02)--FilterExec: __datafusion_extracted_1@0 IS NOT NULL, projection=[id@1, __datafusion_extracted_2@2] -03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/nullable.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id, get_field(s@1, label) as __datafusion_extracted_2], file_type=parquet, predicate=get_field(s@1, value) IS NOT NULL +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/nullable.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id, get_field(s@1, label) as __datafusion_extracted_2], file_type=parquet, predicate=get_field(s@1, value) IS NOT NULL, pruning_predicate=__datafusion_struct_field_2_null_count@1 != row_count@0, required_guarantees=[] # Verify correctness query IT @@ -1460,7 +1460,7 @@ logical_plan physical_plan 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)] 02)--FilterExec: __datafusion_extracted_1@0 > 150, projection=[id@1] -03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id], file_type=parquet, predicate=get_field(s@1, value) > 150 +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id], file_type=parquet, predicate=get_field(s@1, value) > 150, pruning_predicate=__datafusion_struct_field_2_null_count@1 != row_count@2 AND __datafusion_struct_field_2_max@0 > 150, required_guarantees=[] 04)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Verify correctness - id matches and value > 150 @@ -1499,9 +1499,9 @@ logical_plan physical_plan 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)] 02)--FilterExec: __datafusion_extracted_1@0 > 100, projection=[id@1] -03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id], file_type=parquet, predicate=get_field(s@1, value) > 100 +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id], file_type=parquet, predicate=get_field(s@1, value) > 100, pruning_predicate=__datafusion_struct_field_2_null_count@1 != row_count@2 AND __datafusion_struct_field_2_max@0 > 100, required_guarantees=[] 04)--FilterExec: __datafusion_extracted_2@0 > 3, projection=[id@1] -05)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[get_field(s@1, level) as __datafusion_extracted_2, id], file_type=parquet, predicate=get_field(s@1, level) > 3 AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible +05)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[get_field(s@1, level) as __datafusion_extracted_2, id], file_type=parquet, predicate=get_field(s@1, level) > 3 AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible, pruning_predicate=__datafusion_struct_field_2_null_count@1 != row_count@2 AND __datafusion_struct_field_2_max@0 > 3, required_guarantees=[] # Verify correctness - id matches, value > 100, and level > 3 # Matching ids where value > 100: 2(200), 3(150), 4(300), 5(250) @@ -1608,7 +1608,7 @@ physical_plan 02)--HashJoinExec: mode=CollectLeft, join_type=Left, on=[(id@1, id@0)], projection=[id@1, __datafusion_extracted_2@0, __datafusion_extracted_3@3] 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_2, id], file_type=parquet 04)----FilterExec: __datafusion_extracted_1@0 > 5, projection=[id@1, __datafusion_extracted_3@2] -05)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[get_field(s@1, level) as __datafusion_extracted_1, id, get_field(s@1, level) as __datafusion_extracted_3], file_type=parquet, predicate=get_field(s@1, level) > 5 AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible +05)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[get_field(s@1, level) as __datafusion_extracted_1, id, get_field(s@1, level) as __datafusion_extracted_3], file_type=parquet, predicate=get_field(s@1, level) > 5 AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible, pruning_predicate=__datafusion_struct_field_2_null_count@1 != row_count@2 AND __datafusion_struct_field_2_max@0 > 5, required_guarantees=[] # Verify correctness - left join with level > 5 condition # Only join_right rows with level > 5 are matched: id=1 (level=10), id=4 (level=8) @@ -1740,7 +1740,7 @@ logical_plan 05)--------TableScan: simple_struct projection=[id, s], partial_filters=[get_field(simple_struct.s, Utf8("value")) > Int64(200)] physical_plan 01)FilterExec: __datafusion_extracted_1@0 > 200, projection=[id@1] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id], file_type=parquet, predicate=get_field(s@1, value) > 200 +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id], file_type=parquet, predicate=get_field(s@1, value) > 200, pruning_predicate=__datafusion_struct_field_2_null_count@1 != row_count@2 AND __datafusion_struct_field_2_max@0 > 200, required_guarantees=[] # Verify correctness query I