fix(query): avoid serial distribution contagion in joins

src/query/service/src/schedulers/fragments/plan_fragment.rs

@@ -90,80 +90,77 @@ impl PlanFragment {
     ) -> Result<()> {
         let mut fragment_actions = QueryFragmentActions::create(self.fragment_id);
 
-        match &self.fragment_type {
-            FragmentType::Root => {
+        match (&self.fragment_type, self.has_merge_input) {
+            (FragmentType::Root, _) => {
                 let action = QueryFragmentAction::create(
                     Fragmenter::get_local_executor(ctx),
                     self.plan.clone(),
                 );
                 fragment_actions.add_action(action);
             }
-            FragmentType::Intermediate => {
-                if self.has_merge_input {
-                    // Only the coordinator can consume the merge input. Other shuffle
-                    // destinations still need this fragment to receive remote data.
-                    let local_executor = Fragmenter::get_local_executor(ctx);
-                    let action =
-                        QueryFragmentAction::create(local_executor.clone(), self.plan.clone());
+            (FragmentType::Intermediate, false) => {
+                // Otherwise distribute the fragment to all the executors.
+                for executor in Fragmenter::get_executors(ctx) {
+                    let action = QueryFragmentAction::create(executor, self.plan.clone());
                     fragment_actions.add_action(action);
-
-                    if let Some(exchange) = &self.exchange {
-                        let mut empty_plan = self.plan.clone();
-                        let Some(exchange_sink) =
-                            ExchangeSink::from_mut_physical_plan(&mut empty_plan)
-                        else {
-                            return Err(ErrorCode::Internal(
-                                "Intermediate fragment exchange plan has no ExchangeSink",
-                            ));
-                        };
-                        exchange_sink.input = PhysicalPlan::new(ConstantTableScan {
-                            meta: PhysicalPlanMeta::new("ConstantTableScan"),
-                            values: exchange_sink
-                                .schema
-                                .fields()
-                                .iter()
-                                .map(|field| {
-                                    ColumnBuilder::with_capacity(field.data_type(), 0).build()
-                                })
-                                .collect(),
-                            num_rows: 0,
-                            output_schema: exchange_sink.schema.clone(),
-                        });
-
-                        for executor in exchange.get_destinations() {
-                            if executor != local_executor {
-                                fragment_actions.add_action(QueryFragmentAction::create(
-                                    executor,
-                                    empty_plan.clone(),
-                                ));
-                            }
-                        }
-                    }
-                } else {
-                    // Otherwise distribute the fragment to all the executors.
-                    for executor in Fragmenter::get_executors(ctx) {
-                        let action = QueryFragmentAction::create(executor, self.plan.clone());
-                        fragment_actions.add_action(action);
-                    }
                 }
             }
-            FragmentType::Source => {
+            (FragmentType::Source, false) => {
                 // Redistribute partitions
                 self.redistribute_source_fragment(ctx, &mut fragment_actions)?;
             }
-            FragmentType::MutationSource => {
+            (FragmentType::MutationSource, false) => {
                 self.redistribute_mutation_source(ctx, &mut fragment_actions)?;
             }
-            FragmentType::ReplaceInto => {
+            (FragmentType::ReplaceInto, false) => {
                 // Redistribute partitions
                 self.redistribute_replace_into(ctx, &mut fragment_actions)?;
             }
-            FragmentType::Compact => {
+            (FragmentType::Compact, false) => {
                 self.redistribute_compact(ctx, &mut fragment_actions)?;
             }
-            FragmentType::Recluster => {
+            (FragmentType::Recluster, false) => {
                 self.redistribute_recluster(ctx, &mut fragment_actions)?;
             }
+            (_, true) => {
+                // Only the coordinator can consume the merge input. Other exchange
+                // destinations still need this fragment to receive remote data.
+                let local_executor = Fragmenter::get_local_executor(ctx);
+                fragment_actions.add_action(QueryFragmentAction::create(
+                    local_executor.clone(),
+                    self.plan.clone(),
+                ));
+
+                if let Some(exchange) = &self.exchange {
+                    let mut empty_plan = self.plan.clone();
+                    let Some(exchange_sink) = ExchangeSink::from_mut_physical_plan(&mut empty_plan)
+                    else {
+                        return Err(ErrorCode::Internal(
+                            "Merge-input fragment exchange plan has no ExchangeSink",
+                        ));
+                    };
+                    exchange_sink.input = PhysicalPlan::new(ConstantTableScan {
+                        meta: PhysicalPlanMeta::new("ConstantTableScan"),
+                        values: exchange_sink
+                            .schema
+                            .fields()
+                            .iter()
+                            .map(|field| ColumnBuilder::with_capacity(field.data_type(), 0).build())
+                            .collect(),
+                        num_rows: 0,
+                        output_schema: exchange_sink.schema.clone(),
+                    });
+
+                    for executor in exchange.get_destinations() {
+                        if executor != local_executor {
+                            fragment_actions.add_action(QueryFragmentAction::create(
+                                executor,
+                                empty_plan.clone(),
+                            ));
+                        }
+                    }
+                }
+            }
         }
 
         if let Some(ref exchange) = self.exchange {

src/query/sql/src/planner/plans/join.rs

@@ -1053,6 +1053,27 @@ impl Join {
             && !ctx.get_cluster().is_empty()
             && self.spatial_join_candidate(rel_expr)?.is_some())
     }
+
+    fn broadcast_build_is_preferred(ctx: &dyn TableContext, rel_expr: &RelExpr) -> Result<bool> {
+        let settings = ctx.get_settings();
+        if settings.get_enforce_shuffle_join()? {
+            return Ok(false);
+        }
+
+        let left_cardinality = rel_expr.derive_cardinality_child(0)?.cardinality;
+        let right_cardinality = rel_expr.derive_cardinality_child(1)?.cardinality;
+        let broadcast_join_threshold = if settings.get_prefer_broadcast_join()? {
+            ctx.get_cluster().nodes.len().saturating_sub(1) as f64
+        } else {
+            // Use a very large value to prevent broadcast join.
+            1000.0
+        };
+
+        Ok(
+            right_cardinality * broadcast_join_threshold < left_cardinality
+                || settings.get_enforce_broadcast_join()?,
+        )
+    }
 }
 
 fn estimate_anti_join_cardinality(
@@ -1232,7 +1253,41 @@ impl Operator for Join {
             return Ok(required);
         }
 
-        // if join/probe side is Serial or this is a non-equi join, we use Serial distribution
+        // A Serial build can still be redistributed. Broadcasting a small build
+        // keeps the probe distributed instead of propagating Serial to both sides.
+        let has_only_non_equi_conditions =
+            self.equi_conditions.is_empty() && !self.non_equi_conditions.is_empty();
+        if ctx.get_cluster().nodes.len() > 1
+            && build_physical_prop.distribution == Distribution::Serial
+            && probe_physical_prop.distribution != Distribution::Serial
+            && !has_only_non_equi_conditions
+            && !matches!(
+                self.join_type,
+                JoinType::Right
+                    | JoinType::Full
+                    | JoinType::RightAnti
+                    | JoinType::RightSemi
+                    | JoinType::LeftMark
+                    | JoinType::RightSingle
+                    | JoinType::InnerAny
+                    | JoinType::LeftAny
+                    | JoinType::RightAny
+                    | JoinType::Asof
+                    | JoinType::LeftAsof
+                    | JoinType::RightAsof
+                    | JoinType::FullAsof
+            )
+            && Self::broadcast_build_is_preferred(ctx.as_ref(), rel_expr)?
+        {
+            required.distribution = if child_index == 1 {
+                Distribution::Broadcast
+            } else {
+                Distribution::Any
+            };
+            return Ok(required);
+        }
+
+        // If either side remains Serial or this is a non-equi join, use Serial distribution.
         if probe_physical_prop.distribution == Distribution::Serial
             || build_physical_prop.distribution == Distribution::Serial
             || (self.equi_conditions.is_empty() && !self.non_equi_conditions.is_empty())
@@ -1243,7 +1298,6 @@ impl Operator for Join {
         }
 
         // Try to use broadcast join
-        let settings = ctx.get_settings();
         if !matches!(
             self.join_type,
             JoinType::Right
@@ -1258,29 +1312,14 @@ impl Operator for Join {
                 | JoinType::LeftAsof
                 | JoinType::RightAsof
                 | JoinType::FullAsof
-        ) {
-            let left_stat_info = rel_expr.derive_cardinality_child(0)?;
-            let right_stat_info = rel_expr.derive_cardinality_child(1)?;
-            // The broadcast join is cheaper than the hash join when one input is at least (n − 1)× larger than the other
-            // where n is the number of servers in the cluster.
-            let broadcast_join_threshold = if settings.get_prefer_broadcast_join()? {
-                (ctx.get_cluster().nodes.len() - 1) as f64
+        ) && Self::broadcast_build_is_preferred(ctx.as_ref(), rel_expr)?
+        {
+            required.distribution = if child_index == 1 {
+                Distribution::Broadcast
             } else {
-                // Use a very large value to prevent broadcast join.
-                1000.0
+                Distribution::Any
             };
-            if !settings.get_enforce_shuffle_join()?
-                && (right_stat_info.cardinality * broadcast_join_threshold
-                    < left_stat_info.cardinality
-                    || settings.get_enforce_broadcast_join()?)
-            {
-                if child_index == 1 {
-                    required.distribution = Distribution::Broadcast;
-                } else {
-                    required.distribution = Distribution::Any;
-                }
-                return Ok(required);
-            }
+            return Ok(required);
         }
 
         // Otherwise, use hash shuffle

src/query/sql/tests/it/optimizer/distributed_join.rs

@@ -0,0 +1,152 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed 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.
+
+use std::collections::HashMap;
+
+use databend_common_catalog::table_context::TableContextSettings;
+use databend_common_exception::Result;
+
+use super::table_statistics;
+use crate::framework::LiteTableContext;
+use crate::framework::golden::SqlTestCase;
+use crate::framework::golden::open_golden_file;
+use crate::framework::golden::write_case_header;
+
+async fn write_distributed_case(
+    file: &mut impl std::io::Write,
+    case: &SqlTestCase,
+    probe_rows: u64,
+    build_rows: u64,
+) -> Result<()> {
+    let ctx = LiteTableContext::create().await?;
+    ctx.set_cluster_node_num(2);
+    ctx.set_table_warehouse_distribution(true);
+    let settings = ctx.get_settings();
+    settings.set_setting("disable_join_reorder".to_string(), "1".to_string())?;
+    ctx.register_table_sql_with_stats(
+        BIG_TABLE,
+        Some(table_statistics(probe_rows)),
+        HashMap::new(),
+        HashMap::new(),
+    )
+    .await?;
+    ctx.register_table_sql_with_stats(
+        SMALL_TABLE,
+        Some(table_statistics(build_rows)),
+        HashMap::new(),
+        HashMap::new(),
+    )
+    .await?;
+
+    let raw_plan = ctx.bind_sql(case.sql).await?;
+    let optimized_plan = ctx.optimize_plan(raw_plan.clone()).await?;
+
+    write_case_header(file, case)?;
+    writeln!(file, "raw_plan:")?;
+    writeln!(file, "{}", raw_plan.format_indent(Default::default())?)?;
+    writeln!(file, "optimized_plan:")?;
+    writeln!(
+        file,
+        "{}",
+        optimized_plan.format_indent(Default::default())?
+    )?;
+    writeln!(file)?;
+
+    Ok(())
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn test_serial_build_distribution() -> Result<()> {
+    let mut file = open_golden_file("optimizer", "distributed_join.txt")?;
+    let cases = [
+        (
+            SqlTestCase {
+                name: "small_serial_build_is_broadcast",
+                description: "A small scalar aggregate on the build side should be broadcast without merging the distributed probe.",
+                setup_sqls: &[],
+                sql: "SELECT b.k, s.max_k
+FROM big_table AS b
+LEFT JOIN (SELECT max(k) AS max_k FROM small_table) AS s
+ON b.k = s.max_k",
+            },
+            227_000_000,
+            24,
+        ),
+        (
+            SqlTestCase {
+                name: "cross_join_small_serial_build_is_broadcast",
+                description: "A small scalar aggregate should also be broadcast for a cross join without merging the distributed probe.",
+                setup_sqls: &[],
+                sql: "SELECT b.k, s.max_k
+FROM big_table AS b
+CROSS JOIN (SELECT max(k) AS max_k FROM small_table) AS s",
+            },
+            227_000_000,
+            24,
+        ),
+        (
+            SqlTestCase {
+                name: "constant_serial_build_is_broadcast",
+                description: "A constant Serial source should be broadcast while source partitioning keeps it single-produced.",
+                setup_sqls: &[],
+                sql: "SELECT b.k, c.x
+FROM big_table AS b
+CROSS JOIN (SELECT 1 AS x) AS c",
+            },
+            227_000_000,
+            24,
+        ),
+        (
+            SqlTestCase {
+                name: "small_sorted_serial_build_is_broadcast",
+                description: "A small build made Serial by a global window sort should be broadcast without merging the distributed probe.",
+                setup_sqls: &[],
+                sql: "SELECT b.k, s.k
+FROM big_table AS b
+LEFT JOIN (
+    SELECT k, row_number() OVER (ORDER BY k) AS row_num
+    FROM small_table
+) AS s
+ON b.k = s.k AND s.row_num <= 10",
+            },
+            227_000_000,
+            24,
+        ),
+        (
+            SqlTestCase {
+                name: "large_serial_build_remains_serial",
+                description: "A Serial build that is not smaller than the probe should not be broadcast unconditionally.",
+                setup_sqls: &[],
+                sql: "SELECT b.k, s.k
+FROM big_table AS b
+LEFT JOIN (
+    SELECT k, row_number() OVER () AS row_num
+    FROM small_table
+) AS s
+ON b.k = s.k",
+            },
+            1_000,
+            1_000,
+        ),
+    ];
+
+    for (case, probe_rows, build_rows) in &cases {
+        write_distributed_case(&mut file, case, *probe_rows, *build_rows).await?;
+    }
+
+    Ok(())
+}
+
+const BIG_TABLE: &str = "CREATE TABLE big_table (k BIGINT)";
+const SMALL_TABLE: &str = "CREATE TABLE small_table (k BIGINT)";