kosiew commented on code in PR #24490:
URL: https://github.com/apache/datafusion/pull/24490#discussion_r3930651214
##########
datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs:
##########
@@ -4015,6 +4238,533 @@ 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))),
+ ))
+ }
+
+ /// How `FieldPolicyParseIntUdf` shapes its return field.
+ #[derive(Debug, PartialEq, Eq, Hash)]
+ enum FieldPolicy {
+ /// Always tagged with metadata.
+ Tagged,
+ /// Tagged only when the argument is not a scalar.
+ TaggedForArrays,
+ /// Non-nullable, though `f(NULL)` returns NULL.
+ NonNullable,
+ /// Declares Int64 but produces Int32 for anything but `'1'`.
+ Mistyped,
+ }
+
+ /// `ParseIntUdf` whose return field follows `policy`.
+ #[derive(Debug, PartialEq, Eq, Hash)]
+ struct FieldPolicyParseIntUdf {
+ policy: FieldPolicy,
+ }
+
+ impl ScalarUDFImpl for FieldPolicyParseIntUdf {
+ fn name(&self) -> &str {
+ "field_policy_parse_int_udf"
+ }
+
+ fn signature(&self) -> &Signature {
+ static SIGNATURE: LazyLock<Signature> =
+ LazyLock::new(||
Signature::variadic_any(Volatility::Immutable));
+ &SIGNATURE
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn return_field_from_args(&self, args: ReturnFieldArgs) ->
Result<FieldRef> {
+ let tagged = match self.policy {
+ FieldPolicy::Tagged => true,
+ FieldPolicy::TaggedForArrays =>
args.scalar_arguments[0].is_none(),
+ FieldPolicy::NonNullable | FieldPolicy::Mistyped => false,
+ };
+ let nullable = self.policy != FieldPolicy::NonNullable;
+ let field = Field::new(self.name(), DataType::Int64, nullable);
+ Ok(Arc::new(if tagged {
+ field.with_metadata(
+ [("extension".to_string(), "tagged".to_string())].into(),
+ )
+ } else {
+ field
+ }))
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
Result<ColumnarValue> {
+ match args.args.first() {
+ Some(ColumnarValue::Scalar(ScalarValue::Utf8(None)))
+ if self.policy == FieldPolicy::NonNullable =>
+ {
+ Ok(ColumnarValue::Scalar(ScalarValue::Int64(None)))
+ }
+ Some(ColumnarValue::Scalar(ScalarValue::Utf8(Some(s))))
+ if self.policy == FieldPolicy::Mistyped && s != "1" =>
+ {
+ Ok(ColumnarValue::Scalar(ScalarValue::Int32(Some(0))))
+ }
+ _ =>
ParseIntUdf::new(Volatility::Immutable).invoke_with_args(args),
+ }
+ }
+ }
+
+ fn field_policy_parse_int(policy: FieldPolicy, args: Vec<Expr>) -> Expr {
+ Expr::ScalarFunction(ScalarFunction::new_udf(
+ Arc::new(ScalarUDF::new_from_impl(FieldPolicyParseIntUdf { policy
})),
+ args,
+ ))
+ }
+
+ #[test]
+ fn simplify_case_pushdown_declines_metadata_producing_functions() {
+ let expr =
+ field_policy_parse_int(FieldPolicy::Tagged, vec![case_on_c2("1",
Some("2"))]);
+ assert_eq!(simplify(expr.clone()), expr);
+ }
+
+ #[test]
+ fn simplify_case_pushdown_declines_when_original_field_has_metadata() {
+ // Every folded `f(literal)` is untagged, but `f(CASE ...)` is tagged.
+ let expr = field_policy_parse_int(
+ FieldPolicy::TaggedForArrays,
+ vec![case_on_c2("1", Some("2"))],
+ );
+ let field = expr.to_field(expr_test_schema().as_ref()).unwrap().1;
+ assert!(!field.metadata().is_empty());
+ assert_eq!(simplify(expr.clone()), expr);
+ }
+
+ #[test]
+ fn simplify_case_pushdown_declines_nullability_loosening() {
+ // The folded implicit ELSE is a NULL literal: a non-nullable field
+ // must not become nullable.
+ let expr =
+ field_policy_parse_int(FieldPolicy::NonNullable,
vec![case_on_c2("1", None)]);
+ assert_eq!(simplify(expr.clone()), expr);
+
+ // All branches non-null: field stays non-nullable, rewrite applies.
+ assert_eq!(
+ simplify(field_policy_parse_int(
+ FieldPolicy::NonNullable,
+ 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_declines_mistyped_results() {
+ // Only the second branch folds to the wrong type: every branch must
+ // match the declared return type, not just the one that types the
CASE.
+ let expr = field_policy_parse_int(
+ FieldPolicy::Mistyped,
+ 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 branch may
+ // never be taken at runtime, so erroring at plan time would be wrong.
+ 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)))],
Review Comment:
Small optional suggestion: could we add a direct field assertion here for
the `ParseIntUdf(CASE ... THEN '1' ELSE '2')` case to document that the rewrite
intentionally tightens nullability from nullable to non-nullable?
The existing `simplify_case_pushdown_folds_literal_branches` test already
covers the rewrite itself, so I don't think this is needed for correctness. It
would just make the schema rationale explicit and help prevent a future change
from requiring exact nullability equality and accidentally disabling this valid
simplification.
--
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]
