From dd34a26b0a4cef8e6f34fcfded2eb130b9565e21 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Fri, 28 Aug 2026 01:25:48 -0700 Subject: [PATCH 1/3] Count the out-of-map null group in NoDictionarySingleColumnGroupKeyGenerator group counts For primitive stored types with null handling enabled, the null group lives outside the primitive key map but still takes the next dense group id, so getNumKeys() and getCurrentGroupKeyUpperBound() under-counted by one once a null was seen. Since DefaultGroupByExecutor sizes result holders with ensureCapacity(getCurrentGroupKeyUpperBound()), a segment whose group count exceeds the initial holder capacity then wrote one slot past the holder array (ArrayIndexOutOfBoundsException) on the default on-heap path. Object stored types are unaffected (their null key lives inside the map). Adds NoDictionaryNullGroupCountRegressionTest reproducing the AIOOBE through DefaultGroupByExecutor.process() with a shrunk maxInitialResultHolderCapacity for all four primitive stored types, and pinning the counts and the null-group emission. --- .../pinot/core/data/table/TableResizer.java | 2 +- ...ictionaryNullGroupCountRegressionTest.java | 187 ++++++++++++++++++ 2 files changed, 188 insertions(+), 1 deletion(-) create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryNullGroupCountRegressionTest.java 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/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); + } +} From 1347b46739253053dc8ea7a1a74450d0ac7b10e0 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Fri, 28 Aug 2026 01:26:12 -0700 Subject: [PATCH 2/3] Support off-heap group-by key tables and result holders in SSE (config-gated, default off) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an opt-in off-heap storage mode for the SSE per-segment group-by state, targeting high-cardinality group-bys whose on-heap key maps and result holders drive GC pressure: - New pinot-core package o.a.p.core.query.aggregation.groupby.offheap: - OffHeapIntGroupIdMap / OffHeapLongGroupIdMap: open-addressing key->dense-id tables over direct memory (8/16-byte slots, load factor 0.5, linear probing, out-of-band -1/0 key), drop-in replacements for IntGroupIdMap / Long2IntOpenHashMap semantics. - OffHeapBytesGroupIdMap: DuckDB-style two-part table for var-width keys — an 8-byte-entry directory (16-bit salt | 48-bit payload offset) over append-only 256KB payload chunks storing [hash][groupId][keyLength][key bytes]; the stored hash makes directory resize free of key reads. - OffHeapDouble/Long/IntGroupByResultHolder: fixed-width result holders over direct memory with semantics identical to the on-heap holders. - ResourceTrackingGroupKeyGenerator: wraps the generator and owns every off-heap resource, so the existing generator close() call sites release all direct memory (including the shared-generator filtered-aggregation case). - OffHeapGroupByBufferPool: bounded per-thread buffer reuse across queries (mirrors the on-heap thread-local map caching, with an explicit cap and visible accounting), default off. - All hot paths use absolute-indexed direct ByteBuffer views (wrapper fallback beyond 2GB). - Wiring: server config pinot.server.query.executor.groupby.offheap (default false), query option groupByOffHeap, pool cap config groupby.offheap.pool.max.bytes.per.thread (default 0). Off-heap RawKeyHolder variants in DictionaryBasedGroupKeyGenerator (the ARRAY_BASED tier stays on-heap); off-heap modes in both NoDictionary generators; holder mirroring in DefaultGroupByExecutor. Grouping sets stay on-heap. Group ids remain dense ints; no AggregationFunction changes; no wire or storage format changes. - Close-path hardening (also fixes pre-existing on-heap leak windows): exception guards in GroupByOperator/FilteredGroupByOperator/DefaultGroupByExecutor and widened finally coverage in the group-by combine operator. The streaming combine needs no extra plumbing: since #19066 each per-segment result is detached and its generator closed on the producing worker thread, which releases the off-heap state promptly as well. - Tests: differential suites comparing off-heap vs on-heap row-for-row (OffHeapGroupByQueriesTest end-to-end battery with per-query direct-memory leak assertions, OffHeapGroupKeyGeneratorParityTest at generator level incl. null-group id bookkeeping), per-structure unit tests incl. forced wrapper-fallback runs, and buffer pool tests. - Benchmarks (pinot-perf): BenchmarkOffHeapGroupBySSE / -LargeSSE / -HugeSSE and OffHeapGroupByMemoryFootprint. Measured: retained heap for the per-segment state drops to ~0 (e.g. 8.7GB -> 4MB at 100M string groups, with 3.7x faster build); at ~1M+ groups off-heap is faster end-to-end (up to -50%) because it removes the GC pressure that dominates on-heap; at ~80K groups (cache-resident) there is a 10-22% latency premium, which the per-query opt-in avoids. --- .../utils/config/QueryOptionsUtils.java | 6 + .../combine/GroupByCombineOperator.java | 73 +- .../query/FilteredGroupByOperator.java | 18 + .../core/operator/query/GroupByOperator.java | 13 + .../plan/maker/InstancePlanMakerImplV2.java | 20 +- .../groupby/DefaultGroupByExecutor.java | 89 ++- .../DictionaryBasedGroupKeyGenerator.java | 278 ++++++- .../groupby/DoubleGroupByResultHolder.java | 5 + .../groupby/IntGroupByResultHolder.java | 5 + .../groupby/LongGroupByResultHolder.java | 5 + ...ictionaryMultiColumnGroupKeyGenerator.java | 78 +- ...ctionarySingleColumnGroupKeyGenerator.java | 301 +++++++- .../offheap/OffHeapBytesGroupIdMap.java | 521 +++++++++++++ .../OffHeapDoubleGroupByResultHolder.java | 196 +++++ .../offheap/OffHeapGroupByBufferPool.java | 159 ++++ .../groupby/offheap/OffHeapGroupByUtils.java | 122 +++ .../OffHeapIntGroupByResultHolder.java | 196 +++++ .../groupby/offheap/OffHeapIntGroupIdMap.java | 310 ++++++++ .../OffHeapLongGroupByResultHolder.java | 196 +++++ .../offheap/OffHeapLongGroupIdMap.java | 311 ++++++++ .../ResourceTrackingGroupKeyGenerator.java | 109 +++ .../query/request/context/QueryContext.java | 10 + .../StreamingGroupByCombineOperatorTest.java | 89 +++ .../OffHeapGroupKeyGeneratorParityTest.java | 645 ++++++++++++++++ .../offheap/OffHeapBytesGroupIdMapTest.java | 286 +++++++ .../offheap/OffHeapGroupByBufferPoolTest.java | 109 +++ .../OffHeapGroupByResultHolderTest.java | 383 ++++++++++ .../offheap/OffHeapGroupByUtilsTest.java | 145 ++++ .../offheap/OffHeapIntGroupIdMapTest.java | 170 +++++ .../offheap/OffHeapLongGroupIdMapTest.java | 263 +++++++ .../queries/OffHeapGroupByQueriesTest.java | 706 ++++++++++++++++++ .../perf/BenchmarkOffHeapGroupByHugeSSE.java | 213 ++++++ .../perf/BenchmarkOffHeapGroupByLargeSSE.java | 278 +++++++ .../perf/BenchmarkOffHeapGroupBySSE.java | 255 +++++++ .../perf/BenchmarkOffHeapGroupIdMaps.java | 175 +++++ .../perf/OffHeapGroupByMemoryFootprint.java | 301 ++++++++ .../pinot/spi/utils/CommonConstants.java | 24 + 37 files changed, 6991 insertions(+), 72 deletions(-) create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapBytesGroupIdMap.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapDoubleGroupByResultHolder.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByBufferPool.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByUtils.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapIntGroupByResultHolder.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapIntGroupIdMap.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapLongGroupByResultHolder.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapLongGroupIdMap.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/ResourceTrackingGroupKeyGenerator.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/OffHeapGroupKeyGeneratorParityTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapBytesGroupIdMapTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByBufferPoolTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByResultHolderTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapGroupByUtilsTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapIntGroupIdMapTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapLongGroupIdMapTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/queries/OffHeapGroupByQueriesTest.java create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupByHugeSSE.java create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupByLargeSSE.java create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupBySSE.java create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupIdMaps.java create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/OffHeapGroupByMemoryFootprint.java 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/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/groupby/DefaultGroupByExecutor.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java index 8a4015cd84da..f62098863373 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 @@ -37,6 +37,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 +101,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 +170,31 @@ public DefaultGroupByExecutor(QueryContext queryContext, AggregationFunction[] a } } + /// Mirrors fixed-width result holders off-heap. 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 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/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/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/queries/OffHeapGroupByQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/OffHeapGroupByQueriesTest.java new file mode 100644 index 000000000000..9701acef4847 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/queries/OffHeapGroupByQueriesTest.java @@ -0,0 +1,706 @@ +/** + * 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 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.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.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.query.QueryThreadContext; +import org.apache.pinot.spi.utils.BytesUtils; +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"; + + private IndexSegment _mainSegment; + private IndexSegment _nullSegment; + 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(); + _mainSegment = ImmutableSegmentLoader.load(new File(INDEX_DIR, SEGMENT_NAME), ReadMode.mmap); + _nullSegment = ImmutableSegmentLoader.load(new File(INDEX_DIR, NULL_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(); + 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 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; + } + + 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"); + } + } + + @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/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 From 0635c9152dfa5327fbc5ca0be857200034c8f2a3 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Mon, 31 Aug 2026 00:01:33 -0700 Subject: [PATCH 3/3] Support off-heap group-by state for DISTINCTCOUNTULL (config-gated, default off) Adds the first per-function off-heap aggregation state on top of the off-heap group-by SSE feature: DISTINCTCOUNTULL/DISTINCTCOUNTRAWULL keep each group's UltraLogLog register array (2^p bytes, ~4.1KB at the default p=12) in pooled direct-memory chunks instead of one heap object per group. - New optional AggregationFunction#createOffHeapGroupByResultHolder seam (default null = unchanged); DefaultGroupByExecutor consults it first inside the existing off-heap gate and registers the holder on the resource tracker. - OffHeapUltraLogLogGroupByResultHolder: append-only slotId indirection so direct memory grows with actual groups (not the group-count upper bound), 256KB pooled chunks, hash4j 0.30.0 register math vendored verbatim (pinned byte-identical by test), lazy on-heap delegate for the dictionary and pre-serialized-BYTES modes (mode exclusivity enforced hard), untouched groups read back as null, snapshot materialization at extraction. - Validates the user-supplied p literal in [3, 26] at plan time (previously only UltraLogLog.create checked it; the off-heap holder sizes slots as 1 << p, so an unchecked p could over-allocate or corrupt neighbor slots). A bad p now fails planning even when a filter matches zero rows. - Tests: differential holder test vs hash4j across p=3/8/12/18/19 including the buffer-wrapper fallback arm; e2e battery in OffHeapGroupByQueriesTest (raw/dict/MV/explicit-p/RAWULL/filtered/order-by-trim/null-handling, plus a star-tree segment proving the pre-aggregated BYTES path in both modes) with per-query direct-memory leak asserts. - Benchmark BenchmarkOffHeapGroupByUllSSE (10K/200K groups x flag): off-heap segment phase -12%/-32% latency, full query -11%/-13%, with the per-group sketch heap (~42MB/840MB per segment execution) moved off the heap. --- .../function/AggregationFunction.java | 12 + .../DistinctCountULLAggregationFunction.java | 118 ++++--- .../groupby/DefaultGroupByExecutor.java | 13 +- ...OffHeapUltraLogLogGroupByResultHolder.java | 319 ++++++++++++++++++ ...stinctCountULLAggregationFunctionTest.java | 12 + ...eapUltraLogLogGroupByResultHolderTest.java | 178 ++++++++++ .../queries/OffHeapGroupByQueriesTest.java | 123 +++++++ .../perf/BenchmarkOffHeapGroupByUllSSE.java | 273 +++++++++++++++ 8 files changed, 1008 insertions(+), 40 deletions(-) create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapUltraLogLogGroupByResultHolder.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/offheap/OffHeapUltraLogLogGroupByResultHolderTest.java create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOffHeapGroupByUllSSE.java 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 f62098863373..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; @@ -170,13 +171,23 @@ public DefaultGroupByExecutor(QueryContext queryContext, AggregationFunction[] a } } - /// Mirrors fixed-width result holders off-heap. The holder type and default value are discovered through 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) { 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/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/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 index 9701acef4847..82c22ee6cb1a 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/OffHeapGroupByQueriesTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/OffHeapGroupByQueriesTest.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.queries; +import com.dynatrace.hash4j.distinctcount.UltraLogLog; import java.io.File; import java.math.BigDecimal; import java.nio.ByteBuffer; @@ -31,6 +32,7 @@ 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; @@ -41,10 +43,12 @@ 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; @@ -52,6 +56,7 @@ 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; @@ -130,8 +135,17 @@ public class OffHeapGroupByQueriesTest extends BaseQueriesTest { 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; @@ -157,8 +171,10 @@ public void setUp() 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 @@ -174,6 +190,7 @@ public void tearDown() throws Exception { _mainSegment.destroy(); _nullSegment.destroy(); + _starTreeSegment.destroy(); FileUtils.deleteDirectory(INDEX_DIR); } @@ -187,6 +204,11 @@ private void useNullSegment() { _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) @@ -254,6 +276,39 @@ private static Object[] randomMvValues(Random random, int cardinality, String st 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); @@ -454,6 +509,74 @@ public void testSingleColumnGroupBy() { } } + /// 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 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(); + } +}