diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java index 68b32d3b122e..30f07061fab8 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java @@ -668,6 +668,12 @@ public static Integer getNumGroupsWarningLimit(Map queryOptions) return checkedParseIntPositive(QueryOptionKey.NUM_GROUPS_WARNING_LIMIT, numGroupsWarningLimit); } + @Nullable + public static Boolean isGroupByOffHeap(Map queryOptions) { + String groupByOffHeap = queryOptions.get(QueryOptionKey.GROUP_BY_OFF_HEAP); + return groupByOffHeap != null ? Boolean.parseBoolean(groupByOffHeap) : null; + } + @Nullable public static Integer getMaxInitialResultHolderCapacity(Map queryOptions) { String maxInitialResultHolderCapacity = queryOptions.get(QueryOptionKey.MAX_INITIAL_RESULT_HOLDER_CAPACITY); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/table/TableResizer.java b/pinot-core/src/main/java/org/apache/pinot/core/data/table/TableResizer.java index d17a14b2772b..42935748ed21 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/table/TableResizer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/table/TableResizer.java @@ -329,7 +329,7 @@ private Collection getUnsortedTopRecords(Map recordsMap, in /// This method is to be called from individual segment if the intermediate results need to be trimmed. public List sortInSegmentResults(GroupKeyGenerator groupKeyGenerator, GroupByResultHolder[] groupByResultHolders, int size) { - // getNumKeys() does not count nulls + // NOTE: getNumKeys() counts every group, including the null group when null handling is enabled assert groupKeyGenerator.getNumKeys() <= size; Iterator groupKeyIterator = groupKeyGenerator.getGroupKeys(); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java index abc88149a06a..a871ab704f78 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java @@ -106,38 +106,39 @@ protected void processSegments() { ((AcquireReleaseColumnsSegmentOperator) operator).acquire(); } GroupByResultsBlock resultsBlock = (GroupByResultsBlock) operator.nextBlock(); - if (_indexedTable == null) { - synchronized (this) { - if (_indexedTable == null) { - _indexedTable = GroupByUtils.createIndexedTableForCombineOperator(resultsBlock, _queryContext, _numTasks, - _executorService); + // Hold the group-by result so its group key generator (which may own off-heap resources) is always + // released, even when indexed-table creation or the merge below throws + AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult(); + try { + if (_indexedTable == null) { + synchronized (this) { + if (_indexedTable == null) { + _indexedTable = GroupByUtils.createIndexedTableForCombineOperator(resultsBlock, _queryContext, + _numTasks, _executorService); + } } } - } - if (resultsBlock.isGroupsTrimmed()) { - _groupsTrimmed = true; - } - // Set groups limit reached flag. - if (resultsBlock.isNumGroupsLimitReached()) { - _numGroupsLimitReached = true; - } - if (resultsBlock.isNumGroupsWarningLimitReached()) { - _numGroupsWarningLimitReached = true; - } + if (resultsBlock.isGroupsTrimmed()) { + _groupsTrimmed = true; + } + // Set groups limit reached flag. + if (resultsBlock.isNumGroupsLimitReached()) { + _numGroupsLimitReached = true; + } + if (resultsBlock.isNumGroupsWarningLimitReached()) { + _numGroupsWarningLimitReached = true; + } - // Merge aggregation group-by result. - // Iterate over the group-by keys, for each key, update the group-by result in the indexedTable - Collection intermediateRecords = resultsBlock.getIntermediateRecords(); - // Count the number of merged keys - int mergedKeys = 0; - // For now, only GroupBy OrderBy query has pre-constructed intermediate records - if (intermediateRecords == null) { // Merge aggregation group-by result. - AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult(); - if (aggregationGroupByResult != null) { - // Iterate over the group-by keys, for each key, update the group-by result in the indexedTable - try { + // Iterate over the group-by keys, for each key, update the group-by result in the indexedTable + Collection intermediateRecords = resultsBlock.getIntermediateRecords(); + // Count the number of merged keys + int mergedKeys = 0; + // For now, only GroupBy OrderBy query has pre-constructed intermediate records + if (intermediateRecords == null) { + if (aggregationGroupByResult != null) { + // Iterate over the group-by keys, for each key, update the group-by result in the indexedTable Iterator dicGroupKeyIterator = aggregationGroupByResult.getGroupKeyIterator(); while (dicGroupKeyIterator.hasNext()) { QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, EXPLAIN_NAME); @@ -150,16 +151,18 @@ protected void processSegments() { } _indexedTable.upsert(new Key(keys), new Record(values)); } - } finally { - // Release the resources used by the group key generator - aggregationGroupByResult.closeGroupKeyGenerator(); + } + } else { + for (IntermediateRecord intermediateResult : intermediateRecords) { + QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, EXPLAIN_NAME); + //TODO: change upsert api so that it accepts intermediateRecord directly + _indexedTable.upsert(intermediateResult._key, intermediateResult._record); } } - } else { - for (IntermediateRecord intermediateResult : intermediateRecords) { - QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, EXPLAIN_NAME); - //TODO: change upsert api so that it accepts intermediateRecord directly - _indexedTable.upsert(intermediateResult._key, intermediateResult._record); + } finally { + if (aggregationGroupByResult != null) { + // Release the resources used by the group key generator + aggregationGroupByResult.closeGroupKeyGenerator(); } } } catch (RuntimeException e) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java index 166948c106c0..fe29c1e96e8b 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java @@ -136,6 +136,23 @@ protected GroupByResultsBlock getNextBlock() { resultHolderIndexMap.put(_aggregationFunctions[i], i); } + GroupKeyGenerator[] createdGroupKeyGenerator = new GroupKeyGenerator[1]; + try { + return processAndBuildResultsBlock(groupByResultHolders, resultHolderIndexMap, createdGroupKeyGenerator); + } catch (Throwable t) { + // Release group-by resources (including off-heap key tables and result holders) that would otherwise leak. + // Close is idempotent on all generators; on the success path the generator is either closed on the trim/sort + // paths below or handed to the combine operator, which closes it after the merge. + if (createdGroupKeyGenerator[0] != null) { + createdGroupKeyGenerator[0].close(); + } + throw t; + } + } + + private GroupByResultsBlock processAndBuildResultsBlock(GroupByResultHolder[] groupByResultHolders, + IdentityHashMap resultHolderIndexMap, + GroupKeyGenerator[] createdGroupKeyGenerator) { GroupKeyGenerator groupKeyGenerator = null; for (AggregationInfo aggregationInfo : _aggregationInfos) { AggregationFunction[] aggregationFunctions = aggregationInfo.getFunctions(); @@ -162,6 +179,7 @@ protected GroupByResultsBlock getNextBlock() { // GroupByExecutor with a pre-existing GroupKeyGenerator so that the GroupKeyGenerator can be shared across // loop iterations i.e. across all aggs. groupKeyGenerator = groupByExecutor.getGroupKeyGenerator(); + createdGroupKeyGenerator[0] = groupKeyGenerator; int numDocsScanned = 0; ValueBlock valueBlock; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java index e138a8c7973c..c53f85ea05c5 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java @@ -130,6 +130,19 @@ protected GroupByResultsBlock getNextBlock() { } else { groupByExecutor = new DefaultGroupByExecutor(_queryContext, _groupByExpressions, _projectOperator); } + try { + return processAndBuildResultsBlock(groupByExecutor); + } catch (Throwable t) { + // Release group-by resources (including off-heap key tables and result holders) that would otherwise leak. + // On the success path, ownership either ends inside processAndBuildResultsBlock (trim/sort paths close the + // generator there) or moves to the results block consumer (the combine operator closes the generator after + // merging the AggregationGroupByResult). Close is idempotent on all generators. + groupByExecutor.getGroupKeyGenerator().close(); + throw t; + } + } + + private GroupByResultsBlock processAndBuildResultsBlock(GroupByExecutor groupByExecutor) { ValueBlock valueBlock; while ((valueBlock = _projectOperator.nextBlock()) != null) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java index 03199b367cb0..529bb14befc8 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java @@ -46,6 +46,7 @@ import org.apache.pinot.core.plan.StreamingInstanceResponsePlanNode; import org.apache.pinot.core.plan.StreamingSelectionPlanNode; import org.apache.pinot.core.query.aggregation.function.AggregationFunction; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapGroupByBufferPool; import org.apache.pinot.core.query.executor.ResultsBlockStreamer; import org.apache.pinot.core.query.prefetch.FetchPlanner; import org.apache.pinot.core.query.prefetch.FetchPlannerRegistry; @@ -117,6 +118,8 @@ public class InstancePlanMakerImplV2 implements PlanMaker { private int _groupByTrimThreshold = Server.DEFAULT_QUERY_EXECUTOR_GROUPBY_TRIM_THRESHOLD; private AndRestrictionPushdownMode _andRestrictionPushdownMode = Server.DEFAULT_QUERY_EXECUTOR_AND_RESTRICTION_PUSHDOWN_MODE; + // Whether to store group-by key tables and fixed-width result holders in off-heap (direct) memory + private boolean _groupByOffHeap = Server.DEFAULT_QUERY_EXECUTOR_GROUPBY_OFF_HEAP; @Override public void init(PinotConfiguration queryExecutorConfig) { @@ -150,11 +153,16 @@ public void init(PinotConfiguration queryExecutorConfig) { Server.DEFAULT_QUERY_EXECUTOR_AND_RESTRICTION_PUSHDOWN_MODE.name())); Preconditions.checkState(_groupByTrimThreshold > 0, "Invalid configurable: groupByTrimThreshold: %d must be positive", _groupByTrimThreshold); + _groupByOffHeap = + queryExecutorConfig.getProperty(Server.GROUPBY_OFF_HEAP, Server.DEFAULT_QUERY_EXECUTOR_GROUPBY_OFF_HEAP); + OffHeapGroupByBufferPool.setMaxBytesPerThread( + queryExecutorConfig.getProperty(Server.GROUPBY_OFF_HEAP_POOL_MAX_BYTES_PER_THREAD, + Server.DEFAULT_QUERY_EXECUTOR_GROUPBY_OFF_HEAP_POOL_MAX_BYTES_PER_THREAD)); LOGGER.info("Initialized plan maker with maxExecutionThreads: {}, defaultExecutionThreads: {}, " + "maxInitialResultHolderCapacity: {}, numGroupsLimit: {}, minSegmentGroupTrimSize: {}, " - + "minServerGroupTrimSize: {}, groupByTrimThreshold: {}", + + "minServerGroupTrimSize: {}, groupByTrimThreshold: {}, groupByOffHeap: {}", _maxExecutionThreads, _defaultExecutionThreads, _maxInitialResultHolderCapacity, _numGroupsLimit, - _minSegmentGroupTrimSize, _minServerGroupTrimSize, _groupByTrimThreshold); + _minSegmentGroupTrimSize, _minServerGroupTrimSize, _groupByTrimThreshold, _groupByOffHeap); } @VisibleForTesting @@ -207,6 +215,11 @@ public void setGroupByTrimThreshold(int groupByTrimThreshold) { _groupByTrimThreshold = groupByTrimThreshold; } + @VisibleForTesting + public void setGroupByOffHeap(boolean groupByOffHeap) { + _groupByOffHeap = groupByOffHeap; + } + @Override public Plan makeInstancePlan(List segmentContexts, QueryContext queryContext, ExecutorService executorService) { @@ -334,6 +347,9 @@ void applyQueryOptions(QueryContext queryContext) { } else { queryContext.setNumGroupsLimit(_numGroupsLimit); } + // Set groupByOffHeap + Boolean groupByOffHeap = QueryOptionsUtils.isGroupByOffHeap(queryOptions); + queryContext.setGroupByOffHeap(groupByOffHeap != null ? groupByOffHeap : _groupByOffHeap); // Set numGroupsWarningThreshold queryContext.setNumGroupsWarningLimit(_numGroupsWarningLimit); // Set minSegmentGroupTrimSize diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunction.java index 79ca2736e8aa..99224bc8ab5b 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunction.java @@ -150,6 +150,18 @@ public interface AggregationFunction arguments, bo numExpressions); if (arguments.size() == 2) { _p = arguments.get(1).getLiteral().getIntValue(); + // hash4j's UltraLogLog.create(p) bound; validate at plan time so the off-heap holder (which sizes direct + // memory as 1 << p without calling create) can never see an out-of-range p + Preconditions.checkArgument(_p >= 3 && _p <= 26, + "Invalid p for DistinctCountULL: %s, must be in [3, 26]", _p); } else { _p = CommonConstants.Helix.DEFAULT_ULTRALOGLOG_P; } @@ -77,6 +83,11 @@ public GroupByResultHolder createGroupByResultHolder(int initialCapacity, int ma return new ObjectGroupByResultHolder(initialCapacity, maxCapacity); } + @Override + public GroupByResultHolder createOffHeapGroupByResultHolder(int initialCapacity, int maxCapacity) { + return new OffHeapUltraLogLogGroupByResultHolder(_p, initialCapacity, maxCapacity); + } + @Override public void aggregate(int length, AggregationResultHolder aggregationResultHolder, Map blockValSetMap) { @@ -331,8 +342,7 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu int[] intValues = blockValSet.getIntValuesSV(); forEachNotNull(length, blockValSet, (from, to) -> { for (int i = from; i < to; i++) { - UltraLogLogUtils.hashObject(intValues[i]) - .ifPresent(getULL(groupByResultHolder, groupKeyArray[i])::add); + addHash(groupByResultHolder, groupKeyArray[i], UltraLogLogUtils.hashObject(intValues[i])); } }); break; @@ -340,8 +350,7 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu long[] longValues = blockValSet.getLongValuesSV(); forEachNotNull(length, blockValSet, (from, to) -> { for (int i = from; i < to; i++) { - UltraLogLogUtils.hashObject(longValues[i]) - .ifPresent(getULL(groupByResultHolder, groupKeyArray[i])::add); + addHash(groupByResultHolder, groupKeyArray[i], UltraLogLogUtils.hashObject(longValues[i])); } }); break; @@ -349,8 +358,7 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu float[] floatValues = blockValSet.getFloatValuesSV(); forEachNotNull(length, blockValSet, (from, to) -> { for (int i = from; i < to; i++) { - UltraLogLogUtils.hashObject(floatValues[i]) - .ifPresent(getULL(groupByResultHolder, groupKeyArray[i])::add); + addHash(groupByResultHolder, groupKeyArray[i], UltraLogLogUtils.hashObject(floatValues[i])); } }); break; @@ -358,8 +366,7 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu double[] doubleValues = blockValSet.getDoubleValuesSV(); forEachNotNull(length, blockValSet, (from, to) -> { for (int i = from; i < to; i++) { - UltraLogLogUtils.hashObject(doubleValues[i]) - .ifPresent(getULL(groupByResultHolder, groupKeyArray[i])::add); + addHash(groupByResultHolder, groupKeyArray[i], UltraLogLogUtils.hashObject(doubleValues[i])); } }); break; @@ -367,8 +374,7 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu String[] stringValues = blockValSet.getStringValuesSV(); forEachNotNull(length, blockValSet, (from, to) -> { for (int i = from; i < to; i++) { - UltraLogLogUtils.hashObject(stringValues[i]) - .ifPresent(getULL(groupByResultHolder, groupKeyArray[i])::add); + addHash(groupByResultHolder, groupKeyArray[i], UltraLogLogUtils.hashObject(stringValues[i])); } }); break; @@ -376,8 +382,7 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu byte[][] bytesValues = blockValSet.getBytesValuesSV(); forEachNotNull(length, blockValSet, (from, to) -> { for (int i = from; i < to; i++) { - UltraLogLogUtils.hashObject(bytesValues[i]) - .ifPresent(getULL(groupByResultHolder, groupKeyArray[i])::add); + addHash(groupByResultHolder, groupKeyArray[i], UltraLogLogUtils.hashObject(bytesValues[i])); } }); break; @@ -407,9 +412,10 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu int[][] intValues = blockValSet.getIntValuesMV(); forEachNotNull(length, blockValSet, (from, to) -> { for (int i = from; i < to; i++) { - UltraLogLog ull = getULL(groupByResultHolder, groupKeyArray[i]); + int groupKey = groupKeyArray[i]; + touchULL(groupByResultHolder, groupKey); for (int value : intValues[i]) { - UltraLogLogUtils.hashObject(value).ifPresent(ull::add); + addHash(groupByResultHolder, groupKey, UltraLogLogUtils.hashObject(value)); } } }); @@ -418,9 +424,10 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu long[][] longValues = blockValSet.getLongValuesMV(); forEachNotNull(length, blockValSet, (from, to) -> { for (int i = from; i < to; i++) { - UltraLogLog ull = getULL(groupByResultHolder, groupKeyArray[i]); + int groupKey = groupKeyArray[i]; + touchULL(groupByResultHolder, groupKey); for (long value : longValues[i]) { - UltraLogLogUtils.hashObject(value).ifPresent(ull::add); + addHash(groupByResultHolder, groupKey, UltraLogLogUtils.hashObject(value)); } } }); @@ -429,9 +436,10 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu float[][] floatValues = blockValSet.getFloatValuesMV(); forEachNotNull(length, blockValSet, (from, to) -> { for (int i = from; i < to; i++) { - UltraLogLog ull = getULL(groupByResultHolder, groupKeyArray[i]); + int groupKey = groupKeyArray[i]; + touchULL(groupByResultHolder, groupKey); for (float value : floatValues[i]) { - UltraLogLogUtils.hashObject(value).ifPresent(ull::add); + addHash(groupByResultHolder, groupKey, UltraLogLogUtils.hashObject(value)); } } }); @@ -440,9 +448,10 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu double[][] doubleValues = blockValSet.getDoubleValuesMV(); forEachNotNull(length, blockValSet, (from, to) -> { for (int i = from; i < to; i++) { - UltraLogLog ull = getULL(groupByResultHolder, groupKeyArray[i]); + int groupKey = groupKeyArray[i]; + touchULL(groupByResultHolder, groupKey); for (double value : doubleValues[i]) { - UltraLogLogUtils.hashObject(value).ifPresent(ull::add); + addHash(groupByResultHolder, groupKey, UltraLogLogUtils.hashObject(value)); } } }); @@ -451,9 +460,10 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu String[][] stringValues = blockValSet.getStringValuesMV(); forEachNotNull(length, blockValSet, (from, to) -> { for (int i = from; i < to; i++) { - UltraLogLog ull = getULL(groupByResultHolder, groupKeyArray[i]); + int groupKey = groupKeyArray[i]; + touchULL(groupByResultHolder, groupKey); for (String value : stringValues[i]) { - UltraLogLogUtils.hashObject(value).ifPresent(ull::add); + addHash(groupByResultHolder, groupKey, UltraLogLogUtils.hashObject(value)); } } }); @@ -462,9 +472,10 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu byte[][][] bytesValues = blockValSet.getBytesValuesMV(); forEachNotNull(length, blockValSet, (from, to) -> { for (int i = from; i < to; i++) { - UltraLogLog ull = getULL(groupByResultHolder, groupKeyArray[i]); + int groupKey = groupKeyArray[i]; + touchULL(groupByResultHolder, groupKey); for (byte[] value : bytesValues[i]) { - UltraLogLogUtils.hashObject(value).ifPresent(ull::add); + addHash(groupByResultHolder, groupKey, UltraLogLogUtils.hashObject(value)); } } }); @@ -608,9 +619,9 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR for (int i = from; i < to; i++) { int[] intRow = intValues[i]; for (int groupKey : groupKeysArray[i]) { - UltraLogLog ull = getULL(groupByResultHolder, groupKey); + touchULL(groupByResultHolder, groupKey); for (int value : intRow) { - UltraLogLogUtils.hashObject(value).ifPresent(ull::add); + addHash(groupByResultHolder, groupKey, UltraLogLogUtils.hashObject(value)); } } } @@ -622,9 +633,9 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR for (int i = from; i < to; i++) { long[] longRow = longValues[i]; for (int groupKey : groupKeysArray[i]) { - UltraLogLog ull = getULL(groupByResultHolder, groupKey); + touchULL(groupByResultHolder, groupKey); for (long value : longRow) { - UltraLogLogUtils.hashObject(value).ifPresent(ull::add); + addHash(groupByResultHolder, groupKey, UltraLogLogUtils.hashObject(value)); } } } @@ -636,9 +647,9 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR for (int i = from; i < to; i++) { float[] floatRow = floatValues[i]; for (int groupKey : groupKeysArray[i]) { - UltraLogLog ull = getULL(groupByResultHolder, groupKey); + touchULL(groupByResultHolder, groupKey); for (float value : floatRow) { - UltraLogLogUtils.hashObject(value).ifPresent(ull::add); + addHash(groupByResultHolder, groupKey, UltraLogLogUtils.hashObject(value)); } } } @@ -650,9 +661,9 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR for (int i = from; i < to; i++) { double[] doubleRow = doubleValues[i]; for (int groupKey : groupKeysArray[i]) { - UltraLogLog ull = getULL(groupByResultHolder, groupKey); + touchULL(groupByResultHolder, groupKey); for (double value : doubleRow) { - UltraLogLogUtils.hashObject(value).ifPresent(ull::add); + addHash(groupByResultHolder, groupKey, UltraLogLogUtils.hashObject(value)); } } } @@ -664,9 +675,9 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR for (int i = from; i < to; i++) { String[] stringRow = stringValues[i]; for (int groupKey : groupKeysArray[i]) { - UltraLogLog ull = getULL(groupByResultHolder, groupKey); + touchULL(groupByResultHolder, groupKey); for (String value : stringRow) { - UltraLogLogUtils.hashObject(value).ifPresent(ull::add); + addHash(groupByResultHolder, groupKey, UltraLogLogUtils.hashObject(value)); } } } @@ -678,9 +689,9 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR for (int i = from; i < to; i++) { byte[][] bytesRow = bytesValues[i]; for (int groupKey : groupKeysArray[i]) { - UltraLogLog ull = getULL(groupByResultHolder, groupKey); + touchULL(groupByResultHolder, groupKey); for (byte[] value : bytesRow) { - UltraLogLogUtils.hashObject(value).ifPresent(ull::add); + addHash(groupByResultHolder, groupKey, UltraLogLogUtils.hashObject(value)); } } } @@ -811,7 +822,7 @@ protected static RoaringBitmap getDictIdBitmap(GroupByResultHolder groupByResult return dictIdsWrapper._dictIdBitmap; } - /// Returns the HyperLogLogPlus for the given group key or creates a new one if it does not exist. + /// Returns the UltraLogLog for the given group key or creates a new one if it does not exist. protected UltraLogLog getULL(GroupByResultHolder groupByResultHolder, int groupKey) { UltraLogLog ull = groupByResultHolder.getResult(groupKey); if (ull == null) { @@ -821,6 +832,35 @@ protected UltraLogLog getULL(GroupByResultHolder groupByResultHolder, int groupK return ull; } + /// Ensures the per-group ULL state exists (off-heap slot or on-heap object). Mirrors the on-heap path's eager + /// state creation so untouched-vs-empty groups never diverge between the two modes. + protected void touchULL(GroupByResultHolder groupByResultHolder, int groupKey) { + if (groupByResultHolder instanceof OffHeapUltraLogLogGroupByResultHolder) { + ((OffHeapUltraLogLogGroupByResultHolder) groupByResultHolder).touch(groupKey); + } else { + getULL(groupByResultHolder, groupKey); + } + } + + /// Adds the hashed value (when present) into the per-group ULL state, routing register updates straight into + /// direct memory when the query runs with off-heap group-by state. + protected void addHash(GroupByResultHolder groupByResultHolder, int groupKey, Optional hashValue) { + if (groupByResultHolder instanceof OffHeapUltraLogLogGroupByResultHolder) { + OffHeapUltraLogLogGroupByResultHolder offHeapHolder = + (OffHeapUltraLogLogGroupByResultHolder) groupByResultHolder; + if (hashValue.isPresent()) { + // add() creates the slot on first touch, so no separate touch is needed on this (hot) branch + offHeapHolder.add(groupKey, hashValue.get()); + } else { + // State creation still matches the on-heap branch, where getULL runs before the presence check + offHeapHolder.touch(groupKey); + } + } else { + UltraLogLog ull = getULL(groupByResultHolder, groupKey); + hashValue.ifPresent(ull::add); + } + } + /// Helper method to set dictionary id for the given group keys into the result holder. private static void setDictIdForGroupKeys(GroupByResultHolder groupByResultHolder, int[] groupKeys, Dictionary dictionary, int dictId) { @@ -831,9 +871,9 @@ private static void setDictIdForGroupKeys(GroupByResultHolder groupByResultHolde /// Helper method to set value for the given group keys into the result holder. private void setValueForGroupKeys(GroupByResultHolder groupByResultHolder, int[] groupKeys, Object value) { + Optional hashValue = UltraLogLogUtils.hashObject(value); for (int groupKey : groupKeys) { - UltraLogLogUtils.hashObject(value) - .ifPresent(getULL(groupByResultHolder, groupKey)::add); + addHash(groupByResultHolder, groupKey, hashValue); } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java index 8a4015cd84da..b64e4a735f02 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.core.query.aggregation.groupby; +import com.google.common.base.Preconditions; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -37,6 +38,10 @@ import org.apache.pinot.core.plan.DocIdSetPlanNode; import org.apache.pinot.core.query.aggregation.function.AggregationFunction; import org.apache.pinot.core.query.aggregation.function.AggregationFunctionUtils; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapDoubleGroupByResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapIntGroupByResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapLongGroupByResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.offheap.ResourceTrackingGroupKeyGenerator; import org.apache.pinot.core.query.request.context.QueryContext; @@ -97,41 +102,63 @@ public DefaultGroupByExecutor(QueryContext queryContext, AggregationFunction[] a // Initialize group key generator int numGroupsLimit = queryContext.getNumGroupsLimit(); int maxInitialResultHolderCapacity = queryContext.getMaxInitialResultHolderCapacity(); + Map groupByExpressionSizesFromPredicates = + queryContext.isOptimizeMaxInitialResultHolderCapacity() + ? getGroupByExpressionSizesFromPredicates(queryContext, projectOperator) : null; + // Off-heap group-by is not enabled for grouping sets yet: GroupingSetsGroupKeyGenerator keeps its key map and + // on-the-fly dictionaries on heap, so only the fixed-width result holders could move off-heap, and that + // combination is untested. The close plumbing already covers grouping sets (the trim path closes the generator + // in GroupByUtils.buildGroupingSetsResultsBlock, and the combine operators close the AggregationGroupByResult's + // generator), so enabling it later mainly requires off-heap key storage in that generator plus test coverage. + boolean groupByOffHeap = queryContext.isGroupByOffHeap() && !groupingSets; if (groupKeyGenerator != null) { + // Shared generator (filtered aggregations): if the first executor created it in off-heap mode, it is already + // wrapped in a ResourceTrackingGroupKeyGenerator, and this executor registers its holders on the same wrapper _groupKeyGenerator = groupKeyGenerator; - } else if (groupingSets) { - _groupKeyGenerator = - new GroupingSetsGroupKeyGenerator(projectOperator, groupByExpressions, queryContext.getGroupingSets(), - numGroupsLimit, _nullHandlingEnabled); } else { - Map groupByExpressionSizesFromPredicates = - queryContext.isOptimizeMaxInitialResultHolderCapacity() - ? getGroupByExpressionSizesFromPredicates(queryContext, projectOperator) : null; - // Null handling does not steer this choice: every generator below gives a null an id of its own, so the - // encoding of the group-by columns decides on its own which one to use. - if (hasNoDictionaryGroupByExpression) { + GroupKeyGenerator generator; + if (groupingSets) { + generator = new GroupingSetsGroupKeyGenerator(projectOperator, groupByExpressions, + queryContext.getGroupingSets(), numGroupsLimit, _nullHandlingEnabled); + } else if (hasNoDictionaryGroupByExpression) { if (groupByExpressions.length == 1) { - _groupKeyGenerator = + generator = new NoDictionarySingleColumnGroupKeyGenerator(projectOperator, groupByExpressions[0], numGroupsLimit, - _nullHandlingEnabled, groupByExpressionSizesFromPredicates); + _nullHandlingEnabled, groupByExpressionSizesFromPredicates, groupByOffHeap); } else { - _groupKeyGenerator = + generator = new NoDictionaryMultiColumnGroupKeyGenerator(projectOperator, groupByExpressions, numGroupsLimit, - _nullHandlingEnabled, groupByExpressionSizesFromPredicates); + _nullHandlingEnabled, groupByExpressionSizesFromPredicates, groupByOffHeap); } } else { - _groupKeyGenerator = new DictionaryBasedGroupKeyGenerator(projectOperator, groupByExpressions, numGroupsLimit, - maxInitialResultHolderCapacity, _nullHandlingEnabled, groupByExpressionSizesFromPredicates); + generator = new DictionaryBasedGroupKeyGenerator(projectOperator, groupByExpressions, numGroupsLimit, + maxInitialResultHolderCapacity, _nullHandlingEnabled, groupByExpressionSizesFromPredicates, + groupByOffHeap); } + _groupKeyGenerator = groupByOffHeap ? new ResourceTrackingGroupKeyGenerator(generator) : generator; } - // Initialize result holders + // Initialize result holders. In off-heap mode, fixed-width holders are mirrored off-heap and registered on the + // resource-tracking generator so the existing generator close() call sites release them. + ResourceTrackingGroupKeyGenerator offHeapResourceTracker = + _groupKeyGenerator instanceof ResourceTrackingGroupKeyGenerator + ? (ResourceTrackingGroupKeyGenerator) _groupKeyGenerator : null; int maxNumResults = _groupKeyGenerator.getGlobalGroupKeyUpperBound(); int initialCapacity = Math.min(maxNumResults, maxInitialResultHolderCapacity); int numAggregationFunctions = _aggregationFunctions.length; _groupByResultHolders = new GroupByResultHolder[numAggregationFunctions]; - for (int i = 0; i < numAggregationFunctions; i++) { - _groupByResultHolders[i] = _aggregationFunctions[i].createGroupByResultHolder(initialCapacity, maxNumResults); + try { + for (int i = 0; i < numAggregationFunctions; i++) { + _groupByResultHolders[i] = offHeapResourceTracker != null + ? createOffHeapCapableResultHolder(_aggregationFunctions[i], initialCapacity, maxNumResults, + offHeapResourceTracker) + : _aggregationFunctions[i].createGroupByResultHolder(initialCapacity, maxNumResults); + } + } catch (Throwable t) { + // Holder creation failed midway: release the generator (and any off-heap holders already registered on it) + // because the caller never gets an executor reference to clean up. Close is idempotent. + _groupKeyGenerator.close(); + throw t; } // Initialize map from document Id to group key @@ -144,6 +171,41 @@ public DefaultGroupByExecutor(QueryContext queryContext, AggregationFunction[] a } } + /// Mirrors fixed-width result holders off-heap. Functions with a dedicated off-heap state holder (e.g. + /// `DISTINCT_COUNT_ULL`) supply it through [AggregationFunction#createOffHeapGroupByResultHolder] and are + /// consulted first. For the rest, the holder type and default value are discovered through a + /// zero-capacity probe (aggregation functions choose both — createGroupByResultHolder must stay side-effect-free + /// for the probe to be safe), and any non-fixed-width holder (object holders, dummy + /// holders, custom implementations) is recreated on-heap with the real initial capacity. Off-heap holders are + /// registered on the resource tracker, which releases them when the group key generator is closed. + private static GroupByResultHolder createOffHeapCapableResultHolder(AggregationFunction function, + int initialCapacity, int maxCapacity, ResourceTrackingGroupKeyGenerator resourceTracker) { + GroupByResultHolder offHeapStateHolder = function.createOffHeapGroupByResultHolder(initialCapacity, maxCapacity); + if (offHeapStateHolder != null) { + Preconditions.checkState(offHeapStateHolder instanceof AutoCloseable, + "createOffHeapGroupByResultHolder must return an AutoCloseable holder, got: %s", + offHeapStateHolder.getClass().getName()); + resourceTracker.register((AutoCloseable) offHeapStateHolder); + return offHeapStateHolder; + } + GroupByResultHolder probe = function.createGroupByResultHolder(0, maxCapacity); + GroupByResultHolder holder; + if (probe.getClass() == DoubleGroupByResultHolder.class) { + holder = new OffHeapDoubleGroupByResultHolder(initialCapacity, maxCapacity, + ((DoubleGroupByResultHolder) probe).getDefaultValue()); + } else if (probe.getClass() == LongGroupByResultHolder.class) { + holder = new OffHeapLongGroupByResultHolder(initialCapacity, maxCapacity, + ((LongGroupByResultHolder) probe).getDefaultValue()); + } else if (probe.getClass() == IntGroupByResultHolder.class) { + holder = new OffHeapIntGroupByResultHolder(initialCapacity, maxCapacity, + ((IntGroupByResultHolder) probe).getDefaultValue()); + } else { + return function.createGroupByResultHolder(initialCapacity, maxCapacity); + } + resourceTracker.register((AutoCloseable) holder); + return holder; + } + /// Retrieve the sizes of GroupBy expressions from IN an EQ predicates found in the filter context, if available. /// 1. If the filter context is null or lacks GroupBy expressions, return null. /// 2. Ensure the top-level filter context consists solely of AND-type filters; other types for example OR we cannot 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 32fdf9dc214b..e3dcbc3cbcd8 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 @@ -35,6 +35,10 @@ import org.apache.pinot.core.operator.BaseProjectOperator; import org.apache.pinot.core.operator.ColumnContext; import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapBytesGroupIdMap; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapGroupByUtils; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapIntGroupIdMap; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapLongGroupIdMap; import org.apache.pinot.segment.spi.datasource.DataSource; import org.apache.pinot.segment.spi.index.reader.Dictionary; import org.roaringbitmap.PeekableIntIterator; @@ -74,6 +78,9 @@ public class DictionaryBasedGroupKeyGenerator implements GroupKeyGenerator { private static final int INITIAL_MAP_SIZE = (int) ((1 << 9) * 0.75f); private static final int MAX_CACHING_MAP_SIZE = (int) ((1 << 20) * 0.75f); private static final int MAX_DICTIONARY_INTERN_TABLE_SIZE = 10000; + // Initial-size hint for off-heap key tables: bounded so sparse/filtered queries do not over-allocate, while + // high-cardinality queries reach their final size in a few cheap doublings + private static final int OFF_HEAP_INITIAL_ENTRIES_HINT = 8192; @VisibleForTesting static final ThreadLocal THREAD_LOCAL_INT_MAP = ThreadLocal.withInitial(IntGroupIdMap::new); @@ -123,6 +130,14 @@ public class DictionaryBasedGroupKeyGenerator implements GroupKeyGenerator { public DictionaryBasedGroupKeyGenerator(BaseProjectOperator projectOperator, ExpressionContext[] groupByExpressions, int numGroupsLimit, int arrayBasedThreshold, boolean nullHandlingEnabled, @Nullable Map groupByExpressionSizesFromPredicates) { + this(projectOperator, groupByExpressions, numGroupsLimit, arrayBasedThreshold, + nullHandlingEnabled, groupByExpressionSizesFromPredicates, false); + } + + public DictionaryBasedGroupKeyGenerator(BaseProjectOperator projectOperator, + ExpressionContext[] groupByExpressions, int numGroupsLimit, int arrayBasedThreshold, + boolean nullHandlingEnabled, @Nullable Map groupByExpressionSizesFromPredicates, + boolean offHeap) { _groupByExpressions = groupByExpressions; _numGroupByExpressions = groupByExpressions.length; @@ -197,16 +212,24 @@ public DictionaryBasedGroupKeyGenerator(BaseProjectOperator projectOperator, if (longOverflow) { // ArrayMapBasedHolder _globalGroupIdUpperBound = cappedNumGroupsLimit; - Object2IntOpenHashMap groupIdMap = THREAD_LOCAL_INT_ARRAY_MAP.get(); - clearAndTrim(groupIdMap); - _rawKeyHolder = new ArrayMapBasedHolder(groupIdMap); + if (offHeap) { + _rawKeyHolder = new OffHeapArrayMapBasedHolder(); + } else { + Object2IntOpenHashMap groupIdMap = THREAD_LOCAL_INT_ARRAY_MAP.get(); + clearAndTrim(groupIdMap); + _rawKeyHolder = new ArrayMapBasedHolder(groupIdMap); + } } else { if (cardinalityProduct > Integer.MAX_VALUE) { // LongMapBasedHolder _globalGroupIdUpperBound = cappedNumGroupsLimit; - Long2IntOpenHashMap groupIdMap = THREAD_LOCAL_LONG_MAP.get(); - clearAndTrim(groupIdMap); - _rawKeyHolder = new LongMapBasedHolder(groupIdMap); + if (offHeap) { + _rawKeyHolder = new OffHeapLongMapBasedHolder(); + } else { + Long2IntOpenHashMap groupIdMap = THREAD_LOCAL_LONG_MAP.get(); + clearAndTrim(groupIdMap); + _rawKeyHolder = new LongMapBasedHolder(groupIdMap); + } } else { _globalGroupIdUpperBound = (int) Math.min(cardinalityProduct, cappedNumGroupsLimit); // arrayBaseHolder fails with ArrayIndexOutOfBoundsException if numGroupsLimit < cardinalityProduct @@ -219,9 +242,13 @@ public DictionaryBasedGroupKeyGenerator(BaseProjectOperator projectOperator, if (cardinalityProduct > arrayBasedThreshold || numGroupsLimit < cardinalityProduct || optimizedGroupCountUpperBound < cardinalityProduct) { // IntMapBasedHolder - IntGroupIdMap groupIdMap = THREAD_LOCAL_INT_MAP.get(); - groupIdMap.clearAndTrim(); - _rawKeyHolder = new IntMapBasedHolder(groupIdMap); + if (offHeap) { + _rawKeyHolder = new OffHeapIntMapBasedHolder(); + } else { + IntGroupIdMap groupIdMap = THREAD_LOCAL_INT_MAP.get(); + groupIdMap.clearAndTrim(); + _rawKeyHolder = new IntMapBasedHolder(groupIdMap); + } } else { _rawKeyHolder = new ArrayBasedHolder(); } @@ -1077,6 +1104,239 @@ private Object[] getKeys(IntArray rawKey) { return groupKeys; } + /// Off-heap variant of [IntMapBasedHolder]: same int raw keys and 8-byte slot layout, but the key table + /// lives in direct memory. Released through [#close()], never cached across queries. + private class OffHeapIntMapBasedHolder implements RawKeyHolder { + // Unlike the LongMap/ArrayMap tiers (whose upper bound is just numGroupsLimit), this tier's upper bound is the + // exact dict-cardinality product, so pre-sizing to it (capped) is measurement-backed: JMH DICT_INT (80K groups, + // 4 segments) shows the expand chain costs ~2.3ms/query (+22% latency) when starting from the small hint, while + // the sparse-shape downside is one bounded bulk zero-fill (<=8MB, ~0.2ms) — and predicate-derived bounds + // already shrink the upper bound for selective IN/EQ filters. + private final OffHeapIntGroupIdMap _groupIdMap = + new OffHeapIntGroupIdMap(Math.min(_globalGroupIdUpperBound, 1 << 20)); + + @Override + public void processSingleValue(int numDocs, int[] outGroupIds) { + if (_numGroupByExpressions == 1) { + int[] dictIds = _singleValueDictIds[0]; + for (int i = 0; i < numDocs; i++) { + outGroupIds[i] = _groupIdMap.getGroupId(dictIds[i], _globalGroupIdUpperBound); + } + } else { + for (int i = 0; i < numDocs; i++) { + int rawKey = 0; + for (int j = _numGroupByExpressions - 1; j >= 0; j--) { + rawKey = rawKey * _cardinalities[j] + _singleValueDictIds[j][i]; + } + outGroupIds[i] = _groupIdMap.getGroupId(rawKey, _globalGroupIdUpperBound); + } + } + } + + @Override + public void processMultiValue(int numDocs, int[][] outGroupIds) { + for (int i = 0; i < numDocs; i++) { + int[] groupIds = getIntRawKeys(i); + int length = groupIds.length; + for (int j = 0; j < length; j++) { + groupIds[j] = _groupIdMap.getGroupId(groupIds[j], _globalGroupIdUpperBound); + } + outGroupIds[i] = groupIds; + } + } + + @Override + public int getGroupIdUpperBound() { + return _groupIdMap.size(); + } + + @Override + public Iterator getGroupKeys() { + return new Iterator<>() { + private final Iterator _iterator = _groupIdMap.iterator(); + private final GroupKey _groupKey = new GroupKey(); + + @Override + public boolean hasNext() { + return _iterator.hasNext(); + } + + @Override + public GroupKey next() { + OffHeapIntGroupIdMap.Entry entry = _iterator.next(); + _groupKey._groupId = entry._groupId; + _groupKey._keys = getKeys(entry._rawKey); + return _groupKey; + } + + @Override + public void remove() { + throw new UnsupportedOperationException(); + } + }; + } + + @Override + public int getNumKeys() { + return _groupIdMap.size(); + } + + @Override + public void close() { + _groupIdMap.close(); + } + } + + /// Off-heap variant of [LongMapBasedHolder]: same long raw keys, key table in direct memory. + private class OffHeapLongMapBasedHolder implements RawKeyHolder { + private final OffHeapLongGroupIdMap _groupIdMap = + new OffHeapLongGroupIdMap(Math.min(_globalGroupIdUpperBound, OFF_HEAP_INITIAL_ENTRIES_HINT)); + + @Override + public void processSingleValue(int numDocs, int[] outGroupIds) { + for (int i = 0; i < numDocs; i++) { + long rawKey = 0L; + for (int j = _numGroupByExpressions - 1; j >= 0; j--) { + rawKey = rawKey * _cardinalities[j] + _singleValueDictIds[j][i]; + } + outGroupIds[i] = _groupIdMap.getGroupId(rawKey, _globalGroupIdUpperBound); + } + } + + @Override + public void processMultiValue(int numDocs, int[][] outGroupIds) { + for (int i = 0; i < numDocs; i++) { + long[] rawKeys = getLongRawKeys(i); + int length = rawKeys.length; + int[] groupIds = new int[length]; + for (int j = 0; j < length; j++) { + groupIds[j] = _groupIdMap.getGroupId(rawKeys[j], _globalGroupIdUpperBound); + } + outGroupIds[i] = groupIds; + } + } + + @Override + public int getGroupIdUpperBound() { + return _groupIdMap.size(); + } + + @Override + public Iterator getGroupKeys() { + return new Iterator<>() { + private final Iterator _iterator = _groupIdMap.iterator(); + private final GroupKey _groupKey = new GroupKey(); + + @Override + public boolean hasNext() { + return _iterator.hasNext(); + } + + @Override + public GroupKey next() { + OffHeapLongGroupIdMap.Entry entry = _iterator.next(); + _groupKey._groupId = entry._groupId; + _groupKey._keys = getKeys(entry._rawKey); + return _groupKey; + } + + @Override + public void remove() { + throw new UnsupportedOperationException(); + } + }; + } + + @Override + public int getNumKeys() { + return _groupIdMap.size(); + } + + @Override + public void close() { + _groupIdMap.close(); + } + } + + /// Off-heap variant of [ArrayMapBasedHolder]: the dictionary ids of all group-by columns are packed into a + /// fixed-width byte key (4 bytes per column) held in a direct-memory two-part hash table. Unlike the on-heap + /// variant, the per-row lookup is allocation-free (no `IntArray` per row). + private class OffHeapArrayMapBasedHolder implements RawKeyHolder { + private final OffHeapBytesGroupIdMap _groupIdMap = + new OffHeapBytesGroupIdMap(Math.min(_globalGroupIdUpperBound, OFF_HEAP_INITIAL_ENTRIES_HINT)); + private final byte[] _keyScratch = new byte[_numGroupByExpressions * Integer.BYTES]; + private final int[] _dictIdScratch = new int[_numGroupByExpressions]; + + @Override + public void processSingleValue(int numDocs, int[] outGroupIds) { + for (int i = 0; i < numDocs; i++) { + for (int j = 0; j < _numGroupByExpressions; j++) { + _dictIdScratch[j] = _singleValueDictIds[j][i]; + } + int keyLength = OffHeapGroupByUtils.packInts(_dictIdScratch, _numGroupByExpressions, _keyScratch); + outGroupIds[i] = _groupIdMap.getGroupId(_keyScratch, 0, keyLength, _globalGroupIdUpperBound); + } + } + + @Override + public void processMultiValue(int numDocs, int[][] outGroupIds) { + for (int i = 0; i < numDocs; i++) { + IntArray[] rawKeys = getIntArrayRawKeys(i); + int length = rawKeys.length; + int[] groupIds = new int[length]; + for (int j = 0; j < length; j++) { + int keyLength = OffHeapGroupByUtils.packInts(rawKeys[j]._elements, _numGroupByExpressions, _keyScratch); + groupIds[j] = _groupIdMap.getGroupId(_keyScratch, 0, keyLength, _globalGroupIdUpperBound); + } + outGroupIds[i] = groupIds; + } + } + + @Override + public int getGroupIdUpperBound() { + return _groupIdMap.size(); + } + + @Override + public Iterator getGroupKeys() { + return new Iterator<>() { + private final IntArray _rawKey = new IntArray(new int[_numGroupByExpressions]); + private final GroupKey _groupKey = new GroupKey(); + private int _currentGroupId; + + @Override + public boolean hasNext() { + return _currentGroupId < _groupIdMap.size(); + } + + @Override + public GroupKey next() { + _groupIdMap.readKey(_currentGroupId, _keyScratch, 0); + OffHeapGroupByUtils.unpackInts(_keyScratch, _numGroupByExpressions, _rawKey._elements); + _groupKey._groupId = _currentGroupId; + _groupKey._keys = getKeys(_rawKey); + _currentGroupId++; + return _groupKey; + } + + @Override + public void remove() { + throw new UnsupportedOperationException(); + } + }; + } + + @Override + public int getNumKeys() { + return _groupIdMap.size(); + } + + @Override + public void close() { + _groupIdMap.close(); + } + } + /// Fast int-to-int hashmap with [#INVALID_ID] as the default return value. /// /// Different from [it.unimi.dsi.fastutil.ints.Int2IntOpenHashMap], this map uses one single array to store diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DoubleGroupByResultHolder.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DoubleGroupByResultHolder.java index 8cf621bc037e..253415fa6a4b 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DoubleGroupByResultHolder.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DoubleGroupByResultHolder.java @@ -46,6 +46,11 @@ public DoubleGroupByResultHolder(int initialCapacity, int maxCapacity, double de } } + /// Returns the default value for un-initialized results. Used to mirror this holder off-heap. + double getDefaultValue() { + return _defaultValue; + } + @Override public void ensureCapacity(int capacity) { Preconditions.checkArgument(capacity <= _maxCapacity); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/IntGroupByResultHolder.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/IntGroupByResultHolder.java index ae5c138f42a9..00749a9fb62e 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/IntGroupByResultHolder.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/IntGroupByResultHolder.java @@ -47,6 +47,11 @@ public IntGroupByResultHolder(int initialCapacity, int maxCapacity, int defaultV } } + /// Returns the default value for un-initialized results. Used to mirror this holder off-heap. + int getDefaultValue() { + return _defaultValue; + } + @Override public void ensureCapacity(int capacity) { Preconditions.checkArgument(capacity <= _maxCapacity); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/LongGroupByResultHolder.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/LongGroupByResultHolder.java index ea345957d515..5d2ec7d45083 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/LongGroupByResultHolder.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/LongGroupByResultHolder.java @@ -47,6 +47,11 @@ public LongGroupByResultHolder(int initialCapacity, int maxCapacity, long defaul } } + /// Returns the default value for un-initialized results. Used to mirror this holder off-heap. + long getDefaultValue() { + return _defaultValue; + } + @Override public void ensureCapacity(int capacity) { Preconditions.checkArgument(capacity <= _maxCapacity); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryMultiColumnGroupKeyGenerator.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryMultiColumnGroupKeyGenerator.java index 1b6cf547d43c..104e313713c9 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryMultiColumnGroupKeyGenerator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryMultiColumnGroupKeyGenerator.java @@ -31,6 +31,8 @@ import org.apache.pinot.core.operator.BaseProjectOperator; import org.apache.pinot.core.operator.ColumnContext; import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapBytesGroupIdMap; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapGroupByUtils; import org.apache.pinot.core.query.aggregation.groupby.utils.ValueToIdMap; import org.apache.pinot.core.query.aggregation.groupby.utils.ValueToIdMapFactory; import org.apache.pinot.segment.spi.index.reader.Dictionary; @@ -66,9 +68,21 @@ public class NoDictionaryMultiColumnGroupKeyGenerator implements GroupKeyGenerat private final boolean _nullHandlingEnabled; private final int _globalGroupIdUpperBound; + // Off-heap mode (see the offHeap constructor param): the packed per-column key ids (4 bytes per column, + // including ID_FOR_NULL components) go to a direct-memory key table instead of _groupKeyMap, which is null. + private final OffHeapBytesGroupIdMap _offHeapGroupKeyMap; + private final byte[] _keyScratch; + public NoDictionaryMultiColumnGroupKeyGenerator(BaseProjectOperator projectOperator, ExpressionContext[] groupByExpressions, int numGroupsLimit, boolean nullHandlingEnabled, Map groupByExpressionSizesFromPredicates) { + this(projectOperator, groupByExpressions, numGroupsLimit, nullHandlingEnabled, + groupByExpressionSizesFromPredicates, false); + } + + public NoDictionaryMultiColumnGroupKeyGenerator(BaseProjectOperator projectOperator, + ExpressionContext[] groupByExpressions, int numGroupsLimit, boolean nullHandlingEnabled, + Map groupByExpressionSizesFromPredicates, boolean offHeap) { _groupByExpressions = groupByExpressions; _numGroupByExpressions = groupByExpressions.length; _storedTypes = new DataType[_numGroupByExpressions]; @@ -106,8 +120,16 @@ public NoDictionaryMultiColumnGroupKeyGenerator(BaseProjectOperator projectOp _isSingleValueExpressions[i] = columnContext.isSingleValue(); } - _groupKeyMap = new Object2IntOpenHashMap<>(); - _groupKeyMap.defaultReturnValue(INVALID_ID); + if (offHeap) { + _groupKeyMap = null; + _offHeapGroupKeyMap = new OffHeapBytesGroupIdMap(Math.min(numGroupsLimit, 8192)); + _keyScratch = new byte[_numGroupByExpressions * Integer.BYTES]; + } else { + _groupKeyMap = new Object2IntOpenHashMap<>(); + _groupKeyMap.defaultReturnValue(INVALID_ID); + _offHeapGroupKeyMap = null; + _keyScratch = null; + } _numGroupsLimit = numGroupsLimit; _globalGroupIdUpperBound = canOptimizeGroupByUpperBound ? optimizedGroupByUpperBound : numGroupsLimit; } @@ -168,7 +190,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { // note that we are mutating its backing array for memory efficiency FixedIntArray flyweightKey = new FixedIntArray(keyValues); for (int row = 0; row < numDocs; row++) { - int numGroups = _groupKeyMap.size(); + int numGroups = getNumGroupsInternal(); boolean hasInvalidKeyValue = false; if (numGroups < _numGroupsLimit) { for (int col = 0; col < _numGroupByExpressions; col++) { @@ -376,12 +398,23 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { @Override public int getCurrentGroupKeyUpperBound() { - return _groupKeyMap.size(); + return getNumGroupsInternal(); } @Override public Iterator getGroupKeys() { - return new GroupKeyIterator(); + return _offHeapGroupKeyMap != null ? new OffHeapGroupKeyIterator() : new GroupKeyIterator(); + } + + @Override + public void close() { + if (_offHeapGroupKeyMap != null) { + _offHeapGroupKeyMap.close(); + } + } + + private int getNumGroupsInternal() { + return _offHeapGroupKeyMap != null ? _offHeapGroupKeyMap.size() : _groupKeyMap.size(); } /// Helper method to get or create group-id for a group key. @@ -389,6 +422,10 @@ public Iterator getGroupKeys() { /// @param keyList Group key, that is a list of objects to be grouped /// @return Group id private int getGroupIdForKey(FixedIntArray keyList) { + if (_offHeapGroupKeyMap != null) { + int keyLength = OffHeapGroupByUtils.packInts(keyList.elements(), _numGroupByExpressions, _keyScratch); + return _offHeapGroupKeyMap.getGroupId(_keyScratch, 0, keyLength, _numGroupsLimit); + } int numGroups = _groupKeyMap.size(); if (numGroups < _numGroupsLimit) { return _groupKeyMap.computeIfAbsent(keyList, k -> numGroups); @@ -422,7 +459,36 @@ private void getGroupIdsForKeyHelper(int[][] keysList, int[] groupKeyIds, int le @Override public int getNumKeys() { - return _groupKeyMap.size(); + return getNumGroupsInternal(); + } + + /// Iterator over the dense group ids of the off-heap key table: the packed per-column key ids are read back and + /// rebuilt into values through the same dictionaries / on-the-fly dictionaries as the on-heap iterator. + private class OffHeapGroupKeyIterator implements Iterator { + private final GroupKey _groupKey = new GroupKey(); + private final int[] _keyIds = new int[_numGroupByExpressions]; + private final FixedIntArray _flyweightKey = new FixedIntArray(_keyIds); + private int _groupId; + + @Override + public boolean hasNext() { + return _groupId < _offHeapGroupKeyMap.size(); + } + + @Override + public GroupKey next() { + _offHeapGroupKeyMap.readKey(_groupId, _keyScratch, 0); + OffHeapGroupByUtils.unpackInts(_keyScratch, _numGroupByExpressions, _keyIds); + _groupKey._groupId = _groupId; + _groupKey._keys = buildKeysFromIds(_flyweightKey); + _groupId++; + return _groupKey; + } + + @Override + public void remove() { + throw new UnsupportedOperationException(); + } } /// Iterator for [GroupKey]. diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionarySingleColumnGroupKeyGenerator.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionarySingleColumnGroupKeyGenerator.java index 98f83c5f846c..cbf406735ed8 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionarySingleColumnGroupKeyGenerator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionarySingleColumnGroupKeyGenerator.java @@ -30,6 +30,7 @@ import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap; import it.unimi.dsi.fastutil.objects.ObjectIterator; import java.math.BigDecimal; +import java.nio.charset.StandardCharsets; import java.util.Iterator; import java.util.Map; import javax.annotation.Nullable; @@ -38,7 +39,12 @@ import org.apache.pinot.core.operator.BaseProjectOperator; import org.apache.pinot.core.operator.ColumnContext; import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapBytesGroupIdMap; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapGroupByUtils; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapIntGroupIdMap; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapLongGroupIdMap; import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.utils.BigDecimalUtils; import org.apache.pinot.spi.utils.ByteArray; import org.roaringbitmap.RoaringBitmap; @@ -50,6 +56,11 @@ /// null value, which is what a null row is physically stored as. The object-keyed maps hold the null key directly, /// while the primitive-keyed maps cannot, so the null group id is tracked beside them. Both the single-value and the /// multi-value key paths recognize nulls. +/// +/// Off-heap mode caveat for STRING keys: keys are grouped by their UTF-8 encoding (byte-identical to +/// `String#getBytes(UTF_8)`), so two strings that differ only in unpaired surrogates collapse into one group +/// (both encode to `'?'`), and the emitted group key is the re-decoded string — whereas the on-heap +/// `Object2IntOpenHashMap` keeps such (malformed) strings distinct. Valid strings are unaffected. @SuppressWarnings({"rawtypes", "unchecked"}) public class NoDictionarySingleColumnGroupKeyGenerator implements GroupKeyGenerator { private final ExpressionContext _groupByExpression; @@ -62,19 +73,69 @@ public class NoDictionarySingleColumnGroupKeyGenerator implements GroupKeyGenera private Integer _groupIdForNullValue; private int _numGroups; + // Off-heap mode (see the offHeap constructor param): exactly one of the two maps below is non-null and + // _groupKeyMap is null. Null keys never enter the off-heap maps: for every type the null group is tracked via + // _groupIdForNullValue, with _nullGroupIdMapSize recording the map size at the moment the null group was + // assigned, so map-internal ids at/after that point shift up by one to keep the global ids dense in assignment + // order (matching the on-heap _numGroups counter semantics). + private final OffHeapIntGroupIdMap _offHeapIntKeyMap; + private final OffHeapLongGroupIdMap _offHeapLongKeyMap; + private final OffHeapBytesGroupIdMap _offHeapBytesKeyMap; + private int _nullGroupIdMapSize = -1; + private byte[] _stringEncodeScratch = new byte[64]; + public NoDictionarySingleColumnGroupKeyGenerator(BaseProjectOperator projectOperator, ExpressionContext groupByExpression, int numGroupsLimit, boolean nullHandlingEnabled, @Nullable Map groupByExpressionSizesFromPredicates) { + this(projectOperator, groupByExpression, numGroupsLimit, nullHandlingEnabled, + groupByExpressionSizesFromPredicates, false); + } + + public NoDictionarySingleColumnGroupKeyGenerator(BaseProjectOperator projectOperator, + ExpressionContext groupByExpression, int numGroupsLimit, boolean nullHandlingEnabled, + @Nullable Map groupByExpressionSizesFromPredicates, boolean offHeap) { _groupByExpression = groupByExpression; ColumnContext columnContext = projectOperator.getResultColumnContext(groupByExpression); _storedType = columnContext.getDataType().getStoredType(); - _groupKeyMap = createGroupKeyMap(_storedType); if (groupByExpressionSizesFromPredicates != null) { Integer size = groupByExpressionSizesFromPredicates.get(groupByExpression); _globalGroupIdUpperBound = size != null ? Math.min(size, numGroupsLimit) : numGroupsLimit; } else { _globalGroupIdUpperBound = numGroupsLimit; } + if (offHeap) { + _groupKeyMap = null; + switch (_storedType) { + case INT: + case FLOAT: + // 32-bit keys go to the 8-byte-slot int map for better probe locality (FLOAT keys are stored as + // floatToIntBits, which never produces -1 — all NaNs canonicalize — and INT -1 is held out-of-band) + _offHeapIntKeyMap = new OffHeapIntGroupIdMap(Math.min(_globalGroupIdUpperBound, 8192)); + _offHeapLongKeyMap = null; + _offHeapBytesKeyMap = null; + break; + case LONG: + case DOUBLE: + _offHeapIntKeyMap = null; + _offHeapLongKeyMap = new OffHeapLongGroupIdMap(Math.min(_globalGroupIdUpperBound, 8192)); + _offHeapBytesKeyMap = null; + break; + case BIG_DECIMAL: + case STRING: + case BYTES: + _offHeapIntKeyMap = null; + _offHeapLongKeyMap = null; + _offHeapBytesKeyMap = new OffHeapBytesGroupIdMap(Math.min(_globalGroupIdUpperBound, 8192)); + break; + default: + throw new IllegalStateException("Illegal data type for no-dictionary key generator: " + _storedType); + } + } else { + _groupKeyMap = createGroupKeyMap(_storedType); + _offHeapIntKeyMap = null; + _offHeapLongKeyMap = null; + _offHeapBytesKeyMap = null; + } _nullHandlingEnabled = nullHandlingEnabled; _isSingleValueExpression = columnContext.isSingleValue(); } @@ -367,11 +428,23 @@ private static boolean isNull(@Nullable RoaringBitmap nullBitmap, int row) { /// primitive-keyed maps cannot hold a null key, so the group count has to come from the id counter itself. @Override public int getCurrentGroupKeyUpperBound() { + if (_groupKeyMap == null) { + return getOffHeapNumGroups(); + } return _numGroups; } @Override public Iterator getGroupKeys() { + if (_offHeapIntKeyMap != null) { + return new OffHeapIntKeyIterator(); + } + if (_offHeapLongKeyMap != null) { + return new OffHeapLongKeyIterator(); + } + if (_offHeapBytesKeyMap != null) { + return new OffHeapBytesKeyIterator(); + } return switch (_storedType) { case INT -> new IntGroupKeyIterator((Int2IntOpenHashMap) _groupKeyMap, _groupIdForNullValue); case LONG -> new LongGroupKeyIterator((Long2IntOpenHashMap) _groupKeyMap, _groupIdForNullValue); @@ -383,6 +456,9 @@ public Iterator getGroupKeys() { } private int getKeyForNullValue() { + if (_groupKeyMap == null) { + return getOffHeapNullGroupId(); + } if (_groupIdForNullValue != null) { return _groupIdForNullValue; } @@ -395,10 +471,67 @@ private int getKeyForNullValue() { @Override public int getNumKeys() { + if (_groupKeyMap == null) { + return getOffHeapNumGroups(); + } return _numGroups; } + @Override + public void close() { + if (_offHeapIntKeyMap != null) { + _offHeapIntKeyMap.close(); + } + if (_offHeapLongKeyMap != null) { + _offHeapLongKeyMap.close(); + } + if (_offHeapBytesKeyMap != null) { + _offHeapBytesKeyMap.close(); + } + } + + private int getOffHeapMapSize() { + if (_offHeapIntKeyMap != null) { + return _offHeapIntKeyMap.size(); + } + return _offHeapLongKeyMap != null ? _offHeapLongKeyMap.size() : _offHeapBytesKeyMap.size(); + } + + private int getOffHeapNumGroups() { + return getOffHeapMapSize() + (_groupIdForNullValue != null ? 1 : 0); + } + + // Upper bound to pass to the off-heap map: reserve one group id for the null group once it is assigned + private int getOffHeapMapUpperBound() { + return _groupIdForNullValue != null ? _globalGroupIdUpperBound - 1 : _globalGroupIdUpperBound; + } + + // Map-internal ids assigned at/after the null group shift up by one to keep global ids dense in assignment order + private int toGlobalId(int mapId) { + return mapId != INVALID_ID && _nullGroupIdMapSize >= 0 && mapId >= _nullGroupIdMapSize ? mapId + 1 : mapId; + } + + private int getOffHeapNullGroupId() { + if (_groupIdForNullValue != null) { + return _groupIdForNullValue; + } + if (getOffHeapNumGroups() < _globalGroupIdUpperBound) { + _nullGroupIdMapSize = getOffHeapMapSize(); + // The null group takes the next dense id: all existing map ids stay put, later map ids shift up by one + _groupIdForNullValue = _nullGroupIdMapSize; + return _groupIdForNullValue; + } + return INVALID_ID; + } + + private int getOffHeapKeyForBytes(byte[] bytes) { + return toGlobalId(_offHeapBytesKeyMap.getGroupId(bytes, 0, bytes.length, getOffHeapMapUpperBound())); + } + private int getKeyForValue(int value) { + if (_offHeapIntKeyMap != null) { + return toGlobalId(_offHeapIntKeyMap.getGroupId(value, getOffHeapMapUpperBound())); + } Int2IntMap map = (Int2IntMap) _groupKeyMap; int groupId = map.get(value); if (groupId == INVALID_ID && _numGroups < _globalGroupIdUpperBound) { @@ -409,6 +542,9 @@ private int getKeyForValue(int value) { } private int getKeyForValue(long value) { + if (_offHeapLongKeyMap != null) { + return toGlobalId(_offHeapLongKeyMap.getGroupId(value, getOffHeapMapUpperBound())); + } Long2IntMap map = (Long2IntMap) _groupKeyMap; int groupId = map.get(value); if (groupId == INVALID_ID && _numGroups < _globalGroupIdUpperBound) { @@ -419,6 +555,10 @@ private int getKeyForValue(long value) { } private int getKeyForValue(float value) { + if (_offHeapIntKeyMap != null) { + // floatToIntBits (not raw) matches fastutil semantics: all NaNs collapse, +0.0f and -0.0f stay distinct + return toGlobalId(_offHeapIntKeyMap.getGroupId(Float.floatToIntBits(value), getOffHeapMapUpperBound())); + } Float2IntMap map = (Float2IntMap) _groupKeyMap; int groupId = map.get(value); if (groupId == INVALID_ID && _numGroups < _globalGroupIdUpperBound) { @@ -429,6 +569,10 @@ private int getKeyForValue(float value) { } private int getKeyForValue(double value) { + if (_offHeapLongKeyMap != null) { + // doubleToLongBits (not raw) matches fastutil semantics: all NaNs collapse, +0.0 and -0.0 stay distinct + return toGlobalId(_offHeapLongKeyMap.getGroupId(Double.doubleToLongBits(value), getOffHeapMapUpperBound())); + } Double2IntMap map = (Double2IntMap) _groupKeyMap; int groupId = map.get(value); if (groupId == INVALID_ID && _numGroups < _globalGroupIdUpperBound) { @@ -439,6 +583,13 @@ private int getKeyForValue(double value) { } private int getKeyForValue(BigDecimal value) { + if (_offHeapBytesKeyMap != null) { + if (value == null) { + return getOffHeapNullGroupId(); + } + // The serialized form preserves scale and unscaled value, so byte equality == BigDecimal#equals + return getOffHeapKeyForBytes(BigDecimalUtils.serialize(value)); + } Object2IntMap map = (Object2IntMap) _groupKeyMap; int groupId = map.getInt(value); if (groupId == INVALID_ID && _numGroups < _globalGroupIdUpperBound) { @@ -449,6 +600,16 @@ private int getKeyForValue(BigDecimal value) { } private int getKeyForValue(String value) { + if (_offHeapBytesKeyMap != null) { + if (value == null) { + return getOffHeapNullGroupId(); + } + int maxLength = value.length() * 3; + byte[] scratch = OffHeapGroupByUtils.ensureByteCapacity(_stringEncodeScratch, maxLength); + _stringEncodeScratch = scratch; + int length = OffHeapGroupByUtils.encodeUtf8(value, scratch); + return toGlobalId(_offHeapBytesKeyMap.getGroupId(scratch, 0, length, getOffHeapMapUpperBound())); + } Object2IntMap map = (Object2IntMap) _groupKeyMap; int groupId = map.getInt(value); if (groupId == INVALID_ID && _numGroups < _globalGroupIdUpperBound) { @@ -459,6 +620,12 @@ private int getKeyForValue(String value) { } private int getKeyForValue(ByteArray value) { + if (_offHeapBytesKeyMap != null) { + if (value == null) { + return getOffHeapNullGroupId(); + } + return getOffHeapKeyForBytes(value.getBytes()); + } Object2IntMap map = (Object2IntMap) _groupKeyMap; int groupId = map.getInt(value); if (groupId == INVALID_ID && _numGroups < _globalGroupIdUpperBound) { @@ -468,6 +635,138 @@ private int getKeyForValue(ByteArray value) { return groupId; } + /// Iterator for the off-heap int-key map (INT/FLOAT stored types). Emits the null group (if assigned) first, + /// then the map entries with their map-internal ids converted to global ids. + private class OffHeapIntKeyIterator implements Iterator { + private final Iterator _iterator = _offHeapIntKeyMap.iterator(); + private final GroupKey _groupKey = new GroupKey(); + private boolean _nullValuePending = _groupIdForNullValue != null; + + @Override + public boolean hasNext() { + return _nullValuePending || _iterator.hasNext(); + } + + @Override + public GroupKey next() { + if (_nullValuePending) { + _groupKey._groupId = _groupIdForNullValue; + _groupKey._keys = new Object[]{null}; + _nullValuePending = false; + return _groupKey; + } + OffHeapIntGroupIdMap.Entry entry = _iterator.next(); + _groupKey._groupId = toGlobalId(entry._groupId); + _groupKey._keys = new Object[]{decodeIntKey(entry._rawKey)}; + return _groupKey; + } + + @Override + public void remove() { + throw new UnsupportedOperationException(); + } + } + + private Object decodeIntKey(int rawKey) { + switch (_storedType) { + case INT: + return rawKey; + case FLOAT: + return Float.intBitsToFloat(rawKey); + default: + throw new IllegalStateException(); + } + } + + /// Iterator for the off-heap long-key map (LONG/DOUBLE stored types). Emits the null group (if assigned) first, + /// then the map entries with their map-internal ids converted to global ids. + private class OffHeapLongKeyIterator implements Iterator { + private final Iterator _iterator = _offHeapLongKeyMap.iterator(); + private final GroupKey _groupKey = new GroupKey(); + private boolean _nullValuePending = _groupIdForNullValue != null; + + @Override + public boolean hasNext() { + return _nullValuePending || _iterator.hasNext(); + } + + @Override + public GroupKey next() { + if (_nullValuePending) { + _groupKey._groupId = _groupIdForNullValue; + _groupKey._keys = new Object[]{null}; + _nullValuePending = false; + return _groupKey; + } + OffHeapLongGroupIdMap.Entry entry = _iterator.next(); + _groupKey._groupId = toGlobalId(entry._groupId); + _groupKey._keys = new Object[]{decodeLongKey(entry._rawKey)}; + return _groupKey; + } + + @Override + public void remove() { + throw new UnsupportedOperationException(); + } + } + + private Object decodeLongKey(long rawKey) { + switch (_storedType) { + case LONG: + return rawKey; + case DOUBLE: + return Double.longBitsToDouble(rawKey); + default: + throw new IllegalStateException(); + } + } + + /// Iterator for the off-heap bytes-key map (STRING/BYTES/BIG_DECIMAL stored types). Emits the null group (if + /// assigned) first, then the dense map ids converted to global ids. + private class OffHeapBytesKeyIterator implements Iterator { + private final GroupKey _groupKey = new GroupKey(); + private boolean _nullValuePending = _groupIdForNullValue != null; + private int _mapId; + + @Override + public boolean hasNext() { + return _nullValuePending || _mapId < _offHeapBytesKeyMap.size(); + } + + @Override + public GroupKey next() { + if (_nullValuePending) { + _groupKey._groupId = _groupIdForNullValue; + _groupKey._keys = new Object[]{null}; + _nullValuePending = false; + return _groupKey; + } + byte[] keyBytes = _offHeapBytesKeyMap.getKey(_mapId); + _groupKey._groupId = toGlobalId(_mapId); + _groupKey._keys = new Object[]{decodeBytesKey(keyBytes)}; + _mapId++; + return _groupKey; + } + + @Override + public void remove() { + throw new UnsupportedOperationException(); + } + } + + private Object decodeBytesKey(byte[] bytes) { + switch (_storedType) { + case STRING: + return new String(bytes, StandardCharsets.UTF_8); + case BYTES: + return new ByteArray(bytes); + case BIG_DECIMAL: + return BigDecimalUtils.deserialize(bytes); + default: + throw new IllegalStateException(); + } + } + private static class IntGroupKeyIterator implements Iterator { final Iterator _iterator; final GroupKey _groupKey; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapBytesGroupIdMap.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapBytesGroupIdMap.java new file mode 100644 index 000000000000..c2f8eafa0fda --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapBytesGroupIdMap.java @@ -0,0 +1,521 @@ +/** + * 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.offheap; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import javax.annotation.concurrent.NotThreadSafe; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; + + +/// Off-heap hash map from arbitrary `byte[]` keys to dense int group ids, used as a replacement for +/// `Object2IntOpenHashMap` in group-by key generation. +/// +/// The design is modeled on DuckDB's two-part aggregate hash table: +///
    +///
  • Directory: an open-addressing table of 8-byte entries (power-of-two slot count, load factor 0.5, +/// linear probing). An entry of `0` means empty; otherwise the top 16 bits hold a salt (the top 16 bits of +/// the key's 64-bit hash) and the low 48 bits hold `payloadGlobalOffset + 1`. The `+1` keeps payload +/// offset 0 distinguishable from an empty slot. The salt is compared before touching the payload, so most probe +/// misses stay within the directory.
    • +///
  • Payload: append-only chunks (256KB each) storing one record per key: +/// `[long hash][int groupId][int keyLength][key bytes]` (16-byte header). Records never span chunks: a +/// record that does not fit in the remaining space of the current chunk starts a new chunk, and a record larger +/// than the normal chunk size gets a dedicated chunk of exactly the record size (placed at offset 0, with the +/// next record starting a fresh normal chunk). Because the 64-bit hash is stored in the payload, directory resize +/// never rehashes or compares key bytes.
    • +///
  • Id index: a growable off-heap long array mapping group id to payload global offset, supporting +/// dense-id reverse lookup ([#getKey], [#readKey], [#getKeyLength]) without an entry +/// iterator: ids are dense `0..size()-1` assigned in insertion order.
    • +///
+/// +/// All off-heap memory is allocated through [PinotDataBuffer#allocateDirect] and released by +/// [#close()], which is idempotent. Behavior of all other methods after `close()` is undefined. +/// +/// This class is not thread-safe: it is a per-query scratch structure intended to be used from a single +/// thread. +@NotThreadSafe +public class OffHeapBytesGroupIdMap implements AutoCloseable { + // Normal payload chunk size. Records never span chunks, and a record larger than this gets a dedicated chunk. + @VisibleForTesting + static final int CHUNK_SIZE = 1 << 18; + + private static final int MIN_NUM_SLOTS = 1024; + private static final int MAX_NUM_SLOTS = 1 << 30; + // Directory entry: 0 == empty; else (salt << 48) | (payloadGlobalOffset + 1). + private static final long SALT_MASK = 0xFFFF000000000000L; + private static final long OFFSET_MASK = 0x0000FFFFFFFFFFFFL; + // Payload record layout: [long hash][int groupId][int keyLength][key bytes]. + private static final int RECORD_HEADER_SIZE = 16; + private static final int GROUP_ID_OFFSET = 8; + private static final int KEY_LENGTH_OFFSET = 12; + private static final long MURMUR3_SEED = 0x9747b28cL; + private static final String DIRECTORY_DESCRIPTION = "OffHeapBytesGroupIdMap: directory"; + private static final String ID_INDEX_DESCRIPTION = "OffHeapBytesGroupIdMap: id index"; + private static final String CHUNK_DESCRIPTION = "OffHeapBytesGroupIdMap: payload chunk"; + // Reusable zero block for bulk zero-filling freshly allocated directories through the direct view + private static final byte[] ZERO_CHUNK = new byte[8192]; + + private final List _chunks = new ArrayList<>(); + // Absolute-indexed direct views for the per-row hot path (monomorphic, intrinsified ByteBuffer access instead + // of the PinotDataBuffer wrapper); a view is null when its buffer exceeds the 2GB view limit + private final List _chunkViews = new ArrayList<>(); + + private ByteBuffer _directoryView; + private PinotDataBuffer _directory; + private ByteBuffer _idIndexView; + private PinotDataBuffer _idIndex; + private int _numSlots; + private int _size; + // Chunk currently accepting appends; null until the first normal-size record, and reset to null after an + // oversized record so that no record is ever appended after one. + private PinotDataBuffer _currentChunk; + private int _currentChunkIndex; + private int _currentChunkOffset; + // Reusable scratch for bulk key comparison in matchRecord (grown on demand, contents transient) + private byte[] _compareScratch = new byte[64]; + private boolean _closed; + + public OffHeapBytesGroupIdMap(int expectedNumEntries) { + Preconditions.checkArgument(expectedNumEntries >= 0, "Invalid expectedNumEntries: %s", expectedNumEntries); + _numSlots = computeNumSlots(expectedNumEntries); + _directory = allocateZeroFilledDirectory(_numSlots); + _directoryView = OffHeapGroupByUtils.createView(_directory, (long) _numSlots << 3); + try { + // Sized to the directory's max fill (load factor 0.5) so both grow together. + _idIndex = OffHeapGroupByBufferPool.acquire((long) (_numSlots >> 1) << 3, ID_INDEX_DESCRIPTION); + _idIndexView = OffHeapGroupByUtils.createView(_idIndex, (long) (_numSlots >> 1) << 3); + } catch (Throwable t) { + closeBuffer(_directory); + if (_idIndex != null) { + closeBuffer(_idIndex); + } + throw t; + } + } + + /// Returns the dense group id for the given key: + ///
    + ///
  • If the key is already present, always returns its id (even when `size() == groupIdUpperBound`).
    • + ///
  • If the key is absent and `size() < groupIdUpperBound`, assigns the next dense id + /// (`size()`), inserts the key and returns the id.
    • + ///
  • If the key is absent and `size() >= groupIdUpperBound`, returns + /// [GroupKeyGenerator#INVALID_ID] without inserting.
    • + ///
+ public int getGroupId(byte[] key, int offset, int length, int groupIdUpperBound) { + long hash = murmurHash3X64Bit64(key, offset, length); + int mask = _numSlots - 1; + int slot = (int) (hash & mask); + long saltBits = hash & SALT_MASK; + ByteBuffer directoryView = _directoryView; + while (true) { + // While the view exists, slot offsets fit in an int (view size <= Integer.MAX_VALUE) + long entry = directoryView != null ? directoryView.getLong(slot << 3) : _directory.getLong((long) slot << 3); + if (entry == 0L) { + if (_size >= groupIdUpperBound) { + return GroupKeyGenerator.INVALID_ID; + } + return insert(hash, key, offset, length, slot); + } + if ((entry & SALT_MASK) == saltBits) { + int groupId = matchRecord((entry & OFFSET_MASK) - 1, hash, key, offset, length); + if (groupId != GroupKeyGenerator.INVALID_ID) { + return groupId; + } + } + slot = (slot + 1) & mask; + } + } + + /// Convenience variant of [int, int, int)][#getGroupId(byte[],] covering the full key array. + public int getGroupId(byte[] key, int groupIdUpperBound) { + return getGroupId(key, 0, key.length, groupIdUpperBound); + } + + /// Returns the number of keys in the map. Group ids are dense: `0..size()-1`. + public int size() { + return _size; + } + + /// Returns the length in bytes of the key with the given group id. The id must be within `[0, size())`. + public int getKeyLength(int groupId) { + // The id index is unchecked PinotDataBuffer memory: guard the dense-id contract with an assert + assert groupId >= 0 && groupId < _size : "groupId " + groupId + " out of bounds"; + long globalOffset = _idIndex.getLong((long) groupId << 3); + return _chunks.get((int) (globalOffset / CHUNK_SIZE)).getInt((globalOffset % CHUNK_SIZE) + KEY_LENGTH_OFFSET); + } + + /// Copies the key with the given group id into `dest` at `destOffset`. The id must be within + /// `[0, size())`, and `dest` must have at least [#getKeyLength(int)] bytes of room. + public void readKey(int groupId, byte[] dest, int destOffset) { + // See getKeyLength: unchecked buffer access, so guard the dense-id contract with an assert + assert groupId >= 0 && groupId < _size : "groupId " + groupId + " out of bounds"; + long globalOffset = _idIndex.getLong((long) groupId << 3); + PinotDataBuffer chunk = _chunks.get((int) (globalOffset / CHUNK_SIZE)); + long offsetInChunk = globalOffset % CHUNK_SIZE; + int keyLength = chunk.getInt(offsetInChunk + KEY_LENGTH_OFFSET); + chunk.copyTo(offsetInChunk + RECORD_HEADER_SIZE, dest, destOffset, keyLength); + } + + /// Returns a copy of the key with the given group id. The id must be within `[0, size())`. + public byte[] getKey(int groupId) { + byte[] key = new byte[getKeyLength(groupId)]; + readKey(groupId, key, 0); + return key; + } + + /// Returns the total off-heap memory held by the map (directory + id index + payload chunks). + public long getOffHeapMemoryBytes() { + long bytes = _directory.size() + _idIndex.size(); + for (PinotDataBuffer chunk : _chunks) { + bytes += chunk.size(); + } + return bytes; + } + + @Override + public void close() { + if (_closed) { + return; + } + _closed = true; + RuntimeException firstException = closeBufferQuietly(_directory, null); + firstException = closeBufferQuietly(_idIndex, firstException); + for (PinotDataBuffer chunk : _chunks) { + firstException = closeBufferQuietly(chunk, firstException); + } + // Null every buffer and view so any use-after-close (or a second release of a pooled buffer) fails loudly + // with an NPE instead of silently aliasing memory that the pool may have handed to another query + _directory = null; + _directoryView = null; + _idIndex = null; + _idIndexView = null; + _chunks.clear(); + _chunkViews.clear(); + _currentChunk = null; + if (firstException != null) { + throw firstException; + } + } + + @VisibleForTesting + long getPayloadGlobalOffset(int groupId) { + ByteBuffer idIndexView = _idIndexView; + return idIndexView != null ? idIndexView.getLong(groupId << 3) : _idIndex.getLong((long) groupId << 3); + } + + /// Compares the key against the payload record at the given global offset, and returns the record's group id on a + /// full match (stored hash, key length, key bytes), or [GroupKeyGenerator#INVALID_ID] on a mismatch. + private int matchRecord(long globalOffset, long hash, byte[] key, int offset, int length) { + int chunkIndex = (int) (globalOffset / CHUNK_SIZE); + ByteBuffer chunkView = _chunkViews.get(chunkIndex); + if (chunkView == null) { + return matchRecordSlow(chunkIndex, globalOffset % CHUNK_SIZE, hash, key, offset, length); + } + // Record start offsets within a chunk are always < CHUNK_SIZE, so they fit in an int + int offsetInChunk = (int) (globalOffset % CHUNK_SIZE); + if (chunkView.getLong(offsetInChunk) != hash || chunkView.getInt(offsetInChunk + KEY_LENGTH_OFFSET) != length) { + return GroupKeyGenerator.INVALID_ID; + } + if (length > 0) { + // Bulk copy + range equals instead of per-byte buffer reads: this runs on every successful lookup (the + // dominant case — every row after the first per group), and both the bulk get and Arrays.equals are + // intrinsified. The absolute bulk get does not touch the view's position. + byte[] scratch = _compareScratch; + if (scratch.length < length) { + scratch = new byte[Math.max(length, scratch.length << 1)]; + _compareScratch = scratch; + } + chunkView.get(offsetInChunk + RECORD_HEADER_SIZE, scratch, 0, length); + if (!Arrays.equals(scratch, 0, length, key, offset, offset + length)) { + return GroupKeyGenerator.INVALID_ID; + } + } + return chunkView.getInt(offsetInChunk + GROUP_ID_OFFSET); + } + + /// Wrapper-based fallback of [#matchRecord] for the rare chunk without a direct view (larger than 2GB). + private int matchRecordSlow(int chunkIndex, long offsetInChunk, long hash, byte[] key, int offset, int length) { + PinotDataBuffer chunk = _chunks.get(chunkIndex); + if (chunk.getLong(offsetInChunk) != hash || chunk.getInt(offsetInChunk + KEY_LENGTH_OFFSET) != length) { + return GroupKeyGenerator.INVALID_ID; + } + if (length > 0) { + byte[] scratch = _compareScratch; + if (scratch.length < length) { + scratch = new byte[Math.max(length, scratch.length << 1)]; + _compareScratch = scratch; + } + chunk.copyTo(offsetInChunk + RECORD_HEADER_SIZE, scratch, 0, length); + if (!Arrays.equals(scratch, 0, length, key, offset, offset + length)) { + return GroupKeyGenerator.INVALID_ID; + } + } + return chunk.getInt(offsetInChunk + GROUP_ID_OFFSET); + } + + /// Inserts a new key: appends the payload record, records it in the id index, writes the directory entry into the + /// given empty slot, and resizes the directory when it reaches the 0.5 load factor. + private int insert(long hash, byte[] key, int offset, int length, int slot) { + int groupId = _size; + long globalOffset = appendRecord(hash, groupId, key, offset, length); + if (((long) groupId << 3) == _idIndex.size()) { + growIdIndex(); + } + ByteBuffer idIndexView = _idIndexView; + if (idIndexView != null) { + idIndexView.putLong(groupId << 3, globalOffset); + } else { + _idIndex.putLong((long) groupId << 3, globalOffset); + } + _directory.putLong((long) slot << 3, (hash & SALT_MASK) | (globalOffset + 1)); + _size++; + if (_size >= (_numSlots >> 1)) { + resizeDirectory(); + } + return groupId; + } + + /// Appends a payload record and returns its global offset (`chunkIndex * CHUNK_SIZE + offsetInChunk`). + private long appendRecord(long hash, int groupId, byte[] key, int offset, int length) { + long recordSize = RECORD_HEADER_SIZE + (long) length; + PinotDataBuffer chunk; + int chunkIndex; + int recordStart; + if (recordSize > CHUNK_SIZE) { + // Oversized record: dedicated chunk of exactly the record size, record at offset 0. The next record starts + // a fresh normal chunk, so global offset decoding (division by CHUNK_SIZE) stays valid. + chunk = allocateChunk(recordSize); + chunkIndex = _chunks.size() - 1; + recordStart = 0; + _currentChunk = null; + } else { + if (_currentChunk == null || CHUNK_SIZE - _currentChunkOffset < recordSize) { + _currentChunk = allocateChunk(CHUNK_SIZE); + _currentChunkIndex = _chunks.size() - 1; + _currentChunkOffset = 0; + } + chunk = _currentChunk; + chunkIndex = _currentChunkIndex; + recordStart = _currentChunkOffset; + _currentChunkOffset += (int) recordSize; + } + // Record start offsets within a chunk are always < CHUNK_SIZE, which the global offset encoding relies on. + assert recordStart < CHUNK_SIZE; + chunk.putLong(recordStart, hash); + chunk.putInt((long) recordStart + GROUP_ID_OFFSET, groupId); + chunk.putInt((long) recordStart + KEY_LENGTH_OFFSET, length); + if (length > 0) { + chunk.readFrom((long) recordStart + RECORD_HEADER_SIZE, key, offset, length); + } + return (long) chunkIndex * CHUNK_SIZE + recordStart; + } + + private PinotDataBuffer allocateChunk(long sizeBytes) { + PinotDataBuffer chunk = OffHeapGroupByBufferPool.acquire(sizeBytes, CHUNK_DESCRIPTION); + _chunks.add(chunk); + _chunkViews.add(OffHeapGroupByUtils.createView(chunk, sizeBytes)); + return chunk; + } + + /// Doubles the directory. Only the hash stored in each payload record is re-read (via the id index) to recompute + /// the slot and salt; key bytes are never touched. + private void resizeDirectory() { + int newNumSlots = _numSlots << 1; + Preconditions.checkState(newNumSlots > 0 && newNumSlots <= MAX_NUM_SLOTS, "Cannot grow directory beyond %s slots", + MAX_NUM_SLOTS); + PinotDataBuffer newDirectory = allocateZeroFilledDirectory(newNumSlots); + ByteBuffer newDirectoryView = OffHeapGroupByUtils.createView(newDirectory, (long) newNumSlots << 3); + int newMask = newNumSlots - 1; + for (int groupId = 0; groupId < _size; groupId++) { + long globalOffset = getPayloadGlobalOffset(groupId); + int chunkIndex = (int) (globalOffset / CHUNK_SIZE); + ByteBuffer chunkView = _chunkViews.get(chunkIndex); + long hash = chunkView != null ? chunkView.getLong((int) (globalOffset % CHUNK_SIZE)) + : _chunks.get(chunkIndex).getLong(globalOffset % CHUNK_SIZE); + int slot = (int) (hash & newMask); + if (newDirectoryView != null) { + while (newDirectoryView.getLong(slot << 3) != 0L) { + slot = (slot + 1) & newMask; + } + newDirectoryView.putLong(slot << 3, (hash & SALT_MASK) | (globalOffset + 1)); + } else { + while (newDirectory.getLong((long) slot << 3) != 0L) { + slot = (slot + 1) & newMask; + } + newDirectory.putLong((long) slot << 3, (hash & SALT_MASK) | (globalOffset + 1)); + } + } + closeBuffer(_directory); + _directory = newDirectory; + _directoryView = newDirectoryView; + _numSlots = newNumSlots; + } + + private void growIdIndex() { + long oldSizeBytes = _idIndex.size(); + PinotDataBuffer newIdIndex = OffHeapGroupByBufferPool.acquire(oldSizeBytes << 1, ID_INDEX_DESCRIPTION); + _idIndex.copyTo(0, newIdIndex, 0, oldSizeBytes); + closeBuffer(_idIndex); + _idIndex = newIdIndex; + _idIndexView = OffHeapGroupByUtils.createView(newIdIndex, oldSizeBytes << 1); + } + + private static PinotDataBuffer allocateZeroFilledDirectory(int numSlots) { + long sizeBytes = (long) numSlots << 3; + PinotDataBuffer directory = OffHeapGroupByBufferPool.acquire(sizeBytes, DIRECTORY_DESCRIPTION); + // Contents of allocateDirect are undefined, and 0 means an empty slot, so zero-fill explicitly (bulk puts + // through the direct view when available) + ByteBuffer view = OffHeapGroupByUtils.createView(directory, sizeBytes); + if (view != null) { + int size = (int) sizeBytes; + for (int offset = 0; offset < size; offset += ZERO_CHUNK.length) { + view.put(offset, ZERO_CHUNK, 0, Math.min(ZERO_CHUNK.length, size - offset)); + } + } else { + for (long offset = 0; offset < sizeBytes; offset += 8) { + directory.putLong(offset, 0L); + } + } + return directory; + } + + private static int computeNumSlots(int expectedNumEntries) { + long target = Math.max(MIN_NUM_SLOTS, 2L * expectedNumEntries); + long numSlots = Long.highestOneBit(target); + if (numSlots < target) { + numSlots <<= 1; + } + return (int) Math.min(numSlots, MAX_NUM_SLOTS); + } + + private static void closeBuffer(PinotDataBuffer buffer) { + OffHeapGroupByBufferPool.release(buffer); + } + + private static RuntimeException closeBufferQuietly(PinotDataBuffer buffer, RuntimeException firstException) { + try { + closeBuffer(buffer); + return firstException; + } catch (RuntimeException e) { + return firstException != null ? firstException : e; + } + } + + /// Standard MurmurHash3 x64 128 (Austin Appleby), returning the low 64 bits (`h1`). Implemented privately + /// because `MurmurHashFunctions` has no 64-bit variant accepting `(byte[], offset, length)`. + /// Deterministic within a process, which is all this per-query scratch structure needs. + private static long murmurHash3X64Bit64(byte[] data, int offset, int length) { + final long c1 = 0x87c37b91114253d5L; + final long c2 = 0x4cf5ad432745937fL; + long h1 = MURMUR3_SEED; + long h2 = MURMUR3_SEED; + int end = offset + (length & ~15); + for (int i = offset; i < end; i += 16) { + long k1 = getLongLittleEndian(data, i); + long k2 = getLongLittleEndian(data, i + 8); + k1 *= c1; + k1 = Long.rotateLeft(k1, 31); + k1 *= c2; + h1 ^= k1; + h1 = Long.rotateLeft(h1, 27); + h1 += h2; + h1 = h1 * 5 + 0x52dce729L; + k2 *= c2; + k2 = Long.rotateLeft(k2, 33); + k2 *= c1; + h2 ^= k2; + h2 = Long.rotateLeft(h2, 31); + h2 += h1; + h2 = h2 * 5 + 0x38495ab5L; + } + long k1 = 0; + long k2 = 0; + // CHECKSTYLE:OFF: checkstyle:coding + switch (length & 15) { + case 15: + k2 ^= (data[end + 14] & 0xffL) << 48; + case 14: + k2 ^= (data[end + 13] & 0xffL) << 40; + case 13: + k2 ^= (data[end + 12] & 0xffL) << 32; + case 12: + k2 ^= (data[end + 11] & 0xffL) << 24; + case 11: + k2 ^= (data[end + 10] & 0xffL) << 16; + case 10: + k2 ^= (data[end + 9] & 0xffL) << 8; + case 9: + k2 ^= data[end + 8] & 0xffL; + k2 *= c2; + k2 = Long.rotateLeft(k2, 33); + k2 *= c1; + h2 ^= k2; + case 8: + k1 ^= (data[end + 7] & 0xffL) << 56; + case 7: + k1 ^= (data[end + 6] & 0xffL) << 48; + case 6: + k1 ^= (data[end + 5] & 0xffL) << 40; + case 5: + k1 ^= (data[end + 4] & 0xffL) << 32; + case 4: + k1 ^= (data[end + 3] & 0xffL) << 24; + case 3: + k1 ^= (data[end + 2] & 0xffL) << 16; + case 2: + k1 ^= (data[end + 1] & 0xffL) << 8; + case 1: + k1 ^= data[end] & 0xffL; + k1 *= c1; + k1 = Long.rotateLeft(k1, 31); + k1 *= c2; + h1 ^= k1; + } + // CHECKSTYLE:ON: checkstyle:coding + h1 ^= length; + h2 ^= length; + h1 += h2; + h2 += h1; + h1 = fmix64(h1); + h2 = fmix64(h2); + h1 += h2; + return h1; + } + + private static long getLongLittleEndian(byte[] data, int offset) { + return (data[offset] & 0xffL) | ((data[offset + 1] & 0xffL) << 8) | ((data[offset + 2] & 0xffL) << 16) | ( + (data[offset + 3] & 0xffL) << 24) | ((data[offset + 4] & 0xffL) << 32) | ((data[offset + 5] & 0xffL) << 40) | ( + (data[offset + 6] & 0xffL) << 48) | ((data[offset + 7] & 0xffL) << 56); + } + + private static long fmix64(long k) { + k ^= k >>> 33; + k *= 0xff51afd7ed558ccdL; + k ^= k >>> 33; + k *= 0xc4ceb9fe1a85ec53L; + k ^= k >>> 33; + return k; + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapDoubleGroupByResultHolder.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapDoubleGroupByResultHolder.java new file mode 100644 index 000000000000..ff1f7e6b57ca --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapDoubleGroupByResultHolder.java @@ -0,0 +1,196 @@ +/** + * 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.offheap; + +import com.google.common.base.Preconditions; +import java.nio.ByteBuffer; +import javax.annotation.concurrent.NotThreadSafe; +import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; + + +/// Off-heap implementation of [GroupByResultHolder] for double results, backed by a single direct-memory +/// [PinotDataBuffer]. Drop-in replacement for +/// [org.apache.pinot.core.query.aggregation.groupby.DoubleGroupByResultHolder] with identical semantics +/// observed through the interface: each group key indexes a fixed-width 8-byte slot, un-initialized slots return +/// the default value, and [GroupKeyGenerator#INVALID_ID] returns the default value on get and is silently +/// ignored on set. +/// +/// An initial capacity of 0 is legal (a zero-size buffer is allocated); in that state no getter or setter may +/// be called until [#ensureCapacity(int)] grows the holder. +/// +/// [#close()] releases the direct memory and is idempotent; the behavior of all other methods after +/// close is undefined. +/// +/// This class is single-threaded and not thread-safe. +@NotThreadSafe +public class OffHeapDoubleGroupByResultHolder implements GroupByResultHolder, AutoCloseable { + private static final long BYTES_PER_VALUE = Double.BYTES; + private static final int VALUE_SHIFT = 3; + private static final String BUFFER_DESCRIPTION = "OffHeapDoubleGroupByResultHolder"; + + private final int _maxCapacity; + private final double _defaultValue; + + private int _resultHolderCapacity; + private PinotDataBuffer _dataBuffer; + // Absolute-indexed direct view of _dataBuffer for the per-row hot path (monomorphic, intrinsified + // ByteBuffer access instead of the PinotDataBuffer wrapper); null when empty or beyond the 2GB view limit + private ByteBuffer _view; + private boolean _closed; + + /// Constructor for the class. + /// + /// @param initialCapacity Initial capacity of the result holder + /// @param maxCapacity Maximum capacity of the result holder + /// @param defaultValue Default value of un-initialized results + public OffHeapDoubleGroupByResultHolder(int initialCapacity, int maxCapacity, double defaultValue) { + _maxCapacity = maxCapacity; + _defaultValue = defaultValue; + + _resultHolderCapacity = initialCapacity; + _dataBuffer = OffHeapGroupByBufferPool.acquire(initialCapacity * BYTES_PER_VALUE, BUFFER_DESCRIPTION); + _view = createView(_dataBuffer, initialCapacity); + // Direct buffer contents are undefined, so always fill the allocated region with the default value + fillWithDefaultValue(0, initialCapacity); + } + + @Override + public void ensureCapacity(int capacity) { + Preconditions.checkArgument(capacity <= _maxCapacity); + + if (capacity > _resultHolderCapacity) { + int copyLength = _resultHolderCapacity; + // Cap the growth to maximum possible number of group keys. NOTE: _resultHolderCapacity (the bounds-guard + // reference) is updated only after the new buffer is successfully acquired. + int newCapacity = Math.min(Math.max(_resultHolderCapacity * 2, capacity), _maxCapacity); + + PinotDataBuffer current = _dataBuffer; + _dataBuffer = OffHeapGroupByBufferPool.acquire(newCapacity * BYTES_PER_VALUE, BUFFER_DESCRIPTION); + _view = createView(_dataBuffer, newCapacity); + _resultHolderCapacity = newCapacity; + if (copyLength > 0) { + current.copyTo(0, _dataBuffer, 0, copyLength * BYTES_PER_VALUE); + } + // Fill the newly extended region with the default value (direct buffer contents are undefined) + fillWithDefaultValue(copyLength, newCapacity); + closeBuffer(current); + } + } + + @Override + public double getDoubleResult(int groupKey) { + if (groupKey == GroupKeyGenerator.INVALID_ID) { + return _defaultValue; + } else { + // PinotDataBuffer access is unchecked: an out-of-range key would read arbitrary memory instead of the + // on-heap twin's ArrayIndexOutOfBoundsException, so guard the sizing contract with an assert + assert groupKey >= 0 && groupKey < _resultHolderCapacity : "groupKey " + groupKey + " out of bounds"; + ByteBuffer view = _view; + if (view != null) { + return view.getDouble(groupKey << VALUE_SHIFT); + } + return _dataBuffer.getDouble(groupKey * BYTES_PER_VALUE); + } + } + + @Override + public int getIntResult(int groupKey) { + throw new UnsupportedOperationException(); + } + + @Override + public long getLongResult(int groupKey) { + throw new UnsupportedOperationException(); + } + + @Override + public T getResult(int groupKey) { + throw new UnsupportedOperationException(); + } + + @Override + public void setValueForKey(int groupKey, double newValue) { + if (groupKey != GroupKeyGenerator.INVALID_ID) { + // See getDoubleResult: unchecked buffer access means an out-of-range key would corrupt memory, not throw + assert groupKey >= 0 && groupKey < _resultHolderCapacity : "groupKey " + groupKey + " out of bounds"; + ByteBuffer view = _view; + if (view != null) { + view.putDouble(groupKey << VALUE_SHIFT, newValue); + } else { + _dataBuffer.putDouble(groupKey * BYTES_PER_VALUE, newValue); + } + } + } + + @Override + public void setValueForKey(int groupKey, int value) { + throw new UnsupportedOperationException(); + } + + @Override + public void setValueForKey(int groupKey, long value) { + throw new UnsupportedOperationException(); + } + + @Override + public void setValueForKey(int groupKey, Object newValue) { + throw new UnsupportedOperationException(); + } + + public double getDefaultValue() { + return _defaultValue; + } + + @Override + public void close() { + if (_closed) { + return; + } + _closed = true; + closeBuffer(_dataBuffer); + // Null the buffer and view so any use-after-close (or a second release of a pooled buffer) fails loudly with + // an NPE instead of silently aliasing memory that the pool may have handed to another query + _dataBuffer = null; + _view = null; + } + + private void fillWithDefaultValue(int fromIndex, int toIndex) { + ByteBuffer view = _view; + if (view != null) { + // Intrinsified view puts are several times cheaper than the PinotDataBuffer wrapper accessors + for (int i = fromIndex; i < toIndex; i++) { + view.putDouble(i << VALUE_SHIFT, _defaultValue); + } + } else { + for (int i = fromIndex; i < toIndex; i++) { + _dataBuffer.putDouble(i * BYTES_PER_VALUE, _defaultValue); + } + } + } + + private static ByteBuffer createView(PinotDataBuffer dataBuffer, int capacity) { + return OffHeapGroupByUtils.createView(dataBuffer, capacity * BYTES_PER_VALUE); + } + + private static void closeBuffer(PinotDataBuffer buffer) { + OffHeapGroupByBufferPool.release(buffer); + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByBufferPool.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByBufferPool.java new file mode 100644 index 000000000000..ddaf49042e50 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByBufferPool.java @@ -0,0 +1,159 @@ +/** + * 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.offheap; + +import com.google.common.annotations.VisibleForTesting; +import it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap; +import java.io.IOException; +import java.nio.ByteOrder; +import java.util.ArrayDeque; +import java.util.concurrent.atomic.LongAdder; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; + + +/// Bounded per-thread free-list of direct buffers for the off-heap group-by structures. The on-heap group-by maps +/// are thread-local-cached across queries, which spares them per-query allocation and keeps their memory warm; this +/// pool gives the off-heap structures the same steady-state behavior with an explicit bound and full accounting: +///
    +///
  • Disabled by default (`maxBytesPerThread == 0`): [#acquire] allocates and [#release] +/// closes, i.e. exactly the unpooled per-query lifecycle.
    • +///
  • When enabled (server config +/// `pinot.server.query.executor.groupby.offheap.pool.max.bytes.per.thread`), released buffers are kept in +/// an exact-size free-list on the releasing thread, up to the per-thread byte cap; excess buffers are closed. +/// Query shapes repeat, so exact-size reuse hits in steady state.
    • +///
  • Pooled buffers remain open, so they stay visible in [PinotDataBuffer#getDirectBufferUsage()]; the +/// pool-retained portion is additionally tracked by [#getPooledBytes()].
    • +///
  • Buffers are returned dirty: every acquirer must initialize the content it relies on (all off-heap +/// group-by structures already zero-fill / default-fill on construction and growth).
    • +///
+/// +/// The cap is per releasing thread, so the aggregate retention bound is {@code maxBytesPerThread x the number +/// of threads that release group-by buffers} (combine workers plus reduce threads). The pool is intended for +/// long-lived executor threads: buffers pooled by a thread that dies are reclaimed only when its thread-local is +/// garbage-collected, and the usage counters do not observe that, so do not enable it for short-lived thread +/// pools. Setting the cap (back) to 0 drains lazily: each thread closes its pooled buffers on its next +/// [#acquire]. +/// +/// Thread-safety: the free-lists are thread-local; acquire and release may run on different threads (a block's +/// generator can be closed by a combine thread), in which case the buffer simply migrates to the releasing +/// thread's free-list. The global counters use [LongAdder]. +public final class OffHeapGroupByBufferPool { + private OffHeapGroupByBufferPool() { + } + + private static volatile long _maxBytesPerThread = 0; + + private static final LongAdder POOLED_BYTES = new LongAdder(); + private static final ThreadLocal THREAD_POOL = ThreadLocal.withInitial(ThreadPool::new); + + /// Sets the per-thread cap on pooled bytes. 0 (default) disables pooling. Configured once at server startup from + /// the query executor config. + public static void setMaxBytesPerThread(long maxBytesPerThread) { + _maxBytesPerThread = maxBytesPerThread; + } + + /// Returns the total bytes currently retained by the free-lists of all threads. + public static long getPooledBytes() { + return POOLED_BYTES.sum(); + } + + /// Returns a native-order direct buffer of exactly the given size, reusing a pooled buffer when one of the exact + /// size is available on this thread. The content is undefined either way. + public static PinotDataBuffer acquire(long sizeBytes, String description) { + ThreadPool threadPool = THREAD_POOL.get(); + if (_maxBytesPerThread > 0) { + PinotDataBuffer pooled = threadPool.poll(sizeBytes); + if (pooled != null) { + POOLED_BYTES.add(-sizeBytes); + return pooled; + } + } else if (threadPool._pooledBytes > 0) { + // Pooling was disabled (e.g. live config change to 0): drain this thread's retained buffers lazily + POOLED_BYTES.add(-threadPool._pooledBytes); + threadPool.clear(); + } + return PinotDataBuffer.allocateDirect(sizeBytes, ByteOrder.nativeOrder(), description); + } + + /// Returns a buffer obtained from [#acquire] to the pool of the current thread, or closes it when pooling + /// is disabled or the per-thread cap is reached. + public static void release(PinotDataBuffer buffer) { + long maxBytesPerThread = _maxBytesPerThread; + long sizeBytes = buffer.size(); + if (maxBytesPerThread > 0 && sizeBytes > 0 && THREAD_POOL.get().offer(buffer, sizeBytes, maxBytesPerThread)) { + POOLED_BYTES.add(sizeBytes); + return; + } + try { + buffer.close(); + } catch (IOException e) { + throw new RuntimeException("Failed to close PinotDataBuffer", e); + } + } + + /// Closes and drops every buffer pooled by the current thread. Test hook. + @VisibleForTesting + public static void clearCurrentThread() { + ThreadPool threadPool = THREAD_POOL.get(); + POOLED_BYTES.add(-threadPool._pooledBytes); + threadPool.clear(); + } + + private static final class ThreadPool { + private final Long2ObjectOpenHashMap> _freeListsBySize = + new Long2ObjectOpenHashMap<>(); + private long _pooledBytes; + + PinotDataBuffer poll(long sizeBytes) { + ArrayDeque freeList = _freeListsBySize.get(sizeBytes); + if (freeList == null) { + return null; + } + PinotDataBuffer buffer = freeList.poll(); + if (buffer != null) { + _pooledBytes -= sizeBytes; + } + return buffer; + } + + boolean offer(PinotDataBuffer buffer, long sizeBytes, long maxBytes) { + if (_pooledBytes + sizeBytes > maxBytes) { + return false; + } + _freeListsBySize.computeIfAbsent(sizeBytes, k -> new ArrayDeque<>()).offer(buffer); + _pooledBytes += sizeBytes; + return true; + } + + void clear() { + for (ArrayDeque freeList : _freeListsBySize.values()) { + PinotDataBuffer buffer; + while ((buffer = freeList.poll()) != null) { + try { + buffer.close(); + } catch (IOException e) { + throw new RuntimeException("Failed to close pooled PinotDataBuffer", e); + } + } + } + _freeListsBySize.clear(); + _pooledBytes = 0; + } + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByUtils.java new file mode 100644 index 000000000000..8766d22c84e6 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByUtils.java @@ -0,0 +1,122 @@ +/** + * 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.offheap; + +import com.google.common.annotations.VisibleForTesting; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; + + +/// Static helpers for encoding group-by keys into the scratch byte buffers fed to [OffHeapBytesGroupIdMap]. +/// All methods are allocation-free on their fast paths; callers own (and reuse) the scratch arrays. +/// +/// Thread-safety: stateless; the scratch arrays passed in are owned by the single-threaded caller. +public final class OffHeapGroupByUtils { + private OffHeapGroupByUtils() { + } + + // Buffers larger than this get no direct view (ByteBuffer is int-indexed); test hook so the wrapper-based + // fallback arms of every view fast path can be exercised without allocating multi-GB buffers + private static volatile long _viewSizeLimitBytes = Integer.MAX_VALUE; + + @VisibleForTesting + public static void setViewSizeLimitBytes(long viewSizeLimitBytes) { + _viewSizeLimitBytes = viewSizeLimitBytes; + } + + /// Returns an absolute-indexed native-order direct [ByteBuffer] view of the buffer for hot-path access + /// (monomorphic, intrinsified access instead of the [PinotDataBuffer] wrapper), or null when the buffer is + /// empty or exceeds the 2GB view limit — callers must fall back to the wrapper accessors then. + public static ByteBuffer createView(PinotDataBuffer buffer, long sizeBytes) { + return sizeBytes > 0 && sizeBytes <= _viewSizeLimitBytes + ? buffer.toDirectByteBuffer(0, (int) sizeBytes, ByteOrder.nativeOrder()) : null; + } + + /// Returns a scratch array of at least `capacity` bytes, growing (with doubling) if needed. Contents are not + /// preserved on growth. + public static byte[] ensureByteCapacity(byte[] scratch, int capacity) { + if (scratch.length >= capacity) { + return scratch; + } + return new byte[Math.max(capacity, scratch.length << 1)]; + } + + /// Encodes the given string into the scratch buffer as standard UTF-8 and returns the encoded length. The scratch + /// buffer must have capacity of at least `3 * value.length()` bytes (a surrogate pair encodes 2 chars into 4 + /// bytes, so the bound holds for all inputs). + /// + /// This produces byte-for-byte the same encoding as `String.getBytes(StandardCharsets.UTF_8)` for all + /// inputs, including supplementary characters (4-byte sequences) and malformed surrogates, which the JDK encoder + /// replaces with `'?'` — pinned by `OffHeapGroupByUtilsTest`. + public static int encodeUtf8(String value, byte[] scratch) { + int length = value.length(); + int outIndex = 0; + int charIndex = 0; + while (charIndex < length) { + char c = value.charAt(charIndex++); + if (c < 0x80) { + scratch[outIndex++] = (byte) c; + } else if (c < 0x800) { + scratch[outIndex++] = (byte) (0xC0 | (c >> 6)); + scratch[outIndex++] = (byte) (0x80 | (c & 0x3F)); + } else if (c >= Character.MIN_SURROGATE && c <= Character.MAX_SURROGATE) { + if (Character.isHighSurrogate(c) && charIndex < length && Character.isLowSurrogate(value.charAt(charIndex))) { + int codePoint = Character.toCodePoint(c, value.charAt(charIndex++)); + scratch[outIndex++] = (byte) (0xF0 | (codePoint >> 18)); + scratch[outIndex++] = (byte) (0x80 | ((codePoint >> 12) & 0x3F)); + scratch[outIndex++] = (byte) (0x80 | ((codePoint >> 6) & 0x3F)); + scratch[outIndex++] = (byte) (0x80 | (codePoint & 0x3F)); + } else { + // Unpaired surrogate: the JDK UTF-8 encoder replaces it with '?' + scratch[outIndex++] = '?'; + } + } else { + scratch[outIndex++] = (byte) (0xE0 | (c >> 12)); + scratch[outIndex++] = (byte) (0x80 | ((c >> 6) & 0x3F)); + scratch[outIndex++] = (byte) (0x80 | (c & 0x3F)); + } + } + return outIndex; + } + + /// Packs `numValues` ints into the scratch buffer (4 bytes each, big-endian) and returns the packed length. + /// The scratch buffer must have capacity of at least `4 * numValues` bytes. + public static int packInts(int[] values, int numValues, byte[] scratch) { + int outIndex = 0; + for (int i = 0; i < numValues; i++) { + int value = values[i]; + scratch[outIndex++] = (byte) (value >> 24); + scratch[outIndex++] = (byte) (value >> 16); + scratch[outIndex++] = (byte) (value >> 8); + scratch[outIndex++] = (byte) value; + } + return outIndex; + } + + /// Unpacks `numValues` big-endian ints from the scratch buffer written by [#packInts]. + public static void unpackInts(byte[] scratch, int numValues, int[] dest) { + int inIndex = 0; + for (int i = 0; i < numValues; i++) { + dest[i] = ((scratch[inIndex] & 0xFF) << 24) | ((scratch[inIndex + 1] & 0xFF) << 16) + | ((scratch[inIndex + 2] & 0xFF) << 8) | (scratch[inIndex + 3] & 0xFF); + inIndex += 4; + } + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapIntGroupByResultHolder.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapIntGroupByResultHolder.java new file mode 100644 index 000000000000..c9b37b56e148 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapIntGroupByResultHolder.java @@ -0,0 +1,196 @@ +/** + * 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.offheap; + +import com.google.common.base.Preconditions; +import java.nio.ByteBuffer; +import javax.annotation.concurrent.NotThreadSafe; +import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; + + +/// Off-heap implementation of [GroupByResultHolder] for int results, backed by a single direct-memory +/// [PinotDataBuffer]. Drop-in replacement for +/// [org.apache.pinot.core.query.aggregation.groupby.IntGroupByResultHolder] with identical semantics +/// observed through the interface: each group key indexes a fixed-width 4-byte slot, un-initialized slots return +/// the default value, and [GroupKeyGenerator#INVALID_ID] returns the default value on get and is silently +/// ignored on set. +/// +/// An initial capacity of 0 is legal (a zero-size buffer is allocated); in that state no getter or setter may +/// be called until [#ensureCapacity(int)] grows the holder. +/// +/// [#close()] releases the direct memory and is idempotent; the behavior of all other methods after +/// close is undefined. +/// +/// This class is single-threaded and not thread-safe. +@NotThreadSafe +public class OffHeapIntGroupByResultHolder implements GroupByResultHolder, AutoCloseable { + private static final long BYTES_PER_VALUE = Integer.BYTES; + private static final int VALUE_SHIFT = 2; + private static final String BUFFER_DESCRIPTION = "OffHeapIntGroupByResultHolder"; + + private final int _maxCapacity; + private final int _defaultValue; + + private int _resultHolderCapacity; + private PinotDataBuffer _dataBuffer; + // Absolute-indexed direct view of _dataBuffer for the per-row hot path (monomorphic, intrinsified + // ByteBuffer access instead of the PinotDataBuffer wrapper); null when empty or beyond the 2GB view limit + private ByteBuffer _view; + private boolean _closed; + + /// Constructor for the class. + /// + /// @param initialCapacity Initial capacity of the result holder + /// @param maxCapacity Maximum capacity of the result holder + /// @param defaultValue Default value of un-initialized results + public OffHeapIntGroupByResultHolder(int initialCapacity, int maxCapacity, int defaultValue) { + _maxCapacity = maxCapacity; + _defaultValue = defaultValue; + + _resultHolderCapacity = initialCapacity; + _dataBuffer = OffHeapGroupByBufferPool.acquire(initialCapacity * BYTES_PER_VALUE, BUFFER_DESCRIPTION); + _view = createView(_dataBuffer, initialCapacity); + // Direct buffer contents are undefined, so always fill the allocated region with the default value + fillWithDefaultValue(0, initialCapacity); + } + + @Override + public void ensureCapacity(int capacity) { + Preconditions.checkArgument(capacity <= _maxCapacity); + + if (capacity > _resultHolderCapacity) { + int copyLength = _resultHolderCapacity; + // Cap the growth to maximum possible number of group keys. NOTE: _resultHolderCapacity (the bounds-guard + // reference) is updated only after the new buffer is successfully acquired. + int newCapacity = Math.min(Math.max(_resultHolderCapacity * 2, capacity), _maxCapacity); + + PinotDataBuffer current = _dataBuffer; + _dataBuffer = OffHeapGroupByBufferPool.acquire(newCapacity * BYTES_PER_VALUE, BUFFER_DESCRIPTION); + _view = createView(_dataBuffer, newCapacity); + _resultHolderCapacity = newCapacity; + if (copyLength > 0) { + current.copyTo(0, _dataBuffer, 0, copyLength * BYTES_PER_VALUE); + } + // Fill the newly extended region with the default value (direct buffer contents are undefined) + fillWithDefaultValue(copyLength, newCapacity); + closeBuffer(current); + } + } + + @Override + public double getDoubleResult(int groupKey) { + throw new UnsupportedOperationException(); + } + + @Override + public int getIntResult(int groupKey) { + if (groupKey == GroupKeyGenerator.INVALID_ID) { + return _defaultValue; + } else { + // PinotDataBuffer access is unchecked: an out-of-range key would read arbitrary memory instead of the + // on-heap twin's ArrayIndexOutOfBoundsException, so guard the sizing contract with an assert + assert groupKey >= 0 && groupKey < _resultHolderCapacity : "groupKey " + groupKey + " out of bounds"; + ByteBuffer view = _view; + if (view != null) { + return view.getInt(groupKey << VALUE_SHIFT); + } + return _dataBuffer.getInt(groupKey * BYTES_PER_VALUE); + } + } + + @Override + public long getLongResult(int groupKey) { + throw new UnsupportedOperationException(); + } + + @Override + public T getResult(int groupKey) { + throw new UnsupportedOperationException(); + } + + @Override + public void setValueForKey(int groupKey, double newValue) { + throw new UnsupportedOperationException(); + } + + @Override + public void setValueForKey(int groupKey, int newValue) { + if (groupKey != GroupKeyGenerator.INVALID_ID) { + // See getIntResult: unchecked buffer access means an out-of-range key would corrupt memory, not throw + assert groupKey >= 0 && groupKey < _resultHolderCapacity : "groupKey " + groupKey + " out of bounds"; + ByteBuffer view = _view; + if (view != null) { + view.putInt(groupKey << VALUE_SHIFT, newValue); + } else { + _dataBuffer.putInt(groupKey * BYTES_PER_VALUE, newValue); + } + } + } + + @Override + public void setValueForKey(int groupKey, long newValue) { + throw new UnsupportedOperationException(); + } + + @Override + public void setValueForKey(int groupKey, Object newValue) { + throw new UnsupportedOperationException(); + } + + public int getDefaultValue() { + return _defaultValue; + } + + @Override + public void close() { + if (_closed) { + return; + } + _closed = true; + closeBuffer(_dataBuffer); + // Null the buffer and view so any use-after-close (or a second release of a pooled buffer) fails loudly with + // an NPE instead of silently aliasing memory that the pool may have handed to another query + _dataBuffer = null; + _view = null; + } + + private void fillWithDefaultValue(int fromIndex, int toIndex) { + ByteBuffer view = _view; + if (view != null) { + // Intrinsified view puts are several times cheaper than the PinotDataBuffer wrapper accessors + for (int i = fromIndex; i < toIndex; i++) { + view.putInt(i << VALUE_SHIFT, _defaultValue); + } + } else { + for (int i = fromIndex; i < toIndex; i++) { + _dataBuffer.putInt(i * BYTES_PER_VALUE, _defaultValue); + } + } + } + + private static ByteBuffer createView(PinotDataBuffer dataBuffer, int capacity) { + return OffHeapGroupByUtils.createView(dataBuffer, capacity * BYTES_PER_VALUE); + } + + private static void closeBuffer(PinotDataBuffer buffer) { + OffHeapGroupByBufferPool.release(buffer); + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapIntGroupIdMap.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapIntGroupIdMap.java new file mode 100644 index 000000000000..7d3ba4e7147a --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapIntGroupIdMap.java @@ -0,0 +1,310 @@ +/** + * 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.offheap; + +import com.google.common.base.Preconditions; +import it.unimi.dsi.fastutil.HashCommon; +import java.nio.ByteBuffer; +import java.util.Iterator; +import java.util.NoSuchElementException; +import javax.annotation.concurrent.NotThreadSafe; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; + + +/// Off-heap hash table mapping a 32-bit raw group key to a dense int group id — the direct off-heap port of +/// `DictionaryBasedGroupKeyGenerator.IntGroupIdMap`, with the same 8-byte slot layout +/// ([int internalKey][int groupId]) so each probe touches half the memory of [OffHeapLongGroupIdMap]'s +/// 16-byte slots. This is the table behind the most common dictionary-encoded group-by path, and also serves raw +/// INT/FLOAT single-column keys: the one key whose internal form would collide with the empty-slot marker +/// (`rawKey == -1`) is held out-of-band in a field and does not occupy a slot. +/// +/// Group ids are assigned densely in insertion order: 0, 1, 2, ... [int)][#getGroupId(int,] returns the +/// existing id if the key is present (regardless of the upper bound); otherwise it assigns id `size()` if +/// `size() < groupIdUpperBound`, or returns [GroupKeyGenerator#INVALID_ID] without inserting. +/// +/// Implementation: open addressing with linear probing (step +1) over a power-of-two capacity with load factor +/// 0.5. As in the on-heap `IntGroupIdMap`, the internal key is `rawKey + 1` so that 0 marks an empty +/// slot; `rawKey == -1` (internal key 0) is the out-of-band field key. +/// +/// Sizing note: load factor 0.5 over 8-byte slots is 16 bytes of table per entry, versus the on-heap +/// `IntGroupIdMap`'s 0.75 over 8-byte slots (~10.7 bytes/entry) — deliberately trading some direct memory for +/// shorter probe chains, since the memory this feature relieves is the heap. +/// +/// Memory is direct (off-heap) via [PinotDataBuffer]. [#close()] releases it and is idempotent; +/// behavior of all other methods after close is undefined. +/// +/// This class is not thread-safe. +@NotThreadSafe +public class OffHeapIntGroupIdMap implements AutoCloseable { + public static final int INVALID_ID = GroupKeyGenerator.INVALID_ID; + + // Slot layout: [int internalKey][int groupId] = 8 bytes + private static final int SLOT_SHIFT = 3; + private static final int GROUP_ID_OFFSET_IN_SLOT = Integer.BYTES; + private static final int MIN_CAPACITY = 512; + // Largest power-of-two capacity representable as a positive int; expand() past this would overflow + private static final int MAX_CAPACITY = 1 << 30; + // Reusable zero block for bulk zero-filling freshly allocated buffers through the direct view + private static final byte[] ZERO_CHUNK = new byte[8192]; + + private PinotDataBuffer _buffer; + // Absolute-indexed direct view of _buffer for the per-row probe loop (monomorphic, intrinsified ByteBuffer + // access instead of the PinotDataBuffer wrapper); null when the table exceeds the 2GB view limit + private ByteBuffer _view; + private int _capacity; + private int _mask; + // Resize when the number of occupied slots exceeds this (i.e. load factor 0.5) + private int _maxOccupiedSlots; + // Number of occupied slots; the out-of-band -1 key does not occupy a slot and is not counted here + private int _occupiedSlots; + private int _minusOneKeyGroupId = INVALID_ID; + private boolean _closed; + + public OffHeapIntGroupIdMap(int expectedNumEntries) { + Preconditions.checkArgument(expectedNumEntries >= 0, "Invalid expectedNumEntries: %s", expectedNumEntries); + long desiredCapacity = Math.max(MIN_CAPACITY, (long) expectedNumEntries << 1); + _capacity = (int) Math.min(MAX_CAPACITY, Long.highestOneBit((desiredCapacity << 1) - 1)); + _mask = _capacity - 1; + _maxOccupiedSlots = _capacity >>> 1; + long sizeBytes = (long) _capacity << SLOT_SHIFT; + _buffer = allocate(sizeBytes); + _view = OffHeapGroupByUtils.createView(_buffer, sizeBytes); + zeroFill(_buffer, _view, sizeBytes); + } + + /// Returns the number of groups assigned so far, including the group for the raw key -1 if assigned. + public int size() { + return _minusOneKeyGroupId != INVALID_ID ? _occupiedSlots + 1 : _occupiedSlots; + } + + /// Returns the amount of off-heap memory held by this map in bytes. + public long getOffHeapMemoryBytes() { + return (long) _capacity << SLOT_SHIFT; + } + + /// Returns the group id for the given raw key (any int; -1 is held out-of-band). If the key is present, always + /// returns its id (even + /// when `size() >= groupIdUpperBound`). If absent and `size() < groupIdUpperBound`, assigns the next + /// dense id (`size()`) and returns it; otherwise returns [#INVALID_ID] without inserting. + public int getGroupId(int rawKey, int groupIdUpperBound) { + // rawKey == -1 is the one key whose internal form (0) would collide with the empty-slot marker: hold it + // out-of-band. NOTE: rawKey == Integer.MAX_VALUE wraps to Integer.MIN_VALUE, which is fine. + if (rawKey == -1) { + int minusOneKeyGroupId = _minusOneKeyGroupId; + if (minusOneKeyGroupId != INVALID_ID) { + return minusOneKeyGroupId; + } + // The -1 key is not assigned yet, so size() == _occupiedSlots here + int size = size(); + if (size < groupIdUpperBound) { + _minusOneKeyGroupId = size; + return size; + } + return INVALID_ID; + } + int internalKey = rawKey + 1; + ByteBuffer view = _view; + if (view == null) { + return getGroupIdSlow(internalKey, groupIdUpperBound); + } + // While the view exists, slot offsets fit in an int (view size <= Integer.MAX_VALUE) + int slot = HashCommon.mix(internalKey) & _mask; + while (true) { + int slotOffset = slot << SLOT_SHIFT; + int key = view.getInt(slotOffset); + if (key == internalKey) { + return view.getInt(slotOffset + GROUP_ID_OFFSET_IN_SLOT); + } + if (key == 0) { + int size = size(); + if (size >= groupIdUpperBound) { + return INVALID_ID; + } + view.putInt(slotOffset, internalKey); + view.putInt(slotOffset + GROUP_ID_OFFSET_IN_SLOT, size); + if (++_occupiedSlots > _maxOccupiedSlots) { + expand(); + } + return size; + } + slot = (slot + 1) & _mask; + } + } + + private int getGroupIdSlow(int internalKey, int groupIdUpperBound) { + PinotDataBuffer buffer = _buffer; + int slot = HashCommon.mix(internalKey) & _mask; + while (true) { + long slotOffset = (long) slot << SLOT_SHIFT; + int key = buffer.getInt(slotOffset); + if (key == internalKey) { + return buffer.getInt(slotOffset + GROUP_ID_OFFSET_IN_SLOT); + } + if (key == 0) { + int size = size(); + if (size >= groupIdUpperBound) { + return INVALID_ID; + } + buffer.putInt(slotOffset, internalKey); + buffer.putInt(slotOffset + GROUP_ID_OFFSET_IN_SLOT, size); + if (++_occupiedSlots > _maxOccupiedSlots) { + expand(); + } + return size; + } + slot = (slot + 1) & _mask; + } + } + + /// Returns an iterator over all (rawKey, groupId) entries in arbitrary slot order. Yields exactly [#size()] + /// entries. + /// + /// NOTE: The returned [Entry] instance is a flyweight reused across `next()` calls; copy the values + /// out if they need to outlive the next call. + public Iterator iterator() { + return new Iterator<>() { + private final Entry _entry = new Entry(); + private int _slot; + private int _remainingOccupiedSlots = _occupiedSlots; + private boolean _returnMinusOneKey = _minusOneKeyGroupId != INVALID_ID; + + @Override + public boolean hasNext() { + return _remainingOccupiedSlots > 0 || _returnMinusOneKey; + } + + @Override + public Entry next() { + if (_remainingOccupiedSlots > 0) { + int key; + long slotOffset; + do { + slotOffset = (long) _slot << SLOT_SHIFT; + key = _buffer.getInt(slotOffset); + _slot++; + } while (key == 0); + _entry._rawKey = key - 1; + _entry._groupId = _buffer.getInt(slotOffset + GROUP_ID_OFFSET_IN_SLOT); + _remainingOccupiedSlots--; + return _entry; + } + if (_returnMinusOneKey) { + _returnMinusOneKey = false; + _entry._rawKey = -1; + _entry._groupId = _minusOneKeyGroupId; + return _entry; + } + throw new NoSuchElementException(); + } + }; + } + + @Override + public void close() { + if (_closed) { + return; + } + _closed = true; + closeBuffer(_buffer); + // Null the buffer and view so any use-after-close (or a second release of a pooled buffer) fails loudly with + // an NPE instead of silently aliasing memory that the pool may have handed to another query + _buffer = null; + _view = null; + } + + private void expand() { + Preconditions.checkState(_capacity < MAX_CAPACITY, "Cannot expand beyond max capacity: %s", MAX_CAPACITY); + int newCapacity = _capacity << 1; + int newMask = newCapacity - 1; + long newSizeBytes = (long) newCapacity << SLOT_SHIFT; + PinotDataBuffer newBuffer = allocate(newSizeBytes); + ByteBuffer newView = OffHeapGroupByUtils.createView(newBuffer, newSizeBytes); + zeroFill(newBuffer, newView, newSizeBytes); + ByteBuffer oldView = _view; + if (oldView != null && newView != null) { + // Hot path: rehash through the direct views (slot offsets fit in an int while a view exists) + for (int slot = 0; slot < _capacity; slot++) { + int slotOffset = slot << SLOT_SHIFT; + int key = oldView.getInt(slotOffset); + if (key != 0) { + int newSlot = HashCommon.mix(key) & newMask; + int newSlotOffset = newSlot << SLOT_SHIFT; + while (newView.getInt(newSlotOffset) != 0) { + newSlot = (newSlot + 1) & newMask; + newSlotOffset = newSlot << SLOT_SHIFT; + } + newView.putInt(newSlotOffset, key); + newView.putInt(newSlotOffset + GROUP_ID_OFFSET_IN_SLOT, oldView.getInt(slotOffset + GROUP_ID_OFFSET_IN_SLOT)); + } + } + } else { + for (int slot = 0; slot < _capacity; slot++) { + long slotOffset = (long) slot << SLOT_SHIFT; + int key = _buffer.getInt(slotOffset); + if (key != 0) { + int newSlot = HashCommon.mix(key) & newMask; + long newSlotOffset = (long) newSlot << SLOT_SHIFT; + while (newBuffer.getInt(newSlotOffset) != 0) { + newSlot = (newSlot + 1) & newMask; + newSlotOffset = (long) newSlot << SLOT_SHIFT; + } + newBuffer.putInt(newSlotOffset, key); + newBuffer.putInt(newSlotOffset + GROUP_ID_OFFSET_IN_SLOT, + _buffer.getInt(slotOffset + GROUP_ID_OFFSET_IN_SLOT)); + } + } + } + closeBuffer(_buffer); + _buffer = newBuffer; + _capacity = newCapacity; + _mask = newMask; + _maxOccupiedSlots = newCapacity >>> 1; + _view = newView; + } + + private static PinotDataBuffer allocate(long sizeBytes) { + return OffHeapGroupByBufferPool.acquire(sizeBytes, "OffHeapIntGroupIdMap hash table"); + } + + /// Zero-fills a freshly allocated buffer (contents of [PinotDataBuffer#allocateDirect] are undefined, and + /// this map relies on key == 0 marking an empty slot). Uses bulk puts through the direct view when available. + private static void zeroFill(PinotDataBuffer buffer, ByteBuffer view, long sizeBytes) { + if (view != null) { + int size = (int) sizeBytes; + for (int offset = 0; offset < size; offset += ZERO_CHUNK.length) { + view.put(offset, ZERO_CHUNK, 0, Math.min(ZERO_CHUNK.length, size - offset)); + } + } else { + for (long offset = 0; offset < sizeBytes; offset += Long.BYTES) { + buffer.putLong(offset, 0L); + } + } + } + + private static void closeBuffer(PinotDataBuffer buffer) { + OffHeapGroupByBufferPool.release(buffer); + } + + /// Flyweight entry for [#iterator()]. The same instance is reused across `next()` calls. + public static class Entry { + public int _rawKey; + public int _groupId; + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapLongGroupByResultHolder.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapLongGroupByResultHolder.java new file mode 100644 index 000000000000..7f346adc5cc9 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapLongGroupByResultHolder.java @@ -0,0 +1,196 @@ +/** + * 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.offheap; + +import com.google.common.base.Preconditions; +import java.nio.ByteBuffer; +import javax.annotation.concurrent.NotThreadSafe; +import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; + + +/// Off-heap implementation of [GroupByResultHolder] for long results, backed by a single direct-memory +/// [PinotDataBuffer]. Drop-in replacement for +/// [org.apache.pinot.core.query.aggregation.groupby.LongGroupByResultHolder] with identical semantics +/// observed through the interface: each group key indexes a fixed-width 8-byte slot, un-initialized slots return +/// the default value, and [GroupKeyGenerator#INVALID_ID] returns the default value on get and is silently +/// ignored on set. +/// +/// An initial capacity of 0 is legal (a zero-size buffer is allocated); in that state no getter or setter may +/// be called until [#ensureCapacity(int)] grows the holder. +/// +/// [#close()] releases the direct memory and is idempotent; the behavior of all other methods after +/// close is undefined. +/// +/// This class is single-threaded and not thread-safe. +@NotThreadSafe +public class OffHeapLongGroupByResultHolder implements GroupByResultHolder, AutoCloseable { + private static final long BYTES_PER_VALUE = Long.BYTES; + private static final int VALUE_SHIFT = 3; + private static final String BUFFER_DESCRIPTION = "OffHeapLongGroupByResultHolder"; + + private final int _maxCapacity; + private final long _defaultValue; + + private int _resultHolderCapacity; + private PinotDataBuffer _dataBuffer; + // Absolute-indexed direct view of _dataBuffer for the per-row hot path (monomorphic, intrinsified + // ByteBuffer access instead of the PinotDataBuffer wrapper); null when empty or beyond the 2GB view limit + private ByteBuffer _view; + private boolean _closed; + + /// Constructor for the class. + /// + /// @param initialCapacity Initial capacity of the result holder + /// @param maxCapacity Maximum capacity of the result holder + /// @param defaultValue Default value of un-initialized results + public OffHeapLongGroupByResultHolder(int initialCapacity, int maxCapacity, long defaultValue) { + _maxCapacity = maxCapacity; + _defaultValue = defaultValue; + + _resultHolderCapacity = initialCapacity; + _dataBuffer = OffHeapGroupByBufferPool.acquire(initialCapacity * BYTES_PER_VALUE, BUFFER_DESCRIPTION); + _view = createView(_dataBuffer, initialCapacity); + // Direct buffer contents are undefined, so always fill the allocated region with the default value + fillWithDefaultValue(0, initialCapacity); + } + + @Override + public void ensureCapacity(int capacity) { + Preconditions.checkArgument(capacity <= _maxCapacity); + + if (capacity > _resultHolderCapacity) { + int copyLength = _resultHolderCapacity; + // Cap the growth to maximum possible number of group keys. NOTE: _resultHolderCapacity (the bounds-guard + // reference) is updated only after the new buffer is successfully acquired. + int newCapacity = Math.min(Math.max(_resultHolderCapacity * 2, capacity), _maxCapacity); + + PinotDataBuffer current = _dataBuffer; + _dataBuffer = OffHeapGroupByBufferPool.acquire(newCapacity * BYTES_PER_VALUE, BUFFER_DESCRIPTION); + _view = createView(_dataBuffer, newCapacity); + _resultHolderCapacity = newCapacity; + if (copyLength > 0) { + current.copyTo(0, _dataBuffer, 0, copyLength * BYTES_PER_VALUE); + } + // Fill the newly extended region with the default value (direct buffer contents are undefined) + fillWithDefaultValue(copyLength, newCapacity); + closeBuffer(current); + } + } + + @Override + public double getDoubleResult(int groupKey) { + throw new UnsupportedOperationException(); + } + + @Override + public int getIntResult(int groupKey) { + throw new UnsupportedOperationException(); + } + + @Override + public long getLongResult(int groupKey) { + if (groupKey == GroupKeyGenerator.INVALID_ID) { + return _defaultValue; + } else { + // PinotDataBuffer access is unchecked: an out-of-range key would read arbitrary memory instead of the + // on-heap twin's ArrayIndexOutOfBoundsException, so guard the sizing contract with an assert + assert groupKey >= 0 && groupKey < _resultHolderCapacity : "groupKey " + groupKey + " out of bounds"; + ByteBuffer view = _view; + if (view != null) { + return view.getLong(groupKey << VALUE_SHIFT); + } + return _dataBuffer.getLong(groupKey * BYTES_PER_VALUE); + } + } + + @Override + public T getResult(int groupKey) { + throw new UnsupportedOperationException(); + } + + @Override + public void setValueForKey(int groupKey, double newValue) { + throw new UnsupportedOperationException(); + } + + @Override + public void setValueForKey(int groupKey, int newValue) { + throw new UnsupportedOperationException(); + } + + @Override + public void setValueForKey(int groupKey, long newValue) { + if (groupKey != GroupKeyGenerator.INVALID_ID) { + // See getLongResult: unchecked buffer access means an out-of-range key would corrupt memory, not throw + assert groupKey >= 0 && groupKey < _resultHolderCapacity : "groupKey " + groupKey + " out of bounds"; + ByteBuffer view = _view; + if (view != null) { + view.putLong(groupKey << VALUE_SHIFT, newValue); + } else { + _dataBuffer.putLong(groupKey * BYTES_PER_VALUE, newValue); + } + } + } + + @Override + public void setValueForKey(int groupKey, Object newValue) { + throw new UnsupportedOperationException(); + } + + public long getDefaultValue() { + return _defaultValue; + } + + @Override + public void close() { + if (_closed) { + return; + } + _closed = true; + closeBuffer(_dataBuffer); + // Null the buffer and view so any use-after-close (or a second release of a pooled buffer) fails loudly with + // an NPE instead of silently aliasing memory that the pool may have handed to another query + _dataBuffer = null; + _view = null; + } + + private void fillWithDefaultValue(int fromIndex, int toIndex) { + ByteBuffer view = _view; + if (view != null) { + // Intrinsified view puts are several times cheaper than the PinotDataBuffer wrapper accessors + for (int i = fromIndex; i < toIndex; i++) { + view.putLong(i << VALUE_SHIFT, _defaultValue); + } + } else { + for (int i = fromIndex; i < toIndex; i++) { + _dataBuffer.putLong(i * BYTES_PER_VALUE, _defaultValue); + } + } + } + + private static ByteBuffer createView(PinotDataBuffer dataBuffer, int capacity) { + return OffHeapGroupByUtils.createView(dataBuffer, capacity * BYTES_PER_VALUE); + } + + private static void closeBuffer(PinotDataBuffer buffer) { + OffHeapGroupByBufferPool.release(buffer); + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapLongGroupIdMap.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapLongGroupIdMap.java new file mode 100644 index 000000000000..b53e66cbe71a --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapLongGroupIdMap.java @@ -0,0 +1,311 @@ +/** + * 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.offheap; + +import com.google.common.base.Preconditions; +import java.nio.ByteBuffer; +import java.util.Iterator; +import java.util.NoSuchElementException; +import javax.annotation.concurrent.NotThreadSafe; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; + + +/// Off-heap hash table mapping a 64-bit raw group key to a dense int group id, modeled on ClickHouse key64 hash +/// tables and `DictionaryBasedGroupKeyGenerator.IntGroupIdMap`. It is an off-heap replacement for +/// `Long2IntOpenHashMap` style maps in group-by key generation. +/// +/// Group ids are assigned densely in insertion order: 0, 1, 2, ... [int)][#getGroupId(long,] returns the +/// existing id if the key is present (regardless of the upper bound); otherwise it assigns id `size()` if +/// `size() < groupIdUpperBound`, or returns [GroupKeyGenerator#INVALID_ID] without inserting. +/// +/// Implementation: open addressing with linear probing (step +1) over a power-of-two capacity with load factor +/// 0.5. Each slot is 16 bytes: [long key][int groupId][4 bytes unused]. An empty slot is identified by key == 0, so +/// the real key 0 is held out-of-band in a field (the ClickHouse zero-value trick) and does not occupy a slot. +/// Resizing doubles the capacity and rehashes all occupied slots; assigned group ids are never changed by a resize. +/// +/// Memory is direct (off-heap) via [PinotDataBuffer]. [#close()] releases it and is idempotent; +/// behavior of all other methods after close is undefined. +/// +/// This class is not thread-safe. +@NotThreadSafe +public class OffHeapLongGroupIdMap implements AutoCloseable { + public static final int INVALID_ID = GroupKeyGenerator.INVALID_ID; + + // Slot layout: [long key][int groupId][4 bytes unused] = 16 bytes + private static final int SLOT_SHIFT = 4; + private static final int GROUP_ID_OFFSET_IN_SLOT = Long.BYTES; + private static final int MIN_CAPACITY = 512; + // Largest power-of-two capacity representable as a positive int; expand() past this would overflow + private static final int MAX_CAPACITY = 1 << 30; + // Reusable zero block for bulk zero-filling freshly allocated buffers through the direct view + private static final byte[] ZERO_CHUNK = new byte[8192]; + + private PinotDataBuffer _buffer; + // Absolute-indexed direct view of _buffer for the per-row probe loop (monomorphic, intrinsified ByteBuffer + // access instead of the PinotDataBuffer wrapper); null when the table exceeds the 2GB view limit + private ByteBuffer _view; + private int _capacity; + private int _mask; + // Resize when the number of occupied slots exceeds this (i.e. load factor 0.5). The out-of-band zero key does + // not occupy a slot and is not counted here. + private int _maxOccupiedSlots; + private int _occupiedSlots; + private int _zeroKeyGroupId = INVALID_ID; + private boolean _closed; + + public OffHeapLongGroupIdMap(int expectedNumEntries) { + Preconditions.checkArgument(expectedNumEntries >= 0, "Invalid expectedNumEntries: %s", expectedNumEntries); + long desiredCapacity = Math.max(MIN_CAPACITY, (long) expectedNumEntries << 1); + _capacity = (int) Math.min(MAX_CAPACITY, Long.highestOneBit((desiredCapacity << 1) - 1)); + _mask = _capacity - 1; + _maxOccupiedSlots = _capacity >>> 1; + long sizeBytes = (long) _capacity << SLOT_SHIFT; + _buffer = allocate(sizeBytes); + _view = OffHeapGroupByUtils.createView(_buffer, sizeBytes); + zeroFill(_buffer, _view, sizeBytes); + } + + /// Returns the number of groups assigned so far, including the group for the raw key 0 if assigned. + public int size() { + return _zeroKeyGroupId != INVALID_ID ? _occupiedSlots + 1 : _occupiedSlots; + } + + /// Returns the amount of off-heap memory held by this map in bytes. + public long getOffHeapMemoryBytes() { + return (long) _capacity << SLOT_SHIFT; + } + + /// Returns the group id for the given raw key. If the key is present, always returns its id (even when + /// `size() >= groupIdUpperBound`). If absent and `size() < groupIdUpperBound`, assigns the next dense + /// id (`size()`) and returns it; otherwise returns [#INVALID_ID] without inserting. + public int getGroupId(long rawKey, int groupIdUpperBound) { + if (rawKey == 0) { + int zeroKeyGroupId = _zeroKeyGroupId; + if (zeroKeyGroupId != INVALID_ID) { + return zeroKeyGroupId; + } + // Zero key not assigned yet, so size() == _occupiedSlots here + int size = size(); + if (size < groupIdUpperBound) { + _zeroKeyGroupId = size; + return size; + } + return INVALID_ID; + } + ByteBuffer view = _view; + if (view == null) { + return getGroupIdSlow(rawKey, groupIdUpperBound); + } + // While the view exists, slot offsets fit in an int (view size <= Integer.MAX_VALUE) + int slot = (int) (mix(rawKey) & _mask); + while (true) { + int slotOffset = slot << SLOT_SHIFT; + long key = view.getLong(slotOffset); + if (key == rawKey) { + return view.getInt(slotOffset + GROUP_ID_OFFSET_IN_SLOT); + } + if (key == 0) { + int size = size(); + if (size >= groupIdUpperBound) { + return INVALID_ID; + } + view.putLong(slotOffset, rawKey); + view.putInt(slotOffset + GROUP_ID_OFFSET_IN_SLOT, size); + if (++_occupiedSlots > _maxOccupiedSlots) { + expand(); + } + return size; + } + slot = (slot + 1) & _mask; + } + } + + private int getGroupIdSlow(long rawKey, int groupIdUpperBound) { + PinotDataBuffer buffer = _buffer; + int slot = (int) (mix(rawKey) & _mask); + while (true) { + long slotOffset = (long) slot << SLOT_SHIFT; + long key = buffer.getLong(slotOffset); + if (key == rawKey) { + return buffer.getInt(slotOffset + GROUP_ID_OFFSET_IN_SLOT); + } + if (key == 0) { + int size = size(); + if (size >= groupIdUpperBound) { + return INVALID_ID; + } + buffer.putLong(slotOffset, rawKey); + buffer.putInt(slotOffset + GROUP_ID_OFFSET_IN_SLOT, size); + if (++_occupiedSlots > _maxOccupiedSlots) { + expand(); + } + return size; + } + slot = (slot + 1) & _mask; + } + } + + /// Returns an iterator over all (rawKey, groupId) entries in arbitrary slot order, with the zero-key entry (if + /// assigned) yielded last. Yields exactly [#size()] entries. + /// + /// NOTE: The returned [Entry] instance is a flyweight reused across `next()` calls; copy the values + /// out if they need to outlive the next call. + public Iterator iterator() { + return new Iterator<>() { + private final Entry _entry = new Entry(); + private int _slot; + private int _remainingOccupiedSlots = _occupiedSlots; + private boolean _returnZeroKey = _zeroKeyGroupId != INVALID_ID; + + @Override + public boolean hasNext() { + return _remainingOccupiedSlots > 0 || _returnZeroKey; + } + + @Override + public Entry next() { + if (_remainingOccupiedSlots > 0) { + long key; + long slotOffset; + do { + slotOffset = (long) _slot << SLOT_SHIFT; + key = _buffer.getLong(slotOffset); + _slot++; + } while (key == 0); + _entry._rawKey = key; + _entry._groupId = _buffer.getInt(slotOffset + GROUP_ID_OFFSET_IN_SLOT); + _remainingOccupiedSlots--; + return _entry; + } + if (_returnZeroKey) { + _returnZeroKey = false; + _entry._rawKey = 0; + _entry._groupId = _zeroKeyGroupId; + return _entry; + } + throw new NoSuchElementException(); + } + }; + } + + @Override + public void close() { + if (_closed) { + return; + } + _closed = true; + closeBuffer(_buffer); + // Null the buffer and view so any use-after-close (or a second release of a pooled buffer) fails loudly with + // an NPE instead of silently aliasing memory that the pool may have handed to another query + _buffer = null; + _view = null; + } + + private void expand() { + Preconditions.checkState(_capacity < MAX_CAPACITY, "Cannot expand beyond max capacity: %s", MAX_CAPACITY); + int newCapacity = _capacity << 1; + int newMask = newCapacity - 1; + long newSizeBytes = (long) newCapacity << SLOT_SHIFT; + PinotDataBuffer newBuffer = allocate(newSizeBytes); + ByteBuffer newView = OffHeapGroupByUtils.createView(newBuffer, newSizeBytes); + zeroFill(newBuffer, newView, newSizeBytes); + ByteBuffer oldView = _view; + if (oldView != null && newView != null) { + // Hot path: rehash through the direct views (slot offsets fit in an int while a view exists) + for (int slot = 0; slot < _capacity; slot++) { + int slotOffset = slot << SLOT_SHIFT; + long key = oldView.getLong(slotOffset); + if (key != 0) { + int newSlot = (int) (mix(key) & newMask); + int newSlotOffset = newSlot << SLOT_SHIFT; + while (newView.getLong(newSlotOffset) != 0) { + newSlot = (newSlot + 1) & newMask; + newSlotOffset = newSlot << SLOT_SHIFT; + } + newView.putLong(newSlotOffset, key); + newView.putInt(newSlotOffset + GROUP_ID_OFFSET_IN_SLOT, oldView.getInt(slotOffset + GROUP_ID_OFFSET_IN_SLOT)); + } + } + } else { + for (int slot = 0; slot < _capacity; slot++) { + long slotOffset = (long) slot << SLOT_SHIFT; + long key = _buffer.getLong(slotOffset); + if (key != 0) { + int newSlot = (int) (mix(key) & newMask); + long newSlotOffset = (long) newSlot << SLOT_SHIFT; + while (newBuffer.getLong(newSlotOffset) != 0) { + newSlot = (newSlot + 1) & newMask; + newSlotOffset = (long) newSlot << SLOT_SHIFT; + } + newBuffer.putLong(newSlotOffset, key); + newBuffer.putInt(newSlotOffset + GROUP_ID_OFFSET_IN_SLOT, + _buffer.getInt(slotOffset + GROUP_ID_OFFSET_IN_SLOT)); + } + } + } + closeBuffer(_buffer); + _buffer = newBuffer; + _capacity = newCapacity; + _mask = newMask; + _maxOccupiedSlots = newCapacity >>> 1; + _view = newView; + } + + private static PinotDataBuffer allocate(long sizeBytes) { + return OffHeapGroupByBufferPool.acquire(sizeBytes, "OffHeapLongGroupIdMap hash table"); + } + + /// Zero-fills a freshly allocated buffer (contents of [PinotDataBuffer#allocateDirect] are undefined, and + /// this map relies on key == 0 marking an empty slot). Uses bulk puts through the direct view when available. + private static void zeroFill(PinotDataBuffer buffer, ByteBuffer view, long sizeBytes) { + if (view != null) { + int size = (int) sizeBytes; + for (int offset = 0; offset < size; offset += ZERO_CHUNK.length) { + view.put(offset, ZERO_CHUNK, 0, Math.min(ZERO_CHUNK.length, size - offset)); + } + } else { + for (long offset = 0; offset < sizeBytes; offset += Long.BYTES) { + buffer.putLong(offset, 0L); + } + } + } + + private static void closeBuffer(PinotDataBuffer buffer) { + OffHeapGroupByBufferPool.release(buffer); + } + + /// Murmur3 fmix64 finalizer. Bijective over longs, so a non-zero key always hashes deterministically, and the + /// rehash on resize simply recomputes it. + private static long mix(long key) { + long h = key; + h ^= h >>> 33; + h *= 0xff51afd7ed558ccdL; + h ^= h >>> 33; + h *= 0xc4ceb9fe1a85ec53L; + h ^= h >>> 33; + return h; + } + + /// Flyweight entry for [#iterator()]. The same instance is reused across `next()` calls. + public static class Entry { + public long _rawKey; + public int _groupId; + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapUltraLogLogGroupByResultHolder.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapUltraLogLogGroupByResultHolder.java new file mode 100644 index 000000000000..6bff07ceabfd --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapUltraLogLogGroupByResultHolder.java @@ -0,0 +1,319 @@ +/** + * 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.offheap; + +import com.dynatrace.hash4j.distinctcount.UltraLogLog; +import com.google.common.base.Preconditions; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import javax.annotation.concurrent.NotThreadSafe; +import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator; +import org.apache.pinot.core.query.aggregation.groupby.ObjectGroupByResultHolder; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; + + +/// Off-heap implementation of [GroupByResultHolder] for `DISTINCT_COUNT_ULL` group-by state. Each group's +/// [UltraLogLog] register array (`2^p` bytes) lives in direct memory instead of a per-group heap object, so a +/// group-by with many groups carries no per-group heap state and no GC pressure from the sketches. +/// +/// ### Storage layout +/// Slots are assigned append-only on first touch through a `groupKey -> slotId` indirection, so direct memory +/// grows with the number of groups actually seen (like the on-heap holder's lazy per-group allocation), never +/// with the group-count upper bound. Slots live in fixed-size chunks acquired from [OffHeapGroupByBufferPool]; +/// chunks are never resized or moved, so slot addresses are stable. An all-zero slot is exactly the state of an +/// empty [UltraLogLog], so slots are zero-filled on assignment. +/// +/// ### Register update math +/// [#add(int, long)] applies hash4j's `UltraLogLog.add(long)` register update (including `pack`/`unpack`) +/// vendored verbatim from hash4j 0.30.0 (Apache License 2.0), since [UltraLogLog] is final and only operates on +/// heap `byte[]` state. `OffHeapUltraLogLogGroupByResultHolderTest` pins the vendored math byte-identical to the +/// library across precisions. +/// +/// ### Modes +/// The owning aggregation function stores different state types depending on the input column: raw values hash +/// into ULL registers (off-heap slots, via [#touch(int)] / [#add(int, long)]), while dictionary-encoded columns +/// and pre-serialized ULL BYTES columns keep per-group heap objects through the generic +/// [#getResult(int)] / [#setValueForKey(int, Object)] API — those are routed to a lazily-created on-heap +/// [ObjectGroupByResultHolder] delegate. A segment uses exactly one mode per holder (one column encoding per +/// segment); this is asserted, not branched per row. [#getResult(int)] in slot mode materializes a fresh heap +/// [UltraLogLog] copy of the slot (extraction-time only), and returns `null` for untouched groups to match the +/// on-heap holder. +/// +/// [#close()] releases the direct memory and is idempotent; the behavior of all other methods after close is +/// undefined. This class is single-threaded and not thread-safe. +@NotThreadSafe +public class OffHeapUltraLogLogGroupByResultHolder implements GroupByResultHolder, AutoCloseable { + // Matches OffHeapBytesGroupIdMap's chunk size; slots of this size and larger (p >= 18) get one slot per chunk + private static final int TARGET_CHUNK_BYTES = 256 * 1024; + private static final String BUFFER_DESCRIPTION = "OffHeapUltraLogLogGroupByResultHolder"; + + private final int _p; + private final int _slotBytes; + private final int _slotsPerChunkShift; + private final int _slotIndexMask; + private final long _chunkBytes; + // q = 64 - p, hoisted for the vendored register update (see hash4j UltraLogLog.add(long)) + private final int _q; + private final int _maxCapacity; + + private int _resultHolderCapacity; + // groupKey -> slotId; -1 = group never touched (getResult returns null, matching the on-heap holder) + private int[] _slotIds; + private int _numSlots; + private List _chunkBuffers = new ArrayList<>(); + // Absolute-indexed direct views of the chunks for the per-row hot path; a null element means the view limit + // was exceeded (test hook) and accesses fall back to the PinotDataBuffer wrapper + private List _chunkViews = new ArrayList<>(); + private ObjectGroupByResultHolder _delegate; + private boolean _closed; + + /// Constructor for the class. + /// + /// @param p UltraLogLog precision parameter (slot size is `2^p` bytes) + /// @param initialCapacity Initial capacity of the result holder + /// @param maxCapacity Maximum capacity of the result holder + public OffHeapUltraLogLogGroupByResultHolder(int p, int initialCapacity, int maxCapacity) { + // The aggregation function validates p at plan time; re-check here because p sizes the direct-memory slots + // (1 << p) without going through UltraLogLog.create's own bound check, and an out-of-range p would either + // allocate absurd chunks (p up to 30) or let register indexes walk outside the slot (p > 30, int-shift wrap) + Preconditions.checkArgument(p >= 3 && p <= 26, "Invalid UltraLogLog p: %s, must be in [3, 26]", p); + _p = p; + _slotBytes = 1 << p; + int slotsPerChunk = Math.max(1, TARGET_CHUNK_BYTES / _slotBytes); + _slotsPerChunkShift = Integer.numberOfTrailingZeros(slotsPerChunk); + _slotIndexMask = slotsPerChunk - 1; + _chunkBytes = (long) slotsPerChunk * _slotBytes; + _q = 64 - p; + _maxCapacity = maxCapacity; + + _resultHolderCapacity = initialCapacity; + _slotIds = new int[initialCapacity]; + Arrays.fill(_slotIds, GroupKeyGenerator.INVALID_ID); + } + + @Override + public void ensureCapacity(int capacity) { + Preconditions.checkArgument(capacity <= _maxCapacity); + + if (capacity > _resultHolderCapacity) { + int copyLength = _resultHolderCapacity; + int newCapacity = Math.min(Math.max(_resultHolderCapacity * 2, capacity), _maxCapacity); + // _slotIds is null in delegate mode (the delegate tracks its own capacity) + if (_slotIds != null) { + _slotIds = Arrays.copyOf(_slotIds, newCapacity); + Arrays.fill(_slotIds, copyLength, newCapacity, GroupKeyGenerator.INVALID_ID); + } + _resultHolderCapacity = newCapacity; + } + if (_delegate != null) { + _delegate.ensureCapacity(capacity); + } + } + + /// Ensures the group has an (all-zero) slot, mirroring the on-heap path's eager `UltraLogLog.create(p)` on + /// first access so an untouched-vs-empty distinction never diverges between the two modes. + public void touch(int groupKey) { + if (groupKey != GroupKeyGenerator.INVALID_ID) { + slotIdFor(groupKey); + } + } + + /// Adds a 64-bit hash value into the group's off-heap ULL registers. Vendored verbatim from hash4j 0.30.0 + /// `UltraLogLog.add(long)` with the state array replaced by the group's slot. + public void add(int groupKey, long hashValue) { + if (groupKey == GroupKeyGenerator.INVALID_ID) { + return; + } + int slotId = slotIdFor(groupKey); + int chunkIndex = slotId >>> _slotsPerChunkShift; + int slotOffset = (slotId & _slotIndexMask) * _slotBytes; + int idx = (int) (hashValue >>> _q); + int nlz = Long.numberOfLeadingZeros(~(~hashValue << -_q)); // nlz in {0, 1, ..., 64-p} + int registerOffset = slotOffset + idx; + ByteBuffer view = _chunkViews.get(chunkIndex); + if (view != null) { + byte oldRegister = view.get(registerOffset); + long hashPrefix = unpack(oldRegister) | (1L << (nlz + ~_q)); // (nlz + ~q) = (nlz + p - 1) mod 64 + view.put(registerOffset, pack(hashPrefix)); + } else { + PinotDataBuffer chunk = _chunkBuffers.get(chunkIndex); + byte oldRegister = chunk.getByte(registerOffset); + long hashPrefix = unpack(oldRegister) | (1L << (nlz + ~_q)); + chunk.putByte(registerOffset, pack(hashPrefix)); + } + } + + @Override + public double getDoubleResult(int groupKey) { + throw new UnsupportedOperationException(); + } + + @Override + public int getIntResult(int groupKey) { + throw new UnsupportedOperationException(); + } + + @Override + public long getLongResult(int groupKey) { + throw new UnsupportedOperationException(); + } + + /// In delegate mode, returns the delegate's live per-group object. In slot mode, returns a **snapshot**: a + /// fresh heap [UltraLogLog] copy of the slot, so mutating the returned object does NOT update the slot. Slot + /// mode is therefore only correct for extraction-style reads; per-row read-modify-write callers (the + /// pre-serialized-BYTES merge path) always run in delegate mode because their first write goes through + /// [#setValueForKey(int, Object)]. + @Override + @SuppressWarnings("unchecked") + public T getResult(int groupKey) { + if (_delegate != null) { + return _delegate.getResult(groupKey); + } + if (groupKey == GroupKeyGenerator.INVALID_ID) { + return null; + } + // See OffHeapLongGroupByResultHolder: unchecked buffer access means an out-of-range key would read arbitrary + // memory instead of throwing, so guard the sizing contract with an assert + assert groupKey >= 0 && groupKey < _resultHolderCapacity : "groupKey " + groupKey + " out of bounds"; + int slotId = _slotIds[groupKey]; + if (slotId < 0) { + return null; + } + int chunkIndex = slotId >>> _slotsPerChunkShift; + int slotOffset = (slotId & _slotIndexMask) * _slotBytes; + byte[] state = new byte[_slotBytes]; + ByteBuffer view = _chunkViews.get(chunkIndex); + if (view != null) { + view.get(slotOffset, state); + } else { + _chunkBuffers.get(chunkIndex).copyTo(slotOffset, state, 0, _slotBytes); + } + return (T) UltraLogLog.wrap(state); + } + + @Override + public void setValueForKey(int groupKey, double newValue) { + throw new UnsupportedOperationException(); + } + + @Override + public void setValueForKey(int groupKey, int newValue) { + throw new UnsupportedOperationException(); + } + + @Override + public void setValueForKey(int groupKey, long newValue) { + throw new UnsupportedOperationException(); + } + + @Override + public void setValueForKey(int groupKey, Object newValue) { + if (groupKey == GroupKeyGenerator.INVALID_ID) { + return; + } + // Heap-object mode (dictionary id wrappers, pre-serialized ULL merges): route to the on-heap delegate. A + // segment uses exactly one mode per holder; a violation would silently drop state (getResult prefers the + // delegate), so enforce it hard — this runs once per group, not per row. + Preconditions.checkState(_numSlots == 0, + "Off-heap ULL slots and heap-object delegate used on the same holder"); + if (_delegate == null) { + _delegate = new ObjectGroupByResultHolder(_resultHolderCapacity, _maxCapacity); + // Delegate mode never uses the slot indirection; nulling it frees 4 bytes/group and turns any stray slot + // access into a loud NPE instead of a silent divergence + _slotIds = null; + } + _delegate.setValueForKey(groupKey, newValue); + } + + @Override + public void close() { + if (_closed) { + return; + } + _closed = true; + try { + RuntimeException firstFailure = null; + for (PinotDataBuffer chunk : _chunkBuffers) { + // Release every chunk even if one release fails, matching ResourceTrackingGroupKeyGenerator's policy + try { + OffHeapGroupByBufferPool.release(chunk); + } catch (RuntimeException e) { + if (firstFailure == null) { + firstFailure = e; + } + } + } + if (firstFailure != null) { + throw firstFailure; + } + } finally { + // Null the buffers and views so any use-after-close (or a second release of a pooled buffer) fails loudly + // with an NPE instead of silently aliasing memory that the pool may have handed to another query + _chunkBuffers = null; + _chunkViews = null; + _slotIds = null; + _delegate = null; + } + } + + private int slotIdFor(int groupKey) { + // See getResult: unchecked buffer access means an out-of-range key would corrupt memory, not throw + assert groupKey >= 0 && groupKey < _resultHolderCapacity : "groupKey " + groupKey + " out of bounds"; + assert _delegate == null : "off-heap ULL slots and heap-object delegate used on the same holder"; + int slotId = _slotIds[groupKey]; + if (slotId >= 0) { + return slotId; + } + slotId = _numSlots++; + _slotIds[groupKey] = slotId; + int chunkIndex = slotId >>> _slotsPerChunkShift; + if (chunkIndex == _chunkBuffers.size()) { + PinotDataBuffer chunk = OffHeapGroupByBufferPool.acquire(_chunkBytes, BUFFER_DESCRIPTION); + _chunkBuffers.add(chunk); + _chunkViews.add(OffHeapGroupByUtils.createView(chunk, _chunkBytes)); + } + // Pooled buffers come back dirty; an all-zero slot is exactly an empty UltraLogLog state + int slotOffset = (slotId & _slotIndexMask) * _slotBytes; + ByteBuffer view = _chunkViews.get(chunkIndex); + if (view != null) { + for (int i = 0; i < _slotBytes; i += Long.BYTES) { + view.putLong(slotOffset + i, 0L); + } + } else { + PinotDataBuffer chunk = _chunkBuffers.get(chunkIndex); + for (int i = 0; i < _slotBytes; i += Long.BYTES) { + chunk.putLong(slotOffset + i, 0L); + } + } + return slotId; + } + + // pack/unpack vendored verbatim from hash4j 0.30.0 UltraLogLog (Apache License 2.0) + private static long unpack(byte register) { + return (4L | (register & 3)) << ((register >>> 2) - 2); + } + + private static byte pack(long hashPrefix) { + int nlz = Long.numberOfLeadingZeros(hashPrefix) + 1; + return (byte) ((-nlz << 2) | ((hashPrefix << nlz) >>> 62)); + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/ResourceTrackingGroupKeyGenerator.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/ResourceTrackingGroupKeyGenerator.java new file mode 100644 index 000000000000..14a762e9d0f6 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/ResourceTrackingGroupKeyGenerator.java @@ -0,0 +1,109 @@ +/** + * 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.offheap; + +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import javax.annotation.concurrent.NotThreadSafe; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// [GroupKeyGenerator] wrapper that owns every off-heap resource created for one group-by execution: the wrapped +/// generator's own off-heap key table (released through the delegate's `close()`) plus any registered off-heap +/// result holders. The existing operator-level `GroupKeyGenerator.close()` call sites (segment trim/sort paths, +/// combine operators, exception guards) thus release all off-heap memory without knowing about holders. +/// +/// All [GroupKeyGenerator] methods delegate as-is; delegation happens at block granularity, so the extra +/// virtual call is not on the per-row hot path. `close()` is idempotent and closes the delegate first, then every +/// registered resource, attempting all of them even if some fail. +/// +/// Not thread-safe: intended for the single-threaded per-segment group-by execution, mirroring the wrapped +/// generator. In the filtered-aggregation case the same instance is shared sequentially across executors and closed +/// exactly once by the operator. +@NotThreadSafe +public class ResourceTrackingGroupKeyGenerator implements GroupKeyGenerator { + private static final Logger LOGGER = LoggerFactory.getLogger(ResourceTrackingGroupKeyGenerator.class); + + private final GroupKeyGenerator _delegate; + private final List _resources = new ArrayList<>(); + private boolean _closed; + + public ResourceTrackingGroupKeyGenerator(GroupKeyGenerator delegate) { + _delegate = delegate; + } + + /// Registers an off-heap resource to be released when this generator is closed. + public void register(AutoCloseable resource) { + _resources.add(resource); + } + + @Override + public int getGlobalGroupKeyUpperBound() { + return _delegate.getGlobalGroupKeyUpperBound(); + } + + @Override + public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { + _delegate.generateKeysForBlock(valueBlock, groupKeys); + } + + @Override + public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { + _delegate.generateKeysForBlock(valueBlock, groupKeys); + } + + @Override + public int getCurrentGroupKeyUpperBound() { + return _delegate.getCurrentGroupKeyUpperBound(); + } + + @Override + public Iterator getGroupKeys() { + return _delegate.getGroupKeys(); + } + + @Override + public int getNumKeys() { + return _delegate.getNumKeys(); + } + + @Override + public void close() { + if (_closed) { + return; + } + _closed = true; + try { + _delegate.close(); + } catch (Exception e) { + LOGGER.warn("Caught exception while closing group key generator: {}", _delegate.getClass().getName(), e); + } + for (AutoCloseable resource : _resources) { + try { + resource.close(); + } catch (Exception e) { + LOGGER.warn("Caught exception while closing off-heap group-by resource: {}", resource.getClass().getName(), e); + } + } + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java index 83fd2ef62ebb..51aa72ca10f1 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java @@ -115,6 +115,8 @@ public class QueryContext { private int _minInitialIndexedTableCapacity = Server.DEFAULT_QUERY_EXECUTOR_MIN_INITIAL_INDEXED_TABLE_CAPACITY; // Limit of number of groups stored in each segment private int _numGroupsLimit = Server.DEFAULT_QUERY_EXECUTOR_NUM_GROUPS_LIMIT; + // Whether to store group-by key tables and fixed-width result holders in off-heap (direct) memory + private boolean _groupByOffHeap = Server.DEFAULT_QUERY_EXECUTOR_GROUPBY_OFF_HEAP; // Warning threshold of number of groups stored in each segment private int _numGroupsWarningLimit = Server.DEFAULT_QUERY_EXECUTOR_NUM_GROUPS_WARN_LIMIT; // Minimum number of groups to keep per segment when trimming groups for SQL GROUP BY @@ -446,6 +448,14 @@ public void setNumGroupsLimit(int numGroupsLimit) { _numGroupsLimit = numGroupsLimit; } + public boolean isGroupByOffHeap() { + return _groupByOffHeap; + } + + public void setGroupByOffHeap(boolean groupByOffHeap) { + _groupByOffHeap = groupByOffHeap; + } + public int getNumGroupsWarningLimit() { return _numGroupsWarningLimit; } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperatorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperatorTest.java index abb6bcf8e72e..169012509923 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperatorTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperatorTest.java @@ -38,6 +38,7 @@ import org.apache.pinot.core.plan.PlanNode; import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2; import org.apache.pinot.core.plan.maker.PlanMaker; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapGroupByBufferPool; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; @@ -46,6 +47,7 @@ import org.apache.pinot.segment.spi.IndexSegment; import org.apache.pinot.segment.spi.SegmentContext; import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.data.FieldSpec; @@ -477,6 +479,93 @@ private BaseCombineOperator route(String query, int flushThreshold) { return new CombinePlanNode(planNodes, queryContext, EXECUTOR, block -> { }).run(); } + /// With `groupByOffHeap`, every per-segment group-key table and fixed-width result holder lives in direct + /// memory owned by the block's group key generator. On the streaming path the generator is closed on the + /// producing worker thread by [StreamingGroupByCombineOperator#detachFromWorkerThreadState], so after a fully + /// consumed query no direct memory may remain allocated. High-cardinality segments (> the array-based + /// threshold of 10000 groups) force the map-based tier, which is the one the off-heap mode replaces. + @Test + public void testOffHeapStreamingGroupByReleasesDirectMemory() + throws Exception { + runOffHeapStreamingAndAssertNoDirectMemoryLeak(false); + } + + /// The abandonment variant: the consumer stops after the first flushed block, leaving undrained blocks in the + /// hand-off queue and workers mid-stream. Because raw results are detached (and their generators closed) on + /// the worker thread *before* hand-off, abandoned queued blocks hold no off-heap state, and the exception + /// guards in the per-segment operators cover workers interrupted mid-segment — so direct memory must still + /// return to the baseline after stop(). + @Test + public void testOffHeapStreamingGroupByReleasesDirectMemoryOnEarlyStop() + throws Exception { + runOffHeapStreamingAndAssertNoDirectMemoryLeak(true); + } + + private void runOffHeapStreamingAndAssertNoDirectMemoryLeak(boolean abandonAfterFirstBlock) + throws Exception { + int numSegments = 4; + int numGroups = 12_000; + int flushThreshold = 1000; + + File offHeapDir = new File(FileUtils.getTempDirectory(), "StreamingGroupByCombineOperatorTest_offHeap"); + FileUtils.deleteDirectory(offHeapDir); + List segments = new ArrayList<>(numSegments); + try { + for (int i = 0; i < numSegments; i++) { + segments.add(createHighCardinalitySegment(offHeapDir, i, numGroups)); + } + // Disable pooling so every released buffer is freed immediately and the usage returns to the exact baseline + OffHeapGroupByBufferPool.setMaxBytesPerThread(0); + long directBufferBaseline = PinotDataBuffer.getDirectBufferUsage(); + + QueryContext queryContext = QueryContextConverterUtils.getQueryContext( + "SELECT groupColumn, SUM(intColumn) FROM testTable GROUP BY groupColumn LIMIT " + numGroups); + queryContext.setEndTimeMs(System.currentTimeMillis() + Server.DEFAULT_QUERY_EXECUTOR_TIMEOUT_MS); + queryContext.setGroupByOffHeap(true); + + List operators = new ArrayList<>(numSegments); + for (IndexSegment segment : segments) { + operators.add(PLAN_MAKER.makeSegmentPlanNode(new SegmentContext(segment), queryContext).run()); + } + StreamingGroupByCombineOperator combineOperator = + new StreamingGroupByCombineOperator(operators, queryContext, EXECUTOR, flushThreshold); + + Map groupSums = new HashMap<>(); + combineOperator.start(); + try { + BaseResultsBlock block = combineOperator.nextBlock(); + while (!(block instanceof MetadataResultsBlock)) { + assertNull(block.getErrorMessages(), "Expected no errors but got: " + block.getErrorMessages()); + for (Object[] row : ((GroupByResultsBlock) block).getRows()) { + groupSums.merge((int) row[0], ((Number) row[1]).doubleValue(), Double::sum); + } + if (abandonAfterFirstBlock) { + break; + } + block = combineOperator.nextBlock(); + } + } finally { + // stop() joins the worker threads, so once it returns every generator close has happened + combineOperator.stop(); + } + + if (!abandonAfterFirstBlock) { + assertEquals(groupSums.size(), numGroups, "Wrong number of groups"); + for (int g = 0; g < numGroups; g++) { + assertEquals(groupSums.get(g), numSegments * (double) (g + 1), 0.001, "Incorrect sum for group " + g); + } + } + assertEquals(PinotDataBuffer.getDirectBufferUsage(), directBufferBaseline, + "Off-heap group-by state leaked on the streaming path (abandonAfterFirstBlock=" + abandonAfterFirstBlock + + ")"); + } finally { + for (IndexSegment segment : segments) { + segment.destroy(); + } + FileUtils.deleteDirectory(offHeapDir); + } + } + private List buildOperators(QueryContext queryContext) { List operators = new ArrayList<>(NUM_SEGMENTS); for (IndexSegment indexSegment : _indexSegments) { diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunctionTest.java index 0abed9d56903..0460312c6053 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunctionTest.java @@ -38,11 +38,23 @@ import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertThrows; import static org.testng.Assert.assertTrue; public class DistinctCountULLAggregationFunctionTest { + @Test + public void testOutOfRangePRejectedAtPlanTime() { + // p sizes the off-heap holder's direct-memory slots (1 << p), so it must be validated here at plan time + // rather than only inside UltraLogLog.create + for (int p : new int[]{-1, 0, 2, 27, 32}) { + assertThrows(IllegalArgumentException.class, + () -> new DistinctCountULLAggregationFunction(List.of(ExpressionContext.forIdentifier("col"), + ExpressionContext.forLiteral(Literal.intValue(p))), false)); + } + } + @Test public void testCanUseStarTreeDefaultP() { DistinctCountULLAggregationFunction function = new DistinctCountULLAggregationFunction( diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryNullGroupCountRegressionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryNullGroupCountRegressionTest.java new file mode 100644 index 000000000000..ef494b951a79 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryNullGroupCountRegressionTest.java @@ -0,0 +1,187 @@ +/** + * 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; + +import java.io.File; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.core.operator.BaseProjectOperator; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.plan.DocIdSetPlanNode; +import org.apache.pinot.core.plan.ProjectPlanNode; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.SegmentContext; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + +/// Regression test for [NoDictionarySingleColumnGroupKeyGenerator] counting the null group. +/// +/// For primitive stored types (INT/LONG/FLOAT/DOUBLE) with null handling enabled, the null group lives +/// *outside* the primitive key map, but it still takes the next dense group id. `getNumKeys()` and +/// `getCurrentGroupKeyUpperBound()` used to return only the map size, so once a null group was assigned the +/// reported upper bound equaled an already-issued group id. [DefaultGroupByExecutor#process] sizes the result +/// holders with `ensureCapacity(getCurrentGroupKeyUpperBound())`, so on a segment whose distinct-value count +/// exceeds the initial holder capacity the aggregation wrote one slot past the array — +/// an `ArrayIndexOutOfBoundsException` on the default on-heap path. +/// +/// The fixture makes the failure deterministic: the null appears in the very first row (null group id 0), the +/// distinct non-null values exceed the (artificially low) `maxInitialResultHolderCapacity`, and everything fits +/// in one block, so the pre-fix under-count always leaves the highest group id out of the holder. +public class NoDictionaryNullGroupCountRegressionTest { + private static final File TEMP_DIR = + new File(FileUtils.getTempDirectory(), "NoDictionaryNullGroupCountRegressionTest"); + private static final String RAW_TABLE_NAME = "testTable"; + private static final String SEGMENT_NAME = "testSegment"; + + private static final String INT_COLUMN = "nInt"; + private static final String LONG_COLUMN = "nLong"; + private static final String FLOAT_COLUMN = "nFloat"; + private static final String DOUBLE_COLUMN = "nDouble"; + private static final String[] COLUMNS = {INT_COLUMN, LONG_COLUMN, FLOAT_COLUMN, DOUBLE_COLUMN}; + + private static final int NUM_RECORDS = 200; + // 24 distinct non-null values (pool indexes 1..24) plus the null group = 25 groups + private static final int VALUE_POOL_SIZE = 25; + private static final int NUM_GROUPS = VALUE_POOL_SIZE; + // Far below the number of groups, so the holders must grow to exactly the reported upper bound + private static final int MAX_INITIAL_RESULT_HOLDER_CAPACITY = 8; + + private IndexSegment _indexSegment; + + @BeforeClass + public void setUp() + throws Exception { + FileUtils.deleteDirectory(TEMP_DIR); + + Schema schema = new Schema.SchemaBuilder().setSchemaName(RAW_TABLE_NAME) + .addSingleValueDimension(INT_COLUMN, DataType.INT) + .addSingleValueDimension(LONG_COLUMN, DataType.LONG) + .addSingleValueDimension(FLOAT_COLUMN, DataType.FLOAT) + .addSingleValueDimension(DOUBLE_COLUMN, DataType.DOUBLE) + .build(); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME) + .setNoDictionaryColumns(List.of(COLUMNS)).build(); + + List records = new ArrayList<>(NUM_RECORDS); + for (int i = 0; i < NUM_RECORDS; i++) { + // Pool index 0 is the null value, and row 0 uses it, so the null group takes dense group id 0 and every + // later distinct value pushes the maximum issued group id one past the (pre-fix) reported upper bound + int poolIndex = i % VALUE_POOL_SIZE; + GenericRow record = new GenericRow(); + record.putValue(INT_COLUMN, poolIndex == 0 ? null : poolIndex * 3 - 15); + record.putValue(LONG_COLUMN, poolIndex == 0 ? null : poolIndex * 1_000_003L); + record.putValue(FLOAT_COLUMN, poolIndex == 0 ? null : (poolIndex - 5) * 0.25f); + record.putValue(DOUBLE_COLUMN, poolIndex == 0 ? null : (poolIndex - 5) * 0.5d); + records.add(record); + } + + SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(tableConfig, schema); + segmentGeneratorConfig.setTableName(RAW_TABLE_NAME); + segmentGeneratorConfig.setSegmentName(SEGMENT_NAME); + segmentGeneratorConfig.setDefaultNullHandlingEnabled(true); + segmentGeneratorConfig.setOutDir(TEMP_DIR.getPath()); + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(segmentGeneratorConfig, new GenericRowRecordReader(records)); + driver.build(); + _indexSegment = ImmutableSegmentLoader.load(new File(TEMP_DIR, SEGMENT_NAME), ReadMode.mmap); + } + + @DataProvider(name = "primitiveColumns") + public Object[][] primitiveColumns() { + Object[][] result = new Object[COLUMNS.length][]; + for (int i = 0; i < COLUMNS.length; i++) { + result[i] = new Object[]{COLUMNS[i]}; + } + return result; + } + + @Test(dataProvider = "primitiveColumns") + public void testNullGroupCountedInUpperBound(String column) { + QueryContext queryContext = + QueryContextConverterUtils.getQueryContext("SELECT COUNT(*) FROM testTable GROUP BY " + column); + queryContext.setNullHandlingEnabled(true); + queryContext.setMaxInitialResultHolderCapacity(MAX_INITIAL_RESULT_HOLDER_CAPACITY); + ExpressionContext[] groupByExpressions = {ExpressionContext.forIdentifier(column)}; + + ProjectPlanNode projectPlanNode = new ProjectPlanNode(new SegmentContext(_indexSegment), queryContext, + List.of(groupByExpressions), DocIdSetPlanNode.MAX_DOC_PER_CALL); + BaseProjectOperator projectOperator = projectPlanNode.run(); + DefaultGroupByExecutor groupByExecutor = + new DefaultGroupByExecutor(queryContext, groupByExpressions, projectOperator); + + // Pre-fix this throws ArrayIndexOutOfBoundsException: the null group is assigned first (group id 0), the + // 24 distinct values take ids 1..24, but the reported upper bound was 24 (the map size), so the result + // holder never grew to cover group id 24 + ValueBlock valueBlock; + while ((valueBlock = projectOperator.nextBlock()) != null) { + groupByExecutor.process(valueBlock); + } + + GroupKeyGenerator groupKeyGenerator = groupByExecutor.getGroupKeyGenerator(); + assertEquals(groupKeyGenerator.getNumKeys(), NUM_GROUPS, "getNumKeys() must count the out-of-map null group"); + assertEquals(groupKeyGenerator.getCurrentGroupKeyUpperBound(), NUM_GROUPS, + "getCurrentGroupKeyUpperBound() must count the out-of-map null group"); + + // The iterators emit the null group too, and every issued id stays below the reported upper bound + int numKeys = 0; + int numNullKeys = 0; + int maxGroupId = -1; + Iterator groupKeys = groupKeyGenerator.getGroupKeys(); + while (groupKeys.hasNext()) { + GroupKeyGenerator.GroupKey groupKey = groupKeys.next(); + numKeys++; + if (groupKey._keys[0] == null) { + numNullKeys++; + } + maxGroupId = Math.max(maxGroupId, groupKey._groupId); + } + assertEquals(numKeys, NUM_GROUPS); + assertEquals(numNullKeys, 1, "Exactly one null group expected"); + assertTrue(maxGroupId < groupKeyGenerator.getCurrentGroupKeyUpperBound(), + "Issued group id " + maxGroupId + " must stay below the upper bound"); + } + + @AfterClass + public void tearDown() + throws Exception { + _indexSegment.destroy(); + FileUtils.deleteDirectory(TEMP_DIR); + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/OffHeapGroupKeyGeneratorParityTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/OffHeapGroupKeyGeneratorParityTest.java new file mode 100644 index 000000000000..aa0d9ac31692 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/OffHeapGroupKeyGeneratorParityTest.java @@ -0,0 +1,645 @@ +/** + * 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; + +import java.io.File; +import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.TreeSet; +import org.apache.commons.io.FileUtils; +import org.apache.commons.lang3.StringUtils; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.core.operator.BaseProjectOperator; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.plan.DocIdSetPlanNode; +import org.apache.pinot.core.plan.ProjectPlanNode; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.SegmentContext; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.utils.CommonConstants.Server; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + +/// Generator-level differential test: for every group key generator variant, the on-heap and off-heap instances +/// are driven over the same projection blocks and must emit identical group id arrays, identical +/// (groupId -> keys) mappings, and matching counts. +/// +/// Each generator gets its own fresh [BaseProjectOperator] over the same immutable segment: the map-based +/// single-MV-column path overwrites the block's cached dictionary-id arrays in place (in both modes), so two +/// generators must not share one operator's block cache. The blocks produced by the two operators are identical +/// because the segment and query are. +/// +/// Null-group counting: both modes count the null group in `getNumKeys()` / `getCurrentGroupKeyUpperBound()` — +/// the on-heap map-size-based counting used to exclude it for primitive stored types (INT/LONG/FLOAT/DOUBLE, +/// whose null group lives outside the primitive map), which under-sized result holders; that was fixed alongside +/// the off-heap work, so the counts must now match exactly in every mode and for every stored type. +/// +/// Dictionary-based holder selection (arrayBasedThreshold = 10_000, default numGroupsLimit = 100_000), against +/// dict columns `s1..s10` of cardinality 100 and MV columns `m1`/`m2` of cardinality 100: +/// - `s1` -> 100 -> ARRAY_BASED (on-heap in both modes by design) +/// - `s1,s2,s3` -> 10^6 -> INT_MAP_BASED +/// - `s1..s5` -> 10^10 > Integer.MAX_VALUE -> LONG_MAP_BASED +/// - `s1..s10` -> 10^20 > Long.MAX_VALUE -> ARRAY_MAP_BASED +/// - `mHigh` (cardinality ~14_500) -> INT_MAP_BASED for a single MV column (the in-place group id path) +/// +/// Every generator is closed after its run, and the test asserts that +/// [PinotDataBuffer#getDirectBufferUsage()] returns to the pre-test baseline after every comparison and at class +/// end (segments are mmap-loaded, so they do not count as direct memory). +public class OffHeapGroupKeyGeneratorParityTest { + private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "OffHeapGroupKeyGeneratorParityTest"); + private static final String RAW_TABLE_NAME = "testTable"; + private static final String SEGMENT_NAME = "testSegment"; + private static final long RANDOM_SEED = 13; + private static final int NUM_RECORDS = 25_000; + private static final int NUM_GROUPS_LIMIT = Server.DEFAULT_QUERY_EXECUTOR_NUM_GROUPS_LIMIT; + private static final int ARRAY_BASED_THRESHOLD = Server.DEFAULT_QUERY_EXECUTOR_MAX_INITIAL_RESULT_HOLDER_CAPACITY; + private static final int MAX_DOCS_PER_BLOCK = DocIdSetPlanNode.MAX_DOC_PER_CALL; + + // Dict-encoded INT SV columns of cardinality 100 each + private static final String[] DICT_SV_COLUMNS = {"s1", "s2", "s3", "s4", "s5", "s6", "s7", "s8", "s9", "s10"}; + // Dict-encoded INT MV columns of cardinality 100, plus a high-cardinality MV column for the IntMap MV path + private static final String M1 = "m1"; + private static final String M2 = "m2"; + private static final String M_HIGH = "mHigh"; + // Raw (no-dictionary) SV columns + private static final String R_INT = "rInt"; + private static final String R_LONG = "rLong"; + private static final String R_FLOAT = "rFloat"; + private static final String R_DOUBLE = "rDouble"; + private static final String R_STRING = "rString"; + private static final String R_BYTES = "rBytes"; + private static final String R_BIG_DECIMAL = "rBigDecimal"; + private static final String[] RAW_COLUMNS = {R_INT, R_LONG, R_FLOAT, R_DOUBLE, R_STRING, R_BYTES, R_BIG_DECIMAL}; + // Nullable raw columns. "nf" columns have the FIRST row null (plus more sprinkled); "nm" columns have values + // until row 12_000, nulls for rows [12_000, 13_000), then the full pool so new values first appear after the + // null stretch (exercising the off-heap null shift) + private static final String NF_INT = "nfInt"; + private static final String NM_INT = "nmInt"; + private static final String NF_DOUBLE = "nfDouble"; + private static final String NF_STR = "nfStr"; + private static final String NM_STR = "nmStr"; + private static final String NM_BIG_DECIMAL = "nmBigDecimal"; + private static final String[] NULLABLE_COLUMNS = {NF_INT, NM_INT, NF_DOUBLE, NF_STR, NM_STR, NM_BIG_DECIMAL}; + // "nm" columns: 8 distinct values (ids 0-7) appear before the nulls, so the null group takes dense id 8 + private static final int NULL_MID_GROUP_ID = 8; + + private IndexSegment _indexSegment; + private QueryContext _queryContext; + private ExpressionContext[] _projectionExpressions; + private long _directBufferBaseline; + + @BeforeClass + public void setUp() + throws Exception { + FileUtils.deleteDirectory(TEMP_DIR); + + List allColumns = new ArrayList<>(Arrays.asList(DICT_SV_COLUMNS)); + allColumns.addAll(Arrays.asList(M1, M2, M_HIGH)); + allColumns.addAll(Arrays.asList(RAW_COLUMNS)); + allColumns.addAll(Arrays.asList(NULLABLE_COLUMNS)); + + Schema.SchemaBuilder schemaBuilder = new Schema.SchemaBuilder().setSchemaName(RAW_TABLE_NAME); + for (String column : DICT_SV_COLUMNS) { + schemaBuilder.addSingleValueDimension(column, DataType.INT); + } + schemaBuilder.addMultiValueDimension(M1, DataType.INT); + schemaBuilder.addMultiValueDimension(M2, DataType.INT); + schemaBuilder.addMultiValueDimension(M_HIGH, DataType.INT); + schemaBuilder.addSingleValueDimension(R_INT, DataType.INT); + schemaBuilder.addSingleValueDimension(R_LONG, DataType.LONG); + schemaBuilder.addSingleValueDimension(R_FLOAT, DataType.FLOAT); + schemaBuilder.addSingleValueDimension(R_DOUBLE, DataType.DOUBLE); + schemaBuilder.addSingleValueDimension(R_STRING, DataType.STRING); + schemaBuilder.addSingleValueDimension(R_BYTES, DataType.BYTES); + schemaBuilder.addSingleValueDimension(R_BIG_DECIMAL, DataType.BIG_DECIMAL); + schemaBuilder.addSingleValueDimension(NF_INT, DataType.INT); + schemaBuilder.addSingleValueDimension(NM_INT, DataType.INT); + schemaBuilder.addSingleValueDimension(NF_DOUBLE, DataType.DOUBLE); + schemaBuilder.addSingleValueDimension(NF_STR, DataType.STRING); + schemaBuilder.addSingleValueDimension(NM_STR, DataType.STRING); + schemaBuilder.addSingleValueDimension(NM_BIG_DECIMAL, DataType.BIG_DECIMAL); + Schema schema = schemaBuilder.build(); + + List noDictionaryColumns = new ArrayList<>(Arrays.asList(RAW_COLUMNS)); + noDictionaryColumns.addAll(Arrays.asList(NULLABLE_COLUMNS)); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME) + .setNoDictionaryColumns(noDictionaryColumns).build(); + + Random random = new Random(RANDOM_SEED); + List records = new ArrayList<>(NUM_RECORDS); + for (int i = 0; i < NUM_RECORDS; i++) { + GenericRow record = new GenericRow(); + for (int c = 0; c < DICT_SV_COLUMNS.length; c++) { + record.putValue(DICT_SV_COLUMNS[c], c * 1_000_000 + random.nextInt(100)); + } + record.putValue(M1, randomMvValues(random, 100, 20_000_000)); + record.putValue(M2, randomMvValues(random, 100, 21_000_000)); + record.putValue(M_HIGH, randomMvValues(random, 15_000, 22_000_000)); + int rawPoolIndex = random.nextInt(150); + record.putValue(R_INT, rawPoolIndex * 3 - 200); + record.putValue(R_LONG, rawPoolIndex * 1_000_003L); + record.putValue(R_FLOAT, floatPoolValue(rawPoolIndex)); + record.putValue(R_DOUBLE, doublePoolValue(rawPoolIndex)); + record.putValue(R_STRING, rawPoolIndex < 5 ? "s😀" + rawPoolIndex : "str_" + rawPoolIndex); + record.putValue(R_BYTES, new byte[]{ + (byte) rawPoolIndex, (byte) (rawPoolIndex >> 4), (byte) (rawPoolIndex * 3), 42, (byte) i + }); + record.putValue(R_BIG_DECIMAL, BigDecimal.valueOf((rawPoolIndex - 75) * 25L, 2)); + boolean nullFirstIsNull = i % 7 == 0; + int nullFirstPoolIndex = i % 9; + boolean nullMidIsNull = i >= 12_000 && i < 13_000; + int nullMidPoolIndex = i < 12_000 ? i % 8 : i % 11; + record.putValue(NF_INT, nullFirstIsNull ? null : nullFirstPoolIndex * 3 - 15); + record.putValue(NM_INT, nullMidIsNull ? null : nullMidPoolIndex * 3 - 15); + record.putValue(NF_DOUBLE, nullFirstIsNull ? null : (nullFirstPoolIndex - 5) * 0.5d); + record.putValue(NF_STR, nullFirstIsNull ? null : "ns_" + nullFirstPoolIndex); + record.putValue(NM_STR, nullMidIsNull ? null : "ns_" + nullMidPoolIndex); + record.putValue(NM_BIG_DECIMAL, nullMidIsNull ? null : BigDecimal.valueOf((nullMidPoolIndex - 5) * 25L, 2)); + records.add(record); + } + + SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(tableConfig, schema); + segmentGeneratorConfig.setTableName(RAW_TABLE_NAME); + segmentGeneratorConfig.setSegmentName(SEGMENT_NAME); + segmentGeneratorConfig.setDefaultNullHandlingEnabled(true); + segmentGeneratorConfig.setOutDir(TEMP_DIR.getPath()); + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(segmentGeneratorConfig, new GenericRowRecordReader(records)); + driver.build(); + _indexSegment = ImmutableSegmentLoader.load(new File(TEMP_DIR, SEGMENT_NAME), ReadMode.mmap); + + _queryContext = QueryContextConverterUtils.getQueryContext( + "SELECT COUNT(*) FROM testTable GROUP BY " + StringUtils.join(allColumns, ", ")); + _projectionExpressions = getExpressions(allColumns.toArray(new String[0])); + // Bytes-key group-id runs allocate the same warm-up shapes each time; capture the baseline before any + // generator is created + _directBufferBaseline = PinotDataBuffer.getDirectBufferUsage(); + } + + @AfterClass + public void tearDown() + throws Exception { + assertEquals(PinotDataBuffer.getDirectBufferUsage(), _directBufferBaseline, + "Off-heap direct memory leaked by group key generators"); + _indexSegment.destroy(); + FileUtils.deleteDirectory(TEMP_DIR); + } + + private static Object[] randomMvValues(Random random, int cardinality, int base) { + int numValues = 1 + random.nextInt(3); + Object[] values = new Object[numValues]; + for (int i = 0; i < numValues; i++) { + values[i] = base + random.nextInt(cardinality); + } + return values; + } + + private static float floatPoolValue(int poolIndex) { + if (poolIndex == 0) { + return -0.0f; + } + if (poolIndex == 1) { + return 0.0f; + } + return (poolIndex - 75) * 0.25f; + } + + private static double doublePoolValue(int poolIndex) { + if (poolIndex == 0) { + return -0.0d; + } + if (poolIndex == 1) { + return 0.0d; + } + return (poolIndex - 75) * 0.5d; + } + + private static ExpressionContext[] getExpressions(String[] columns) { + ExpressionContext[] expressions = new ExpressionContext[columns.length]; + for (int i = 0; i < columns.length; i++) { + expressions[i] = ExpressionContext.forIdentifier(columns[i]); + } + return expressions; + } + + private BaseProjectOperator createProjectOperator() { + return new ProjectPlanNode(new SegmentContext(_indexSegment), _queryContext, + Arrays.asList(_projectionExpressions), MAX_DOCS_PER_BLOCK).run(); + } + + // --------------------------------------------------------------------------------------------- + // Run harness + // --------------------------------------------------------------------------------------------- + + private interface GeneratorFactory { + GroupKeyGenerator create(BaseProjectOperator projectOperator); + } + + private static class RunResult { + // Single-value group ids per block (null for MV runs) + final List _svBlockGroupIds = new ArrayList<>(); + // Multi-value group ids per block (null for SV runs) + final List _mvBlockGroupIds = new ArrayList<>(); + final List _numKeysPerBlock = new ArrayList<>(); + final List _upperBoundPerBlock = new ArrayList<>(); + Map> _groupKeys; + int _numKeys; + int _upperBound; + int _globalUpperBound; + // Highest direct-buffer usage observed while the generator was open (guards against the off-heap flag being + // silently ignored, which would make every parity assertion pass vacuously) + long _peakDirectUsage; + } + + private RunResult run(GeneratorFactory factory, boolean multiValue) { + RunResult result = new RunResult(); + BaseProjectOperator projectOperator = createProjectOperator(); + GroupKeyGenerator generator = factory.create(projectOperator); + try { + result._peakDirectUsage = PinotDataBuffer.getDirectBufferUsage(); + result._globalUpperBound = generator.getGlobalGroupKeyUpperBound(); + int[] svBuffer = multiValue ? null : new int[MAX_DOCS_PER_BLOCK]; + int[][] mvBuffer = multiValue ? new int[MAX_DOCS_PER_BLOCK][] : null; + ValueBlock block; + while ((block = projectOperator.nextBlock()) != null) { + int numDocs = block.getNumDocs(); + if (multiValue) { + generator.generateKeysForBlock(block, mvBuffer); + int[][] blockGroupIds = new int[numDocs][]; + for (int i = 0; i < numDocs; i++) { + // Deep-copy: the buffer rows may alias (and the map-based single-MV path mutates) block cache arrays + blockGroupIds[i] = mvBuffer[i].clone(); + } + result._mvBlockGroupIds.add(blockGroupIds); + } else { + generator.generateKeysForBlock(block, svBuffer); + result._svBlockGroupIds.add(Arrays.copyOf(svBuffer, numDocs)); + } + result._numKeysPerBlock.add(generator.getNumKeys()); + result._upperBoundPerBlock.add(generator.getCurrentGroupKeyUpperBound()); + result._peakDirectUsage = Math.max(result._peakDirectUsage, PinotDataBuffer.getDirectBufferUsage()); + } + result._groupKeys = collectGroupKeys(generator); + result._numKeys = generator.getNumKeys(); + result._upperBound = generator.getCurrentGroupKeyUpperBound(); + } finally { + generator.close(); + } + return result; + } + + private static Map> collectGroupKeys(GroupKeyGenerator generator) { + Map> groupKeys = new HashMap<>(); + Iterator iterator = generator.getGroupKeys(); + while (iterator.hasNext()) { + GroupKeyGenerator.GroupKey groupKey = iterator.next(); + // The GroupKey is a reused flyweight; copy the keys out + assertNull(groupKeys.put(groupKey._groupId, Arrays.asList(groupKey._keys.clone())), + "Iterator yielded duplicate group id: " + groupKey._groupId); + } + return groupKeys; + } + + // --------------------------------------------------------------------------------------------- + // Comparison helpers + // --------------------------------------------------------------------------------------------- + + /// Compares an on-heap run against an off-heap run. + /// + /// @param primitiveNullDivergence historical name: marks runs over a single no-dict column of a primitive + /// stored type with null handling enabled. Since the on-heap null-group counting fix landed alongside + /// the off-heap work, both modes count the null group and every count must match exactly — the flag is + /// retained only to document which runs carry a null group outside the map. + /// @param denseIds whether the generator assigns dense ids `0..numKeys-1` (all map-based variants; false only + /// for the dict ARRAY_BASED holder, whose ids are raw cardinality-products). + private void compareRuns(RunResult onHeap, RunResult offHeap, boolean primitiveNullDivergence, boolean denseIds, + String context) { + assertEquals(offHeap._globalUpperBound, onHeap._globalUpperBound, context + ": globalGroupKeyUpperBound"); + boolean multiValue = !onHeap._mvBlockGroupIds.isEmpty() || !offHeap._mvBlockGroupIds.isEmpty(); + int numBlocks = onHeap._numKeysPerBlock.size(); + assertEquals(offHeap._numKeysPerBlock.size(), numBlocks, context + ": block count"); + + // Group ids emitted per block must be identical arrays, and the per-block counts must match exactly (both + // modes count the primitive-type null group since the on-heap counting fix) + for (int b = 0; b < numBlocks; b++) { + if (multiValue) { + int[][] onHeapIds = onHeap._mvBlockGroupIds.get(b); + int[][] offHeapIds = offHeap._mvBlockGroupIds.get(b); + assertEquals(offHeapIds.length, onHeapIds.length, context + ": numDocs in block " + b); + for (int i = 0; i < onHeapIds.length; i++) { + assertTrue(Arrays.equals(offHeapIds[i], onHeapIds[i]), + context + ": MV group ids differ in block " + b + " at doc " + i + ": expected " + + Arrays.toString(onHeapIds[i]) + " but got " + Arrays.toString(offHeapIds[i])); + } + } else { + int[] onHeapIds = onHeap._svBlockGroupIds.get(b); + int[] offHeapIds = offHeap._svBlockGroupIds.get(b); + assertTrue(Arrays.equals(offHeapIds, onHeapIds), + context + ": SV group ids differ in block " + b + " at doc " + firstMismatch(onHeapIds, offHeapIds)); + } + assertEquals(offHeap._numKeysPerBlock.get(b), onHeap._numKeysPerBlock.get(b), + context + ": getNumKeys after block " + b); + assertEquals(offHeap._upperBoundPerBlock.get(b), onHeap._upperBoundPerBlock.get(b), + context + ": getCurrentGroupKeyUpperBound after block " + b); + } + assertEquals(offHeap._numKeys, onHeap._numKeys, context + ": final getNumKeys"); + assertEquals(offHeap._upperBound, onHeap._upperBound, context + ": final getCurrentGroupKeyUpperBound"); + + // The iterators must yield the same (groupId -> keys) mapping (iterator order may differ) + assertEquals(offHeap._groupKeys, onHeap._groupKeys, context + ": group key mapping"); + + if (denseIds) { + assertDenseIds(offHeap, context + " (off-heap)"); + // On-heap map-based variants are dense too; verifying both pins the shared contract + assertDenseIds(onHeap, context + " (on-heap)"); + // The dense (map-based) variants are exactly the ones that must go off-heap: assert the off-heap run + // actually held direct memory while open, so a silently ignored offHeap flag cannot pass this test + assertTrue(offHeap._peakDirectUsage > _directBufferBaseline, + context + ": off-heap generator never allocated direct memory"); + } else { + // The dict ARRAY_BASED (T0) holder stays on-heap by design even in off-heap mode + assertEquals(offHeap._peakDirectUsage, _directBufferBaseline, + context + ": ARRAY_BASED holder unexpectedly allocated direct memory"); + } + assertEquals(PinotDataBuffer.getDirectBufferUsage(), _directBufferBaseline, + context + ": off-heap direct memory leaked"); + } + + private static int firstMismatch(int[] expected, int[] actual) { + for (int i = 0; i < Math.min(expected.length, actual.length); i++) { + if (expected[i] != actual[i]) { + return i; + } + } + return -1; + } + + /// Returns the group id mapped to a single null key, or null if no null group exists. + private static Integer findNullGroupId(Map> groupKeys) { + Integer nullGroupId = null; + for (Map.Entry> entry : groupKeys.entrySet()) { + List keys = entry.getValue(); + if (keys.size() == 1 && keys.get(0) == null) { + assertNull(nullGroupId, "Multiple null groups found: " + nullGroupId + " and " + entry.getKey()); + nullGroupId = entry.getKey(); + } + } + return nullGroupId; + } + + /// Asserts the iterator emitted dense ids `0..numKeys-1` with no gaps or duplicates. + private static void assertDenseIds(RunResult result, String context) { + TreeSet ids = new TreeSet<>(result._groupKeys.keySet()); + assertEquals(ids.size(), result._groupKeys.size(), context + ": duplicate ids"); + if (!ids.isEmpty()) { + assertEquals((int) ids.first(), 0, context + ": ids must start at 0"); + assertEquals((int) ids.last(), ids.size() - 1, context + ": ids must be dense (no gaps)"); + } + } + + // --------------------------------------------------------------------------------------------- + // DictionaryBasedGroupKeyGenerator + // --------------------------------------------------------------------------------------------- + + private void compareDictionary(String[] columns, int numGroupsLimit, boolean multiValue, boolean denseIds) { + String context = "Dictionary" + (multiValue ? " MV " : " SV ") + Arrays.toString(columns) + " limit " + + numGroupsLimit; + RunResult onHeap = run(op -> new DictionaryBasedGroupKeyGenerator(op, getExpressions(columns), numGroupsLimit, + ARRAY_BASED_THRESHOLD, false, null, false), multiValue); + RunResult offHeap = run(op -> new DictionaryBasedGroupKeyGenerator(op, getExpressions(columns), numGroupsLimit, + ARRAY_BASED_THRESHOLD, false, null, true), multiValue); + compareRuns(onHeap, offHeap, false, denseIds, context); + } + + @Test + public void testDictionarySingleValueVariants() { + // ARRAY_BASED (product 100): the T0 path stays on-heap in both modes by design + compareDictionary(new String[]{"s1"}, NUM_GROUPS_LIMIT, false, false); + // INT_MAP_BASED (product 10^6) + compareDictionary(new String[]{"s1", "s2", "s3"}, NUM_GROUPS_LIMIT, false, true); + // LONG_MAP_BASED (product 10^10 > Integer.MAX_VALUE) + compareDictionary(new String[]{"s1", "s2", "s3", "s4", "s5"}, NUM_GROUPS_LIMIT, false, true); + // ARRAY_MAP_BASED (product 10^20 > Long.MAX_VALUE) + compareDictionary(DICT_SV_COLUMNS, NUM_GROUPS_LIMIT, false, true); + } + + @Test + public void testDictionaryMultiValueVariants() { + // ARRAY_BASED MV (product 100) + compareDictionary(new String[]{M1}, NUM_GROUPS_LIMIT, true, false); + // INT_MAP_BASED single MV column (~14_500 > arrayBasedThreshold): the in-place group id path + compareDictionary(new String[]{M_HIGH}, NUM_GROUPS_LIMIT, true, true); + // INT_MAP_BASED MV (product 10^6) + compareDictionary(new String[]{M1, "s1", "s2"}, NUM_GROUPS_LIMIT, true, true); + // LONG_MAP_BASED MV (product 10^10) + compareDictionary(new String[]{M1, M2, "s1", "s2", "s3"}, NUM_GROUPS_LIMIT, true, true); + // ARRAY_MAP_BASED MV (product 10^24) + String[] arrayMapColumns = new String[DICT_SV_COLUMNS.length + 2]; + arrayMapColumns[0] = M1; + arrayMapColumns[1] = M2; + System.arraycopy(DICT_SV_COLUMNS, 0, arrayMapColumns, 2, DICT_SV_COLUMNS.length); + compareDictionary(arrayMapColumns, NUM_GROUPS_LIMIT, true, true); + } + + @Test + public void testDictionaryCapSemantics() { + // numGroupsLimit < cardinality product forces the map-based holders; group ids (including INVALID_ID + // positions) must be identical + compareDictionary(new String[]{"s1"}, 7, false, true); + compareDictionary(new String[]{"s1", "s2", "s3", "s4", "s5"}, 7, false, true); + compareDictionary(DICT_SV_COLUMNS, 7, false, true); + compareDictionary(new String[]{M1}, 7, true, true); + } + + // --------------------------------------------------------------------------------------------- + // NoDictionarySingleColumnGroupKeyGenerator + // --------------------------------------------------------------------------------------------- + + private RunResult[] compareNoDictionarySingle(String column, int numGroupsLimit, boolean nullHandlingEnabled, + boolean multiValue, boolean primitiveNullDivergence) { + String context = "NoDictionarySingle " + column + " limit " + numGroupsLimit + + (nullHandlingEnabled ? " nullHandling" : ""); + ExpressionContext expression = ExpressionContext.forIdentifier(column); + RunResult onHeap = run(op -> new NoDictionarySingleColumnGroupKeyGenerator(op, expression, numGroupsLimit, + nullHandlingEnabled, null, false), multiValue); + RunResult offHeap = run(op -> new NoDictionarySingleColumnGroupKeyGenerator(op, expression, numGroupsLimit, + nullHandlingEnabled, null, true), multiValue); + compareRuns(onHeap, offHeap, primitiveNullDivergence, true, context); + return new RunResult[]{onHeap, offHeap}; + } + + @Test + public void testNoDictionarySingleColumn() { + for (String column : RAW_COLUMNS) { + compareNoDictionarySingle(column, NUM_GROUPS_LIMIT, false, false, false); + } + // Dict-encoded column routed through the no-dict generator (as the executor does when null handling is on) + compareNoDictionarySingle("s1", NUM_GROUPS_LIMIT, false, false, false); + } + + @Test + public void testNoDictionarySingleColumnMultiValue() { + compareNoDictionarySingle(M1, NUM_GROUPS_LIMIT, false, true, false); + compareNoDictionarySingle(M_HIGH, NUM_GROUPS_LIMIT, false, true, false); + } + + @Test + public void testNoDictionarySingleColumnCapSemantics() { + compareNoDictionarySingle(R_INT, 7, false, false, false); + compareNoDictionarySingle(R_STRING, 7, false, false, false); + compareNoDictionarySingle(R_BIG_DECIMAL, 7, false, false, false); + } + + @Test + public void testNoDictionarySingleColumnNullHandling() { + // Primitive stored types: off-heap counts must exceed on-heap by exactly 1 once the null group is assigned + RunResult[] runs = compareNoDictionarySingle(NF_INT, NUM_GROUPS_LIMIT, true, false, true); + // First row is null, so the null group must take dense id 0 + assertEquals(findNullGroupId(runs[1]._groupKeys), Integer.valueOf(0), "nfInt null group id"); + runs = compareNoDictionarySingle(NM_INT, NUM_GROUPS_LIMIT, true, false, true); + // 8 distinct values (ids 0-7) precede the null stretch, so the null group must take dense id 8; the values + // first appearing after the null stretch then shift to ids 9+ + assertEquals(findNullGroupId(runs[1]._groupKeys), Integer.valueOf(NULL_MID_GROUP_ID), "nmInt null group id"); + compareNoDictionarySingle(NF_DOUBLE, NUM_GROUPS_LIMIT, true, false, true); + // Object stored types: the on-heap map holds the null key, so the counts must match exactly + runs = compareNoDictionarySingle(NF_STR, NUM_GROUPS_LIMIT, true, false, false); + assertEquals(findNullGroupId(runs[1]._groupKeys), Integer.valueOf(0), "nfStr null group id"); + runs = compareNoDictionarySingle(NM_STR, NUM_GROUPS_LIMIT, true, false, false); + assertEquals(findNullGroupId(runs[1]._groupKeys), Integer.valueOf(NULL_MID_GROUP_ID), "nmStr null group id"); + compareNoDictionarySingle(NM_BIG_DECIMAL, NUM_GROUPS_LIMIT, true, false, false); + } + + @Test + public void testNoDictionarySingleColumnNullCapSemantics() { + // Null group is the id that hits the cap: 8 value groups precede the null stretch, cap 9 -> null gets id 8 + // and every value first appearing after the null stretch gets INVALID_ID + RunResult[] runs = compareNoDictionarySingle(NM_INT, NULL_MID_GROUP_ID + 1, true, false, true); + assertEquals(findNullGroupId(runs[1]._groupKeys), Integer.valueOf(NULL_MID_GROUP_ID), + "nmInt null group id at cap boundary"); + assertEquals(runs[1]._groupKeys.size(), NULL_MID_GROUP_ID + 1, "nmInt group count at cap boundary"); + // Cap hit before any null appears: 8 value groups fill a cap of 3 long before row 12_000, so the null group + // must never be assigned (getKeyForNullValue returns INVALID_ID in both modes) + runs = compareNoDictionarySingle(NM_INT, 3, true, false, true); + assertNull(findNullGroupId(runs[1]._groupKeys), "nmInt cap-before-null must not assign a null group"); + assertEquals(runs[1]._groupKeys.size(), 3, "nmInt group count under cap 3"); + // Null in the very first row with a cap: null takes id 0, later new values are cut off by the cap + runs = compareNoDictionarySingle(NF_INT, 7, true, false, true); + assertEquals(findNullGroupId(runs[1]._groupKeys), Integer.valueOf(0), "nfInt null group id under cap"); + // Same cap scenarios for an object stored type + runs = compareNoDictionarySingle(NM_STR, NULL_MID_GROUP_ID + 1, true, false, false); + assertEquals(findNullGroupId(runs[1]._groupKeys), Integer.valueOf(NULL_MID_GROUP_ID), + "nmStr null group id at cap boundary"); + runs = compareNoDictionarySingle(NM_STR, 3, true, false, false); + assertNull(findNullGroupId(runs[1]._groupKeys), "nmStr cap-before-null must not assign a null group"); + } + + // --------------------------------------------------------------------------------------------- + // NoDictionaryMultiColumnGroupKeyGenerator + // --------------------------------------------------------------------------------------------- + + private RunResult[] compareNoDictionaryMulti(String[] columns, int numGroupsLimit, boolean nullHandlingEnabled, + boolean multiValue) { + String context = "NoDictionaryMulti " + Arrays.toString(columns) + " limit " + numGroupsLimit + + (nullHandlingEnabled ? " nullHandling" : ""); + RunResult onHeap = run(op -> new NoDictionaryMultiColumnGroupKeyGenerator(op, getExpressions(columns), + numGroupsLimit, nullHandlingEnabled, null, false), multiValue); + RunResult offHeap = run(op -> new NoDictionaryMultiColumnGroupKeyGenerator(op, getExpressions(columns), + numGroupsLimit, nullHandlingEnabled, null, true), multiValue); + // The multi-column generator counts groups from the key map in both modes (null components are ID_FOR_NULL + // inside the composite key), so there is never a counting divergence + compareRuns(onHeap, offHeap, false, true, context); + return new RunResult[]{onHeap, offHeap}; + } + + @Test + public void testNoDictionaryMultiColumn() { + compareNoDictionaryMulti(new String[]{R_INT, R_STRING}, NUM_GROUPS_LIMIT, false, false); + compareNoDictionaryMulti(new String[]{R_FLOAT, R_DOUBLE}, NUM_GROUPS_LIMIT, false, false); + compareNoDictionaryMulti(new String[]{R_LONG, R_BYTES, R_BIG_DECIMAL}, NUM_GROUPS_LIMIT, false, false); + // Hybrid: dict-encoded column + raw column + compareNoDictionaryMulti(new String[]{R_STRING, "s1"}, NUM_GROUPS_LIMIT, false, false); + } + + @Test + public void testNoDictionaryMultiColumnNullHandling() { + RunResult[] runs = compareNoDictionaryMulti(new String[]{NF_INT, NM_STR}, NUM_GROUPS_LIMIT, true, false); + // Sanity: composite groups with a null component must exist and contain nulls in the key positions + boolean sawNullComponent = false; + for (List keys : runs[1]._groupKeys.values()) { + if (keys.get(0) == null || keys.get(1) == null) { + sawNullComponent = true; + break; + } + } + assertTrue(sawNullComponent, "Expected composite groups with null components"); + compareNoDictionaryMulti(new String[]{NM_INT, NM_BIG_DECIMAL}, NUM_GROUPS_LIMIT, true, false); + } + + @Test + public void testNoDictionaryMultiColumnCapSemantics() { + compareNoDictionaryMulti(new String[]{R_INT, R_STRING}, 7, false, false); + compareNoDictionaryMulti(new String[]{NF_INT, NM_STR}, 5, true, false); + } + + @Test + public void testNoDictionaryMultiColumnMultiValue() { + compareNoDictionaryMulti(new String[]{M1, R_INT}, NUM_GROUPS_LIMIT, false, true); + compareNoDictionaryMulti(new String[]{M1, M2}, NUM_GROUPS_LIMIT, false, true); + } + + @Test + public void testNoDictionarySingleColumnDenseIdsWithNullPresent() { + // Focused null-shift structural check: dense ids, no gaps, no duplicates, and the null group present exactly + // once in the iterator, for both null layouts + for (String column : new String[]{NF_INT, NM_INT, NF_STR, NM_STR}) { + ExpressionContext expression = ExpressionContext.forIdentifier(column); + RunResult offHeap = run(op -> new NoDictionarySingleColumnGroupKeyGenerator(op, expression, NUM_GROUPS_LIMIT, + true, null, true), false); + assertDenseIds(offHeap, column); + assertNotNull(findNullGroupId(offHeap._groupKeys), column + ": null group missing from iterator"); + assertEquals(PinotDataBuffer.getDirectBufferUsage(), _directBufferBaseline, + column + ": off-heap direct memory leaked"); + } + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapBytesGroupIdMapTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapBytesGroupIdMapTest.java new file mode 100644 index 000000000000..4325f576c948 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapBytesGroupIdMapTest.java @@ -0,0 +1,286 @@ +/** + * 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.offheap; + +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Random; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + + +public class OffHeapBytesGroupIdMapTest { + private static final int NO_BOUND = Integer.MAX_VALUE; + + @Test + public void testDifferentialAgainstReferenceMap() { + Random random = new Random(42); + Map referenceMap = new HashMap<>(); + List generatedKeys = new ArrayList<>(); + try (OffHeapBytesGroupIdMap map = new OffHeapBytesGroupIdMap(1024)) { + for (int i = 0; i < 200_000; i++) { + byte[] key; + if (generatedKeys.isEmpty() || random.nextBoolean()) { + key = new byte[random.nextInt(65)]; + random.nextBytes(key); + generatedKeys.add(key); + } else { + // Re-submit a previously generated key to exercise the duplicate path + key = generatedKeys.get(random.nextInt(generatedKeys.size())); + } + // ISO-8859-1 maps each byte to a unique char, so the String is a faithful reference key + String referenceKey = new String(key, StandardCharsets.ISO_8859_1); + int expectedId = referenceMap.computeIfAbsent(referenceKey, k -> referenceMap.size()); + assertEquals(map.getGroupId(key, NO_BOUND), expectedId); + } + assertEquals(map.size(), referenceMap.size()); + + // getKey/readKey/getKeyLength round-trip for every 1000th id + Map idToKey = new HashMap<>(); + referenceMap.forEach((key, id) -> idToKey.put(id, key)); + for (int groupId = 0; groupId < map.size(); groupId += 1000) { + byte[] expectedKey = idToKey.get(groupId).getBytes(StandardCharsets.ISO_8859_1); + assertEquals(map.getKeyLength(groupId), expectedKey.length); + assertEquals(map.getKey(groupId), expectedKey); + int destOffset = 3; + byte[] dest = new byte[expectedKey.length + destOffset + 4]; + map.readKey(groupId, dest, destOffset); + assertEquals(Arrays.copyOfRange(dest, destOffset, destOffset + expectedKey.length), expectedKey); + } + } + } + + @Test + public void testEmptyKey() { + try (OffHeapBytesGroupIdMap map = new OffHeapBytesGroupIdMap(16)) { + byte[] emptyKey = new byte[0]; + assertEquals(map.getGroupId(emptyKey, NO_BOUND), 0); + assertEquals(map.getGroupId(new byte[0], NO_BOUND), 0); + assertEquals(map.size(), 1); + assertEquals(map.getKeyLength(0), 0); + assertEquals(map.getKey(0), emptyKey); + // The empty key is distinct from a single 0x00 byte + assertEquals(map.getGroupId(new byte[1], NO_BOUND), 1); + assertEquals(map.size(), 2); + } + } + + @Test + public void testPrefixKeys() { + byte[] fullKey = "abcdefgh".getBytes(StandardCharsets.UTF_8); + try (OffHeapBytesGroupIdMap map = new OffHeapBytesGroupIdMap(16)) { + for (int length = 0; length <= fullKey.length; length++) { + assertEquals(map.getGroupId(fullKey, 0, length, NO_BOUND), length); + } + assertEquals(map.size(), fullKey.length + 1); + for (int length = 0; length <= fullKey.length; length++) { + assertEquals(map.getGroupId(fullKey, 0, length, NO_BOUND), length); + assertEquals(map.getKey(length), Arrays.copyOfRange(fullKey, 0, length)); + } + } + } + + @Test + public void testAllZeroKeys() { + try (OffHeapBytesGroupIdMap map = new OffHeapBytesGroupIdMap(16)) { + // All-zero keys live in the payload and must not be confused with empty directory slots + for (int length = 1; length <= 16; length++) { + assertEquals(map.getGroupId(new byte[length], NO_BOUND), length - 1); + } + assertEquals(map.size(), 16); + for (int length = 1; length <= 16; length++) { + assertEquals(map.getGroupId(new byte[length], NO_BOUND), length - 1); + assertEquals(map.getKey(length - 1), new byte[length]); + } + } + } + + @Test + public void testGroupIdUpperBound() { + try (OffHeapBytesGroupIdMap map = new OffHeapBytesGroupIdMap(16)) { + int upperBound = 10; + for (int i = 0; i < upperBound; i++) { + assertEquals(map.getGroupId(key(i), upperBound), i); + } + assertEquals(map.size(), upperBound); + // Existing keys always resolve, even at the cap + for (int i = 0; i < upperBound; i++) { + assertEquals(map.getGroupId(key(i), upperBound), i); + } + // New keys are rejected at the cap, and the size stays frozen + assertEquals(map.getGroupId(key(10), upperBound), GroupKeyGenerator.INVALID_ID); + assertEquals(map.getGroupId(key(11), upperBound), GroupKeyGenerator.INVALID_ID); + assertEquals(map.getGroupId(key(10), upperBound), GroupKeyGenerator.INVALID_ID); + assertEquals(map.size(), upperBound); + // A rejected key was not inserted: raising the bound assigns it the next dense id + assertEquals(map.getGroupId(key(10), upperBound + 1), upperBound); + assertEquals(map.size(), upperBound + 1); + } + } + + @Test + public void testResizePreservesIds() { + int numKeys = 100_000; + // Start with the minimum directory (1024 slots) to force many directory resizes + try (OffHeapBytesGroupIdMap map = new OffHeapBytesGroupIdMap(0)) { + for (int i = 0; i < numKeys; i++) { + assertEquals(map.getGroupId(key(i), NO_BOUND), i); + } + assertEquals(map.size(), numKeys); + for (int i = 0; i < numKeys; i += 1000) { + assertEquals(map.getGroupId(key(i), NO_BOUND), i); + assertEquals(map.getKey(i), key(i)); + } + assertEquals(map.size(), numKeys); + } + } + + @Test + public void testOversizedKey() { + Random random = new Random(42); + byte[] oversizedKey = new byte[300_000]; + random.nextBytes(oversizedKey); + try (OffHeapBytesGroupIdMap map = new OffHeapBytesGroupIdMap(16)) { + for (int i = 0; i < 100; i++) { + assertEquals(map.getGroupId(key(i), NO_BOUND), i); + } + int oversizedId = map.getGroupId(oversizedKey, NO_BOUND); + assertEquals(oversizedId, 100); + // Normal inserts continue after the oversized record + for (int i = 100; i < 200; i++) { + assertEquals(map.getGroupId(key(i), NO_BOUND), i + 1); + } + assertEquals(map.getGroupId(oversizedKey, NO_BOUND), oversizedId); + assertEquals(map.getKeyLength(oversizedId), oversizedKey.length); + assertEquals(map.getKey(oversizedId), oversizedKey); + // The oversized record sits at offset 0 of its dedicated chunk + assertEquals(map.getPayloadGlobalOffset(oversizedId) % OffHeapBytesGroupIdMap.CHUNK_SIZE, 0); + verifyRecordOffsets(map); + } + } + + @Test + public void testChunkBoundaryKey() { + // A record of exactly CHUNK_SIZE bytes (16-byte header + key) exactly fills one normal chunk + byte[] boundaryKey = new byte[OffHeapBytesGroupIdMap.CHUNK_SIZE - 16]; + new Random(42).nextBytes(boundaryKey); + try (OffHeapBytesGroupIdMap map = new OffHeapBytesGroupIdMap(16)) { + assertEquals(map.getGroupId(boundaryKey, NO_BOUND), 0); + assertEquals(map.getGroupId(boundaryKey, NO_BOUND), 0); + assertEquals(map.getKeyLength(0), boundaryKey.length); + assertEquals(map.getKey(0), boundaryKey); + assertEquals(map.getPayloadGlobalOffset(0) % OffHeapBytesGroupIdMap.CHUNK_SIZE, 0); + // The next record starts a new chunk + assertEquals(map.getGroupId(key(1), NO_BOUND), 1); + assertEquals(map.getKey(1), key(1)); + verifyRecordOffsets(map); + } + } + + @Test + public void testOffsetLengthVariant() { + Random random = new Random(42); + byte[] outerArray = new byte[64]; + random.nextBytes(outerArray); + int offset = 13; + int length = 21; + byte[] slice = Arrays.copyOfRange(outerArray, offset, offset + length); + try (OffHeapBytesGroupIdMap map = new OffHeapBytesGroupIdMap(16)) { + int groupId = map.getGroupId(outerArray, offset, length, NO_BOUND); + assertEquals(groupId, 0); + // The same bytes submitted as a standalone array resolve to the same id + assertEquals(map.getGroupId(slice, NO_BOUND), groupId); + assertEquals(map.size(), 1); + assertEquals(map.getKey(groupId), slice); + } + } + + /// Exercises the wrapper-based fallback arms (directory probe, matchRecordSlow, id-index reads, resize, + /// zero-fill) that normally only run for buffers beyond the 2GB view limit. + @Test + public void testDifferentialWithoutViews() { + OffHeapGroupByUtils.setViewSizeLimitBytes(0); + try { + testDifferentialAgainstReferenceMap(); + testResizePreservesIds(); + testOversizedKey(); + } finally { + OffHeapGroupByUtils.setViewSizeLimitBytes(Integer.MAX_VALUE); + } + } + + @Test + public void testCloseIsIdempotent() { + OffHeapBytesGroupIdMap map = new OffHeapBytesGroupIdMap(16); + assertEquals(map.getGroupId(key(0), NO_BOUND), 0); + map.close(); + // Second close is a no-op + map.close(); + } + + @Test + public void testNoDirectMemoryLeak() { + long baseline = PinotDataBuffer.getDirectBufferUsage(); + OffHeapBytesGroupIdMap map = new OffHeapBytesGroupIdMap(16); + try { + Random random = new Random(42); + // Enough keys to force directory resizes, id index growth and multiple payload chunks + for (int i = 0; i < 50_000; i++) { + byte[] key = new byte[random.nextInt(65)]; + random.nextBytes(key); + map.getGroupId(key, NO_BOUND); + } + byte[] oversizedKey = new byte[300_000]; + random.nextBytes(oversizedKey); + map.getGroupId(oversizedKey, NO_BOUND); + assertTrue(map.getOffHeapMemoryBytes() > 0); + assertTrue(PinotDataBuffer.getDirectBufferUsage() > baseline); + } finally { + map.close(); + } + assertEquals(PinotDataBuffer.getDirectBufferUsage(), baseline); + } + + /// Verifies the global offset encoding invariants: every record start offset within a chunk is smaller than + /// CHUNK_SIZE, normal records fit entirely within their chunk, and oversized records start at offset 0. + private static void verifyRecordOffsets(OffHeapBytesGroupIdMap map) { + for (int groupId = 0; groupId < map.size(); groupId++) { + long offsetInChunk = map.getPayloadGlobalOffset(groupId) % OffHeapBytesGroupIdMap.CHUNK_SIZE; + assertTrue(offsetInChunk < OffHeapBytesGroupIdMap.CHUNK_SIZE); + long recordSize = 16L + map.getKeyLength(groupId); + if (recordSize <= OffHeapBytesGroupIdMap.CHUNK_SIZE) { + assertTrue(offsetInChunk + recordSize <= OffHeapBytesGroupIdMap.CHUNK_SIZE); + } else { + assertEquals(offsetInChunk, 0); + } + } + } + + private static byte[] key(int i) { + return ("key-" + i).getBytes(StandardCharsets.UTF_8); + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByBufferPoolTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByBufferPoolTest.java new file mode 100644 index 000000000000..ce6b1b808f46 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByBufferPoolTest.java @@ -0,0 +1,109 @@ +/** + * 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.offheap; + +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotSame; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertTrue; + + +public class OffHeapGroupByBufferPoolTest { + + @AfterMethod + public void resetPool() { + OffHeapGroupByBufferPool.clearCurrentThread(); + OffHeapGroupByBufferPool.setMaxBytesPerThread(0); + } + + @Test + public void testDisabledPoolIsPassThrough() { + long baseline = PinotDataBuffer.getDirectBufferUsage(); + PinotDataBuffer buffer = OffHeapGroupByBufferPool.acquire(4096, "test"); + assertTrue(PinotDataBuffer.getDirectBufferUsage() > baseline); + OffHeapGroupByBufferPool.release(buffer); + // Disabled pool closes on release: usage returns to baseline and nothing is retained + assertEquals(PinotDataBuffer.getDirectBufferUsage(), baseline); + assertEquals(OffHeapGroupByBufferPool.getPooledBytes(), 0); + } + + @Test + public void testEnabledPoolReusesExactSize() { + OffHeapGroupByBufferPool.setMaxBytesPerThread(1 << 20); + long baseline = PinotDataBuffer.getDirectBufferUsage(); + PinotDataBuffer first = OffHeapGroupByBufferPool.acquire(8192, "test"); + OffHeapGroupByBufferPool.release(first); + // Pooled buffer stays open and accounted + assertEquals(PinotDataBuffer.getDirectBufferUsage(), baseline + 8192); + assertEquals(OffHeapGroupByBufferPool.getPooledBytes(), 8192); + // Exact-size acquire reuses the same buffer instance; a different size allocates fresh + PinotDataBuffer reused = OffHeapGroupByBufferPool.acquire(8192, "test"); + assertSame(reused, first); + assertEquals(OffHeapGroupByBufferPool.getPooledBytes(), 0); + PinotDataBuffer other = OffHeapGroupByBufferPool.acquire(4096, "test"); + assertNotSame(other, first); + OffHeapGroupByBufferPool.release(reused); + OffHeapGroupByBufferPool.release(other); + assertEquals(OffHeapGroupByBufferPool.getPooledBytes(), 8192 + 4096); + // clearCurrentThread closes everything + OffHeapGroupByBufferPool.clearCurrentThread(); + assertEquals(OffHeapGroupByBufferPool.getPooledBytes(), 0); + assertEquals(PinotDataBuffer.getDirectBufferUsage(), baseline); + } + + @Test + public void testPerThreadCapEvicts() { + OffHeapGroupByBufferPool.setMaxBytesPerThread(10_000); + long baseline = PinotDataBuffer.getDirectBufferUsage(); + PinotDataBuffer first = OffHeapGroupByBufferPool.acquire(8192, "test"); + PinotDataBuffer second = OffHeapGroupByBufferPool.acquire(8192, "test"); + OffHeapGroupByBufferPool.release(first); + // Second release would exceed the 10_000-byte cap: the buffer is closed instead of pooled + OffHeapGroupByBufferPool.release(second); + assertEquals(OffHeapGroupByBufferPool.getPooledBytes(), 8192); + assertEquals(PinotDataBuffer.getDirectBufferUsage(), baseline + 8192); + OffHeapGroupByBufferPool.clearCurrentThread(); + assertEquals(PinotDataBuffer.getDirectBufferUsage(), baseline); + } + + @Test + public void testStructuresRunCorrectlyWithPoolEnabled() { + OffHeapGroupByBufferPool.setMaxBytesPerThread(16 << 20); + long baseline = PinotDataBuffer.getDirectBufferUsage(); + // Run a map through two full lifecycles: the second run reuses dirty pooled buffers, so any missing + // re-initialization (zero-fill) would corrupt its results + for (int run = 0; run < 2; run++) { + try (OffHeapIntGroupIdMap map = new OffHeapIntGroupIdMap(0)) { + for (int key = 0; key < 50_000; key++) { + assertEquals(map.getGroupId(key, Integer.MAX_VALUE), key, "run " + run + " key " + key); + } + for (int key = 0; key < 50_000; key++) { + assertEquals(map.getGroupId(key, Integer.MAX_VALUE), key, "run " + run + " lookup " + key); + } + } + } + assertTrue(OffHeapGroupByBufferPool.getPooledBytes() > 0); + OffHeapGroupByBufferPool.clearCurrentThread(); + assertEquals(PinotDataBuffer.getDirectBufferUsage(), baseline); + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByResultHolderTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByResultHolderTest.java new file mode 100644 index 000000000000..d0729617689a --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByResultHolderTest.java @@ -0,0 +1,383 @@ +/** + * 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.offheap; + +import java.util.Random; +import org.apache.pinot.core.query.aggregation.groupby.DoubleGroupByResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator; +import org.apache.pinot.core.query.aggregation.groupby.IntGroupByResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.LongGroupByResultHolder; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; +import org.testng.Assert; +import org.testng.annotations.Test; + + +/// Tests for the off-heap fixed-width `GroupByResultHolder` implementations, asserting semantic equivalence +/// with their on-heap twins. +public class OffHeapGroupByResultHolderTest { + private static final long RANDOM_SEED = 42; + private static final int NUM_SLOTS = 50_000; + private static final int NUM_OPERATIONS = 100_000; + private static final int INITIAL_CAPACITY = 100; + private static final int MAX_GROWTH_STEP = 5000; + + @Test + public void testDifferentialDouble() { + Random random = new Random(RANDOM_SEED); + DoubleGroupByResultHolder onHeap = new DoubleGroupByResultHolder(INITIAL_CAPACITY, NUM_SLOTS, -1.0); + try (OffHeapDoubleGroupByResultHolder offHeap = + new OffHeapDoubleGroupByResultHolder(INITIAL_CAPACITY, NUM_SLOTS, -1.0)) { + int capacity = INITIAL_CAPACITY; + for (int i = 0; i < NUM_OPERATIONS; i++) { + int op = random.nextInt(10); + if (op == 0 && capacity < NUM_SLOTS) { + int newCapacity = Math.min(capacity + 1 + random.nextInt(MAX_GROWTH_STEP), NUM_SLOTS); + onHeap.ensureCapacity(newCapacity); + offHeap.ensureCapacity(newCapacity); + capacity = newCapacity; + } else if (op < 6) { + int groupKey = nextGroupKey(random, capacity); + double value = random.nextDouble(); + onHeap.setValueForKey(groupKey, value); + offHeap.setValueForKey(groupKey, value); + } else { + int groupKey = nextGroupKey(random, capacity); + assertDoubleEquals(offHeap.getDoubleResult(groupKey), onHeap.getDoubleResult(groupKey)); + } + } + onHeap.ensureCapacity(NUM_SLOTS); + offHeap.ensureCapacity(NUM_SLOTS); + for (int groupKey = 0; groupKey < NUM_SLOTS; groupKey++) { + assertDoubleEquals(offHeap.getDoubleResult(groupKey), onHeap.getDoubleResult(groupKey)); + } + } + } + + @Test + public void testDifferentialLong() { + Random random = new Random(RANDOM_SEED); + LongGroupByResultHolder onHeap = new LongGroupByResultHolder(INITIAL_CAPACITY, NUM_SLOTS, -1L); + try (OffHeapLongGroupByResultHolder offHeap = + new OffHeapLongGroupByResultHolder(INITIAL_CAPACITY, NUM_SLOTS, -1L)) { + int capacity = INITIAL_CAPACITY; + for (int i = 0; i < NUM_OPERATIONS; i++) { + int op = random.nextInt(10); + if (op == 0 && capacity < NUM_SLOTS) { + int newCapacity = Math.min(capacity + 1 + random.nextInt(MAX_GROWTH_STEP), NUM_SLOTS); + onHeap.ensureCapacity(newCapacity); + offHeap.ensureCapacity(newCapacity); + capacity = newCapacity; + } else if (op < 6) { + int groupKey = nextGroupKey(random, capacity); + long value = random.nextLong(); + onHeap.setValueForKey(groupKey, value); + offHeap.setValueForKey(groupKey, value); + } else { + int groupKey = nextGroupKey(random, capacity); + Assert.assertEquals(offHeap.getLongResult(groupKey), onHeap.getLongResult(groupKey)); + } + } + onHeap.ensureCapacity(NUM_SLOTS); + offHeap.ensureCapacity(NUM_SLOTS); + for (int groupKey = 0; groupKey < NUM_SLOTS; groupKey++) { + Assert.assertEquals(offHeap.getLongResult(groupKey), onHeap.getLongResult(groupKey)); + } + } + } + + @Test + public void testDifferentialInt() { + Random random = new Random(RANDOM_SEED); + IntGroupByResultHolder onHeap = new IntGroupByResultHolder(INITIAL_CAPACITY, NUM_SLOTS, -1); + try (OffHeapIntGroupByResultHolder offHeap = + new OffHeapIntGroupByResultHolder(INITIAL_CAPACITY, NUM_SLOTS, -1)) { + int capacity = INITIAL_CAPACITY; + for (int i = 0; i < NUM_OPERATIONS; i++) { + int op = random.nextInt(10); + if (op == 0 && capacity < NUM_SLOTS) { + int newCapacity = Math.min(capacity + 1 + random.nextInt(MAX_GROWTH_STEP), NUM_SLOTS); + onHeap.ensureCapacity(newCapacity); + offHeap.ensureCapacity(newCapacity); + capacity = newCapacity; + } else if (op < 6) { + int groupKey = nextGroupKey(random, capacity); + int value = random.nextInt(); + onHeap.setValueForKey(groupKey, value); + offHeap.setValueForKey(groupKey, value); + } else { + int groupKey = nextGroupKey(random, capacity); + Assert.assertEquals(offHeap.getIntResult(groupKey), onHeap.getIntResult(groupKey)); + } + } + onHeap.ensureCapacity(NUM_SLOTS); + offHeap.ensureCapacity(NUM_SLOTS); + for (int groupKey = 0; groupKey < NUM_SLOTS; groupKey++) { + Assert.assertEquals(offHeap.getIntResult(groupKey), onHeap.getIntResult(groupKey)); + } + } + } + + @Test + public void testDefaultValueVisibilityDouble() { + double[] defaultValues = {0.0, 3.25, Double.NEGATIVE_INFINITY, Double.NaN}; + for (double defaultValue : defaultValues) { + try (OffHeapDoubleGroupByResultHolder holder = new OffHeapDoubleGroupByResultHolder(8, 1000, defaultValue)) { + assertDoubleEquals(holder.getDefaultValue(), defaultValue); + for (int i = 0; i < 8; i += 2) { + holder.setValueForKey(i, (double) i); + } + // Multiple growths; the extended tail must be visible as the default value after each one + holder.ensureCapacity(20); + holder.ensureCapacity(100); + holder.ensureCapacity(1000); + for (int i = 0; i < 8; i += 2) { + assertDoubleEquals(holder.getDoubleResult(i), i); + } + for (int i = 1; i < 8; i += 2) { + assertDoubleEquals(holder.getDoubleResult(i), defaultValue); + } + for (int i = 8; i < 1000; i++) { + assertDoubleEquals(holder.getDoubleResult(i), defaultValue); + } + assertDoubleEquals(holder.getDoubleResult(GroupKeyGenerator.INVALID_ID), defaultValue); + } + } + } + + @Test + public void testDefaultValueVisibilityLong() { + long[] defaultValues = {0L, -42L, Long.MIN_VALUE}; + for (long defaultValue : defaultValues) { + try (OffHeapLongGroupByResultHolder holder = new OffHeapLongGroupByResultHolder(8, 1000, defaultValue)) { + Assert.assertEquals(holder.getDefaultValue(), defaultValue); + for (int i = 0; i < 8; i += 2) { + holder.setValueForKey(i, (long) i); + } + holder.ensureCapacity(20); + holder.ensureCapacity(100); + holder.ensureCapacity(1000); + for (int i = 0; i < 8; i += 2) { + Assert.assertEquals(holder.getLongResult(i), i); + } + for (int i = 1; i < 8; i += 2) { + Assert.assertEquals(holder.getLongResult(i), defaultValue); + } + for (int i = 8; i < 1000; i++) { + Assert.assertEquals(holder.getLongResult(i), defaultValue); + } + Assert.assertEquals(holder.getLongResult(GroupKeyGenerator.INVALID_ID), defaultValue); + } + } + } + + @Test + public void testDefaultValueVisibilityInt() { + int[] defaultValues = {0, -42, Integer.MIN_VALUE}; + for (int defaultValue : defaultValues) { + try (OffHeapIntGroupByResultHolder holder = new OffHeapIntGroupByResultHolder(8, 1000, defaultValue)) { + Assert.assertEquals(holder.getDefaultValue(), defaultValue); + for (int i = 0; i < 8; i += 2) { + holder.setValueForKey(i, i + 1000); + } + holder.ensureCapacity(20); + holder.ensureCapacity(100); + holder.ensureCapacity(1000); + for (int i = 0; i < 8; i += 2) { + Assert.assertEquals(holder.getIntResult(i), i + 1000); + } + for (int i = 1; i < 8; i += 2) { + Assert.assertEquals(holder.getIntResult(i), defaultValue); + } + for (int i = 8; i < 1000; i++) { + Assert.assertEquals(holder.getIntResult(i), defaultValue); + } + Assert.assertEquals(holder.getIntResult(GroupKeyGenerator.INVALID_ID), defaultValue); + } + } + } + + @Test + public void testEnsureCapacityBeyondMaxCapacityThrows() { + try (OffHeapDoubleGroupByResultHolder doubleHolder = new OffHeapDoubleGroupByResultHolder(10, 100, 0.0); + OffHeapLongGroupByResultHolder longHolder = new OffHeapLongGroupByResultHolder(10, 100, 0L); + OffHeapIntGroupByResultHolder intHolder = new OffHeapIntGroupByResultHolder(10, 100, 0)) { + Assert.assertThrows(IllegalArgumentException.class, () -> doubleHolder.ensureCapacity(101)); + Assert.assertThrows(IllegalArgumentException.class, () -> longHolder.ensureCapacity(101)); + Assert.assertThrows(IllegalArgumentException.class, () -> intHolder.ensureCapacity(101)); + } + + // On-heap twins must behave identically + DoubleGroupByResultHolder onHeapDouble = new DoubleGroupByResultHolder(10, 100, 0.0); + LongGroupByResultHolder onHeapLong = new LongGroupByResultHolder(10, 100, 0L); + IntGroupByResultHolder onHeapInt = new IntGroupByResultHolder(10, 100, 0); + Assert.assertThrows(IllegalArgumentException.class, () -> onHeapDouble.ensureCapacity(101)); + Assert.assertThrows(IllegalArgumentException.class, () -> onHeapLong.ensureCapacity(101)); + Assert.assertThrows(IllegalArgumentException.class, () -> onHeapInt.ensureCapacity(101)); + } + + @Test + public void testGrowthClampsToMaxCapacity() { + // Doubling 10 -> 20 must clamp to maxCapacity 15; the clamped tail must be default-initialized and writable + try (OffHeapDoubleGroupByResultHolder holder = new OffHeapDoubleGroupByResultHolder(10, 15, -1.0)) { + holder.ensureCapacity(12); + for (int i = 10; i < 15; i++) { + assertDoubleEquals(holder.getDoubleResult(i), -1.0); + } + holder.setValueForKey(14, 42.0); + assertDoubleEquals(holder.getDoubleResult(14), 42.0); + holder.ensureCapacity(15); + assertDoubleEquals(holder.getDoubleResult(14), 42.0); + Assert.assertThrows(IllegalArgumentException.class, () -> holder.ensureCapacity(16)); + } + try (OffHeapLongGroupByResultHolder holder = new OffHeapLongGroupByResultHolder(10, 15, -1L)) { + holder.ensureCapacity(12); + for (int i = 10; i < 15; i++) { + Assert.assertEquals(holder.getLongResult(i), -1L); + } + holder.setValueForKey(14, 42L); + Assert.assertEquals(holder.getLongResult(14), 42L); + holder.ensureCapacity(15); + Assert.assertEquals(holder.getLongResult(14), 42L); + Assert.assertThrows(IllegalArgumentException.class, () -> holder.ensureCapacity(16)); + } + try (OffHeapIntGroupByResultHolder holder = new OffHeapIntGroupByResultHolder(10, 15, -1)) { + holder.ensureCapacity(12); + for (int i = 10; i < 15; i++) { + Assert.assertEquals(holder.getIntResult(i), -1); + } + holder.setValueForKey(14, 42); + Assert.assertEquals(holder.getIntResult(14), 42); + holder.ensureCapacity(15); + Assert.assertEquals(holder.getIntResult(14), 42); + Assert.assertThrows(IllegalArgumentException.class, () -> holder.ensureCapacity(16)); + } + } + + @Test + public void testUnsupportedTypedMethodsThrow() { + try (OffHeapDoubleGroupByResultHolder doubleHolder = new OffHeapDoubleGroupByResultHolder(10, 100, 0.0); + OffHeapLongGroupByResultHolder longHolder = new OffHeapLongGroupByResultHolder(10, 100, 0L); + OffHeapIntGroupByResultHolder intHolder = new OffHeapIntGroupByResultHolder(10, 100, 0)) { + Assert.assertThrows(UnsupportedOperationException.class, () -> doubleHolder.getIntResult(0)); + Assert.assertThrows(UnsupportedOperationException.class, () -> doubleHolder.getLongResult(0)); + Assert.assertThrows(UnsupportedOperationException.class, () -> doubleHolder.getResult(0)); + Assert.assertThrows(UnsupportedOperationException.class, () -> doubleHolder.setValueForKey(0, 1)); + Assert.assertThrows(UnsupportedOperationException.class, () -> doubleHolder.setValueForKey(0, 1L)); + Assert.assertThrows(UnsupportedOperationException.class, () -> doubleHolder.setValueForKey(0, (Object) "v")); + + Assert.assertThrows(UnsupportedOperationException.class, () -> longHolder.getDoubleResult(0)); + Assert.assertThrows(UnsupportedOperationException.class, () -> longHolder.getIntResult(0)); + Assert.assertThrows(UnsupportedOperationException.class, () -> longHolder.getResult(0)); + Assert.assertThrows(UnsupportedOperationException.class, () -> longHolder.setValueForKey(0, 1.0)); + Assert.assertThrows(UnsupportedOperationException.class, () -> longHolder.setValueForKey(0, 1)); + Assert.assertThrows(UnsupportedOperationException.class, () -> longHolder.setValueForKey(0, (Object) "v")); + + Assert.assertThrows(UnsupportedOperationException.class, () -> intHolder.getDoubleResult(0)); + Assert.assertThrows(UnsupportedOperationException.class, () -> intHolder.getLongResult(0)); + Assert.assertThrows(UnsupportedOperationException.class, () -> intHolder.getResult(0)); + Assert.assertThrows(UnsupportedOperationException.class, () -> intHolder.setValueForKey(0, 1.0)); + Assert.assertThrows(UnsupportedOperationException.class, () -> intHolder.setValueForKey(0, 1L)); + Assert.assertThrows(UnsupportedOperationException.class, () -> intHolder.setValueForKey(0, (Object) "v")); + } + } + + @Test + public void testZeroInitialCapacity() { + try (OffHeapDoubleGroupByResultHolder holder = new OffHeapDoubleGroupByResultHolder(0, 100, -1.0)) { + holder.ensureCapacity(10); + assertDoubleEquals(holder.getDoubleResult(5), -1.0); + holder.setValueForKey(5, 3.5); + assertDoubleEquals(holder.getDoubleResult(5), 3.5); + } + try (OffHeapLongGroupByResultHolder holder = new OffHeapLongGroupByResultHolder(0, 100, -1L)) { + holder.ensureCapacity(10); + Assert.assertEquals(holder.getLongResult(5), -1L); + holder.setValueForKey(5, 3L); + Assert.assertEquals(holder.getLongResult(5), 3L); + } + try (OffHeapIntGroupByResultHolder holder = new OffHeapIntGroupByResultHolder(0, 100, -1)) { + holder.ensureCapacity(10); + Assert.assertEquals(holder.getIntResult(5), -1); + holder.setValueForKey(5, 3); + Assert.assertEquals(holder.getIntResult(5), 3); + } + } + + /// Exercises the wrapper-based fallback arms of the holder accessors and fills that normally only run for + /// buffers beyond the 2GB view limit. + @Test + public void testDifferentialWithoutViews() { + OffHeapGroupByUtils.setViewSizeLimitBytes(0); + try { + testDifferentialDouble(); + testDifferentialLong(); + testDifferentialInt(); + testDefaultValueVisibilityDouble(); + } finally { + OffHeapGroupByUtils.setViewSizeLimitBytes(Integer.MAX_VALUE); + } + } + + @Test + public void testCloseIsIdempotent() { + OffHeapDoubleGroupByResultHolder doubleHolder = new OffHeapDoubleGroupByResultHolder(10, 100, 0.0); + doubleHolder.close(); + doubleHolder.close(); + + OffHeapLongGroupByResultHolder longHolder = new OffHeapLongGroupByResultHolder(10, 100, 0L); + longHolder.close(); + longHolder.close(); + + OffHeapIntGroupByResultHolder intHolder = new OffHeapIntGroupByResultHolder(10, 100, 0); + intHolder.close(); + intHolder.close(); + } + + @Test + public void testNoDirectMemoryLeak() { + long baseline = PinotDataBuffer.getDirectBufferUsage(); + try (OffHeapDoubleGroupByResultHolder doubleHolder = new OffHeapDoubleGroupByResultHolder(128, 4096, -1.0); + OffHeapLongGroupByResultHolder longHolder = new OffHeapLongGroupByResultHolder(128, 4096, -1L); + OffHeapIntGroupByResultHolder intHolder = new OffHeapIntGroupByResultHolder(128, 4096, -1)) { + // Grow several times so intermediate buffers are allocated and released along the way + for (int capacity : new int[]{256, 1000, 4096}) { + doubleHolder.ensureCapacity(capacity); + longHolder.ensureCapacity(capacity); + intHolder.ensureCapacity(capacity); + } + for (int i = 0; i < 4096; i++) { + doubleHolder.setValueForKey(i, (double) i); + longHolder.setValueForKey(i, (long) i); + intHolder.setValueForKey(i, i); + } + Assert.assertTrue(PinotDataBuffer.getDirectBufferUsage() > baseline); + } + Assert.assertEquals(PinotDataBuffer.getDirectBufferUsage(), baseline); + } + + private static int nextGroupKey(Random random, int capacity) { + // Roughly 5% of the accesses target INVALID_ID to exercise the guard paths + return random.nextInt(20) == 0 ? GroupKeyGenerator.INVALID_ID : random.nextInt(capacity); + } + + private static void assertDoubleEquals(double actual, double expected) { + // Bit-wise comparison so that NaN default values are asserted correctly + Assert.assertEquals(Double.doubleToLongBits(actual), Double.doubleToLongBits(expected)); + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByUtilsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByUtilsTest.java new file mode 100644 index 000000000000..c011c173cb37 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByUtilsTest.java @@ -0,0 +1,145 @@ +/** + * 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.offheap; + +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Random; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; + + +public class OffHeapGroupByUtilsTest { + + @Test + public void testEncodeUtf8MatchesJdkEncoder() { + // Deliberate boundary and malformed cases. The encoder must be byte-for-byte identical to + // String#getBytes(StandardCharsets.UTF_8), including '?' replacement of unpaired surrogates. + String[] cases = { + "", + "ascii only", + " ", + "", // 1-byte upper bound + "€", // 2-byte lower bound + "߿", // 2-byte upper bound + "ࠀ", // 3-byte lower bound + "퟿", // last char before the surrogate range + "", // first char after the surrogate range + "￿", // 3-byte upper bound + "café 你好", // mixed 1/2/3-byte + "😀", // valid surrogate pair (emoji, 4-byte) + "a😀b🎉c", // pairs embedded in text + "\uD800", // unpaired high surrogate at end + "\uDC00", // unpaired low surrogate + "\uD800a", // high surrogate followed by a normal char + "\uD800𐀀", // unpaired high followed by a valid pair + "\uDC00\uD800", // low then high (both unpaired) + "x\uD800", // trailing unpaired high + }; + for (String value : cases) { + assertEncodeMatches(value); + } + + // Randomized: arbitrary char sequences (freely mixing valid text and surrogate salad) + Random random = new Random(42); + for (int i = 0; i < 10_000; i++) { + int length = random.nextInt(32); + char[] chars = new char[length]; + int charIndex = 0; + while (charIndex < length) { + switch (random.nextInt(5)) { + case 0: + chars[charIndex++] = (char) random.nextInt(0x80); + break; + case 1: + chars[charIndex++] = (char) random.nextInt(0x800); + break; + case 2: + chars[charIndex++] = (char) random.nextInt(0x10000); + break; + case 3: + chars[charIndex++] = (char) (Character.MIN_SURROGATE + random.nextInt( + Character.MAX_SURROGATE - Character.MIN_SURROGATE + 1)); + break; + default: + // Frequently emit valid pairs so the 4-byte path is well covered + if (charIndex + 1 < length) { + int codePoint = 0x10000 + random.nextInt(0x100000); + chars[charIndex++] = Character.highSurrogate(codePoint); + chars[charIndex++] = Character.lowSurrogate(codePoint); + } else { + chars[charIndex++] = 'z'; + } + break; + } + } + assertEncodeMatches(new String(chars)); + } + } + + private static void assertEncodeMatches(String value) { + byte[] scratch = new byte[value.length() * 3 + 1]; + int length = OffHeapGroupByUtils.encodeUtf8(value, scratch); + byte[] expected = value.getBytes(StandardCharsets.UTF_8); + assertEquals(Arrays.copyOf(scratch, length), expected, + "encodeUtf8 mismatch for chars: " + Arrays.toString(value.chars().toArray())); + } + + @Test + public void testPackUnpackIntsRoundTrip() { + Random random = new Random(42); + for (int iteration = 0; iteration < 1000; iteration++) { + int numValues = 1 + random.nextInt(8); + int[] values = new int[numValues]; + for (int i = 0; i < numValues; i++) { + // Include negative sentinels like ID_FOR_NULL (-2) and extremes + switch (random.nextInt(4)) { + case 0: + values[i] = random.nextInt(); + break; + case 1: + values[i] = -2; + break; + case 2: + values[i] = Integer.MIN_VALUE; + break; + default: + values[i] = random.nextInt(100); + break; + } + } + byte[] scratch = new byte[numValues * Integer.BYTES]; + int length = OffHeapGroupByUtils.packInts(values, numValues, scratch); + assertEquals(length, numValues * Integer.BYTES); + int[] unpacked = new int[numValues]; + OffHeapGroupByUtils.unpackInts(scratch, numValues, unpacked); + assertEquals(unpacked, values); + } + } + + @Test + public void testEnsureByteCapacity() { + byte[] scratch = new byte[8]; + assertEquals(OffHeapGroupByUtils.ensureByteCapacity(scratch, 8), scratch); + assertEquals(OffHeapGroupByUtils.ensureByteCapacity(scratch, 4), scratch); + assertEquals(OffHeapGroupByUtils.ensureByteCapacity(scratch, 9).length, 16); + assertEquals(OffHeapGroupByUtils.ensureByteCapacity(scratch, 100).length, 100); + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapIntGroupIdMapTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapIntGroupIdMapTest.java new file mode 100644 index 000000000000..d85ccedbcc7a --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapIntGroupIdMapTest.java @@ -0,0 +1,170 @@ +/** + * 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.offheap; + +import it.unimi.dsi.fastutil.ints.Int2IntOpenHashMap; +import java.util.HashMap; +import java.util.Iterator; +import java.util.Map; +import java.util.Random; +import org.apache.pinot.core.query.aggregation.groupby.DictionaryBasedGroupKeyGenerator.IntGroupIdMap; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + + +public class OffHeapIntGroupIdMapTest { + private static final int GROUP_ID_UPPER_BOUND = Integer.MAX_VALUE; + + @Test + public void testDifferentialAgainstOnHeapIntGroupIdMap() { + try (OffHeapIntGroupIdMap offHeapMap = new OffHeapIntGroupIdMap(0)) { + IntGroupIdMap onHeapMap = new IntGroupIdMap(); + Random random = new Random(42); + for (int i = 0; i < 200_000; i++) { + // ~50% duplicates; keys include 0 and Integer.MAX_VALUE + int rawKey = random.nextBoolean() ? random.nextInt(50_000) : switch (random.nextInt(3)) { + case 0 -> 0; + case 1 -> Integer.MAX_VALUE; + default -> random.nextInt(Integer.MAX_VALUE); + }; + assertEquals(offHeapMap.getGroupId(rawKey, GROUP_ID_UPPER_BOUND), + onHeapMap.getGroupId(rawKey, GROUP_ID_UPPER_BOUND), "Mismatch for key: " + rawKey + " at op " + i); + } + assertEquals(offHeapMap.size(), onHeapMap.size()); + + // Iterator parity as sets of (rawKey -> groupId) + Map offHeapEntries = new HashMap<>(); + Iterator offHeapIterator = offHeapMap.iterator(); + while (offHeapIterator.hasNext()) { + OffHeapIntGroupIdMap.Entry entry = offHeapIterator.next(); + offHeapEntries.put(entry._rawKey, entry._groupId); + } + Map onHeapEntries = new HashMap<>(); + Iterator onHeapIterator = onHeapMap.iterator(); + while (onHeapIterator.hasNext()) { + IntGroupIdMap.Entry entry = onHeapIterator.next(); + onHeapEntries.put(entry._rawKey, entry._groupId); + } + assertEquals(offHeapEntries, onHeapEntries); + } + } + + @Test + public void testCapSemantics() { + try (OffHeapIntGroupIdMap map = new OffHeapIntGroupIdMap(0)) { + Int2IntOpenHashMap reference = new Int2IntOpenHashMap(); + for (int key = 0; key < 150; key++) { + int groupId = map.getGroupId(key * 31, 100); + if (key < 100) { + assertEquals(groupId, key); + reference.put(key * 31, groupId); + } else { + assertEquals(groupId, OffHeapIntGroupIdMap.INVALID_ID); + } + } + assertEquals(map.size(), 100); + // Existing keys still resolve at cap; rejected keys were not inserted + for (int key = 0; key < 100; key++) { + assertEquals(map.getGroupId(key * 31, 100), reference.get(key * 31)); + } + assertEquals(map.getGroupId(149 * 31, 100), OffHeapIntGroupIdMap.INVALID_ID); + // Raising the bound assigns the next dense id + assertEquals(map.getGroupId(149 * 31, 101), 100); + } + } + + @Test + public void testGrowthAcrossResizes() { + try (OffHeapIntGroupIdMap map = new OffHeapIntGroupIdMap(0)) { + for (int key = 0; key < 100_000; key++) { + assertEquals(map.getGroupId(key, GROUP_ID_UPPER_BOUND), key); + } + // All ids stable after many resizes + for (int key = 0; key < 100_000; key++) { + assertEquals(map.getGroupId(key, GROUP_ID_UPPER_BOUND), key); + } + // 100K entries at load factor 0.5 over 8-byte slots: 262144 slots * 8 bytes + assertEquals(map.getOffHeapMemoryBytes(), 262_144L * 8); + } + } + + @Test + public void testMinusOneKeyOutOfBand() { + // -1 first + try (OffHeapIntGroupIdMap map = new OffHeapIntGroupIdMap(0)) { + assertEquals(map.getGroupId(-1, GROUP_ID_UPPER_BOUND), 0); + assertEquals(map.getGroupId(7, GROUP_ID_UPPER_BOUND), 1); + assertEquals(map.getGroupId(-1, GROUP_ID_UPPER_BOUND), 0); + assertEquals(map.size(), 2); + Map entries = new HashMap<>(); + Iterator iterator = map.iterator(); + while (iterator.hasNext()) { + OffHeapIntGroupIdMap.Entry entry = iterator.next(); + entries.put(entry._rawKey, entry._groupId); + } + assertEquals(entries, Map.of(-1, 0, 7, 1)); + } + // -1 mid-stream + try (OffHeapIntGroupIdMap map = new OffHeapIntGroupIdMap(0)) { + assertEquals(map.getGroupId(10, GROUP_ID_UPPER_BOUND), 0); + assertEquals(map.getGroupId(-1, GROUP_ID_UPPER_BOUND), 1); + assertEquals(map.getGroupId(20, GROUP_ID_UPPER_BOUND), 2); + assertEquals(map.size(), 3); + } + // -1 rejected at the cap while existing keys still resolve + try (OffHeapIntGroupIdMap map = new OffHeapIntGroupIdMap(0)) { + assertEquals(map.getGroupId(10, 1), 0); + assertEquals(map.getGroupId(-1, 1), OffHeapIntGroupIdMap.INVALID_ID); + assertEquals(map.getGroupId(10, 1), 0); + assertEquals(map.size(), 1); + // Raising the bound assigns the next dense id + assertEquals(map.getGroupId(-1, 2), 1); + } + } + + /// Exercises the wrapper-based fallback arms (probe, expand, zero-fill) that normally only run for buffers + /// beyond the 2GB view limit. + @Test + public void testDifferentialWithoutViews() { + OffHeapGroupByUtils.setViewSizeLimitBytes(0); + try { + testDifferentialAgainstOnHeapIntGroupIdMap(); + testGrowthAcrossResizes(); + testMinusOneKeyOutOfBand(); + } finally { + OffHeapGroupByUtils.setViewSizeLimitBytes(Integer.MAX_VALUE); + } + } + + @Test + public void testCloseTwiceAndNoLeak() { + long baseline = PinotDataBuffer.getDirectBufferUsage(); + OffHeapIntGroupIdMap map = new OffHeapIntGroupIdMap(0); + for (int key = 0; key < 10_000; key++) { + map.getGroupId(key, GROUP_ID_UPPER_BOUND); + } + assertTrue(PinotDataBuffer.getDirectBufferUsage() > baseline); + map.close(); + map.close(); + assertEquals(PinotDataBuffer.getDirectBufferUsage(), baseline); + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapLongGroupIdMapTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapLongGroupIdMapTest.java new file mode 100644 index 000000000000..89e69a38d11d --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapLongGroupIdMapTest.java @@ -0,0 +1,263 @@ +/** + * 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.offheap; + +import it.unimi.dsi.fastutil.longs.Long2IntOpenHashMap; +import java.util.HashMap; +import java.util.Iterator; +import java.util.Map; +import java.util.Random; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + + +/// Unit test for [OffHeapLongGroupIdMap], using [Long2IntOpenHashMap] as the on-heap reference +/// implementation of the same dense group-id assignment contract. +public class OffHeapLongGroupIdMapTest { + private static final int INVALID_ID = OffHeapLongGroupIdMap.INVALID_ID; + + private static Long2IntOpenHashMap newReference() { + Long2IntOpenHashMap reference = new Long2IntOpenHashMap(); + reference.defaultReturnValue(INVALID_ID); + return reference; + } + + /// Reference implementation of the dense-id contract: existing keys always resolve; new keys get id = size() when + /// under the upper bound, otherwise INVALID_ID without insertion. + private static int referenceGetGroupId(Long2IntOpenHashMap reference, long rawKey, int groupIdUpperBound) { + int groupId = reference.get(rawKey); + if (groupId != INVALID_ID) { + return groupId; + } + if (reference.size() < groupIdUpperBound) { + groupId = reference.size(); + reference.put(rawKey, groupId); + return groupId; + } + return INVALID_ID; + } + + private static void verifyIteratorMatchesReference(OffHeapLongGroupIdMap map, Long2IntOpenHashMap reference) { + Map actual = new HashMap<>(); + long lastRawKey = Long.MIN_VALUE; + Iterator iterator = map.iterator(); + while (iterator.hasNext()) { + OffHeapLongGroupIdMap.Entry entry = iterator.next(); + // The Entry is a reused flyweight, so copy the values out + assertNull(actual.put(entry._rawKey, entry._groupId), "Iterator yielded duplicate key: " + entry._rawKey); + lastRawKey = entry._rawKey; + } + assertEquals(actual.size(), reference.size(), "Iterator must yield exactly size() entries"); + for (Map.Entry entry : actual.entrySet()) { + assertEquals(entry.getValue().intValue(), reference.get(entry.getKey().longValue()), + "Mismatch for key: " + entry.getKey()); + } + if (reference.containsKey(0L)) { + assertEquals(lastRawKey, 0L, "Zero-key entry must be yielded last"); + } + } + + @Test + public void testDifferentialAgainstReference() { + Random random = new Random(42); + long[] keys = new long[200_000]; + for (int i = 0; i < keys.length; i++) { + if (i > 0 && random.nextBoolean()) { + // Duplicate of an earlier key + keys[i] = keys[random.nextInt(i)]; + } else { + keys[i] = random.nextLong(); + } + } + // Force the zero key into the stream (random longs will essentially never produce it) + keys[1_000] = 0; + keys[2_000] = 0; + Long2IntOpenHashMap reference = newReference(); + try (OffHeapLongGroupIdMap map = new OffHeapLongGroupIdMap(1024)) { + for (long key : keys) { + assertEquals(map.getGroupId(key, Integer.MAX_VALUE), referenceGetGroupId(reference, key, Integer.MAX_VALUE), + "Mismatch for key: " + key); + } + assertEquals(map.size(), reference.size()); + verifyIteratorMatchesReference(map, reference); + } + } + + @Test + public void testZeroKeyFirst() { + try (OffHeapLongGroupIdMap map = new OffHeapLongGroupIdMap(16)) { + assertEquals(map.getGroupId(0, 10), 0); + assertEquals(map.getGroupId(0, 10), 0); + assertEquals(map.getGroupId(42, 10), 1); + assertEquals(map.getGroupId(0, 10), 0); + assertEquals(map.size(), 2); + } + } + + @Test + public void testZeroKeyMid() { + try (OffHeapLongGroupIdMap map = new OffHeapLongGroupIdMap(16)) { + for (int i = 0; i < 10; i++) { + assertEquals(map.getGroupId(i + 1, 100), i); + } + assertEquals(map.getGroupId(0, 100), 10); + assertEquals(map.getGroupId(0, 100), 10); + assertEquals(map.getGroupId(11, 100), 11); + assertEquals(map.size(), 12); + } + } + + @Test + public void testZeroKeyAtCap() { + try (OffHeapLongGroupIdMap map = new OffHeapLongGroupIdMap(16)) { + for (int i = 0; i < 5; i++) { + assertEquals(map.getGroupId(i + 1, 5), i); + } + // At the cap: the zero key must be rejected and not inserted + assertEquals(map.getGroupId(0, 5), INVALID_ID); + assertEquals(map.size(), 5); + // With a larger bound it gets the next dense id, proving the rejection did not insert it + assertEquals(map.getGroupId(0, 6), 5); + // Present keys always resolve, even when size() >= upper bound + assertEquals(map.getGroupId(0, 5), 5); + assertEquals(map.size(), 6); + } + } + + @Test + public void testNegativeKeys() { + long[] keys = {Long.MIN_VALUE, -1, Long.MAX_VALUE, -123_456_789L}; + Long2IntOpenHashMap reference = newReference(); + try (OffHeapLongGroupIdMap map = new OffHeapLongGroupIdMap(16)) { + for (int i = 0; i < keys.length; i++) { + assertEquals(map.getGroupId(keys[i], 100), i); + referenceGetGroupId(reference, keys[i], 100); + } + for (int i = 0; i < keys.length; i++) { + assertEquals(map.getGroupId(keys[i], 100), i); + } + assertEquals(map.size(), keys.length); + verifyIteratorMatchesReference(map, reference); + } + } + + @Test + public void testCapSemantics() { + try (OffHeapLongGroupIdMap map = new OffHeapLongGroupIdMap(256)) { + for (int i = 0; i < 150; i++) { + assertEquals(map.getGroupId(i + 1, 100), i < 100 ? i : INVALID_ID); + } + assertEquals(map.size(), 100); + // Existing keys still resolve at the cap; rejected keys were never inserted + for (int i = 0; i < 150; i++) { + assertEquals(map.getGroupId(i + 1, 100), i < 100 ? i : INVALID_ID); + } + assertEquals(map.size(), 100); + Iterator iterator = map.iterator(); + int numEntries = 0; + while (iterator.hasNext()) { + OffHeapLongGroupIdMap.Entry entry = iterator.next(); + assertEquals(entry._groupId, (int) entry._rawKey - 1); + numEntries++; + } + assertEquals(numEntries, 100); + } + } + + @Test + public void testUpperBoundSmallerThanCurrentSize() { + try (OffHeapLongGroupIdMap map = new OffHeapLongGroupIdMap(64)) { + for (int i = 0; i < 50; i++) { + assertEquals(map.getGroupId(i + 1, 1000), i); + } + assertEquals(map.getGroupId(0, 1000), 50); + // An upper bound smaller than the current size never breaks existing lookups + for (int i = 0; i < 50; i++) { + assertEquals(map.getGroupId(i + 1, 1), i); + } + assertEquals(map.getGroupId(0, 1), 50); + // But it rejects new keys + assertEquals(map.getGroupId(9999, 1), INVALID_ID); + assertEquals(map.size(), 51); + } + } + + @Test + public void testGrowthAcrossMultipleResizes() { + int numKeys = 100_000; + try (OffHeapLongGroupIdMap map = new OffHeapLongGroupIdMap(10)) { + // Initial capacity is max(512, ceilPow2(expected * 2)) = 512 slots of 16 bytes + assertEquals(map.getOffHeapMemoryBytes(), 512L * 16); + for (int i = 0; i < numKeys; i++) { + assertEquals(map.getGroupId(i + 1, Integer.MAX_VALUE), i); + } + assertEquals(map.size(), numKeys); + // Load factor 0.5: the smallest power-of-two capacity with capacity / 2 >= 100_000 is 262144 + assertEquals(map.getOffHeapMemoryBytes(), 262_144L * 16); + // Ids are untouched by resize: every key still resolves to its original id + for (int i = 0; i < numKeys; i++) { + assertEquals(map.getGroupId(i + 1, Integer.MAX_VALUE), i); + } + assertEquals(map.size(), numKeys); + } + } + + /// Exercises the wrapper-based fallback arms (probe, expand, zero-fill) that normally only run for buffers + /// beyond the 2GB view limit. + @Test + public void testDifferentialWithoutViews() { + OffHeapGroupByUtils.setViewSizeLimitBytes(0); + try { + testDifferentialAgainstReference(); + testGrowthAcrossMultipleResizes(); + } finally { + OffHeapGroupByUtils.setViewSizeLimitBytes(Integer.MAX_VALUE); + } + } + + @Test + public void testCloseTwiceIsSafe() { + OffHeapLongGroupIdMap map = new OffHeapLongGroupIdMap(16); + assertEquals(map.getGroupId(123, 10), 0); + map.close(); + map.close(); + } + + @Test + public void testNoDirectMemoryLeak() { + long baseline = PinotDataBuffer.getDirectBufferUsage(); + OffHeapLongGroupIdMap map = new OffHeapLongGroupIdMap(10); + try { + // Force multiple resizes; each resize must close the old buffer eagerly + for (int i = 0; i < 10_000; i++) { + map.getGroupId(i + 1, Integer.MAX_VALUE); + } + assertTrue(PinotDataBuffer.getDirectBufferUsage() > baseline); + assertEquals(PinotDataBuffer.getDirectBufferUsage() - baseline, map.getOffHeapMemoryBytes(), + "Only the current hash table buffer should be alive after resizes"); + } finally { + map.close(); + } + assertEquals(PinotDataBuffer.getDirectBufferUsage(), baseline, "Direct buffer usage must return to baseline"); + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapUltraLogLogGroupByResultHolderTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapUltraLogLogGroupByResultHolderTest.java new file mode 100644 index 000000000000..edbc48cd8a45 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapUltraLogLogGroupByResultHolderTest.java @@ -0,0 +1,178 @@ +/** + * 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.offheap; + +import com.dynatrace.hash4j.distinctcount.UltraLogLog; +import java.util.HashMap; +import java.util.Map; +import java.util.Random; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; +import org.testng.Assert; +import org.testng.annotations.Test; + + +/// Tests for [OffHeapUltraLogLogGroupByResultHolder], pinning the vendored register-update math byte-identical +/// to hash4j's [UltraLogLog] across precisions and asserting on-heap-equivalent holder semantics (untouched +/// groups, growth, delegate mode, close). +public class OffHeapUltraLogLogGroupByResultHolderTest { + private static final long RANDOM_SEED = 42; + private static final int NUM_SLOTS = 2000; + private static final int INITIAL_CAPACITY = 16; + private static final int NUM_OPERATIONS = 200_000; + // Register-update edge cases on top of random hashes + private static final long[] EDGE_HASHES = {0L, -1L, 1L, Long.MIN_VALUE, Long.MAX_VALUE, 0x8000000000000001L}; + + @Test + public void testDifferentialAddAcrossPrecisions() { + // p=3 is the library minimum (8B slots), 12 is the Pinot default (4KB), 18 fills a whole chunk, 19 exceeds + // the target chunk size (one slot per chunk) + for (int p : new int[]{3, 8, 12, 18, 19}) { + runDifferential(p, p >= 18 ? 40 : NUM_SLOTS, p >= 18 ? 20_000 : NUM_OPERATIONS); + } + } + + @Test + public void testDifferentialWithoutViews() { + // Force the PinotDataBuffer wrapper fallback arm of every view fast path + OffHeapGroupByUtils.setViewSizeLimitBytes(0); + try { + runDifferential(12, 500, 50_000); + } finally { + OffHeapGroupByUtils.setViewSizeLimitBytes(Integer.MAX_VALUE); + } + } + + private void runDifferential(int p, int numSlots, int numOperations) { + Random random = new Random(RANDOM_SEED + p); + Map reference = new HashMap<>(); + try (OffHeapUltraLogLogGroupByResultHolder offHeap = + new OffHeapUltraLogLogGroupByResultHolder(p, INITIAL_CAPACITY, numSlots)) { + int capacity = INITIAL_CAPACITY; + for (int i = 0; i < numOperations; i++) { + int op = random.nextInt(20); + if (op == 0 && capacity < numSlots) { + int newCapacity = Math.min(capacity + 1 + random.nextInt(numSlots / 4), numSlots); + offHeap.ensureCapacity(newCapacity); + capacity = newCapacity; + } else if (op == 1) { + int groupKey = random.nextInt(capacity); + offHeap.touch(groupKey); + reference.computeIfAbsent(groupKey, k -> UltraLogLog.create(p)); + } else { + int groupKey = random.nextInt(capacity); + long hashValue = op < 5 ? EDGE_HASHES[random.nextInt(EDGE_HASHES.length)] : random.nextLong(); + offHeap.add(groupKey, hashValue); + reference.computeIfAbsent(groupKey, k -> UltraLogLog.create(p)).add(hashValue); + } + } + offHeap.ensureCapacity(numSlots); + for (int groupKey = 0; groupKey < numSlots; groupKey++) { + UltraLogLog expected = reference.get(groupKey); + UltraLogLog actual = offHeap.getResult(groupKey); + if (expected == null) { + Assert.assertNull(actual, "untouched group " + groupKey + " must materialize as null"); + } else { + Assert.assertNotNull(actual, "touched group " + groupKey + " must not materialize as null"); + Assert.assertEquals(actual.getState(), expected.getState(), + "state bytes diverged from hash4j for group " + groupKey + " at p=" + p); + Assert.assertEquals(actual.getDistinctCountEstimate(), expected.getDistinctCountEstimate()); + } + } + } + } + + @Test + public void testTouchCreatesEmptyState() { + try (OffHeapUltraLogLogGroupByResultHolder holder = + new OffHeapUltraLogLogGroupByResultHolder(12, INITIAL_CAPACITY, NUM_SLOTS)) { + Assert.assertNull(holder.getResult(0)); + holder.touch(0); + UltraLogLog materialized = holder.getResult(0); + Assert.assertNotNull(materialized); + Assert.assertEquals(materialized.getState(), UltraLogLog.create(12).getState()); + // Materialization returns a copy: mutating it must not touch the slot + materialized.add(12345L); + Assert.assertEquals(((UltraLogLog) holder.getResult(0)).getState(), UltraLogLog.create(12).getState()); + } + } + + @Test + public void testDelegateMode() { + try (OffHeapUltraLogLogGroupByResultHolder holder = + new OffHeapUltraLogLogGroupByResultHolder(12, INITIAL_CAPACITY, NUM_SLOTS)) { + Assert.assertNull(holder.getResult(3)); + Object dictWrapper = new Object(); + holder.setValueForKey(3, dictWrapper); + Assert.assertSame(holder.getResult(3), dictWrapper); + Assert.assertNull(holder.getResult(4)); + // Growth must apply to the delegate as well + holder.ensureCapacity(NUM_SLOTS); + holder.setValueForKey(NUM_SLOTS - 1, dictWrapper); + Assert.assertSame(holder.getResult(NUM_SLOTS - 1), dictWrapper); + } + } + + @Test + public void testInvalidId() { + try (OffHeapUltraLogLogGroupByResultHolder holder = + new OffHeapUltraLogLogGroupByResultHolder(12, INITIAL_CAPACITY, NUM_SLOTS)) { + holder.touch(GroupKeyGenerator.INVALID_ID); + holder.add(GroupKeyGenerator.INVALID_ID, 123L); + holder.setValueForKey(GroupKeyGenerator.INVALID_ID, new Object()); + Assert.assertNull(holder.getResult(GroupKeyGenerator.INVALID_ID)); + // None of the INVALID_ID calls may have created state or a delegate + Assert.assertNull(holder.getResult(0)); + } + } + + @Test + public void testCloseReleasesDirectMemoryAndIsIdempotent() { + long baseline = PinotDataBuffer.getDirectBufferUsage(); + OffHeapUltraLogLogGroupByResultHolder holder = + new OffHeapUltraLogLogGroupByResultHolder(12, INITIAL_CAPACITY, NUM_SLOTS); + holder.add(0, 42L); + Assert.assertTrue(PinotDataBuffer.getDirectBufferUsage() > baseline, + "adding a value must allocate a direct-memory chunk"); + holder.close(); + Assert.assertEquals(PinotDataBuffer.getDirectBufferUsage(), baseline, + "close must release all direct memory"); + holder.close(); + Assert.assertEquals(PinotDataBuffer.getDirectBufferUsage(), baseline); + } + + @Test + public void testOutOfRangePrecisionRejected() { + // p sizes direct memory as 1 << p without going through UltraLogLog.create's own bound check, so the holder + // must reject out-of-range p itself (p in [27, 30] would allocate up to 1GB per group; p > 30 would let + // register indexes walk outside the slot after int-shift wrap) + for (int p : new int[]{Integer.MIN_VALUE, -1, 0, 2, 27, 30, 32, 40, Integer.MAX_VALUE}) { + Assert.assertThrows(IllegalArgumentException.class, + () -> new OffHeapUltraLogLogGroupByResultHolder(p, INITIAL_CAPACITY, NUM_SLOTS)); + } + } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void testEnsureCapacityBeyondMaxThrows() { + try (OffHeapUltraLogLogGroupByResultHolder holder = + new OffHeapUltraLogLogGroupByResultHolder(12, INITIAL_CAPACITY, NUM_SLOTS)) { + holder.ensureCapacity(NUM_SLOTS + 1); + } + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/OffHeapGroupByQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/OffHeapGroupByQueriesTest.java new file mode 100644 index 000000000000..82c22ee6cb1a --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/queries/OffHeapGroupByQueriesTest.java @@ -0,0 +1,829 @@ +/** + * 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.queries; + +import com.dynatrace.hash4j.distinctcount.UltraLogLog; +import java.io.File; +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Random; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.response.broker.BrokerResponseNative; +import org.apache.pinot.common.response.broker.ResultTable; +import org.apache.pinot.core.common.ObjectSerDeUtils; +import org.apache.pinot.core.operator.BaseProjectOperator; +import org.apache.pinot.core.plan.DocIdSetPlanNode; +import org.apache.pinot.core.plan.ProjectPlanNode; +import org.apache.pinot.core.query.aggregation.groupby.DefaultGroupByExecutor; +import org.apache.pinot.core.query.aggregation.groupby.offheap.ResourceTrackingGroupKeyGenerator; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.local.utils.UltraLogLogUtils; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.SegmentContext; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; +import org.apache.pinot.spi.config.table.StarTreeIndexConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.query.QueryThreadContext; +import org.apache.pinot.spi.utils.BytesUtils; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.CommonConstants.Server; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.apache.pinot.sql.parsers.CalciteSqlParser; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + +/// End-to-end differential test for off-heap SSE GROUP BY (`SET groupByOffHeap=true`). +/// +/// Every query in the battery is executed twice over the same segments — once on-heap and once with the +/// `groupByOffHeap` query option — and the result rows must match (as unordered multisets for un-ordered +/// queries, exact ordered lists for ORDER BY queries). Numeric exactness: all FLOAT/DOUBLE/BIG_DECIMAL values in +/// the fixtures are dyadic rationals (multiples of 0.25/0.5), so double aggregation is exact and independent of +/// combine order, allowing exact value comparison. +/// +/// The fixture cardinalities are chosen against the holder-selection logic in the +/// [DictionaryBasedGroupKeyGenerator][org.apache.pinot.core.query.aggregation.groupby.DictionaryBasedGroupKeyGenerator] +/// constructor (arrayBasedThreshold = 10_000, default numGroupsLimit = 100_000): +/// - `dLowStr(10) * dInt(1000) = 10_000` — not > threshold -> ARRAY_BASED (stays on-heap by design) +/// - `dHighStr(20_000)` alone or with others up to `4 * 10^8` -> INT_MAP_BASED +/// - `dHighStr * dHighInt * dLowStr = 4 * 10^9 > Integer.MAX_VALUE` -> LONG_MAP_BASED +/// - `dHighStr * dHighInt * dHighInt2 * dHighLong * dInt = 1.6 * 10^20 > Long.MAX_VALUE` -> ARRAY_MAP_BASED +/// +/// After every off-heap execution the test asserts that [PinotDataBuffer#getDirectBufferUsage()] returns to the +/// baseline captured after a warm-up query, proving the close path releases all direct memory. (The baseline is +/// captured after one warm-up off-heap query instead of right after segment load so that any lazily allocated +/// direct buffer elsewhere in the query stack cannot move the baseline mid-test; segments are loaded with mmap, +/// which is not counted as direct usage.) +/// +/// Null handling runs against a second small segment fixture (built in the same class) whose columns carry nulls +/// in three patterns: first row null, null only mid-stream (with new values appearing after the null), and all +/// nulls — exercising the off-heap "null shift" logic in NoDictionarySingleColumnGroupKeyGenerator. Since the +/// on-heap null-group counting fix (both modes now count the primitive-type null group in getNumKeys()), the +/// numGroupsLimitReached flag is compared between the two modes on every query, null fixtures included. +public class OffHeapGroupByQueriesTest extends BaseQueriesTest { + private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "OffHeapGroupByQueriesTest"); + private static final String RAW_TABLE_NAME = "testTable"; + private static final String SEGMENT_NAME = "testSegment"; + private static final String NULL_SEGMENT_NAME = "testNullSegment"; + private static final long RANDOM_SEED = 42; + private static final int NUM_RECORDS = 50_000; + private static final int NUM_NULL_RECORDS = 1_200; + private static final String OFF_HEAP_PREFIX = "SET groupByOffHeap=true; "; + + // Dict-encoded SV columns + private static final String D_LOW_STR = "dLowStr"; // cardinality 10 + private static final String D_INT = "dInt"; // cardinality 1000 + private static final String D_HIGH_STR = "dHighStr"; // cardinality 20_000 + private static final String D_HIGH_INT = "dHighInt"; // cardinality 20_000 + private static final String D_HIGH_INT2 = "dHighInt2"; // cardinality 20_000 + private static final String D_HIGH_LONG = "dHighLong"; // cardinality 20_000 + // Raw (no-dictionary) SV columns + private static final String RAW_INT = "rawInt"; + private static final String RAW_LONG = "rawLong"; + private static final String RAW_FLOAT = "rawFloat"; + private static final String RAW_DOUBLE = "rawDouble"; + private static final String RAW_STRING = "rawString"; + private static final String RAW_BYTES = "rawBytes"; + private static final String RAW_BIG_DECIMAL = "rawBigDecimal"; + // Dict-encoded MV columns + private static final String MV_INT = "mvInt"; // cardinality 50 + private static final String MV_STR = "mvStr"; // cardinality 30 + private static final String MV_HIGH_INT = "mvHighInt"; // cardinality ~15_000 (forces the IntMap MV path) + // Metric + private static final String METRIC = "metric"; + + // Null-segment columns: "nf" = first row is null, "nm" = nulls only mid-stream (new values appear after the + // null stretch), "an" = all rows null + private static final String[] NULL_TYPE_SUFFIXES = {"Int", "Long", "Float", "Double", "Str", "Bytes", "BigDecimal"}; + private static final String AN_INT = "anInt"; + private static final String AN_STR = "anStr"; + private static final String NS_METRIC = "nsMetric"; + + // Star-tree segment columns: pre-aggregated DISTINCTCOUNTULL(stValInt) per stDim, plus a raw BYTES column + // carrying pre-serialized ULLs (the input shape the star-tree feeds through StarTreeGroupByExecutor) + private static final String ST_SEGMENT_NAME = "testStarTreeSegment"; + private static final int NUM_ST_RECORDS = 5_000; + private static final String ST_DIM = "stDim"; // cardinality 20 + private static final String ST_VAL_INT = "stValInt"; // cardinality 500 + private static final String ST_ULL_BYTES = "stUllBytes"; + + private IndexSegment _mainSegment; + private IndexSegment _nullSegment; + private IndexSegment _starTreeSegment; + private IndexSegment _indexSegment; + private List _indexSegments; + private long _directBufferBaseline; + + @Override + protected String getFilter() { + return ""; + } + + @Override + protected IndexSegment getIndexSegment() { + return _indexSegment; + } + + @Override + protected List getIndexSegments() { + return _indexSegments; + } + + @BeforeClass + public void setUp() + throws Exception { + FileUtils.deleteDirectory(INDEX_DIR); + buildMainSegment(); + buildNullSegment(); + buildStarTreeSegment(); + _mainSegment = ImmutableSegmentLoader.load(new File(INDEX_DIR, SEGMENT_NAME), ReadMode.mmap); + _nullSegment = ImmutableSegmentLoader.load(new File(INDEX_DIR, NULL_SEGMENT_NAME), ReadMode.mmap); + _starTreeSegment = ImmutableSegmentLoader.load(new File(INDEX_DIR, ST_SEGMENT_NAME), ReadMode.mmap); + useMainSegment(); + + // Warm-up off-heap query, then capture the direct-buffer baseline (see class doc for why the baseline is + // captured after the warm-up instead of right after segment load) + BrokerResponseNative warmUp = + getBrokerResponse(OFF_HEAP_PREFIX + "SELECT dLowStr, COUNT(*) FROM testTable GROUP BY dLowStr LIMIT 100"); + assertTrue(warmUp.getExceptions().isEmpty(), "Warm-up query failed: " + warmUp.getExceptions()); + _directBufferBaseline = PinotDataBuffer.getDirectBufferUsage(); + } + + @AfterClass + public void tearDown() + throws Exception { + _mainSegment.destroy(); + _nullSegment.destroy(); + _starTreeSegment.destroy(); + FileUtils.deleteDirectory(INDEX_DIR); + } + + private void useMainSegment() { + _indexSegment = _mainSegment; + _indexSegments = Arrays.asList(_mainSegment, _mainSegment); + } + + private void useNullSegment() { + _indexSegment = _nullSegment; + _indexSegments = Arrays.asList(_nullSegment, _nullSegment); + } + + private void useStarTreeSegment() { + _indexSegment = _starTreeSegment; + _indexSegments = Arrays.asList(_starTreeSegment, _starTreeSegment); + } + + private void buildMainSegment() + throws Exception { + Schema schema = new Schema.SchemaBuilder().setSchemaName(RAW_TABLE_NAME) + .addSingleValueDimension(D_LOW_STR, DataType.STRING) + .addSingleValueDimension(D_INT, DataType.INT) + .addSingleValueDimension(D_HIGH_STR, DataType.STRING) + .addSingleValueDimension(D_HIGH_INT, DataType.INT) + .addSingleValueDimension(D_HIGH_INT2, DataType.INT) + .addSingleValueDimension(D_HIGH_LONG, DataType.LONG) + .addSingleValueDimension(RAW_INT, DataType.INT) + .addSingleValueDimension(RAW_LONG, DataType.LONG) + .addSingleValueDimension(RAW_FLOAT, DataType.FLOAT) + .addSingleValueDimension(RAW_DOUBLE, DataType.DOUBLE) + .addSingleValueDimension(RAW_STRING, DataType.STRING) + .addSingleValueDimension(RAW_BYTES, DataType.BYTES) + .addSingleValueDimension(RAW_BIG_DECIMAL, DataType.BIG_DECIMAL) + .addMultiValueDimension(MV_INT, DataType.INT) + .addMultiValueDimension(MV_STR, DataType.STRING) + .addMultiValueDimension(MV_HIGH_INT, DataType.INT) + .addMetric(METRIC, DataType.DOUBLE) + .build(); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME) + .setNoDictionaryColumns(Arrays.asList(RAW_INT, RAW_LONG, RAW_FLOAT, RAW_DOUBLE, RAW_STRING, RAW_BYTES, + RAW_BIG_DECIMAL)) + .build(); + + Random random = new Random(RANDOM_SEED); + List records = new ArrayList<>(NUM_RECORDS); + for (int i = 0; i < NUM_RECORDS; i++) { + GenericRow record = new GenericRow(); + record.putValue(D_LOW_STR, "low_" + random.nextInt(10)); + record.putValue(D_INT, random.nextInt(1000)); + record.putValue(D_HIGH_STR, "high_" + (i % 20_000)); + record.putValue(D_HIGH_INT, (i * 7 + 13) % 20_000); + record.putValue(D_HIGH_INT2, (i * 11 + 5) % 20_000); + record.putValue(D_HIGH_LONG, (i * 13L + 3) % 20_000); + record.putValue(RAW_INT, random.nextInt(500) - 250); + record.putValue(RAW_LONG, (random.nextInt(500) - 250) * 1_000_003L); + // FLOAT/DOUBLE: dyadic values with identical values across rows, plus both -0.0 and +0.0 + int floatStep = random.nextInt(201) - 100; + record.putValue(RAW_FLOAT, floatStep == 0 && random.nextBoolean() ? -0.0f : floatStep * 0.25f); + int doubleStep = random.nextInt(201) - 100; + record.putValue(RAW_DOUBLE, doubleStep == 0 && random.nextBoolean() ? -0.0d : doubleStep * 0.5d); + // A small fraction of strings carry surrogate pairs to exercise the 4-byte UTF-8 encoding path + int stringId = random.nextInt(300); + record.putValue(RAW_STRING, stringId < 10 ? "raw_😀_" + stringId : "raw_" + stringId); + record.putValue(RAW_BYTES, ByteBuffer.allocate(8).putLong(random.nextInt(150) * 0x9E3779B97F4AL).array()); + record.putValue(RAW_BIG_DECIMAL, BigDecimal.valueOf((random.nextInt(400) - 200) * 25L, 2)); + record.putValue(MV_INT, randomMvValues(random, 50, null)); + record.putValue(MV_STR, randomMvValues(random, 30, "mv_")); + record.putValue(MV_HIGH_INT, randomMvValues(random, 15_000, null)); + record.putValue(METRIC, (random.nextInt(2001) - 1000) * 0.25d); + records.add(record); + } + buildSegment(tableConfig, schema, records, SEGMENT_NAME, false); + } + + private static Object[] randomMvValues(Random random, int cardinality, String stringPrefix) { + int numValues = 1 + random.nextInt(3); + Object[] values = new Object[numValues]; + for (int i = 0; i < numValues; i++) { + int value = random.nextInt(cardinality); + values[i] = stringPrefix != null ? stringPrefix + value : value; + } + return values; + } + + /// Builds a small segment with a star-tree on `stDim -> DISTINCTCOUNTULL(stValInt)` (stored pre-aggregated as + /// serialized ULL BYTES) plus a raw BYTES column carrying per-row serialized ULLs. + private void buildStarTreeSegment() + throws Exception { + Schema schema = new Schema.SchemaBuilder().setSchemaName(RAW_TABLE_NAME) + .addSingleValueDimension(ST_DIM, DataType.STRING) + .addSingleValueDimension(ST_VAL_INT, DataType.INT) + .addSingleValueDimension(ST_ULL_BYTES, DataType.BYTES) + .build(); + StarTreeIndexConfig starTreeIndexConfig = + new StarTreeIndexConfig(List.of(ST_DIM), null, List.of("distinctCountULL__" + ST_VAL_INT), null, 1); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME) + .setNoDictionaryColumns(List.of(ST_ULL_BYTES)) + .setStarTreeIndexConfigs(List.of(starTreeIndexConfig)) + .build(); + + Random random = new Random(RANDOM_SEED); + List records = new ArrayList<>(NUM_ST_RECORDS); + for (int i = 0; i < NUM_ST_RECORDS; i++) { + GenericRow record = new GenericRow(); + record.putValue(ST_DIM, "st_" + random.nextInt(20)); + record.putValue(ST_VAL_INT, random.nextInt(500)); + UltraLogLog ull = UltraLogLog.create(CommonConstants.Helix.DEFAULT_ULTRALOGLOG_P); + int numValues = 1 + random.nextInt(3); + for (int j = 0; j < numValues; j++) { + UltraLogLogUtils.hashObject((long) random.nextInt(2000)).ifPresent(ull::add); + } + record.putValue(ST_ULL_BYTES, ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.serialize(ull)); + records.add(record); + } + buildSegment(tableConfig, schema, records, ST_SEGMENT_NAME, false); + } + + private void buildNullSegment() + throws Exception { + Schema.SchemaBuilder schemaBuilder = new Schema.SchemaBuilder().setSchemaName(RAW_TABLE_NAME); + List noDictionaryColumns = new ArrayList<>(); + DataType[] dataTypes = { + DataType.INT, DataType.LONG, DataType.FLOAT, DataType.DOUBLE, DataType.STRING, DataType.BYTES, + DataType.BIG_DECIMAL + }; + for (int i = 0; i < NULL_TYPE_SUFFIXES.length; i++) { + String nullFirstColumn = "nf" + NULL_TYPE_SUFFIXES[i]; + String nullMidColumn = "nm" + NULL_TYPE_SUFFIXES[i]; + schemaBuilder.addSingleValueDimension(nullFirstColumn, dataTypes[i]); + schemaBuilder.addSingleValueDimension(nullMidColumn, dataTypes[i]); + noDictionaryColumns.add(nullFirstColumn); + noDictionaryColumns.add(nullMidColumn); + } + schemaBuilder.addSingleValueDimension(AN_INT, DataType.INT); + schemaBuilder.addSingleValueDimension(AN_STR, DataType.STRING); + noDictionaryColumns.add(AN_INT); + noDictionaryColumns.add(AN_STR); + schemaBuilder.addMetric(NS_METRIC, DataType.DOUBLE); + Schema schema = schemaBuilder.build(); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME) + .setNoDictionaryColumns(noDictionaryColumns).build(); + + List records = new ArrayList<>(NUM_NULL_RECORDS); + for (int i = 0; i < NUM_NULL_RECORDS; i++) { + GenericRow record = new GenericRow(); + // "nf" columns: nulls sprinkled from row 0 on; values from an 9-value pool + boolean nullFirstIsNull = i % 5 == 0; + int nullFirstPoolIndex = i % 9; + // "nm" columns: values only (8-value pool prefix) until row 600, nulls for rows [600, 700), then the full + // 11-value pool so that pool indexes 7-10 first appear after the null stretch (700 % 11 == 7) + boolean nullMidIsNull = i >= 600 && i < 700; + int nullMidPoolIndex = i < 600 ? i % 7 : i % 11; + for (int t = 0; t < NULL_TYPE_SUFFIXES.length; t++) { + record.putValue("nf" + NULL_TYPE_SUFFIXES[t], nullFirstIsNull ? null : nullPoolValue(t, nullFirstPoolIndex)); + record.putValue("nm" + NULL_TYPE_SUFFIXES[t], nullMidIsNull ? null : nullPoolValue(t, nullMidPoolIndex)); + } + record.putValue(AN_INT, null); + record.putValue(AN_STR, null); + record.putValue(NS_METRIC, (i % 41 - 20) * 0.25d); + records.add(record); + } + buildSegment(tableConfig, schema, records, NULL_SEGMENT_NAME, true); + } + + private static Object nullPoolValue(int typeIndex, int poolIndex) { + switch (typeIndex) { + case 0: + return poolIndex * 3 - 15; + case 1: + return poolIndex * 997L; + case 2: + return (poolIndex - 5) * 0.25f; + case 3: + return (poolIndex - 5) * 0.5d; + case 4: + return "ns_" + poolIndex; + case 5: + return new byte[]{(byte) poolIndex, (byte) (poolIndex + 1), (byte) (poolIndex * 2), 7}; + case 6: + return BigDecimal.valueOf((poolIndex - 5) * 25L, 2); + default: + throw new IllegalStateException(); + } + } + + private void buildSegment(TableConfig tableConfig, Schema schema, List records, String segmentName, + boolean nullHandlingEnabled) + throws Exception { + SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(tableConfig, schema); + segmentGeneratorConfig.setTableName(RAW_TABLE_NAME); + segmentGeneratorConfig.setSegmentName(segmentName); + segmentGeneratorConfig.setDefaultNullHandlingEnabled(nullHandlingEnabled); + segmentGeneratorConfig.setOutDir(INDEX_DIR.getPath()); + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(segmentGeneratorConfig, new GenericRowRecordReader(records)); + driver.build(); + } + + // --------------------------------------------------------------------------------------------- + // Comparison helpers + // --------------------------------------------------------------------------------------------- + + private BrokerResponseNative runQuery(String query) { + BrokerResponseNative response = getBrokerResponse(query); + assertTrue(response.getExceptions().isEmpty(), + "Query failed: " + query + " with exceptions: " + response.getExceptions()); + return response; + } + + /// Runs the query on-heap and off-heap, asserts direct memory returns to the baseline after the off-heap run, + /// and returns both responses for row comparison. + private BrokerResponseNative[] runBothModes(String query) { + BrokerResponseNative onHeap = runQuery(query); + BrokerResponseNative offHeap = runQuery(OFF_HEAP_PREFIX + query); + assertEquals(PinotDataBuffer.getDirectBufferUsage(), _directBufferBaseline, + "Off-heap direct memory leaked by query: " + query); + return new BrokerResponseNative[]{onHeap, offHeap}; + } + + /// Unordered comparison: rows compared as multisets (canonical string form, sorted). The numGroupsLimitReached + /// flag must match between the two modes for every query (both modes count groups identically, including the + /// primitive-type null group since the on-heap counting fix). + private void testQuery(String query) { + BrokerResponseNative[] responses = runBothModes(query); + assertEquals(responses[1].isNumGroupsLimitReached(), responses[0].isNumGroupsLimitReached(), + "numGroupsLimitReached mismatch between modes for query: " + query); + List onHeapRows = canonicalRows(responses[0].getResultTable(), query); + List offHeapRows = canonicalRows(responses[1].getResultTable(), query); + Collections.sort(onHeapRows); + Collections.sort(offHeapRows); + assertEquals(offHeapRows, onHeapRows, "Result mismatch (unordered) for query: " + query); + } + + /// Ordered comparison for ORDER BY queries: exact ordered row lists. + private void testOrderedQuery(String query) { + BrokerResponseNative[] responses = runBothModes(query); + List onHeapRows = canonicalRows(responses[0].getResultTable(), query); + List offHeapRows = canonicalRows(responses[1].getResultTable(), query); + assertEquals(offHeapRows, onHeapRows, "Result mismatch (ordered) for query: " + query); + } + + /// Unordered comparison plus an assertion of the expected numGroupsLimitReached value in both modes. + private void testCappedQuery(String query, boolean expectLimitReached) { + BrokerResponseNative[] responses = runBothModes(query); + assertEquals(responses[0].isNumGroupsLimitReached(), expectLimitReached, + "Unexpected on-heap numGroupsLimitReached for query: " + query); + assertEquals(responses[1].isNumGroupsLimitReached(), expectLimitReached, + "Unexpected off-heap numGroupsLimitReached for query: " + query); + List onHeapRows = canonicalRows(responses[0].getResultTable(), query); + List offHeapRows = canonicalRows(responses[1].getResultTable(), query); + Collections.sort(onHeapRows); + Collections.sort(offHeapRows); + assertEquals(offHeapRows, onHeapRows, "Result mismatch (capped) for query: " + query); + } + + private static List canonicalRows(ResultTable resultTable, String query) { + assertTrue(resultTable != null, "Missing result table for query: " + query); + List rows = resultTable.getRows(); + // Every query in the battery matches at least one group; an empty result would make the differential + // comparison pass vacuously + assertFalse(rows.isEmpty(), "Empty result rows for query: " + query); + List canonicalRows = new ArrayList<>(rows.size()); + StringBuilder builder = new StringBuilder(); + for (Object[] row : rows) { + builder.setLength(0); + for (Object cell : row) { + builder.append(canonicalCell(cell)).append('|'); + } + canonicalRows.add(builder.toString()); + } + return canonicalRows; + } + + private static String canonicalCell(Object cell) { + if (cell == null) { + return "null"; + } + if (cell instanceof byte[]) { + return "bytes:" + BytesUtils.toHexString((byte[]) cell); + } + if (cell instanceof Object[]) { + return "array:" + Arrays.deepToString((Object[]) cell); + } + if (cell instanceof int[]) { + return "ints:" + Arrays.toString((int[]) cell); + } + if (cell instanceof long[]) { + return "longs:" + Arrays.toString((long[]) cell); + } + if (cell instanceof double[]) { + return "doubles:" + Arrays.toString((double[]) cell); + } + // Include the type so that a value/type flip between modes cannot cancel out in the string form + return cell.getClass().getSimpleName() + ':' + cell; + } + + private static String aggregations(String metricColumn) { + return "COUNT(*), SUM(" + metricColumn + "), MIN(" + metricColumn + "), MAX(" + metricColumn + "), AVG(" + + metricColumn + "), DISTINCTCOUNT(dInt)"; + } + + // --------------------------------------------------------------------------------------------- + // Query battery + // --------------------------------------------------------------------------------------------- + + @Test + public void testSingleColumnGroupBy() { + String[] columns = { + D_LOW_STR, D_INT, D_HIGH_STR, D_HIGH_INT, D_HIGH_INT2, D_HIGH_LONG, RAW_INT, RAW_LONG, RAW_FLOAT, RAW_DOUBLE, + RAW_STRING, RAW_BYTES, RAW_BIG_DECIMAL + }; + for (String column : columns) { + testQuery("SELECT " + column + ", " + aggregations(METRIC) + " FROM testTable GROUP BY " + column + + " LIMIT 100000"); + } + } + + /// DISTINCTCOUNTULL group-by state moves off-heap (OffHeapUltraLogLogGroupByResultHolder) when the input is a + /// raw (no-dictionary) column; dictionary-encoded inputs keep the on-heap dict-id bitmap through the holder's + /// delegate. Both modes must produce byte-identical ULL states, so even the serialized RAWULL output compares + /// exactly. + @Test + public void testDistinctCountULL() { + // Raw input values -> off-heap register slots, across group-key generator variants + testQuery("SELECT dLowStr, DISTINCTCOUNTULL(rawInt) FROM testTable GROUP BY dLowStr LIMIT 100000"); + testQuery("SELECT dHighStr, DISTINCTCOUNTULL(rawString) FROM testTable GROUP BY dHighStr LIMIT 100000"); + testQuery("SELECT rawInt, DISTINCTCOUNTULL(rawDouble) FROM testTable GROUP BY rawInt LIMIT 100000"); + // NOTE: no rawBytes input — a single-value BYTES input is always interpreted as pre-serialized ULLs + testQuery("SELECT rawString, DISTINCTCOUNTULL(rawLong), DISTINCTCOUNTULL(rawFloat) FROM testTable" + + " GROUP BY rawString LIMIT 100000"); + // Explicit precision argument (non-default slot sizes) + testQuery("SELECT dLowStr, DISTINCTCOUNTULL(rawInt, 8), DISTINCTCOUNTULL(rawInt, 14) FROM testTable" + + " GROUP BY dLowStr LIMIT 100000"); + // Dictionary-encoded input -> heap dict-id bitmap through the holder's delegate + testQuery("SELECT dLowStr, DISTINCTCOUNTULL(dInt), DISTINCTCOUNTULL(dHighStr) FROM testTable GROUP BY dLowStr" + + " LIMIT 100000"); + // MV input values (dict-encoded -> delegate; the raw MV value path is covered by the unit tests) + testQuery("SELECT dLowStr, DISTINCTCOUNTULL(mvInt) FROM testTable GROUP BY dLowStr LIMIT 100000"); + // MV group keys: SV raw input over int[] group keys (setValueForGroupKeys), and MV input over MV keys + testQuery("SELECT mvInt, DISTINCTCOUNTULL(rawInt) FROM testTable GROUP BY mvInt LIMIT 100000"); + testQuery("SELECT mvInt, DISTINCTCOUNTULL(mvStr) FROM testTable GROUP BY mvInt LIMIT 100000"); + // Serialized (RAWULL) final results compare exactly because the states are byte-identical + testQuery("SELECT dLowStr, DISTINCTCOUNTRAWULL(rawString) FROM testTable GROUP BY dLowStr LIMIT 100000"); + // Order-by + trim path extracts every group's ULL through TableResizer + testOrderedQuery("SELECT dHighStr, DISTINCTCOUNTULL(rawInt) FROM testTable GROUP BY dHighStr" + + " ORDER BY DISTINCTCOUNTULL(rawInt) DESC, dHighStr LIMIT 10"); + // Filtered aggregation shares the generator between the filtered and unfiltered holders + testQuery("SELECT dLowStr, DISTINCTCOUNTULL(rawInt) FILTER (WHERE dInt < 100), COUNT(*) FROM testTable" + + " GROUP BY dLowStr LIMIT 100000"); + // Null handling: forEachNotNull must skip null stretches identically in both modes + useNullSegment(); + try { + testQuery("SET enableNullHandling=true; SELECT nfStr, DISTINCTCOUNTULL(nmInt) FROM testTable GROUP BY nfStr" + + " LIMIT 1000"); + testQuery("SET enableNullHandling=true; SELECT nmInt, DISTINCTCOUNTULL(nfLong), DISTINCTCOUNTULL(anStr)" + + " FROM testTable GROUP BY nmInt LIMIT 1000"); + } finally { + useMainSegment(); + } + } + + /// Pre-serialized-ULL BYTES input: the one group-by mode where the function must read back a LIVE mutable + /// object per row (get-merge or adopt via setValueForKey), which the off-heap holder serves through its heap + /// delegate. Covered twice: through a raw BYTES column of serialized ULLs, and through a star-tree whose + /// pre-aggregated `distinctCountULL__stValInt` column feeds StarTreeGroupByExecutor with the same shape. + @Test + public void testDistinctCountULLSerializedBytesAndStarTree() { + useStarTreeSegment(); + try { + // Plain BYTES branch through the ordinary executor + testQuery("SELECT stDim, DISTINCTCOUNTULL(stUllBytes) FROM testTable GROUP BY stDim LIMIT 1000"); + // Star-tree path: pre-aggregated ULL BYTES through StarTreeGroupByExecutor in both modes + String starTreeQuery = "SELECT stDim, DISTINCTCOUNTULL(stValInt) FROM testTable GROUP BY stDim LIMIT 1000"; + testQuery(starTreeQuery); + // Prove the star-tree actually served the query (fully pre-aggregated at maxLeafRecords=1, so far fewer + // docs than rows are scanned) — otherwise the star-tree leg of this test is vacuous + for (BrokerResponseNative response : runBothModes(starTreeQuery)) { + assertTrue(response.getNumDocsScanned() < NUM_ST_RECORDS, + "Expected the star-tree to serve the query, but numDocsScanned=" + response.getNumDocsScanned()); + } + } finally { + useMainSegment(); + } + } + + @Test + public void testDictionaryMultiColumnVariants() { + // ARRAY_BASED: 10 * 1000 = 10_000, not above the array-based threshold + testQuery("SELECT dLowStr, dInt, COUNT(*), SUM(metric) FROM testTable GROUP BY dLowStr, dInt LIMIT 100000"); + // INT_MAP_BASED: 20_000 * 10 = 200_000 + testQuery("SELECT dHighStr, dLowStr, COUNT(*), SUM(metric) FROM testTable GROUP BY dHighStr, dLowStr" + + " LIMIT 100000"); + // INT_MAP_BASED with a large product: 20_000 * 20_000 = 4 * 10^8 < Integer.MAX_VALUE + testQuery("SELECT dHighStr, dHighInt, COUNT(*), SUM(metric) FROM testTable GROUP BY dHighStr, dHighInt" + + " LIMIT 100000"); + // LONG_MAP_BASED: 20_000 * 20_000 * 10 = 4 * 10^9 > Integer.MAX_VALUE + testQuery("SELECT dHighStr, dHighInt, dLowStr, COUNT(*), SUM(metric) FROM testTable" + + " GROUP BY dHighStr, dHighInt, dLowStr LIMIT 100000"); + // LONG_MAP_BASED: 20_000^3 * 1000 = 8 * 10^15 < Long.MAX_VALUE + testQuery("SELECT dHighInt, dHighInt2, dHighLong, dInt, COUNT(*), SUM(metric) FROM testTable" + + " GROUP BY dHighInt, dHighInt2, dHighLong, dInt LIMIT 100000"); + // ARRAY_MAP_BASED: 20_000^4 * 1000 = 1.6 * 10^20 > Long.MAX_VALUE (long overflow) + testQuery("SELECT dHighStr, dHighInt, dHighInt2, dHighLong, dInt, COUNT(*), SUM(metric) FROM testTable" + + " GROUP BY dHighStr, dHighInt, dHighInt2, dHighLong, dInt LIMIT 100000"); + } + + @Test + public void testRawAndMixedMultiColumn() { + testQuery("SELECT rawInt, rawString, COUNT(*), SUM(metric) FROM testTable GROUP BY rawInt, rawString" + + " LIMIT 100000"); + testQuery("SELECT rawFloat, rawDouble, COUNT(*), SUM(metric) FROM testTable GROUP BY rawFloat, rawDouble" + + " LIMIT 100000"); + testQuery("SELECT rawLong, rawBytes, rawBigDecimal, COUNT(*), SUM(metric) FROM testTable" + + " GROUP BY rawLong, rawBytes, rawBigDecimal LIMIT 100000"); + // Mixed raw + dict + testQuery("SELECT rawString, dLowStr, COUNT(*), SUM(metric) FROM testTable GROUP BY rawString, dLowStr" + + " LIMIT 100000"); + testQuery("SELECT dHighStr, rawInt, COUNT(*), SUM(metric) FROM testTable GROUP BY dHighStr, rawInt" + + " LIMIT 200000"); + } + + @Test + public void testMultiValueGroupBy() { + testQuery("SELECT mvInt, " + aggregations(METRIC) + " FROM testTable GROUP BY mvInt LIMIT 100000"); + testQuery("SELECT mvStr, " + aggregations(METRIC) + " FROM testTable GROUP BY mvStr LIMIT 100000"); + // Single high-cardinality MV column: 15_000 > 10_000 forces the IntMap MV path + testQuery("SELECT mvHighInt, COUNT(*), SUM(metric) FROM testTable GROUP BY mvHighInt LIMIT 100000"); + // MV + MV (ARRAY_BASED: 50 * 30 = 1500) + testQuery("SELECT mvInt, mvStr, COUNT(*), SUM(metric) FROM testTable GROUP BY mvInt, mvStr LIMIT 100000"); + // MV + SV dict combos across holder variants + testQuery("SELECT mvInt, dLowStr, COUNT(*), SUM(metric) FROM testTable GROUP BY mvInt, dLowStr LIMIT 100000"); + // INT_MAP_BASED MV: 50 * 20_000 = 10^6 + testQuery("SELECT mvInt, dHighStr, COUNT(*), SUM(metric) FROM testTable GROUP BY mvInt, dHighStr LIMIT 200000"); + // LONG_MAP_BASED MV: 50 * 20_000 * 20_000 = 2 * 10^10 + testQuery("SELECT mvInt, dHighStr, dHighInt, COUNT(*), SUM(metric) FROM testTable" + + " GROUP BY mvInt, dHighStr, dHighInt LIMIT 200000"); + // ARRAY_MAP_BASED MV: 50 * 20_000^4 * 1000 = 8 * 10^21 (long overflow) + testQuery("SELECT mvInt, dHighStr, dHighInt, dHighInt2, dHighLong, dInt, COUNT(*), SUM(metric) FROM testTable" + + " GROUP BY mvInt, dHighStr, dHighInt, dHighInt2, dHighLong, dInt LIMIT 200000"); + // MV + raw SV -> NoDictionaryMultiColumn MV path + testQuery("SELECT mvStr, rawInt, COUNT(*), SUM(metric) FROM testTable GROUP BY mvStr, rawInt LIMIT 200000"); + } + + @Test + public void testOrderByTrimPath() { + testOrderedQuery("SELECT dHighStr, COUNT(*) FROM testTable GROUP BY dHighStr" + + " ORDER BY COUNT(*) DESC, dHighStr LIMIT 10"); + testOrderedQuery("SELECT rawString, SUM(metric) FROM testTable GROUP BY rawString" + + " ORDER BY SUM(metric) DESC, rawString LIMIT 10"); + testOrderedQuery("SELECT dLowStr, dInt, MAX(metric) FROM testTable GROUP BY dLowStr, dInt" + + " ORDER BY MAX(metric) DESC, dLowStr, dInt LIMIT 10"); + testOrderedQuery("SELECT rawFloat, COUNT(*) FROM testTable GROUP BY rawFloat ORDER BY rawFloat LIMIT 20"); + testOrderedQuery("SELECT mvInt, SUM(metric) FROM testTable GROUP BY mvInt ORDER BY SUM(metric) DESC, mvInt" + + " LIMIT 10"); + testOrderedQuery("SELECT dHighInt, rawBigDecimal, MIN(metric) FROM testTable GROUP BY dHighInt, rawBigDecimal" + + " ORDER BY MIN(metric), dHighInt, rawBigDecimal LIMIT 15"); + } + + @Test + public void testFilteredGroupBy() { + String[] queries = { + "SELECT dInt, COUNT(*), SUM(metric) FROM testTable WHERE dInt IN (1, 2, 3, 4, 5) GROUP BY dInt LIMIT 1000", + "SELECT rawInt, COUNT(*), SUM(metric) FROM testTable WHERE rawInt IN (0, 1, 2, 3) GROUP BY rawInt" + + " LIMIT 1000", + "SELECT dLowStr, dInt, COUNT(*) FROM testTable WHERE dLowStr IN ('low_1', 'low_2') AND dInt = 7" + + " GROUP BY dLowStr, dInt LIMIT 1000", + "SELECT dHighStr, COUNT(*) FROM testTable WHERE dHighStr IN ('high_1', 'high_2', 'high_3')" + + " GROUP BY dHighStr LIMIT 1000" + }; + for (String query : queries) { + testQuery(query); + } + } + + @Test + public void testFilteredGroupByWithOptimizedResultHolderCapacity() { + // Predicate-based upper-bound optimization variants. NOTE: a pre-existing (upstream, on-heap) bug limits what + // can be exercised here for dictionary columns: with optimizeMaxInitialResultHolderCapacity the predicate + // sizes shrink cardinalityProduct/_globalGroupIdUpperBound to the IN-list size, which selects the + // ArrayBasedHolder — whose group ids are raw dictionary-id products. Any matching dictionary id >= the + // optimized bound then throws ArrayIndexOutOfBoundsException in BOTH modes (the T0 ArrayBased path is on-heap + // by design even when groupByOffHeap is set). e.g. `WHERE dInt IN (1, 2, 3, 4, 5)` fails with + // "Index 5 out of bounds for length 5" on this fixture with or without the off-heap option. The dict variant + // below therefore uses IN (0..4), whose dictionary ids (0..4 — dInt values are 0..999, sorted) fit inside the + // optimized bound; raw columns are unaffected (the no-dict generators only cap their key maps). + testQuery("SET optimizeMaxInitialResultHolderCapacity=true; SELECT dInt, COUNT(*), SUM(metric) FROM testTable" + + " WHERE dInt IN (0, 1, 2, 3, 4) GROUP BY dInt LIMIT 1000"); + testQuery("SET optimizeMaxInitialResultHolderCapacity=true; SELECT rawInt, COUNT(*), SUM(metric)" + + " FROM testTable WHERE rawInt IN (0, 1, 2, 3) GROUP BY rawInt LIMIT 1000"); + testQuery("SET optimizeMaxInitialResultHolderCapacity=true; SELECT rawString, COUNT(*) FROM testTable" + + " WHERE rawString IN ('raw_11', 'raw_12') GROUP BY rawString LIMIT 1000"); + testQuery("SET optimizeMaxInitialResultHolderCapacity=true; SELECT rawInt, rawString, COUNT(*) FROM testTable" + + " WHERE rawInt IN (0, 1, 2, 3) AND rawString IN ('raw_11', 'raw_12') GROUP BY rawInt, rawString" + + " LIMIT 1000"); + } + + @Test + public void testFilteredAggregations() { + // FILTER clauses share one group key generator across executors; with off-heap enabled all executors register + // their off-heap holders on the same resource-tracking wrapper, closed once + testQuery("SELECT dLowStr, COUNT(*) FILTER (WHERE dInt > 500), SUM(metric) FILTER (WHERE rawInt > 0)," + + " COUNT(*) FROM testTable GROUP BY dLowStr LIMIT 1000"); + testQuery("SELECT rawString, COUNT(*) FILTER (WHERE rawDouble > 0), MIN(metric) FILTER (WHERE dInt < 100)" + + " FROM testTable GROUP BY rawString LIMIT 100000"); + testQuery("SELECT dHighStr, SUM(metric) FILTER (WHERE dLowStr = 'low_3'), COUNT(*) FROM testTable" + + " GROUP BY dHighStr LIMIT 100000"); + } + + @Test + public void testNumGroupsLimitCap() { + // Cap semantics parity: group ids assign in row order, so the kept groups (and rows) must be identical, and + // the numGroupsLimitReached flag must match (no null groups on this fixture, so no counting divergence) + testCappedQuery("SET numGroupsLimit=100; SELECT dHighStr, COUNT(*), SUM(metric) FROM testTable" + + " GROUP BY dHighStr LIMIT 100000", true); + testCappedQuery("SET numGroupsLimit=100; SELECT rawString, COUNT(*), SUM(metric) FROM testTable" + + " GROUP BY rawString LIMIT 100000", true); + testCappedQuery("SET numGroupsLimit=100; SELECT dHighStr, dHighInt, COUNT(*) FROM testTable" + + " GROUP BY dHighStr, dHighInt LIMIT 100000", true); + testCappedQuery("SET numGroupsLimit=100; SELECT rawInt, rawString, COUNT(*) FROM testTable" + + " GROUP BY rawInt, rawString LIMIT 100000", true); + testCappedQuery("SET numGroupsLimit=100; SELECT mvHighInt, COUNT(*) FROM testTable GROUP BY mvHighInt" + + " LIMIT 100000", true); + // Limit not reached: flag must be false in both modes + testCappedQuery("SET numGroupsLimit=100; SELECT dLowStr, COUNT(*) FROM testTable GROUP BY dLowStr LIMIT 1000", + false); + } + + @Test + public void testNullHandlingEnabledWithoutNulls() { + // enableNullHandling routes all group-bys through the NoDictionary generators even for dict columns; this + // fixture has no null rows, so results must match the null-disabled semantics-compatible comparison + String[] columns = {D_INT, D_LOW_STR, D_HIGH_STR, RAW_INT, RAW_STRING, RAW_FLOAT, RAW_BYTES, RAW_BIG_DECIMAL}; + for (String column : columns) { + testQuery("SET enableNullHandling=true; SELECT " + column + ", COUNT(*), SUM(metric) FROM testTable GROUP BY " + + column + " LIMIT 100000"); + } + testQuery("SET enableNullHandling=true; SELECT rawInt, rawString, COUNT(*), SUM(metric) FROM testTable" + + " GROUP BY rawInt, rawString LIMIT 100000"); + testQuery("SET enableNullHandling=true; SELECT dLowStr, rawDouble, COUNT(*), SUM(metric) FROM testTable" + + " GROUP BY dLowStr, rawDouble LIMIT 100000"); + // MV group-by with null handling -> NoDictionarySingleColumn MV path + testQuery("SET enableNullHandling=true; SELECT mvInt, COUNT(*), SUM(metric) FROM testTable GROUP BY mvInt" + + " LIMIT 100000"); + testQuery("SET enableNullHandling=true; SELECT mvInt, rawString, COUNT(*) FROM testTable" + + " GROUP BY mvInt, rawString LIMIT 200000"); + } + + /// Guards the whole differential battery against passing vacuously: proves the `groupByOffHeap` query option + /// actually reaches the QueryContext through the plan maker, and that DefaultGroupByExecutor reacts to the + /// flag by wrapping the generator in the off-heap resource tracker (allocating direct memory) — so a broken + /// option plumbing cannot silently turn every off-heap run in this test into an on-heap run. + @Test + public void testGroupByOffHeapOptionPlumbing() { + // Query option -> QueryContext (InstancePlanMakerImplV2#applyQueryOptions) + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery( + OFF_HEAP_PREFIX + "SELECT dHighStr, COUNT(*) FROM testTable GROUP BY dHighStr"); + QueryContext queryContext = QueryContextConverterUtils.getQueryContext(pinotQuery); + queryContext.setEndTimeMs(System.currentTimeMillis() + Server.DEFAULT_QUERY_EXECUTOR_TIMEOUT_MS); + try (QueryThreadContext ignore = QueryThreadContext.openForSseTest()) { + PLAN_MAKER.makeInstancePlan(List.of(new SegmentContext(_mainSegment)), queryContext, EXECUTOR_SERVICE); + } + assertTrue(queryContext.isGroupByOffHeap(), "groupByOffHeap query option did not reach the QueryContext"); + + // QueryContext flag -> off-heap executor wiring + ExpressionContext[] groupByExpressions = {ExpressionContext.forIdentifier(D_HIGH_STR)}; + QueryContext offHeapContext = + QueryContextConverterUtils.getQueryContext("SELECT COUNT(*) FROM testTable GROUP BY dHighStr"); + offHeapContext.setGroupByOffHeap(true); + BaseProjectOperator projectOperator = new ProjectPlanNode(new SegmentContext(_mainSegment), offHeapContext, + Arrays.asList(groupByExpressions), DocIdSetPlanNode.MAX_DOC_PER_CALL).run(); + DefaultGroupByExecutor executor = new DefaultGroupByExecutor(offHeapContext, groupByExpressions, projectOperator); + try { + assertTrue(executor.getGroupKeyGenerator() instanceof ResourceTrackingGroupKeyGenerator, + "Off-heap group-by must wrap the generator in ResourceTrackingGroupKeyGenerator"); + assertTrue(PinotDataBuffer.getDirectBufferUsage() > _directBufferBaseline, + "Off-heap group-by executor did not allocate direct memory"); + } finally { + executor.getGroupKeyGenerator().close(); + } + assertEquals(PinotDataBuffer.getDirectBufferUsage(), _directBufferBaseline, + "Off-heap executor leaked direct memory after close"); + + // Negative: with the flag off, the generator must not be wrapped + QueryContext onHeapContext = + QueryContextConverterUtils.getQueryContext("SELECT COUNT(*) FROM testTable GROUP BY dHighStr"); + BaseProjectOperator onHeapProjectOperator = new ProjectPlanNode(new SegmentContext(_mainSegment), + onHeapContext, Arrays.asList(groupByExpressions), DocIdSetPlanNode.MAX_DOC_PER_CALL).run(); + DefaultGroupByExecutor onHeapExecutor = + new DefaultGroupByExecutor(onHeapContext, groupByExpressions, onHeapProjectOperator); + try { + assertFalse(onHeapExecutor.getGroupKeyGenerator() instanceof ResourceTrackingGroupKeyGenerator, + "On-heap group-by must not wrap the generator in ResourceTrackingGroupKeyGenerator"); + } finally { + onHeapExecutor.getGroupKeyGenerator().close(); + } + } + + @Test + public void testNullHandlingWithNulls() { + useNullSegment(); + try { + String aggs = "COUNT(*), SUM(nsMetric), MIN(nsMetric), MAX(nsMetric)"; + // Single-column group-by on every type with the FIRST row null and with nulls only mid-stream + for (String suffix : NULL_TYPE_SUFFIXES) { + testQuery("SET enableNullHandling=true; SELECT nf" + suffix + ", " + aggs + " FROM testTable GROUP BY nf" + + suffix + " LIMIT 1000"); + testQuery("SET enableNullHandling=true; SELECT nm" + suffix + ", " + aggs + " FROM testTable GROUP BY nm" + + suffix + " LIMIT 1000"); + } + // All-null columns + testQuery("SET enableNullHandling=true; SELECT anInt, " + aggs + " FROM testTable GROUP BY anInt LIMIT 1000"); + testQuery("SET enableNullHandling=true; SELECT anStr, " + aggs + " FROM testTable GROUP BY anStr LIMIT 1000"); + // Multi-column group-bys with nulls + testQuery("SET enableNullHandling=true; SELECT nfInt, nmStr, " + aggs + " FROM testTable GROUP BY nfInt, nmStr" + + " LIMIT 1000"); + testQuery("SET enableNullHandling=true; SELECT nmInt, nmLong, " + aggs + " FROM testTable" + + " GROUP BY nmInt, nmLong LIMIT 1000"); + testQuery("SET enableNullHandling=true; SELECT anInt, nfStr, " + aggs + " FROM testTable GROUP BY anInt, nfStr" + + " LIMIT 1000"); + // Ordered null query (null key ordering is deterministic in both modes) + testOrderedQuery("SET enableNullHandling=true; SELECT nmInt, COUNT(*) FROM testTable GROUP BY nmInt" + + " ORDER BY COUNT(*) DESC, nmInt LIMIT 5"); + // Null + tiny numGroupsLimit: rows and the cross-mode numGroupsLimitReached flag are compared (both modes + // count the null group since the on-heap counting fix) + testQuery("SET enableNullHandling=true; SET numGroupsLimit=5; SELECT nmInt, COUNT(*) FROM testTable" + + " GROUP BY nmInt LIMIT 1000"); + testQuery("SET enableNullHandling=true; SET numGroupsLimit=5; SELECT nfStr, COUNT(*) FROM testTable" + + " GROUP BY nfStr LIMIT 1000"); + testQuery("SET enableNullHandling=true; SET numGroupsLimit=3; SELECT nfInt, nmStr, COUNT(*) FROM testTable" + + " GROUP BY nfInt, nmStr LIMIT 1000"); + } finally { + useMainSegment(); + } + } +} diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupByHugeSSE.java b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupByHugeSSE.java new file mode 100644 index 000000000000..6aa3c7253eba --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupByHugeSSE.java @@ -0,0 +1,213 @@ +/** + * 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.perf; + +import java.io.File; +import java.util.HashMap; +import java.util.Map; +import java.util.Random; +import java.util.concurrent.TimeUnit; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock; +import org.apache.pinot.core.operator.query.GroupByOperator; +import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2; +import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapGroupByBufferPool; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.SegmentContext; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +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.data.readers.GenericRow; +import org.apache.pinot.spi.data.readers.RecordReader; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.apache.pinot.sql.parsers.CalciteSqlParser; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.runner.Runner; +import org.openjdk.jmh.runner.options.OptionsBuilder; + + +/// Huge-group SSE benchmark: ~10M distinct groups in a single 12M-row segment (numGroupsLimit raised to 21M), +/// measuring the per-segment group-by phase only — the phase the `groupByOffHeap` feature changes. The +/// cross-segment combine is deliberately excluded: merging 10M groups into the (mode-independent, still on-heap) +/// IndexedTable dominates and GC-thrashes both arms identically; it is the Milestone-4 work item. +/// +/// Run with a large fixed heap and explicit direct-memory ceiling, e.g. +/// `-jvmArgs '-Xms12g -Xmx12g -XX:MaxDirectMemorySize=8g'`, and `-prof gc`: score (ms/op), +/// gc.alloc.rate.norm and gc.time are the interesting metrics. Pair with [OffHeapGroupByMemoryFootprint] +/// (which goes to 100M groups) for retained-memory numbers. +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@Fork(1) +@Warmup(iterations = 2, time = 10) +@Measurement(iterations = 4, time = 10) +@State(Scope.Benchmark) +public class BenchmarkOffHeapGroupByHugeSSE { + private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "BenchmarkOffHeapGroupByHugeSSE"); + private static final String TABLE_NAME = "MyTable"; + private static final int NUM_ROWS = 12_000_000; + private static final int CARDINALITY = 10_000_000; + private static final int RAISED_NUM_GROUPS_LIMIT = 21_000_000; + + private static final String DICT_INT_HUGE = "DICT_INT_HUGE"; + private static final String RAW_STRING_HUGE = "RAW_STRING_HUGE"; + private static final String METRIC = "METRIC"; + + private static final Map QUERIES = Map.of( + "DICT_INT", "SELECT DICT_INT_HUGE, COUNT(*), SUM(METRIC) FROM MyTable GROUP BY DICT_INT_HUGE LIMIT 10", + "RAW_STRING", + "SELECT RAW_STRING_HUGE, COUNT(*), SUM(METRIC) FROM MyTable GROUP BY RAW_STRING_HUGE LIMIT 10"); + + private static final TableConfig TABLE_CONFIG = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(TABLE_NAME) + .setNoDictionaryColumns(java.util.List.of(RAW_STRING_HUGE)) + .build(); + + private static final Schema SCHEMA = new Schema.SchemaBuilder() + .setSchemaName(TABLE_NAME) + .addSingleValueDimension(DICT_INT_HUGE, FieldSpec.DataType.INT) + .addSingleValueDimension(RAW_STRING_HUGE, FieldSpec.DataType.STRING) + .addMetric(METRIC, FieldSpec.DataType.LONG) + .build(); + + @Param({"DICT_INT", "RAW_STRING"}) + private String _scenario; + @Param({"false", "true"}) + private String _groupByOffHeap; + + private InstancePlanMakerImplV2 _planMaker; + private IndexSegment _indexSegment; + private String _query; + + @Setup + public void setUp() + throws Exception { + OffHeapGroupByBufferPool.setMaxBytesPerThread(2L << 30); + _planMaker = new InstancePlanMakerImplV2(); + _planMaker.init(new PinotConfiguration(Map.of( + CommonConstants.Server.CONFIG_OF_QUERY_EXECUTOR_NUM_GROUPS_LIMIT, RAISED_NUM_GROUPS_LIMIT, + CommonConstants.Server.CONFIG_OF_QUERY_EXECUTOR_NUM_GROUPS_WARN_LIMIT, RAISED_NUM_GROUPS_LIMIT))); + FileUtils.deleteQuietly(INDEX_DIR); + buildSegment(); + _indexSegment = + ImmutableSegmentLoader.load(new File(INDEX_DIR, "testSegment"), new IndexLoadingConfig(TABLE_CONFIG, SCHEMA)); + _query = QUERIES.get(_scenario); + } + + @TearDown + public void tearDown() { + _indexSegment.destroy(); + FileUtils.deleteQuietly(INDEX_DIR); + } + + private void buildSegment() + throws Exception { + Random random = new Random(42); + LazyDataGenerator rows = new LazyDataGenerator() { + @Override + public int size() { + return NUM_ROWS; + } + + @Override + public GenericRow next(GenericRow row, int i) { + int group = random.nextInt(CARDINALITY); + row.putValue(DICT_INT_HUGE, group); + row.putValue(RAW_STRING_HUGE, makeKey(group)); + row.putValue(METRIC, (long) random.nextInt(1000)); + return null; + } + + @Override + public void rewind() { + random.setSeed(42); + } + }; + SegmentGeneratorConfig config = new SegmentGeneratorConfig(TABLE_CONFIG, SCHEMA); + config.setOutDir(INDEX_DIR.getPath()); + config.setTableName(TABLE_NAME); + config.setSegmentName("testSegment"); + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + try (RecordReader recordReader = new GeneratedDataRecordReader(rows)) { + driver.init(config, recordReader); + driver.build(); + } + } + + private static String makeKey(int i) { + char[] chars = {'k', 'e', 'y', '-', '0', '0', '0', '0', '0', '0', '0', '0', '0', '-', + 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'}; + int value = i; + for (int position = 12; position >= 4 && value > 0; position--) { + chars[position] = (char) ('0' + (value % 10)); + value /= 10; + } + return new String(chars); + } + + /// Single-segment group-by over ~10M distinct groups; the result block's generator (owning the off-heap state) + /// is closed after each invocation, mirroring the combine operator's contract. + @Benchmark + public GroupByResultsBlock segmentGroupBy() { + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(_query); + pinotQuery.setQueryOptions(new HashMap<>()); + QueryContext queryContext = QueryContextConverterUtils.getQueryContext(pinotQuery); + queryContext.setEndTimeMs( + System.currentTimeMillis() + CommonConstants.Server.DEFAULT_QUERY_EXECUTOR_TIMEOUT_MS); + queryContext.setNumGroupsLimit(RAISED_NUM_GROUPS_LIMIT); + queryContext.setNumGroupsWarningLimit(RAISED_NUM_GROUPS_LIMIT); + queryContext.setGroupByOffHeap(Boolean.parseBoolean(_groupByOffHeap)); + GroupByOperator groupByOperator = + (GroupByOperator) _planMaker.makeSegmentPlanNode(new SegmentContext(_indexSegment), queryContext).run(); + GroupByResultsBlock resultsBlock = groupByOperator.nextBlock(); + AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult(); + if (aggregationGroupByResult != null) { + aggregationGroupByResult.closeGroupKeyGenerator(); + } + return resultsBlock; + } + + public static void main(String[] args) + throws Exception { + new Runner( + new OptionsBuilder().include(BenchmarkOffHeapGroupByHugeSSE.class.getSimpleName()).addProfiler("gc") + .build()).run(); + } +} diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupByLargeSSE.java b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupByLargeSSE.java new file mode 100644 index 000000000000..66123f4634a0 --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupByLargeSSE.java @@ -0,0 +1,278 @@ +/** + * 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.perf; + +import java.io.File; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.datatable.DataTable; +import org.apache.pinot.common.datatable.DataTableFactory; +import org.apache.pinot.common.metrics.BrokerMetrics; +import org.apache.pinot.common.request.BrokerRequest; +import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.common.response.broker.BrokerResponseNative; +import org.apache.pinot.core.operator.blocks.InstanceResponseBlock; +import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock; +import org.apache.pinot.core.operator.query.GroupByOperator; +import org.apache.pinot.core.plan.Plan; +import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2; +import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapGroupByBufferPool; +import org.apache.pinot.core.query.reduce.BrokerReduceService; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.apache.pinot.core.transport.ServerRoutingInstance; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.SegmentContext; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +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.data.readers.GenericRow; +import org.apache.pinot.spi.data.readers.RecordReader; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.apache.pinot.sql.parsers.CalciteSqlCompiler; +import org.apache.pinot.sql.parsers.CalciteSqlParser; +import org.mockito.Mockito; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.runner.Runner; +import org.openjdk.jmh.runner.options.OptionsBuilder; + + +/// Very-large-group SSE benchmark: ~1M distinct groups per segment (numGroupsLimit raised to 2.1M), 2 segments x +/// 2M rows, on-heap vs off-heap group-by state. Two measurements per configuration: +///
    +///
  • [#query()]: the full flow (plan, 8-thread combine, serialize, broker reduce). NOTE: the combine +/// phase merges ~1M groups into the (mode-independent, on-heap) IndexedTable in both arms, so it dilutes the +/// per-segment difference — kept for the honest end-to-end picture.
    • +///
  • [#segmentGroupBy()]: a single segment's GroupByOperator only — isolates the phase the off-heap +/// feature changes. The block's group key generator is closed after each invocation, mirroring the combine +/// operator's contract.
    • +///
+/// Run with `-prof gc`: score (ms/op), gc.alloc.rate.norm and gc.count are the interesting metrics; pair +/// with [OffHeapGroupByMemoryFootprint] for retained-heap numbers. +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@Fork(1) +@Warmup(iterations = 2, time = 5) +@Measurement(iterations = 5, time = 5) +@State(Scope.Benchmark) +public class BenchmarkOffHeapGroupByLargeSSE { + private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "BenchmarkOffHeapGroupByLargeSSE"); + private static final String TABLE_NAME = "MyTable"; + private static final int NUM_SEGMENTS = 2; + private static final int NUM_ROWS_PER_SEGMENT = 2_000_000; + private static final int CARDINALITY = 1_000_000; + private static final int RAISED_NUM_GROUPS_LIMIT = 2_100_000; + + private static final String DICT_INT_LARGE = "DICT_INT_LARGE"; + private static final String RAW_STRING_LARGE = "RAW_STRING_LARGE"; + private static final String METRIC = "METRIC"; + + private static final Map QUERIES = Map.of( + "DICT_INT", "SELECT DICT_INT_LARGE, COUNT(*), SUM(METRIC) FROM MyTable GROUP BY DICT_INT_LARGE LIMIT 10", + "RAW_STRING", + "SELECT RAW_STRING_LARGE, COUNT(*), SUM(METRIC) FROM MyTable GROUP BY RAW_STRING_LARGE LIMIT 10"); + + private static final TableConfig TABLE_CONFIG = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(TABLE_NAME) + .setNoDictionaryColumns(List.of(RAW_STRING_LARGE)) + .build(); + + private static final Schema SCHEMA = new Schema.SchemaBuilder() + .setSchemaName(TABLE_NAME) + .addSingleValueDimension(DICT_INT_LARGE, FieldSpec.DataType.INT) + .addSingleValueDimension(RAW_STRING_LARGE, FieldSpec.DataType.STRING) + .addMetric(METRIC, FieldSpec.DataType.LONG) + .build(); + + private static final BrokerMetrics BROKER_METRICS = Mockito.mock(BrokerMetrics.class); + + @Param({"DICT_INT", "RAW_STRING"}) + private String _scenario; + @Param({"false", "true"}) + private String _groupByOffHeap; + + private InstancePlanMakerImplV2 _planMaker; + private ExecutorService _executorService; + private BrokerReduceService _brokerReduceService; + private List _indexSegments; + private String _query; + + @Setup + public void setUp() + throws Exception { + // Recommended production configuration for off-heap group-by: pool buffers per thread + OffHeapGroupByBufferPool.setMaxBytesPerThread(256L << 20); + // Raise the group limits so ~1M groups per segment are never capped (and never spam the warn log) + _planMaker = new InstancePlanMakerImplV2(); + _planMaker.init(new PinotConfiguration(Map.of( + CommonConstants.Server.CONFIG_OF_QUERY_EXECUTOR_NUM_GROUPS_LIMIT, RAISED_NUM_GROUPS_LIMIT, + CommonConstants.Server.CONFIG_OF_QUERY_EXECUTOR_NUM_GROUPS_WARN_LIMIT, RAISED_NUM_GROUPS_LIMIT))); + _executorService = Executors.newFixedThreadPool(8); + _brokerReduceService = new BrokerReduceService( + new PinotConfiguration(Map.of(CommonConstants.Broker.CONFIG_OF_MAX_REDUCE_THREADS_PER_QUERY, 2))); + FileUtils.deleteQuietly(INDEX_DIR); + _indexSegments = new ArrayList<>(NUM_SEGMENTS); + IndexLoadingConfig indexLoadingConfig = new IndexLoadingConfig(TABLE_CONFIG, SCHEMA); + for (int i = 0; i < NUM_SEGMENTS; i++) { + String segmentName = "testSegment" + i; + buildSegment(segmentName, i); + _indexSegments.add(ImmutableSegmentLoader.load(new File(INDEX_DIR, segmentName), indexLoadingConfig)); + } + _query = QUERIES.get(_scenario); + } + + @TearDown + public void tearDown() { + for (IndexSegment indexSegment : _indexSegments) { + indexSegment.destroy(); + } + FileUtils.deleteQuietly(INDEX_DIR); + _executorService.shutdownNow(); + _brokerReduceService.shutDown(); + } + + private void buildSegment(String segmentName, int segmentIndex) + throws Exception { + Random random = new Random(42 + segmentIndex); + LazyDataGenerator rows = new LazyDataGenerator() { + @Override + public int size() { + return NUM_ROWS_PER_SEGMENT; + } + + @Override + public GenericRow next(GenericRow row, int i) { + int group = random.nextInt(CARDINALITY); + row.putValue(DICT_INT_LARGE, group); + row.putValue(RAW_STRING_LARGE, String.format("key-%08d-abcdefgh", group)); + row.putValue(METRIC, (long) random.nextInt(1000)); + return null; + } + + @Override + public void rewind() { + random.setSeed(42 + segmentIndex); + } + }; + SegmentGeneratorConfig config = new SegmentGeneratorConfig(TABLE_CONFIG, SCHEMA); + config.setOutDir(INDEX_DIR.getPath()); + config.setTableName(TABLE_NAME); + config.setSegmentName(segmentName); + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + try (RecordReader recordReader = new GeneratedDataRecordReader(rows)) { + driver.init(config, recordReader); + driver.build(); + } + } + + private QueryContext buildQueryContext() { + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(_query); + Map queryOptions = new HashMap<>(); + queryOptions.put("groupByOffHeap", _groupByOffHeap); + queryOptions.put("numGroupsLimit", String.valueOf(RAISED_NUM_GROUPS_LIMIT)); + pinotQuery.setQueryOptions(queryOptions); + QueryContext queryContext = QueryContextConverterUtils.getQueryContext(pinotQuery); + queryContext.setEndTimeMs( + System.currentTimeMillis() + CommonConstants.Server.DEFAULT_QUERY_EXECUTOR_TIMEOUT_MS); + return queryContext; + } + + /// Full end-to-end flow over both segments (combine + serialize + reduce are mode-independent). + @Benchmark + public BrokerResponseNative query() + throws TimeoutException { + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(_query); + Map queryOptions = new HashMap<>(); + queryOptions.put("groupByOffHeap", _groupByOffHeap); + queryOptions.put("numGroupsLimit", String.valueOf(RAISED_NUM_GROUPS_LIMIT)); + pinotQuery.setQueryOptions(queryOptions); + QueryContext queryContext = QueryContextConverterUtils.getQueryContext(pinotQuery); + queryContext.setEndTimeMs( + System.currentTimeMillis() + CommonConstants.Server.DEFAULT_QUERY_EXECUTOR_TIMEOUT_MS); + List segmentContexts = new ArrayList<>(_indexSegments.size()); + _indexSegments.forEach(s -> segmentContexts.add(new SegmentContext(s))); + Plan plan = _planMaker.makeInstancePlan(segmentContexts, queryContext, _executorService); + InstanceResponseBlock instanceResponse = plan.execute(); + Map dataTableMap = new HashMap<>(); + try { + byte[] serializedResponse = instanceResponse.toDataTable().toBytes(); + dataTableMap.put(new ServerRoutingInstance("localhost", 1234, TableType.OFFLINE), + DataTableFactory.getDataTable(serializedResponse)); + } catch (Exception e) { + throw new RuntimeException(e); + } + BrokerRequest brokerRequest = CalciteSqlCompiler.convertToBrokerRequest(pinotQuery); + return _brokerReduceService.reduceOnDataTable(brokerRequest, brokerRequest, dataTableMap, + CommonConstants.Broker.DEFAULT_BROKER_TIMEOUT_MS, BROKER_METRICS); + } + + /// Single-segment group-by only: isolates the phase the off-heap feature changes. The result block's group key + /// generator (owning the off-heap state) is closed after each invocation, mirroring the combine operator. + @Benchmark + public GroupByResultsBlock segmentGroupBy() { + QueryContext queryContext = buildQueryContext(); + // Apply the group-by limits the instance plan would apply + queryContext.setNumGroupsLimit(RAISED_NUM_GROUPS_LIMIT); + queryContext.setNumGroupsWarningLimit(RAISED_NUM_GROUPS_LIMIT); + queryContext.setGroupByOffHeap(Boolean.parseBoolean(_groupByOffHeap)); + GroupByOperator groupByOperator = + (GroupByOperator) _planMaker.makeSegmentPlanNode(new SegmentContext(_indexSegments.get(0)), queryContext) + .run(); + GroupByResultsBlock resultsBlock = groupByOperator.nextBlock(); + AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult(); + if (aggregationGroupByResult != null) { + aggregationGroupByResult.closeGroupKeyGenerator(); + } + return resultsBlock; + } + + public static void main(String[] args) + throws Exception { + new Runner( + new OptionsBuilder().include(BenchmarkOffHeapGroupByLargeSSE.class.getSimpleName()).addProfiler("gc") + .build()).run(); + } +} diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupBySSE.java b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupBySSE.java new file mode 100644 index 000000000000..d18b114a1159 --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupBySSE.java @@ -0,0 +1,255 @@ +/** + * 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.perf; + +import java.io.File; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.datatable.DataTable; +import org.apache.pinot.common.datatable.DataTableFactory; +import org.apache.pinot.common.metrics.BrokerMetrics; +import org.apache.pinot.common.request.BrokerRequest; +import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.common.response.broker.BrokerResponseNative; +import org.apache.pinot.core.operator.blocks.InstanceResponseBlock; +import org.apache.pinot.core.plan.Plan; +import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2; +import org.apache.pinot.core.plan.maker.PlanMaker; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapGroupByBufferPool; +import org.apache.pinot.core.query.reduce.BrokerReduceService; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.apache.pinot.core.transport.ServerRoutingInstance; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.SegmentContext; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +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.data.readers.GenericRow; +import org.apache.pinot.spi.data.readers.RecordReader; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.apache.pinot.sql.parsers.CalciteSqlCompiler; +import org.apache.pinot.sql.parsers.CalciteSqlParser; +import org.mockito.Mockito; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.runner.Runner; +import org.openjdk.jmh.runner.options.OptionsBuilder; + + +/// End-to-end SSE group-by benchmark comparing on-heap vs off-heap group-by state (the `groupByOffHeap` query +/// option) over the generator tiers the off-heap feature covers: +///
    +///
  • `DICT_INT`: single dict-encoded high-cardinality INT — DictionaryBased IntMap tier
    • +///
  • `DICT_TWO_COLS`: two dict-encoded columns whose cardinality product exceeds Integer.MAX_VALUE — +/// DictionaryBased LongMap tier
    • +///
  • `RAW_INT`: single raw INT — NoDictionarySingleColumn long-key tier
    • +///
  • `RAW_STRING`: single raw STRING — NoDictionarySingleColumn bytes tier
    • +///
  • `RAW_MULTI`: raw INT + raw STRING — NoDictionaryMultiColumn packed-bytes tier
    • +///
+/// 4 segments x 300K rows, ~80K/30K distinct groups (below the default numGroupsLimit, so no capping). LIMIT 10 +/// without ORDER BY keeps the (mode-independent) combine/reduce phases cheap relative to the per-segment phase this +/// feature changes. Run with `-prof gc`: the interesting metrics are score (ms/op) and gc.alloc.rate.norm. +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@Fork(1) +@Warmup(iterations = 3, time = 3) +@Measurement(iterations = 5, time = 3) +@State(Scope.Benchmark) +public class BenchmarkOffHeapGroupBySSE { + private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "BenchmarkOffHeapGroupBySSE"); + private static final String TABLE_NAME = "MyTable"; + private static final int NUM_SEGMENTS = 4; + private static final int NUM_ROWS_PER_SEGMENT = 300_000; + private static final int INT_HIGH_CARDINALITY = 80_000; + private static final int INT_MED_CARDINALITY = 30_000; + private static final int RAW_INT_CARDINALITY = 80_000; + private static final int RAW_STRING_CARDINALITY = 30_000; + + private static final String INT_HIGH = "INT_HIGH"; + private static final String INT_MED = "INT_MED"; + private static final String RAW_INT = "RAW_INT"; + private static final String RAW_STRING = "RAW_STRING"; + private static final String METRIC = "METRIC"; + + private static final Map QUERIES = Map.of( + "DICT_INT", "SELECT INT_HIGH, COUNT(*), SUM(METRIC) FROM MyTable GROUP BY INT_HIGH LIMIT 10", + "DICT_TWO_COLS", + "SELECT INT_HIGH, INT_MED, COUNT(*), SUM(METRIC) FROM MyTable GROUP BY INT_HIGH, INT_MED LIMIT 10", + "RAW_INT", "SELECT RAW_INT, COUNT(*), SUM(METRIC) FROM MyTable GROUP BY RAW_INT LIMIT 10", + "RAW_STRING", "SELECT RAW_STRING, COUNT(*), SUM(METRIC) FROM MyTable GROUP BY RAW_STRING LIMIT 10", + "RAW_MULTI", + "SELECT RAW_INT, RAW_STRING, COUNT(*), SUM(METRIC) FROM MyTable GROUP BY RAW_INT, RAW_STRING LIMIT 10"); + + private static final TableConfig TABLE_CONFIG = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(TABLE_NAME) + .setNoDictionaryColumns(List.of(RAW_INT, RAW_STRING)) + .build(); + + private static final Schema SCHEMA = new Schema.SchemaBuilder() + .setSchemaName(TABLE_NAME) + .addSingleValueDimension(INT_HIGH, FieldSpec.DataType.INT) + .addSingleValueDimension(INT_MED, FieldSpec.DataType.INT) + .addSingleValueDimension(RAW_INT, FieldSpec.DataType.INT) + .addSingleValueDimension(RAW_STRING, FieldSpec.DataType.STRING) + .addMetric(METRIC, FieldSpec.DataType.LONG) + .build(); + + private static final PlanMaker PLAN_MAKER = new InstancePlanMakerImplV2(); + private static final BrokerMetrics BROKER_METRICS = Mockito.mock(BrokerMetrics.class); + + @Param({"DICT_INT", "DICT_TWO_COLS", "RAW_INT", "RAW_STRING", "RAW_MULTI"}) + private String _scenario; + @Param({"false", "true"}) + private String _groupByOffHeap; + + private ExecutorService _executorService; + private List _indexSegments; + private String _query; + + @Setup + public void setUp() + throws Exception { + // Recommended production configuration for off-heap group-by: pool buffers per thread, mirroring the + // on-heap thread-local map caching (bench measures steady-state reuse in both modes) + OffHeapGroupByBufferPool.setMaxBytesPerThread(64L << 20); + _executorService = Executors.newFixedThreadPool(8); + FileUtils.deleteQuietly(INDEX_DIR); + _indexSegments = new ArrayList<>(NUM_SEGMENTS); + IndexLoadingConfig indexLoadingConfig = new IndexLoadingConfig(TABLE_CONFIG, SCHEMA); + for (int i = 0; i < NUM_SEGMENTS; i++) { + String segmentName = "testSegment" + i; + buildSegment(segmentName, i); + _indexSegments.add(ImmutableSegmentLoader.load(new File(INDEX_DIR, segmentName), indexLoadingConfig)); + } + _query = QUERIES.get(_scenario); + } + + @TearDown + public void tearDown() { + for (IndexSegment indexSegment : _indexSegments) { + indexSegment.destroy(); + } + FileUtils.deleteQuietly(INDEX_DIR); + _executorService.shutdownNow(); + } + + private void buildSegment(String segmentName, int segmentIndex) + throws Exception { + Random random = new Random(42 + segmentIndex); + LazyDataGenerator rows = new LazyDataGenerator() { + @Override + public int size() { + return NUM_ROWS_PER_SEGMENT; + } + + @Override + public GenericRow next(GenericRow row, int i) { + row.putValue(INT_HIGH, random.nextInt(INT_HIGH_CARDINALITY)); + row.putValue(INT_MED, random.nextInt(INT_MED_CARDINALITY)); + row.putValue(RAW_INT, random.nextInt(RAW_INT_CARDINALITY)); + row.putValue(RAW_STRING, String.format("key-%08d-abcdefgh", random.nextInt(RAW_STRING_CARDINALITY))); + row.putValue(METRIC, (long) random.nextInt(1000)); + return null; + } + + @Override + public void rewind() { + random.setSeed(42 + segmentIndex); + } + }; + SegmentGeneratorConfig config = new SegmentGeneratorConfig(TABLE_CONFIG, SCHEMA); + config.setOutDir(INDEX_DIR.getPath()); + config.setTableName(TABLE_NAME); + config.setSegmentName(segmentName); + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + try (RecordReader recordReader = new GeneratedDataRecordReader(rows)) { + driver.init(config, recordReader); + driver.build(); + } + } + + @Benchmark + public BrokerResponseNative query() + throws TimeoutException { + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(_query); + Map queryOptions = new HashMap<>(); + queryOptions.put("groupByOffHeap", _groupByOffHeap); + pinotQuery.setQueryOptions(queryOptions); + QueryContext queryContext = QueryContextConverterUtils.getQueryContext(pinotQuery); + queryContext.setGroupByOffHeap(Boolean.parseBoolean(_groupByOffHeap)); + queryContext.setEndTimeMs( + System.currentTimeMillis() + CommonConstants.Server.DEFAULT_QUERY_EXECUTOR_TIMEOUT_MS); + + // Server side + List segmentContexts = new ArrayList<>(_indexSegments.size()); + _indexSegments.forEach(s -> segmentContexts.add(new SegmentContext(s))); + Plan plan = PLAN_MAKER.makeInstancePlan(segmentContexts, queryContext, _executorService); + InstanceResponseBlock instanceResponse = plan.execute(); + + // Broker side + Map dataTableMap = new HashMap<>(); + try { + byte[] serializedResponse = instanceResponse.toDataTable().toBytes(); + dataTableMap.put(new ServerRoutingInstance("localhost", 1234, TableType.OFFLINE), + DataTableFactory.getDataTable(serializedResponse)); + } catch (Exception e) { + throw new RuntimeException(e); + } + BrokerRequest brokerRequest = CalciteSqlCompiler.convertToBrokerRequest(pinotQuery); + BrokerReduceService brokerReduceService = new BrokerReduceService( + new PinotConfiguration(Map.of(CommonConstants.Broker.CONFIG_OF_MAX_REDUCE_THREADS_PER_QUERY, 2))); + BrokerResponseNative brokerResponse = brokerReduceService.reduceOnDataTable(brokerRequest, brokerRequest, + dataTableMap, CommonConstants.Broker.DEFAULT_BROKER_TIMEOUT_MS, BROKER_METRICS); + brokerReduceService.shutDown(); + return brokerResponse; + } + + public static void main(String[] args) + throws Exception { + new Runner( + new OptionsBuilder().include(BenchmarkOffHeapGroupBySSE.class.getSimpleName()).addProfiler("gc").build()) + .run(); + } +} diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupByUllSSE.java b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupByUllSSE.java new file mode 100644 index 000000000000..2276935cfe86 --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupByUllSSE.java @@ -0,0 +1,273 @@ +/** + * 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.perf; + +import java.io.File; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.datatable.DataTable; +import org.apache.pinot.common.datatable.DataTableFactory; +import org.apache.pinot.common.metrics.BrokerMetrics; +import org.apache.pinot.common.request.BrokerRequest; +import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.common.response.broker.BrokerResponseNative; +import org.apache.pinot.core.operator.blocks.InstanceResponseBlock; +import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock; +import org.apache.pinot.core.operator.query.GroupByOperator; +import org.apache.pinot.core.plan.Plan; +import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2; +import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapGroupByBufferPool; +import org.apache.pinot.core.query.reduce.BrokerReduceService; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.apache.pinot.core.transport.ServerRoutingInstance; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.SegmentContext; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +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.data.readers.GenericRow; +import org.apache.pinot.spi.data.readers.RecordReader; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.apache.pinot.sql.parsers.CalciteSqlCompiler; +import org.apache.pinot.sql.parsers.CalciteSqlParser; +import org.mockito.Mockito; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.runner.Runner; +import org.openjdk.jmh.runner.options.OptionsBuilder; + + +/// SSE benchmark for off-heap `DISTINCTCOUNTULL` group-by state (`OffHeapUltraLogLogGroupByResultHolder`). +/// 2 segments x 2M rows; the group column is a dict-encoded INT whose cardinality is the `numGroups` param, the +/// ULL input is a raw (no-dictionary) LONG so the per-group state is the ULL register array itself — the exact +/// state this feature moves off-heap (~4.1KB per group at the default p=12: 10K groups ≈ 42MB, 200K groups ≈ +/// 840MB of on-heap sketch state per segment execution). +/// +/// Two measurements per configuration, mirroring [BenchmarkOffHeapGroupByLargeSSE]: +///
    +///
  • [#query()]: the full flow (plan, 8-thread combine, serialize, broker reduce); the combine phase merges +/// heap ULLs in both arms.
    • +///
  • [#segmentGroupBy()]: a single segment's GroupByOperator only — isolates the phase the off-heap holder +/// changes.
    • +///
+/// Run with `-prof gc` and at least `-Xmx10g` (the 200K-group on-heap arm retains ~1.7GB of sketches during the +/// combine): score (ms/op), gc.alloc.rate.norm and gc.count are the interesting metrics. +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@Fork(1) +@Warmup(iterations = 2, time = 5) +@Measurement(iterations = 5, time = 5) +@State(Scope.Benchmark) +public class BenchmarkOffHeapGroupByUllSSE { + private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "BenchmarkOffHeapGroupByUllSSE"); + private static final String TABLE_NAME = "MyTable"; + private static final int NUM_SEGMENTS = 2; + private static final int NUM_ROWS_PER_SEGMENT = 2_000_000; + private static final int VALUE_CARDINALITY = 1_000_000; + private static final int RAISED_NUM_GROUPS_LIMIT = 2_100_000; + + private static final String DICT_INT_GROUP = "DICT_INT_GROUP"; + private static final String RAW_LONG_VALUE = "RAW_LONG_VALUE"; + + private static final Schema SCHEMA = new Schema.SchemaBuilder() + .setSchemaName(TABLE_NAME) + .addSingleValueDimension(DICT_INT_GROUP, FieldSpec.DataType.INT) + .addSingleValueDimension(RAW_LONG_VALUE, FieldSpec.DataType.LONG) + .build(); + + private static final TableConfig TABLE_CONFIG = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(TABLE_NAME) + .setNoDictionaryColumns(List.of(RAW_LONG_VALUE)) + .build(); + + private static final String QUERY = + "SELECT DICT_INT_GROUP, DISTINCTCOUNTULL(RAW_LONG_VALUE) FROM MyTable GROUP BY DICT_INT_GROUP LIMIT 10"; + + private static final BrokerMetrics BROKER_METRICS = Mockito.mock(BrokerMetrics.class); + + @Param({"10000", "200000"}) + private int _numGroups; + @Param({"false", "true"}) + private String _groupByOffHeap; + + private InstancePlanMakerImplV2 _planMaker; + private ExecutorService _executorService; + private BrokerReduceService _brokerReduceService; + private List _indexSegments; + + @Setup + public void setUp() + throws Exception { + // Recommended production configuration for off-heap group-by: pool buffers per thread + OffHeapGroupByBufferPool.setMaxBytesPerThread(256L << 20); + _planMaker = new InstancePlanMakerImplV2(); + _planMaker.init(new PinotConfiguration(Map.of( + CommonConstants.Server.CONFIG_OF_QUERY_EXECUTOR_NUM_GROUPS_LIMIT, RAISED_NUM_GROUPS_LIMIT, + CommonConstants.Server.CONFIG_OF_QUERY_EXECUTOR_NUM_GROUPS_WARN_LIMIT, RAISED_NUM_GROUPS_LIMIT))); + _executorService = Executors.newFixedThreadPool(8); + _brokerReduceService = new BrokerReduceService( + new PinotConfiguration(Map.of(CommonConstants.Broker.CONFIG_OF_MAX_REDUCE_THREADS_PER_QUERY, 2))); + FileUtils.deleteQuietly(INDEX_DIR); + _indexSegments = new ArrayList<>(NUM_SEGMENTS); + IndexLoadingConfig indexLoadingConfig = new IndexLoadingConfig(TABLE_CONFIG, SCHEMA); + for (int i = 0; i < NUM_SEGMENTS; i++) { + String segmentName = "testSegment" + i; + buildSegment(segmentName, i); + _indexSegments.add(ImmutableSegmentLoader.load(new File(INDEX_DIR, segmentName), indexLoadingConfig)); + } + } + + @TearDown + public void tearDown() { + // Restore the (static) pool config so the benchmark leaves no global state behind + OffHeapGroupByBufferPool.setMaxBytesPerThread(0); + for (IndexSegment indexSegment : _indexSegments) { + indexSegment.destroy(); + } + FileUtils.deleteQuietly(INDEX_DIR); + _executorService.shutdownNow(); + _brokerReduceService.shutDown(); + } + + private void buildSegment(String segmentName, int segmentIndex) + throws Exception { + Random random = new Random(42 + segmentIndex); + LazyDataGenerator rows = new LazyDataGenerator() { + @Override + public int size() { + return NUM_ROWS_PER_SEGMENT; + } + + @Override + public GenericRow next(GenericRow row, int i) { + row.putValue(DICT_INT_GROUP, random.nextInt(_numGroups)); + row.putValue(RAW_LONG_VALUE, (long) random.nextInt(VALUE_CARDINALITY)); + return null; + } + + @Override + public void rewind() { + random.setSeed(42 + segmentIndex); + } + }; + SegmentGeneratorConfig config = new SegmentGeneratorConfig(TABLE_CONFIG, SCHEMA); + config.setOutDir(INDEX_DIR.getPath()); + config.setTableName(TABLE_NAME); + config.setSegmentName(segmentName); + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + try (RecordReader recordReader = new GeneratedDataRecordReader(rows)) { + driver.init(config, recordReader); + driver.build(); + } + } + + private QueryContext buildQueryContext() { + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(QUERY); + Map queryOptions = new HashMap<>(); + queryOptions.put("groupByOffHeap", _groupByOffHeap); + queryOptions.put("numGroupsLimit", String.valueOf(RAISED_NUM_GROUPS_LIMIT)); + pinotQuery.setQueryOptions(queryOptions); + QueryContext queryContext = QueryContextConverterUtils.getQueryContext(pinotQuery); + queryContext.setEndTimeMs( + System.currentTimeMillis() + CommonConstants.Server.DEFAULT_QUERY_EXECUTOR_TIMEOUT_MS); + return queryContext; + } + + /// Full end-to-end flow over both segments (combine + serialize + reduce merge heap ULLs in both arms). + @Benchmark + public BrokerResponseNative query() + throws TimeoutException { + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(QUERY); + Map queryOptions = new HashMap<>(); + queryOptions.put("groupByOffHeap", _groupByOffHeap); + queryOptions.put("numGroupsLimit", String.valueOf(RAISED_NUM_GROUPS_LIMIT)); + pinotQuery.setQueryOptions(queryOptions); + QueryContext queryContext = QueryContextConverterUtils.getQueryContext(pinotQuery); + queryContext.setEndTimeMs( + System.currentTimeMillis() + CommonConstants.Server.DEFAULT_QUERY_EXECUTOR_TIMEOUT_MS); + List segmentContexts = new ArrayList<>(_indexSegments.size()); + _indexSegments.forEach(s -> segmentContexts.add(new SegmentContext(s))); + Plan plan = _planMaker.makeInstancePlan(segmentContexts, queryContext, _executorService); + InstanceResponseBlock instanceResponse = plan.execute(); + Map dataTableMap = new HashMap<>(); + try { + byte[] serializedResponse = instanceResponse.toDataTable().toBytes(); + dataTableMap.put(new ServerRoutingInstance("localhost", 1234, TableType.OFFLINE), + DataTableFactory.getDataTable(serializedResponse)); + } catch (Exception e) { + throw new RuntimeException(e); + } + BrokerRequest brokerRequest = CalciteSqlCompiler.convertToBrokerRequest(pinotQuery); + return _brokerReduceService.reduceOnDataTable(brokerRequest, brokerRequest, dataTableMap, + CommonConstants.Broker.DEFAULT_BROKER_TIMEOUT_MS, BROKER_METRICS); + } + + /// Single-segment group-by only: isolates the phase the off-heap ULL holder changes. The result block's group + /// key generator (owning the off-heap state) is closed after each invocation, mirroring the combine operator. + @Benchmark + public GroupByResultsBlock segmentGroupBy() { + QueryContext queryContext = buildQueryContext(); + queryContext.setNumGroupsLimit(RAISED_NUM_GROUPS_LIMIT); + queryContext.setNumGroupsWarningLimit(RAISED_NUM_GROUPS_LIMIT); + queryContext.setGroupByOffHeap(Boolean.parseBoolean(_groupByOffHeap)); + GroupByOperator groupByOperator = + (GroupByOperator) _planMaker.makeSegmentPlanNode(new SegmentContext(_indexSegments.get(0)), queryContext) + .run(); + GroupByResultsBlock resultsBlock = groupByOperator.nextBlock(); + AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult(); + if (aggregationGroupByResult != null) { + aggregationGroupByResult.closeGroupKeyGenerator(); + } + return resultsBlock; + } + + public static void main(String[] args) + throws Exception { + new Runner( + new OptionsBuilder().include(BenchmarkOffHeapGroupByUllSSE.class.getSimpleName()).addProfiler("gc") + .build()).run(); + } +} diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupIdMaps.java b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupIdMaps.java new file mode 100644 index 000000000000..401a526d0227 --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupIdMaps.java @@ -0,0 +1,175 @@ +/** + * 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.perf; + +import it.unimi.dsi.fastutil.longs.Long2IntOpenHashMap; +import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap; +import java.util.Random; +import java.util.concurrent.TimeUnit; +import org.apache.pinot.core.query.aggregation.groupby.DictionaryBasedGroupKeyGenerator.IntGroupIdMap; +import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapBytesGroupIdMap; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapGroupByUtils; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapIntGroupIdMap; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapLongGroupIdMap; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.runner.Runner; +import org.openjdk.jmh.runner.options.OptionsBuilder; + + +/// Micro-benchmark comparing the off-heap group-by key tables against the on-heap structures they replace, over the +/// per-segment group-by access pattern: a stream of raw keys with duplicates mapped to dense group ids, with the map +/// created and released per pass (mirroring the per-query lifecycle — the on-heap maps are thread-local-cached in +/// production, so the on-heap numbers here are slightly pessimistic on construction cost, while the on-heap string +/// map benefits from cached String hash codes that the real per-block string materialization does not have). +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@Fork(1) +@Warmup(iterations = 3, time = 2) +@Measurement(iterations = 5, time = 3) +@State(Scope.Benchmark) +public class BenchmarkOffHeapGroupIdMaps { + private static final int NUM_OPS = 2_000_000; + private static final int GROUP_ID_UPPER_BOUND = Integer.MAX_VALUE; + + @Param({"10000", "100000", "1000000"}) + private int _numDistinct; + + private int[] _intKeys; + private long[] _longKeys; + private String[] _stringKeys; + private byte[] _encodeScratch; + + @Setup + public void setUp() { + Random random = new Random(42); + _intKeys = new int[NUM_OPS]; + _longKeys = new long[NUM_OPS]; + _stringKeys = new String[NUM_OPS]; + // Distinct string values reused by reference (interned per distinct id) so the on-heap map sees cached + // hash codes — a deliberate bias in favor of the on-heap baseline + String[] distinctStrings = new String[Math.min(_numDistinct, 1_000_000)]; + for (int i = 0; i < distinctStrings.length; i++) { + distinctStrings[i] = String.format("key-%08d-abcdefgh", i); + } + for (int i = 0; i < NUM_OPS; i++) { + int distinct = random.nextInt(_numDistinct); + _intKeys[i] = distinct; + _longKeys[i] = ((long) distinct << 20) | (distinct & 0xFFFFF); + _stringKeys[i] = distinctStrings[distinct % distinctStrings.length]; + } + int maxKeyLength = 0; + for (String key : _stringKeys) { + maxKeyLength = Math.max(maxKeyLength, key.length()); + } + _encodeScratch = new byte[maxKeyLength * 3]; + } + + @Benchmark + public long onHeapIntMap() { + IntGroupIdMap map = new IntGroupIdMap(); + long sum = 0; + for (int key : _intKeys) { + sum += map.getGroupId(key, GROUP_ID_UPPER_BOUND); + } + map.clearAndTrim(); + return sum; + } + + @Benchmark + public long offHeapIntMap() { + long sum = 0; + try (OffHeapIntGroupIdMap map = new OffHeapIntGroupIdMap(0)) { + for (int key : _intKeys) { + sum += map.getGroupId(key, GROUP_ID_UPPER_BOUND); + } + } + return sum; + } + + @Benchmark + public long onHeapLongMap() { + Long2IntOpenHashMap map = new Long2IntOpenHashMap(); + map.defaultReturnValue(GroupKeyGenerator.INVALID_ID); + long sum = 0; + for (long key : _longKeys) { + int numGroups = map.size(); + int groupId = map.putIfAbsent(key, numGroups); + sum += groupId == GroupKeyGenerator.INVALID_ID ? numGroups : groupId; + } + return sum; + } + + @Benchmark + public long offHeapLongMap() { + long sum = 0; + try (OffHeapLongGroupIdMap map = new OffHeapLongGroupIdMap(0)) { + for (long key : _longKeys) { + sum += map.getGroupId(key, GROUP_ID_UPPER_BOUND); + } + } + return sum; + } + + @Benchmark + public long onHeapStringMap() { + Object2IntOpenHashMap map = new Object2IntOpenHashMap<>(); + map.defaultReturnValue(GroupKeyGenerator.INVALID_ID); + long sum = 0; + for (String key : _stringKeys) { + int groupId = map.getInt(key); + if (groupId == GroupKeyGenerator.INVALID_ID) { + groupId = map.size(); + map.put(key, groupId); + } + sum += groupId; + } + return sum; + } + + @Benchmark + public long offHeapBytesMap() { + long sum = 0; + byte[] scratch = _encodeScratch; + try (OffHeapBytesGroupIdMap map = new OffHeapBytesGroupIdMap(0)) { + for (String key : _stringKeys) { + int length = OffHeapGroupByUtils.encodeUtf8(key, scratch); + sum += map.getGroupId(scratch, 0, length, GROUP_ID_UPPER_BOUND); + } + } + return sum; + } + + public static void main(String[] args) + throws Exception { + new Runner( + new OptionsBuilder().include(BenchmarkOffHeapGroupIdMaps.class.getSimpleName()).addProfiler("gc").build()) + .run(); + } +} diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/OffHeapGroupByMemoryFootprint.java b/pinot-perf/src/main/java/org/apache/pinot/perf/OffHeapGroupByMemoryFootprint.java new file mode 100644 index 000000000000..b954d3ba4faf --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/OffHeapGroupByMemoryFootprint.java @@ -0,0 +1,301 @@ +/** + * 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.perf; + +import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap; +import java.lang.management.GarbageCollectorMXBean; +import java.lang.management.ManagementFactory; +import java.lang.management.MemoryMXBean; +import java.util.ArrayList; +import java.util.List; +import org.apache.pinot.core.query.aggregation.groupby.DictionaryBasedGroupKeyGenerator.IntGroupIdMap; +import org.apache.pinot.core.query.aggregation.groupby.DoubleGroupByResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapBytesGroupIdMap; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapDoubleGroupByResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapGroupByUtils; +import org.apache.pinot.core.query.aggregation.groupby.offheap.OffHeapIntGroupIdMap; +import org.apache.pinot.segment.spi.memory.PinotDataBuffer; + + +/// Deterministic memory-footprint and throughput measurement (not a JMH benchmark): builds the per-segment +/// group-by state — one key table plus two result holders, mirroring a `GROUP BY k ... COUNT(*), SUM(m)` — +/// at the requested group counts, forces GC, and reports: +///
    +///
  • retained JVM heap vs direct (off-heap) memory while the state is live (the bytes `groupByOffHeap` +/// moves off the heap),
    • +///
  • build time (one insert per distinct key), lookup time (a full second all-hits pass), and the GC time +/// accumulated during the build.
    • +///
+/// Keys are generated on the fly, so only state genuinely retained by the structures is measured (the on-heap +/// string map retains the String keys, exactly like the on-heap no-dictionary generator does). +/// +/// Usage: {@code java -XmxSIZE -XX:MaxDirectMemorySize=SIZE -cp benchmarks.jar +/// org.apache.pinot.perf.OffHeapGroupByMemoryFootprint [int|string|all] [count...]} +/// (defaults: all 100_000 1_000_000 4_000_000). At 100M groups use ~12GB heap for the int tier and ~20GB heap for +/// the on-heap string tier (100M retained Strings), with MaxDirectMemorySize of at least 10GB. +public final class OffHeapGroupByMemoryFootprint { + private OffHeapGroupByMemoryFootprint() { + } + + private static final MemoryMXBean MEMORY_MX_BEAN = ManagementFactory.getMemoryMXBean(); + + public static void main(String[] args) + throws Exception { + String tier = args.length > 0 ? args[0] : "all"; + int[] groupCounts; + if (args.length > 1) { + groupCounts = new int[args.length - 1]; + for (int i = 1; i < args.length; i++) { + groupCounts[i - 1] = Integer.parseInt(args[i].replace("_", "")); + } + } else { + groupCounts = new int[]{100_000, 1_000_000, 4_000_000}; + } + System.out.printf("%-24s %12s %11s %11s %10s %10s %12s%n", + "configuration", "numGroups", "heap MB", "direct MB", "build ms", "lookup ms", "gc build ms"); + for (int numGroups : groupCounts) { + if (!"string".equals(tier)) { + measure("onHeap int tier", numGroups, new OnHeapIntState(numGroups)); + measure("offHeap int tier", numGroups, new OffHeapIntState(numGroups)); + } + if (!"int".equals(tier)) { + measure("onHeap string tier", numGroups, new OnHeapStringState(numGroups)); + measure("offHeap string tier", numGroups, new OffHeapStringState(numGroups)); + } + System.out.println(); + } + } + + /// One measured configuration: build inserts every distinct key once, lookup re-resolves every key (all hits). + private abstract static class TierState { + final int _numGroups; + final List _closeables = new ArrayList<>(); + + TierState(int numGroups) { + _numGroups = numGroups; + } + + abstract void build(); + + abstract void lookup(); + + void close() + throws Exception { + for (AutoCloseable closeable : _closeables) { + closeable.close(); + } + _closeables.clear(); + } + + GroupByResultHolder buildHolder(boolean offHeap) { + GroupByResultHolder holder = offHeap + ? new OffHeapDoubleGroupByResultHolder(Math.min(_numGroups, 10_000), _numGroups, 0.0) + : new DoubleGroupByResultHolder(Math.min(_numGroups, 10_000), _numGroups, 0.0); + if (holder instanceof AutoCloseable) { + _closeables.add((AutoCloseable) holder); + } + holder.ensureCapacity(_numGroups); + return holder; + } + } + + private static final class OnHeapIntState extends TierState { + private IntGroupIdMap _map; + private GroupByResultHolder _holder1; + private GroupByResultHolder _holder2; + + OnHeapIntState(int numGroups) { + super(numGroups); + } + + @Override + void build() { + _map = new IntGroupIdMap(); + for (int i = 0; i < _numGroups; i++) { + _map.getGroupId(i * 31, Integer.MAX_VALUE); + } + _holder1 = buildHolder(false); + _holder2 = buildHolder(false); + } + + @Override + void lookup() { + long sum = 0; + for (int i = 0; i < _numGroups; i++) { + sum += _map.getGroupId(i * 31, Integer.MAX_VALUE); + } + consume(sum); + } + } + + private static final class OffHeapIntState extends TierState { + private OffHeapIntGroupIdMap _map; + private GroupByResultHolder _holder1; + private GroupByResultHolder _holder2; + + OffHeapIntState(int numGroups) { + super(numGroups); + } + + @Override + void build() { + _map = new OffHeapIntGroupIdMap(0); + _closeables.add(_map); + for (int i = 0; i < _numGroups; i++) { + _map.getGroupId(i * 31, Integer.MAX_VALUE); + } + _holder1 = buildHolder(true); + _holder2 = buildHolder(true); + } + + @Override + void lookup() { + long sum = 0; + for (int i = 0; i < _numGroups; i++) { + sum += _map.getGroupId(i * 31, Integer.MAX_VALUE); + } + consume(sum); + } + } + + private static final class OnHeapStringState extends TierState { + private Object2IntOpenHashMap _map; + private GroupByResultHolder _holder1; + private GroupByResultHolder _holder2; + + OnHeapStringState(int numGroups) { + super(numGroups); + } + + @Override + void build() { + _map = new Object2IntOpenHashMap<>(); + _map.defaultReturnValue(-1); + for (int i = 0; i < _numGroups; i++) { + // The map retains the String keys, exactly like the on-heap no-dictionary generator does + _map.putIfAbsent(makeKey(i), _map.size()); + } + _holder1 = buildHolder(false); + _holder2 = buildHolder(false); + } + + @Override + void lookup() { + long sum = 0; + for (int i = 0; i < _numGroups; i++) { + // Fresh String per lookup, mirroring per-block string materialization + sum += _map.getInt(makeKey(i)); + } + consume(sum); + } + } + + private static final class OffHeapStringState extends TierState { + private OffHeapBytesGroupIdMap _map; + private final byte[] _scratch = new byte[128]; + private GroupByResultHolder _holder1; + private GroupByResultHolder _holder2; + + OffHeapStringState(int numGroups) { + super(numGroups); + } + + @Override + void build() { + _map = new OffHeapBytesGroupIdMap(0); + _closeables.add(_map); + for (int i = 0; i < _numGroups; i++) { + int length = OffHeapGroupByUtils.encodeUtf8(makeKey(i), _scratch); + _map.getGroupId(_scratch, 0, length, Integer.MAX_VALUE); + } + _holder1 = buildHolder(true); + _holder2 = buildHolder(true); + } + + @Override + void lookup() { + long sum = 0; + for (int i = 0; i < _numGroups; i++) { + int length = OffHeapGroupByUtils.encodeUtf8(makeKey(i), _scratch); + sum += _map.getGroupId(_scratch, 0, length, Integer.MAX_VALUE); + } + consume(sum); + } + } + + /// Fast distinct-key builder (String.format is ~1us/call, far too slow for 100M keys): "key-" + 9 digits + + /// "-abcdefgh", 22 chars, a fresh String per call like per-block string materialization. + private static String makeKey(int i) { + char[] chars = {'k', 'e', 'y', '-', '0', '0', '0', '0', '0', '0', '0', '0', '0', '-', + 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'}; + int value = i; + for (int position = 12; position >= 4 && value > 0; position--) { + chars[position] = (char) ('0' + (value % 10)); + value /= 10; + } + return new String(chars); + } + + private static volatile long _sink; + + private static void consume(long value) { + _sink = value; + } + + private static void measure(String label, int numGroups, TierState state) + throws Exception { + forceGc(); + long heapBefore = MEMORY_MX_BEAN.getHeapMemoryUsage().getUsed(); + long directBefore = PinotDataBuffer.getDirectBufferUsage(); + long gcBefore = totalGcTimeMs(); + + long buildStartNs = System.nanoTime(); + state.build(); + long buildMs = (System.nanoTime() - buildStartNs) / 1_000_000; + long gcBuildMs = totalGcTimeMs() - gcBefore; + + long lookupStartNs = System.nanoTime(); + state.lookup(); + long lookupMs = (System.nanoTime() - lookupStartNs) / 1_000_000; + + forceGc(); + double heapMb = (MEMORY_MX_BEAN.getHeapMemoryUsage().getUsed() - heapBefore) / 1048576.0; + double directMb = (PinotDataBuffer.getDirectBufferUsage() - directBefore) / 1048576.0; + System.out.printf("%-24s %,12d %11.1f %11.1f %,10d %,10d %,12d%n", + label, numGroups, heapMb, directMb, buildMs, lookupMs, gcBuildMs); + state.close(); + } + + private static long totalGcTimeMs() { + long total = 0; + for (GarbageCollectorMXBean gcBean : ManagementFactory.getGarbageCollectorMXBeans()) { + total += Math.max(0, gcBean.getCollectionTime()); + } + return total; + } + + private static void forceGc() + throws InterruptedException { + for (int i = 0; i < 3; i++) { + System.gc(); + Thread.sleep(100); + } + } +} diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index 4b31ac6e3463..05150b183748 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -1025,6 +1025,8 @@ public static class QueryOptionKey { public static final String NUM_GROUPS_LIMIT = "numGroupsLimit"; // Not actually accepted as Query Option but faked as one during MSE public static final String NUM_GROUPS_WARNING_LIMIT = "numGroupsWarningLimit"; + /// Store SSE group-by key tables and fixed-width result holders in off-heap (direct) memory. + public static final String GROUP_BY_OFF_HEAP = "groupByOffHeap"; public static final String MAX_INITIAL_RESULT_HOLDER_CAPACITY = "maxInitialResultHolderCapacity"; public static final String MIN_INITIAL_INDEXED_TABLE_CAPACITY = "minInitialIndexedTableCapacity"; public static final String MSE_MAX_INITIAL_RESULT_HOLDER_CAPACITY = "mseMaxInitialResultHolderCapacity"; @@ -1639,6 +1641,28 @@ public enum AndRestrictionPushdownMode { public static final String CONFIG_OF_QUERY_EXECUTOR_GROUPBY_TRIM_THRESHOLD = QUERY_EXECUTOR_CONFIG_PREFIX + "." + GROUPBY_TRIM_THRESHOLD; public static final int DEFAULT_QUERY_EXECUTOR_GROUPBY_TRIM_THRESHOLD = 1_000_000; + // Store SSE group-by key tables and fixed-width aggregation result holders in off-heap (direct) memory. + // NOTE: Off-heap group-by memory draws from -XX:MaxDirectMemorySize (shared with segment buffers and Netty) + // and is not yet visible to the per-query resource accountant or bounded by a per-query byte budget — the + // only bound is numGroupsLimit per table. Size direct memory accordingly before enabling. The per-query + // option can override this config in either direction (matching the numGroupsLimit precedent). + // On the streaming combine (MSE leaf stages), each per-segment result is materialized into on-heap records + // at hand-off and the off-heap state is released immediately, so there the mode only relieves the + // segment-execution phase itself — it does not reduce the heap footprint of the streamed results. + public static final String GROUPBY_OFF_HEAP = "groupby.offheap"; + public static final String CONFIG_OF_QUERY_EXECUTOR_GROUPBY_OFF_HEAP = + QUERY_EXECUTOR_CONFIG_PREFIX + "." + GROUPBY_OFF_HEAP; + public static final boolean DEFAULT_QUERY_EXECUTOR_GROUPBY_OFF_HEAP = false; + // Per-thread cap on direct buffers pooled for reuse across queries by the off-heap group-by structures + // (mirrors the on-heap thread-local map caching); 0 disables pooling. Pooled bytes stay visible in the + // direct-buffer usage accounting; the aggregate retention bound is this cap times the number of threads that + // release group-by buffers (combine workers plus reduce threads). Changing the config to 0 drains each + // thread's retained buffers lazily on its next group-by. TODO: export the pooled bytes as a server gauge. + public static final String GROUPBY_OFF_HEAP_POOL_MAX_BYTES_PER_THREAD = + "groupby.offheap.pool.max.bytes.per.thread"; + public static final String CONFIG_OF_QUERY_EXECUTOR_GROUPBY_OFF_HEAP_POOL_MAX_BYTES_PER_THREAD = + QUERY_EXECUTOR_CONFIG_PREFIX + "." + GROUPBY_OFF_HEAP_POOL_MAX_BYTES_PER_THREAD; + public static final long DEFAULT_QUERY_EXECUTOR_GROUPBY_OFF_HEAP_POOL_MAX_BYTES_PER_THREAD = 0; // Do sort-aggregation when LIMIT is below this threshold public static final int DEFAULT_SORT_AGGREGATE_LIMIT_THRESHOLD = 10_000; // Use sequential instead of pair-wise combine for sort-aggr when numSegments is below this threshold