Use DataFusion CoalesceFunc::simplify instead of duplicating CASE rewrite.

Add assign.rs integration tests for coalesce via extend (OPL + KQL).
Adjust fixtures so MultiJoin row alignment matches explicit-null attrs.

Signed-off-by: Parameshwaran Krishnasamy <Parameshwaran.K@ibm.com>

Author: Paramesh324

rust/otap-dataflow/crates/query-engine/src/pipeline/assign.rs

@@ -5979,6 +5979,86 @@ mod test {
         test_update_attr_to_upper_case_function_call::<KqlParser>().await
     }
 
+    /// `extend` with `coalesce` across two attribute keys (see #2905).
+    ///
+    /// The second log uses explicit null `attr1` so the `attr1` attribute batch still has a row
+    /// for that log (MultiJoin follows the first attribute argument's row set).
+    ///
+    /// The third log uses explicit null `attr1` and `attr2` so the join still has rows before
+    /// the literal default `"foo"`.
+    async fn test_update_attr_coalesce_function_call<P: Parser>() {
+        let logs_data = to_logs_data(vec![
+            LogRecord::build()
+                .attributes(vec![
+                    KeyValue::new("attr1", AnyValue::new_string("X")),
+                    KeyValue::new("attr2", AnyValue::new_string("Y")),
+                ])
+                .finish(),
+            LogRecord::build()
+                .attributes(vec![
+                    KeyValue::new("attr1", AnyValue { value: None }),
+                    KeyValue::new("attr2", AnyValue::new_string("Z")),
+                ])
+                .finish(),
+            LogRecord::build()
+                .attributes(vec![
+                    KeyValue::new("attr1", AnyValue { value: None }),
+                    KeyValue::new("attr2", AnyValue { value: None }),
+                ])
+                .finish(),
+        ]);
+
+        let query = r#"logs | extend attributes["attr3"] = coalesce(attributes["attr1"], attributes["attr2"], "foo")"#;
+        let pipeline_expr = P::parse_with_options(query, default_parser_options())
+            .unwrap()
+            .pipeline;
+        let mut pipeline = Pipeline::new(pipeline_expr);
+
+        let input = otlp_to_otap(&OtlpProtoMessage::Logs(logs_data));
+        let result = pipeline.execute(input).await.unwrap();
+        let OtlpProtoMessage::Logs(result_logs_data) = otap_to_otlp(&result) else {
+            panic!("invalid signal type");
+        };
+        let log_0 = &result_logs_data.resource_logs[0].scope_logs[0].log_records[0];
+        assert_eq!(
+            log_0.attributes,
+            vec![
+                KeyValue::new("attr1", AnyValue::new_string("X")),
+                KeyValue::new("attr2", AnyValue::new_string("Y")),
+                KeyValue::new("attr3", AnyValue::new_string("X")),
+            ]
+        );
+        let log_1 = &result_logs_data.resource_logs[0].scope_logs[0].log_records[1];
+        assert_eq!(
+            log_1.attributes,
+            vec![
+                KeyValue::new("attr1", AnyValue { value: None }),
+                KeyValue::new("attr2", AnyValue::new_string("Z")),
+                KeyValue::new("attr3", AnyValue::new_string("Z")),
+            ]
+        );
+
+        let log_2 = &result_logs_data.resource_logs[0].scope_logs[0].log_records[2];
+        assert_eq!(
+            log_2.attributes,
+            vec![
+                KeyValue::new("attr1", AnyValue { value: None }),
+                KeyValue::new("attr2", AnyValue { value: None }),
+                KeyValue::new("attr3", AnyValue::new_string("foo")),
+            ]
+        );
+    }
+
+    #[tokio::test]
+    async fn test_update_attr_coalesce_function_call_opl_parser() {
+        test_update_attr_coalesce_function_call::<OplParser>().await
+    }
+
+    #[tokio::test]
+    async fn test_update_attr_coalesce_function_call_kql_parser() {
+        test_update_attr_coalesce_function_call::<KqlParser>().await
+    }
+
     async fn test_update_attr_to_lower_case_function_call<P: Parser>() {
         let logs_data = to_logs_data(vec![
             LogRecord::build()

rust/otap-dataflow/crates/query-engine/src/pipeline/expr.rs

@@ -58,6 +58,7 @@ use data_engine_expressions::{
     StaticScalarExpression, StringScalarExpression, StringValue, TextScalarExpression,
 };
 use datafusion::common::DFSchema;
+use datafusion::functions::core::coalesce::CoalesceFunc;
 use datafusion::functions::core::expr_ext::FieldAccessor;
 use datafusion::functions::crypto::{md5, sha256, sha512};
 use datafusion::functions::datetime::to_char;
@@ -66,8 +67,9 @@ use datafusion::functions::encoding::encode;
 use datafusion::functions::math::log10;
 use datafusion::functions::string::{concat, concat_ws, lower, ltrim, replace, rtrim, upper, uuid};
 use datafusion::logical_expr::expr::ScalarFunction;
+use datafusion::logical_expr::simplify::{ExprSimplifyResult, SimplifyContext};
 use datafusion::logical_expr::{
-    BinaryExpr, ColumnarValue, Expr, Operator, ScalarUDF, cast, col, lit, when,
+    BinaryExpr, ColumnarValue, Expr, Operator, ScalarUDF, ScalarUDFImpl, cast, col, lit,
 };
 use datafusion::physical_expr::{PhysicalExprRef, create_physical_expr};
 use datafusion::prelude::SessionContext;
@@ -426,10 +428,21 @@ impl ExprLogicalPlanner {
         let (df_args, source_scope, _requires_dict_downcast) =
             self.plan_function_args(coalesce_expr.get_expressions().iter(), functions)?;
 
-        // DataFusion's `coalesce` UDF is simplified to CASE during logical optimization and does not
-        // support direct physical evaluation. Mirror that rewrite here (same shape as
-        // `CoalesceFunc::simplify` in Apache DataFusion).
-        let case_expr = Self::coalesce_args_to_case_expr(df_args)?;
+        // DataFusion's `coalesce` UDF does not support direct physical evaluation; the optimizer
+        // rewrites it via `CoalesceFunc::simplify`. Reuse that implementation here.
+        let coalesce_func = CoalesceFunc::new();
+        let simplify_result = coalesce_func
+            .simplify(df_args, &SimplifyContext::default())
+            .map_err(Error::from)?;
+        let case_expr = match simplify_result {
+            ExprSimplifyResult::Simplified(expr) => expr,
+            ExprSimplifyResult::Original(_) => {
+                return Err(Error::InvalidPipelineError {
+                    cause: "expected coalesce simplify to produce a single expression".into(),
+                    query_location: None,
+                });
+            }
+        };
 
         Ok(ScopedLogicalExpr {
             logical_expr: case_expr,
@@ -441,37 +454,6 @@ impl ExprLogicalPlanner {
         })
     }
 
-    /// Rewrites `coalesce(e1, …, eN)` to chained `WHEN eK IS NOT NULL THEN eK … ELSE eN`,
-    /// matching DataFusion's `coalesce` simplification.
-    fn coalesce_args_to_case_expr(mut args: Vec<Expr>) -> Result<Expr> {
-        match args.len() {
-            0 => Err(Error::InvalidPipelineError {
-                cause: "coalesce requires at least one argument".into(),
-                query_location: None,
-            }),
-            1 => Ok(args.pop().ok_or_else(|| Error::InvalidPipelineError {
-                cause: "coalesce internal error: empty args after len check".into(),
-                query_location: None,
-            })?),
-            _ => {
-                let last = args.pop().ok_or_else(|| Error::InvalidPipelineError {
-                    cause: "coalesce internal error: missing last argument".into(),
-                    query_location: None,
-                })?;
-                let first = args.first().ok_or_else(|| Error::InvalidPipelineError {
-                    cause: "coalesce internal error: missing first argument".into(),
-                    query_location: None,
-                })?;
-                let first = first.clone();
-                let mut builder = when(first.clone().is_not_null(), first);
-                for a in args.into_iter().skip(1) {
-                    builder = builder.when(a.clone().is_not_null(), a);
-                }
-                builder.otherwise(last).map_err(Error::from)
-            }
-        }
-    }
-
     fn plan_binary_math_expr(
         &self,
         binary_math_expr: &BinaryMathematicalScalarExpression,