This is an automated email from the ASF dual-hosted git repository.
jackie 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 a7357d836bc Fix key map cleanup in group key generator (#16982)
a7357d836bc is described below
commit a7357d836bcfec91b9b5f3733132d5c3ac256fc5
Author: Xiaotian (Jackie) Jiang <[email protected]>
AuthorDate: Thu Oct 9 11:57:55 2025 -0700
Fix key map cleanup in group key generator (#16982)
---
.../groupby/DictionaryBasedGroupKeyGenerator.java | 134 ++++++++-------------
1 file changed, 52 insertions(+), 82 deletions(-)
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGenerator.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGenerator.java
index 7fbc0bdb266..cfe70b785bb 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGenerator.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGenerator.java
@@ -29,6 +29,7 @@ import java.util.Arrays;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
+import java.util.function.ToIntFunction;
import javax.annotation.Nullable;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.pinot.common.request.context.ExpressionContext;
@@ -147,16 +148,21 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
cardinalityProduct = Math.min(optimizedCardinality.getRight(),
cardinalityProduct);
}
}
+ // NOTE: We need to clean up the thread-local map before using it in case
RawKeyHolder.close() is not called
+ // for the previous segment
+ // TODO: Ensure RawKeyHolder.close()
if (longOverflow) {
// ArrayMapBasedHolder
_globalGroupIdUpperBound = numGroupsLimit;
Object2IntOpenHashMap<IntArray> groupIdMap =
THREAD_LOCAL_INT_ARRAY_MAP.get();
+ clearAndTrim(groupIdMap);
_rawKeyHolder = new ArrayMapBasedHolder(groupIdMap);
} else {
if (cardinalityProduct > Integer.MAX_VALUE) {
// LongMapBasedHolder
_globalGroupIdUpperBound = numGroupsLimit;
Long2IntOpenHashMap groupIdMap = THREAD_LOCAL_LONG_MAP.get();
+ clearAndTrim(groupIdMap);
_rawKeyHolder = new LongMapBasedHolder(groupIdMap);
} else {
_globalGroupIdUpperBound = Math.min((int) cardinalityProduct,
numGroupsLimit);
@@ -165,6 +171,7 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
if (cardinalityProduct > arrayBasedThreshold || numGroupsLimit <
cardinalityProduct) {
// IntMapBasedHolder
IntGroupIdMap groupIdMap = THREAD_LOCAL_INT_MAP.get();
+ groupIdMap.clearAndTrim();
_rawKeyHolder = new IntMapBasedHolder(groupIdMap);
} else {
_rawKeyHolder = new ArrayBasedHolder();
@@ -189,6 +196,26 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
return Pair.of(true, maxInitialResultHolderCapacity);
}
+ private static void clearAndTrim(Long2IntOpenHashMap map) {
+ int size = map.size();
+ if (size > 0) {
+ map.clear();
+ if (size > MAX_CACHING_MAP_SIZE) {
+ map.trim();
+ }
+ }
+ }
+
+ private static void clearAndTrim(Object2IntOpenHashMap<IntArray> map) {
+ int size = map.size();
+ if (size > 0) {
+ map.clear();
+ if (size > MAX_CACHING_MAP_SIZE) {
+ map.trim();
+ }
+ }
+ }
+
@Override
public int getGlobalGroupKeyUpperBound() {
return _globalGroupIdUpperBound;
@@ -283,10 +310,6 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
private final boolean[] _flags = new boolean[_globalGroupIdUpperBound];
private int _numKeys = 0;
- @Override
- public void close() {
- }
-
@Override
public void processSingleValue(int numDocs, int[] outGroupIds) {
switch (_numGroupByExpressions) {
@@ -375,7 +398,7 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
@Override
public Iterator<GroupKey> getGroupKeys() {
- return new Iterator<GroupKey>() {
+ return new Iterator<>() {
private int _currentGroupId;
private final GroupKey _groupKey = new GroupKey();
@@ -412,15 +435,14 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
public int getNumKeys() {
return _numKeys;
}
- }
-
- private class IntMapBasedHolder implements RawKeyHolder {
- private final IntGroupIdMap _groupIdMap;
@Override
public void close() {
- _groupIdMap.clearAndTrim();
}
+ }
+
+ private class IntMapBasedHolder implements RawKeyHolder {
+ private final IntGroupIdMap _groupIdMap;
public IntMapBasedHolder(IntGroupIdMap groupIdMap) {
_groupIdMap = groupIdMap;
@@ -470,7 +492,7 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
@Override
public Iterator<GroupKey> getGroupKeys() {
- return new Iterator<GroupKey>() {
+ return new Iterator<>() {
private final Iterator<IntGroupIdMap.Entry> _iterator =
_groupIdMap.iterator();
private final GroupKey _groupKey = new GroupKey();
@@ -498,6 +520,11 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
public int getNumKeys() {
return _groupIdMap.size();
}
+
+ @Override
+ public void close() {
+ _groupIdMap.clearAndTrim();
+ }
}
/**
@@ -611,27 +638,6 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
return rawValue;
}
- /**
- * Helper method to get the string key from the raw key.
- */
- private String getStringKey(int rawKey) {
- // Specialize single group-by column case
- if (_numGroupByExpressions == 1) {
- return _dictionaries[0].getStringValue(rawKey);
- } else {
- int cardinality = _cardinalities[0];
- StringBuilder groupKeyBuilder = new
StringBuilder(_dictionaries[0].getStringValue(rawKey % cardinality));
- rawKey /= cardinality;
- for (int i = 1; i < _numGroupByExpressions; i++) {
- groupKeyBuilder.append(GroupKeyGenerator.DELIMITER);
- cardinality = _cardinalities[i];
- groupKeyBuilder.append(_dictionaries[i].getStringValue(rawKey %
cardinality));
- rawKey /= cardinality;
- }
- return groupKeyBuilder.toString();
- }
- }
-
private class LongMapBasedHolder implements RawKeyHolder {
private final Long2IntOpenHashMap _groupIdMap;
@@ -639,15 +645,6 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
_groupIdMap = groupIdMap;
}
- @Override
- public void close() {
- int size = _groupIdMap.size();
- _groupIdMap.clear();
- if (size > MAX_CACHING_MAP_SIZE) {
- _groupIdMap.trim();
- }
- }
-
@Override
public void processSingleValue(int numDocs, int[] outGroupIds) {
for (int i = 0; i < numDocs; i++) {
@@ -689,7 +686,7 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
@Override
public Iterator<GroupKey> getGroupKeys() {
- return new Iterator<GroupKey>() {
+ return new Iterator<>() {
private final ObjectIterator<Long2IntMap.Entry> _iterator =
_groupIdMap.long2IntEntrySet().fastIterator();
private final GroupKey _groupKey = new GroupKey();
@@ -717,6 +714,11 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
public int getNumKeys() {
return _groupIdMap.size();
}
+
+ @Override
+ public void close() {
+ clearAndTrim(_groupIdMap);
+ }
}
/**
@@ -805,22 +807,6 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
return groupKeys;
}
- /**
- * Helper method to get the string key from the raw key.
- */
- private String getStringKey(long rawKey) {
- int cardinality = _cardinalities[0];
- StringBuilder groupKeyBuilder = new
StringBuilder(_dictionaries[0].getStringValue((int) (rawKey % cardinality)));
- rawKey /= cardinality;
- for (int i = 1; i < _numGroupByExpressions; i++) {
- groupKeyBuilder.append(GroupKeyGenerator.DELIMITER);
- cardinality = _cardinalities[i];
- groupKeyBuilder.append(_dictionaries[i].getStringValue((int) (rawKey %
cardinality)));
- rawKey /= cardinality;
- }
- return groupKeyBuilder.toString();
- }
-
private class ArrayMapBasedHolder implements RawKeyHolder {
private final Object2IntOpenHashMap<IntArray> _groupIdMap;
@@ -828,15 +814,6 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
_groupIdMap = groupIdMap;
}
- @Override
- public void close() {
- int size = _groupIdMap.size();
- _groupIdMap.clear();
- if (size > MAX_CACHING_MAP_SIZE) {
- _groupIdMap.trim();
- }
- }
-
@Override
public void processSingleValue(int numDocs, int[] outGroupIds) {
for (int i = 0; i < numDocs; i++) {
@@ -864,7 +841,7 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
private int getGroupId(IntArray rawKey) {
int numGroups = _groupIdMap.size();
if (numGroups < _globalGroupIdUpperBound) {
- return _groupIdMap.computeIntIfAbsent(rawKey, k -> numGroups);
+ return _groupIdMap.computeIfAbsent(rawKey, (ToIntFunction<? super
IntArray>) k -> numGroups);
} else {
return _groupIdMap.getInt(rawKey);
}
@@ -877,7 +854,7 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
@Override
public Iterator<GroupKey> getGroupKeys() {
- return new Iterator<GroupKey>() {
+ return new Iterator<>() {
private final ObjectIterator<Object2IntMap.Entry<IntArray>> _iterator =
_groupIdMap.object2IntEntrySet().fastIterator();
private final GroupKey _groupKey = new GroupKey();
@@ -906,6 +883,11 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
public int getNumKeys() {
return _groupIdMap.size();
}
+
+ @Override
+ public void close() {
+ clearAndTrim(_groupIdMap);
+ }
}
/**
@@ -996,18 +978,6 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
return groupKeys;
}
- /**
- * Helper method to get the string key from the raw key.
- */
- private String getStringKey(IntArray rawKey) {
- StringBuilder groupKeyBuilder = new
StringBuilder(_dictionaries[0].getStringValue(rawKey._elements[0]));
- for (int i = 1; i < _numGroupByExpressions; i++) {
- groupKeyBuilder.append(GroupKeyGenerator.DELIMITER);
-
groupKeyBuilder.append(_dictionaries[i].getStringValue(rawKey._elements[i]));
- }
- return groupKeyBuilder.toString();
- }
-
/**
* Fast int-to-int hashmap with {@link #INVALID_ID} as the default return
value.
* <p>Different from {@link it.unimi.dsi.fastutil.ints.Int2IntOpenHashMap},
this map uses one single array to store
@@ -1108,7 +1078,7 @@ public class DictionaryBasedGroupKeyGenerator implements
GroupKeyGenerator {
}
public Iterator<Entry> iterator() {
- return new Iterator<Entry>() {
+ return new Iterator<>() {
private final Entry _entry = new Entry();
private int _index;
private int _numRemainingEntries = _size;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
