[analytics-engine] Fix CHECKED_LONG_SUM conversion

sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/checked_long_sum.rs

@@ -0,0 +1,159 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+//! Distinctly named analytics binding for SQL's `CHECKED_LONG_SUM`.
+//!
+//! Analytics routing intentionally keeps DataFusion's native SUM semantics. This wrapper delegates
+//! every SUM execution path while retaining the `checked_long_sum` name, preventing collisions when
+//! a distributed intermediate schema contains both SUM and CHECKED_LONG_SUM over the same field.
+
+use std::sync::Arc;
+
+use datafusion::arrow::datatypes::{DataType, FieldRef};
+use datafusion::common::Result;
+use datafusion::execution::context::SessionContext;
+use datafusion::functions_aggregate::sum::sum_udaf;
+use datafusion::logical_expr::expr::AggregateFunction;
+use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion::logical_expr::utils::AggregateOrderSensitivity;
+use datafusion::logical_expr::{
+    Accumulator, AggregateUDF, AggregateUDFImpl, Documentation, Expr, GroupsAccumulator, Operator,
+    ReversedUDAF, SetMonotonicity, Signature,
+};
+
+pub fn register_all(ctx: &SessionContext) {
+    ctx.register_udaf(AggregateUDF::from(CheckedLongSum::new()));
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct CheckedLongSum {
+    native_sum: Arc<AggregateUDF>,
+}
+
+impl CheckedLongSum {
+    fn new() -> Self {
+        Self {
+            native_sum: sum_udaf(),
+        }
+    }
+}
+
+impl AggregateUDFImpl for CheckedLongSum {
+    fn name(&self) -> &str {
+        "checked_long_sum"
+    }
+
+    fn signature(&self) -> &Signature {
+        self.native_sum.inner().signature()
+    }
+
+    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+        self.native_sum.inner().return_type(arg_types)
+    }
+
+    fn accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+        self.native_sum.inner().accumulator(args)
+    }
+
+    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
+        self.native_sum.inner().state_fields(args)
+    }
+
+    fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
+        self.native_sum.inner().groups_accumulator_supported(args)
+    }
+
+    fn create_groups_accumulator(
+        &self,
+        args: AccumulatorArgs,
+    ) -> Result<Box<dyn GroupsAccumulator>> {
+        self.native_sum.inner().create_groups_accumulator(args)
+    }
+
+    fn create_sliding_accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+        self.native_sum.inner().create_sliding_accumulator(args)
+    }
+
+    fn reverse_expr(&self) -> ReversedUDAF {
+        self.native_sum.inner().reverse_expr()
+    }
+
+    fn order_sensitivity(&self) -> AggregateOrderSensitivity {
+        self.native_sum.inner().order_sensitivity()
+    }
+
+    fn documentation(&self) -> Option<&Documentation> {
+        self.native_sum.inner().documentation()
+    }
+
+    fn set_monotonicity(&self, data_type: &DataType) -> SetMonotonicity {
+        self.native_sum.inner().set_monotonicity(data_type)
+    }
+
+    fn simplify_expr_op_literal(
+        &self,
+        aggregate: &AggregateFunction,
+        arg: &Expr,
+        op: Operator,
+        literal: &Expr,
+        arg_is_left: bool,
+    ) -> Result<Option<Expr>> {
+        self.native_sum
+            .inner()
+            .simplify_expr_op_literal(aggregate, arg, op, literal, arg_is_left)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use datafusion::arrow::array::{Array, Int64Array, RecordBatch};
+    use datafusion::arrow::datatypes::{Field, Schema};
+
+    #[tokio::test]
+    async fn native_and_checked_sum_keep_distinct_names() {
+        let ctx = SessionContext::new();
+        register_all(&ctx);
+        let batch = RecordBatch::try_new(
+            Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, false)])),
+            vec![Arc::new(Int64Array::from(vec![1, 2, 3]))],
+        )
+        .unwrap();
+        ctx.register_batch("t", batch).unwrap();
+
+        let batches = ctx
+            .sql("SELECT SUM(x), CHECKED_LONG_SUM(x) FROM t")
+            .await
+            .unwrap()
+            .collect()
+            .await
+            .unwrap();
+        let result = &batches[0];
+
+        assert_eq!(result.schema().field(0).name(), "sum(t.x)");
+        assert_eq!(result.schema().field(1).name(), "checked_long_sum(t.x)");
+        assert_eq!(
+            result
+                .column(0)
+                .as_any()
+                .downcast_ref::<Int64Array>()
+                .unwrap()
+                .value(0),
+            6
+        );
+        assert_eq!(
+            result
+                .column(1)
+                .as_any()
+                .downcast_ref::<Int64Array>()
+                .unwrap()
+                .value(0),
+            6
+        );
+    }
+}

sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/mod.rs

@@ -13,12 +13,14 @@
 use datafusion::execution::context::SessionContext;
 
 pub mod approx_distinct_safe;
+pub mod checked_long_sum;
 pub mod internal_pattern;
 pub mod list_merge;
 pub mod os_count_distinct;
 pub mod take;
 
 pub fn register_all(ctx: &SessionContext) {
+    checked_long_sum::register_all(ctx);
     take::register_all(ctx);
     list_merge::register_all(ctx);
     internal_pattern::register_all(ctx);

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java

@@ -28,6 +28,8 @@
 import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexOver;
+import org.apache.calcite.rex.RexWindow;
 import org.apache.calcite.schema.ColumnStrategy;
 import org.apache.calcite.sql.SqlAggFunction;
 import org.apache.calcite.sql.SqlFunction;
@@ -70,6 +72,7 @@
 import io.substrait.isthmus.TypeConverter;
 import io.substrait.isthmus.expression.AggregateFunctionConverter;
 import io.substrait.isthmus.expression.FunctionMappings;
+import io.substrait.isthmus.expression.RexExpressionConverter;
 import io.substrait.isthmus.expression.ScalarFunctionConverter;
 import io.substrait.isthmus.expression.WindowFunctionConverter;
 import io.substrait.plan.Plan;
@@ -431,6 +434,31 @@ public RexNode normaliseLiteralArg(int argIndex, RexLiteral lit, RexBuilder rexB
     ) {
     };
 
+    /**
+     * Analytics-engine binding for SQL's reflective {@code CHECKED_LONG_SUM}. The runtime
+     * implementation delegates to DataFusion's native SUM but keeps this distinct function name
+     * so plans containing both SUM and CHECKED_LONG_SUM do not produce duplicate Arrow field names.
+     */
+    static final SqlAggFunction LOCAL_CHECKED_LONG_SUM_OP = new SqlAggFunction(
+        "checked_long_sum",
+        null,
+        SqlKind.SUM,
+        ReturnTypes.BIGINT_NULLABLE,
+        null,
+        OperandTypes.NUMERIC,
+        SqlFunctionCategory.USER_DEFINED_FUNCTION,
+        false,
+        false,
+        Optionality.FORBIDDEN
+    ) {
+    };
+
+    static boolean isUnboundCheckedLongSum(SqlAggFunction operator) {
+        return operator != LOCAL_CHECKED_LONG_SUM_OP
+            && operator.getKind() == SqlKind.SUM
+            && "CHECKED_LONG_SUM".equalsIgnoreCase(operator.getName());
+    }
+
     private static final List<FunctionMappings.Sig> ADDITIONAL_AGGREGATE_SIGS = List.of(
         FunctionMappings.s(SqlStdOperatorTable.APPROX_COUNT_DISTINCT, "approx_distinct"),
         FunctionMappings.s(LOCAL_TAKE_OP, "take"),
@@ -441,14 +469,16 @@ public RexNode normaliseLiteralArg(int argIndex, RexLiteral lit, RexBuilder rexB
         FunctionMappings.s(LOCAL_LIST_MERGE_DISTINCT_OP, "list_merge_distinct"),
         FunctionMappings.s(LOCAL_PERCENTILE_APPROX_OP, "approx_percentile_cont"),
         FunctionMappings.s(LOCAL_INTERNAL_PATTERN_OP, "internal_pattern"),
-        FunctionMappings.s(LOCAL_OS_COUNT_DISTINCT_OP, "os_count_distinct")
+        FunctionMappings.s(LOCAL_OS_COUNT_DISTINCT_OP, "os_count_distinct"),
+        FunctionMappings.s(LOCAL_CHECKED_LONG_SUM_OP, "checked_long_sum")
     );
 
     private static final List<FunctionMappings.Sig> ADDITIONAL_WINDOW_SIGS = List.of(
         FunctionMappings.s(LOCAL_INTERNAL_PATTERN_WINDOW_OP, "internal_pattern"),
         // Mirror ADDITIONAL_AGGREGATE_SIGS: rename APPROX_COUNT_DISTINCT to DataFusion's `approx_distinct`.
         FunctionMappings.s(SqlStdOperatorTable.APPROX_COUNT_DISTINCT, "approx_distinct"),
-        FunctionMappings.s(LOCAL_OS_COUNT_DISTINCT_OP, "os_count_distinct")
+        FunctionMappings.s(LOCAL_OS_COUNT_DISTINCT_OP, "os_count_distinct"),
+        FunctionMappings.s(LOCAL_CHECKED_LONG_SUM_OP, "checked_long_sum")
     );
 
     /**
@@ -713,7 +743,8 @@ public Optional<AggregateFunctionInvocation> convert(
                 AggregateCall call,
                 Function<RexNode, Expression> rexConverter
             ) {
-                Optional<AggregateFunctionInvocation> bound = super.convert(input, inputType, call, rexConverter);
+                AggregateCall substraitCall = bindCheckedLongSum(call);
+                Optional<AggregateFunctionInvocation> bound = super.convert(input, inputType, substraitCall, rexConverter);
                 if (bound.isEmpty()) {
                     return bound;
                 }
@@ -746,6 +777,26 @@ public Optional<AggregateFunctionInvocation> convert(
                 if (rewritten == null) return bound;
                 return Optional.of(ImmutableAggregateFunctionInvocation.builder().from(fn).arguments(rewritten).build());
             }
+
+            private AggregateCall bindCheckedLongSum(AggregateCall call) {
+                if (!isUnboundCheckedLongSum(call.getAggregation())) {
+                    return call;
+                }
+                return AggregateCall.create(
+                    call.getParserPosition(),
+                    LOCAL_CHECKED_LONG_SUM_OP,
+                    call.isDistinct(),
+                    call.isApproximate(),
+                    call.ignoreNulls(),
+                    call.rexList,
+                    call.getArgList(),
+                    call.filterArg,
+                    call.distinctKeys,
+                    call.collation,
+                    call.getType(),
+                    call.getName()
+                );
+            }
         };
         // Same APPROX_COUNT_DISTINCT filter as aggConverter — let our `approx_distinct` entry win.
         WindowFunctionConverter windowConverter = new WindowFunctionConverter(
@@ -760,6 +811,37 @@ protected ImmutableList<FunctionMappings.Sig> getSigs() {
                     .filter(sig -> sig.operator != SqlStdOperatorTable.APPROX_COUNT_DISTINCT)
                     .collect(ImmutableList.toImmutableList());
             }
+
+            @Override
+            public Optional<Expression.WindowFunctionInvocation> convert(
+                RexOver call,
+                Function<RexNode, Expression> rexConverter,
+                RexExpressionConverter rexExpressionConverter
+            ) {
+                return super.convert(bindCheckedLongSum(call), rexConverter, rexExpressionConverter);
+            }
+
+            private RexOver bindCheckedLongSum(RexOver call) {
+                if (!isUnboundCheckedLongSum(call.getAggOperator())) {
+                    return call;
+                }
+                RexWindow window = call.getWindow();
+                return (RexOver) new RexBuilder(typeFactory).makeOver(
+                    call.getType(),
+                    LOCAL_CHECKED_LONG_SUM_OP,
+                    call.getOperands(),
+                    window.partitionKeys,
+                    window.orderKeys,
+                    window.getLowerBound(),
+                    window.getUpperBound(),
+                    window.getExclude(),
+                    window.isRows(),
+                    true,
+                    false,
+                    call.isDistinct(),
+                    call.ignoreNulls()
+                );
+            }
         };
         ConverterProvider converterProvider = new ConverterProvider(
             typeFactory,

sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_aggregate_functions.yaml

@@ -2,6 +2,19 @@
 ---
 urn: extension:org.opensearch:aggregate_functions
 aggregate_functions:
+  - name: checked_long_sum
+    description: >-
+      Analytics-engine binding for SQL's CHECKED_LONG_SUM. The DataFusion runtime
+      delegates this function to its native SUM implementation while retaining a
+      distinct output name for intermediate schemas.
+    impls:
+      - args:
+          - name: x
+            value: any
+        nullability: DECLARED_OUTPUT
+        decomposable: MANY
+        intermediate: i64?
+        return: i64?
   - name: approx_distinct
     description: >-
       Approximate distinct count using HyperLogLog. Maps to DataFusion's

sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_window_functions.yaml

@@ -2,6 +2,18 @@
 ---
 urn: extension:org.opensearch:window_functions
 window_functions:
+  - name: checked_long_sum
+    description: >-
+      Window form of the analytics-engine CHECKED_LONG_SUM binding. Execution
+      delegates to DataFusion's native SUM window accumulator.
+    impls:
+      - args:
+          - name: x
+            value: any
+        decomposable: MANY
+        intermediate: i64?
+        return: i64?
+        window_type: STREAMING
   - name: os_count_distinct
     description: >-
       Window form of `os_count_distinct(x)`. Emitted as the substrait window