This is an automated email from the ASF dual-hosted git repository.
Jackie-Jiang 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 4b71d907c54 Reuse the group key flyweight in MultiKeysGroupIdGenerator
(#19834)
4b71d907c54 is described below
commit 4b71d907c5439dc22587a676ad8152ab5acfd754
Author: Deepak kumar <[email protected]>
AuthorDate: Fri Oct 16 17:15:10 2026 -0700
Reuse the group key flyweight in MultiKeysGroupIdGenerator (#19834)
---
.../groupby/MultiKeysGroupIdGenerator.java | 27 +++-
.../groupby/MultiKeysGroupIdGeneratorTest.java | 142 +++++++++++++++++++++
2 files changed, 164 insertions(+), 5 deletions(-)
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/groupby/MultiKeysGroupIdGenerator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/groupby/MultiKeysGroupIdGenerator.java
index 62adb089add..6c9943f1865 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/groupby/MultiKeysGroupIdGenerator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/groupby/MultiKeysGroupIdGenerator.java
@@ -28,6 +28,11 @@ import
org.apache.pinot.core.query.aggregation.groupby.utils.ValueToIdMapFactory
import org.apache.pinot.spi.utils.FixedIntArray;
+/// [GroupIdGenerator] for 3 or more key columns, where each key value is
converted into a key id and the resulting
+/// key ids are wrapped into a [FixedIntArray] used as the map key.
+///
+/// This class is not thread-safe. It keeps a reusable [FixedIntArray] to look
up the group id without allocating on
+/// every row, so a single instance must be accessed by a single thread at a
time.
public class MultiKeysGroupIdGenerator implements GroupIdGenerator {
private final Object2IntOpenHashMap<FixedIntArray> _groupIdMap;
/// A function to generate the next group ID based on the current size of
the map.
@@ -36,6 +41,10 @@ public class MultiKeysGroupIdGenerator implements
GroupIdGenerator {
private final ValueToIdMap[] _keyToIdMaps;
private final int _numGroupsLimit;
+ /// Reusable key holding the key ids of the row being processed. It is
replaced with a new instance whenever it is
+ /// inserted into [#_groupIdMap], which takes ownership of the key.
+ private FixedIntArray _currentKey;
+
public MultiKeysGroupIdGenerator(ColumnDataType[] keyTypes, int
numKeyColumns,
int numGroupsLimit, int initialCapacity) {
_groupIdMap = new Object2IntOpenHashMap<>(initialCapacity);
@@ -47,19 +56,27 @@ public class MultiKeysGroupIdGenerator implements
GroupIdGenerator {
_keyToIdMaps[i] = ValueToIdMapFactory.get(keyTypes[i].toDataType());
}
_numGroupsLimit = numGroupsLimit;
+ _currentKey = new FixedIntArray(new int[numKeyColumns]);
}
@Override
public int getGroupId(Object key) {
Object[] keyValues = (Object[]) key;
- int numKeyColumns = keyValues.length;
- int[] keyIds = new int[numKeyColumns];
- if (_groupIdMap.size() < _numGroupsLimit) {
+ int numKeyColumns = _keyToIdMaps.length;
+ FixedIntArray currentKey = _currentKey;
+ int[] keyIds = currentKey.elements();
+ int numGroups = _groupIdMap.size();
+ if (numGroups < _numGroupsLimit) {
for (int i = 0; i < numKeyColumns; i++) {
Object keyValue = keyValues[i];
keyIds[i] = keyValue != null ? _keyToIdMaps[i].put(keyValue) : NULL_ID;
}
- return _groupIdMap.computeIfAbsent(new FixedIntArray(keyIds),
_groupIdGenerator);
+ int groupId = _groupIdMap.computeIfAbsent(currentKey, _groupIdGenerator);
+ if (_groupIdMap.size() != numGroups) {
+ // A new group is created and the map holds a reference to the current
key, so allocate a new one
+ _currentKey = new FixedIntArray(new int[numKeyColumns]);
+ }
+ return groupId;
} else {
for (int i = 0; i < numKeyColumns; i++) {
Object keyValue = keyValues[i];
@@ -73,7 +90,7 @@ public class MultiKeysGroupIdGenerator implements
GroupIdGenerator {
keyIds[i] = keyId;
}
}
- return _groupIdMap.getInt(new FixedIntArray(keyIds));
+ return _groupIdMap.getInt(currentKey);
}
}
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/groupby/MultiKeysGroupIdGeneratorTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/groupby/MultiKeysGroupIdGeneratorTest.java
new file mode 100644
index 00000000000..d05f636d40e
--- /dev/null
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/groupby/MultiKeysGroupIdGeneratorTest.java
@@ -0,0 +1,142 @@
+/**
+ * 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.groupby;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import
org.apache.pinot.query.runtime.operator.groupby.GroupIdGenerator.GroupKey;
+import org.testng.annotations.Test;
+
+import static
org.apache.pinot.query.runtime.operator.groupby.GroupIdGenerator.INVALID_ID;
+import static org.testng.Assert.assertEquals;
+
+
+public class MultiKeysGroupIdGeneratorTest {
+ private static final ColumnDataType[] KEY_TYPES = {ColumnDataType.STRING,
ColumnDataType.INT, ColumnDataType.LONG};
+ private static final int NUM_KEY_COLUMNS = KEY_TYPES.length;
+ // Number of columns in the output row, which includes the key columns and
one aggregation column
+ private static final int NUM_COLUMNS = NUM_KEY_COLUMNS + 1;
+
+ @Test
+ public void testGroupIds() {
+ GroupIdGenerator generator = new MultiKeysGroupIdGenerator(KEY_TYPES,
NUM_KEY_COLUMNS, 100, 10);
+ // The caller reuses the same array for all the rows
+ Object[] key = new Object[NUM_KEY_COLUMNS];
+
+ assertEquals(getGroupId(generator, key, "a", 1, 10L), 0);
+ assertEquals(getGroupId(generator, key, "b", 2, 20L), 1);
+ // Same keys should map to the same group ids
+ assertEquals(getGroupId(generator, key, "a", 1, 10L), 0);
+ assertEquals(getGroupId(generator, key, "b", 2, 20L), 1);
+ // Keys sharing a prefix with an existing group should get new group ids
+ assertEquals(getGroupId(generator, key, "a", 1, 20L), 2);
+ assertEquals(getGroupId(generator, key, "a", 2, 10L), 3);
+ assertEquals(generator.getNumGroups(), 4);
+
+ assertEquals(getGroupKeys(generator), Map.of(
+ 0, Arrays.asList("a", 1, 10L, null),
+ 1, Arrays.asList("b", 2, 20L, null),
+ 2, Arrays.asList("a", 1, 20L, null),
+ 3, Arrays.asList("a", 2, 10L, null)));
+ }
+
+ @Test
+ public void testNullKeyValues() {
+ GroupIdGenerator generator = new MultiKeysGroupIdGenerator(KEY_TYPES,
NUM_KEY_COLUMNS, 100, 10);
+ Object[] key = new Object[NUM_KEY_COLUMNS];
+
+ assertEquals(getGroupId(generator, key, null, 1, 10L), 0);
+ assertEquals(getGroupId(generator, key, "a", null, 10L), 1);
+ assertEquals(getGroupId(generator, key, null, 1, 10L), 0);
+ assertEquals(getGroupId(generator, key, null, null, null), 2);
+ assertEquals(generator.getNumGroups(), 3);
+
+ assertEquals(getGroupKeys(generator), Map.of(
+ 0, Arrays.asList(null, 1, 10L, null),
+ 1, Arrays.asList("a", null, 10L, null),
+ 2, Arrays.asList(null, null, null, null)));
+ }
+
+ @Test
+ public void testNumGroupsLimit() {
+ GroupIdGenerator generator = new MultiKeysGroupIdGenerator(KEY_TYPES,
NUM_KEY_COLUMNS, 2, 10);
+ Object[] key = new Object[NUM_KEY_COLUMNS];
+
+ assertEquals(getGroupId(generator, key, "a", 1, 10L), 0);
+ assertEquals(getGroupId(generator, key, "b", 2, 20L), 1);
+
+ // Limit is reached, new keys should not be added
+ assertEquals(getGroupId(generator, key, "c", 3, 30L), INVALID_ID);
+ // The first key value is known, but the second one is not
+ assertEquals(getGroupId(generator, key, "a", 3, 30L), INVALID_ID);
+ // All the key values are known, but the combination is not an existing
group
+ assertEquals(getGroupId(generator, key, "a", 2, 20L), INVALID_ID);
+ // Existing groups should still be resolved after the failed lookups
+ assertEquals(getGroupId(generator, key, "a", 1, 10L), 0);
+ assertEquals(getGroupId(generator, key, "b", 2, 20L), 1);
+ assertEquals(generator.getNumGroups(), 2);
+
+ assertEquals(getGroupKeys(generator), Map.of(
+ 0, Arrays.asList("a", 1, 10L, null),
+ 1, Arrays.asList("b", 2, 20L, null)));
+ }
+
+ /// Every row creates a new group, which forces the map to grow while the
keys are generated from the same reusable
+ /// array.
+ @Test
+ public void testManyGroups() {
+ int numGroups = 1000;
+ GroupIdGenerator generator = new MultiKeysGroupIdGenerator(KEY_TYPES,
NUM_KEY_COLUMNS, numGroups, 10);
+ Object[] key = new Object[NUM_KEY_COLUMNS];
+
+ for (int i = 0; i < numGroups; i++) {
+ assertEquals(getGroupId(generator, key, "k" + i, i, (long) i), i);
+ }
+ // Replay the same rows in reverse order
+ for (int i = numGroups - 1; i >= 0; i--) {
+ assertEquals(getGroupId(generator, key, "k" + i, i, (long) i), i);
+ }
+ assertEquals(generator.getNumGroups(), numGroups);
+
+ Map<Integer, List<Object>> groupKeys = getGroupKeys(generator);
+ assertEquals(groupKeys.size(), numGroups);
+ for (int i = 0; i < numGroups; i++) {
+ assertEquals(groupKeys.get(i), Arrays.asList("k" + i, i, (long) i,
null));
+ }
+ }
+
+ private static int getGroupId(GroupIdGenerator generator, Object[] key,
Object... keyValues) {
+ System.arraycopy(keyValues, 0, key, 0, NUM_KEY_COLUMNS);
+ return generator.getGroupId(key);
+ }
+
+ private static Map<Integer, List<Object>> getGroupKeys(GroupIdGenerator
generator) {
+ Map<Integer, List<Object>> groupKeys = new HashMap<>();
+ Iterator<GroupKey> iterator = generator.getGroupKeyIterator(NUM_COLUMNS);
+ while (iterator.hasNext()) {
+ GroupKey groupKey = iterator.next();
+ groupKeys.put(groupKey._groupId, Arrays.asList(groupKey._row));
+ }
+ return groupKeys;
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
