diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java index 1fa5dec3da471..a89c5b8a724a9 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java +++ b/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) { diff --git a/sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/spi/AggregateFunctionTests.java b/sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/spi/AggregateFunctionTests.java index f91dfe5cc00ba..b46172eb1ce44 100644 --- a/sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/spi/AggregateFunctionTests.java +++ b/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() { diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java index 4e265ff0d09b6..e1e61c9be53b5 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java +++ b/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"), diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java index bc71c033508b3..a1a1f876c0687 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java +++ b/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>; substrait can't carry ANY. targetOp = DataFusionFragmentConvertor.LOCAL_INTERNAL_PATTERN_OP; diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionFragmentConvertorTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionFragmentConvertorTests.java index 3c5e4b723822d..6c54f7aedbd73 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionFragmentConvertorTests.java +++ b/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()}.