This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch release-1.6.0-rc
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/release-1.6.0-rc by this push:
new 15aeb95fdbd [1.6.0] Backport MSE approximate aggregation CASE/null fix
(#19977)
15aeb95fdbd is described below
commit 15aeb95fdbd8a98305e0aab8e49be4481edb4275
Author: Xiang Fu <[email protected]>
AuthorDate: Thu Oct 22 15:40:42 2026 -0700
[1.6.0] Backport MSE approximate aggregation CASE/null fix (#19977)
Run the MSE approximate aggregate rewrite after CASE-to-FILTER conversion,
preserving conditional distinct-count NULL handling.
(cherry picked from commit 9fbd72ae42a892f1638b732a02d7e7946f7c9b6b)
---
...ApproximateFunctionOverrideIntegrationTest.java | 44 +++++++++++++++++++---
.../PinotApproximateAggregateRewriteRule.java | 10 ++++-
.../calcite/rel/rules/PinotQueryRuleSets.java | 11 +++---
.../query/ApproximateAggregateRewriteTest.java | 44 ++++++++++++++++++++++
4 files changed, 97 insertions(+), 12 deletions(-)
diff --git
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ApproximateFunctionOverrideIntegrationTest.java
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ApproximateFunctionOverrideIntegrationTest.java
index 73bef16c83c..60880c7cc69 100644
---
a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ApproximateFunctionOverrideIntegrationTest.java
+++
b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ApproximateFunctionOverrideIntegrationTest.java
@@ -22,6 +22,7 @@ import com.fasterxml.jackson.databind.JsonNode;
import java.io.File;
import java.util.List;
import java.util.Locale;
+import java.util.Map;
import org.apache.commons.io.FileUtils;
import org.apache.helix.model.HelixConfigScope;
import org.apache.helix.model.builder.HelixConfigScopeBuilder;
@@ -35,6 +36,7 @@ import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
/// End-to-end test of the cluster config that rewrites exact aggregations
into their approximate counterparts. It
@@ -48,6 +50,12 @@ public class ApproximateFunctionOverrideIntegrationTest
extends BaseClusterInteg
"SELECT Carrier, COUNT(DISTINCT Origin) FROM mytable GROUP BY Carrier
ORDER BY Carrier LIMIT 100";
private static final String PERCENTILE_QUERY =
"SELECT Carrier, PERCENTILE(ArrDelay, 90) FROM mytable GROUP BY Carrier
ORDER BY Carrier LIMIT 100";
+ // The last predicate matches no row, so the exact multi-stage count is 0
and a counted NULL would show up as 1.
+ private static final String CONDITIONAL_DISTINCT_COUNT_QUERY = "SELECT
Carrier, "
+ + "COUNT(DISTINCT CASE WHEN ArrDelay > 0 THEN Origin ELSE NULL END), "
+ + "DISTINCTCOUNT(CASE WHEN ArrDelay > 0 THEN Origin ELSE NULL END), "
+ + "COUNT(DISTINCT CASE WHEN ArrDelay > 1000000 THEN Origin ELSE NULL
END) "
+ + "FROM mytable GROUP BY Carrier ORDER BY Carrier LIMIT 100";
private HelixConfigScope _clusterScope;
@@ -152,6 +160,31 @@ public class ApproximateFunctionOverrideIntegrationTest
extends BaseClusterInteg
}
}
+ /// The multi-stage engine turns `COUNT(DISTINCT CASE WHEN c THEN x END)`
into `COUNT(DISTINCT x) FILTER (WHERE c)`,
+ /// and the rewrite must not stop that, or the NULL of every row that fails
`c` counts as one more value. The
+ /// single-stage engine has no such conversion, but its answer must not
change either.
+ @Test(dataProvider = "useBothQueryEngines")
+ public void testConditionalDistinctCountMatchesExact(boolean
useMultiStageQueryEngine)
+ throws Exception {
+ setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+ JsonNode exact = postQuery(CONDITIONAL_DISTINCT_COUNT_QUERY);
+ assertEquals(exact.get("exceptions").size(), 0, exact.toString());
+ assertFalse(rows(exact).isEmpty(), exact.toString());
+
+ // The threshold keeps the rewritten functions exact, so the flag stays
off.
+ setApproximateFunctionOverride("threshold=1000000");
+ waitForOverride(true);
+ JsonNode approximate = postQuery(CONDITIONAL_DISTINCT_COUNT_QUERY);
+ assertEquals(approximate.get("exceptions").size(), 0,
approximate.toString());
+ assertApproximated(approximate, false);
+ assertSameResult(exact, approximate);
+ if (useMultiStageQueryEngine) {
+ for (JsonNode row : rows(approximate)) {
+ assertEquals(row.get(3).asLong(), 0L, approximate.toString());
+ }
+ }
+ }
+
/// Pins the one visible difference the rewrite makes to a single-stage
response, so that changing it is a
/// deliberate decision. The multi-stage engine keeps the Calcite-derived
column name.
@Test
@@ -194,8 +227,8 @@ public class ApproximateFunctionOverrideIntegrationTest
extends BaseClusterInteg
}
/// The rewrite is already on, so [#waitForOverride] would return before the
new threshold arrives. Waits until both
- /// engines report an approximated percentile instead. The percentile params
are the last key that
- /// [#setApproximateFunctionOverride] writes, so the distinct count params
have arrived too by then.
+ /// engines report an approximated percentile instead.
[#setApproximateFunctionOverride] writes all keys at once, so
+ /// the distinct count params have arrived too by then.
private void waitForApproximatedResult() {
boolean useMultiStageQueryEngine = useMultiStageQueryEngine();
try {
@@ -215,10 +248,11 @@ public class ApproximateFunctionOverrideIntegrationTest
extends BaseClusterInteg
}
}
+ /// Writes all keys at once, so the broker cannot see the flag on before the
parameters arrive.
private void setApproximateFunctionOverride(String params) {
- _helixManager.getConfigAccessor().set(_clusterScope,
Broker.USE_APPROXIMATE_FUNCTION, "true");
- _helixManager.getConfigAccessor().set(_clusterScope,
Broker.APPROXIMATE_FUNCTION_DISTINCT_COUNT_PARAMS, params);
- _helixManager.getConfigAccessor().set(_clusterScope,
Broker.APPROXIMATE_FUNCTION_PERCENTILE_PARAMS, params);
+ _helixManager.getConfigAccessor().set(_clusterScope,
Map.of(Broker.USE_APPROXIMATE_FUNCTION, "true",
+ Broker.APPROXIMATE_FUNCTION_DISTINCT_COUNT_PARAMS, params,
Broker.APPROXIMATE_FUNCTION_PERCENTILE_PARAMS,
+ params));
}
/// Removes the keys, so the broker falls back to its own conf.
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotApproximateAggregateRewriteRule.java
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotApproximateAggregateRewriteRule.java
index 449bef9add7..21a23a5dc76 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotApproximateAggregateRewriteRule.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotApproximateAggregateRewriteRule.java
@@ -55,6 +55,11 @@ import
org.apache.pinot.spi.utils.CommonConstants.Broker.PlannerRuleNames;
/// It must run before [PinotAggregateExchangeNodeInsertRule], which derives
the leaf-to-final intermediate result
/// format from the function name, hence `Phase.BASIC` rather than that rule's
`POST_LOGICAL`. That is also why the
/// runtime is not a valid place to do this rewrite.
+///
+/// It must also run after `AggregateCaseToFilterRule`, which only matches a
single-argument call. With the parameter
+/// argument appended first, `COUNT(DISTINCT CASE WHEN c THEN x END)` keeps
its CASE and counts the NULL as a value,
+/// instead of becoming `COUNT(DISTINCT x) FILTER (WHERE c)`. The rewrite
keeps the call's filter, so a converted CASE
+/// carries over.
public class PinotApproximateAggregateRewriteRule extends RelOptRule {
public static final PinotApproximateAggregateRewriteRule INSTANCE =
new
PinotApproximateAggregateRewriteRule(PinotRuleUtils.PINOT_REL_FACTORY);
@@ -92,8 +97,9 @@ public class PinotApproximateAggregateRewriteRule extends
RelOptRule {
String percentileParams = hasPercentile ?
envConfig.approximateFunctionPercentileParams() : "";
// The parameters must be input fields, because an aggregate call holds
field indices rather than inline literals,
- // so they are projected underneath the aggregate. With no parameters
configured the plan keeps the shape it had
- // before, with no extra project.
+ // so they are projected underneath the aggregate. They go after the input
fields, so the argument and filter
+ // indices of every call stay valid. With no parameters configured the
plan keeps the shape it had before, with no
+ // extra project.
//
// Invariant: this project has to stay directly beneath the aggregate.
Both PinotAggregateExchangeNodeInsertRule
// and AggregatePushdownRule inline the literal only through
findImmediateProjects(input), so anything that puts a
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java
index f62294c58ef..0d61c3a62cb 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java
@@ -175,14 +175,15 @@ public class PinotQueryRuleSets {
PinotAggregateFunctionRewriteRule
.instanceWithDescription(PlannerRuleNames.AGGREGATE_FUNCTION_REWRITE),
- // Must stay ahead of PinotAggregateExchangeNodeInsertRule
(POST_LOGICAL), which fixes the leaf-to-final
- // intermediate result format from the function name.
- PinotApproximateAggregateRewriteRule.INSTANCE,
-
// convert CASE-style filtered aggregates into true filtered aggregates
// put it after AGGREGATE_REDUCE_FUNCTIONS where SUM is converted to SUM0
+ // put it before APPROXIMATE_AGGREGATE_REWRITE, see
PinotApproximateAggregateRewriteRule
AggregateCaseToFilterRule.Config.DEFAULT
- .withDescription(PlannerRuleNames.AGGREGATE_CASE_TO_FILTER).toRule()
+ .withDescription(PlannerRuleNames.AGGREGATE_CASE_TO_FILTER).toRule(),
+
+ // Must run after AGGREGATE_CASE_TO_FILTER and before
PinotAggregateExchangeNodeInsertRule (POST_LOGICAL). The
+ // rule's class doc says why.
+ PinotApproximateAggregateRewriteRule.INSTANCE
);
// Filter pushdown rules run using a RuleCollection since we want to push
down a filter as much as possible in a
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/ApproximateAggregateRewriteTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/ApproximateAggregateRewriteTest.java
index 7a93ab1966e..a220377a1a9 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/ApproximateAggregateRewriteTest.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/ApproximateAggregateRewriteTest.java
@@ -18,6 +18,7 @@
*/
package org.apache.pinot.query;
+import java.util.List;
import java.util.Map;
import org.apache.pinot.core.routing.MockRoutingManagerFactory;
import org.apache.pinot.core.routing.RoutingManager;
@@ -26,6 +27,7 @@ import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.utils.CommonConstants;
import org.testng.annotations.Test;
+import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertTrue;
@@ -108,6 +110,43 @@ public class ApproximateAggregateRewriteTest extends
QueryEnvironmentTestBase {
assertTrue(subquery.contains("threshold=23"), subquery);
}
+ /// `AggregateCaseToFilterRule` turns `COUNT(DISTINCT CASE WHEN c THEN x
END)` into `COUNT(DISTINCT x) FILTER (WHERE
+ /// c)`, so the rows that fail `c` add nothing to the count. It only matches
a single-argument call, so the rewrite
+ /// has to run after it: with the parameters appended first, the CASE stays
and its NULL counts as one more value.
+ @Test
+ public void testKeepsCaseToFilterConversion() {
+ QueryEnvironment exactEnv = newQueryEnvironment(false, "", "");
+ QueryEnvironment approximateEnv = newQueryEnvironment(true, "", "");
+ QueryEnvironment approximateWithParamsEnv = newQueryEnvironment(true,
"threshold=17", "");
+ QueryEnvironment physicalOptimizerEnv = newQueryEnvironment(true,
"threshold=17", "", true);
+ for (String query : List.of(
+ "SELECT COUNT(DISTINCT CASE WHEN col3 > 0 THEN col1 ELSE NULL END)
FROM a",
+ "SELECT col2, COUNT(DISTINCT CASE WHEN col3 > 0 THEN col1 ELSE NULL
END) FROM a GROUP BY col2",
+ "SELECT DISTINCTCOUNT(CASE WHEN col3 > 0 THEN col1 ELSE NULL END) FROM
a",
+ "SELECT col2, DISTINCTCOUNT(CASE WHEN col3 > 0 THEN col1 END) FROM a
GROUP BY col2",
+ // An explicit FILTER is combined with the CASE condition, and must
survive the rewrite too.
+ "SELECT COUNT(DISTINCT CASE WHEN col3 > 0 THEN col1 END) FILTER (WHERE
col2 = 'x') FROM a")) {
+ String exact = explain(exactEnv, query);
+ assertFilteredWithoutCase(exact);
+
+ // Without parameters the only difference from the exact plan is the
function name.
+ String approximate = explain(approximateEnv, query);
+ assertRewritten(approximate, "distinctcountsmarthll");
+ assertEquals(approximate.replace("DISTINCTCOUNTSMARTHLL",
"DISTINCTCOUNT"), exact);
+
+ // This is the case that kept the CASE, because the parameter makes it a
two-argument call.
+ String approximateWithParams = explain(approximateWithParamsEnv, query);
+ assertRewritten(approximateWithParams, "distinctcountsmarthll");
+ assertTrue(approximateWithParams.contains("threshold=17"),
approximateWithParams);
+ assertFilteredWithoutCase(approximateWithParams);
+
+ String physicalOptimizer = explain(physicalOptimizerEnv, query);
+ assertRewritten(physicalOptimizer, "distinctcountsmarthll");
+ assertTrue(physicalOptimizer.contains("threshold=17"),
physicalOptimizer);
+ assertFilteredWithoutCase(physicalOptimizer);
+ }
+ }
+
/// The physical optimizer splits aggregates in its own rule, which also
keys the leaf-to-final intermediate format
/// off the function name, so the rewrite has to hold on that path too.
@Test
@@ -130,6 +169,11 @@ public class ApproximateAggregateRewriteTest extends
QueryEnvironmentTestBase {
assertFalse(plan.toLowerCase().contains("smart"), "Unexpected rewrite
in:\n" + plan);
}
+ private static void assertFilteredWithoutCase(String plan) {
+ assertTrue(plan.contains(") FILTER $"), "Expected a filtered aggregate
in:\n" + plan);
+ assertFalse(plan.contains("CASE"), "Expected the CASE to be converted to a
filter in:\n" + plan);
+ }
+
private static String explain(QueryEnvironment env, String query) {
return env.explainQuery("EXPLAIN PLAN FOR " + query,
RANDOM_REQUEST_ID_GEN.nextLong());
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
