This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new d22a30034d3 [perf] Reduce allocations in multi-stage grouping and
result conversion (#19601)
d22a30034d3 is described below
commit d22a30034d385fc60a325db26c77fa17c80ca117
Author: Xiang Fu <[email protected]>
AuthorDate: Mon Sep 21 12:58:39 2026 -0700
[perf] Reduce allocations in multi-stage grouping and result conversion
(#19601)
* Reduce allocations in multi-stage group-by merging
* Address allocation optimization review feedback
---
.../aggregation/groupby/utils/DoubleToIdMap.java | 5 +-
.../aggregation/groupby/utils/FloatToIdMap.java | 5 +-
.../aggregation/groupby/utils/IntToIdMap.java | 5 +-
.../aggregation/groupby/utils/LongToIdMap.java | 5 +-
.../aggregation/groupby/utils/ObjectToIdMap.java | 5 +-
.../groupby/utils/ValueToIdMapTest.java | 85 ++++++++++
.../operator/MultistageGroupByExecutor.java | 56 ++++---
.../runtime/operator/groupby/GroupIdGenerator.java | 3 +-
.../query/runtime/operator/utils/TypeUtils.java | 8 +-
.../operator/MultistageGroupByExecutorTest.java | 173 +++++++++++++++++++++
.../runtime/operator/utils/TypeUtilsTest.java | 87 +++++++++++
11 files changed, 404 insertions(+), 33 deletions(-)
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/DoubleToIdMap.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/DoubleToIdMap.java
index 4ad1a2d17aa..74a3dac7acd 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/DoubleToIdMap.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/DoubleToIdMap.java
@@ -36,9 +36,10 @@ public class DoubleToIdMap implements ValueToIdMap {
@Override
public int put(double value) {
int numValues = _valueToIdMap.size();
- int id = _valueToIdMap.computeIfAbsent(value, k -> numValues);
- if (id == numValues) {
+ int id = _valueToIdMap.putIfAbsent(value, numValues);
+ if (id == INVALID_KEY) {
_idToValueMap.add(value);
+ return numValues;
}
return id;
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/FloatToIdMap.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/FloatToIdMap.java
index ad74990f69b..3050d0015bb 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/FloatToIdMap.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/FloatToIdMap.java
@@ -36,9 +36,10 @@ public class FloatToIdMap implements ValueToIdMap {
@Override
public int put(float value) {
int numValues = _valueToIdMap.size();
- int id = _valueToIdMap.computeIfAbsent(value, k -> numValues);
- if (id == numValues) {
+ int id = _valueToIdMap.putIfAbsent(value, numValues);
+ if (id == INVALID_KEY) {
_idToValueMap.add(value);
+ return numValues;
}
return id;
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/IntToIdMap.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/IntToIdMap.java
index 383012a17e8..9ec5bfbcf2d 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/IntToIdMap.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/IntToIdMap.java
@@ -36,9 +36,10 @@ public class IntToIdMap implements ValueToIdMap {
@Override
public int put(int value) {
int numValues = _valueToIdMap.size();
- int id = _valueToIdMap.computeIfAbsent(value, k -> numValues);
- if (id == numValues) {
+ int id = _valueToIdMap.putIfAbsent(value, numValues);
+ if (id == INVALID_KEY) {
_idToValueMap.add(value);
+ return numValues;
}
return id;
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/LongToIdMap.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/LongToIdMap.java
index 889c848828b..ced3687ad3c 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/LongToIdMap.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/LongToIdMap.java
@@ -36,9 +36,10 @@ public class LongToIdMap implements ValueToIdMap {
@Override
public int put(long value) {
int numValues = _valueToIdMap.size();
- int id = _valueToIdMap.computeIfAbsent(value, k -> numValues);
- if (id == numValues) {
+ int id = _valueToIdMap.putIfAbsent(value, numValues);
+ if (id == INVALID_KEY) {
_idToValueMap.add(value);
+ return numValues;
}
return id;
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/ObjectToIdMap.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/ObjectToIdMap.java
index b05722c94e4..218fc3aeac0 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/ObjectToIdMap.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/ObjectToIdMap.java
@@ -36,9 +36,10 @@ public class ObjectToIdMap implements ValueToIdMap {
@Override
public int put(Object value) {
int numValues = _valueToIdMap.size();
- int id = _valueToIdMap.computeIntIfAbsent(value, k -> numValues);
- if (id == numValues) {
+ int id = _valueToIdMap.putIfAbsent(value, numValues);
+ if (id == INVALID_KEY) {
_idToValueMap.add(value);
+ return numValues;
}
return id;
}
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/utils/ValueToIdMapTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/utils/ValueToIdMapTest.java
new file mode 100644
index 00000000000..73406b4e2e4
--- /dev/null
+++
b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/utils/ValueToIdMapTest.java
@@ -0,0 +1,85 @@
+/**
+ * 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.core.query.aggregation.groupby.utils;
+
+import java.math.BigDecimal;
+import java.nio.ByteBuffer;
+import java.util.function.IntFunction;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.utils.ByteArray;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+
+
+/// Verifies key identity and contiguous IDs across insertion, repetition, and
map growth.
+public class ValueToIdMapTest {
+ @DataProvider(name = "maps")
+ public Object[][] maps() {
+ return new Object[][]{
+ {new IntToIdMap(), new Object[]{Integer.MIN_VALUE, -1, 0, 1,
Integer.MAX_VALUE},
+ (IntFunction<Object>) Integer::valueOf},
+ {new LongToIdMap(), new Object[]{Long.MIN_VALUE, -1L, 0L, 1L,
Long.MAX_VALUE, 9007199254740993L},
+ (IntFunction<Object>) Long::valueOf},
+ // fastutil preserves raw floating-point bits: signed zeros and
different NaN payloads are distinct keys.
+ {new FloatToIdMap(), new Object[]{Float.NEGATIVE_INFINITY,
-Float.MAX_VALUE, -0.0f, 0.0f, Float.MIN_VALUE,
+ Float.MAX_VALUE, Float.POSITIVE_INFINITY, Float.NaN,
Float.intBitsToFloat(0x7fc00001)},
+ (IntFunction<Object>) Float::valueOf},
+ {new DoubleToIdMap(), new Object[]{Double.NEGATIVE_INFINITY,
-Double.MAX_VALUE, -0.0d, 0.0d, Double.MIN_VALUE,
+ Double.MAX_VALUE, Double.POSITIVE_INFINITY, Double.NaN,
Double.longBitsToDouble(0x7ff8000000000001L)},
+ (IntFunction<Object>) Double::valueOf},
+ {ValueToIdMapFactory.get(DataType.STRING), new Object[]{null, "", "a",
"b"},
+ (IntFunction<Object>) String::valueOf},
+ {ValueToIdMapFactory.get(DataType.BYTES),
+ new Object[]{new ByteArray(new byte[0]), new ByteArray(new
byte[]{1})},
+ (IntFunction<Object>) value -> new
ByteArray(ByteBuffer.allocate(Integer.BYTES).putInt(value).array())},
+ {ValueToIdMapFactory.get(DataType.BIG_DECIMAL), new
Object[]{BigDecimal.ZERO, new BigDecimal("1.0"),
+ new BigDecimal("1.00")}, (IntFunction<Object>) BigDecimal::valueOf}
+ };
+ }
+
+ @Test(dataProvider = "maps")
+ public void testKeys(ValueToIdMap map, Object[] values, IntFunction<Object>
valueFactory) {
+ for (int i = 0; i < values.length; i++) {
+ assertEquals(map.getId(values[i]), ValueToIdMap.INVALID_KEY);
+ assertMapping(map, values[i], i);
+ }
+ // These values are disjoint from the initial keys and force several
rehashes.
+ for (int i = 0; i < 4096; i++) {
+ Object value = valueFactory.apply(1000 + i);
+ assertEquals(map.getId(value), ValueToIdMap.INVALID_KEY);
+ assertMapping(map, value, values.length + i);
+ }
+ for (int i = 0; i < values.length; i++) {
+ assertMapping(map, values[i], i);
+ }
+ for (int i = 0; i < 4096; i++) {
+ assertMapping(map, valueFactory.apply(1000 + i), values.length + i);
+ }
+ }
+
+ private static void assertMapping(ValueToIdMap map, Object value, int
expectedId) {
+ assertEquals(map.put(value), expectedId);
+ assertEquals(map.put(value), expectedId);
+ assertEquals(map.getId(value), expectedId);
+ assertEquals(map.get(expectedId), value);
+ assertEquals(map.getId(map.get(expectedId)), expectedId);
+ }
+}
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java
index 02fd6baf9e3..5409aa7c463 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java
@@ -372,7 +372,7 @@ public class MultistageGroupByExecutor {
int[] groupByKeys = generateGroupByKeys(block);
int numRows = groupByKeys.length;
int numFunctions = _aggFunctions.length;
- Object[][] intermediateResults = new Object[numFunctions][numRows];
+ Object[][] intermediateResults = new Object[numFunctions][];
for (int i = 0; i < numFunctions; i++) {
intermediateResults[i] =
AggregateOperator.getIntermediateResults(_aggFunctions[i], block);
}
@@ -446,16 +446,26 @@ public class MultistageGroupByExecutor {
}
private int[] generateGroupByKeys(DataBlock dataBlock) {
- Object[] keys;
- if (_groupKeyIds.length == 1) {
- keys = DataBlockExtractUtils.extractKey(dataBlock, _groupKeyIds[0]);
- } else {
- keys = DataBlockExtractUtils.extractKeys(dataBlock, _groupKeyIds);
- }
- int numRows = keys.length;
+ int numRows = dataBlock.getNumberOfRows();
int[] intKeys = new int[numRows];
- for (int i = 0; i < numRows; i++) {
- intKeys[i] = _groupIdGenerator.getGroupId(keys[i]);
+ int numKeys = _groupKeyIds.length;
+ if (numKeys == 1) {
+ Object[] keys = DataBlockExtractUtils.extractKey(dataBlock,
_groupKeyIds[0]);
+ for (int i = 0; i < numRows; i++) {
+ intKeys[i] = _groupIdGenerator.getGroupId(keys[i]);
+ }
+ } else {
+ Object[][] columns = new Object[numKeys][];
+ for (int i = 0; i < numKeys; i++) {
+ columns[i] = DataBlockExtractUtils.extractKey(dataBlock,
_groupKeyIds[i]);
+ }
+ Object[] key = new Object[numKeys];
+ for (int rowId = 0; rowId < numRows; rowId++) {
+ for (int i = 0; i < numKeys; i++) {
+ key[i] = columns[i][rowId];
+ }
+ intKeys[rowId] = _groupIdGenerator.getGroupId(key);
+ }
}
return intKeys;
}
@@ -492,15 +502,25 @@ public class MultistageGroupByExecutor {
}
private int[] generateGroupByKeys(DataBlock dataBlock, int numMatchedRows,
RoaringBitmap matchedBitmap) {
- Object[] keys;
- if (_groupKeyIds.length == 1) {
- keys = DataBlockExtractUtils.extractKey(dataBlock, _groupKeyIds[0],
numMatchedRows, matchedBitmap);
- } else {
- keys = DataBlockExtractUtils.extractKeys(dataBlock, _groupKeyIds,
numMatchedRows, matchedBitmap);
- }
int[] intKeys = new int[numMatchedRows];
- for (int i = 0; i < numMatchedRows; i++) {
- intKeys[i] = _groupIdGenerator.getGroupId(keys[i]);
+ int numKeys = _groupKeyIds.length;
+ if (numKeys == 1) {
+ Object[] keys = DataBlockExtractUtils.extractKey(dataBlock,
_groupKeyIds[0], numMatchedRows, matchedBitmap);
+ for (int i = 0; i < numMatchedRows; i++) {
+ intKeys[i] = _groupIdGenerator.getGroupId(keys[i]);
+ }
+ } else {
+ Object[][] columns = new Object[numKeys][];
+ for (int i = 0; i < numKeys; i++) {
+ columns[i] = DataBlockExtractUtils.extractKey(dataBlock,
_groupKeyIds[i], numMatchedRows, matchedBitmap);
+ }
+ Object[] key = new Object[numKeys];
+ for (int rowId = 0; rowId < numMatchedRows; rowId++) {
+ for (int i = 0; i < numKeys; i++) {
+ key[i] = columns[i][rowId];
+ }
+ intKeys[rowId] = _groupIdGenerator.getGroupId(key);
+ }
}
return intKeys;
}
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/groupby/GroupIdGenerator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/groupby/GroupIdGenerator.java
index 65eef96b83d..ec3a20d27c6 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/groupby/GroupIdGenerator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/groupby/GroupIdGenerator.java
@@ -28,7 +28,8 @@ public interface GroupIdGenerator {
/// Returns the group id for the given key. When a new key is encountered,
it assigns a new group id to it before
/// reaching the groups limit, or returns [#INVALID_ID] when the limit is
reached.
/// For single key column, the input is a single Object. For multi key
columns, the input is an Object\[\] containing
- /// the values for each key column.
+ /// the values for each key column. The multi-column array is scratch
storage and may be mutated immediately after
+ /// this method returns. Implementations must not retain the array.
///
/// This method is called once per input row, so the implementation should
be as fast as possible and reduce memory
/// allocation.
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/TypeUtils.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/TypeUtils.java
index d95b60e0c22..f7a5ebe9c03 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/TypeUtils.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/utils/TypeUtils.java
@@ -39,13 +39,13 @@ public class TypeUtils {
public static Object convert(Object value, ColumnDataType storedType) {
switch (storedType) {
case INT:
- return ((Number) value).intValue();
+ return value instanceof Integer ? value : ((Number) value).intValue();
case LONG:
- return ((Number) value).longValue();
+ return value instanceof Long ? value : ((Number) value).longValue();
case FLOAT:
- return ((Number) value).floatValue();
+ return value instanceof Float ? value : ((Number) value).floatValue();
case DOUBLE:
- return ((Number) value).doubleValue();
+ return value instanceof Double ? value : ((Number)
value).doubleValue();
case BIG_DECIMAL:
return value instanceof BigDecimal ? value :
BigDecimal.valueOf(((Number) value).doubleValue());
// For AggregationFunctions that return serialized custom object, e.g.
DistinctCountRawHLLAggregationFunction
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutorTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutorTest.java
new file mode 100644
index 00000000000..6a9956e6be1
--- /dev/null
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutorTest.java
@@ -0,0 +1,173 @@
+/**
+ * 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.query.runtime.operator;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
+import
org.apache.pinot.core.query.aggregation.function.CountAggregationFunction;
+import org.apache.pinot.query.planner.plannode.AggregateNode.AggType;
+import org.apache.pinot.query.runtime.blocks.MseBlock;
+import
org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertTrue;
+
+
+/// Exercises serialized composite keys across merging and filtered
aggregation. Each test owns its executor and blocks.
+public class MultistageGroupByExecutorTest {
+ private static final DataSchema INPUT_SCHEMA = new DataSchema(new
String[]{"weight", "count", "tv", "tag"},
+ new ColumnDataType[]{ColumnDataType.DOUBLE, ColumnDataType.LONG,
ColumnDataType.INT, ColumnDataType.STRING});
+
+ @DataProvider
+ public Object[][] mergeModes() {
+ return new Object[][]{{2, false}, {2, true}, {3, false}, {3, true}};
+ }
+
+ @Test(dataProvider = "mergeModes")
+ public void testSerializedKeysAcrossBlocks(int numKeys, boolean
leafReturnFinalResult) {
+ MultistageGroupByExecutor executor = newExecutor(numKeys,
leafReturnFinalResult, 100);
+ MseBlock.Data first = OperatorTestUtil.block(INPUT_SCHEMA,
+ new Object[]{1.5, 2L, 1000, "Aa"},
+ new Object[]{null, 3L, 1000, "BB"},
+ new Object[]{1.5, 5L, null, "Aa"},
+ new Object[]{null, 7L, null, null},
+ new Object[]{0.0, 11L, 1000, "Aa"},
+ new Object[]{-0.0, 13L, 1000, "BB"},
+ new Object[]{Double.NaN, 17L, 1000, "Aa"},
+ new Object[]{Double.POSITIVE_INFINITY, 19L, 1000,
null}).asSerialized();
+ executor.processBlock(first);
+ executor.processBlock(OperatorTestUtil.block(INPUT_SCHEMA).asSerialized());
+ executor.processBlock(OperatorTestUtil.block(INPUT_SCHEMA,
+ new Object[]{null, 23L, null, null},
+ new Object[]{Double.NaN, 29L, 1000, "Aa"},
+ new Object[]{1.5, 31L, 1000, "Aa"}).asSerialized());
+
+ List<Object[]> expected = List.of(
+ new Object[]{1000, 1.5, "Aa", 33L},
+ new Object[]{1000, null, "BB", 3L},
+ new Object[]{null, 1.5, "Aa", 5L},
+ new Object[]{null, null, null, 30L},
+ new Object[]{1000, 0.0, "Aa", 11L},
+ new Object[]{1000, -0.0, "BB", 13L},
+ new Object[]{1000, Double.NaN, "Aa", 46L},
+ new Object[]{1000, Double.POSITIVE_INFINITY, null, 19L});
+ assertEquals(asMap(executor.getResult(100), numKeys),
expectedMap(expected, numKeys));
+ assertEquals(executor.getNumGroups(), 8);
+ }
+
+ @Test(dataProvider = "mergeModes")
+ public void testExistingKeysStillMergeAtGroupLimit(int numKeys, boolean
leafReturnFinalResult) {
+ MultistageGroupByExecutor executor = newExecutor(numKeys,
leafReturnFinalResult, 2);
+ executor.processBlock(OperatorTestUtil.block(INPUT_SCHEMA,
+ new Object[]{1.5, 2L, 1000, "Aa"},
+ new Object[]{null, 3L, null, null},
+ new Object[]{2.5, 100L, 2000, "BB"},
+ new Object[]{1.5, 5L, 1000, "Aa"},
+ new Object[]{null, 7L, null, null}).asSerialized());
+ executor.processBlock(OperatorTestUtil.block(INPUT_SCHEMA,
+ new Object[]{1.5, 100L, null, "Aa"},
+ new Object[]{1.5, 11L, 1000, "Aa"}).asSerialized());
+
+ assertEquals(asMap(executor.getResult(100), numKeys), expectedMap(List.of(
+ new Object[]{1000, 1.5, "Aa", 18L}, new Object[]{null, null, null,
10L}), numKeys));
+ assertTrue(executor.isNumGroupsLimitReached());
+ }
+
+ @DataProvider
+ public Object[][] filteredModes() {
+ return new Object[][]{{2, false}, {2, true}, {3, false}, {3, true}};
+ }
+
+ @Test(dataProvider = "filteredModes")
+ public void testFilteredSerializedKeysAcrossBlocks(int numKeys, boolean
skipEmptyGroups) {
+ DataSchema inputSchema = new DataSchema(new String[]{"weight", "count",
"tv", "tag", "filter"},
+ new ColumnDataType[]{ColumnDataType.DOUBLE, ColumnDataType.LONG,
ColumnDataType.INT, ColumnDataType.STRING,
+ ColumnDataType.BOOLEAN});
+ MultistageGroupByExecutor executor = newExecutor(numKeys, false, 100, 4,
skipEmptyGroups);
+ executor.processBlock(OperatorTestUtil.block(inputSchema,
+ new Object[]{1.5, 2L, 1000, "Aa", 1},
+ new Object[]{9.5, 3L, 9000, "unmatched", 0},
+ new Object[]{null, 5L, null, null, 1},
+ new Object[]{1.5, 7L, 1000, "Aa", 0},
+ new Object[]{2.5, 11L, 2000, "BB", 1}).asSerialized());
+ executor.processBlock(OperatorTestUtil.block(inputSchema).asSerialized());
+ executor.processBlock(OperatorTestUtil.block(inputSchema,
+ new Object[]{9.5, 13L, 9000, "unmatched", 0}).asSerialized());
+ executor.processBlock(OperatorTestUtil.block(inputSchema,
+ new Object[]{null, 17L, null, null, 1},
+ new Object[]{2.5, 19L, 2000, "BB", 0},
+ new Object[]{1.5, 23L, 1000, "Aa", 1}).asSerialized());
+
+ Map<List<Object>, Long> expected = expectedMap(List.of(
+ new Object[]{1000, 1.5, "Aa", 2L},
+ new Object[]{null, null, null, 2L},
+ new Object[]{2000, 2.5, "BB", 1L}), numKeys);
+ if (!skipEmptyGroups) {
+ expected.put(Arrays.asList(Arrays.copyOf(new Object[]{9000, 9.5,
"unmatched"}, numKeys)), 0L);
+ }
+ assertEquals(asMap(executor.getResult(100), numKeys), expected);
+ assertEquals(executor.getNumGroups(), expected.size());
+ }
+
+ private static MultistageGroupByExecutor newExecutor(int numKeys, boolean
leafReturnFinalResult, int groupLimit) {
+ return newExecutor(numKeys, leafReturnFinalResult, groupLimit, -1, false);
+ }
+
+ private static MultistageGroupByExecutor newExecutor(int numKeys, boolean
leafReturnFinalResult, int groupLimit,
+ int filterArgId, boolean skipEmptyGroups) {
+ int[] groupKeys = numKeys == 2 ? new int[]{2, 0} : new int[]{2, 0, 3};
+ DataSchema resultSchema = numKeys == 2
+ ? new DataSchema(new String[]{"tv", "weight", "count"},
+ new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.DOUBLE,
ColumnDataType.LONG})
+ : new DataSchema(new String[]{"tv", "weight", "tag", "count"},
+ new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.DOUBLE,
ColumnDataType.STRING,
+ ColumnDataType.LONG});
+ AggregationFunction<?, ?>[] functions = {
+ new
CountAggregationFunction(List.of(ExpressionContext.forIdentifier("$1")), true)};
+ return new MultistageGroupByExecutor(groupKeys, functions, new
int[]{filterArgId}, filterArgId,
+ filterArgId < 0 ? AggType.FINAL : AggType.DIRECT,
leafReturnFinalResult, resultSchema,
+ Map.of(QueryOptionKey.NUM_GROUPS_LIMIT, Integer.toString(groupLimit),
+ QueryOptionKey.FILTERED_AGGREGATIONS_SKIP_EMPTY_GROUPS,
Boolean.toString(skipEmptyGroups)), null);
+ }
+
+ private static Map<List<Object>, Long> expectedMap(List<Object[]> rows, int
numKeys) {
+ Map<List<Object>, Long> result = new HashMap<>();
+ for (Object[] row : rows) {
+ result.put(Arrays.asList(Arrays.copyOf(row, numKeys)), (Long) row[3]);
+ }
+ return result;
+ }
+
+ private static Map<List<Object>, Long> asMap(List<Object[]> rows, int
numKeys) {
+ Map<List<Object>, Long> result = new HashMap<>();
+ for (Object[] row : rows) {
+ assertNull(result.put(Arrays.asList(Arrays.copyOf(row, numKeys)), (Long)
row[numKeys]), "Duplicate output group");
+ }
+ return result;
+ }
+}
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/utils/TypeUtilsTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/utils/TypeUtilsTest.java
new file mode 100644
index 00000000000..653f08c699a
--- /dev/null
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/utils/TypeUtilsTest.java
@@ -0,0 +1,87 @@
+/**
+ * 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.query.runtime.operator.utils;
+
+import java.math.BigDecimal;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertThrows;
+
+
+/// Verifies identity conversion and numeric coercions at the single-stage to
multi-stage boundary.
+public class TypeUtilsTest {
+ @DataProvider(name = "numericValues")
+ public Object[][] numericValues() {
+ return new Object[][]{
+ {Integer.valueOf(1024)}, {Long.valueOf(9007199254740993L)},
{Float.valueOf(-0.0f)},
+ {Double.valueOf(-0.0d)}, {Float.intBitsToFloat(0x7fc00001)},
+ {Double.longBitsToDouble(0x7ff8000000000001L)},
{Double.POSITIVE_INFINITY},
+ {Double.NEGATIVE_INFINITY}, {Double.MIN_VALUE}, {new
BigDecimal("123456789.125")},
+ {Byte.valueOf((byte) -1)}, {Short.valueOf((short) 1024)}
+ };
+ }
+
+ @Test(dataProvider = "numericValues")
+ public void testNumericConversions(Number value) {
+ Object converted = TypeUtils.convert(value, ColumnDataType.INT);
+ assertEquals(converted, Integer.valueOf(value.intValue()));
+ if (value instanceof Integer) {
+ assertSame(converted, value);
+ }
+ converted = TypeUtils.convert(value, ColumnDataType.LONG);
+ assertEquals(converted, Long.valueOf(value.longValue()));
+ if (value instanceof Long) {
+ assertSame(converted, value);
+ }
+ converted = TypeUtils.convert(value, ColumnDataType.FLOAT);
+ assertEquals(Float.floatToRawIntBits((Float) converted),
Float.floatToRawIntBits(value.floatValue()));
+ if (value instanceof Float) {
+ assertSame(converted, value);
+ }
+ converted = TypeUtils.convert(value, ColumnDataType.DOUBLE);
+ assertEquals(Double.doubleToRawLongBits((Double) converted),
Double.doubleToRawLongBits(value.doubleValue()));
+ if (value instanceof Double) {
+ assertSame(converted, value);
+ }
+ }
+
+ @Test
+ public void testConvertRowWithNullAndMixedTypes() {
+ Object[] row = {Integer.valueOf(1024), Double.valueOf(-0.0d), null,
Long.valueOf(9007199254740993L),
+ Double.valueOf(9.75d), Integer.valueOf(1024)};
+ Object[] original = row.clone();
+ TypeUtils.convertRow(row, new ColumnDataType[]{ColumnDataType.INT,
ColumnDataType.DOUBLE, ColumnDataType.FLOAT,
+ ColumnDataType.LONG, ColumnDataType.INT, ColumnDataType.DOUBLE});
+ assertSame(row[0], original[0]);
+ assertSame(row[1], original[1]);
+ assertNull(row[2]);
+ assertSame(row[3], original[3]);
+ assertEquals(row[4], Integer.valueOf(9));
+ assertEquals(row[5], Double.valueOf(1024d));
+ for (ColumnDataType type : new ColumnDataType[]{ColumnDataType.INT,
ColumnDataType.LONG, ColumnDataType.FLOAT,
+ ColumnDataType.DOUBLE}) {
+ assertThrows(NullPointerException.class, () -> TypeUtils.convert(null,
type));
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
