wirybeaver commented on code in PR #19469: URL: https://github.com/apache/pinot/pull/19469#discussion_r4102819968
########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregationSpillManager.java: ########## @@ -0,0 +1,428 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.operator; + +import com.google.common.annotations.VisibleForTesting; +import java.io.BufferedInputStream; +import java.io.DataInputStream; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.file.FileVisitResult; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.SimpleFileVisitor; +import java.nio.file.StandardOpenOption; +import java.nio.file.attribute.BasicFileAttributes; +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.function.Consumer; +import org.apache.pinot.common.datablock.DataBlock; +import org.apache.pinot.common.datablock.DataBlockUtils; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.core.query.aggregation.function.AggregationFunction; +import org.apache.pinot.query.runtime.blocks.MseBlock; +import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock; +import org.apache.pinot.query.runtime.blocks.SerializedDataBlock; +import org.apache.pinot.spi.query.QueryThreadContext; +import org.apache.pinot.spi.utils.CommonConstants.Server; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/// Manages hash-partitioned aggregation spill files for one operator. The caller must finish reading before calling +/// [#close()], which recursively removes the operator-scoped directory. This class is not thread-safe. +@SuppressWarnings("rawtypes") +class AggregationSpillManager implements AutoCloseable { + private static final Logger LOGGER = LoggerFactory.getLogger(AggregationSpillManager.class); + private static final String SPILL_FILE_PREFIX = "partition-"; + private static final String SPILL_FILE_SUFFIX = ".spill"; + private static final String SPILL_SCOPE = "AggregationSpillManager#spill"; + private static final String RESTORE_SCOPE = "AggregationSpillManager#consumePartition"; + private static final int MAX_BUFFERED_ROWS = 1024; + private static final int MAX_BUFFERED_PARTITIONS = 8; + + private final int _numPartitions; + private final int _numGroupKeys; + private final DataSchema _spillSchema; + private final AggregationFunction[] _aggFunctions; + private final Path _spillDirectory; + private final Map<Integer, FileChannel> _spillWriters = new HashMap<>(); + private final ByteBuffer _recordLengthBuffer = ByteBuffer.allocate(Integer.BYTES); + + AggregationSpillManager(int numPartitions, int numGroupKeys, DataSchema spillSchema, + AggregationFunction[] aggFunctions) { + if (numPartitions <= 0 || numPartitions > Server.MAX_MSE_AGGREGATION_SPILL_PARTITIONS) { + throw new IllegalArgumentException( + "Number of spill partitions must be between 1 and " + Server.MAX_MSE_AGGREGATION_SPILL_PARTITIONS); + } + if (numGroupKeys < 0 || numGroupKeys > spillSchema.size()) { + throw new IllegalArgumentException("Invalid number of group keys: " + numGroupKeys); + } + _numPartitions = numPartitions; + _numGroupKeys = numGroupKeys; + _spillSchema = spillSchema; + _aggFunctions = aggFunctions; + try { + _spillDirectory = Files.createTempDirectory("pinot-aggregation-spill-"); Review Comment: Done in 2e40e9402d3069aa86e22cf7e6858b904795b57f. The spill root is server-configurable and defaults beneath the instance data directory; an instance-scoped startup sweep removes orphaned spill directories. Per-query (1 GiB default) and process-wide (8 GiB default) byte reservations stop writes before either configured budget is exceeded. The root must be on a dedicated local filesystem, not tmpfs. [addressed by agent] ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java: ########## @@ -268,12 +367,87 @@ private MseBlock.Eos consumeGroupBy() { MseBlock block = _input.nextBlock(); while (block.isData()) { _groupByExecutor.processBlock((MseBlock.Data) block); + if (_spillThreshold > 0 && _groupByExecutor.getNumGroups() >= _spillThreshold) { + spillCurrentGroups(); + } checkTerminationAndSampleUsage(); block = _input.nextBlock(); } return (MseBlock.Eos) block; } + private void spillCurrentGroups() { + assert _groupByExecutor != null; + if (_groupByExecutor.getNumGroups() == 0) { + return; + } + if (_spillManager == null) { + _spillManager = + new AggregationSpillManager(_spillPartitions, _groupKeyIds.length, getSpillSchema(), _aggFunctions); + _spillDirectory = _spillManager.getSpillDirectory(); + } + AggregationSpillManager.SpillResult spillResult = + _spillManager.spill(_groupByExecutor.getIntermediateResultIterator()); + _statMap.merge(StatKey.SPILL_COUNT, 1L); + _statMap.merge(StatKey.SPILLED_ROWS, spillResult.getRows()); + _statMap.merge(StatKey.SPILLED_BYTES, spillResult.getBytes()); + _groupByExecutor = newInputGroupByExecutor(); + } + + private DataSchema getSpillSchema() { + String[] columnNames = _resultSchema.getColumnNames().clone(); + ColumnDataType[] columnDataTypes = _resultSchema.getColumnDataTypes().clone(); + int numKeys = _groupKeyIds.length; + for (int i = 0; i < _aggFunctions.length; i++) { + columnDataTypes[numKeys + i] = _aggFunctions[i].getType() == AggregationFunctionType.ANYVALUE + ? ColumnDataType.OBJECT : _aggFunctions[i].getIntermediateResultColumnType(); + } + return new DataSchema(columnNames, columnDataTypes); + } + + @Nullable + private MseBlock.Data restoreSpillPartition(int partitionId) { + assert _spillManager != null; + if (!_spillManager.hasPartition(partitionId)) { + return null; + } + MultistageGroupByExecutor executor = newSpillMergeGroupByExecutor(); + _spillManager.consumePartition(partitionId, block -> { + executor.processSpillBlock(block); + if (executor.getNumGroups() > _spillThreshold) { + throw QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED.asException( Review Comment: Done in 2e40e9402d3069aa86e22cf7e6858b904795b57f. Restore now admits groups up to the existing numGroupsLimit rather than throwing when a hash partition exceeds the lower spill trigger. The established truncate/NUM_GROUPS_LIMIT_REACHED behavior remains the default; errorOnNumGroupsLimit still raises an error. The PR description now documents this behavior. [addressed by agent] ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java: ########## @@ -268,12 +367,87 @@ private MseBlock.Eos consumeGroupBy() { MseBlock block = _input.nextBlock(); while (block.isData()) { _groupByExecutor.processBlock((MseBlock.Data) block); + if (_spillThreshold > 0 && _groupByExecutor.getNumGroups() >= _spillThreshold) { + spillCurrentGroups(); + } checkTerminationAndSampleUsage(); block = _input.nextBlock(); } return (MseBlock.Eos) block; } + private void spillCurrentGroups() { + assert _groupByExecutor != null; + if (_groupByExecutor.getNumGroups() == 0) { + return; + } + if (_spillManager == null) { + _spillManager = + new AggregationSpillManager(_spillPartitions, _groupKeyIds.length, getSpillSchema(), _aggFunctions); + _spillDirectory = _spillManager.getSpillDirectory(); + } + AggregationSpillManager.SpillResult spillResult = + _spillManager.spill(_groupByExecutor.getIntermediateResultIterator()); + _statMap.merge(StatKey.SPILL_COUNT, 1L); + _statMap.merge(StatKey.SPILLED_ROWS, spillResult.getRows()); + _statMap.merge(StatKey.SPILLED_BYTES, spillResult.getBytes()); + _groupByExecutor = newInputGroupByExecutor(); + } + + private DataSchema getSpillSchema() { + String[] columnNames = _resultSchema.getColumnNames().clone(); + ColumnDataType[] columnDataTypes = _resultSchema.getColumnDataTypes().clone(); + int numKeys = _groupKeyIds.length; + for (int i = 0; i < _aggFunctions.length; i++) { + columnDataTypes[numKeys + i] = _aggFunctions[i].getType() == AggregationFunctionType.ANYVALUE Review Comment: Done in 2e40e9402d3069aa86e22cf7e6858b904795b57f. ANY_VALUE now reports OBJECT as its intermediate result column type, including when its input type has not been resolved on a merge-only executor. AggregateOperator uses the function contract directly instead of a function-name special case; spill tests cover BYTES, UUID, and other intermediate values. [addressed by agent] ########## pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java: ########## @@ -309,7 +312,7 @@ private Object deserializeValue(ByteBuffer buffer) { case BIG_DECIMAL: return new BigDecimal(new String(deserializeVariableBytes(buffer), StandardCharsets.UTF_8)); case BYTES: - return deserializeVariableBytes(buffer); + return new ByteArray(deserializeVariableBytes(buffer)); Review Comment: The independent, directly tested BYTES fix is in #19667 (910d0a391295bd82a468c6276cdf5a52b608e373). Its regression test covers raw byte[] serialization, ByteArray deserialization, extractFinalResult, and reserialization. The original feature commit in #19469 still contains identical source lines for now; I have marked #19667 as a prerequisite and will remove the duplicate from this PR after the independent fix lands. I am leaving this thread open until that cleanup. [addressed by agent] ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregationSpillManager.java: ########## @@ -0,0 +1,428 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.operator; + +import com.google.common.annotations.VisibleForTesting; +import java.io.BufferedInputStream; +import java.io.DataInputStream; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.file.FileVisitResult; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.SimpleFileVisitor; +import java.nio.file.StandardOpenOption; +import java.nio.file.attribute.BasicFileAttributes; +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.function.Consumer; +import org.apache.pinot.common.datablock.DataBlock; +import org.apache.pinot.common.datablock.DataBlockUtils; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.core.query.aggregation.function.AggregationFunction; +import org.apache.pinot.query.runtime.blocks.MseBlock; +import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock; +import org.apache.pinot.query.runtime.blocks.SerializedDataBlock; +import org.apache.pinot.spi.query.QueryThreadContext; +import org.apache.pinot.spi.utils.CommonConstants.Server; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/// Manages hash-partitioned aggregation spill files for one operator. The caller must finish reading before calling +/// [#close()], which recursively removes the operator-scoped directory. This class is not thread-safe. +@SuppressWarnings("rawtypes") +class AggregationSpillManager implements AutoCloseable { + private static final Logger LOGGER = LoggerFactory.getLogger(AggregationSpillManager.class); + private static final String SPILL_FILE_PREFIX = "partition-"; + private static final String SPILL_FILE_SUFFIX = ".spill"; + private static final String SPILL_SCOPE = "AggregationSpillManager#spill"; + private static final String RESTORE_SCOPE = "AggregationSpillManager#consumePartition"; + private static final int MAX_BUFFERED_ROWS = 1024; + private static final int MAX_BUFFERED_PARTITIONS = 8; + + private final int _numPartitions; + private final int _numGroupKeys; + private final DataSchema _spillSchema; + private final AggregationFunction[] _aggFunctions; + private final Path _spillDirectory; + private final Map<Integer, FileChannel> _spillWriters = new HashMap<>(); + private final ByteBuffer _recordLengthBuffer = ByteBuffer.allocate(Integer.BYTES); + + AggregationSpillManager(int numPartitions, int numGroupKeys, DataSchema spillSchema, + AggregationFunction[] aggFunctions) { + if (numPartitions <= 0 || numPartitions > Server.MAX_MSE_AGGREGATION_SPILL_PARTITIONS) { + throw new IllegalArgumentException( + "Number of spill partitions must be between 1 and " + Server.MAX_MSE_AGGREGATION_SPILL_PARTITIONS); + } + if (numGroupKeys < 0 || numGroupKeys > spillSchema.size()) { + throw new IllegalArgumentException("Invalid number of group keys: " + numGroupKeys); + } + _numPartitions = numPartitions; + _numGroupKeys = numGroupKeys; + _spillSchema = spillSchema; + _aggFunctions = aggFunctions; + try { + _spillDirectory = Files.createTempDirectory("pinot-aggregation-spill-"); + } catch (IOException e) { + throw new UncheckedIOException("Failed to create aggregation spill directory", e); + } + } + + SpillResult spill(Iterator<Object[]> rows) { + List<Object[]>[] partitions = createPartitions(); + boolean[] touched = new boolean[_numPartitions]; + int[] touchedPartitions = new int[_numPartitions]; + int numTouchedPartitions = 0; + int numRows = 0; + int numBufferedRows = 0; + long serializedBytes = 0; + int numRowsProcessed = 0; + int maxBufferedRows = MAX_BUFFERED_ROWS * Math.min(_numPartitions, MAX_BUFFERED_PARTITIONS); + while (rows.hasNext()) { + QueryThreadContext.checkTerminationAndSampleUsagePeriodically(numRowsProcessed++, SPILL_SCOPE); + Object[] row = rows.next(); + int partitionId = getPartition(row); + List<Object[]> partition = partitions[partitionId]; + if (partition == null) { + partition = new ArrayList<>(); + partitions[partitionId] = partition; + } + if (!touched[partitionId]) { + touched[partitionId] = true; + touchedPartitions[numTouchedPartitions++] = partitionId; + } + partition.add(row); + numRows++; + numBufferedRows++; + if (partition.size() == MAX_BUFFERED_ROWS) { + serializedBytes += appendPartition(partitionId, partition); + partition.clear(); + numBufferedRows -= MAX_BUFFERED_ROWS; + } + if (numBufferedRows >= maxBufferedRows) { + serializedBytes += flushPartitions(partitions, touched, touchedPartitions, numTouchedPartitions); + numTouchedPartitions = 0; + numBufferedRows = 0; + } + } + if (numTouchedPartitions > 0) { + serializedBytes += flushPartitions(partitions, touched, touchedPartitions, numTouchedPartitions); + } + return new SpillResult(numRows, serializedBytes); + } + + boolean hasPartition(int partitionId) { + checkPartitionId(partitionId); + return Files.exists(getSpillFile(partitionId)); + } + + /// Consumes all records from a partition and deletes its file after the attempt, including when reading or + /// processing fails. + void consumePartition(int partitionId, Consumer<MseBlock.Data> consumer) { + checkPartitionId(partitionId); + Path spillFile = getSpillFile(partitionId); + if (!Files.exists(spillFile)) { + return; + } + + RuntimeException failure = null; + closeSpillWriter(partitionId); + try (DataInputStream input = + new DataInputStream(new BufferedInputStream(Files.newInputStream(spillFile)))) { + long remainingBytes = Files.size(spillFile); + int numRecordsRead = 0; + while (remainingBytes > 0) { + QueryThreadContext.checkTerminationAndSampleUsagePeriodically(numRecordsRead++, RESTORE_SCOPE); + if (remainingBytes < Integer.BYTES) { + throw new IOException("Truncated spill record length in: " + spillFile); + } + int recordLength = input.readInt(); + remainingBytes -= Integer.BYTES; + if (recordLength < 0 || recordLength > remainingBytes) { + throw new IOException("Invalid spill record length " + recordLength + " in: " + spillFile); + } + byte[] bytes = input.readNBytes(recordLength); + if (bytes.length != recordLength) { + throw new IOException("Truncated spill record in: " + spillFile); + } + remainingBytes -= recordLength; + ByteBuffer buffer = ByteBuffer.wrap(bytes); + DataBlock dataBlock = DataBlockUtils.readFrom(buffer); + if (buffer.hasRemaining()) { + throw new IOException("Trailing bytes in aggregation spill record: " + spillFile); + } + consumer.accept(new SerializedDataBlock(dataBlock)); + } + } catch (IOException e) { + failure = new UncheckedIOException("Failed to read aggregation spill partition: " + partitionId, e); + throw failure; + } catch (RuntimeException e) { + failure = e; + throw e; + } finally { + try { + deleteSpillFile(spillFile); + } catch (RuntimeException e) { + if (failure != null) { + failure.addSuppressed(e); + } else { + LOGGER.warn("Failed to delete consumed aggregation spill partition; close will retry: {}", spillFile, e); + } + } + } + } + + int getNumPartitions() { + return _numPartitions; + } + + Path getSpillDirectory() { + return _spillDirectory; + } + + @VisibleForTesting + int getNumOpenSpillWriters() { + return _spillWriters.size(); + } + + @Override + public void close() { + if (!Files.exists(_spillDirectory)) { + return; + } + RuntimeException failure = null; + try { + closeSpillWriters(); + } catch (RuntimeException e) { + failure = e; + } + try { + Files.walkFileTree(_spillDirectory, new SimpleFileVisitor<>() { + @Override + public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) + throws IOException { + Files.delete(file); + return FileVisitResult.CONTINUE; + } + + @Override + public FileVisitResult postVisitDirectory(Path directory, IOException exception) + throws IOException { + if (exception != null) { + throw exception; + } + Files.delete(directory); + return FileVisitResult.CONTINUE; + } + }); + } catch (IOException e) { + RuntimeException cleanupFailure = + new UncheckedIOException("Failed to delete aggregation spill directory", e); + if (failure != null) { + failure.addSuppressed(cleanupFailure); + } else { + failure = cleanupFailure; + } + } + if (failure != null) { + throw failure; + } + } + + @SuppressWarnings("unchecked") + private List<Object[]>[] createPartitions() { + return new List[_numPartitions]; + } + + private long flushPartitions(List<Object[]>[] partitions, boolean[] touched, int[] touchedPartitions, + int numTouchedPartitions) { + long serializedBytes = 0; + for (int i = 0; i < numTouchedPartitions; i++) { + QueryThreadContext.checkTerminationAndSampleUsagePeriodically(i, SPILL_SCOPE); + int partitionId = touchedPartitions[i]; + List<Object[]> partition = partitions[partitionId]; + if (!partition.isEmpty()) { + serializedBytes += appendPartition(partitionId, partition); + partition.clear(); + } + touched[partitionId] = false; + } + return serializedBytes; + } + + private int getPartition(Object[] row) { Review Comment: Done in 2e40e9402d3069aa86e22cf7e6858b904795b57f. The partitioner now documents the hash/equality invariant across the DataBlock round trip. An array-valued group key is tested through multiple spill rounds; OneObjectKeyGroupIdGenerator groups equivalent array contents instead of their identities, so restore produces one group. [addressed by agent] ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java: ########## @@ -131,23 +157,58 @@ public AggregateOperator(OpChainExecutionContext context, MultiStageOperator inp _comparator = comparator; _errorOnNumGroupsLimit = getErrorOnNumGroupsLimit(node.getNodeHint(), context.getOpChainMetadata()); + _numGroupsLimit = MultistageGroupByExecutor.getNumGroupsLimit(_opChainMetadata, _nodeHint); + int spillThreshold = 0; + int spillPartitions = Server.DEFAULT_MSE_AGGREGATION_SPILL_PARTITIONS; + boolean spillEligible = !groupKeys.isEmpty() && !_leafReturnFinalResult Review Comment: Done in 2e40e9402d3069aa86e22cf7e6858b904795b57f. When spill is requested but ineligible, AggregateOperator logs the reason, including a pushed-down LIMIT enabling group trimming; the PR description calls this exclusion out. Spill partitions are validated independently even if no group-count trigger is set. [addressed by agent] ########## pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryRunnerTest.java: ########## @@ -187,6 +190,28 @@ public void testSelfStatsAreNotNegative() { Assert.assertTrue(checked > 0, "expected some self stats to check, got: " + statsTree); } + @Test + public void testMSEAggregationSpill() { Review Comment: Done in 2e40e9402d3069aa86e22cf7e6858b904795b57f. A two-server end-to-end DISTINCTCOUNT query now compares spilled results with the non-spilled baseline; the existing SUM/AVG end-to-end test additionally checks normal spill-directory cleanup. A focused operator-level test exercises array group keys across spill rounds and verifies one merged result. The array case is at operator level rather than a distributed MV-column query. [addressed by agent] ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/SendStatsPredicate.java: ########## @@ -153,6 +165,11 @@ public boolean isSendStats() { return _sendStats; } + @Override + public boolean isClusterVersionCompatible() { Review Comment: Done in 2e40e9402d3069aa86e22cf7e6858b904795b57f. Removed the spill-specific homogeneous-cluster/SAFE gate. StatMap rejects an unknown ordinal with IllegalArgumentException and MultiStageStatsTreeDecoder converts unchecked decode failures to DecodeFailedException, so the stream path degrades like the existing mailbox path. Both failure boundaries have regression tests. [addressed by agent] ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java: ########## @@ -131,23 +157,58 @@ public AggregateOperator(OpChainExecutionContext context, MultiStageOperator inp _comparator = comparator; _errorOnNumGroupsLimit = getErrorOnNumGroupsLimit(node.getNodeHint(), context.getOpChainMetadata()); + _numGroupsLimit = MultistageGroupByExecutor.getNumGroupsLimit(_opChainMetadata, _nodeHint); + int spillThreshold = 0; + int spillPartitions = Server.DEFAULT_MSE_AGGREGATION_SPILL_PARTITIONS; + boolean spillEligible = !groupKeys.isEmpty() && !_leafReturnFinalResult + && groupTrimSize == Integer.MAX_VALUE && QueryOptionsUtils.isMSEAggregationSpillEnabled(_opChainMetadata); + if (spillEligible) { + Integer configuredSpillThreshold = QueryOptionsUtils.getMSEAggregationSpillThreshold(_opChainMetadata); + if (configuredSpillThreshold != null) { + spillThreshold = configuredSpillThreshold; + Integer configuredSpillPartitions = QueryOptionsUtils.getMSEAggregationSpillPartitions(_opChainMetadata); + if (configuredSpillPartitions != null) { + spillPartitions = configuredSpillPartitions; + } + } + } + _spillThreshold = spillThreshold; + _spillPartitions = spillPartitions; Review Comment: Done in 2e40e9402d3069aa86e22cf7e6858b904795b57f. Spill input now retains numGroupsLimit in its GroupIdGenerator instead of replacing it with Integer.MAX_VALUE. The spill predicate triggers at min(mseAggregationSpillMaxGroups, numGroupsLimit), so an oversized configured group trigger cannot remove the pre-existing per-table memory backstop. [addressed by agent] ########## pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java: ########## @@ -131,23 +157,58 @@ public AggregateOperator(OpChainExecutionContext context, MultiStageOperator inp _comparator = comparator; _errorOnNumGroupsLimit = getErrorOnNumGroupsLimit(node.getNodeHint(), context.getOpChainMetadata()); + _numGroupsLimit = MultistageGroupByExecutor.getNumGroupsLimit(_opChainMetadata, _nodeHint); + int spillThreshold = 0; + int spillPartitions = Server.DEFAULT_MSE_AGGREGATION_SPILL_PARTITIONS; + boolean spillEligible = !groupKeys.isEmpty() && !_leafReturnFinalResult + && groupTrimSize == Integer.MAX_VALUE && QueryOptionsUtils.isMSEAggregationSpillEnabled(_opChainMetadata); + if (spillEligible) { + Integer configuredSpillThreshold = QueryOptionsUtils.getMSEAggregationSpillThreshold(_opChainMetadata); + if (configuredSpillThreshold != null) { + spillThreshold = configuredSpillThreshold; + Integer configuredSpillPartitions = QueryOptionsUtils.getMSEAggregationSpillPartitions(_opChainMetadata); + if (configuredSpillPartitions != null) { + spillPartitions = configuredSpillPartitions; + } + } + } + _spillThreshold = spillThreshold; + _spillPartitions = spillPartitions; // Initialize the appropriate executor. - AggregateNode.AggType aggType = node.getAggType(); // TODO: Allow leaf return final result for non-group-by queries - boolean leafReturnFinalResult = node.isLeafReturnFinalResult(); if (groupKeys.isEmpty()) { _aggregationExecutor = - new MultistageAggregationExecutor(_aggFunctions, filterArgIds, maxFilterArgId, aggType, _resultSchema); + new MultistageAggregationExecutor(_aggFunctions, _filterArgIds, _maxFilterArgId, _aggType, _resultSchema); _groupByExecutor = null; } else { - _groupByExecutor = - new MultistageGroupByExecutor(getGroupKeyIds(groupKeys), _aggFunctions, filterArgIds, maxFilterArgId, aggType, - leafReturnFinalResult, _resultSchema, context.getOpChainMetadata(), node.getNodeHint()); + _groupByExecutor = newInputGroupByExecutor(); _aggregationExecutor = null; } } + private MultistageGroupByExecutor newInputGroupByExecutor() { + if (_spillThreshold > 0) { + return MultistageGroupByExecutor.forSpillInput(_groupKeyIds, _aggFunctions, _filterArgIds, _maxFilterArgId, + _aggType, _resultSchema, _opChainMetadata, _nodeHint, _spillThreshold); + } + return new MultistageGroupByExecutor(_groupKeyIds, _aggFunctions, _filterArgIds, _maxFilterArgId, _aggType, + _leafReturnFinalResult, _resultSchema, _opChainMetadata, _nodeHint); + } + + private MultistageGroupByExecutor newSpillMergeGroupByExecutor() { + AggregateNode.AggType mergeAggType = + _aggType.isOutputIntermediateFormat() ? AggregateNode.AggType.INTERMEDIATE : AggregateNode.AggType.FINAL; + int[] spillGroupKeyIds = new int[_groupKeyIds.length]; + for (int i = 0; i < spillGroupKeyIds.length; i++) { + spillGroupKeyIds[i] = i; + } + int expectedPartitionGroups = Review Comment: Done in 2e40e9402d3069aa86e22cf7e6858b904795b57f. The former _spillThreshold is split into the input trigger, a restore-side group ceiling, and an initial merge-table capacity hint. The hint is documented as an estimate, not a bound. The restore ceiling follows numGroupsLimit to preserve the previous truncation semantics. [addressed by agent] -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
