From f7b202d89adef4f69b46f7e54f93be78d16c0cd1 Mon Sep 17 00:00:00 2001 From: Varun Bansal Date: Fri, 18 Sep 2026 21:06:49 +0000 Subject: [PATCH] fix(analytics): Fix multi-shard GROUP BY on multi_value keys Implicit GROUP BY on a multi_value keyword field failed on any index with more than one shard, in two layers on the coordinator reduce stage. First, attachFragmentOnTop round-trips the shard (PARTIAL) fragment through the stock substrait-java ProtoPlanConverter, which maps the opensearch://analytics/multi_value_expand/v1 ExtensionSingle to an EmptyDetail with an empty record type. The grouping key that the expansion appended was then out of range, failing plan assembly with "Field reference offset (N) must be less than number of fields in struct (0)". MultiValueExpandDetail gains a fromProto inverse and derives its record type from the input (append or replace the LIST column with its nullable element), and decodePlan uses a ProtoPlanConverter whose detailFromExtensionSingleRel recognizes the extension type URL. Second, the FINAL aggregate's StageInputTableScan still carried the pre-expansion LIST type while the shards stream the already-expanded scalar key, so DataFusion rejected the reduce ReadRel with "Field 'tags' in Substrait schema has a different type (List(Utf8)) than the corresponding field in the table schema (Utf8View)". OpenSearchAggregate now tags FINAL-mode aggregates with a RelHint when stripping annotations, and MultiValueRelRewriter retypes LIST group keys on a FINAL aggregate over a stage input to their element type instead of re-expanding them. Unit tests cover the decode round trip and the FINAL retype path; the three MultiValueAggregationIT cases pinned on the issue are un-skipped. Resolves #23057 Signed-off-by: Varun Bansal --- .../DataFusionFragmentConvertor.java | 111 ++++++++++++++- .../be/datafusion/MultiValueRelRewriter.java | 70 ++++++++-- .../DataFusionFragmentConvertorTests.java | 126 ++++++++++++++++++ .../planner/rel/OpenSearchAggregate.java | 17 ++- .../analytics/qa/MultiValueAggregationIT.java | 12 +- 5 files changed, 312 insertions(+), 24 deletions(-) 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 75cb99ba73d89..189e5ea3b7d4d 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 @@ -897,25 +897,98 @@ private static MultiValueExpandSpec explicitMultiValueExpand(Correlate correlate ); } + /** + * Wire-level description of a multi-value expand. {@code outputType} is the Calcite-derived + * record type when the spec is built from a {@link RelNode}; it is {@code null} when the spec + * was decoded from serialized bytes, in which case the record type is derived from the input. + */ private record MultiValueExpandSpec(int fieldIndex, Integer limit, boolean append, boolean distinct, Type.Struct outputType) { } - private static final class MultiValueExpandDetail implements Extension.SingleRelDetail { - private static final String TYPE_URL = "opensearch://analytics/multi_value_expand/v1"; + /** + * {@link ExtensionSingle} detail for the multi-value expand, mirrored by + * {@code substrait_consumer.rs} on the Rust side. The 16-byte payload carries + * {@code fieldIndex, limit, append, distinct} as big-endian i32s; the output schema is not on + * the wire and is re-derived from the input on decode ({@link #deriveRecordType(Rel)}). + */ + static final class MultiValueExpandDetail implements Extension.SingleRelDetail { + static final String TYPE_URL = "opensearch://analytics/multi_value_expand/v1"; + private static final int PAYLOAD_BYTES = 16; private final MultiValueExpandSpec spec; private MultiValueExpandDetail(MultiValueExpandSpec spec) { this.spec = spec; } + /** Whether {@code any} carries a multi-value expand payload this class can decode. */ + static boolean matches(Any any) { + return any != null && TYPE_URL.equals(any.getTypeUrl()); + } + + /** Inverse of {@link #toProto}; the record type is derived lazily from the input. */ + static MultiValueExpandDetail fromProto(Any any) { + if (!matches(any)) { + throw new IllegalArgumentException("Not a multi-value expand extension: " + (any == null ? null : any.getTypeUrl())); + } + ByteString value = any.getValue(); + if (value.size() != PAYLOAD_BYTES) { + throw new IllegalArgumentException( + "Malformed multi-value expand payload: expected " + PAYLOAD_BYTES + " bytes, got " + value.size() + ); + } + ByteBuffer payload = value.asReadOnlyByteBuffer(); + int fieldIndex = payload.getInt(); + int limit = payload.getInt(); + int append = payload.getInt(); + int distinct = payload.getInt(); + if (fieldIndex < 0 || limit < -1 || (append != 0 && append != 1) || (distinct != 0 && distinct != 1)) { + throw new IllegalArgumentException( + "Malformed multi-value expand payload: fieldIndex=" + + fieldIndex + + " limit=" + + limit + + " append=" + + append + + " distinct=" + + distinct + ); + } + return new MultiValueExpandDetail( + new MultiValueExpandSpec(fieldIndex, limit < 0 ? null : limit, append == 1, distinct == 1, null) + ); + } + @Override public Type.Struct deriveRecordType(Rel input) { - return spec.outputType(); + if (spec.outputType() != null) { + return spec.outputType(); + } + // Decoded from the wire: rebuild the same shape substrait_consumer.rs produces. The + // expanded column is the LIST's element type, nullable (unnest preserves nulls); in + // append mode it is added after the input columns, otherwise it replaces the LIST column. + List inputFields = input.getRecordType().fields(); + if (spec.fieldIndex() >= inputFields.size()) { + throw new IllegalArgumentException( + "multi-value expand field index " + spec.fieldIndex() + " is outside " + inputFields.size() + " input columns" + ); + } + Type source = inputFields.get(spec.fieldIndex()); + if (!(source instanceof Type.ListType list)) { + throw new IllegalArgumentException("multi-value expand field " + spec.fieldIndex() + " is not a LIST: " + source); + } + Type element = list.elementType().withNullable(true); + List outputFields = new ArrayList<>(inputFields); + if (spec.append()) { + outputFields.add(element); + } else { + outputFields.set(spec.fieldIndex(), element); + } + return Type.Struct.builder().nullable(input.getRecordType().nullable()).addAllFields(outputFields).build(); } @Override public Any toProto(io.substrait.relation.RelProtoConverter converter) { - ByteBuffer payload = ByteBuffer.allocate(16); + ByteBuffer payload = ByteBuffer.allocate(PAYLOAD_BYTES); payload.putInt(spec.fieldIndex()); payload.putInt(spec.limit() == null ? -1 : spec.limit()); payload.putInt(spec.append() ? 1 : 0); @@ -924,6 +997,32 @@ public Any toProto(io.substrait.relation.RelProtoConverter converter) { } } + /** + * {@link ProtoPlanConverter} that understands this backend's own {@link ExtensionSingle} + * details. The stock converter maps every unknown extension to an {@code EmptyDetail} whose + * record type is an empty struct, so any field reference above it (e.g. the GROUP BY key of a + * PARTIAL aggregate over the expanded column) fails with + * "Field reference offset (N) must be less than number of fields in struct (0)" when a + * shard-stage plan is decoded on the coordinator to splice the reduce fragment on top. + */ + private static final class OpenSearchProtoPlanConverter extends ProtoPlanConverter { + OpenSearchProtoPlanConverter(SimpleExtension.ExtensionCollection extensions) { + super(extensions); + } + + @Override + protected io.substrait.relation.ProtoRelConverter getProtoRelConverter(io.substrait.extension.ExtensionLookup functionLookup) { + return new io.substrait.relation.ProtoRelConverter(functionLookup, extensionCollection, protoExtensionConverter) { + @Override + protected Extension.SingleRelDetail detailFromExtensionSingleRel(Any any) { + return MultiValueExpandDetail.matches(any) + ? MultiValueExpandDetail.fromProto(any) + : super.detailFromExtensionSingleRel(any); + } + }; + } + } + /** * Adds an {@code is_not_null} {@code preMeasureFilter} to each measure whose {@link LocalAggOp} * declares {@link LocalAggOp#filtersNullArgs} — so the converter stays generic and only the op @@ -1102,11 +1201,11 @@ public boolean filtersNullArgs(AggregateCall call) { // ── Plan serde helpers ────────────────────────────────────────────────────── - /** Decodes serialized Substrait bytes into a model-level {@link Plan}. */ + /** Decodes serialized Substrait bytes into a model-level {@link Plan}, including this backend's own extension rels. */ private Plan decodePlan(byte[] bytes) { try { io.substrait.proto.Plan proto = io.substrait.proto.Plan.parseFrom(bytes); - return new ProtoPlanConverter(extensions).from(proto); + return new OpenSearchProtoPlanConverter(extensions).from(proto); } catch (InvalidProtocolBufferException e) { throw new IllegalArgumentException("Failed to decode Substrait plan bytes", e); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MultiValueRelRewriter.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MultiValueRelRewriter.java index 3e3277c62539a..90412bb830555 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MultiValueRelRewriter.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MultiValueRelRewriter.java @@ -12,8 +12,12 @@ import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.core.Aggregate; 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.RexNode; import org.apache.calcite.util.ImmutableBitSet; +import org.opensearch.analytics.planner.rel.OpenSearchAggregate; import java.util.ArrayList; import java.util.HashMap; @@ -25,20 +29,21 @@ * *

This rewriter appends a scalar unnested column per LIST GROUP BY key and remaps * the grouping set indices so that Substrait serialisation sees scalar types throughout. - * It currently runs at Substrait-conversion time (inside + * It runs at Substrait-conversion time (inside * {@code DataFusionFragmentConvertor.preprocessForSubstrait}), which is after * {@code PlannerImpl.decomposeAggregates} has already split the aggregate into - * PARTIAL/FINAL halves. In a multi-shard query the FINAL fragment therefore receives - * a {@code List(Utf8)} GROUP BY key where Calcite expects the element type - * ({@code Utf8View}), causing a type mismatch 500. + * PARTIAL/FINAL halves. The PARTIAL half is expanded here; the FINAL half reads the + * already-expanded rows from its stage input, so its LIST keys are retyped to the element + * type instead of being expanded again (see {@link #retypeExpandedStageInputKeys}). * - *

Known issue: the correct fix is to move this expansion into + *

Known issue: the cleaner fix is to move this expansion into * {@code PlannerImpl.runAllOptimizations} before the * {@code decomposeAggregates} call so that Calcite propagates element types into both - * the PARTIAL and FINAL fragments. That requires relocating - * {@link MultiValueExpandRel} into a module that {@code analytics-engine} can depend - * on (the dependency currently flows the other way: {@code analytics-backend-datafusion} - * extends {@code analytics-engine}). This is tracked for a follow-up PR. + * the PARTIAL and FINAL fragments and the FINAL-side retyping becomes unnecessary. That + * requires relocating {@link MultiValueExpandRel} into a module that {@code analytics-engine} + * can depend on (the dependency currently flows the other way: + * {@code analytics-backend-datafusion} extends {@code analytics-engine}). This is tracked + * for a follow-up PR. */ final class MultiValueRelRewriter { @@ -56,6 +61,9 @@ public RelNode visit(RelNode other) { private static RelNode rewriteAggregate(Aggregate aggregate) { RelNode input = aggregate.getInput(); + if (OpenSearchAggregate.isFinalMode(aggregate) && input instanceof DataFusionFragmentConvertor.StageInputTableScan stageInput) { + return retypeExpandedStageInputKeys(aggregate, stageInput); + } Map expandedGroupFields = new HashMap<>(); for (int fieldIndex : aggregate.getGroupSet()) { if (input.getRowType().getFieldList().get(fieldIndex).getType().getComponentType() != null) { @@ -88,6 +96,50 @@ private static RelNode rewriteAggregate(Aggregate aggregate) { return LogicalProject.create(rewritten, List.of(), projects, aggregate.getRowType().getFieldNames()); } + /** + * FINAL-side counterpart of the expansion above. The PARTIAL fragment already replaced each + * LIST GROUP BY key with its expanded scalar element (see {@link #rewriteAggregate}), so the + * rows arriving on the reduce stage's input partition carry the element type. Calcite, + * however, still types the stage input from the pre-expansion aggregate row type, so the + * FINAL aggregate would (a) expand an already-scalar column a second time and (b) declare + * the partition's Substrait {@code base_schema} as {@code List(Utf8)} where the registered + * stream is {@code Utf8View}, which DataFusion rejects at the ReadRel. Retype the affected + * keys on the stage input to the nullable element type instead; the aggregate's row type + * re-derives from the input, so no expansion or reordering Project is needed. + */ + private static RelNode retypeExpandedStageInputKeys(Aggregate aggregate, DataFusionFragmentConvertor.StageInputTableScan stageInput) { + RelDataTypeFactory typeFactory = aggregate.getCluster().getTypeFactory(); + List fields = stageInput.getRowType().getFieldList(); + RelDataTypeFactory.Builder builder = typeFactory.builder(); + boolean changed = false; + for (int fieldIndex = 0; fieldIndex < fields.size(); fieldIndex++) { + RelDataTypeField field = fields.get(fieldIndex); + RelDataType elementType = field.getType().getComponentType(); + if (elementType != null && aggregate.getGroupSet().get(fieldIndex)) { + builder.add(field.getName(), typeFactory.createTypeWithNullability(elementType, true)); + changed = true; + } else { + builder.add(field.getName(), field.getType()); + } + } + if (!changed) { + return aggregate; + } + RelNode retyped = new DataFusionFragmentConvertor.StageInputTableScan( + stageInput.getCluster(), + stageInput.getTraitSet(), + stageInput.getTable().getQualifiedName().getFirst(), + builder.build() + ); + return aggregate.copy( + aggregate.getTraitSet(), + retyped, + aggregate.getGroupSet(), + aggregate.getGroupSets(), + aggregate.getAggCallList() + ); + } + private static ImmutableBitSet remap(ImmutableBitSet fields, Map replacements) { ImmutableBitSet.Builder builder = ImmutableBitSet.builder(); for (int field : fields) { 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 3e06dd5c70c6c..3e2a8aad7133b 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 @@ -19,6 +19,7 @@ import org.apache.calcite.rel.core.CorrelationId; import org.apache.calcite.rel.core.JoinRelType; import org.apache.calcite.rel.core.Uncollect; +import org.apache.calcite.rel.hint.RelHint; import org.apache.calcite.rel.logical.LogicalAggregate; import org.apache.calcite.rel.logical.LogicalCorrelate; import org.apache.calcite.rel.logical.LogicalFilter; @@ -41,6 +42,7 @@ import org.apache.calcite.sql.type.SqlTypeTransforms; import org.apache.calcite.util.ImmutableBitSet; import org.apache.calcite.util.Optionality; +import org.opensearch.analytics.planner.rel.OpenSearchAggregate; import org.opensearch.analytics.planner.rel.OpenSearchStageInputScan; import org.opensearch.analytics.spi.DelegatedPredicateFunction; import org.opensearch.test.OpenSearchTestCase; @@ -811,6 +813,130 @@ public void testListGroupByAddsDistinctAppendExpansion() throws Exception { assertEquals(List.of("tags", "count"), plan.getRelations(0).getRoot().getNamesList()); } + /** + * Multi-shard LIST GROUP BY: the coordinator decodes the shard-stage bytes (which carry the + * multi-value expand as an {@code ExtensionSingleRel}) to splice the FINAL aggregate on top. + * The stock proto converter turns unknown extensions into an empty-struct detail, so the + * PARTIAL aggregate's GROUP BY field reference fails with + * "Field reference offset (N) must be less than number of fields in struct (0)". The decoder + * must rebuild the extension with its real output schema so the round trip succeeds and the + * extension survives re-serialization unchanged. + */ + public void testAttachFragmentOnTopDecodesMultiValueExpandExtension() throws Exception { + RelNode scan = buildListTableScan("test_index"); + AggregateCall count = AggregateCall.create( + SqlStdOperatorTable.COUNT, + false, + List.of(), + -1, + typeFactory.createSqlType(SqlTypeName.BIGINT), + "count" + ); + LogicalAggregate partial = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(0), null, List.of(count)); + byte[] innerBytes = newConvertor().convertFragment(partial); + assertTrue( + "shard fragment must carry the expand extension", + rootRel(decodeSubstrait(innerBytes)).getProject().getInput().getAggregate().getInput().hasExtensionSingle() + ); + + // FINAL half: SUM(count) grouped by the expanded (now scalar VARCHAR) tags column. + RelDataType element = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.VARCHAR), true); + RelDataType partialRowType = typeFactory.builder() + .add("tags", element) + .add("count", typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true)) + .build(); + RelNode stageInput = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 1, partialRowType, List.of("datafusion"), List.of()); + AggregateCall sumCounts = AggregateCall.create( + SqlStdOperatorTable.SUM, + false, + List.of(1), + -1, + typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true), + "count" + ); + LogicalAggregate finalAgg = LogicalAggregate.create(stageInput, List.of(), ImmutableBitSet.of(0), null, List.of(sumCounts)); + + byte[] combined = newConvertor().attachFragmentOnTop(finalAgg, innerBytes); + + Rel root = rootRel(decodeSubstrait(combined)); + assertTrue("root must be the FINAL aggregate", root.hasAggregate()); + Rel inner = root.getAggregate().getInput(); + assertTrue("FINAL aggregate must sit on the shard fragment's output projection", inner.hasProject()); + Rel partialAgg = inner.getProject().getInput(); + assertTrue(partialAgg.hasAggregate()); + Rel expanded = partialAgg.getAggregate().getInput(); + assertTrue("expand extension must survive decode + re-encode", expanded.hasExtensionSingle()); + assertEquals("opensearch://analytics/multi_value_expand/v1", expanded.getExtensionSingle().getDetail().getTypeUrl()); + java.nio.ByteBuffer payload = expanded.getExtensionSingle().getDetail().getValue().asReadOnlyByteBuffer(); + assertEquals(0, payload.getInt()); + assertEquals(-1, payload.getInt()); + assertEquals(1, payload.getInt()); + assertEquals(1, payload.getInt()); + assertEquals( + "PARTIAL GROUP BY key must still reference the appended expanded column", + 1, + partialAgg.getAggregate() + .getGroupings(0) + .getGroupingExpressions(0) + .getSelection() + .getDirectReference() + .getStructField() + .getField() + ); + } + + /** + * FINAL half of a multi-shard LIST GROUP BY. The stage input is still typed with the + * pre-expansion ARRAY key, but the PARTIAL fragment already emits the scalar element, so the + * FINAL aggregate must not expand again and its Read must declare the key as a string — + * otherwise DataFusion rejects the plan with "Field 'tags' in Substrait schema has a + * different type (List(Utf8)) than the corresponding field in the table schema (Utf8View)". + */ + public void testFinalAggregateRetypesExpandedStageInputKeyInsteadOfExpanding() throws Exception { + RelDataType element = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.VARCHAR), true); + RelDataType list = typeFactory.createTypeWithNullability(typeFactory.createArrayType(element, -1), true); + RelDataType partialRowType = typeFactory.builder() + .add("tags", list) + .add("count", typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true)) + .build(); + RelNode stageInput = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 1, partialRowType, List.of("datafusion"), List.of()); + AggregateCall sumCounts = AggregateCall.create( + SqlStdOperatorTable.SUM, + false, + List.of(1), + -1, + typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true), + "count" + ); + LogicalAggregate finalAgg = LogicalAggregate.create( + stageInput, + List.of(RelHint.builder(OpenSearchAggregate.FINAL_MODE_HINT).build()), + ImmutableBitSet.of(0), + null, + List.of(sumCounts) + ); + + Rel root = rootRel(decodeSubstrait(newConvertor().convertFragment(finalAgg))); + assertTrue("FINAL aggregate must be emitted without an expansion Project", root.hasAggregate()); + Rel input = root.getAggregate().getInput(); + assertTrue("FINAL aggregate must read the stage input directly, not a second expand", input.hasRead()); + io.substrait.proto.Type keyType = input.getRead().getBaseSchema().getStruct().getTypes(0); + assertTrue("stage-input key must be declared with the element type: " + keyType, keyType.hasString()); + assertEquals("input-1", input.getRead().getNamedTable().getNames(0)); + + // Without the FINAL hint the same shape is a PARTIAL over a worker stage input and must still expand. + LogicalAggregate partialOverStageInput = LogicalAggregate.create( + stageInput, + List.of(), + ImmutableBitSet.of(0), + null, + List.of(sumCounts) + ); + Rel partialRoot = rootRel(decodeSubstrait(newConvertor().convertFragment(partialOverStageInput))); + assertTrue(partialRoot.hasProject()); + assertTrue(partialRoot.getProject().getInput().getAggregate().getInput().hasExtensionSingle()); + } + public void testExplicitMvExpandCorrelateEmitsAppendExtensionWithLimit() throws Exception { RelNode left = buildListTableScan("test_index"); CorrelationId correlationId = cluster.createCorrel(); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java index 05cddb93912c0..5235bc3e89e39 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java @@ -18,6 +18,7 @@ import org.apache.calcite.rel.RelWriter; import org.apache.calcite.rel.core.Aggregate; import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.hint.RelHint; import org.apache.calcite.rel.logical.LogicalAggregate; import org.apache.calcite.rel.metadata.RelMetadataQuery; import org.apache.calcite.rex.RexLiteral; @@ -469,9 +470,23 @@ public RelNode copyResolved(String backend, List children, List FINAL_MODE_HINT.equals(hint.hintName)); + } + @Override public RelNode stripAnnotations(List strippedChildren) { // Annotations live out-of-band; the aggCall list passes through unchanged. - return LogicalAggregate.create(strippedChildren.getFirst(), List.of(), getGroupSet(), getGroupSets(), getAggCallList()); + List hints = mode == AggregateMode.FINAL ? List.of(RelHint.builder(FINAL_MODE_HINT).build()) : List.of(); + return LogicalAggregate.create(strippedChildren.getFirst(), hints, getGroupSet(), getGroupSets(), getAggCallList()); } } diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/MultiValueAggregationIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/MultiValueAggregationIT.java index d8cbc52ba4c8d..7469dbea62001 100644 --- a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/MultiValueAggregationIT.java +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/MultiValueAggregationIT.java @@ -28,13 +28,12 @@ * three rows. An absent field yields a single {@code null} row; an empty array yields none. * * - *

Multi-shard grouping by a LIST key, {@code list()}/{@code values()} over a LIST input, and - * {@code dc()} over a LIST are known gaps and are pinned with {@link AwaitsFix} so the suite - * documents them until the linked issues are fixed. + *

{@code list()}/{@code values()} over a LIST input and {@code dc()} over a LIST are known + * gaps and are pinned with {@link AwaitsFix} so the suite documents them until the linked + * issues are fixed. */ public class MultiValueAggregationIT extends MultiValueRestTestCase { - private static final String MULTI_SHARD_GROUP_BY_ISSUE = "https://github.com/opensearch-project/OpenSearch/issues/23057"; private static final String LIST_VALUES_OVER_LIST_ISSUE = "https://github.com/opensearch-project/OpenSearch/issues/23058"; private static final String DC_OVER_LIST_ISSUE = "https://github.com/opensearch-project/OpenSearch/issues/23059"; @@ -131,9 +130,8 @@ public void testMvexpandThenStatsByOneShardKeepsDuplicates() throws Exception { ); } - // ---- known gaps: multi-shard GROUP BY on a LIST key -------------------------------------- + // ---- multi-shard GROUP BY on a LIST key (coordinator reduce stage) ---------------------- - @AwaitsFix(bugUrl = MULTI_SHARD_GROUP_BY_ISSUE) public void testCountByListFieldTwoShards() throws Exception { Map groups = groups( executePpl("source = " + TWO_SHARD_INDEX + " | stats count() as cnt by tags"), @@ -143,7 +141,6 @@ public void testCountByListFieldTwoShards() throws Exception { assertEquals(EXPECTED_TAG_COUNTS, nonNullCounts(groups)); } - @AwaitsFix(bugUrl = MULTI_SHARD_GROUP_BY_ISSUE) public void testSqlGroupByListAndScalarTwoShards() throws Exception { Map groups = groups( executeSql("SELECT tags, region, SUM(latency) AS total FROM " + TWO_SHARD_INDEX + " GROUP BY tags, region"), @@ -157,7 +154,6 @@ public void testSqlGroupByListAndScalarTwoShards() throws Exception { ); } - @AwaitsFix(bugUrl = MULTI_SHARD_GROUP_BY_ISSUE) public void testOneShardAndTwoShardAggregatesAgree() throws Exception { List queries = List.of( " | stats count() as v by tags",