diff --git a/datafusion/core/tests/parquet/filter_pushdown.rs b/datafusion/core/tests/parquet/filter_pushdown.rs index d337979e5fd00..da9f008d7ce21 100644 --- a/datafusion/core/tests/parquet/filter_pushdown.rs +++ b/datafusion/core/tests/parquet/filter_pushdown.rs @@ -809,3 +809,235 @@ async fn pushed_down_predicate_reports_the_original_error() { "expected the original cast error, got {root:?}" ); } + +/// Uses a different UDF type and argument order from get_field. It deliberately +/// does not simplify to get_field, so the physical consumers must use the capability. +#[derive(Debug, PartialEq, Eq, Hash)] +struct FieldAt { + declare_access: bool, + signature: datafusion_expr::Signature, +} + +impl datafusion_expr::ScalarUDFImpl for FieldAt { + fn name(&self) -> &str { + "field_at" + } + + fn signature(&self) -> &datafusion_expr::Signature { + &self.signature + } + + fn return_type( + &self, + _: &[arrow::datatypes::DataType], + ) -> datafusion_common::Result { + unreachable!("return_field_from_args is implemented") + } + + fn return_field_from_args( + &self, + args: datafusion_expr::ReturnFieldArgs, + ) -> datafusion_common::Result { + datafusion_functions::core::get_field().return_field_from_args( + datafusion_expr::ReturnFieldArgs { + arg_fields: &[args.arg_fields[1].clone(), args.arg_fields[0].clone()], + scalar_arguments: &[args.scalar_arguments[1], args.scalar_arguments[0]], + }, + ) + } + + fn invoke_with_args( + &self, + mut args: datafusion_expr::ScalarFunctionArgs, + ) -> datafusion_common::Result { + args.args.swap(0, 1); + args.arg_fields.swap(0, 1); + datafusion_functions::core::get_field().invoke_with_args(args) + } + + fn struct_field_access( + &self, + literals: &[Option], + ) -> Option { + if !self.declare_access { + return None; + } + Some(datafusion_expr::StructFieldAccess { + source_arg: 1, + field_path: vec![ + literals + .first()? + .as_ref()? + .try_as_str() + .flatten()? + .to_owned(), + ], + }) + } + + fn placement( + &self, + args: &[datafusion_expr::ExpressionPlacement], + ) -> datafusion_expr::ExpressionPlacement { + use datafusion_expr::ExpressionPlacement; + if args[0] == ExpressionPlacement::Literal && args[1].should_push_to_leaves() { + ExpressionPlacement::MoveTowardsLeafNodes + } else { + ExpressionPlacement::KeepInPlace + } + } +} + +#[tokio::test] +async fn custom_struct_accessor_pushdown_and_schema_adaptation() { + use arrow::array::{Array, StructArray}; + use arrow::buffer::NullBuffer; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_expr::{ScalarUDF, Signature, Volatility}; + + let values = Arc::new(Int32Array::from(vec![Some(0), Some(10), None, Some(99)])); + let deep = StructArray::from(vec![( + Arc::new(Field::new("value", DataType::Int32, true)), + values.clone() as ArrayRef, + )]); + let fields = vec![ + Arc::new(Field::new("pad", DataType::Utf8, false)), + Arc::new(Field::new("value", DataType::Int32, true)), + Arc::new(Field::new("deep", deep.data_type().clone(), true)), + Arc::new(Field::new("dotted.name", DataType::Int32, true)), + ]; + let s = StructArray::new( + fields.into(), + vec![ + Arc::new(StringArray::from(vec!["pad"; 4])), + values.clone(), + Arc::new(deep), + values, + ], + Some(NullBuffer::from(vec![true, true, true, false])), + ); + let batch = RecordBatch::try_from_iter(vec![ + ("s", Arc::new(s) as ArrayRef), + ( + "id", + Arc::new(Int32Array::from(vec![0, 1, 2, 3])) as ArrayRef, + ), + ]) + .unwrap(); + let dir = TempDir::new().unwrap(); + let path = dir.path().join("custom-access.parquet"); + let mut writer = + ArrowWriter::try_new(File::create(&path).unwrap(), batch.schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let evolved_schema = Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "s", + DataType::Struct( + vec![ + Arc::new(Field::new( + "deep", + DataType::Struct( + vec![Arc::new(Field::new("value", DataType::Int64, true))] + .into(), + ), + true, + )), + Arc::new(Field::new("missing", DataType::Int32, true)), + Arc::new(Field::new("dotted.name", DataType::Int64, true)), + Arc::new(Field::new("value", DataType::Int64, true)), + ] + .into(), + ), + true, + ), + ]); + for evolved in [false, true] { + for declare_access in [false, true] { + for pushdown in [false, true] { + let mut config = SessionConfig::new().with_target_partitions(1); + config.options_mut().execution.parquet.pushdown_filters = pushdown; + let ctx = SessionContext::new_with_config(config); + ctx.register_udf( + ScalarUDF::from(FieldAt { + declare_access, + signature: Signature::any(2, Volatility::Immutable), + }) + .with_aliases(["alias_field_at"]), + ); + let mut options = ParquetReadOptions::default(); + if evolved { + options = options.schema(&evolved_schema); + } + ctx.register_parquet("t", path.to_str().unwrap(), options) + .await + .unwrap(); + for predicate in [ + "field_at('value', s) > 5", + "alias_field_at('value', s) > 5", + "field_at('value', field_at('deep', s)) > 5", + "field_at('dotted.name', s) > 5", + ] { + let plan = ctx + .sql(&format!("SELECT id FROM t WHERE {predicate} ORDER BY id")) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let results = collect(plan.clone(), ctx.task_ctx()).await.unwrap(); + datafusion_common::assert_batches_eq!( + ["+----+", "| id |", "+----+", "| 1 |", "+----+"], + &results + ); + let metrics = TestParquetFile::parquet_metrics(&plan).unwrap(); + assert_eq!( + get_value(&metrics, "pushdown_rows_pruned"), + if pushdown && declare_access { 3 } else { 0 }, + "evolved={evolved}, capability={declare_access}, pushdown={pushdown}, {predicate}\n{}", + displayable(plan.as_ref()).indent(false) + ); + } + // An explicit cast must still evaluate an unselected sibling. + // Narrowing it to just `value` would silently hide the bad pad cast. + if !evolved && declare_access { + let accessor = ScalarUDF::from(FieldAt { + declare_access: true, + signature: Signature::any(2, Volatility::Immutable), + }); + let target = DataType::Struct( + vec![ + Arc::new(Field::new("pad", DataType::Int32, true)), + Arc::new(Field::new("value", DataType::Int32, true)), + ] + .into(), + ); + let predicate = accessor + .call(vec![lit("value"), datafusion_expr::cast(col("s"), target)]) + .gt(lit(5)); + let error = ctx + .table("t") + .await + .unwrap() + .filter(predicate) + .unwrap() + .collect() + .await + .unwrap_err(); + assert!(error.to_string().contains("pad"), "{error}"); + } + if evolved { + let results = ctx.sql("SELECT id FROM t WHERE field_at('missing', s) IS NULL ORDER BY id").await.unwrap().collect().await.unwrap(); + datafusion_common::assert_batches_eq!( + [ + "+----+", "| id |", "+----+", "| 0 |", "| 1 |", "| 2 |", + "| 3 |", "+----+" + ], + &results + ); + } + } + } + } +} diff --git a/datafusion/datasource-parquet/src/projection_read_plan.rs b/datafusion/datasource-parquet/src/projection_read_plan.rs index 9c7028cb773b1..6d0edc069a342 100644 --- a/datafusion/datasource-parquet/src/projection_read_plan.rs +++ b/datafusion/datasource-parquet/src/projection_read_plan.rs @@ -38,9 +38,8 @@ use datafusion_common::Result; use datafusion_common::nested_struct::requires_nested_struct_cast; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor}; use datafusion_functions::core::file_row_index::FileRowIndexFunc; -use datafusion_functions::core::getfield::GetFieldFunc; -use datafusion_physical_expr::expressions::{CastExpr, Column, Literal}; -use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr::expressions::{CastExpr, Column}; +use datafusion_physical_expr::utils::{collect_columns, reassign_expr_columns}; use datafusion_physical_expr::{PhysicalExpr, ScalarFunctionExpr}; use crate::nested_schema_pruning::{ @@ -64,7 +63,7 @@ pub(crate) struct ParquetReadPlan { pub projected_schema: SchemaRef, } -/// Records a struct field access via `get_field(struct_col, 'field1', 'field2', ...)`. +/// Records a nested input required by an expression, including UDF requirements. /// /// This allows the row filter to project only the specific Parquet leaf columns /// needed by the filter, rather than all leaves of the struct. @@ -272,15 +271,13 @@ impl<'schema> PushdownChecker<'schema> { /// conversions still run. fn check_cast_struct_field_access( &mut self, - func: &ScalarFunctionExpr, + source: &Arc, + field_path: &[String], + return_type: &DataType, ) -> Option { if !self.allow_struct_casts { return None; } - let (source, field_names) = func.args().split_first()?; - if field_names.is_empty() { - return None; - } let cast = source.downcast_ref::()?; let column = cast.expr().downcast_ref::()?; let index = self.file_schema.index_of(column.name()).ok()?; @@ -290,7 +287,6 @@ impl<'schema> PushdownChecker<'schema> { ) { return None; } - let return_type = func.return_type(); if DataType::is_nested(return_type) && !self.is_nested_type_supported(return_type) { return None; @@ -298,21 +294,7 @@ impl<'schema> PushdownChecker<'schema> { // Every key must resolve through Struct fields in the cast target. // In particular, a key following a Map field is a runtime lookup. - let mut data_type = cast.cast_type(); - for field_name in field_names { - let name = field_name - .downcast_ref::()? - .value() - .try_as_str() - .flatten()?; - let DataType::Struct(fields) = data_type else { - return None; - }; - data_type = fields - .iter() - .find(|field| field.name() == name)? - .data_type(); - } + resolve_struct_field_type(cast.cast_type(), field_path)?; self.cast_accesses.push(CastColumnAccess { root_index: index, @@ -394,90 +376,112 @@ impl TreeNodeVisitor<'_> for PushdownChecker<'_> { type Node = Arc; fn f_down(&mut self, node: &Self::Node) -> Result { - // Handle struct field access like `s['foo']['bar'] > 10`. - // - // DataFusion represents nested field access as `get_field(Column("s"), "foo")` - // (or chained: `get_field(get_field(Column("s"), "foo"), "bar")`). - // - // We intercept the outermost `get_field` on the way *down* the tree so - // the visitor never reaches the raw `Column("s")` node. Without this, - // `check_single_column` would see that `s` is a Struct and reject it. - // - // The strategy: - // 1. Match `get_field` whose first arg is a `Column` (the struct root). - // 2. Check that the *resolved* return type is primitive — meaning we've - // drilled all the way to a leaf (e.g. `s['foo']` → Utf8). - // 3. Record the root column index via `check_struct_field_column` and - // return `Jump` to skip visiting the children (the Column and the - // literal field-name args), since we've already handled them. - // - // If the return type is still nested (e.g. `s['nested_struct']` → Struct), - // we fall through and let normal traversal continue, which will - // eventually reject the expression when it hits the struct Column. - if let Some(func) = - ScalarFunctionExpr::try_downcast_func::(node.as_ref()) - { - if let Some(recursion) = self.check_cast_struct_field_access(func) { + // Resolve capability-declaring accessors, including chains with + // different UDFs and argument layouts. Do not look through casts. + let mut source = node; + let mut paths = Vec::new(); + while let Some(function) = source.downcast_ref::() { + let Some(access) = function.struct_field_access() else { + break; + }; + paths.push(access.field_path); + source = &function.args()[access.source_arg]; + } + let field_path = paths.into_iter().rev().flatten().collect::>(); + if !field_path.is_empty() { + let return_type = node.data_type(self.file_schema)?; + if let Some(recursion) = + self.check_cast_struct_field_access(source, &field_path, &return_type) + { return Ok(recursion); } - let args = func.args(); - - if let Some(column) = args.first().and_then(|a| a.downcast_ref::()) { - // for Map columns, get_field performs a runtime key lookup rather than a - // schema-level field access so the entire Map column must be read, - // we skip the struct field optimization and defer to normal Column traversal - let is_map_column = self + if let Some(column) = source.downcast_ref::() { + // Resolve by name: physical column indices may still refer to + // an unprojected schema. Map/List paths must remain opaque. + let leaf_type = self .file_schema - .index_of(column.name()) + .field_with_name(column.name()) .ok() - .map(|idx| { - matches!( - self.file_schema.field(idx).data_type(), - DataType::Map(_, _) - ) - }) - .unwrap_or(false); - - let return_type = func.return_type(); - - if !is_map_column - && (!DataType::is_nested(return_type) - || self.is_nested_type_supported(return_type)) + .and_then(|root| { + resolve_struct_field_type(root.data_type(), &field_path) + }); + if leaf_type.is_some() + && (!return_type.is_nested() + || self.is_nested_type_supported(&return_type)) { - // if any field name argument is not a string literal we cannot - // determine the exact leaf path, so we fall back to reading the - // entire struct root column - let field_path = args[1..] - .iter() - .map(|arg| { - arg.downcast_ref::().and_then(|lit| { - lit.value().try_as_str().flatten().map(|s| s.to_string()) - }) + if let Some(recursion) = + self.check_struct_field_column(column.name(), field_path) + { + return Ok(recursion); + } + return Ok(TreeNodeRecursion::Jump); + } + } + } + + if let Some(function) = node.downcast_ref::() + && let Some(requirements) = function.required_input_fields(self.file_schema) + && !requirements.is_empty() + // Declaring dependencies cannot bypass the List/Map pushdown policy. + && requirements.iter().all(|requirement| { + let argument = &function.args()[requirement.arg_index]; + let data_type = if let Some(column) = argument.downcast_ref::() { + // Column indices can still refer to an earlier schema. + self.file_schema + .field_with_name(column.name()) + .map(|field| field.data_type().clone()) + .ok() + } else { + reassign_expr_columns(Arc::clone(argument), self.file_schema) + .and_then(|argument| argument.data_type(self.file_schema)) + .ok() + }; + data_type.is_some_and(|data_type| { + requirement.field_paths.iter().all(|path| { + resolve_struct_field_type(&data_type, path).is_some_and(|leaf| { + matches!(leaf, DataType::Struct(_)) + || !leaf.is_nested() + || self.is_nested_type_supported(leaf) }) - .collect(); - - match field_path { - Some(path) => { - if let Some(recursion) = - self.check_struct_field_column(column.name(), path) - { - return Ok(recursion); - } - } - None => { - // Could not resolve field path — fall back to - // reading the entire struct root column. - if let Some(recursion) = - self.check_single_column(column.name()) - { - return Ok(recursion); - } + }) + }) + }) + { + for (index, argument) in function.args().iter().enumerate() { + if let Some(requirement) = + requirements.iter().find(|r| r.arg_index == index) + { + if let Some(column) = argument.downcast_ref::() { + for path in &requirement.field_paths { + self.check_struct_field_column(column.name(), path.clone()); } + continue; + } + // A dependency declaration cannot remove conversions from + // an argument. Runtime schema adaptation may have inserted a + // cast; retain its entire target and evaluate it unchanged. + if self.allow_struct_casts + && let Some(cast) = argument.downcast_ref::() + && let Some(column) = cast.expr().downcast_ref::() + && let Ok(root_index) = self.file_schema.index_of(column.name()) + && matches!( + self.file_schema.field(root_index).data_type(), + DataType::Struct(_) + ) + && matches!(cast.cast_type(), DataType::Struct(_)) + { + self.cast_accesses.push(CastColumnAccess { + root_index, + target_type: cast.cast_type().clone(), + }); + continue; } - - return Ok(TreeNodeRecursion::Jump); } + // Unspecified arguments and arbitrary argument expressions must + // still be evaluated, including columns and errors they depend on. + argument.visit(self)?; } + return Ok(TreeNodeRecursion::Jump); } // Handle whole-column casts to a narrower nested type, e.g. @@ -521,6 +525,21 @@ impl TreeNodeVisitor<'_> for PushdownChecker<'_> { } } +/// Resolve literal names through structs, rejecting missing or ambiguous fields. +fn resolve_struct_field_type<'a>( + data_type: &'a DataType, + path: &[String], +) -> Option<&'a DataType> { + path.iter().try_fold(data_type, |data_type, name| { + let DataType::Struct(fields) = data_type else { + return None; + }; + let mut matches = fields.iter().filter(|field| field.name() == name); + let field = matches.next()?; + matches.next().is_none().then_some(field.data_type()) + }) +} + /// Result of checking which columns are required for filter pushdown. #[derive(Debug)] pub(crate) struct PushdownColumns { @@ -1138,6 +1157,457 @@ mod test { use std::collections::HashMap; use tempfile::NamedTempFile; + #[derive(Debug, PartialEq, Eq, Hash)] + struct CustomStructLabel; + + /// Computes a result from two fields and another argument, rather than + /// returning a field unchanged. Only the struct argument is restricted. + #[derive(Debug, PartialEq, Eq, Hash)] + struct LabelScore { + requirements: Option>, + } + + impl datafusion_expr::ScalarUDFImpl for LabelScore { + fn name(&self) -> &str { + "label_score" + } + + fn signature(&self) -> &datafusion_expr::Signature { + static SIGNATURE: std::sync::LazyLock = + std::sync::LazyLock::new(|| { + datafusion_expr::Signature::any( + 2, + datafusion_expr::Volatility::Immutable, + ) + }); + &SIGNATURE + } + + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Int32) + } + + fn required_input_fields( + &self, + args: datafusion_expr::ReturnFieldArgs, + ) -> Option> { + // The source schema is available to the downstream implementation. + assert!(matches!( + args.arg_fields[1].data_type(), + DataType::Struct(_) + )); + self.requirements.clone() + } + + fn invoke_with_args( + &self, + args: datafusion_expr::ScalarFunctionArgs, + ) -> Result { + let id = args.args[0].to_array(args.number_rows)?; + let s = args.args[1].to_array(args.number_rows)?; + let id = id.as_any().downcast_ref::().unwrap(); + let s = s.as_any().downcast_ref::().unwrap(); + let values = s + .column_by_name("value") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let labels = s + .column_by_name("label") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let scores = (0..s.len()) + .map(|i| { + (!s.is_null(i) + && !id.is_null(i) + && !values.is_null(i) + && !labels.is_null(i)) + .then(|| id.value(i) + values.value(i) + labels.value(i).len() as i32) + }) + .collect::(); + Ok(datafusion_expr::ColumnarValue::Array(Arc::new(scores))) + } + } + + fn score_requirement() -> datafusion_expr::InputFieldRequirement { + datafusion_expr::InputFieldRequirement { + arg_index: 1, + field_paths: vec![vec!["value".into()], vec!["label".into()]], + } + } + + #[test] + fn udf_input_requirements_respect_nested_pushdown_policy() { + #[derive(Debug, PartialEq, Eq, Hash)] + struct RequiredFieldsIsNull(Vec); + + impl datafusion_expr::ScalarUDFImpl for RequiredFieldsIsNull { + fn name(&self) -> &str { + "required_fields_is_null" + } + + fn signature(&self) -> &datafusion_expr::Signature { + static SIGNATURE: std::sync::LazyLock = + std::sync::LazyLock::new(|| { + datafusion_expr::Signature::any( + 1, + datafusion_expr::Volatility::Immutable, + ) + }); + &SIGNATURE + } + + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Boolean) + } + + fn required_input_fields( + &self, + _: datafusion_expr::ReturnFieldArgs, + ) -> Option> { + Some(vec![datafusion_expr::InputFieldRequirement { + arg_index: 0, + field_paths: vec![self.0.clone()], + }]) + } + + fn invoke_with_args( + &self, + args: datafusion_expr::ScalarFunctionArgs, + ) -> Result { + Ok(datafusion_expr::ColumnarValue::Array(Arc::new( + arrow::compute::is_null( + args.args[0].to_array(args.number_rows)?.as_ref(), + )?, + ))) + } + } + + let item = Arc::new(Field::new("item", DataType::Int32, true)); + let map = DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int32, true), + ] + .into(), + ), + false, + )), + false, + ); + for data_type in [ + DataType::Int32, + DataType::Struct(vec![Field::new("value", DataType::Int32, true)].into()), + DataType::List(Arc::clone(&item)), + DataType::LargeList(Arc::clone(&item)), + DataType::FixedSizeList(item, 2), + map, + ] { + for nested in [false, true] { + let (input_type, path) = if nested { + ( + DataType::Struct( + vec![Field::new("selected", data_type.clone(), true)].into(), + ), + vec!["selected".into()], + ) + } else { + (data_type.clone(), vec![]) + }; + let schema = Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("s", input_type, true), + ]); + let expr = logical2physical( + &datafusion_expr::ScalarUDF::from(RequiredFieldsIsNull(path)) + .call(vec![col("s")]), + &schema, + ); + for index in [0, 1] { + // Index 0 simulates an expression from an earlier schema. + let expr = Arc::clone(&expr) + .with_new_children(vec![Arc::new(PhysicalColumn::new( + "s", index, + ))]) + .unwrap(); + for allow_lists in [false, true] { + let mut checker = + PushdownChecker::new(&schema, allow_lists, false); + expr.visit(&mut checker).unwrap(); + let blocked = match &data_type { + DataType::Map(_, _) => true, + DataType::List(_) + | DataType::LargeList(_) + | DataType::FixedSizeList(_, _) => !allow_lists, + _ => false, + }; + assert_eq!( + checker.prevents_pushdown(), + blocked, + "{data_type:?}, nested={nested}, index={index}, allow_lists={allow_lists}" + ); + } + } + } + } + } + + #[test] + fn udf_input_requirements_prune_multiple_fields_and_preserve_other_arguments() { + let (file, schema, metadata) = write_id_struct_file_with_handle(); + let descriptor = metadata.file_metadata().schema_descr(); + for requirements in [ + None, + Some(vec![score_requirement()]), + Some(vec![ + datafusion_expr::InputFieldRequirement { + arg_index: 0, + field_paths: vec![vec![]], + }, + score_requirement(), + ]), + ] { + let declare = requirements.is_some(); + let udf = datafusion_expr::ScalarUDF::from(LabelScore { requirements }) + .with_aliases(["score_alias"]); + let expr = logical2physical(&udf.call(vec![col("id"), col("s")]), &schema); + let plan = + build_projection_read_plan(vec![Arc::clone(&expr)], &schema, descriptor); + assert_eq!( + plan.projection_mask, + ProjectionMask::leaves( + descriptor, + if declare { + vec![0, 1, 2] + } else { + vec![0, 1, 2, 3] + } + ) + ); + let mut checker = PushdownChecker::new(&schema, false, false); + expr.visit(&mut checker).unwrap(); + assert_eq!(checker.prevents_pushdown(), !declare); + let mut reader = + ParquetRecordBatchReaderBuilder::try_new(file.reopen().unwrap()) + .unwrap() + .with_projection(plan.projection_mask) + .build() + .unwrap(); + let batch = reader.next().unwrap().unwrap(); + let result = expr + .evaluate(&batch) + .unwrap() + .into_array(batch.num_rows()) + .unwrap(); + assert_eq!( + result.as_ref(), + &Int32Array::from(vec![12, 23, 34]) as &dyn Array + ); + let stale = Arc::clone(&expr) + .with_new_children(vec![ + Arc::new(datafusion_physical_expr::expressions::BinaryExpr::new( + Arc::new(PhysicalColumn::new("id", 1)), + datafusion_expr::Operator::Plus, + Arc::new(datafusion_physical_expr::expressions::Literal::new( + ScalarValue::Int32(Some(0)), + )), + )), + Arc::new(PhysicalColumn::new("s", 0)), + ]) + .unwrap(); + let stale_plan = build_projection_read_plan(vec![stale], &schema, descriptor); + assert_eq!(stale_plan.projected_schema, plan.projected_schema); + // A second consumer of the whole struct overrides narrower requirements. + let plan = build_projection_read_plan( + vec![expr, Arc::new(PhysicalColumn::new("s", 1))], + &schema, + descriptor, + ); + assert_eq!( + plan.projection_mask, + ProjectionMask::leaves(descriptor, [0, 1, 2, 3]) + ); + } + } + + #[test] + fn invalid_udf_input_requirements_keep_full_inputs() { + let (schema, metadata) = write_id_struct_file(); + let descriptor = metadata.file_metadata().schema_descr(); + let requirement = score_requirement(); + for requirements in [ + vec![datafusion_expr::InputFieldRequirement { + arg_index: 2, + ..requirement.clone() + }], + vec![requirement.clone(), requirement.clone()], + vec![datafusion_expr::InputFieldRequirement { + field_paths: vec![], + ..requirement.clone() + }], + vec![datafusion_expr::InputFieldRequirement { + field_paths: vec![vec!["missing".into()]], + ..requirement.clone() + }], + vec![datafusion_expr::InputFieldRequirement { + field_paths: vec![vec!["value".into(), "not_a_struct".into()]], + ..requirement + }], + ] { + let udf = datafusion_expr::ScalarUDF::from(LabelScore { + requirements: Some(requirements), + }); + let expr = logical2physical(&udf.call(vec![col("id"), col("s")]), &schema); + let plan = build_projection_read_plan(vec![expr], &schema, descriptor); + assert_eq!( + plan.projection_mask, + ProjectionMask::leaves(descriptor, [0, 1, 2, 3]) + ); + } + } + + #[test] + fn udf_input_requirements_preserve_argument_cast_errors() { + let (file, schema, metadata) = write_id_struct_file_with_handle(); + let target = DataType::Struct( + vec![ + Arc::new(Field::new("value", DataType::Int32, false)), + Arc::new(Field::new("label", DataType::Utf8, false)), + Arc::new(Field::new("pad", DataType::Int32, false)), + ] + .into(), + ); + let expr = logical2physical( + &datafusion_expr::ScalarUDF::from(LabelScore { + requirements: Some(vec![score_requirement()]), + }) + .call(vec![col("id"), datafusion_expr::cast(col("s"), target)]), + &schema, + ); + // The UDF does not use pad, but its argument's explicit conversion does. + for allow_casts in [false, true] { + let mut checker = PushdownChecker::new(&schema, false, allow_casts); + expr.visit(&mut checker).unwrap(); + assert_eq!(checker.prevents_pushdown(), !allow_casts); + } + let (plan, _) = + crate::row_filter::build_parquet_read_plan(&expr, &schema, &metadata) + .unwrap() + .unwrap(); + assert_eq!( + plan.projection_mask, + ProjectionMask::leaves(metadata.file_metadata().schema_descr(), [0, 1, 2, 3]) + ); + let mut reader = ParquetRecordBatchReaderBuilder::try_new(file.reopen().unwrap()) + .unwrap() + .with_projection(plan.projection_mask) + .build() + .unwrap(); + let batch = reader.next().unwrap().unwrap(); + let error = expr.evaluate(&batch).unwrap_err(); + assert!(error.to_string().contains("pad"), "{error}"); + } + + impl datafusion_expr::ScalarUDFImpl for CustomStructLabel { + fn name(&self) -> &str { + "custom_struct_label" + } + + fn signature(&self) -> &datafusion_expr::Signature { + static SIGNATURE: std::sync::LazyLock = + std::sync::LazyLock::new(|| { + datafusion_expr::Signature::any( + 1, + datafusion_expr::Volatility::Immutable, + ) + }); + &SIGNATURE + } + + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) + } + + fn struct_field_access( + &self, + _: &[Option], + ) -> Option { + Some(datafusion_expr::StructFieldAccess { + source_arg: 0, + field_path: vec!["label".into()], + }) + } + + fn invoke_with_args( + &self, + mut args: datafusion_expr::ScalarFunctionArgs, + ) -> Result { + args.args + .push(datafusion_expr::ColumnarValue::Scalar(ScalarValue::Utf8( + Some("label".into()), + ))); + args.arg_fields + .push(Arc::new(Field::new("key", DataType::Utf8, false))); + get_field().invoke_with_args(args) + } + } + + #[test] + fn custom_struct_accessor_prunes_leaves_and_allows_row_filter() { + let (schema, metadata) = write_id_struct_file(); + let expr = logical2physical( + &datafusion_expr::ScalarUDF::from(CustomStructLabel).call(vec![col("s")]), + &schema, + ); + let schema_descr = metadata.file_metadata().schema_descr(); + let plan = build_projection_read_plan(vec![expr.clone()], &schema, schema_descr); + assert_eq!( + plan.projection_mask, + ProjectionMask::leaves(schema_descr, [2]) + ); + let mut checker = PushdownChecker::new(&schema, false, false); + expr.visit(&mut checker).unwrap(); + assert!(!checker.prevents_pushdown()); + } + + #[test] + fn custom_struct_accessor_does_not_prune_map_entries() { + let map = DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Arc::new(Field::new("key", DataType::Utf8, false)), + Arc::new(Field::new("value", DataType::Utf8, true)), + ] + .into(), + ), + false, + )), + false, + ); + let schema = Arc::new(Schema::new(vec![Field::new("m", map, true)])); + let parquet_schema = ArrowSchemaConverter::new().convert(&schema).unwrap(); + let expr = logical2physical( + &datafusion_expr::ScalarUDF::from(CustomStructLabel).call(vec![col("m")]), + &schema, + ); + let mut checker = PushdownChecker::new(&schema, false, false); + expr.visit(&mut checker).unwrap(); + assert!(checker.prevents_pushdown()); + let plan = build_projection_read_plan(vec![expr], &schema, &parquet_schema); + assert_eq!( + plan.projection_mask, + ProjectionMask::leaves(&parquet_schema, [0, 1]) + ); + } + #[test] fn projection_read_plan_preserves_full_struct() { // Schema: id (Int32), s (Struct{value: Int32, label: Utf8}) diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index a904422989942..9a346ac975f08 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -148,7 +148,8 @@ pub use udaf::{ udaf_default_window_function_display_name, udaf_default_window_function_schema_name, }; pub use udf::{ - ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, StructFieldMapping, + InputFieldRequirement, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, + StructFieldAccess, StructFieldMapping, }; pub use udwf::{LimitEffect, ReversedUDWF, WindowUDF, WindowUDFImpl}; pub use window_frame::{WindowFrame, WindowFrameBound, WindowFrameUnits}; diff --git a/datafusion/expr/src/udf.rs b/datafusion/expr/src/udf.rs index 5caccc87121ef..1db8a8a9c646b 100644 --- a/datafusion/expr/src/udf.rs +++ b/datafusion/expr/src/udf.rs @@ -57,6 +57,30 @@ pub struct StructFieldMapping { pub fields: Vec<(Vec, usize)>, } +/// A struct field read by a scalar function. Field names are separate path +/// components: a literal dot in a name is not a path separator. +/// +/// See [`ScalarUDFImpl::struct_field_access`] for the semantic contract. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StructFieldAccess { + /// Argument containing the struct. + pub source_arg: usize, + /// Non-empty path through struct fields beneath the source argument. + pub field_path: Vec, +} + +/// Nested fields sufficient to evaluate one argument of a scalar function. +/// See [`ScalarUDFImpl::required_input_fields`] for the projection contract. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct InputFieldRequirement { + /// Index of the argument whose struct fields may be pruned. + pub arg_index: usize, + /// Non-empty list of paths through struct fields. An empty path selects + /// the entire argument; a path ending at a nested field selects its whole + /// subtree. Names are literal components, including any dots. + pub field_paths: Vec>, +} + /// Logical representation of a Scalar User Defined Function. /// /// A scalar function produces a single row output for each row of input. This @@ -320,6 +344,22 @@ impl ScalarUDF { self.inner.evaluate_bounds(inputs) } + /// See [`ScalarUDFImpl::struct_field_access`] for more details. + pub fn struct_field_access( + &self, + literal_args: &[Option], + ) -> Option { + self.inner.struct_field_access(literal_args) + } + + /// See [`ScalarUDFImpl::required_input_fields`]. + pub fn required_input_fields( + &self, + args: ReturnFieldArgs, + ) -> Option> { + self.inner.required_input_fields(args) + } + /// See [`ScalarUDFImpl::struct_field_mapping`] for more details. pub fn struct_field_mapping( &self, @@ -998,6 +1038,60 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any { not_impl_err!("Function {} does not implement coerce_types", self.name()) } + /// Describe this call as an exact struct-field extraction, if possible. + /// + /// `literal_args[i]` contains argument `i` when it is a known literal. + /// Returning `Some` lets readers decode only the selected field and lets + /// schema adapters narrow struct casts to that field. All arguments other + /// than `source_arg` must be literals. The source may itself be an expression. + /// + /// The function must return the selected field unchanged, including its + /// type, metadata, and nulls inherited from every struct ancestor. It must + /// work with reordered or narrowed source structs, resolving fields by name, + /// and cannot depend on siblings. It must not perform additional conversions, + /// introduce errors, or replace nulls with defaults. Readers validate that + /// the path traverses structs; this does not describe runtime Map lookups or + /// repeated List elements. Return `None` when these guarantees do not hold. + /// + /// This is independent of [`Self::placement`], which describes where an + /// expression should execute rather than which field it reads. + fn struct_field_access( + &self, + _literal_args: &[Option], + ) -> Option { + None + } + + /// Describe which nested input fields suffice to evaluate this call. + /// + /// The supplied argument fields describe the schema at the point of use, + /// and `scalar_arguments` exposes known literals. Each returned entry may + /// restrict one argument to the union of its field paths. Arguments without + /// an entry are evaluated normally. Return `None` if no restriction is known. + /// Each argument index must occur at most once and each path must exist in + /// the supplied schema, traversing only Struct fields. A path can select an + /// entire List or Map, but cannot select individual elements or entries. + /// + /// The function must produce the same values, output field and errors when + /// unselected fields are removed, retaining selected fields and their + /// ancestors' metadata and validity. It must resolve fields by name and + /// accept any union of these requirements with other consumers' fields. + /// Required metadata, fallback values and fields needed for validation must + /// all be included. The declaration does not permit skipping evaluation of + /// argument expressions or their conversions. + /// + /// Readers may prune inputs and still evaluate the original function. This + /// does not assert that its output equals a field (see + /// [`Self::struct_field_access`]), provide output statistics, or authorize + /// moving the function across arbitrary operators (see [`Self::placement`]). + /// Existing implementations and unknown layouts retain their full inputs. + fn required_input_fields( + &self, + _args: ReturnFieldArgs, + ) -> Option> { + None + } + /// For struct-producing functions, return how output fields map to input /// arguments. This enables the optimizer to propagate orderings through /// struct projections. @@ -1163,6 +1257,13 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl { self.inner.propagate_constraints(interval, inputs) } + fn struct_field_access( + &self, + literal_args: &[Option], + ) -> Option { + self.inner.struct_field_access(literal_args) + } + fn struct_field_mapping( &self, literal_args: &[Option], @@ -1170,6 +1271,13 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl { self.inner.struct_field_mapping(literal_args) } + fn required_input_fields( + &self, + args: ReturnFieldArgs, + ) -> Option> { + self.inner.required_input_fields(args) + } + fn output_ordering(&self, inputs: &[ExprProperties]) -> Result { self.inner.output_ordering(inputs) } diff --git a/datafusion/functions/src/core/getfield.rs b/datafusion/functions/src/core/getfield.rs index 388daa32bf067..562f91fad4cb7 100644 --- a/datafusion/functions/src/core/getfield.rs +++ b/datafusion/functions/src/core/getfield.rs @@ -642,6 +642,24 @@ impl ScalarUDFImpl for GetFieldFunc { self.doc() } + fn struct_field_access( + &self, + literal_args: &[Option], + ) -> Option { + let (_, keys) = literal_args.split_first()?; + if keys.is_empty() { + return None; + } + let field_path = keys + .iter() + .map(|key| Some(key.as_ref()?.try_as_str().flatten()?.to_owned())) + .collect::>>()?; + Some(datafusion_expr::StructFieldAccess { + source_arg: 0, + field_path, + }) + } + fn placement(&self, args: &[ExpressionPlacement]) -> ExpressionPlacement { // get_field can be pushed to leaves if: // 1. The base (first arg) is a column or already placeable at leaves diff --git a/datafusion/physical-expr-adapter/src/schema_rewriter.rs b/datafusion/physical-expr-adapter/src/schema_rewriter.rs index 2548ffc6fb1b7..3f85e6e48692c 100644 --- a/datafusion/physical-expr-adapter/src/schema_rewriter.rs +++ b/datafusion/physical-expr-adapter/src/schema_rewriter.rs @@ -32,6 +32,7 @@ use datafusion_common::{ nested_struct::{requires_nested_struct_cast, validate_data_type_compatibility}, tree_node::{Transformed, TransformedResult, TreeNode}, }; +#[cfg(test)] use datafusion_functions::core::getfield::GetFieldFunc; use datafusion_physical_expr::PhysicalExprSimplifier; use datafusion_physical_expr::expressions::Literal; @@ -450,15 +451,13 @@ impl DefaultPhysicalExprAdapterRewriter { &self, expr: &Arc, ) -> Result>> { - let Some(get_field_expr) = - ScalarFunctionExpr::try_downcast_func::(expr.as_ref()) - else { + let Some(get_field_expr) = expr.downcast_ref::() else { return Ok(None); }; - let Some((source_expr, field_name_exprs)) = get_field_expr.args().split_first() - else { + let Some(access) = get_field_expr.struct_field_access() else { return Ok(None); }; + let source_expr = &get_field_expr.args()[access.source_arg]; let Some(cast) = source_expr.downcast_ref::() else { return Ok(None); }; @@ -469,19 +468,11 @@ impl DefaultPhysicalExprAdapterRewriter { return Ok(None); } - // Every key has to be a string literal, otherwise the leaf field - // cannot be resolved statically. - let mut field_path = Vec::with_capacity(field_name_exprs.len()); - for field_name_expr in field_name_exprs { - let Some(field_name) = field_name_expr - .downcast_ref::() - .and_then(|lit| lit.value().try_as_str().flatten()) - else { - return Ok(None); - }; - field_path.push(field_name); - } - // A `get_field` with no keys is not a field access we can narrow. + let field_path = access + .field_path + .iter() + .map(String::as_str) + .collect::>(); let Some((first_key, rest_keys)) = field_path.split_first() else { return Ok(None); }; @@ -559,7 +550,7 @@ impl DefaultPhysicalExprAdapterRewriter { return Ok(None); }; let mut args = get_field_expr.args().to_vec(); - args[0] = Arc::new(CastExpr::new_with_target_field( + args[access.source_arg] = Arc::new(CastExpr::new_with_target_field( Arc::clone(inner), target_field, Some(cast.cast_options().clone()), @@ -569,9 +560,8 @@ impl DefaultPhysicalExprAdapterRewriter { // Rebuild `get_field` over the uncast struct so its return field is // recomputed from the physical field type. - let mut args = Vec::with_capacity(get_field_expr.args().len()); - args.push(Arc::clone(inner)); - args.extend(field_name_exprs.iter().map(Arc::clone)); + let mut args = get_field_expr.args().to_vec(); + args[access.source_arg] = Arc::clone(inner); let extracted = Arc::new(ScalarFunctionExpr::try_new( Arc::new(get_field_expr.fun().clone()), args, @@ -594,75 +584,57 @@ impl DefaultPhysicalExprAdapterRewriter { )))) } - /// Attempt to rewrite struct field access expressions to return null if the field does not exist in the physical schema. - /// Note that this does *not* handle nested struct fields, only top-level struct field access. - /// See for more details. + /// Replace a field access with null when its path exists in the logical + /// schema but is missing from the physical struct. fn try_rewrite_struct_field_access( &self, expr: &Arc, ) -> Result>> { - let Some(get_field_expr) = - ScalarFunctionExpr::try_downcast_func::(expr.as_ref()) - else { + let Some(function) = expr.downcast_ref::() else { return Ok(None); }; - - let Some(source_expr) = get_field_expr.args().first() else { - return Ok(None); - }; - - let Some(field_name_expr) = get_field_expr.args().get(1) else { - return Ok(None); - }; - - let Some(lit) = field_name_expr.downcast_ref::() else { - return Ok(None); - }; - - let Some(field_name) = lit.value().try_as_str().flatten() else { + let Some(access) = function.struct_field_access() else { return Ok(None); }; - - let Some(column) = source_expr.downcast_ref::() else { + let Some(column) = function.args()[access.source_arg].downcast_ref::() + else { return Ok(None); }; - let Ok(physical_field) = self.physical_file_schema.field_with_name(column.name()) else { return Ok(None); }; - - let DataType::Struct(physical_struct_fields) = physical_field.data_type() else { - return Ok(None); - }; - - if physical_struct_fields - .iter() - .any(|f| f.name() == field_name) - { - return Ok(None); - } - let Ok(logical_field) = self.logical_file_schema.field_with_name(column.name()) else { return Ok(None); }; - - let DataType::Struct(logical_struct_fields) = logical_field.data_type() else { + let (DataType::Struct(physical_fields), DataType::Struct(logical_fields)) = + (physical_field.data_type(), logical_field.data_type()) + else { return Ok(None); }; - - let Some(logical_struct_field) = logical_struct_fields + let path = access + .field_path .iter() - .find(|f| f.name() == field_name) - else { + .map(String::as_str) + .collect::>(); + let Some((first, rest)) = path.split_first() else { return Ok(None); }; - - let null_value = ScalarValue::Null.cast_to(logical_struct_field.data_type())?; + if !matches!( + resolve_field_path(logical_fields, first, rest), + FieldPathResolution::Found(_) + ) || !matches!( + resolve_field_path(physical_fields, first, rest), + FieldPathResolution::Missing + ) { + return Ok(None); + } + let return_field = expr.return_field(&self.logical_file_schema)?; + let null_value = ScalarValue::Null.cast_to(return_field.data_type())?; Ok(Some(Arc::new(Literal::new_with_metadata( null_value, - Some(FieldMetadata::from(logical_struct_field.as_ref())), + Some(FieldMetadata::from(return_field.as_ref())), )))) } diff --git a/datafusion/physical-expr/src/scalar_function.rs b/datafusion/physical-expr/src/scalar_function.rs index 6a5ab219aa8dd..a832883a1132e 100644 --- a/datafusion/physical-expr/src/scalar_function.rs +++ b/datafusion/physical-expr/src/scalar_function.rs @@ -34,7 +34,7 @@ use std::hash::{Hash, Hasher}; use std::sync::Arc; use crate::PhysicalExpr; -use crate::expressions::Literal; +use crate::expressions::{Column, Literal}; use arrow::array::{Array, RecordBatch}; use arrow::datatypes::{DataType, FieldRef, Schema}; @@ -158,6 +158,90 @@ impl ScalarFunctionExpr { &self.config_options } + /// Describe this call's struct-field access, if the UDF supports it. + /// Callers must also validate the path against their source schema; a + /// syntactically identical call may perform a Map lookup instead. + pub fn struct_field_access(&self) -> Option { + let literals = self + .args + .iter() + .map(|arg| { + arg.downcast_ref::() + .map(|literal| literal.value().clone()) + }) + .collect::>(); + let access = self.fun.struct_field_access(&literals)?; + if access.source_arg >= self.args.len() + || access.field_path.is_empty() + || literals + .iter() + .enumerate() + .any(|(index, literal)| index != access.source_arg && literal.is_none()) + { + return None; + } + Some(access) + } + + /// Ask the UDF for input requirements against this schema, validating every + /// argument index and struct path. Invalid declarations retain full inputs. + /// Column names are resolved again because projection analysis can receive + /// expressions with indices from an earlier schema. + pub fn required_input_fields( + &self, + schema: &Schema, + ) -> Option> { + let fields = self + .args + .iter() + .map(|arg| { + if let Some(column) = arg.downcast_ref::() { + Some(Arc::clone( + schema.fields().get(schema.index_of(column.name()).ok()?)?, + )) + } else { + crate::utils::reassign_expr_columns(Arc::clone(arg), schema) + .ok()? + .return_field(schema) + .ok() + } + }) + .collect::>>()?; + let literals = self + .args + .iter() + .map(|arg| arg.downcast_ref::().map(Literal::value)) + .collect::>(); + let requirements = self.fun.required_input_fields(ReturnFieldArgs { + arg_fields: &fields, + scalar_arguments: &literals, + })?; + let mut seen = vec![false; self.args.len()]; + for requirement in &requirements { + let field = fields.get(requirement.arg_index)?; + if std::mem::replace(&mut seen[requirement.arg_index], true) + || requirement.field_paths.is_empty() + { + return None; + } + for path in &requirement.field_paths { + let mut data_type = field.data_type(); + for name in path { + let DataType::Struct(children) = data_type else { + return None; + }; + let mut matches = + children.iter().filter(|child| child.name() == name); + data_type = matches.next()?.data_type(); + if matches.next().is_some() { + return None; + } + } + } + } + Some(requirements) + } + /// Given an arbitrary PhysicalExpr attempt to downcast it to a ScalarFunctionExpr /// and verify that its inner function is of type T. /// If the downcast fails, or the function is not of type T, returns `None`. @@ -385,6 +469,40 @@ mod tests { } } + #[test] + fn struct_field_access_requires_literal_keys() { + let fun = datafusion_functions::core::get_field(); + let source = Arc::new(Column::new("s", 0)) as Arc; + let key = Arc::new(Literal::new(ScalarValue::Utf8(Some("a.b".into())))) + as Arc; + let make_expr = |args| { + ScalarFunctionExpr::new( + "get_field", + Arc::clone(&fun), + args, + Arc::new(Field::new("result", DataType::Int32, true)), + Arc::new(ConfigOptions::default()), + ) + }; + assert_eq!( + make_expr(vec![Arc::clone(&source), key]).struct_field_access(), + Some(datafusion_expr::StructFieldAccess { + source_arg: 0, + field_path: vec!["a.b".into()], + }) + ); + assert!( + make_expr(vec![Arc::clone(&source)]) + .struct_field_access() + .is_none() + ); + assert!( + make_expr(vec![Arc::clone(&source), source]) + .struct_field_access() + .is_none() + ); + } + #[test] fn test_scalar_function_volatile_node() { // Create a volatile UDF