adriangb commented on code in PR #22239:
URL: https://github.com/apache/datafusion/pull/22239#discussion_r3293847823
##########
datafusion/functions/src/core/getfield.rs:
##########
@@ -828,4 +924,181 @@ mod tests {
let args = vec![ExpressionPlacement::Literal,
ExpressionPlacement::Literal];
assert_eq!(func.placement(&args), ExpressionPlacement::KeepInPlace);
}
+
+ // --- get_field over struct constructor simplification
--------------------
+
+ use datafusion_common::Column;
+ use datafusion_expr::simplify::SimplifyContext;
+
+ /// A non-empty string literal expression.
+ fn lit_str(s: &str) -> Expr {
+ Expr::Literal(ScalarValue::Utf8(Some(s.to_string())), None)
+ }
+
+ /// A column reference expression.
+ fn col(name: &str) -> Expr {
+ Expr::Column(Column::from_name(name))
+ }
+
+ fn scalar_fn(udf: ScalarUDF, args: Vec<Expr>) -> Expr {
+ Expr::ScalarFunction(ScalarFunction::new_udf(Arc::new(udf), args))
+ }
+
+ /// `named_struct(name1, value1, name2, value2, ...)`.
+ fn named_struct(pairs: Vec<(&str, Expr)>) -> Expr {
+ let args = pairs
+ .into_iter()
+ .flat_map(|(name, value)| [lit_str(name), value])
+ .collect();
+ scalar_fn(ScalarUDF::new_from_impl(NamedStructFunc::new()), args)
+ }
+
+ /// `struct(value0, value1, ...)`.
+ fn struct_fn(values: Vec<Expr>) -> Expr {
+ scalar_fn(ScalarUDF::new_from_impl(StructFunc::new()), values)
+ }
+
+ /// `get_field(args...)`.
+ fn get_field(args: Vec<Expr>) -> Expr {
+ scalar_fn(ScalarUDF::new_from_impl(GetFieldFunc::new()), args)
+ }
+
+ /// Run `GetFieldFunc::simplify` once and return the rewritten expression,
+ /// panicking if the input was left unchanged.
+ fn simplified(args: Vec<Expr>) -> Expr {
+ match GetFieldFunc::new()
+ .simplify(args, &SimplifyContext::default())
+ .unwrap()
+ {
+ ExprSimplifyResult::Simplified(expr) => expr,
+ ExprSimplifyResult::Original(args) => {
+ panic!("expected the expression to be simplified, got
{args:?}")
+ }
+ }
+ }
+
+ /// Assert that `GetFieldFunc::simplify` leaves the arguments unchanged.
+ fn assert_not_simplified(args: Vec<Expr>) {
+ match GetFieldFunc::new()
+ .simplify(args.clone(), &SimplifyContext::default())
+ .unwrap()
+ {
+ ExprSimplifyResult::Original(unchanged) => assert_eq!(unchanged,
args),
+ ExprSimplifyResult::Simplified(expr) => {
+ panic!("expected no simplification, got {expr:?}")
+ }
+ }
+ }
+
+ #[test]
+ fn simplify_get_field_named_struct_returns_matching_value() {
+ // get_field(named_struct('min', a, 'max', b), 'max') => b
+ let expr = get_field(vec![
+ named_struct(vec![("min", col("a")), ("max", col("b"))]),
+ lit_str("max"),
+ ]);
+ let Expr::ScalarFunction(ScalarFunction { args, .. }) = expr else {
+ unreachable!()
+ };
+ assert_eq!(simplified(args), col("b"));
+ }
+
Review Comment:
Done in 11a390ea0d — the test now builds `args` directly and passes it to
`simplified(args)`, matching the neighboring tests.
##########
datafusion/functions/src/core/getfield.rs:
##########
@@ -249,6 +252,96 @@ fn extract_single_field(base: ColumnarValue, name:
ScalarValue) -> Result<Column
}
}
+/// Try to simplify a `get_field` call whose base is an inline struct
+/// constructor by resolving the field access at plan time.
+///
+/// Handles both struct constructors:
+/// * `named_struct('a', x, 'b', y)` — fields are looked up by name.
+/// * `struct(x, y)` — fields are positional and named `c0`, `c1`, ...
+///
+/// For example:
+/// * `get_field(named_struct('min', a, 'max', b), 'max')` => `b`
+/// * `get_field(struct(a, b), 'c1')` => `b`
+///
+/// `args` is the (already flattened) argument list of the `get_field` call:
+/// `[base, field_name, rest_of_path...]`. When extra path elements remain
+/// after resolving the first one (`get_field(named_struct('s', inner), 's',
'k')`),
+/// the resolved value is re-wrapped in a `get_field` call for the remaining
+/// path so the simplifier can recurse into it on the next pass.
+///
+/// Returns `None` — leaving the expression untouched — whenever the rewrite
+/// cannot be proven safe, e.g. a non-literal field name, a `named_struct`
+/// with a non-literal field name (which might shadow the requested field at
+/// runtime), or a field the constructor does not produce.
+fn simplify_get_field_over_struct_constructor(args: &[Expr]) -> Option<Expr> {
+ let [base, field_name, rest @ ..] = args else {
+ return None;
+ };
+
+ // The accessed field name must be a non-empty string literal.
+ let Expr::Literal(field_name, _) = field_name else {
+ return None;
+ };
+ let field_name = field_name
+ .try_as_str()
+ .flatten()
+ .filter(|s| !s.is_empty())?;
+
Review Comment:
Done in b2e0978382 — added `simplify_get_field_null_field_name_left_alone`
covering the `ScalarValue::Utf8(None)` case, asserting the expression is left
unchanged.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
