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 33072d5fd4..b68cf99a9a 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 @@ -143,6 +143,9 @@ protected GroupByResultsBlock getNextBlock() { if (groupByExecutor.getNumGroups() > trimSize) { TableResizer tableResizer = new TableResizer(_dataSchema, _queryContext); Collection intermediateRecords = groupByExecutor.trimGroupByResult(trimSize, tableResizer); + // intermediateRecords no longer references the holders; safe to return their arrays + // to the ThreadLocal cache for the next query. + groupByExecutor.releaseHolders(); GroupByResultsBlock resultsBlock = new GroupByResultsBlock(_dataSchema, intermediateRecords, _queryContext); resultsBlock.setNumGroupsLimitReached(numGroupsLimitReached); resultsBlock.setNumGroupsWarningLimitReached(numGroupsWarningLimitReached); 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 1881bce800..e6ff1ea74d 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 @@ -241,4 +241,11 @@ public GroupKeyGenerator getGroupKeyGenerator() { public GroupByResultHolder[] getGroupByResultHolders() { return _groupByResultHolders; } + + @Override + public void releaseHolders() { + for (GroupByResultHolder holder : _groupByResultHolders) { + holder.release(); + } + } } 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 60e1b8bb47..5f06c40643 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 @@ -19,6 +19,7 @@ package org.apache.pinot.core.query.aggregation.groupby; import com.google.common.base.Preconditions; +import java.lang.ref.SoftReference; import java.util.Arrays; @@ -26,6 +27,11 @@ * Result Holder implemented using DoubleArray. */ public class DoubleGroupByResultHolder implements GroupByResultHolder { + // Don't cache arrays above this size; outlier queries let them GC normally + // instead of pinning per-worker-thread allocations forever. + private static final int MAX_CACHEABLE_CAPACITY = 1_000_000; + private static final ThreadLocal> THREAD_LOCAL_ARRAY = new ThreadLocal<>(); + private final int _maxCapacity; private final double _defaultValue; @@ -55,22 +61,51 @@ public void ensureCapacity(int capacity) { Preconditions.checkArgument(capacity <= _maxCapacity); if (capacity > _resultHolderCapacity) { + SoftReference ref = THREAD_LOCAL_ARRAY.get(); + double[] cached = (ref != null) ? ref.get() : null; + // Don't use the ThreadLocal array if it's smaller than the capacity + if (cached != null && cached.length < capacity) { + cached = null; + } int copyLength = _resultHolderCapacity; - _resultHolderCapacity = Math.max(_resultHolderCapacity * 2, capacity); + if (cached != null) { + _resultHolderCapacity = cached.length; + } else { + _resultHolderCapacity = Math.max(_resultHolderCapacity * 2, capacity); - // Cap the growth to maximum possible number of group keys - _resultHolderCapacity = Math.min(_resultHolderCapacity, _maxCapacity); + // Cap the growth to maximum possible number of group keys + _resultHolderCapacity = Math.min(_resultHolderCapacity, _maxCapacity); + } double[] current = _resultArray; - _resultArray = new double[_resultHolderCapacity]; + if (cached != null) { + _resultArray = cached; + } else { + _resultArray = new double[_resultHolderCapacity]; + } + System.arraycopy(current, 0, _resultArray, 0, copyLength); - if (_defaultValue != 0.0) { + // Reused arrays may contain stale data beyond copyLength; fresh arrays only need + // filling when the default value is non-zero. + if (cached != null || _defaultValue != 0.0) { Arrays.fill(_resultArray, copyLength, _resultHolderCapacity, _defaultValue); } } } + @Override + public void release() { + if (_resultArray.length <= MAX_CACHEABLE_CAPACITY) { + SoftReference ref = THREAD_LOCAL_ARRAY.get(); + double[] cached = (ref != null) ? ref.get() : null; + if (cached == null || _resultArray.length > cached.length) { + THREAD_LOCAL_ARRAY.set(new SoftReference<>(_resultArray)); + } + } + _resultArray = null; + } + @Override public double getDoubleResult(int groupKey) { if (groupKey == GroupKeyGenerator.INVALID_ID) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DummyGroupByResultHolder.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DummyGroupByResultHolder.java index f76e8967cf..aa00cd7bd7 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DummyGroupByResultHolder.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DummyGroupByResultHolder.java @@ -53,4 +53,8 @@ public T getResult(int groupKey) { @Override public void ensureCapacity(int capacity) { } + + @Override + public void release() { + } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupByExecutor.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupByExecutor.java index 8c5524ec45..6ae88d8738 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupByExecutor.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupByExecutor.java @@ -60,4 +60,14 @@ public interface GroupByExecutor { GroupKeyGenerator getGroupKeyGenerator(); GroupByResultHolder[] getGroupByResultHolders(); + + /** + * Releases the per-aggregation result holders, allowing each holder to return its backing + * array to the ThreadLocal cache for reuse by the next query. + *

Must only be called once the holders' data is no longer needed (i.e. after + * {@link #trimGroupByResult} or after the {@link AggregationGroupByResult} returned by + * {@link #getResult} has been fully consumed). Accessing the holders after this call is + * undefined. + */ + void releaseHolders(); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupByResultHolder.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupByResultHolder.java index 64017ec97a..69877984f1 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupByResultHolder.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupByResultHolder.java @@ -86,4 +86,9 @@ public interface GroupByResultHolder { * @param capacity */ void ensureCapacity(int capacity); + + /** + * Reset the threadlocal array. + */ + void release(); } 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 3ae29f4580..50962a282e 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 @@ -20,17 +20,22 @@ package org.apache.pinot.core.query.aggregation.groupby; import com.google.common.base.Preconditions; +import java.lang.ref.SoftReference; import java.util.Arrays; public class IntGroupByResultHolder implements GroupByResultHolder { + // Don't cache arrays above this size; outlier queries let them GC normally + // instead of pinning per-worker-thread allocations forever. + private static final int MAX_CACHEABLE_CAPACITY = 1_000_000; + private static final ThreadLocal> THREAD_LOCAL_ARRAY = new ThreadLocal<>(); + private final int _maxCapacity; private final int _defaultValue; private int _resultHolderCapacity; private int[] _resultArray; - /** * Constructor for the class. * @@ -54,22 +59,51 @@ public void ensureCapacity(int capacity) { Preconditions.checkArgument(capacity <= _maxCapacity); if (capacity > _resultHolderCapacity) { + SoftReference ref = THREAD_LOCAL_ARRAY.get(); + int[] cached = (ref != null) ? ref.get() : null; + // Don't use the ThreadLocal array if it's smaller than the capacity + if (cached != null && cached.length < capacity) { + cached = null; + } + int copyLength = _resultHolderCapacity; - _resultHolderCapacity = Math.max(_resultHolderCapacity * 2, capacity); + if (cached != null) { + _resultHolderCapacity = cached.length; + } else { + _resultHolderCapacity = Math.max(_resultHolderCapacity * 2, capacity); - // Cap the growth to maximum possible number of group keys - _resultHolderCapacity = Math.min(_resultHolderCapacity, _maxCapacity); + // Cap the growth to maximum possible number of group keys + _resultHolderCapacity = Math.min(_resultHolderCapacity, _maxCapacity); + } int[] current = _resultArray; - _resultArray = new int[_resultHolderCapacity]; + if (cached != null) { + _resultArray = cached; + } else { + _resultArray = new int[_resultHolderCapacity]; + } System.arraycopy(current, 0, _resultArray, 0, copyLength); - if (_defaultValue != 0) { + // Reused arrays may contain stale data beyond copyLength; fresh arrays only need + // filling when the default value is non-zero. + if (cached != null || _defaultValue != 0) { Arrays.fill(_resultArray, copyLength, _resultHolderCapacity, _defaultValue); } } } + @Override + public void release() { + if (_resultArray.length <= MAX_CACHEABLE_CAPACITY) { + SoftReference ref = THREAD_LOCAL_ARRAY.get(); + int[] cached = (ref != null) ? ref.get() : null; + if (cached == null || _resultArray.length > cached.length) { + THREAD_LOCAL_ARRAY.set(new SoftReference<>(_resultArray)); + } + } + _resultArray = null; + } + @Override public double getDoubleResult(int groupKey) { throw new UnsupportedOperationException(); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/ObjectGroupByResultHolder.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/ObjectGroupByResultHolder.java index ebf5ec7b60..db990e8287 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/ObjectGroupByResultHolder.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/ObjectGroupByResultHolder.java @@ -19,12 +19,19 @@ package org.apache.pinot.core.query.aggregation.groupby; import com.google.common.base.Preconditions; +import java.lang.ref.SoftReference; +import java.util.Arrays; /** * Result Holder implemented using ObjectArray. */ public class ObjectGroupByResultHolder implements GroupByResultHolder { + // Don't cache arrays above this size; outlier queries let them GC normally + // instead of pinning per-worker-thread allocations forever. + private static final int MAX_CACHEABLE_CAPACITY = 1_000_000; + private static final ThreadLocal> THREAD_LOCAL_ARRAY = new ThreadLocal<>(); + private final int _maxCapacity; private int _resultHolderCapacity; @@ -43,21 +50,51 @@ public ObjectGroupByResultHolder(int initialCapacity, int maxCapacity) { _resultArray = new Object[initialCapacity]; } + @Override public void ensureCapacity(int capacity) { Preconditions.checkArgument(capacity <= _maxCapacity); if (capacity > _resultHolderCapacity) { + SoftReference ref = THREAD_LOCAL_ARRAY.get(); + Object[] cached = (ref != null) ? ref.get() : null; + if (cached != null && cached.length < capacity) { + cached = null; + } int copyLength = _resultHolderCapacity; - _resultHolderCapacity = Math.max(_resultHolderCapacity * 2, capacity); + if (cached != null) { + _resultHolderCapacity = cached.length; + } else { + _resultHolderCapacity = Math.max(_resultHolderCapacity * 2, capacity); - // Cap the growth to maximum possible number of group keys - _resultHolderCapacity = Math.min(_resultHolderCapacity, _maxCapacity); + // Cap the growth to maximum possible number of group keys + _resultHolderCapacity = Math.min(_resultHolderCapacity, _maxCapacity); + } Object[] current = _resultArray; - _resultArray = new Object[_resultHolderCapacity]; + if (cached != null) { + _resultArray = cached; + } else { + _resultArray = new Object[_resultHolderCapacity]; + } System.arraycopy(current, 0, _resultArray, 0, copyLength); + // release() nulls slots before caching, so reused arrays start clean past copyLength. + } + } + + @Override + public void release() { + // Null slots so per-query aggregation results can be GC'd, regardless of whether + // we end up caching the array for reuse. + Arrays.fill(_resultArray, 0, _resultHolderCapacity, null); + if (_resultArray.length <= MAX_CACHEABLE_CAPACITY) { + SoftReference ref = THREAD_LOCAL_ARRAY.get(); + Object[] cached = (ref != null) ? ref.get() : null; + if (cached == null || _resultArray.length > cached.length) { + THREAD_LOCAL_ARRAY.set(new SoftReference<>(_resultArray)); + } } + _resultArray = null; } @Override diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/GroupByAllocationBenchmark.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/GroupByAllocationBenchmark.java new file mode 100644 index 0000000000..0a9f33016f --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/GroupByAllocationBenchmark.java @@ -0,0 +1,212 @@ +/** + * 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.lang.management.ManagementFactory; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock; +import org.apache.pinot.core.operator.combine.GroupByCombineOperator; +import org.apache.pinot.core.plan.GroupByPlanNode; +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.TableType; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.MetricFieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.utils.CommonConstants; +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; + + +/** + * Allocation benchmark for the ThreadLocal aggregate-holder cache. + * + * Runs a fixed sequence of high-cardinality group-by queries on a single worker thread and + * prints the aggregate bytes allocated across all threads (HotSpot's per-thread allocation + * counter). Diff master vs branch for the savings number. + * + * Defaults are sized for a first comparison (peak heap well under 1 GB, runs in seconds). + * For a louder signal, scale up: + * NUM_ROWS 1_000_000 -> 5_000_000 + * NUM_QUERIES 20 -> 100 + * That targets ~2-4 GB peak heap; run with -Xmx4g. + * + * Not part of the default test suite (class name doesn't end in "Test"). Invoke explicitly: + * ./mvnw -pl pinot-core -Dcheckstyle.skip=true \ + * -Dtest=GroupByAllocationBenchmark test + */ +public class GroupByAllocationBenchmark { + // ----- tune these for the run ----- + private static final int NUM_ROWS = 5_000_000; + private static final int NUM_QUERIES = 20; + private static final int WARMUP_QUERIES = 3; + // 10 rows per group -> 100k distinct group keys for 1M rows, well above the 5k trim threshold. + private static final int GROUP_CARDINALITY = NUM_ROWS / 10; + // ---------------------------------- + + private static final File INDEX_DIR = + new File(FileUtils.getTempDirectory(), "GroupByAllocationBenchmark"); + private static final String SEGMENT_NAME = "benchSegment"; + + // Double (MAX/SUM/MIN) and Int (COUNT) holders only. DISTINCTCOUNT was removed because its + // per-group state objects allocate per query regardless of the array cache, masking the win. + private static final String QUERY = + "SELECT group_col," + + " MAX(metric_a), SUM(metric_b), MIN(metric_c)," + + " COUNT(metric_d) " + + "FROM t GROUP BY group_col " + + "ORDER BY MAX(metric_a) DESC LIMIT 100"; + + // Single thread keeps the ThreadLocal cache state deterministic across queries. + private final ExecutorService _executor = Executors.newSingleThreadExecutor(); + private IndexSegment _segment; + + @BeforeClass + public void setUp() + throws Exception { + FileUtils.deleteQuietly(INDEX_DIR); + buildSegment(); + } + + @AfterClass + public void tearDown() { + if (_segment != null) { + _segment.destroy(); + } + _executor.shutdown(); + FileUtils.deleteQuietly(INDEX_DIR); + } + + @Test + public void benchmarkAllocation() + throws Exception { + com.sun.management.ThreadMXBean tmx = + (com.sun.management.ThreadMXBean) ManagementFactory.getThreadMXBean(); + tmx.setThreadAllocatedMemoryEnabled(true); + + // Warm up: populate the ThreadLocal cache and prime the JIT so the measured run reflects + // steady-state allocation, not first-query overhead. + long warmupStartNs = System.nanoTime(); + for (int i = 0; i < WARMUP_QUERIES; i++) { + runQuery(); + } + long warmupElapsedNs = System.nanoTime() - warmupStartNs; + + long before = sumAllThreadAllocations(tmx); + long measuredStartNs = System.nanoTime(); + for (int i = 0; i < NUM_QUERIES; i++) { + runQuery(); + } + long measuredElapsedNs = System.nanoTime() - measuredStartNs; + long after = sumAllThreadAllocations(tmx); + + long total = after - before; + double mb = total / (1024.0 * 1024.0); + double perQueryMb = mb / NUM_QUERIES; + double measuredMs = measuredElapsedNs / 1_000_000.0; + double perQueryMs = measuredMs / NUM_QUERIES; + double warmupPerQueryMs = (warmupElapsedNs / 1_000_000.0) / WARMUP_QUERIES; + System.out.printf("%n=== GroupByAllocationBenchmark ===%n"); + System.out.printf("rows=%d queries=%d warmup=%d groupCardinality=%d%n", + NUM_ROWS, NUM_QUERIES, WARMUP_QUERIES, GROUP_CARDINALITY); + System.out.printf("Aggregate bytes allocated across all threads: %,d (%,.1f MB)%n", + total, mb); + System.out.printf("Per-query average allocation: %,.1f MB%n", perQueryMb); + System.out.printf("Measured wall-clock time: %,.1f ms total, %,.1f ms/query%n", + measuredMs, perQueryMs); + System.out.printf("Warmup per-query wall-clock time: %,.1f ms (for reference)%n", + warmupPerQueryMs); + } + + private void runQuery() { + QueryContext ctx = QueryContextConverterUtils.getQueryContext(QUERY); + ctx.setEndTimeMs(System.currentTimeMillis() + CommonConstants.Server.DEFAULT_QUERY_EXECUTOR_TIMEOUT_MS); + // Make sure the trim path triggers (groupCardinality >> trimSize), so releaseHolders() runs + // and the ThreadLocal cache gets populated. + ctx.setMinSegmentGroupTrimSize(5000); + Operator op = + new GroupByPlanNode(new SegmentContext(_segment), ctx).run(); + GroupByCombineOperator combine = + new GroupByCombineOperator(Collections.singletonList(op), ctx, _executor); + combine.nextBlock(); + } + + private long sumAllThreadAllocations(com.sun.management.ThreadMXBean tmx) { + long total = 0; + for (long id : tmx.getAllThreadIds()) { + long bytes = tmx.getThreadAllocatedBytes(id); + if (bytes > 0) { + total += bytes; + } + } + return total; + } + + private void buildSegment() + throws Exception { + Schema schema = new Schema(); + schema.addField(new DimensionFieldSpec("group_col", FieldSpec.DataType.INT, true)); + schema.addField(new MetricFieldSpec("metric_a", FieldSpec.DataType.DOUBLE)); + schema.addField(new MetricFieldSpec("metric_b", FieldSpec.DataType.DOUBLE)); + schema.addField(new MetricFieldSpec("metric_c", FieldSpec.DataType.DOUBLE)); + schema.addField(new MetricFieldSpec("metric_d", FieldSpec.DataType.INT)); + schema.addField(new DimensionFieldSpec("metric_e", FieldSpec.DataType.STRING, true)); + + SegmentGeneratorConfig config = new SegmentGeneratorConfig( + new TableConfigBuilder(TableType.OFFLINE).setTableName("t").build(), schema); + config.setSegmentName(SEGMENT_NAME); + config.setOutDir(INDEX_DIR.getAbsolutePath()); + + List rows = new ArrayList<>(NUM_ROWS); + for (int i = 0; i < NUM_ROWS; i++) { + GenericRow row = new GenericRow(); + row.putValue("group_col", i % GROUP_CARDINALITY); + row.putValue("metric_a", (double) i); + row.putValue("metric_b", (double) (i * 2)); + row.putValue("metric_c", (double) (i + 7)); + row.putValue("metric_d", i % 1000); + row.putValue("metric_e", "v_" + (i % 50_000)); + rows.add(row); + } + + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(config, new GenericRowRecordReader(rows)); + driver.build(); + + _segment = ImmutableSegmentLoader.load(new File(INDEX_DIR, driver.getSegmentName()), ReadMode.heap); + } +}