This is an automated email from the ASF dual-hosted git repository. BiteTheDDDDt pushed a commit to branch opt_perf_4.1 in repository https://gitbox.apache.org/repos/asf/doris.git
commit d888436c3d16bfc9524cf303a0889f83d7bd5ed9 Author: BiteTheDDDDt <[email protected]> AuthorDate: Mon Jun 22 16:00:16 2026 +0800 [improvement](runtime-filter) Limit broadcast runtime filter producers Issue Number: None Related PR: None Problem Summary: Broadcast join runtime filters are globally merged by accepting the first producer, but every build-side BE still creates and sends the same broadcast runtime filter. This wastes runtime filter build CPU and network traffic. Add a Nereids-side producer selection controlled by runtime_filter_broadcast_join_producer_num so only up to the configured number of BEs receive broadcast runtime filter producer descriptors with remote targets. When the selected runtime-filter m [...] Add session variable runtime_filter_broadcast_join_producer_num to limit Nereids broadcast join runtime filter producers. - Test: - Unit Test: ./run-fe-ut.sh --run org.apache.doris.qe.runtime.RuntimeFiltersThriftBuilderTest - Unit Test: ./run-fe-ut.sh --run org.apache.doris.qe.runtime.RuntimeFiltersThriftBuilderTest,org.apache.doris.qe.SessionVariablesTest - Manual Check: git diff --check upstream/master - Behavior changed: Yes. Nereids broadcast join runtime filter producers with remote targets are limited by runtime_filter_broadcast_join_producer_num, defaulting to 3, and prefer the runtime-filter merge coordinator when it is also a build worker. - Does this need documentation: No (cherry picked from commit f1e0e560772c56f86e30ca355295fb4c9b3edcff) --- .../java/org/apache/doris/qe/SessionVariable.java | 19 ++ .../qe/runtime/RuntimeFiltersThriftBuilder.java | 120 +++++++++- .../doris/qe/runtime/ThriftPlansBuilder.java | 13 +- .../org/apache/doris/qe/SessionVariablesTest.java | 16 ++ .../runtime/RuntimeFiltersThriftBuilderTest.java | 250 +++++++++++++++++++++ 5 files changed, 407 insertions(+), 11 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java index 5efec5d589d..e3752257797 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java @@ -257,6 +257,9 @@ public class SessionVariable implements Serializable, Writable { // if the right table is greater than this value in the hash join, we will ignore IN filter public static final String RUNTIME_FILTER_MAX_IN_NUM = "runtime_filter_max_in_num"; + public static final String RUNTIME_FILTER_BROADCAST_JOIN_PRODUCER_NUM = + "runtime_filter_broadcast_join_producer_num"; + public static final String ENABLE_SYNC_RUNTIME_FILTER_SIZE = "enable_sync_runtime_filter_size"; public static final String ENABLE_PARALLEL_RESULT_SINK = "enable_parallel_result_sink"; @@ -1738,6 +1741,14 @@ public class SessionVariable implements Serializable, Writable { @VariableMgr.VarAttr(name = ENABLE_SYNC_RUNTIME_FILTER_SIZE, needForward = true, fuzzy = true) private boolean enableSyncRuntimeFilterSize = true; + @VariableMgr.VarAttr(name = RUNTIME_FILTER_BROADCAST_JOIN_PRODUCER_NUM, needForward = true, + description = {"控制 Nereids 分布式规划中每个 broadcast join runtime filter 的生产 BE 数量。" + + "设置为小于等于 0 时不限制。Legacy Coordinator 路径保持原行为。", + "Controls the number of producer BEs for each broadcast join runtime filter in " + + "the Nereids distributed planner. Values less than or equal to 0 disable the limit. " + + "The legacy Coordinator path keeps the existing behavior."}) + private int runtimeFilterBroadcastJoinProducerNum = 3; + @VariableMgr.VarAttr(name = "runtime_filter_max_build_row_count", needForward = true, fuzzy = false) public long runtimeFilterMaxBuildRowCount = 64L * 1024L * 1024L; @@ -4724,6 +4735,14 @@ public class SessionVariable implements Serializable, Writable { this.runtimeFilterMaxInNum = runtimeFilterMaxInNum; } + public int getRuntimeFilterBroadcastJoinProducerNum() { + return runtimeFilterBroadcastJoinProducerNum; + } + + public void setRuntimeFilterBroadcastJoinProducerNum(int runtimeFilterBroadcastJoinProducerNum) { + this.runtimeFilterBroadcastJoinProducerNum = runtimeFilterBroadcastJoinProducerNum; + } + public void setEnableLocalShuffle(boolean enableLocalShuffle) { this.enableLocalShuffle = enableLocalShuffle; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/RuntimeFiltersThriftBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/RuntimeFiltersThriftBuilder.java index fca5461fe1c..33182b6a912 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/RuntimeFiltersThriftBuilder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/RuntimeFiltersThriftBuilder.java @@ -26,16 +26,23 @@ import org.apache.doris.planner.RuntimeFilter; import org.apache.doris.planner.RuntimeFilterId; import org.apache.doris.system.Backend; import org.apache.doris.thrift.TNetworkAddress; +import org.apache.doris.thrift.TPlanFragment; +import org.apache.doris.thrift.TPlanNode; +import org.apache.doris.thrift.TRuntimeFilterDesc; import org.apache.doris.thrift.TRuntimeFilterParams; import org.apache.doris.thrift.TRuntimeFilterTargetParamsV2; +import com.google.common.base.Preconditions; import com.google.common.collect.Maps; import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.ThreadLocalRandom; import java.util.stream.Collectors; /** RuntimeFiltersThriftBuilder */ @@ -47,18 +54,58 @@ public class RuntimeFiltersThriftBuilder { private final Set<Integer> broadcastRuntimeFilterIds; private final Map<RuntimeFilterId, List<RuntimeFilterTarget>> ridToTargets; private final Map<RuntimeFilterId, Integer> ridToBuilderNum; + private final boolean limitBroadcastRuntimeFilterProducers; + private final Map<Long, List<Integer>> workerIdToBroadcastRuntimeFilterIds; + private final Map<Integer, Integer> broadcastRuntimeFilterIdToBuilderNodeId; private RuntimeFiltersThriftBuilder( TNetworkAddress mergeAddress, List<RuntimeFilter> runtimeFilters, AssignedJob mergeInstance, Set<Integer> broadcastRuntimeFilterIds, Map<RuntimeFilterId, List<RuntimeFilterTarget>> ridToTargets, - Map<RuntimeFilterId, Integer> ridToBuilderNum) { + Map<RuntimeFilterId, Integer> ridToBuilderNum, + boolean limitBroadcastRuntimeFilterProducers, + Map<Long, List<Integer>> workerIdToBroadcastRuntimeFilterIds, + Map<Integer, Integer> broadcastRuntimeFilterIdToBuilderNodeId) { this.mergeAddress = mergeAddress; this.runtimeFilters = runtimeFilters; this.mergeInstance = mergeInstance; this.broadcastRuntimeFilterIds = broadcastRuntimeFilterIds; this.ridToTargets = ridToTargets; this.ridToBuilderNum = ridToBuilderNum; + this.limitBroadcastRuntimeFilterProducers = limitBroadcastRuntimeFilterProducers; + this.workerIdToBroadcastRuntimeFilterIds = workerIdToBroadcastRuntimeFilterIds; + this.broadcastRuntimeFilterIdToBuilderNodeId = broadcastRuntimeFilterIdToBuilderNodeId; + } + + public void pruneBroadcastRuntimeFilterProducers( + TPlanFragment planFragment, DistributedPlanWorker worker) { + if (!limitBroadcastRuntimeFilterProducers) { + return; + } + Set<Integer> producerFilterIds = new HashSet<>( + workerIdToBroadcastRuntimeFilterIds.getOrDefault(worker.id(), Collections.emptyList())); + if (planFragment.isSetPlan()) { + for (TPlanNode node : planFragment.getPlan().getNodes()) { + if (node.isSetRuntimeFilters()) { + node.setRuntimeFilters(pruneRuntimeFilterDescs( + node.getNodeId(), node.getRuntimeFilters(), producerFilterIds)); + } + } + } + } + + private List<TRuntimeFilterDesc> pruneRuntimeFilterDescs( + int nodeId, List<TRuntimeFilterDesc> runtimeFilterDescs, Set<Integer> producerFilterIds) { + List<TRuntimeFilterDesc> selectedRuntimeFilterDescs = new ArrayList<>(runtimeFilterDescs.size()); + for (TRuntimeFilterDesc desc : runtimeFilterDescs) { + Integer builderNodeId = broadcastRuntimeFilterIdToBuilderNodeId.get(desc.filter_id); + if (!desc.is_broadcast_join || !desc.has_remote_targets + || builderNodeId == null || builderNodeId != nodeId + || producerFilterIds.contains(desc.filter_id)) { + selectedRuntimeFilterDescs.add(desc); + } + } + return selectedRuntimeFilterDescs; } public boolean isMergeRuntimeFilterInstance(AssignedJob instance) { @@ -104,11 +151,22 @@ public class RuntimeFiltersThriftBuilder { public static RuntimeFiltersThriftBuilder compute( List<RuntimeFilter> runtimeFilters, List<PipelineDistributedPlan> distributedPlans) { + return compute(runtimeFilters, distributedPlans, 0); + } + + public static RuntimeFiltersThriftBuilder compute( + List<RuntimeFilter> runtimeFilters, List<PipelineDistributedPlan> distributedPlans, + int broadcastRuntimeFilterProducerNum) { PipelineDistributedPlan topMostPlan = distributedPlans.get(distributedPlans.size() - 1); AssignedJob mergeInstance = topMostPlan.getInstanceJobs().get(0); BackendWorker worker = (BackendWorker) mergeInstance.getAssignedWorker(); TNetworkAddress mergeAddress = new TNetworkAddress(worker.host(), worker.brpcPort()); + Map<Integer, RuntimeFilter> idToRuntimeFilter = runtimeFilters + .stream() + .collect(Collectors.toMap(r -> r.getFilterId().asInt(), r -> r, (left, right) -> left, + LinkedHashMap::new)); + Set<Integer> broadcastRuntimeFilterIds = runtimeFilters .stream() .filter(RuntimeFilter::isBroadcast) @@ -117,6 +175,10 @@ public class RuntimeFiltersThriftBuilder { Map<RuntimeFilterId, List<RuntimeFilterTarget>> ridToTargetParam = Maps.newLinkedHashMap(); Map<RuntimeFilterId, Integer> ridToBuilderNum = Maps.newLinkedHashMap(); + Map<Integer, List<BackendWorker>> builderNodeToProducerWorkers = Maps.newLinkedHashMap(); + Map<Long, List<Integer>> workerIdToBroadcastRuntimeFilterIds = Maps.newLinkedHashMap(); + Map<Integer, Integer> broadcastRuntimeFilterIdToBuilderNodeId = Maps.newLinkedHashMap(); + boolean limitBroadcastRuntimeFilterProducers = broadcastRuntimeFilterProducerNum > 0; for (PipelineDistributedPlan plan : distributedPlans) { PlanFragment fragment = plan.getFragmentJob().getFragment(); // Transform <fragment, runtimeFilterId> to <runtimeFilterId, fragment> @@ -132,19 +194,61 @@ public class RuntimeFiltersThriftBuilder { } } + List<BackendWorker> builderWorkers = collectDistinctBackendWorkers(plan.getInstanceJobs()); + int distinctWorkerNum = builderWorkers.size(); for (RuntimeFilterId rid : fragment.getBuilderRuntimeFilterIds()) { - int distinctWorkerNum = (int) plan.getInstanceJobs() - .stream() - .map(AssignedJob::getAssignedWorker) - .map(DistributedPlanWorker::id) - .distinct() - .count(); ridToBuilderNum.merge(rid, distinctWorkerNum, Integer::sum); + RuntimeFilter rf = idToRuntimeFilter.get(rid.asInt()); + if (limitBroadcastRuntimeFilterProducers + && rf != null && rf.isBroadcast() && rf.hasRemoteTargets()) { + int builderNodeId = rf.getBuilderNode().getId().asInt(); + broadcastRuntimeFilterIdToBuilderNodeId.put(rid.asInt(), builderNodeId); + List<BackendWorker> producerWorkers = builderNodeToProducerWorkers.computeIfAbsent( + builderNodeId, + id -> selectBroadcastRuntimeFilterProducerWorkers( + builderWorkers, broadcastRuntimeFilterProducerNum, worker)); + for (BackendWorker producerWorker : producerWorkers) { + workerIdToBroadcastRuntimeFilterIds.computeIfAbsent( + producerWorker.id(), id -> new ArrayList<>()).add(rid.asInt()); + } + } } } return new RuntimeFiltersThriftBuilder( mergeAddress, runtimeFilters, mergeInstance, - broadcastRuntimeFilterIds, ridToTargetParam, ridToBuilderNum); + broadcastRuntimeFilterIds, ridToTargetParam, ridToBuilderNum, + limitBroadcastRuntimeFilterProducers, workerIdToBroadcastRuntimeFilterIds, + broadcastRuntimeFilterIdToBuilderNodeId); + } + + static List<BackendWorker> collectDistinctBackendWorkers(List<AssignedJob> instanceJobs) { + Map<Long, BackendWorker> workerMap = Maps.newLinkedHashMap(); + for (AssignedJob instanceJob : instanceJobs) { + BackendWorker worker = (BackendWorker) instanceJob.getAssignedWorker(); + workerMap.putIfAbsent(worker.id(), worker); + } + return new ArrayList<>(workerMap.values()); + } + + static List<BackendWorker> selectBroadcastRuntimeFilterProducerWorkers( + List<BackendWorker> workers, int producerNum, BackendWorker preferredWorker) { + Preconditions.checkArgument(producerNum > 0, + "broadcast runtime filter producer num must be positive"); + if (workers.size() <= producerNum) { + return workers; + } + List<BackendWorker> selectedWorkers = new ArrayList<>(producerNum); + for (BackendWorker worker : workers) { + if (worker.equals(preferredWorker)) { + selectedWorkers.add(worker); + break; + } + } + List<BackendWorker> remainingWorkers = new ArrayList<>(workers); + remainingWorkers.removeAll(selectedWorkers); + Collections.shuffle(remainingWorkers, ThreadLocalRandom.current()); + selectedWorkers.addAll(remainingWorkers.subList(0, producerNum - selectedWorkers.size())); + return selectedWorkers; } public static class RuntimeFilterTarget { diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java index eba3343bd03..7af7556d89c 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/ThriftPlansBuilder.java @@ -113,8 +113,12 @@ public class ThriftPlansBuilder { // we should set runtime predicate first, then we can use heap sort and to thrift setRuntimePredicateIfNeed(coordinatorContext.scanNodes); + int broadcastRuntimeFilterProducerNum = coordinatorContext.connectContext == null + ? 0 + : coordinatorContext.connectContext.getSessionVariable() + .getRuntimeFilterBroadcastJoinProducerNum(); RuntimeFiltersThriftBuilder runtimeFiltersThriftBuilder = RuntimeFiltersThriftBuilder.compute( - coordinatorContext.runtimeFilters, distributedPlans); + coordinatorContext.runtimeFilters, distributedPlans, broadcastRuntimeFilterProducerNum); Supplier<List<TTopnFilterDesc>> topNFilterThriftSupplier = topNFilterToThrift(coordinatorContext.topnFilters); @@ -134,7 +138,8 @@ public class ThriftPlansBuilder { TPipelineFragmentParams currentFragmentParam = fragmentToThriftIfAbsent( currentFragmentPlan, instanceJob, workerToCurrentFragment, instancesPerWorker, exchangeSenderNum, sharedFileScanRangeParams, - workerProcessInstanceNum, fragmentToNotifyClose, coordinatorContext); + workerProcessInstanceNum, fragmentToNotifyClose, coordinatorContext, + runtimeFiltersThriftBuilder); TPipelineInstanceParams instanceParam = instanceToThrift( currentFragmentParam, instanceJob, currentInstanceIndex++); @@ -344,7 +349,8 @@ public class ThriftPlansBuilder { Map<Integer, TFileScanRangeParams> fileScanRangeParamsMap, Multiset<DistributedPlanWorker> workerProcessInstanceNum, Set<Integer> fragmentToNotifyClose, - CoordinatorContext coordinatorContext) { + CoordinatorContext coordinatorContext, + RuntimeFiltersThriftBuilder runtimeFiltersThriftBuilder) { DistributedPlanWorker worker = assignedJob.getAssignedWorker(); return workerToFragmentParams.computeIfAbsent(worker, w -> { PlanFragment fragment = fragmentPlan.getFragmentJob().getFragment(); @@ -399,6 +405,7 @@ public class ThriftPlansBuilder { params.setSendQueryStatisticsWithEveryBatch(fragment.isTransferQueryStatisticsWithEveryBatch()); TPlanFragment planThrift = fragment.toThrift(); + runtimeFiltersThriftBuilder.pruneBroadcastRuntimeFilterProducers(planThrift, worker); planThrift.query_cache_param = fragment.queryCacheParam; params.setFragment(planThrift); params.setLocalParams(Lists.newArrayList()); diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java index e5ab3243118..93252963f22 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java @@ -167,6 +167,22 @@ public class SessionVariablesTest extends TestWithFeService { () -> sessionVar.checkInsertVisibleTimeoutReturnMode("")); } + @Test + public void testRuntimeFilterBroadcastJoinProducerNumDescription() throws Exception { + SessionVariable sessionVar = new SessionVariable(); + Assertions.assertEquals(3, sessionVar.getRuntimeFilterBroadcastJoinProducerNum()); + + Field field = SessionVariable.class.getDeclaredField("runtimeFilterBroadcastJoinProducerNum"); + VariableMgr.VarAttr varAttr = field.getAnnotation(VariableMgr.VarAttr.class); + Assertions.assertArrayEquals(new String[] { + "控制 Nereids 分布式规划中每个 broadcast join runtime filter 的生产 BE 数量。" + + "设置为小于等于 0 时不限制。Legacy Coordinator 路径保持原行为。", + "Controls the number of producer BEs for each broadcast join runtime filter in " + + "the Nereids distributed planner. Values less than or equal to 0 disable the limit. " + + "The legacy Coordinator path keeps the existing behavior." + }, varAttr.description()); + } + @Test public void testForceEagerAggHintParseWhenSetSessionVariable() throws Exception { SessionVariable sessionVar = new SessionVariable(); diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/RuntimeFiltersThriftBuilderTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/RuntimeFiltersThriftBuilderTest.java new file mode 100644 index 00000000000..850c216a840 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/RuntimeFiltersThriftBuilderTest.java @@ -0,0 +1,250 @@ +// 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.doris.qe.runtime; + +import org.apache.doris.common.IdGenerator; +import org.apache.doris.nereids.StatementContext; +import org.apache.doris.nereids.trees.AbstractTreeNode; +import org.apache.doris.nereids.trees.plans.distribute.DistributeContext; +import org.apache.doris.nereids.trees.plans.distribute.PipelineDistributedPlan; +import org.apache.doris.nereids.trees.plans.distribute.worker.BackendWorker; +import org.apache.doris.nereids.trees.plans.distribute.worker.job.AssignedJob; +import org.apache.doris.nereids.trees.plans.distribute.worker.job.DefaultScanSource; +import org.apache.doris.nereids.trees.plans.distribute.worker.job.UnassignedJob; +import org.apache.doris.planner.ExchangeNode; +import org.apache.doris.planner.PlanFragment; +import org.apache.doris.planner.PlanFragmentId; +import org.apache.doris.planner.PlanNode; +import org.apache.doris.planner.PlanNodeId; +import org.apache.doris.planner.RuntimeFilter; +import org.apache.doris.planner.RuntimeFilterId; +import org.apache.doris.planner.ScanNode; +import org.apache.doris.system.Backend; +import org.apache.doris.thrift.TPlan; +import org.apache.doris.thrift.TPlanFragment; +import org.apache.doris.thrift.TPlanNode; +import org.apache.doris.thrift.TRuntimeFilterDesc; +import org.apache.doris.thrift.TUniqueId; + +import com.google.common.collect.ArrayListMultimap; +import com.google.common.collect.ImmutableSetMultimap; +import com.google.common.collect.ListMultimap; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; + +public class RuntimeFiltersThriftBuilderTest { + @Test + public void testSelectBroadcastRuntimeFilterProducerWorkersPreferMergeWorker() { + BackendWorker worker1 = newBackendWorker(1); + BackendWorker worker2 = newBackendWorker(2); + BackendWorker worker3 = newBackendWorker(3); + + List<BackendWorker> selectedWorkers = RuntimeFiltersThriftBuilder.selectBroadcastRuntimeFilterProducerWorkers( + Arrays.asList(worker1, worker2, worker3), 1, worker2); + + Assertions.assertEquals(Collections.singletonList(worker2), selectedWorkers); + } + + @Test + public void testSelectBroadcastRuntimeFilterProducerWorkersIgnoreNonCandidatePreferredWorker() { + BackendWorker worker1 = newBackendWorker(1); + BackendWorker worker2 = newBackendWorker(2); + BackendWorker worker3 = newBackendWorker(3); + + List<BackendWorker> selectedWorkers = RuntimeFiltersThriftBuilder.selectBroadcastRuntimeFilterProducerWorkers( + Arrays.asList(worker1, worker2), 1, worker3); + + Assertions.assertEquals(1, selectedWorkers.size()); + Assertions.assertTrue(Arrays.asList(worker1, worker2).contains(selectedWorkers.get(0))); + Assertions.assertFalse(selectedWorkers.contains(worker3)); + } + + @Test + public void testPruneBroadcastRuntimeFilterProducers() { + BackendWorker worker1 = newBackendWorker(1); + BackendWorker worker2 = newBackendWorker(2); + BackendWorker worker3 = newBackendWorker(3); + + IdGenerator<RuntimeFilterId> idGenerator = RuntimeFilterId.createGenerator(); + RuntimeFilterId broadcastRid1 = idGenerator.getNextId(); + RuntimeFilterId broadcastRid2 = idGenerator.getNextId(); + RuntimeFilterId localBroadcastRid = idGenerator.getNextId(); + RuntimeFilterId shuffleRid = idGenerator.getNextId(); + RuntimeFilter broadcastRf1 = newRuntimeFilter(broadcastRid1, 10, true, true); + RuntimeFilter broadcastRf2 = newRuntimeFilter(broadcastRid2, 10, true, true); + RuntimeFilter localBroadcastRf = newRuntimeFilter(localBroadcastRid, 10, true, false); + RuntimeFilter shuffleRf = newRuntimeFilter(shuffleRid, 11, false, true); + + PlanFragment fragment = newFragment(0, broadcastRid1, broadcastRid2, localBroadcastRid, shuffleRid); + PipelineDistributedPlan distributedPlan = newDistributedPlan(fragment, worker1, worker2, worker3); + RuntimeFiltersThriftBuilder builder = RuntimeFiltersThriftBuilder.compute( + Arrays.asList(broadcastRf1, broadcastRf2, localBroadcastRf, shuffleRf), + Collections.singletonList(distributedPlan), 1); + + int selectedWorkerNum = 0; + for (BackendWorker worker : Arrays.asList(worker1, worker2, worker3)) { + TPlanFragment planFragment = newPlanFragment( + newPlanNode(10, + newRuntimeFilterDesc(broadcastRid1, true, true), + newRuntimeFilterDesc(broadcastRid2, true, true), + newRuntimeFilterDesc(localBroadcastRid, true, false)), + newPlanNode(11, + newRuntimeFilterDesc(shuffleRid, false, true)), + newPlanNode(20, + newRuntimeFilterDesc(broadcastRid1, true, true), + newRuntimeFilterDesc(broadcastRid2, true, true))); + builder.pruneBroadcastRuntimeFilterProducers(planFragment, worker); + + List<Integer> builderFilterIds = planFragment.getPlan().getNodes().get(0).getRuntimeFilters() + .stream() + .map(desc -> desc.filter_id) + .collect(Collectors.toList()); + List<Integer> shuffleFilterIds = planFragment.getPlan().getNodes().get(1).getRuntimeFilters() + .stream() + .map(desc -> desc.filter_id) + .collect(Collectors.toList()); + List<Integer> targetFilterIds = planFragment.getPlan().getNodes().get(2).getRuntimeFilters() + .stream() + .map(desc -> desc.filter_id) + .collect(Collectors.toList()); + Assertions.assertEquals(Collections.singletonList(shuffleRid.asInt()), shuffleFilterIds); + Assertions.assertEquals(new HashSet<>(Arrays.asList( + broadcastRid1.asInt(), broadcastRid2.asInt())), new HashSet<>(targetFilterIds)); + if (builderFilterIds.contains(broadcastRid1.asInt())) { + selectedWorkerNum++; + Assertions.assertEquals(new HashSet<>(Arrays.asList( + broadcastRid1.asInt(), broadcastRid2.asInt(), localBroadcastRid.asInt())), + new HashSet<>(builderFilterIds)); + } else { + Assertions.assertFalse(builderFilterIds.contains(broadcastRid2.asInt())); + Assertions.assertEquals(Collections.singletonList(localBroadcastRid.asInt()), builderFilterIds); + } + } + Assertions.assertEquals(1, selectedWorkerNum); + } + + private BackendWorker newBackendWorker(long id) { + Backend backend = new Backend(id, "host" + id, (int) (9000 + id)); + backend.setBePort((int) (8000 + id)); + backend.setBrpcPort((int) (7000 + id)); + return new BackendWorker(0, backend); + } + + private PipelineDistributedPlan newDistributedPlan(PlanFragment fragment, BackendWorker... workers) { + TestUnassignedJob unassignedJob = new TestUnassignedJob(); + unassignedJob.fragment = fragment; + List<AssignedJob> assignedJobs = Arrays.stream(workers) + .map(worker -> unassignedJob.assignWorkerAndDataSources( + 0, new TUniqueId(), worker, DefaultScanSource.empty())) + .collect(Collectors.toList()); + return new PipelineDistributedPlan(unassignedJob, assignedJobs, ImmutableSetMultimap.of()); + } + + private RuntimeFilter newRuntimeFilter(RuntimeFilterId rid, int builderNodeId, + boolean isBroadcast, boolean hasRemoteTargets) { + PlanNode builderNode = Mockito.mock(PlanNode.class); + Mockito.when(builderNode.getId()).thenReturn(new PlanNodeId(builderNodeId)); + + RuntimeFilter runtimeFilter = Mockito.mock(RuntimeFilter.class); + Mockito.when(runtimeFilter.getFilterId()).thenReturn(rid); + Mockito.when(runtimeFilter.getBuilderNode()).thenReturn(builderNode); + Mockito.when(runtimeFilter.isBroadcast()).thenReturn(isBroadcast); + Mockito.when(runtimeFilter.hasRemoteTargets()).thenReturn(hasRemoteTargets); + return runtimeFilter; + } + + private TPlanFragment newPlanFragment(TPlanNode... nodes) { + TPlan plan = new TPlan(); + plan.setNodes(Arrays.asList(nodes)); + TPlanFragment fragment = new TPlanFragment(); + fragment.setPlan(plan); + return fragment; + } + + private TPlanNode newPlanNode(int nodeId, TRuntimeFilterDesc... runtimeFilterDescs) { + TPlanNode node = new TPlanNode(); + node.setNodeId(nodeId); + node.setRuntimeFilters(Arrays.asList(runtimeFilterDescs)); + return node; + } + + private TRuntimeFilterDesc newRuntimeFilterDesc( + RuntimeFilterId rid, boolean isBroadcast, boolean hasRemoteTargets) { + TRuntimeFilterDesc desc = new TRuntimeFilterDesc(); + desc.setFilterId(rid.asInt()); + desc.setIsBroadcastJoin(isBroadcast); + desc.setHasRemoteTargets(hasRemoteTargets); + return desc; + } + + private PlanFragment newFragment(int fragmentId, RuntimeFilterId... builderRuntimeFilterIds) { + PlanFragment fragment = Mockito.mock(PlanFragment.class); + Mockito.when(fragment.getFragmentId()).thenReturn(new PlanFragmentId(fragmentId)); + Set<RuntimeFilterId> builderIds = new HashSet<>(Arrays.asList(builderRuntimeFilterIds)); + Mockito.when(fragment.getBuilderRuntimeFilterIds()).thenReturn(builderIds); + Mockito.when(fragment.getTargetRuntimeFilterIds()).thenReturn(Collections.emptySet()); + return fragment; + } + + private static final class TestUnassignedJob extends AbstractTreeNode<UnassignedJob> implements UnassignedJob { + private PlanFragment fragment; + + private TestUnassignedJob() { + super(Collections.emptyList()); + } + + @Override + public StatementContext getStatementContext() { + return null; + } + + @Override + public PlanFragment getFragment() { + return fragment; + } + + @Override + public List<ScanNode> getScanNodes() { + return Collections.emptyList(); + } + + @Override + public ListMultimap<ExchangeNode, UnassignedJob> getExchangeToChildJob() { + return ArrayListMultimap.create(); + } + + @Override + public List<AssignedJob> computeAssignedJobs( + DistributeContext distributeContext, ListMultimap<ExchangeNode, AssignedJob> inputJobs) { + return Collections.emptyList(); + } + + @Override + public UnassignedJob withChildren(List<UnassignedJob> children) { + return this; + } + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
