mrhhsg commented on code in PR #68651:
URL: https://github.com/apache/doris/pull/68651#discussion_r4141254450
##########
fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java:
##########
@@ -3216,6 +3200,262 @@ private PlanFragment connectJoinNode(HashJoinNode
hashJoinNode, PlanFragment lef
return leftFragment;
}
+ /**
+ * Check whether the one-phase GLOBAL hash aggregate can be fused with its
+ * distribute child into a BucketedAggregationNode. This eliminates
exchange
+ * overhead on single-BE deployments by using in-memory per-bucket merging.
+ */
+ private boolean shouldUseBucketedFusion(PhysicalHashAggregate<? extends
Plan> aggregate,
+ PlanTranslatorContext context) {
+ // Shared eligibility (also used by the regulator, the output property
deriver
+ // and the cost model): session var, single-BE, GROUP BY, spill /
query cache
+ // off, smooth upgrade, no UDAF, one-phase GLOBAL INPUT_TO_RESULT, no
partial
+ // (buffer-producing) function, two-phase capable functions, no pushed
TopN.
+ if (!AggregateUtils.isBucketedHashAggFusible(aggregate)) {
+ return false;
+ }
+ // Child must be PhysicalDistribute with hash distribution matching
group keys
+ Plan child = aggregate.child(0);
+ if (!(child instanceof PhysicalDistribute)) {
+ return false;
+ }
+ // Bucketed fusion bypasses the distribute/exchange and builds
directly on the
+ // child fragment. When the child subtree contains a CTE consumer
(materialized
+ // multicast CTE), the child fragment is the MultiCastPlanFragment; a
parent
+ // distribute would then treat the aggregate output slots as consumer
slots and
+ // fail with "Required producer slot ... doesn't exist". Fall back to
the
+ // regular one-phase path (which keeps the exchange) for such plans.
+ if (containsCTEConsumer(child)) {
+ return false;
+ }
+ // The distribute's child subtree must be a unary pipeline over
exactly one
+ // olap scan. Fusing an aggregate whose input contains a join / set-op
/ CTE
+ // subtree would leave multiple olap scans in a single fragment
(rejected by
+ // UnassignedJobBuilder: "Not supported multiple scan multiple
OlapTable but
+ // not contains colocate join or bucket shuffle join"), and fusing
over a
+ // nested aggregate would break the bucket alignment between stages.
+ if (!isSingleOlapScanPipeline(aggregate.child(0).child(0))) {
+ return false;
+ }
+ // The parent is a fragment-merging node (join / set-op) that consumes
this
+ // fragment without an exchange boundary: fusing removes the exchange
that
+ // keeps the scan in its own fragment, so multiple scans would end up
in the
+ // same fragment and the scan-assignment would fail. Only fuse when the
+ // parent chain keeps an exchange boundary (e.g. a top-level
aggregate).
+ if (context.isInFragmentMergeChild()) {
+ return false;
+ }
+ DistributionSpec distSpec = ((PhysicalDistribute<?>)
child).getDistributionSpec();
+ if (!(distSpec instanceof DistributionSpecHash)) {
+ return false;
+ }
+ List<ExprId> distKeys = ((DistributionSpecHash)
distSpec).getOrderedShuffledColumns();
+ List<ExprId> groupByKeys = aggregate.getGroupByExpressions().stream()
+ .filter(SlotReference.class::isInstance)
+ .map(SlotReference.class::cast)
+ .map(SlotReference::getExprId)
+ .collect(Collectors.toList());
+ return distKeys.equals(groupByKeys);
+ }
+
+ /** Returns true if the plan subtree contains a physical CTE consumer. */
+ private boolean containsCTEConsumer(Plan plan) {
+ if (plan instanceof PhysicalCTEConsumer) {
+ return true;
+ }
+ for (Plan child : plan.children()) {
+ if (containsCTEConsumer(child)) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ /**
+ * Returns true if the plan subtree is a unary pipeline over exactly one
olap
+ * scan, i.e. it translates into a single-scan fragment that bucketed
fusion
+ * can safely build upon. Subtrees containing fragment-merging or
+ * distribution-changing nodes (join / set-op / CTE / nested aggregate /
+ * storage-layer aggregate) are rejected.
+ */
+ private boolean isSingleOlapScanPipeline(Plan plan) {
+ if (plan instanceof PhysicalOlapScan) {
+ return true;
+ }
+ if (plan instanceof PhysicalHashJoin
+ || plan instanceof PhysicalNestedLoopJoin
+ || plan instanceof PhysicalSetOperation
+ || plan instanceof PhysicalCTEConsumer
+ || plan instanceof PhysicalCTEAnchor
+ || plan instanceof PhysicalHashAggregate
+ || plan instanceof PhysicalStorageLayerAggregate) {
+ return false;
+ }
+ if (plan.children().size() == 1) {
+ return isSingleOlapScanPipeline(plan.child(0));
+ }
+ return false;
+ }
+
+ /**
+ * Fuse a one-phase GLOBAL hash aggregate and its PhysicalDistribute child
+ * into a BucketedAggregationNode, skipping the exchange node entirely.
+ * Visits the distribute's child directly to keep everything in one
fragment.
+ */
+ private PlanFragment visitBucketedFusion(
+ PhysicalHashAggregate<? extends Plan> aggregate,
+ PlanTranslatorContext context) {
+ // Visit the distribute's direct child, bypassing the distribute
entirely.
+ // This avoids creating an ExchangeNode that bucketed agg does not
need.
+ Plan distributeChild = aggregate.child(0).child(0);
+ PlanFragment inputPlanFragment = distributeChild.accept(this, context);
Review Comment:
Fixed in e86c5594e2d. The single-BE decision is now carried through to scan
placement:
- `AggregateUtils.getBucketedHashAggBackend` returns the only alive backend
(the gate uses it). `visitPhysicalHashAggregate` reads it again right before
the fused fragment is built and falls back to the regular aggregate + exchange
if the alive set changed in between.
- `visitBucketedFusion` puts that backend into the translator context, and
`visitPhysicalOlapScan` pins the scan (`OlapScanNode.setPinnedBackendId`)
*before* `init()` builds the scan range locations (they are built there, not in
`finalizeForNereids`). `addScanRangeLocations` keeps only the replicas on the
pinned backend, and `LoadBalanceScanWorkerSelector` only ever chooses among a
tablet's locations, so a backend that recovers after the gate cannot receive
any tablet of the fused fragment. If the pinned backend has no queryable
replica of a tablet, the query fails with a retry hint instead of splitting a
group across backends. Non-scan fragments above a fused subtree take their
workers from their child fragments, so they follow the pin as well.
- `checkBucketedFusionScanPinned` asserts that exactly one olap scan was
translated below the aggregate and that it carries the pin.
Test:
`BucketedAggregateMultiBackendTest.testFusedScanIsPinnedToTheBackendSeenBySingleBackendGate`
uses a 2-BE cluster and a `replication_num=2` table, mocks
`getAllBackendByCurrentCluster` so the gate sees one alive backend while both
are alive when the locations are built, and asserts the plan is fused and every
tablet offers exactly the pinned backend (the control plan offers both). The
test failed while the pin was applied after the subtree was translated, so it
does catch the ordering.
##########
fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriver.java:
##########
@@ -215,12 +216,35 @@ public PhysicalProperties visitPhysicalHashAggregate(
case GLOBAL:
case DISTINCT_LOCAL:
case DISTINCT_GLOBAL:
+ // Bucketed hash agg fusion: when the one-phase GLOBAL
aggregate
+ // will be fused with its distribute child into
BucketedAggregationNode,
+ // the output is NOT hash-distributed (256-bucket internal
hash is
+ // not shuffle-compatible). Advertise ANY to prevent parent
operators
+ // from incorrectly skipping exchanges.
+ if (AggregateUtils.isBucketedHashAggFusible(agg)
+ &&
isShuffleCompatible(childOutputProperty.getDistributionSpec())) {
+ return PhysicalProperties.ANY;
Review Comment:
Partially addressed in e86c5594e2d; the exact shape described here does not
hit the ANY branch on this branch.
The distribute the enforcer inserts below the aggregate carries
`ShuffleType.EXECUTION_BUCKETED`
(`EnforceMissingPropertiesHelper.enforceDistribution`), and
`isShuffleCompatible` only accepts `REQUIRE`. So for `Join <- GlobalAgg(k) <-
Distribute(HASH(k)) <- Scan` the deriver keeps the child's hash property, the
join consumes the aggregate directly (physical plan: `hashJoin[bucketShuffle]`
with no distribute between them), no exchange is enforced above it and the
translator keeps the regular aggregate. There is no redundant exchange; the new
UT `testJoinOnGroupKeyConsumesRegularAggregateWithoutExtraExchange` pins that
(the aggregation node is the join's direct child in the same fragment).
What was inconsistent is the other case: when the join key is not the group
key (or the parent requests GATHER, as a recursive union does), the enforcer
adds a distribute *above* the aggregate, and the translator still refused to
fuse because the join had set the fragment-merge context, although that
exchange keeps the fused fragment apart from the join. The cost model had
already granted the discount for a plan that ended up shipping raw rows below a
regular aggregate plus the enforcer exchange. `visitPhysicalDistribute` now
clears the fragment-merge context for its subtree (an exchange is a fragment
boundary), so those aggregates fuse into `Join <- Exchange <- BucketedAgg <-
Scan`, matching the deriver/cost decision. Tests:
`testAggregateBelowJoinExchangeIsFusedIntoBucketedAggregation`,
`RecursiveUnionFragmentMergeContextTest` (now asserts fusion below the gather
exchange, a separate fragment and one olap scan per fragment), regression
`bucketed_hash_agg` Test 9/10.
The remaining gap is an aggregate the merging node consumes without any
exchange (a union child, the probe side of a broadcast join) or an aggregate
over a non-single-scan subtree: the deriver only sees memo groups, not the
physical subtree, so it cannot know that the translator keeps the exchange
there. In those cases the reported property is conservative (ANY can only add
an exchange the parent asked for) and the discount is a cost bias, not a
correctness issue; that matches master and I would rather fix it upstream first
than teach the deriver the translator's subtree walk.
##########
be/src/exec/operator/bucketed_aggregation_source_operator.cpp:
##########
@@ -0,0 +1,774 @@
+// 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.
+
+#include "exec/operator/bucketed_aggregation_source_operator.h"
+
+#include <memory>
+#include <string>
+
+#include "common/exception.h"
+#include "core/column/column_vector.h"
+#include "exec/common/hash_table/hash.h"
+#include "exec/common/util.hpp"
+#include "exec/operator/bucketed_aggregation_sink_operator.h"
+#include "exec/operator/operator.h"
+#include "exprs/vectorized_agg_fn.h"
+#include "runtime/runtime_profile.h"
+#include "runtime/thread_context.h"
+
+namespace doris {
+
+// Helper to set/get null key data on hash tables that support it
(DataWithNullKey).
+// For hash tables without nullable key support (PHHashMap), these are no-ops.
+// This is needed because in nested std::visit lambdas, the outer hash table
type is already
+// resolved and doesn't depend on the inner template parameter, so `if
constexpr` inside the
+// inner lambda cannot suppress compilation of code that accesses
has_null_key_data() on the
+// outer (non-dependent) type.
+template <typename HashTable>
+constexpr bool has_nullable_key_v =
+
std::is_assignable_v<decltype(std::declval<HashTable&>().has_null_key_data()),
bool>;
+
+template <typename HashTable>
+void set_null_key_flag(HashTable& ht, bool val) {
+ if constexpr (has_nullable_key_v<HashTable>) {
+ ht.has_null_key_data() = val;
+ }
+}
+
+template <typename HashTable>
+bool get_null_key_flag(const HashTable& ht) {
+ if constexpr (has_nullable_key_v<HashTable>) {
+ return ht.has_null_key_data();
+ } else {
+ return false;
+ }
+}
+
+template <typename HashTable>
+AggregateDataPtr get_null_key_agg_data(HashTable& ht) {
+ if constexpr (has_nullable_key_v<HashTable>) {
+ return ht.template get_null_key_data<AggregateDataPtr>();
+ } else {
+ return nullptr;
+ }
+}
+
+template <typename HashTable>
+void set_null_key_agg_data(HashTable& ht, AggregateDataPtr val) {
+ if constexpr (has_nullable_key_v<HashTable>) {
+ ht.template get_null_key_data<AggregateDataPtr>() = val;
+ }
+}
+
+// Returns a REFERENCE to the null key's AggregateDataPtr slot.
+// Critical for simple_count merge: writing through a copy would lose the
update (Bug #30).
+template <typename HashTable>
+AggregateDataPtr& get_null_key_agg_data_ref(HashTable& ht) {
+ static_assert(has_nullable_key_v<HashTable>,
+ "get_null_key_agg_data_ref requires a nullable hash table");
+ return ht.template get_null_key_data<AggregateDataPtr>();
+}
+
+// Helper for emplace that works with PHHashMap (3-arg).
+template <typename HashTable, typename Key>
+auto hash_table_emplace(HashTable& ht, const Key& key, typename
HashTable::LookupResult& it,
+ bool& inserted) -> decltype(ht.emplace(key, it,
inserted), void()) {
+ ht.emplace(key, it, inserted);
+}
+
+/// Merge src aggregate state into dst_ref (a reference to the mapped slot).
+/// For simple_count, adds UInt64 counters directly via the reference.
+/// For regular aggregates, calls merge() on each function then destroys src
state.
+/// After return, src is consumed and must not be used.
+static void merge_agg_states(AggregateDataPtr& dst_ref, AggregateDataPtr src,
bool use_simple_count,
+ const std::vector<AggFnEvaluator*>& evaluators,
const Sizes& offsets,
+ Arena& arena) {
+ if (use_simple_count) {
+ // simple_count: mapped slots hold UInt64 counters. MUST use reference
+ // to write back correctly.
+ reinterpret_cast<UInt64&>(dst_ref) += reinterpret_cast<UInt64>(src);
+ } else {
+ const size_t num_fns = evaluators.size();
+ for (size_t i = 0; i < num_fns; ++i) {
+ evaluators[i]->function()->merge(dst_ref + offsets[i], src +
offsets[i], arena);
+ }
+ for (size_t i = 0; i < num_fns; ++i) {
+ evaluators[i]->function()->destroy(src + offsets[i]);
+ }
+ }
+}
+
+/// Merge a source null key into a destination null key slot. Handles three
cases:
+/// 1. Dst has no null key yet: move src's null key to dst (no merge needed).
+/// 2. Dst already has a null key: merge src into dst using merge_agg_states.
+/// 3. Src has no null key: no-op.
+/// After merge, clears the src null key slot.
+template <typename HashTable>
+static void merge_null_key(HashTable& dst_data, HashTable& src_data, bool
use_simple_count,
+ const std::vector<AggFnEvaluator*>& evaluators,
const Sizes& offsets,
+ Arena& arena) {
+ if constexpr (has_nullable_key_v<HashTable>) {
+ if (!get_null_key_flag(src_data)) {
+ return;
+ }
+ auto src_null = get_null_key_agg_data(src_data);
+ if (!src_null) {
+ return;
+ }
+ if (!get_null_key_flag(dst_data)) {
+ // Dst has no null key yet — move src's null key to dst.
+ set_null_key_flag(dst_data, true);
+ set_null_key_agg_data(dst_data, src_null);
+ } else {
+ // Both have null keys — merge src into dst.
+ auto& dst_null_ref = get_null_key_agg_data_ref(dst_data);
+ merge_agg_states(dst_null_ref, src_null, use_simple_count,
evaluators, offsets, arena);
+ }
+ set_null_key_agg_data(src_data, nullptr);
+ set_null_key_flag(src_data, false);
+ }
+}
+
+BucketedAggLocalState::BucketedAggLocalState(RuntimeState* state,
OperatorXBase* parent)
+ : Base(state, parent) {}
+
+Status BucketedAggLocalState::init(RuntimeState* state, LocalStateInfo& info) {
+ RETURN_IF_ERROR(Base::init(state, info));
+ SCOPED_TIMER(exec_time_counter());
+ SCOPED_TIMER(_init_timer);
+
+ _task_idx = info.task_idx;
+
+ _get_results_timer = ADD_TIMER(custom_profile(), "GetResultsTime");
+ _hash_table_iterate_timer = ADD_TIMER(custom_profile(),
"HashTableIterateTime");
+ _insert_keys_to_column_timer = ADD_TIMER(custom_profile(),
"InsertKeysToColumnTime");
+ _insert_values_to_column_timer = ADD_TIMER(custom_profile(),
"InsertValuesToColumnTime");
+ _merge_timer = ADD_TIMER(custom_profile(), "MergeTime");
+
+ return Status::OK();
+}
+
+Status BucketedAggLocalState::close(RuntimeState* state) {
+ SCOPED_TIMER(exec_time_counter());
+ SCOPED_TIMER(_close_timer);
+ if (_closed) {
+ return Status::OK();
+ }
+
+ // Release any held per-bucket CAS lock. This can happen when the source
+ // is closed prematurely (e.g., LIMIT reached via reached_limit() while
+ // we were mid-output on a bucket). Without this, the other source instance
+ // would spin forever trying to acquire this bucket's lock.
+ if (_current_output_bucket >= 0) {
+ auto& bs = _shared_state->bucket_states[_current_output_bucket];
+ bs.output_done.store(true, std::memory_order_release);
+ bs.merge_in_progress.store(false, std::memory_order_release);
+ _current_output_bucket = -1;
+ _shared_state->state_generation.fetch_add(1,
std::memory_order_release);
+ _wake_up_other_sources();
+ }
+
+ return Base::close(state);
+}
+
+void BucketedAggLocalState::_make_nullable_output_key(Block* block) {
+ if (block->rows() != 0) {
+ for (auto cid : _shared_state->make_nullable_keys) {
+ block->get_by_position(cid).column =
make_nullable(block->get_by_position(cid).column);
+ block->get_by_position(cid).type =
make_nullable(block->get_by_position(cid).type);
+ }
+ }
+}
+
+void BucketedAggLocalState::_wake_up_other_sources() {
+ auto& shared_state = *_shared_state;
+ for (int i = 0; i < static_cast<int>(shared_state.source_deps.size());
++i) {
+ shared_state.source_deps[i]->set_ready();
+ }
+}
+
+int BucketedAggLocalState::_merge_bucket(int bucket, int merge_target) {
+ SCOPED_TIMER(_merge_timer);
+ auto& shared_state = *_shared_state;
+ auto& bs = shared_state.bucket_states[bucket];
+ // Other source instances may merge other buckets at the same time, so
aggregate
+ // function merges must allocate from this source instance's own arena.
+ DCHECK_LT(_task_idx, shared_state.source_merge_arenas.size());
+ auto& merge_arena = *shared_state.source_merge_arenas[_task_idx];
+
+ // Merge target's bucket is the destination.
+ auto& dst_agg_data =
*shared_state.per_instance_data[merge_target].bucket_agg_data[bucket];
+ int merged_count = 0;
+
+ std::visit(
+ Overload {
+ [&](std::monostate& arg) -> void {
+ // uninited — no data to merge
+ },
+ [&](auto& dst_method) -> void {
+ using AggMethodType =
std::decay_t<decltype(dst_method)>;
+ auto& dst_data = *dst_method.hash_table;
+
+ // Merge all finished sink instances (except
merge_target itself)
+ // into the merge target's bucket.
+ for (int inst_idx = 0; inst_idx <
shared_state.num_sink_instances;
+ ++inst_idx) {
+ if (inst_idx == merge_target) {
+ continue;
+ }
+ // Skip instances already merged for this bucket.
+ if (bs.merged_instances[inst_idx]) {
+ continue;
+ }
+ // Only merge sinks that have finished.
+ if (!shared_state.sink_finished[inst_idx].load(
+ std::memory_order_acquire)) {
+ continue;
+ }
+
+ auto& src_inst =
shared_state.per_instance_data[inst_idx];
+ auto& src_agg_data =
*src_inst.bucket_agg_data[bucket];
+
+ std::visit(
+ Overload {
+ [&](std::monostate& arg) -> void {
+ // Mark as merged even if
monostate (no data).
+ bs.merged_instances[inst_idx]
= true;
+ },
+ [&](auto& src_method) -> void {
+ using SrcMethodType =
+
std::decay_t<decltype(src_method)>;
+ if constexpr
(std::is_same_v<SrcMethodType,
+
AggMethodType>) {
+ auto& src_data =
*src_method.hash_table;
+
+ ++merged_count;
+
+ // Direct merge: iterate
source hash table
+ // entries, emplace into
destination, and null
+ // out source entries in
one pass. This avoids
+ // allocating intermediate
vectors (keys,
+ // mappeds, hashes) and
eliminates the separate
+ // null-out traversal.
+ // The source slot is
cleared only after the
+ // state is owned by dst:
emplace and merge can
+ // throw, and then
shared-state cleanup must
+ // still find and destroy
the source state.
+ const bool
use_simple_count =
+
shared_state.use_simple_count;
+
src_data.for_each([&](const auto& key,
+
auto& mapped) {
+ if (!mapped) {
+ return;
+ }
+
+ typename
std::remove_reference_t<
+
decltype(dst_data)>::LookupResult
+ dst_it;
+ bool inserted = false;
+
hash_table_emplace(dst_data, key, dst_it,
Review Comment:
Fixed in e86c5594e2d. `BucketedAggSourceOperatorX::get_block_impl` now runs
under `SCOPED_PEAK_MEM(&estimate_memory_usage())`, the same mechanism
`AggSourceOperatorX` uses, so `OperatorX::get_reserve_mem_size` reserves at
least what the previous merge round allocated instead of only the minimum
operator memory. `_merge_bucket` also accounts for what it adds: the growth of
the merge target's bucket tables (buffer size before/after the emplaces) and
this source instance's merge arena go into the new
`MemoryUsageMergedHashTables` / `MemoryUsageMergeArena` counters and the
operator's `MemoryUsage` (verified in a profile: the source's `MemoryUsage` and
`MemoryUsageMergedHashTables` now show the merge growth; before they were
always 0).
For the record, the reservation only affects scheduling under a workload
group: bucketed aggregation is only planned with spill disabled, and
`PipelineTask::_try_to_reserve_memory` disables reservation and lets the task
run when a reserve fails for a non-spill query, so this is an
accounting/observability fix rather than a correctness one.
--
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]
