Fix LIST/VALUES aggregate type-mismatch

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

@@ -19,14 +19,19 @@
 import org.apache.calcite.rel.logical.LogicalProject;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.rel.type.RelDataTypeField;
 import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.sql.SqlAggFunction;
 import org.apache.calcite.sql.type.SqlTypeName;
+import org.opensearch.analytics.schema.BinaryType;
+import org.opensearch.analytics.schema.IpType;
 
 import java.util.ArrayList;
+import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 /**
@@ -50,6 +55,10 @@ final class PplAggregateCallRewriter {
         DataFusionFragmentConvertor.LOCAL_INTERNAL_PATTERN_OP
     );
 
+    private static final String LIST = "LIST";
+    private static final String VALUES = "VALUES";
+    private static final String STR_SUFFIX = "$str";
+
     private PplAggregateCallRewriter() {}
 
     static RelNode rewrite(RelNode root) {
@@ -69,22 +78,139 @@ public RelNode visit(RelNode other) {
     }
 
     private static RelNode rewriteAggregate(Aggregate agg) {
-        List<AggregateCall> oldCalls = agg.getAggCallList();
+        Aggregate lifted = liftListValuesOperandsToVarchar(agg);
+        List<AggregateCall> oldCalls = lifted.getAggCallList();
         List<AggregateCall> newCalls = new ArrayList<>(oldCalls.size());
-        boolean changed = false;
+        boolean changed = lifted != agg;
         for (AggregateCall call : oldCalls) {
-            AggregateCall rewritten = rewriteCall(agg, call);
-            if (rewritten == call) {
-                newCalls.add(call);
-            } else {
-                newCalls.add(rewritten);
-                changed = true;
-            }
+            AggregateCall rewritten = rewriteCall(lifted, call);
+            newCalls.add(rewritten);
+            changed |= rewritten != call;
         }
         if (!changed) {
             return agg;
         }
-        return agg.copy(agg.getTraitSet(), agg.getInput(), agg.getGroupSet(), agg.getGroupSets(), newCalls);
+        return lifted.copy(lifted.getTraitSet(), lifted.getInput(), lifted.getGroupSet(), lifted.getGroupSets(), newCalls);
+    }
+
+    /**
+     * Lifts each scalar operand of a {@code LIST}/{@code VALUES} call into a VARCHAR column
+     * via a {@link LogicalProject} inserted above {@code agg.getInput()}, so the aggregator
+     * produces an {@code ARRAY<VARCHAR>}. Skips ARRAY operands (the partial→final merge path)
+     * and operands that are already VARCHAR. Other aggregate calls and group keys are
+     * untouched.
+     */
+    private static Aggregate liftListValuesOperandsToVarchar(Aggregate agg) {
+        Map<Integer, Integer> castMap = collectListValuesScalarOperands(agg);
+        if (castMap.isEmpty()) {
+            return agg;
+        }
+        RelNode lifted = buildLiftingProject(agg, castMap);
+        List<AggregateCall> rewired = rewireListValuesCalls(agg, lifted, castMap);
+        return agg.copy(agg.getTraitSet(), lifted, agg.getGroupSet(), agg.getGroupSets(), rewired);
+    }
+
+    /**
+     * Returns a map from original-input-column-index to new-projected-column-index for every
+     * scalar (non-VARCHAR, non-ARRAY) operand of a LIST/VALUES call. Insertion order preserved
+     * so cast slots end up contiguous in the lifted Project.
+     */
+    private static Map<Integer, Integer> collectListValuesScalarOperands(Aggregate agg) {
+        List<RelDataTypeField> origFields = agg.getInput().getRowType().getFieldList();
+        int origFieldCount = origFields.size();
+        Map<Integer, Integer> castMap = new LinkedHashMap<>();
+        for (AggregateCall call : agg.getAggCallList()) {
+            if (!isListOrValuesCall(call) || call.getArgList().isEmpty()) {
+                continue;
+            }
+            int argIdx = call.getArgList().get(0);
+            RelDataType argType = origFields.get(argIdx).getType();
+            if (argType.getComponentType() != null || argType.getSqlTypeName() == SqlTypeName.VARCHAR) {
+                continue;
+            }
+            castMap.putIfAbsent(argIdx, origFieldCount + castMap.size());
+        }
+        return castMap;
+    }
+
+    /**
+     * Builds the lifting Project: passes through every original column, then appends one
+     * VARCHAR column per entry in {@code castMap}. {@link IpType} routes to
+     * {@code ip_to_string} and {@link BinaryType} to {@code binary_to_base64} directly —
+     * {@link IpBinaryCastFunctionAdapter} would normally rewrite a CAST, but that pass has
+     * already run by the time this rewriter fires. Everything else uses a plain CAST.
+     */
+    private static RelNode buildLiftingProject(Aggregate agg, Map<Integer, Integer> castMap) {
+        RelDataTypeFactory typeFactory = agg.getCluster().getTypeFactory();
+        RexBuilder rexBuilder = agg.getCluster().getRexBuilder();
+        RelDataType varcharNullable = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.VARCHAR), true);
+        List<RelDataTypeField> origFields = agg.getInput().getRowType().getFieldList();
+        List<RexNode> projects = new ArrayList<>(origFields.size() + castMap.size());
+        List<String> names = new ArrayList<>(origFields.size() + castMap.size());
+        for (RelDataTypeField field : origFields) {
+            projects.add(rexBuilder.makeInputRef(field.getType(), field.getIndex()));
+            names.add(field.getName());
+        }
+        for (int srcIdx : castMap.keySet()) {
+            RelDataType srcType = origFields.get(srcIdx).getType();
+            RexNode srcRef = rexBuilder.makeInputRef(srcType, srcIdx);
+            projects.add(toVarchar(srcType, srcRef, varcharNullable, rexBuilder));
+            names.add(origFields.get(srcIdx).getName() + STR_SUFFIX);
+        }
+        return LogicalProject.create(agg.getInput(), List.of(), projects, names, Set.of());
+    }
+
+    private static RexNode toVarchar(RelDataType srcType, RexNode srcRef, RelDataType varcharNullable, RexBuilder rexBuilder) {
+        if (srcType instanceof IpType) {
+            return rexBuilder.makeCall(varcharNullable, IpBinaryCastFunctionAdapter.IP_TO_STRING_OP, List.of(srcRef));
+        }
+        if (srcType instanceof BinaryType) {
+            return rexBuilder.makeCall(varcharNullable, IpBinaryCastFunctionAdapter.BINARY_TO_BASE64_OP, List.of(srcRef));
+        }
+        return rexBuilder.makeCast(varcharNullable, srcRef);
+    }
+
+    /**
+     * Rewires each LIST/VALUES call whose operand was lifted to point at its new VARCHAR
+     * column. Other aggregate calls keep their original column indices because the lifting
+     * Project preserves the original prefix. The explicitReturnType is left null so the
+     * downstream LIST/VALUES dispatch in {@link #rewriteCall} recomputes it for the rewired
+     * input.
+     */
+    private static List<AggregateCall> rewireListValuesCalls(Aggregate agg, RelNode lifted, Map<Integer, Integer> castMap) {
+        List<AggregateCall> rewired = new ArrayList<>(agg.getAggCallList().size());
+        for (AggregateCall call : agg.getAggCallList()) {
+            Integer newIdx = isListOrValuesCall(call) && !call.getArgList().isEmpty() ? castMap.get(call.getArgList().get(0)) : null;
+            if (newIdx == null) {
+                rewired.add(call);
+                continue;
+            }
+            List<Integer> newArgList = new ArrayList<>(call.getArgList());
+            newArgList.set(0, newIdx);
+            rewired.add(
+                AggregateCall.create(
+                    call.getAggregation(),
+                    call.isDistinct(),
+                    call.isApproximate(),
+                    call.ignoreNulls(),
+                    call.rexList,
+                    newArgList,
+                    call.filterArg,
+                    call.distinctKeys,
+                    call.collation,
+                    agg.getGroupCount(),
+                    lifted,
+                    null,
+                    call.getName()
+                )
+            );
+        }
+        return rewired;
+    }
+
+    private static boolean isListOrValuesCall(AggregateCall call) {
+        String name = call.getAggregation().getName();
+        return LIST.equalsIgnoreCase(name) || VALUES.equalsIgnoreCase(name);
     }
 
     /** Replace any RexLiteral{SymbolFlag} in {@code project}'s projection list with a VARCHAR literal of the symbol's name. */
@@ -176,7 +302,11 @@ private static AggregateCall rewriteCall(Aggregate agg, AggregateCall call) {
                 } else {
                     targetOp = DataFusionFragmentConvertor.LOCAL_ARRAY_AGG_OP;
                     targetDistinct = isValues;
-                    explicitReturnType = agg.getCluster().getTypeFactory().createArrayType(arg0Type, -1);
+                    // Match LOCAL_ARRAY_AGG_OP's nullable ARRAY inference; the 2-arg
+                    // createArrayType overload defaults to NOT NULL and trips Calcite's
+                    // typeMatchesInferred check.
+                    RelDataType arrayType = agg.getCluster().getTypeFactory().createArrayType(arg0Type, -1);
+                    explicitReturnType = agg.getCluster().getTypeFactory().createTypeWithNullability(arrayType, true);
                 }
             }
             case "PATTERN" -> {

sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CoordinatorReduceIT.java

@@ -417,7 +417,7 @@ public void testListSingleShard() throws Exception {
         java.util.Set<Integer> seen = new java.util.HashSet<>();
         for (Object v : listed) {
             assertNotNull("list(value) elements must not be null", v);
-            seen.add(((Number) v).intValue());
+            seen.add(Integer.parseInt((String) v));
         }
         java.util.Set<Integer> expected = new java.util.HashSet<>();
         for (int i = 1; i <= DOCS_PER_SHARD; i++) {
@@ -456,7 +456,7 @@ public void testListAcrossShards() throws Exception {
         java.util.Set<Integer> seen = new java.util.HashSet<>();
         for (Object v : listed) {
             assertNotNull("list(value) elements must not be null", v);
-            seen.add(((Number) v).intValue());
+            seen.add(Integer.parseInt((String) v));
         }
         java.util.Set<Integer> expected = new java.util.HashSet<>();
         for (int i = 1; i <= totalDocs; i++) {
@@ -494,7 +494,7 @@ public void testValuesSingleShard() throws Exception {
         java.util.Set<Integer> seen = new java.util.HashSet<>();
         for (Object v : got) {
             assertNotNull("values(value) elements must not be null", v);
-            seen.add(((Number) v).intValue());
+            seen.add(Integer.parseInt((String) v));
         }
         java.util.Set<Integer> expected = new java.util.HashSet<>();
         for (int i = 1; i <= 5; i++) {
@@ -532,7 +532,7 @@ public void testValuesAcrossShards() throws Exception {
         java.util.Set<Integer> seen = new java.util.HashSet<>();
         for (Object v : got) {
             assertNotNull("values(value) elements must not be null", v);
-            seen.add(((Number) v).intValue());
+            seen.add(Integer.parseInt((String) v));
         }
         java.util.Set<Integer> expected = new java.util.HashSet<>();
         for (int i = 1; i <= 10; i++) {