Preserve element type for mvcombine ARRAY_AGG on the DataFusion route

sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java

@@ -139,6 +139,11 @@ public static AggregateFunction fromSqlKind(SqlKind kind) {
 
     /** Case-insensitive name lookup; throws if not recognized. */
     public static AggregateFunction fromNameOrError(String name) {
+        // mvcombine's ARRAY_AGG uses the LIST decomposition; element type is preserved in
+        // PplAggregateCallRewriter (LOCAL_ARRAY_AGG_TYPED_OP).
+        if (name.equalsIgnoreCase("ARRAY_AGG")) {
+            return LIST;
+        }
         try {
             return valueOf(name.toUpperCase(java.util.Locale.ROOT));
         } catch (IllegalArgumentException e) {

sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/spi/AggregateFunctionTests.java

@@ -182,6 +182,12 @@ public void testPercentileApproxResolvesByName() {
         assertSame(AggregateFunction.PERCENTILE_APPROX, AggregateFunction.fromNameOrError("PERCENTILE_APPROX"));
     }
 
+    // ── ARRAY_AGG (PPL mvcombine) resolves to the LIST family ──
+    public void testArrayAggResolvesToList() {
+        assertSame(LIST, AggregateFunction.fromNameOrError("ARRAY_AGG"));
+        assertSame(LIST, AggregateFunction.fromNameOrError("array_agg"));
+    }
+
     // ── fromSqlKind still works ──
 
     public void testFromSqlKindResolvesExistingEntries() {

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

@@ -315,6 +315,21 @@ private static RexNode castToVarchar(RexNode arg, RexBuilder rexBuilder, RelData
         return rexBuilder.makeCast(varcharNullable, arg);
     }
 
+    /**
+     * mvcombine's ARRAY_AGG: type-preserving array collect. Unlike {@link #LOCAL_ARRAY_AGG_OP}
+     * (VARCHAR-cast for list()/values()), the element type is kept. Drops null elements.
+     */
+    static final SqlAggFunction LOCAL_ARRAY_AGG_TYPED_OP = new LocalAggOp("array_agg", SqlKind.OTHER_FUNCTION, opBinding -> {
+        RelDataTypeFactory tf = opBinding.getTypeFactory();
+        RelDataType elem = tf.createTypeWithNullability(opBinding.getOperandType(0), true);
+        return tf.createTypeWithNullability(tf.createArrayType(elem, -1), true);
+    }, OperandTypes.ANY) {
+        @Override
+        public boolean filtersNullArgs(AggregateCall call) {
+            return true;
+        }
+    };
+
     /** FINAL-side merge for LIST; un-nests per-shard list states. */
     static final SqlAggFunction LOCAL_LIST_MERGE_OP = new SqlAggFunction(
         "list_merge",
@@ -437,6 +452,7 @@ public RexNode normaliseLiteralArg(int argIndex, RexLiteral lit, RexBuilder rexB
         FunctionMappings.s(LOCAL_FIRST_OP, "first_value"),
         FunctionMappings.s(LOCAL_LAST_OP, "last_value"),
         FunctionMappings.s(LOCAL_ARRAY_AGG_OP, "array_agg"),
+        FunctionMappings.s(LOCAL_ARRAY_AGG_TYPED_OP, "array_agg"),
         FunctionMappings.s(LOCAL_LIST_MERGE_OP, "list_merge"),
         FunctionMappings.s(LOCAL_LIST_MERGE_DISTINCT_OP, "list_merge_distinct"),
         FunctionMappings.s(LOCAL_PERCENTILE_APPROX_OP, "approx_percentile_cont"),

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

@@ -181,6 +181,40 @@ private static AggregateCall rewriteCall(Aggregate agg, AggregateCall call) {
                     explicitReturnType = typeFactory.createTypeWithNullability(arrayType, true);
                 }
             }
+            case "ARRAY_AGG" -> {
+                // mvcombine: type-preserving collect; typed op drops nulls, so filterArg=-1 below.
+                if (call.getArgList().isEmpty()) {
+                    return call;
+                }
+                RelDataType arrayAggArg0 = agg.getInput().getRowType().getFieldList().get(call.getArgList().get(0)).getType();
+                RelDataTypeFactory arrayAggTf = agg.getCluster().getTypeFactory();
+                SqlAggFunction arrayAggOp;
+                RelDataType arrayAggType;
+                if (arrayAggArg0.getComponentType() != null) {
+                    arrayAggOp = DataFusionFragmentConvertor.LOCAL_LIST_MERGE_OP;
+                    arrayAggType = arrayAggArg0;
+                } else {
+                    // nullable element: NOT NULL trips Calcite typeMatchesInferred
+                    arrayAggOp = DataFusionFragmentConvertor.LOCAL_ARRAY_AGG_TYPED_OP;
+                    RelDataType arrayAggElem = arrayAggTf.createTypeWithNullability(arrayAggArg0, true);
+                    arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggTf.createArrayType(arrayAggElem, -1), true);
+                }
+                return AggregateCall.create(
+                    arrayAggOp,
+                    false,
+                    call.isApproximate(),
+                    call.ignoreNulls(),
+                    call.rexList,
+                    call.getArgList(),
+                    -1,
+                    call.distinctKeys,
+                    call.collation,
+                    agg.getGroupCount(),
+                    agg.getInput(),
+                    arrayAggType,
+                    call.getName()
+                );
+            }
             case "PATTERN" -> {
                 // PPL declares ARRAY<MAP<VARCHAR, ANY>>; substrait can't carry ANY.
                 targetOp = DataFusionFragmentConvertor.LOCAL_INTERNAL_PATTERN_OP;

sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionFragmentConvertorTests.java

@@ -338,6 +338,86 @@ public void testAttachFragmentOnTop_NoGroupListOverScalar_MeasureArgIsVarcharCas
         );
     }
 
+    /**
+     * ARRAY_AGG (mvcombine) must ride the ORIGINAL field with no VARCHAR cast, so the collected
+     * array keeps the source element type, unlike list()/values() (see the test above).
+     */
+    public void testAttachFragmentOnTop_ArrayAggOverScalar_MeasureArgIsRawFieldNotVarcharCast() throws Exception {
+        DataFusionFragmentConvertor convertor = newConvertor();
+
+        // Inner fragment: Project that outputs a single INTEGER column (the gathered reduce input).
+        OpenSearchStageInputScan innerStage = new OpenSearchStageInputScan(
+            cluster,
+            cluster.traitSet(),
+            0,
+            rowType("c0", "c1", "c2", "c3", "c4"),
+            List.of("datafusion"),
+            List.of()
+        );
+        org.apache.calcite.rel.logical.LogicalProject innerProject = org.apache.calcite.rel.logical.LogicalProject.create(
+            innerStage,
+            List.of(),
+            List.of(rexBuilder.makeInputRef(innerStage, 4)),
+            List.of("picked"),
+            java.util.Set.of()
+        );
+        byte[] innerBytes = convertor.convertFragment(innerProject);
+
+        // Wrapper: ARRAY_AGG(picked) with NO group-by, over a 1-column placeholder (inner output shape).
+        OpenSearchStageInputScan aggLeaf = new OpenSearchStageInputScan(
+            cluster,
+            cluster.traitSet(),
+            -1,
+            innerProject.getRowType(),
+            List.of("datafusion"),
+            List.of()
+        );
+        SqlAggFunction arrayAggOp = new SqlAggFunction(
+            "ARRAY_AGG",
+            null,
+            SqlKind.OTHER_FUNCTION,
+            ReturnTypes.TO_ARRAY.andThen(SqlTypeTransforms.FORCE_NULLABLE),
+            null,
+            OperandTypes.ANY,
+            SqlFunctionCategory.USER_DEFINED_FUNCTION,
+            false,
+            false,
+            Optionality.FORBIDDEN
+        ) {
+        };
+        RelDataType nullableInt = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.INTEGER), true);
+        RelDataType arrayType = typeFactory.createTypeWithNullability(typeFactory.createArrayType(nullableInt, -1), true);
+        AggregateCall arrayAggCall = AggregateCall.create(
+            arrayAggOp,
+            false,
+            false,
+            false,
+            List.of(),
+            List.of(0),
+            -1,
+            null,
+            org.apache.calcite.rel.RelCollations.EMPTY,
+            0,
+            aggLeaf,
+            arrayType,
+            "a"
+        );
+        LogicalAggregate agg = LogicalAggregate.create(aggLeaf, List.of(), ImmutableBitSet.of(), null, List.of(arrayAggCall));
+
+        byte[] bytes = convertor.attachFragmentOnTop(agg, innerBytes);
+        Plan plan = decodeSubstrait(bytes);
+        Rel root = rootRel(plan);
+        assertTrue("root must be an AggregateRel", root.hasAggregate());
+        Expression arg = root.getAggregate().getMeasures(0).getMeasure().getArguments(0).getValue();
+        assertFalse("ARRAY_AGG(scalar) measure arg must preserve type: raw field, NOT a VARCHAR cast", arg.hasCast());
+        assertTrue("ARRAY_AGG(scalar) measure arg must be a direct field selection", arg.hasSelection());
+        assertEquals(
+            "the selection must reference the ORIGINAL input field (index 0)",
+            0,
+            arg.getSelection().getDirectReference().getStructField().getField()
+        );
+    }
+
     /**
      * Attaching a {@link LogicalSort} on top of inner bytes yields
      * {@code SortRel(<inner>)}.