kosiew commented on code in PR #24490:
URL: https://github.com/apache/datafusion/pull/24490#discussion_r3911034500
##########
datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs:
##########
@@ -1455,6 +1543,112 @@ impl TreeNodeRewriter for Simplifier<'_> {
}))
}
+ // Push an immutable scalar function into the branches of a CASE
with
+ // literal outputs, eagerly folding each branch; if any branch
fails
+ // to fold, the expression is left untouched:
+ // f(CASE WHEN X THEN "a" ELSE "b" END) --> CASE WHEN X THEN
f("a") ELSE f("b") END
+ Expr::ScalarFunction(ScalarFunction { func, args })
+ if func.signature().volatility == Volatility::Immutable =>
+ {
+ // The function's own simplify() gets the first shot
(arrow_cast
+ // must rewrite itself into a Cast); push down only if it
+ // declines. A third-party simplify() may hand back rewritten
+ // args despite Original's keep-unmodified contract, so the
+ // pushdown target is located on whatever comes back instead of
+ // being assumed from a pre-check.
+ let mut args = match func.simplify(args, info)? {
+ ExprSimplifyResult::Simplified(expr) => {
+ return Ok(Transformed::yes(expr));
+ }
+ ExprSimplifyResult::Original(args) => args,
+ };
+ let Some(case_position) = case_pushdown_target(&args) else {
+ return
Ok(Transformed::no(Expr::ScalarFunction(ScalarFunction {
+ func,
+ args,
+ })));
+ };
+ // The scan above matched a CASE at this position; move it out,
+ // leaving a placeholder `fold_case_branch` never reads.
+ let Expr::Case(case) = std::mem::replace(
+ &mut args[case_position],
+ Expr::Literal(ScalarValue::Null, None),
+ ) else {
+ return internal_err!("case pushdown target is not a CASE
argument");
+ };
+
+ // Types the untyped-NULL branches and the implicit-NULL ELSE;
+ // all-NULL outputs leave nothing to type from.
+ let output_type = case
+ .when_then_expr
+ .iter()
+ .map(|(_, then)| then.as_ref())
+ .chain(case.else_expr.as_deref())
+ .find_map(|e| match e {
+ Expr::Literal(s, _) if s.data_type() != DataType::Null
=> {
+ Some(s.data_type())
+ }
+ _ => None,
+ });
+ let Some(output_type) = output_type else {
+ args[case_position] = Expr::Case(case);
+ return
Ok(Transformed::no(Expr::ScalarFunction(ScalarFunction {
+ func,
+ args,
+ })));
+ };
+
+ let mut const_evaluator =
+
ConstEvaluator::try_new(Some(Arc::clone(info.config_options())))?;
+
+ let mut folded = Vec::with_capacity(case.when_then_expr.len()
+ 1);
+ let implicit_null =
+ Expr::Literal(ScalarValue::try_new_null(&output_type)?,
None);
+ let branches = case
+ .when_then_expr
+ .iter()
+ .map(|(_, then)| then.as_ref())
+ .chain(std::iter::once(
+ case.else_expr.as_deref().unwrap_or(&implicit_null),
+ ));
+ let mut fold_failed = false;
+ for branch in branches {
+ match fold_case_branch(
+ &mut const_evaluator,
+ &func,
+ &args,
+ case_position,
+ &output_type,
+ branch,
+ )? {
+ Some(lit) => folded.push(lit),
+ None => {
+ fold_failed = true;
+ break;
+ }
+ }
+ }
+ if fold_failed {
+ args[case_position] = Expr::Case(case);
+ return
Ok(Transformed::no(Expr::ScalarFunction(ScalarFunction {
+ func,
+ args,
+ })));
+ }
+
+ let folded_else = folded.pop().map(Box::new);
+ Transformed::yes(Expr::Case(Case {
Review Comment:
I think this can silently lose metadata from the original scalar-function
field. `ReturnFieldArgs` allows a UDF to choose its field based on whether an
argument is scalar. For example, an immutable UDF could attach metadata when it
receives the original non-literal CASE (`scalar_arguments[i] == None`), but
return no metadata for each literal substituted by `fold_case_branch`.
In that case, every folded branch passes the current metadata check, but the
resulting CASE no longer has the metadata from the original function call. That
changes the output schema contract.
Could we compute the field of the original function call using `info`'s
schema and only accept the rewrite when the resulting CASE preserves its type,
nullability, and metadata? Alternatively, we could preserve those attributes
explicitly if there is a suitable way to do that. It would also be good to add
a regression UDF whose metadata is present only when the CASE argument is
non-scalar.
##########
datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs:
##########
@@ -4015,6 +4209,450 @@ mod tests {
);
}
+ /// Parses Utf8 to Int64; -1 for NULL (not null-propagating), error on
junk.
+ #[derive(Debug, PartialEq, Eq, Hash)]
+ struct ParseIntUdf {
+ signature: Signature,
+ }
+
+ impl ParseIntUdf {
+ fn new(volatility: Volatility) -> Self {
+ Self {
+ signature: Signature::variadic_any(volatility),
+ }
+ }
+ }
+
+ impl ScalarUDFImpl for ParseIntUdf {
+ fn name(&self) -> &str {
+ "parse_int_udf"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
Result<ColumnarValue> {
+ match args.args.first() {
+ Some(ColumnarValue::Scalar(ScalarValue::Utf8(None))) => {
+ Ok(ColumnarValue::Scalar(ScalarValue::Int64(Some(-1))))
+ }
+ Some(ColumnarValue::Scalar(ScalarValue::Utf8(Some(s)))) => {
+ match s.parse::<i64>() {
+ Ok(v) =>
Ok(ColumnarValue::Scalar(ScalarValue::Int64(Some(v)))),
+ Err(_) => datafusion_common::exec_err!(
+ "parse_int_udf: invalid integer {s:?}"
+ ),
+ }
+ }
+ _ => {
+ datafusion_common::exec_err!("parse_int_udf: expected a
Utf8 scalar")
+ }
+ }
+ }
+ }
+
+ fn parse_int(args: Vec<Expr>) -> Expr {
+ Expr::ScalarFunction(ScalarFunction::new_udf(
+ Arc::new(ScalarUDF::new_from_impl(ParseIntUdf::new(
+ Volatility::Immutable,
+ ))),
+ args,
+ ))
+ }
+
+ /// `CASE WHEN c2_non_null THEN then ELSE els END` (no ELSE when `els` is
None).
+ fn case_on_c2(then: &str, els: Option<&str>) -> Expr {
+ Expr::Case(Case::new(
+ None,
+ vec![(Box::new(col("c2_non_null")), Box::new(lit(then)))],
+ els.map(|e| Box::new(lit(e))),
+ ))
+ }
+
+ /// `ParseIntUdf` that tags its output field with metadata.
+ #[derive(Debug, PartialEq, Eq, Hash)]
+ struct MetadataParseIntUdf {
+ inner: ParseIntUdf,
+ }
+
+ impl ScalarUDFImpl for MetadataParseIntUdf {
+ fn name(&self) -> &str {
+ "metadata_parse_int_udf"
+ }
+
+ fn signature(&self) -> &Signature {
+ self.inner.signature()
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ self.inner.return_type(arg_types)
+ }
+
+ fn return_field_from_args(&self, _args: ReturnFieldArgs) ->
Result<FieldRef> {
+ Ok(Arc::new(
+ Field::new(self.name(), DataType::Int64, true).with_metadata(
+ [("extension".to_string(), "tagged".to_string())].into(),
+ ),
+ ))
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
Result<ColumnarValue> {
+ self.inner.invoke_with_args(args)
+ }
+ }
+
+ #[test]
+ fn simplify_case_pushdown_declines_metadata_producing_functions() {
+ // CASE never surfaces branch-literal metadata in its schema field, so
+ // folding would silently drop the tag; the rewrite must decline.
+ let expr = Expr::ScalarFunction(ScalarFunction::new_udf(
+ Arc::new(ScalarUDF::new_from_impl(MetadataParseIntUdf {
+ inner: ParseIntUdf::new(Volatility::Immutable),
+ })),
+ vec![case_on_c2("1", Some("2"))],
+ ));
+ assert_eq!(simplify(expr.clone()), expr);
+ }
+
+ /// Breaches `ExprSimplifyResult::Original`'s keep-args-unmodified
+ /// contract: swaps the CASE argument for a plain literal.
+ #[derive(Debug, PartialEq, Eq, Hash)]
+ struct ContractBreachingUdf {
+ inner: ParseIntUdf,
+ }
+
+ impl ScalarUDFImpl for ContractBreachingUdf {
+ fn name(&self) -> &str {
+ "contract_breaching_udf"
+ }
+
+ fn signature(&self) -> &Signature {
+ self.inner.signature()
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ self.inner.return_type(arg_types)
+ }
+
+ fn simplify(
+ &self,
+ args: Vec<Expr>,
+ _info: &SimplifyContext,
+ ) -> Result<ExprSimplifyResult> {
+ Ok(ExprSimplifyResult::Original(
+ args.into_iter().map(|_| lit("9")).collect(),
+ ))
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
Result<ColumnarValue> {
+ self.inner.invoke_with_args(args)
+ }
+ }
+
+ #[test]
+ fn simplify_case_pushdown_tolerates_contract_breaching_simplify() {
+ // A simplify() that hands back rewritten args under Original must not
+ // abort planning: the pushdown re-locates its target on what came
+ // back and degrades to a no-op rewrite when none is found.
+ let udf = Arc::new(ScalarUDF::new_from_impl(ContractBreachingUdf {
+ inner: ParseIntUdf::new(Volatility::Immutable),
+ }));
+ let expr = Expr::ScalarFunction(ScalarFunction::new_udf(
+ Arc::clone(&udf),
+ vec![case_on_c2("1", Some("2"))],
+ ));
+ // The rewritten args stand, with no pushdown and no planning error.
+ assert_eq!(
+ simplify(expr),
+ Expr::ScalarFunction(ScalarFunction::new_udf(udf, vec![lit("9")]))
+ );
+ }
+
+ #[test]
+ fn simplify_case_pushdown_folds_literal_branches() {
+ // f(CASE WHEN c2 THEN '1' ELSE '2' END)
+ // --> CASE WHEN c2 THEN f('1') ELSE f('2') END
+ // --> CASE WHEN c2 THEN 1 ELSE 2 END
+ assert_eq!(
+ simplify(parse_int(vec![case_on_c2("1", Some("2"))])),
+ Expr::Case(Case::new(
+ None,
+ vec![(Box::new(col("c2_non_null")), Box::new(lit(1i64)))],
+ Some(Box::new(lit(2i64))),
+ ))
+ );
+ }
+
+ #[test]
+ fn simplify_case_pushdown_bails_out_when_a_branch_fails_to_fold() {
+ // A failing fold must leave the expression untouched: the physical
+ // CASE evaluates branch expressions even on empty batches.
+ let expr = parse_int(vec![case_on_c2("3", Some("garbage"))]);
+ assert_eq!(simplify(expr.clone()), expr);
+ }
+
+ #[test]
+ fn simplify_case_pushdown_requires_homogeneous_literal_types() {
+ // Mixed-type branch literals would change what type `f` observes.
+ let case = Expr::Case(Case::new(
+ None,
+ vec![(Box::new(col("c2_non_null")), Box::new(lit("1")))],
+ Some(Box::new(lit(2i64))),
+ ));
+ let expr = parse_int(vec![case]);
+ assert_eq!(simplify(expr.clone()), expr);
+ }
+
+ #[test]
+ fn simplify_case_pushdown_types_untyped_null_branches() {
+ // Untyped NULL THEN and the missing ELSE fold as f(typed NULL) = -1.
+ let case = Expr::Case(Case::new(
+ None,
+ vec![
+ (
+ Box::new(col("c2_non_null")),
+ Box::new(lit(ScalarValue::Null)),
+ ),
+ (Box::new(col("c2")), Box::new(lit("5"))),
+ ],
+ None,
+ ));
+ assert_eq!(
+ simplify(parse_int(vec![case])),
+ Expr::Case(Case::new(
+ None,
+ vec![
+ (Box::new(col("c2_non_null")), Box::new(lit(-1i64))),
+ (Box::new(col("c2")), Box::new(lit(5i64))),
+ ],
+ Some(Box::new(lit(-1i64))),
+ ))
+ );
+ }
+
+ #[test]
+ fn simplify_case_pushdown_skips_literals_with_metadata() {
+ // Literals carrying field metadata could change what `f` observes.
+ let metadata = FieldMetadata::from(std::collections::BTreeMap::from([(
+ "k".to_string(),
+ "v".to_string(),
+ )]));
+ let case = Expr::Case(Case::new(
+ None,
+ vec![(
+ Box::new(col("c2_non_null")),
+ Box::new(Expr::Literal(ScalarValue::from("1"),
Some(metadata))),
+ )],
+ Some(Box::new(lit("2"))),
+ ));
+ let expr = parse_int(vec![case]);
+ assert_eq!(simplify(expr.clone()), expr);
+ }
+
+ #[test]
+ fn simplify_case_pushdown_materializes_implicit_null_else() {
+ // No ELSE: the implicit NULL branch becomes f(NULL) = -1.
+ assert_eq!(
+ simplify(parse_int(vec![case_on_c2("5", None)])),
+ Expr::Case(Case::new(
+ None,
+ vec![(Box::new(col("c2_non_null")), Box::new(lit(5i64)))],
+ Some(Box::new(lit(-1i64))),
+ ))
+ );
+ }
+
+ #[test]
+ fn simplify_case_pushdown_skips_volatile_functions() {
+ let volatile = Expr::ScalarFunction(ScalarFunction::new_udf(
+ Arc::new(ScalarUDF::new_from_impl(ParseIntUdf::new(
+ Volatility::Volatile,
+ ))),
+ vec![case_on_c2("1", Some("2"))],
+ ));
+ assert_eq!(simplify(volatile.clone()), volatile);
+ }
+
+ #[test]
+ fn simplify_case_pushdown_multi_argument_guards() {
+ // Extra literal argument: still pushed and folded.
+ assert_eq!(
+ simplify(parse_int(vec![case_on_c2("1", Some("2")), lit("9")])),
+ Expr::Case(Case::new(
+ None,
+ vec![(Box::new(col("c2_non_null")), Box::new(lit(1i64)))],
+ Some(Box::new(lit(2i64))),
+ ))
+ );
+ // Non-literal sibling argument: not rewritten.
+ let with_column = parse_int(vec![case_on_c2("1", Some("2")),
col("c1")]);
+ assert_eq!(simplify(with_column.clone()), with_column);
+ // Two CASE arguments: not rewritten.
+ let with_two_cases =
+ parse_int(vec![case_on_c2("1", Some("2")), case_on_c2("3",
Some("4"))]);
+ assert_eq!(simplify(with_two_cases.clone()), with_two_cases);
+ }
+
+ #[test]
+ fn simplify_case_pushdown_leaves_non_literal_branches() {
+ // A branch referencing a column is not a literal output: no rewrite.
+ let case = Expr::Case(Case::new(
+ None,
+ vec![(Box::new(col("c2_non_null")), Box::new(col("c1")))],
+ Some(Box::new(lit("2"))),
+ ));
+ let expr = parse_int(vec![case]);
+ assert_eq!(simplify(expr.clone()), expr);
+ }
+
+ #[test]
+ fn simplify_case_pushdown_after_inner_simplification() {
+ // WHEN-true collapses the CASE first, then the function folds
outright.
+ let case = Expr::Case(Case::new(
+ None,
+ vec![(Box::new(lit(true)), Box::new(lit("7")))],
+ Some(Box::new(lit("8"))),
+ ));
+ assert_eq!(simplify(parse_int(vec![case])), lit(7i64));
+ }
+
+ #[test]
+ fn simplify_case_pushdown_interacts_with_eq_rule() {
+ // The folded CASE feeds the Eq/NotEq pushdown within the same pass.
+ let expr = binary_expr(
+ parse_int(vec![case_on_c2("1", Some("2"))]),
+ Operator::Eq,
+ lit(1i64),
+ );
+ assert_eq!(simplify(expr), col("c2_non_null"));
+ }
+
+ #[test]
+ fn simplify_case_pushdown_does_not_apply_to_cast() {
+ // Cast is outside the rule: folding CAST('garbage') fails fast at
plan time.
+ let expr = Expr::Cast(Cast::new(
+ Box::new(case_on_c2("1", Some("garbage"))),
+ DataType::Int64,
+ ));
+ let simplified = simplify(expr);
+ assert_contains!(format!("{simplified}"), "CASE");
+ assert_contains!(format!("{simplified}"), "garbage");
+ }
+
+ /// Counts how many input rows the UDF is invoked over.
+ #[derive(Debug)]
+ struct RowCountingUdf {
+ signature: Signature,
+ rows_seen: Arc<std::sync::atomic::AtomicUsize>,
+ }
+
+ impl PartialEq for RowCountingUdf {
+ fn eq(&self, other: &Self) -> bool {
+ self.signature == other.signature
+ }
+ }
+
+ impl Eq for RowCountingUdf {}
+
+ impl Hash for RowCountingUdf {
+ fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+ self.signature.hash(state);
+ }
+ }
+
+ impl ScalarUDFImpl for RowCountingUdf {
+ fn name(&self) -> &str {
+ "row_counting_udf"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
Result<ColumnarValue> {
+ use std::sync::atomic::Ordering;
+ match &args.args[0] {
+ ColumnarValue::Scalar(_) => {
+ self.rows_seen.fetch_add(args.number_rows,
Ordering::SeqCst);
+ Ok(ColumnarValue::Scalar(ScalarValue::Int64(Some(0))))
+ }
+ ColumnarValue::Array(array) => {
+ self.rows_seen.fetch_add(array.len(), Ordering::SeqCst);
+ Ok(ColumnarValue::Array(Arc::new(
+ arrow::array::Int64Array::from(vec![0i64;
array.len()]),
+ )))
+ }
+ }
+ }
+ }
+
+ #[test]
+ fn simplify_case_pushdown_eliminates_per_row_invocations() {
Review Comment:
Could we also add an end-to-end optimizer or SLT regression for the
motivating `to_timestamp(CASE ...)` query? The custom UDF tests cover the
rewrite mechanism well, but an SQL-level plan assertion would additionally
cover built-in function registration, coercion, and the specific use case from
the issue. This one is non-blocking.
--
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]
