This is an automated email from the ASF dual-hosted git repository.
mrhhsg pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new fa6da72465f [fix](fe) Keep Java and Python UDAFs out of bucketed hash
aggregation (#68565)
fa6da72465f is described below
commit fa6da72465fb65771579517808e7bcb87ad7e79f
Author: Jerry Hu <[email protected]>
AuthorDate: Tue Sep 29 20:46:33 2026 +0800
[fix](fe) Keep Java and Python UDAFs out of bucketed hash aggregation
(#68565)
### What problem does this PR solve?
Issue Number: None
Related PR: #61495, #65024
Problem Summary: On a single-BE cluster, bucketed hash aggregation is on
by
default and the translator fuses a one-phase GLOBAL aggregate with its
distribute child into a BucketedAggregationNode. The source side of that
operator merges the live aggregate states built by different sink
instances
directly, instead of serializing them and deserializing them with the
merging
evaluator as the two-phase plan does. Java and Python UDAFs rely on the
latter:
- Java UDAF: the extra evaluator clone used by the bucketed source never
calls
create(), so its _exec_place stays null and merge()/insert_result_into()
dereference a null state. Reproduced locally with a Java UDAF
(`SELECT k, my_udaf(v) FROM t GROUP BY k`): UBSan reports "reference
binding
to null pointer of type AggregateJavaUdafData" in
AggregateJavaUdaf::merge
and the query fails / the BE goes down.
- Python UDAF: merge() builds the rhs state from serialize_data, which
is
only filled on the deserialize path, so the rhs contribution is dropped
or
the Python server RPC fails.
None of the FE gates excluded UDAFs. Add the check to the shared gate
AggregateUtils.isBucketedHashAggEnabled, which now takes the aggregate
and
returns false when any aggregate function is a Udf (JavaUdaf /
PythonUdaf).
The translator, ChildrenPropertiesRegulator, ChildOutputPropertyDeriver
and
CostModel all go through this gate, so the optimizer also stops
preferring
the one-phase plan for these aggregates and they keep the regular
aggregation path.
### Release note
Fix BE crash / wrong result when a Java or Python UDAF is used with
GROUP BY
on a single-BE cluster with bucketed hash aggregation enabled.
### Check List (For Author)
- Test:
- Unit Test: BucketedAggregateTranslatorTest (new Python UDAF case under
agg_phase=0 and agg_phase=1, fails
without the fix), BucketedAggregateTest, ChildOutputPropertyDeriverTest,
ChildrenPropertiesRegulatorTest, CostModelV1Test
- Regression test: query_p0/javaudf/test_javaudaf_bucketed_agg (default
and agg_phase=1 plans; fails on
the old FE with BUCKETED AGGREGATE in the plan and a BE null deref when
executed), plus bucketed_hash_agg and percentile_bucketed_agg_merge
- Behavior changed: Yes (aggregates containing Java/Python UDAFs no
longer
use bucketed hash aggregation)
- Does this need documentation: No
---
.../org/apache/doris/nereids/cost/CostModel.java | 3 +-
.../glue/translator/PhysicalPlanTranslator.java | 4 +-
.../properties/ChildOutputPropertyDeriver.java | 3 +-
.../properties/ChildrenPropertiesRegulator.java | 2 +-
.../apache/doris/nereids/util/AggregateUtils.java | 16 +++-
.../BucketedAggregateTranslatorTest.java | 66 ++++++++++++++--
.../javaudf/test_javaudaf_bucketed_agg.out | 16 ++++
.../javaudf/test_javaudaf_bucketed_agg.groovy | 89 ++++++++++++++++++++++
8 files changed, 182 insertions(+), 17 deletions(-)
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/nereids/cost/CostModel.java
b/fe/fe-core/src/main/java/org/apache/doris/nereids/cost/CostModel.java
index 4c9d2a0b4ba..aa88b013c9f 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/cost/CostModel.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/cost/CostModel.java
@@ -364,8 +364,7 @@ class CostModel extends PlanVisitor<Cost, PlanContext> {
// gates are enforced by ChildrenPropertiesRegulator), apply a
discount
// to prefer this path over two-phase aggregation.
if (aggregate.getAggMode() == AggMode.INPUT_TO_RESULT
- && AggregateUtils.isBucketedHashAggEnabled(
- aggregate.getGroupByExpressions().size())) {
+ && AggregateUtils.isBucketedHashAggEnabled(aggregate)) {
rowCost *= BUCKETED_AGG_COST_DISCOUNT;
}
return Cost.of(context.getCostWeight(),
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java
b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java
index 67de840bf92..04d78a69983 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java
@@ -3330,8 +3330,8 @@ public class PhysicalPlanTranslator extends
DefaultPlanVisitor<PlanFragment, Pla
*/
private boolean shouldUseBucketedFusion(PhysicalHashAggregate<? extends
Plan> aggregate,
PlanTranslatorContext context) {
- // Shared eligibility: session var, single-BE, GROUP BY, smooth upgrade
- if
(!AggregateUtils.isBucketedHashAggEnabled(aggregate.getGroupByExpressions().size()))
{
+ // Shared eligibility: session var, single-BE, GROUP BY, smooth
upgrade, no UDAF
+ if (!AggregateUtils.isBucketedHashAggEnabled(aggregate)) {
return false;
}
// Must be one-phase: GLOBAL + INPUT_TO_RESULT
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriver.java
b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriver.java
index 2df7723a7ab..a1568491a9d 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriver.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriver.java
@@ -199,8 +199,7 @@ public class ChildOutputPropertyDeriver extends
PlanVisitor<PhysicalProperties,
// from incorrectly skipping exchanges.
if (agg.getAggPhase().isGlobal()
&& agg.getAggMode() == AggMode.INPUT_TO_RESULT
- && AggregateUtils.isBucketedHashAggEnabled(
- agg.getGroupByExpressions().size())
+ && AggregateUtils.isBucketedHashAggEnabled(agg)
&&
isShuffleCompatible(childOutputProperty.getDistributionSpec())) {
return PhysicalProperties.ANY;
}
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildrenPropertiesRegulator.java
b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildrenPropertiesRegulator.java
index 9741d1de68a..483fe9cab8a 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildrenPropertiesRegulator.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildrenPropertiesRegulator.java
@@ -164,7 +164,7 @@ public class ChildrenPropertiesRegulator extends
PlanVisitor<List<List<PhysicalP
// pattern so the translator can fuse it into
BucketedAggregationNode.
// Gate with data-volume checks using group-level statistics to
avoid
// generating this pattern when bucketed agg is unsuitable.
- if
(AggregateUtils.isBucketedHashAggEnabled(aggregate.getGroupByExpressions().size()))
{
+ if (AggregateUtils.isBucketedHashAggEnabled(aggregate)) {
return !bucketedDataVolumeGatesPass(aggregate);
}
return true;
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/nereids/util/AggregateUtils.java
b/fe/fe-core/src/main/java/org/apache/doris/nereids/util/AggregateUtils.java
index b16320beeaa..d10c1805b62 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/util/AggregateUtils.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/util/AggregateUtils.java
@@ -23,6 +23,7 @@ import org.apache.doris.nereids.trees.expressions.Cast;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.IsNull;
import org.apache.doris.nereids.trees.expressions.NamedExpression;
+import org.apache.doris.nereids.trees.expressions.functions.Udf;
import
org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateParam;
import org.apache.doris.nereids.trees.expressions.functions.agg.Count;
@@ -232,9 +233,10 @@ public class AggregateUtils {
* discount), and PhysicalPlanTranslator (for fusion into
BucketedAggregationNode).
*
* @return true if the session variable is enabled, there is exactly one
alive BE,
- * no smooth upgrade is in progress, and the aggregate has GROUP
BY keys.
+ * no smooth upgrade is in progress, the aggregate has GROUP BY
keys and
+ * contains no user-defined aggregate function.
*/
- public static boolean isBucketedHashAggEnabled(int groupByExprCount) {
+ public static boolean isBucketedHashAggEnabled(Aggregate<? extends Plan>
aggregate) {
ConnectContext ctx = ConnectContext.get();
if (ctx == null) {
return false;
@@ -243,7 +245,7 @@ public class AggregateUtils {
return false;
}
// Must have GROUP BY keys (without-key aggregation not supported)
- if (groupByExprCount == 0) {
+ if (aggregate.getGroupByExpressions().isEmpty()) {
return false;
}
// Correctness gate: single-BE only (cross-BE in-memory merge is
impossible).
@@ -265,6 +267,14 @@ public class AggregateUtils {
return false;
}
}
+ // Bucketed agg merges the live states built by different sink
instances
+ // directly, without serializing them. Java / Python UDAFs can only
merge a
+ // state that was deserialized by the merging evaluator (the Java UDAF
+ // executor place and the Python UDAF serialized buffer are only set
up on
+ // that path), so they must stay on the regular aggregation path.
+ if
(aggregate.getAggregateFunctions().stream().anyMatch(Udf.class::isInstance)) {
+ return false;
+ }
return true;
}
}
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/nereids/glue/translator/BucketedAggregateTranslatorTest.java
b/fe/fe-core/src/test/java/org/apache/doris/nereids/glue/translator/BucketedAggregateTranslatorTest.java
index 479d55e380b..14457d2ba05 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/nereids/glue/translator/BucketedAggregateTranslatorTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/nereids/glue/translator/BucketedAggregateTranslatorTest.java
@@ -41,6 +41,47 @@ public class BucketedAggregateTranslatorTest extends
TestWithFeService {
+ "kint INT NOT NULL, kbint INT NOT NULL, kstr STRING NOT
NULL) "
+ "DISTRIBUTED BY HASH(kint) BUCKETS 4 "
+ "PROPERTIES('replication_num' = '1')");
+ createFunction("CREATE AGGREGATE FUNCTION
bucketed_aggregate_translator_test.py_udaf_sum(INT) "
+ + "RETURNS BIGINT PROPERTIES('type'='PYTHON_UDF',
'symbol'='SumUdaf', "
+ + "'runtime_version'='3.10.2')");
+ }
+
+ @Test
+ public void testPythonUdafIsNotFusedIntoBucketedAggregation() throws
Exception {
+ SessionVariable sessionVariable = connectContext.getSessionVariable();
+ int oldAggPhase = sessionVariable.aggPhase;
+ int oldBeNumberForTest = sessionVariable.getBeNumberForTest();
+ long oldBucketedAggMinInputRows =
sessionVariable.bucketedAggMinInputRows;
+ long oldBucketedAggMaxGroupKeys =
sessionVariable.bucketedAggMaxGroupKeys;
+ double oldBucketedAggHighCardThreshold =
sessionVariable.bucketedAggHighCardThreshold;
+ boolean oldEnableBucketedHashAgg =
sessionVariable.enableBucketedHashAgg;
+ try {
+ sessionVariable.setBeNumberForTest(1);
+ sessionVariable.bucketedAggMinInputRows = 0;
+ sessionVariable.bucketedAggMaxGroupKeys = 0;
+ sessionVariable.bucketedAggHighCardThreshold = 1.0;
+ sessionVariable.enableBucketedHashAgg = true;
+
+ // agg_phase=0 lets the optimizer choose the plan; agg_phase=1
forces the
+ // one-phase plan, so only the translator fusion gate can reject
the UDAF.
+ for (int aggPhase : new int[] {0, 1}) {
+ sessionVariable.aggPhase = aggPhase;
+ // A builtin aggregate on the same shape is fused, so the UDAF
cases below
+ // are rejected because of the UDAF rather than the plan shape.
+
Assertions.assertFalse(collectBucketedAggregationNodes("sum(kint)").isEmpty());
+ assertUsesRegularAggregation(
+
"bucketed_aggregate_translator_test.py_udaf_sum(kint)");
+ assertUsesRegularAggregation(
+ "sum(kint),
bucketed_aggregate_translator_test.py_udaf_sum(kint)");
+ }
+ } finally {
+ sessionVariable.aggPhase = oldAggPhase;
+ sessionVariable.setBeNumberForTest(oldBeNumberForTest);
+ sessionVariable.bucketedAggMinInputRows =
oldBucketedAggMinInputRows;
+ sessionVariable.bucketedAggMaxGroupKeys =
oldBucketedAggMaxGroupKeys;
+ sessionVariable.bucketedAggHighCardThreshold =
oldBucketedAggHighCardThreshold;
+ sessionVariable.enableBucketedHashAgg = oldEnableBucketedHashAgg;
+ }
}
@Test
@@ -78,18 +119,29 @@ public class BucketedAggregateTranslatorTest extends
TestWithFeService {
}
private void assertUsesRegularAggregation(String aggregateFunction) throws
Exception {
- Planner planner = getSQLPlanner("SELECT " + aggregateFunction
+ Planner planner = planAggregate(aggregateFunction);
+ Assertions.assertTrue(collectNodes(planner,
BucketedAggregationNode.class).isEmpty());
+ Assertions.assertFalse(collectNodes(planner,
AggregationNode.class).isEmpty());
+ }
+
+ private List<BucketedAggregationNode>
collectBucketedAggregationNodes(String aggregateFunction)
+ throws Exception {
+ return collectNodes(planAggregate(aggregateFunction),
BucketedAggregationNode.class);
+ }
+
+ private Planner planAggregate(String aggregateFunction) throws Exception {
+ return getSQLPlanner("SELECT " + aggregateFunction
+ " FROM
bucketed_aggregate_translator_test.agg_group_concat_table GROUP BY kbint");
- List<BucketedAggregationNode> bucketedAggregationNodes =
Lists.newArrayList();
- List<AggregationNode> aggregationNodes = Lists.newArrayList();
+ }
+
+ private <T extends PlanNode> List<T> collectNodes(Planner planner,
Class<T> nodeClass) {
+ List<T> nodes = Lists.newArrayList();
for (PlanFragment fragment : planner.getFragments()) {
PlanNode root = fragment.getPlanRoot();
if (root != null) {
- root.collect(BucketedAggregationNode.class,
bucketedAggregationNodes);
- root.collect(AggregationNode.class, aggregationNodes);
+ root.collect(nodeClass, nodes);
}
}
- Assertions.assertTrue(bucketedAggregationNodes.isEmpty());
- Assertions.assertFalse(aggregationNodes.isEmpty());
+ return nodes;
}
}
diff --git
a/regression-test/data/query_p0/javaudf/test_javaudaf_bucketed_agg.out
b/regression-test/data/query_p0/javaudf/test_javaudaf_bucketed_agg.out
new file mode 100644
index 00000000000..025b79b7d76
--- /dev/null
+++ b/regression-test/data/query_p0/javaudf/test_javaudaf_bucketed_agg.out
@@ -0,0 +1,16 @@
+-- This file is automatically generated. You should know what you did if you
want to edit this
+-- !udaf --
+0 1683
+1 1617
+2 1650
+
+-- !udaf_with_builtin --
+0 1683 1683
+1 1617 1617
+2 1650 1650
+
+-- !udaf_one_phase --
+0 1683
+1 1617
+2 1650
+
diff --git
a/regression-test/suites/query_p0/javaudf/test_javaudaf_bucketed_agg.groovy
b/regression-test/suites/query_p0/javaudf/test_javaudaf_bucketed_agg.groovy
new file mode 100644
index 00000000000..c7f829fabae
--- /dev/null
+++ b/regression-test/suites/query_p0/javaudf/test_javaudaf_bucketed_agg.groovy
@@ -0,0 +1,89 @@
+// 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.
+
+// Bucketed hash aggregation merges the live states of different sink instances
+// without serializing them, which Java UDAFs cannot handle. A Java UDAF must
+// therefore keep the regular aggregation plan even when bucketed agg applies.
+suite("test_javaudaf_bucketed_agg") {
+ def jarPath =
"""${context.file.parent}/../../javaudf_p0/jars/java-udf-case-jar-with-dependencies.jar"""
+ scp_udf_file_to_all_be(jarPath)
+
+ sql "set be_number_for_test=1"
+ sql "set enable_bucketed_hash_agg=true"
+ sql "set bucketed_agg_min_input_rows=0"
+ sql "set bucketed_agg_max_group_keys=0"
+ sql "set bucketed_agg_high_card_threshold=1.0"
+ sql "set parallel_pipeline_task_num=2"
+
+ sql "DROP TABLE IF EXISTS test_javaudaf_bucketed_agg_tbl"
+ sql """
+ CREATE TABLE test_javaudaf_bucketed_agg_tbl (
+ id INT NOT NULL,
+ k INT NOT NULL,
+ v INT NOT NULL
+ )
+ DISTRIBUTED BY HASH(id) BUCKETS 4
+ PROPERTIES("replication_num" = "1")
+ """
+ // The same group key spreads over all tablets, so several sink instances
+ // build a state for it and the source side has to merge them.
+ sql """
+ INSERT INTO test_javaudaf_bucketed_agg_tbl
+ SELECT number, number % 3, number FROM numbers("number" = "100")
+ """
+
+ sql "DROP FUNCTION IF EXISTS test_javaudaf_bucketed_agg_sum(int)"
+ sql """ CREATE AGGREGATE FUNCTION test_javaudaf_bucketed_agg_sum(int)
RETURNS BigInt PROPERTIES (
+ "file"="file://${jarPath}",
+ "symbol"="org.apache.doris.udf.MySumInt",
+ "always_nullable"="false",
+ "type"="JAVA_UDF"
+ ); """
+
+ // A builtin aggregate on the same shape still uses bucketed agg.
+ explain {
+ sql "SELECT k, sum(v) FROM test_javaudaf_bucketed_agg_tbl GROUP BY k"
+ contains("BUCKETED AGGREGATE")
+ }
+ explain {
+ sql "SELECT k, test_javaudaf_bucketed_agg_sum(v) FROM
test_javaudaf_bucketed_agg_tbl GROUP BY k"
+ notContains("BUCKETED AGGREGATE")
+ }
+ explain {
+ sql """SELECT k, sum(v), test_javaudaf_bucketed_agg_sum(v)
+ FROM test_javaudaf_bucketed_agg_tbl GROUP BY k"""
+ notContains("BUCKETED AGGREGATE")
+ }
+
+ order_qt_udaf """
+ SELECT k, test_javaudaf_bucketed_agg_sum(v) FROM
test_javaudaf_bucketed_agg_tbl GROUP BY k
+ """
+ order_qt_udaf_with_builtin """
+ SELECT k, sum(v), test_javaudaf_bucketed_agg_sum(v) FROM
test_javaudaf_bucketed_agg_tbl GROUP BY k
+ """
+
+ // Force the one-phase plan so that only the translator fusion gate can
keep
+ // the Java UDAF away from bucketed agg.
+ sql "set agg_phase=1"
+ explain {
+ sql "SELECT k, test_javaudaf_bucketed_agg_sum(v) FROM
test_javaudaf_bucketed_agg_tbl GROUP BY k"
+ notContains("BUCKETED AGGREGATE")
+ }
+ order_qt_udaf_one_phase """
+ SELECT k, test_javaudaf_bucketed_agg_sum(v) FROM
test_javaudaf_bucketed_agg_tbl GROUP BY k
+ """
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
