From 620ccb7cc892dede9778f96f46ce1e70f2ec1f86 Mon Sep 17 00:00:00 2001 From: David Yang Date: Tue, 12 May 2026 09:45:57 -0700 Subject: [PATCH 1/8] Fix CountAggregationFunction.extractFinalResult returning null instead of 0 --- .../query/aggregation/function/CountAggregationFunction.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/CountAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/CountAggregationFunction.java index 7285ba0b5c71..8bbdda5fcec2 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/CountAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/CountAggregationFunction.java @@ -21,6 +21,7 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import javax.annotation.Nullable; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.core.common.BlockValSet; @@ -204,8 +205,8 @@ public ColumnDataType getFinalResultColumnType() { } @Override - public Long extractFinalResult(Long intermediateResult) { - return intermediateResult; + public Long extractFinalResult(@Nullable Long intermediateResult) { + return intermediateResult != null ? intermediateResult : 0L; } @Override From 176ce09d8ccd2f6eee11c1b646e65b61c896c297 Mon Sep 17 00:00:00 2001 From: David Yang Date: Tue, 12 May 2026 09:46:01 -0700 Subject: [PATCH 2/8] Add TODO for MSE broker-side short-circuit when all segments are pruned --- .../broker/requesthandler/MultiStageBrokerRequestHandler.java | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java index 846f6acb3c83..fda12c4ef2d4 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java @@ -588,6 +588,9 @@ private BrokerResponse query(QueryEnvironment.CompiledQuery query, long requestI DispatchableSubPlan dispatchableSubPlan = queryPlanResult.getQueryPlan(); + // TODO: Short-circuit here if all segments are pruned like SSE does. Currently fully-pruned queries still + // dispatch intermediate stages to servers. + // Optionally set ignoreMissingSegments query option based on broker config if not already set. if (_config.getProperty(CommonConstants.Broker.CONFIG_OF_IGNORE_MISSING_SEGMENTS, CommonConstants.Broker.DEFAULT_IGNORE_MISSING_SEGMENTS)) { From 1d1d7be002a3e5f505a7878fc50f9b635fee9388 Mon Sep 17 00:00:00 2001 From: David Yang Date: Tue, 12 May 2026 09:46:05 -0700 Subject: [PATCH 3/8] Pre-initialize MSE merge result holder with identity values for pruned-segment safety --- .../MultistageAggregationExecutor.java | 4 + .../operator/AggregateOperatorTest.java | 76 +++++++++++++++++++ 2 files changed, 80 insertions(+) diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageAggregationExecutor.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageAggregationExecutor.java index 14b76c7c6198..4e44e90de3a5 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageAggregationExecutor.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageAggregationExecutor.java @@ -65,6 +65,10 @@ public MultistageAggregationExecutor(AggregationFunction[] aggFunctions, int[] f _mergeResultHolder = null; } else { _mergeResultHolder = new Object[numFunctions]; + for (int i = 0; i < numFunctions; i++) { + _mergeResultHolder[i] = aggFunctions[i].extractAggregationResult( + aggFunctions[i].createAggregationResultHolder()); + } _aggregateResultHolder = null; } } diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java index 60ab3e503564..8d392d97a39e 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java @@ -49,6 +49,7 @@ import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.BOOLEAN; import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.DOUBLE; import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.INT; +import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.LONG; import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.STRING; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -362,6 +363,81 @@ private AggregateOperator getAggregateOperator(OpChainExecutionContext context, collations, limit)); } + // ---- Merge-mode identity value tests ---- + // The merge result holder is pre-initialized with identity values so that when no input blocks arrive + // (e.g., all segments pruned by broker), the correct result is returned. This affects both FINAL and + // INTERMEDIATE modes since both use the merge path (isInputIntermediateFormat=true). + + // FINAL-mode tests: identity values are converted to final form via extractFinalResult. + + @Test + public void testFinalCountWithNoInputReturnsZero() { + assertMergeAggIdentityValue("COUNT", LONG, AggType.FINAL, 0L); + } + + @Test + public void testFinalSumWithNoInputReturnsNull() { + assertMergeAggIdentityValue("SUM", DOUBLE, AggType.FINAL, null); + } + + @Test + public void testFinalMinWithNoInputReturnsNull() { + assertMergeAggIdentityValue("MIN", DOUBLE, AggType.FINAL, null); + } + + @Test + public void testFinalMaxWithNoInputReturnsNull() { + assertMergeAggIdentityValue("MAX", DOUBLE, AggType.FINAL, null); + } + + @Test + public void testFinalDistinctCountWithNoInputReturnsZero() { + assertMergeAggIdentityValue("DISTINCTCOUNT", INT, AggType.FINAL, 0); + } + + @Test + public void testFinalAvgWithNoInputReturnsNull() { + assertMergeAggIdentityValue("AVG", DOUBLE, AggType.FINAL, null); + } + + @Test + public void testFinalDistinctCountHllWithNoInputReturnsZero() { + assertMergeAggIdentityValue("DISTINCTCOUNTHLL", LONG, AggType.FINAL, 0L); + } + + // INTERMEDIATE-mode tests: identity values are returned as-is (intermediate form). + + @Test + public void testIntermediateCountWithNoInputReturnsZero() { + assertMergeAggIdentityValue("COUNT", LONG, AggType.INTERMEDIATE, 0L); + } + + @Test + public void testIntermediateSumWithNoInputReturnsNull() { + // INTERMEDIATE returns the raw intermediate value, not converted to final form + assertMergeAggIdentityValue("SUM", DOUBLE, AggType.INTERMEDIATE, null); + } + + private void assertMergeAggIdentityValue(String functionName, ColumnDataType resultType, AggType aggType, + Object expectedValue) { + MseBlock block = runMergeAggWithNoInput(functionName, resultType, aggType); + assertTrue(block.isData(), "Expected a data block for " + functionName + " " + aggType); + List rows = ((MseBlock.Data) block).asRowHeap().getRows(); + assertEquals(rows.size(), 1, functionName + " should produce exactly 1 row"); + assertEquals(rows.get(0)[0], expectedValue, functionName + " " + aggType + " identity value mismatch"); + } + + private MseBlock runMergeAggWithNoInput(String functionName, ColumnDataType resultType, AggType aggType) { + RexExpression.FunctionCall aggCall = + new RexExpression.FunctionCall(resultType, functionName, List.of(new RexExpression.InputRef(0))); + when(_input.nextBlock()).thenReturn(SuccessMseBlock.INSTANCE); + DataSchema resultSchema = new DataSchema(new String[]{"result"}, new ColumnDataType[]{resultType}); + AggregateOperator operator = new AggregateOperator(OperatorTestUtil.getContext(Map.of()), _input, + new AggregateNode(-1, resultSchema, PlanNode.NodeHint.EMPTY, List.of(), List.of(aggCall), List.of(-1), + List.of(), aggType, false, null, 0)); + return operator.nextBlock(); + } + @Test public void shouldRecordNumGroupsBelowLimit() { // Given: 1 distinct group key, limit = 2 — below limit, no overflow From cfba92cac14974c19fe80bda6991719ad2bee002 Mon Sep 17 00:00:00 2001 From: David Yang Date: Tue, 12 May 2026 09:57:23 -0700 Subject: [PATCH 4/8] Add null guards to Covariance, IdSet, and AnyValue extractFinalResult --- .../aggregation/function/AnyValueAggregationFunction.java | 3 ++- .../aggregation/function/CovarianceAggregationFunction.java | 6 +++++- .../aggregation/function/IdSetAggregationFunction.java | 6 +++++- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java index 267233478747..63af9d6a8d31 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java @@ -24,6 +24,7 @@ import java.util.List; import java.util.Map; import java.util.function.Consumer; +import javax.annotation.Nullable; import org.apache.pinot.common.CustomObject; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; @@ -102,7 +103,7 @@ public Object extractGroupByResult(GroupByResultHolder groupByResultHolder, int } @Override - public Comparable extractFinalResult(Object intermediateResult) { + public Comparable extractFinalResult(@Nullable Object intermediateResult) { return (Comparable) intermediateResult; } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/CovarianceAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/CovarianceAggregationFunction.java index a3881c3df70d..4a8900885bdd 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/CovarianceAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/CovarianceAggregationFunction.java @@ -22,6 +22,7 @@ import java.util.ArrayList; import java.util.List; import java.util.Map; +import javax.annotation.Nullable; import org.apache.pinot.common.CustomObject; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; @@ -196,7 +197,10 @@ public ColumnDataType getFinalResultColumnType() { } @Override - public Double extractFinalResult(CovarianceTuple covarianceTuple) { + public Double extractFinalResult(@Nullable CovarianceTuple covarianceTuple) { + if (covarianceTuple == null) { + return null; + } long count = covarianceTuple.getCount(); if (count == 0L) { return DEFAULT_FINAL_RESULT; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/IdSetAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/IdSetAggregationFunction.java index 3613b42831f4..b51946c6c5a3 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/IdSetAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/IdSetAggregationFunction.java @@ -22,6 +22,7 @@ import java.io.IOException; import java.util.List; import java.util.Map; +import javax.annotation.Nullable; import org.apache.commons.lang3.StringUtils; import org.apache.pinot.common.CustomObject; import org.apache.pinot.common.request.context.ExpressionContext; @@ -492,7 +493,10 @@ public ColumnDataType getFinalResultColumnType() { } @Override - public String extractFinalResult(IdSet intermediateResult) { + public String extractFinalResult(@Nullable IdSet intermediateResult) { + if (intermediateResult == null) { + return null; + } try { return intermediateResult.toBase64String(); } catch (IOException e) { From 9513da56ebd009a0734f35f3e4581e17b206adc3 Mon Sep 17 00:00:00 2001 From: David Yang Date: Tue, 12 May 2026 10:02:39 -0700 Subject: [PATCH 5/8] Revert "Pre-initialize MSE merge result holder with identity values for pruned-segment safety" --- .../MultistageAggregationExecutor.java | 4 - .../operator/AggregateOperatorTest.java | 76 ------------------- 2 files changed, 80 deletions(-) diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageAggregationExecutor.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageAggregationExecutor.java index 4e44e90de3a5..14b76c7c6198 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageAggregationExecutor.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageAggregationExecutor.java @@ -65,10 +65,6 @@ public MultistageAggregationExecutor(AggregationFunction[] aggFunctions, int[] f _mergeResultHolder = null; } else { _mergeResultHolder = new Object[numFunctions]; - for (int i = 0; i < numFunctions; i++) { - _mergeResultHolder[i] = aggFunctions[i].extractAggregationResult( - aggFunctions[i].createAggregationResultHolder()); - } _aggregateResultHolder = null; } } diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java index 8d392d97a39e..60ab3e503564 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java @@ -49,7 +49,6 @@ import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.BOOLEAN; import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.DOUBLE; import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.INT; -import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.LONG; import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.STRING; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -363,81 +362,6 @@ private AggregateOperator getAggregateOperator(OpChainExecutionContext context, collations, limit)); } - // ---- Merge-mode identity value tests ---- - // The merge result holder is pre-initialized with identity values so that when no input blocks arrive - // (e.g., all segments pruned by broker), the correct result is returned. This affects both FINAL and - // INTERMEDIATE modes since both use the merge path (isInputIntermediateFormat=true). - - // FINAL-mode tests: identity values are converted to final form via extractFinalResult. - - @Test - public void testFinalCountWithNoInputReturnsZero() { - assertMergeAggIdentityValue("COUNT", LONG, AggType.FINAL, 0L); - } - - @Test - public void testFinalSumWithNoInputReturnsNull() { - assertMergeAggIdentityValue("SUM", DOUBLE, AggType.FINAL, null); - } - - @Test - public void testFinalMinWithNoInputReturnsNull() { - assertMergeAggIdentityValue("MIN", DOUBLE, AggType.FINAL, null); - } - - @Test - public void testFinalMaxWithNoInputReturnsNull() { - assertMergeAggIdentityValue("MAX", DOUBLE, AggType.FINAL, null); - } - - @Test - public void testFinalDistinctCountWithNoInputReturnsZero() { - assertMergeAggIdentityValue("DISTINCTCOUNT", INT, AggType.FINAL, 0); - } - - @Test - public void testFinalAvgWithNoInputReturnsNull() { - assertMergeAggIdentityValue("AVG", DOUBLE, AggType.FINAL, null); - } - - @Test - public void testFinalDistinctCountHllWithNoInputReturnsZero() { - assertMergeAggIdentityValue("DISTINCTCOUNTHLL", LONG, AggType.FINAL, 0L); - } - - // INTERMEDIATE-mode tests: identity values are returned as-is (intermediate form). - - @Test - public void testIntermediateCountWithNoInputReturnsZero() { - assertMergeAggIdentityValue("COUNT", LONG, AggType.INTERMEDIATE, 0L); - } - - @Test - public void testIntermediateSumWithNoInputReturnsNull() { - // INTERMEDIATE returns the raw intermediate value, not converted to final form - assertMergeAggIdentityValue("SUM", DOUBLE, AggType.INTERMEDIATE, null); - } - - private void assertMergeAggIdentityValue(String functionName, ColumnDataType resultType, AggType aggType, - Object expectedValue) { - MseBlock block = runMergeAggWithNoInput(functionName, resultType, aggType); - assertTrue(block.isData(), "Expected a data block for " + functionName + " " + aggType); - List rows = ((MseBlock.Data) block).asRowHeap().getRows(); - assertEquals(rows.size(), 1, functionName + " should produce exactly 1 row"); - assertEquals(rows.get(0)[0], expectedValue, functionName + " " + aggType + " identity value mismatch"); - } - - private MseBlock runMergeAggWithNoInput(String functionName, ColumnDataType resultType, AggType aggType) { - RexExpression.FunctionCall aggCall = - new RexExpression.FunctionCall(resultType, functionName, List.of(new RexExpression.InputRef(0))); - when(_input.nextBlock()).thenReturn(SuccessMseBlock.INSTANCE); - DataSchema resultSchema = new DataSchema(new String[]{"result"}, new ColumnDataType[]{resultType}); - AggregateOperator operator = new AggregateOperator(OperatorTestUtil.getContext(Map.of()), _input, - new AggregateNode(-1, resultSchema, PlanNode.NodeHint.EMPTY, List.of(), List.of(aggCall), List.of(-1), - List.of(), aggType, false, null, 0)); - return operator.nextBlock(); - } - @Test public void shouldRecordNumGroupsBelowLimit() { // Given: 1 distinct group key, limit = 2 — below limit, no overflow From 9371346addded87fa97e9b9f7f003a33600bb825 Mon Sep 17 00:00:00 2001 From: David Yang Date: Tue, 19 May 2026 21:54:24 -0700 Subject: [PATCH 6/8] revert to only the count agg change --- .../requesthandler/MultiStageBrokerRequestHandler.java | 3 --- .../aggregation/function/AnyValueAggregationFunction.java | 3 +-- .../aggregation/function/CovarianceAggregationFunction.java | 6 +----- .../aggregation/function/IdSetAggregationFunction.java | 6 +----- 4 files changed, 3 insertions(+), 15 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java index fda12c4ef2d4..846f6acb3c83 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java @@ -588,9 +588,6 @@ private BrokerResponse query(QueryEnvironment.CompiledQuery query, long requestI DispatchableSubPlan dispatchableSubPlan = queryPlanResult.getQueryPlan(); - // TODO: Short-circuit here if all segments are pruned like SSE does. Currently fully-pruned queries still - // dispatch intermediate stages to servers. - // Optionally set ignoreMissingSegments query option based on broker config if not already set. if (_config.getProperty(CommonConstants.Broker.CONFIG_OF_IGNORE_MISSING_SEGMENTS, CommonConstants.Broker.DEFAULT_IGNORE_MISSING_SEGMENTS)) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java index 63af9d6a8d31..267233478747 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java @@ -24,7 +24,6 @@ import java.util.List; import java.util.Map; import java.util.function.Consumer; -import javax.annotation.Nullable; import org.apache.pinot.common.CustomObject; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; @@ -103,7 +102,7 @@ public Object extractGroupByResult(GroupByResultHolder groupByResultHolder, int } @Override - public Comparable extractFinalResult(@Nullable Object intermediateResult) { + public Comparable extractFinalResult(Object intermediateResult) { return (Comparable) intermediateResult; } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/CovarianceAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/CovarianceAggregationFunction.java index 4a8900885bdd..a3881c3df70d 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/CovarianceAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/CovarianceAggregationFunction.java @@ -22,7 +22,6 @@ import java.util.ArrayList; import java.util.List; import java.util.Map; -import javax.annotation.Nullable; import org.apache.pinot.common.CustomObject; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; @@ -197,10 +196,7 @@ public ColumnDataType getFinalResultColumnType() { } @Override - public Double extractFinalResult(@Nullable CovarianceTuple covarianceTuple) { - if (covarianceTuple == null) { - return null; - } + public Double extractFinalResult(CovarianceTuple covarianceTuple) { long count = covarianceTuple.getCount(); if (count == 0L) { return DEFAULT_FINAL_RESULT; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/IdSetAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/IdSetAggregationFunction.java index b51946c6c5a3..3613b42831f4 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/IdSetAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/IdSetAggregationFunction.java @@ -22,7 +22,6 @@ import java.io.IOException; import java.util.List; import java.util.Map; -import javax.annotation.Nullable; import org.apache.commons.lang3.StringUtils; import org.apache.pinot.common.CustomObject; import org.apache.pinot.common.request.context.ExpressionContext; @@ -493,10 +492,7 @@ public ColumnDataType getFinalResultColumnType() { } @Override - public String extractFinalResult(@Nullable IdSet intermediateResult) { - if (intermediateResult == null) { - return null; - } + public String extractFinalResult(IdSet intermediateResult) { try { return intermediateResult.toBase64String(); } catch (IOException e) { From 2158f68e639302cdc9d6d39d312a1f83fdecbeb4 Mon Sep 17 00:00:00 2001 From: David Yang Date: Tue, 19 May 2026 22:02:08 -0700 Subject: [PATCH 7/8] add test --- .../function/CountAggregationFunctionTest.java | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/CountAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/CountAggregationFunctionTest.java index 360233b0a204..cffce5dce688 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/CountAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/CountAggregationFunctionTest.java @@ -19,12 +19,16 @@ package org.apache.pinot.core.query.aggregation.function; +import java.util.List; +import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.queries.FluentQueryTest; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.Schema; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; +import static org.testng.Assert.assertEquals; + public class CountAggregationFunctionTest extends AbstractAggregationFunctionTest { @@ -234,6 +238,14 @@ public void countGroupByMV() { ); } + @Test + public void testExtractFinalResultReturnsZeroForNull() { + CountAggregationFunction function = + new CountAggregationFunction(List.of(ExpressionContext.forIdentifier("col")), false); + assertEquals(function.extractFinalResult(null), Long.valueOf(0L)); + assertEquals(function.extractFinalResult(5L), Long.valueOf(5L)); + } + @DataProvider(name = "nullHandlingEnabled") public Object[][] nullHandlingEnabled() { return new Object[][]{ From b59d560b6b09f19bf2c264119f7d1d610e3b4884 Mon Sep 17 00:00:00 2001 From: David Yang Date: Wed, 20 May 2026 15:47:43 -0700 Subject: [PATCH 8/8] empty commit to trigger CI