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 cdcfe5eabfcf0..8d190ddcd4bf8 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 @@ -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 oldCalls = agg.getAggCallList(); + Aggregate lifted = liftListValuesOperandsToVarchar(agg); + List oldCalls = lifted.getAggCallList(); List 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}. 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 castMap = collectListValuesScalarOperands(agg); + if (castMap.isEmpty()) { + return agg; + } + RelNode lifted = buildLiftingProject(agg, castMap); + List 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 collectListValuesScalarOperands(Aggregate agg) { + List origFields = agg.getInput().getRowType().getFieldList(); + int origFieldCount = origFields.size(); + Map 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 castMap) { + RelDataTypeFactory typeFactory = agg.getCluster().getTypeFactory(); + RexBuilder rexBuilder = agg.getCluster().getRexBuilder(); + RelDataType varcharNullable = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.VARCHAR), true); + List origFields = agg.getInput().getRowType().getFieldList(); + List projects = new ArrayList<>(origFields.size() + castMap.size()); + List 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 rewireListValuesCalls(Aggregate agg, RelNode lifted, Map castMap) { + List 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 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" -> { diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CoordinatorReduceIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CoordinatorReduceIT.java index f32fad59a990b..22d8a41a432ae 100644 --- a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CoordinatorReduceIT.java +++ b/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 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 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 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 expected = new java.util.HashSet<>(); for (int i = 1; i <= totalDocs; i++) { @@ -494,7 +494,7 @@ public void testValuesSingleShard() throws Exception { java.util.Set 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 expected = new java.util.HashSet<>(); for (int i = 1; i <= 5; i++) { @@ -532,7 +532,7 @@ public void testValuesAcrossShards() throws Exception { java.util.Set 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 expected = new java.util.HashSet<>(); for (int i = 1; i <= 10; i++) { diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ListAggregateMultiTypeIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ListAggregateMultiTypeIT.java new file mode 100644 index 0000000000000..7f38c93ba6151 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ListAggregateMultiTypeIT.java @@ -0,0 +1,184 @@ +/* + * 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. + */ + +package org.opensearch.analytics.qa; + +import org.opensearch.client.Request; + +import java.util.List; +import java.util.Map; + +/** + * Verifies {@code stats list()} round-trips and renders correctly for every supported + * element type — boolean, byte, short, integer, long, float, double, keyword, text, date, + * date_nanos, ip, binary. One method per type asserts the returned array contains the indexed + * value in its canonical form (IP as dotted-quad, binary as Base64, dates as the configured + * format). + */ +public class ListAggregateMultiTypeIT extends AnalyticsRestTestCase { + + private static final String INDEX = "list_agg_multi_type"; + + public void testListBoolean() throws Exception { + provision(); + assertSingleElement("boolean_value", "true"); + } + + public void testListByte() throws Exception { + provision(); + assertSingleElement("byte_value", "4"); + } + + public void testListShort() throws Exception { + provision(); + assertSingleElement("short_value", "3"); + } + + public void testListInteger() throws Exception { + provision(); + assertSingleElement("integer_value", "2"); + } + + public void testListLong() throws Exception { + provision(); + assertSingleElement("long_value", "1"); + } + + public void testListFloat() throws Exception { + provision(); + assertSingleElement("float_value", "6.2"); + } + + public void testListDouble() throws Exception { + provision(); + assertSingleElement("double_value", "5.1"); + } + + public void testListKeyword() throws Exception { + provision(); + assertSingleElement("keyword_value", "keyword"); + } + + public void testListDate() throws Exception { + provision(); + // DataFusion's CAST(TIMESTAMP AS VARCHAR) emits ISO-8601 'T' between date and time. + assertSingleElement("date_value", "2020-10-13T13:00:00"); + } + + public void testListDateNanos() throws Exception { + provision(); + // DataFusion's CAST(TIMESTAMP_NS AS VARCHAR) emits ISO-8601 'T' between date and time. + assertSingleElement("date_nanos_value", "2019-03-24T01:34:46.123456789"); + } + + public void testListText() throws Exception { + provision(); + assertSingleElement("text_value", "text"); + } + + public void testListIp() throws Exception { + provision(); + assertSingleElement("ip_value", "127.0.0.1"); + } + + public void testListBinary() throws Exception { + provision(); + assertSingleElement("binary_value", "U29tZSBiaW5hcnkgYmxvYg=="); + } + + private void assertSingleElement(String field, Object expected) throws Exception { + List listed = runListQuery(field); + assertEquals("list(" + field + ") must return exactly 1 element", 1, listed.size()); + assertEquals("list(" + field + ")[0]", expected, listed.get(0)); + } + + @SuppressWarnings("unchecked") + private List runListQuery(String field) throws Exception { + Map result = executePpl("source = " + INDEX + " | stats list(" + field + ") as l"); + + List columns = extractColumnNames(result); + assertNotNull("schema must not be null", columns); + assertTrue("columns must contain 'l', got " + columns, columns.contains("l")); + + List> rows = (List>) result.get("datarows"); + assertNotNull("rows must not be null", rows); + assertEquals("scalar agg must return exactly 1 row", 1, rows.size()); + + Object cell = rows.get(0).get(columns.indexOf("l")); + assertNotNull("cell for 'l' must not be null", cell); + assertTrue("list() must return a List, got " + cell.getClass(), cell instanceof List); + return (List) cell; + } + + private void provision() throws Exception { + try { + client().performRequest(new Request("DELETE", "/" + INDEX)); + } catch (Exception ignored) {} + + String mapping = "{" + + "\"settings\": {" + + " \"number_of_shards\": 1," + + " \"number_of_replicas\": 0," + + " \"index.pluggable.dataformat.enabled\": true," + + " \"index.pluggable.dataformat\": \"composite\"," + + " \"index.composite.primary_data_format\": \"parquet\"," + + " \"index.composite.secondary_data_formats\": [\"lucene\"]" + + "}," + + "\"mappings\": {" + + " \"properties\": {" + + " \"boolean_value\": { \"type\": \"boolean\" }," + + " \"byte_value\": { \"type\": \"byte\" }," + + " \"short_value\": { \"type\": \"short\" }," + + " \"integer_value\": { \"type\": \"integer\" }," + + " \"long_value\": { \"type\": \"long\" }," + + " \"float_value\": { \"type\": \"float\" }," + + " \"double_value\": { \"type\": \"double\" }," + + " \"keyword_value\": { \"type\": \"keyword\" }," + + " \"text_value\": { \"type\": \"text\" }," + + " \"binary_value\": { \"type\": \"binary\" }," + + " \"date_value\": { \"type\": \"date\", \"format\": \"yyyy-MM-dd HH:mm:ss\" }," + + " \"date_nanos_value\": { \"type\": \"date_nanos\" }," + + " \"ip_value\": { \"type\": \"ip\" }" + + " }" + + "}" + + "}"; + + Request create = new Request("PUT", "/" + INDEX); + create.setJsonEntity(mapping); + Map response = assertOkAndParse(client().performRequest(create), "create " + INDEX); + assertEquals("index creation must be acknowledged", true, response.get("acknowledged")); + + Request health = new Request("GET", "/_cluster/health/" + INDEX); + health.addParameter("wait_for_status", "green"); + health.addParameter("timeout", "30s"); + client().performRequest(health); + + String bulk = "{\"index\":{\"_id\":\"1\"}}\n" + + "{" + + "\"boolean_value\":true," + + "\"byte_value\":4," + + "\"short_value\":3," + + "\"integer_value\":2," + + "\"long_value\":1," + + "\"float_value\":6.2," + + "\"double_value\":5.1," + + "\"keyword_value\":\"keyword\"," + + "\"text_value\":\"text\"," + + "\"binary_value\":\"U29tZSBiaW5hcnkgYmxvYg==\"," + + "\"date_value\":\"2020-10-13 13:00:00\"," + + "\"date_nanos_value\":\"2019-03-23T21:34:46.123456789-04:00\"," + + "\"ip_value\":\"127.0.0.1\"" + + "}\n"; + + Request bulkRequest = new Request("POST", "/" + INDEX + "/_bulk"); + bulkRequest.setJsonEntity(bulk); + bulkRequest.addParameter("refresh", "true"); + client().performRequest(bulkRequest); + client().performRequest(new Request("POST", "/" + INDEX + "/_flush?force=true")); + } +}