From 86d9b6a834930a51ca311cc37f2d3c66ae2ece4b Mon Sep 17 00:00:00 2001 From: Glen Matsushita Date: Thu, 24 Sep 2026 16:42:51 -0700 Subject: [PATCH 1/2] [hubspot] Fix SSE multi-key GROUP BY corruption in sort-aggregate combine OrderByComparatorFactory.getGroupByExpressionIndexMap keyed group-by expressions by getIdentifier(), which is null for any transform expression (e.g. dateTrunc(...), json_extract_index(...)). With two or more transform-based GROUP BY keys, they all collapsed onto the same null map entry, so every ORDER BY expression resolved to the same column index. This corrupted the SSE sort-aggregate combine (SortedGroupByCombineOperator / SequentialSortedGroupByCombineOperator) when merging per-segment sorted results across segments. Key on the whole ExpressionContext instead, and add regression tests covering multi-transform-key comparison, the group-key generator at the groups limit with null handling enabled, and end-to-end multi-key JSON group-by queries against offline and realtime tables. --- .../query/utils/OrderByComparatorFactory.java | 21 +- .../NoDictionaryGroupKeyGeneratorTest.java | 43 +++ .../utils/OrderByComparatorFactoryTest.java | 38 +++ .../JsonExtractIndexGroupByRealtimeTest.java | 64 +++++ .../custom/JsonExtractIndexGroupByTest.java | 257 ++++++++++++++++++ 5 files changed, 416 insertions(+), 7 deletions(-) create mode 100644 pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByRealtimeTest.java create mode 100644 pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByTest.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/utils/OrderByComparatorFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/query/utils/OrderByComparatorFactory.java index 6c044df5606e..61c9faec10d1 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/utils/OrderByComparatorFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/utils/OrderByComparatorFactory.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.core.query.utils; +import com.google.common.base.Preconditions; import java.util.ArrayList; import java.util.Comparator; import java.util.HashMap; @@ -70,11 +71,15 @@ public static Comparator getRecordKeyComparator(List valueComparator.compare(k1.getValues(), k2.getValues()); } - private static Map getGroupByExpressionIndexMap(List groupByExpressions) { - Map groupByExpressionIndexMap = new HashMap<>(); + private static Map getGroupByExpressionIndexMap( + List groupByExpressions) { + Map groupByExpressionIndexMap = new HashMap<>(); int numGroupByExpressions = groupByExpressions.size(); for (int i = 0; i < numGroupByExpressions; i++) { - groupByExpressionIndexMap.put(groupByExpressions.get(i).getIdentifier(), i); + // NOTE: Key on the whole expression, not on getIdentifier(). getIdentifier() is null for anything that is not a + // plain column reference, so keying on it collapses every transform group-by key onto a single null entry + // and makes all ORDER BY expressions resolve to the same column index. + groupByExpressionIndexMap.put(groupByExpressions.get(i), i); } return groupByExpressionIndexMap; } @@ -93,12 +98,14 @@ private static class OrderByExpressionWithIndex { /// Add an index for each orderby expression with respect to its position in the group keys private static List getGroupKeyOrderByExpressionFromRowOrderByExpressions( List rowOrderByExpressions, List groupByExpressions) { - Map groupByExpressionIndexMap = getGroupByExpressionIndexMap(groupByExpressions); + Map groupByExpressionIndexMap = getGroupByExpressionIndexMap(groupByExpressions); List result = new ArrayList<>(); // get index wrt group key for each order by expression - rowOrderByExpressions.forEach(expr -> - result.add( - new OrderByExpressionWithIndex(expr, groupByExpressionIndexMap.get(expr.getExpression().getIdentifier())))); + rowOrderByExpressions.forEach(expr -> { + Integer index = groupByExpressionIndexMap.get(expr.getExpression()); + Preconditions.checkState(index != null, "ORDER BY expression: %s is not a GROUP BY key", expr.getExpression()); + result.add(new OrderByExpressionWithIndex(expr, index)); + }); return result; } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryGroupKeyGeneratorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryGroupKeyGeneratorTest.java index 8bb0015cdb63..a4c03f4006af 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryGroupKeyGeneratorTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryGroupKeyGeneratorTest.java @@ -22,9 +22,11 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; +import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; import java.util.List; +import java.util.Map; import java.util.Random; import java.util.Set; import org.apache.commons.io.FileUtils; @@ -186,6 +188,47 @@ public void testMultiColumnHybridGroupKeyGenerator() { } } + /// When the group limit is reached, the multi-column generator must still look up each row's own key. The test + /// data contains no null values, so enabling null handling must not change the generated group keys at all. + @Test + public void testMultiColumnGroupKeyGeneratorAtGroupsLimitWithNullHandling() { + int numGroupsLimit = 5; + ExpressionContext[] groupByExpressions = new ExpressionContext[]{ + ExpressionContext.forIdentifier(INT_COLUMN), ExpressionContext.forIdentifier(LONG_COLUMN) + }; + + int[] groupKeysWithoutNullHandling = new int[NUM_RECORDS]; + new NoDictionaryMultiColumnGroupKeyGenerator(_projectOperator, groupByExpressions, numGroupsLimit, false, + null).generateKeysForBlock(_valueBlock, groupKeysWithoutNullHandling); + + int[] groupKeysWithNullHandling = new int[NUM_RECORDS]; + new NoDictionaryMultiColumnGroupKeyGenerator(_projectOperator, groupByExpressions, numGroupsLimit, true, + null).generateKeysForBlock(_valueBlock, groupKeysWithNullHandling); + + assertEquals(groupKeysWithNullHandling, groupKeysWithoutNullHandling); + + // Sanity check the absolute values too, derived from the block's own contents rather than from an assumed row + // layout: the first 'numGroupsLimit' distinct keys encountered get sequential ids, and every key first seen after + // the limit is reached stays invalid forever. + int numDocs = _valueBlock.getNumDocs(); + assertEquals(numDocs, NUM_RECORDS); + int[] intValues = _valueBlock.getBlockValueSet(groupByExpressions[0]).getIntValuesSV(); + long[] longValues = _valueBlock.getBlockValueSet(groupByExpressions[1]).getLongValuesSV(); + Map, Integer> keyToId = new HashMap<>(); + int[] expectedGroupKeys = new int[NUM_RECORDS]; + int nextId = 0; + for (int row = 0; row < numDocs; row++) { + List key = List.of(intValues[row], longValues[row]); + Integer id = keyToId.get(key); + if (id == null) { + id = nextId < numGroupsLimit ? nextId++ : GroupKeyGenerator.INVALID_ID; + keyToId.put(key, id); + } + expectedGroupKeys[row] = id; + } + assertEquals(groupKeysWithNullHandling, expectedGroupKeys); + } + private void testGroupKeyGenerator(int[] groupByColumnIndexes) { int numGroupByColumns = groupByColumnIndexes.length; GroupKeyGenerator groupKeyGenerator; diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/utils/OrderByComparatorFactoryTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/utils/OrderByComparatorFactoryTest.java index d0deaf5bbc5c..0e88f5a5ee01 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/utils/OrderByComparatorFactoryTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/utils/OrderByComparatorFactoryTest.java @@ -20,13 +20,17 @@ package org.apache.pinot.core.query.utils; import java.util.Arrays; +import java.util.Comparator; import java.util.List; import java.util.stream.Collectors; import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.request.context.FunctionContext; import org.apache.pinot.common.request.context.OrderByExpressionContext; +import org.apache.pinot.core.data.table.Record; import org.testng.annotations.Test; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; public class OrderByComparatorFactoryTest { @@ -104,4 +108,38 @@ public void testTwoNullsCompareNextColumn() { assertEquals(extractColumn(_rows, COLUMN2_INDEX), Arrays.asList(1, 2, 3)); } + + private static ExpressionContext function(String name, String arg) { + return ExpressionContext.forFunction( + new FunctionContext(FunctionContext.Type.TRANSFORM, name, List.of(ExpressionContext.forIdentifier(arg)))); + } + + /// getRecordKeyComparator maps each ORDER BY expression to its position in the GROUP BY list. When several + /// group-by keys are transform expressions rather than plain identifiers, every expression must still resolve to + /// its own column, otherwise SortedRecordsMerger sees unequal groups as equal and merges them. + @Test + public void testRecordKeyComparatorWithMultipleTransformGroupByKeys() { + ExpressionContext key0 = function("datetrunc", "tsColumn"); + ExpressionContext key1 = function("jsonextractindex", "jsonColumn"); + List groupByExpressions = List.of(key0, key1); + List orderBys = + List.of(new OrderByExpressionContext(key0, ASC, NULLS_LAST), new OrderByExpressionContext(key1, ASC, + NULLS_LAST)); + + Comparator comparator = + OrderByComparatorFactory.getRecordKeyComparator(orderBys, groupByExpressions, false); + + // Same second key, different first key: these are distinct groups and must not compare equal. + Record a = new Record(new Object[]{1L, "x", 10.0}); + Record b = new Record(new Object[]{2L, "x", 20.0}); + assertTrue(comparator.compare(a, b) < 0, "rows differing only in the first group key compared equal"); + assertTrue(comparator.compare(b, a) > 0); + + // Same first key, different second key: also distinct groups. + Record c = new Record(new Object[]{1L, "y", 30.0}); + assertTrue(comparator.compare(a, c) < 0, "rows differing only in the second group key compared equal"); + + // Identical keys are the same group. + assertEquals(comparator.compare(a, new Record(new Object[]{1L, "x", 99.0})), 0); + } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByRealtimeTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByRealtimeTest.java new file mode 100644 index 000000000000..8449306d5c1b --- /dev/null +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByRealtimeTest.java @@ -0,0 +1,64 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.integration.tests.custom; + +import java.io.File; +import java.util.List; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.Test; + + +/** + * Same as {@link JsonExtractIndexGroupByTest} but against a realtime table, so the JSON index being read is the + * mutable (consuming segment) implementation. The production table in PINOT-489 is realtime. + */ +@Test(suiteName = "CustomClusterIntegrationTest") +public class JsonExtractIndexGroupByRealtimeTest extends JsonExtractIndexGroupByTest { + private static final String TABLE_NAME = "RTJsonExtractIndexGroupByTest"; + + @Override + public String getTableName() { + return TABLE_NAME; + } + + @Override + public boolean isRealtimeTable() { + return true; + } + + @Override + protected int getRealtimeSegmentFlushSize() { + // Keep everything in consuming (mutable) segments so the mutable JSON index is exercised. + return 1_000_000; + } + + @Override + protected TableConfig createRealtimeTableConfig(File sampleAvroFile) { + AvroFileSchemaKafkaAvroMessageDecoder._avroFile = sampleAvroFile; + return new TableConfigBuilder(TableType.REALTIME).setTableName(getTableName()) + .setStreamConfigs(getStreamConfigs()) + .setTimeColumnName(getTimeColumnName()) + .setJsonIndexColumns(List.of(PROPERTIES_FIELD)) + .setNoDictionaryColumns(List.of(PROPERTIES_FIELD)) + .setNumReplicas(getNumReplicas()) + .build(); + } +} diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByTest.java new file mode 100644 index 000000000000..f7665445fa5d --- /dev/null +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByTest.java @@ -0,0 +1,257 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.integration.tests.custom; + +import com.fasterxml.jackson.databind.JsonNode; +import java.io.File; +import java.math.BigDecimal; +import java.time.LocalDate; +import java.time.ZoneOffset; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.avro.file.DataFileWriter; +import org.apache.avro.generic.GenericData; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.JsonUtils; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; + + +/** + * Reproduction attempt for PINOT-489 / pinot-planning#491: + * {@code dateTrunc(json_extract_scalar(...))} combined with a second group-by key computed by + * {@code json_extract_index(...)} is reported to return non-deterministic per-bucket attribution on the + * multistage engine, while the single-key variant and the single-stage engine are correct. + */ +@Test(suiteName = "CustomClusterIntegrationTest") +public class JsonExtractIndexGroupByTest extends CustomDataQueryClusterIntegrationTest { + private static final String DEFAULT_TABLE_NAME = "JsonExtractIndexGroupByTest"; + protected static final String PROPERTIES_FIELD = "properties"; + protected static final String BUCKET_FIELD = "bucket"; + + // Large enough to span several 10k doc-id blocks per segment. + // Several segments per server matters: the SSE sort-aggregate combine only pair-wise merges segment results when a + // server holds more than one segment, and that merge is where PINOT-489 corrupts group keys. + protected static final int NUM_DOCS_PER_SEGMENT = 12_000; + protected static final int NUM_AVRO_FILES = 8; + protected static final int NUM_MONTHS = 12; + protected static final int NUM_BUCKETS = 10; + protected static final int SELECTED_BUCKET = 3; + protected static final String[] MOVEMENT_TYPES = {"new", "expansion", "contraction", "churn", "reactivation"}; + + protected static final String MONTH_KEY = + "dateTrunc('month', json_extract_scalar(properties, '$.hs_revenue_month', 'Long', 0), 'MILLISECONDS', 'UTC')"; + protected static final String TYPE_KEY_INDEX = + "json_extract_index(properties, '$.hs_mrr_movement_type', 'String', 'null')"; + protected static final String TYPE_KEY_SCALAR = + "json_extract_scalar(properties, '$.hs_mrr_movement_type', 'String', 'null')"; + protected static final String SUM_EXPR = + "sumprecision(json_extract_scalar(properties, '$.hs_mrr_in_company_currency', 'Double', 0))"; + + /** month start epoch millis -> movement type -> expected sum, for all rows */ + protected final Map> _expected = new LinkedHashMap<>(); + /** month start epoch millis -> movement type -> expected sum, restricted to bucket = SELECTED_BUCKET */ + protected final Map> _expectedFiltered = new LinkedHashMap<>(); + protected final Map _expectedByMonth = new LinkedHashMap<>(); + + @Override + public int getNumAvroFiles() { + return NUM_AVRO_FILES; + } + + @Override + protected long getCountStarResult() { + long numDocsPerSegment = NUM_DOCS_PER_SEGMENT; + return numDocsPerSegment * NUM_AVRO_FILES; + } + + @Override + public String getTableName() { + return DEFAULT_TABLE_NAME; + } + + @Override + public Schema createSchema() { + return new Schema.SchemaBuilder().setSchemaName(getTableName()) + .addSingleValueDimension(PROPERTIES_FIELD, FieldSpec.DataType.STRING) + .addSingleValueDimension(BUCKET_FIELD, FieldSpec.DataType.INT) + .addDateTime(TIMESTAMP_FIELD_NAME, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + } + + @Override + public TableConfig createOfflineTableConfig() { + return new TableConfigBuilder(TableType.OFFLINE).setTableName(getTableName()) + .setTimeColumnName(TIMESTAMP_FIELD_NAME) + .setJsonIndexColumns(List.of(PROPERTIES_FIELD)) + .setNoDictionaryColumns(List.of(PROPERTIES_FIELD)) + .build(); + } + + @Override + public List createAvroFiles() + throws Exception { + org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false); + avroSchema.setFields(List.of( + new org.apache.avro.Schema.Field(PROPERTIES_FIELD, + org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING), null, null), + new org.apache.avro.Schema.Field(BUCKET_FIELD, + org.apache.avro.Schema.create(org.apache.avro.Schema.Type.INT), null, null), + new org.apache.avro.Schema.Field(TIMESTAMP_FIELD_NAME, + org.apache.avro.Schema.create(org.apache.avro.Schema.Type.LONG), null, null))); + + long[] monthStarts = new long[NUM_MONTHS]; + for (int m = 0; m < NUM_MONTHS; m++) { + monthStarts[m] = LocalDate.of(2024, m + 1, 1).atStartOfDay(ZoneOffset.UTC).toInstant().toEpochMilli(); + } + + try (AvroFilesAndWriters avroFilesAndWriters = createAvroFilesAndWriters(avroSchema)) { + List> writers = avroFilesAndWriters.getWriters(); + for (int fileId = 0; fileId < writers.size(); fileId++) { + DataFileWriter writer = writers.get(fileId); + for (int i = 0; i < NUM_DOCS_PER_SEGMENT; i++) { + // Give every segment a DIFFERENT, partially-overlapping subset of (month, type) pairs. If all segments + // produce the same sorted key sequence, a positional merge lines up by accident and hides key-comparison + // bugs in the sort-aggregate combine, so the subsets must genuinely diverge. + int monthIdx = (fileId + (i / 2) % 6) % NUM_MONTHS; + String type = MOVEMENT_TYPES[(fileId + (i % 2)) % MOVEMENT_TYPES.length]; + int bucket = (i * 7 + fileId) % NUM_BUCKETS; + long monthStart = monthStarts[monthIdx]; + long revenueMonth = monthStart + (i % 27) * 86_400_000L; + BigDecimal amount = BigDecimal.valueOf((i % 100) + 1); + + Map properties = new HashMap<>(); + properties.put("hs_revenue_month", revenueMonth); + properties.put("hs_mrr_movement_type", type); + properties.put("hs_mrr_in_company_currency", amount.doubleValue()); + properties.put("filler", "padding-" + i + "-" + fileId); + + GenericData.Record record = new GenericData.Record(avroSchema); + record.put(PROPERTIES_FIELD, JsonUtils.objectToString(properties)); + record.put(BUCKET_FIELD, bucket); + record.put(TIMESTAMP_FIELD_NAME, revenueMonth); + writer.append(record); + + _expected.computeIfAbsent(monthStart, k -> new LinkedHashMap<>()).merge(type, amount, BigDecimal::add); + _expectedByMonth.merge(monthStart, amount, BigDecimal::add); + if (bucket == SELECTED_BUCKET) { + _expectedFiltered.computeIfAbsent(monthStart, k -> new LinkedHashMap<>()) + .merge(type, amount, BigDecimal::add); + } + } + } + return avroFilesAndWriters.getAvroFiles(); + } + } + + @Test(dataProvider = "useBothQueryEngines") + public void testSingleKeyGroupBy(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + String query = String.format("SELECT %s AS m, %s AS s FROM %s GROUP BY %s ORDER BY %s LIMIT 500", MONTH_KEY, + SUM_EXPR, getTableName(), MONTH_KEY, MONTH_KEY); + for (int run = 0; run < 5; run++) { + Map actual = new LinkedHashMap<>(); + JsonNode rows = postQuery(query).get("resultTable").get("rows"); + for (JsonNode row : rows) { + actual.put(row.get(0).asLong(), new BigDecimal(row.get(1).asText())); + } + assertEquals(actual.keySet(), _expectedByMonth.keySet(), "run " + run + " month keys"); + for (Map.Entry entry : _expectedByMonth.entrySet()) { + assertEquals(actual.get(entry.getKey()).compareTo(entry.getValue()), 0, + "run " + run + " month " + entry.getKey()); + } + } + } + + @Test(dataProvider = "useBothQueryEngines") + public void testTwoKeyGroupByWithJsonExtractIndex(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + assertTwoKeyGroupBy(TYPE_KEY_INDEX, null, _expected); + } + + @Test(dataProvider = "useBothQueryEngines") + public void testTwoKeyGroupByWithJsonExtractScalar(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + assertTwoKeyGroupBy(TYPE_KEY_SCALAR, null, _expected); + } + + /** + * Same as above but with a sparse filter, so the doc-id sets handed to the JSON index reader are a small, + * scattered subset of each block. This is the shape the production query has. + */ + @Test(dataProvider = "useBothQueryEngines") + public void testTwoKeyGroupByWithSparseFilter(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + assertTwoKeyGroupBy(TYPE_KEY_INDEX, BUCKET_FIELD + " = " + SELECTED_BUCKET, _expectedFiltered); + } + + protected String buildQuery(String typeKey, @Nullable String whereClause) { + String where = whereClause == null ? "" : " WHERE " + whereClause; + return String.format("SELECT %s AS m, %s AS t, %s AS s FROM %s%s GROUP BY %s, %s ORDER BY %s, %s LIMIT 500", + MONTH_KEY, typeKey, SUM_EXPR, getTableName(), where, MONTH_KEY, typeKey, MONTH_KEY, typeKey); + } + + protected void assertTwoKeyGroupBy(String typeKey, @Nullable String whereClause, + Map> expected) + throws Exception { + String query = buildQuery(typeKey, whereClause); + List failures = new ArrayList<>(); + for (int run = 0; run < 5; run++) { + Map> actual = new LinkedHashMap<>(); + JsonNode rows = postQuery(query).get("resultTable").get("rows"); + for (JsonNode row : rows) { + actual.computeIfAbsent(row.get(0).asLong(), k -> new LinkedHashMap<>()) + .merge(row.get(1).asText(), new BigDecimal(row.get(2).asText()), BigDecimal::add); + } + if (!actual.keySet().equals(expected.keySet())) { + failures.add("run " + run + ": month keys " + actual.keySet() + " != " + expected.keySet()); + continue; + } + for (Map.Entry> monthEntry : expected.entrySet()) { + Map actualForMonth = actual.get(monthEntry.getKey()); + if (!actualForMonth.keySet().equals(monthEntry.getValue().keySet())) { + failures.add("run " + run + " month " + monthEntry.getKey() + ": types " + actualForMonth.keySet() + " != " + + monthEntry.getValue().keySet()); + continue; + } + for (Map.Entry typeEntry : monthEntry.getValue().entrySet()) { + BigDecimal actualSum = actualForMonth.get(typeEntry.getKey()); + if (actualSum.compareTo(typeEntry.getValue()) != 0) { + failures.add("run " + run + " (" + monthEntry.getKey() + ", " + typeEntry.getKey() + "): " + actualSum + + " != " + typeEntry.getValue()); + } + } + } + } + assertEquals(failures, List.of(), "mismatches for query: " + query + "\n" + String.join("\n", failures)); + } +} From 4f91189369d62ac26ac2eabce3d27ae372acb36b Mon Sep 17 00:00:00 2001 From: Glen Matsushita Date: Thu, 24 Sep 2026 17:04:45 -0700 Subject: [PATCH 2/2] [hubspot] Convert Javadoc blocks to /// doc comments in new integration tests Checkstyle in this repo rejects /** */ Javadoc in favor of /// markdown doc comments (JEP 467); the linter CI job on PR #19663 failed with 5 violations in the two new JsonExtractIndexGroupByTest files for this reason. --- .../JsonExtractIndexGroupByRealtimeTest.java | 6 ++---- .../custom/JsonExtractIndexGroupByTest.java | 20 ++++++++----------- 2 files changed, 10 insertions(+), 16 deletions(-) diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByRealtimeTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByRealtimeTest.java index 8449306d5c1b..1fd8de0192d6 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByRealtimeTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByRealtimeTest.java @@ -26,10 +26,8 @@ import org.testng.annotations.Test; -/** - * Same as {@link JsonExtractIndexGroupByTest} but against a realtime table, so the JSON index being read is the - * mutable (consuming segment) implementation. The production table in PINOT-489 is realtime. - */ +/// Same as [JsonExtractIndexGroupByTest] but against a realtime table, so the JSON index being read is the mutable +/// (consuming segment) implementation. The production table in PINOT-489 is realtime. @Test(suiteName = "CustomClusterIntegrationTest") public class JsonExtractIndexGroupByRealtimeTest extends JsonExtractIndexGroupByTest { private static final String TABLE_NAME = "RTJsonExtractIndexGroupByTest"; diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByTest.java index f7665445fa5d..e4167fc094e8 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByTest.java @@ -42,12 +42,10 @@ import static org.testng.Assert.assertEquals; -/** - * Reproduction attempt for PINOT-489 / pinot-planning#491: - * {@code dateTrunc(json_extract_scalar(...))} combined with a second group-by key computed by - * {@code json_extract_index(...)} is reported to return non-deterministic per-bucket attribution on the - * multistage engine, while the single-key variant and the single-stage engine are correct. - */ +/// Reproduction attempt for PINOT-489 / pinot-planning#491: +/// `dateTrunc(json_extract_scalar(...))` combined with a second group-by key computed by +/// `json_extract_index(...)` is reported to return non-deterministic per-bucket attribution on the +/// multistage engine, while the single-key variant and the single-stage engine are correct. @Test(suiteName = "CustomClusterIntegrationTest") public class JsonExtractIndexGroupByTest extends CustomDataQueryClusterIntegrationTest { private static final String DEFAULT_TABLE_NAME = "JsonExtractIndexGroupByTest"; @@ -73,9 +71,9 @@ public class JsonExtractIndexGroupByTest extends CustomDataQueryClusterIntegrati protected static final String SUM_EXPR = "sumprecision(json_extract_scalar(properties, '$.hs_mrr_in_company_currency', 'Double', 0))"; - /** month start epoch millis -> movement type -> expected sum, for all rows */ + /// month start epoch millis -> movement type -> expected sum, for all rows protected final Map> _expected = new LinkedHashMap<>(); - /** month start epoch millis -> movement type -> expected sum, restricted to bucket = SELECTED_BUCKET */ + /// month start epoch millis -> movement type -> expected sum, restricted to bucket = SELECTED_BUCKET protected final Map> _expectedFiltered = new LinkedHashMap<>(); protected final Map _expectedByMonth = new LinkedHashMap<>(); @@ -203,10 +201,8 @@ public void testTwoKeyGroupByWithJsonExtractScalar(boolean useMultiStageQueryEng assertTwoKeyGroupBy(TYPE_KEY_SCALAR, null, _expected); } - /** - * Same as above but with a sparse filter, so the doc-id sets handed to the JSON index reader are a small, - * scattered subset of each block. This is the shape the production query has. - */ + /// Same as above but with a sparse filter, so the doc-id sets handed to the JSON index reader are a small, + /// scattered subset of each block. This is the shape the production query has. @Test(dataProvider = "useBothQueryEngines") public void testTwoKeyGroupByWithSparseFilter(boolean useMultiStageQueryEngine) throws Exception {