feat: optimize broadcast join (#11365)

* feat: optimize broadcast join * add threshold * fix test
databendlabs · May 10, 2023 · 4e66b31 · 4e66b31 · vercel · May 10, 2023
1 parent c965372
commit 4e66b31
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 6 deletions.
diff --git a/src/query/sql/src/planner/optimizer/property/enforcer.rs b/src/query/sql/src/planner/optimizer/property/enforcer.rs
@@ -47,6 +47,15 @@ pub fn require_property(
             // If the child is join probe side and join type is broadcast join
             // We should wrap the child with Random exchange to make it partition to all nodes
             if index == 0 && required.distribution == Distribution::Broadcast {
+                if optimized_expr
+                    .child(0)?
+                    .children()
+                    .iter()
+                    .any(check_partition)
+                {
+                    children.push(optimized_expr.child(index)?.clone());
+                    continue;
+                }
                 let enforced_child =
                     enforce_property(optimized_expr.child(index)?, &RequiredProperty {
                         distribution: Distribution::Any,
@@ -130,3 +139,17 @@ fn check_merge(s_expr: &SExpr) -> bool {
     }
     false
 }
+
+fn check_partition(s_expr: &SExpr) -> bool {
+    if let RelOperator::Exchange(op) = &s_expr.plan {
+        if matches!(op, Exchange::Random | Exchange::Hash(_)) {
+            return true;
+        }
+    }
+    for child in s_expr.children() {
+        if check_partition(child) {
+            return true;
+        }
+    }
+    false
+}
diff --git a/src/query/sql/src/planner/plans/join.rs b/src/query/sql/src/planner/plans/join.rs
@@ -40,6 +40,8 @@ use crate::plans::RelOp;
 use crate::plans::ScalarExpr;
 use crate::IndexType;
 
+const BROADCAST_JOIN_THRESHOLD: f64 = 20.0;
+
 #[derive(Clone, Debug, Eq, PartialEq, Hash, serde::Serialize, serde::Deserialize)]
 pub enum JoinType {
     Inner,
@@ -342,6 +344,11 @@ impl Operator for Join {
             (Distribution::Random, _) => Ok(PhysicalProperty {
                 distribution: build_prop.distribution.clone(),
             }),
+            // If both sides are broadcast, which means broadcast join is enabled, to make sure the current join is broadcast, should return Random.
+            // Then required proper is broadcast, and the join will be broadcast.
+            (Distribution::Broadcast, Distribution::Broadcast) => Ok(PhysicalProperty {
+                distribution: Distribution::Random,
+            }),
             // Otherwise pass through probe side.
             _ => Ok(PhysicalProperty {
                 distribution: probe_prop.distribution.clone(),
@@ -414,6 +421,7 @@ impl Operator for Join {
         {
             // TODO(leiysky): we can enforce redistribution here
             required.distribution = Distribution::Serial;
+            return Ok(required);
         } else if ctx.get_settings().get_prefer_broadcast_join()?
             && !matches!(
                 self.join_type,
@@ -423,11 +431,15 @@ impl Operator for Join {
                     | JoinType::RightSemi
                     | JoinType::RightMark
             )
-            && probe_physical_prop.distribution != Distribution::Broadcast
-            && build_physical_prop.distribution != Distribution::Broadcast
         {
-            required.distribution = Distribution::Broadcast;
-        } else if child_index == 0 {
+            let left_stat_info = rel_expr.derive_cardinality_child(0)?;
+            let right_stat_info = rel_expr.derive_cardinality_child(1)?;
+            if right_stat_info.cardinality * BROADCAST_JOIN_THRESHOLD < left_stat_info.cardinality {
+                required.distribution = Distribution::Broadcast;
+                return Ok(required);
+            }
+        }
+        if child_index == 0 {
             required.distribution = Distribution::Hash(self.left_conditions.clone());
         } else {
             required.distribution = Distribution::Hash(self.right_conditions.clone());

diff --git a/tests/sqllogictests/suites/mode/cluster/04_0002_explain_v2.test b/tests/sqllogictests/suites/mode/cluster/04_0002_explain_v2.test
@@ -254,7 +254,7 @@ Exchange
         ├── filters: []
         ├── estimated rows: 0.00
         ├── Exchange(Build)
-        │   ├── exchange type: Broadcast
+        │   ├── exchange type: Hash(t2.a (#2))
         │   └── TableScan
         │       ├── table: default.default.t2
         │       ├── read rows: 0
@@ -264,7 +264,7 @@ Exchange
         │       ├── push downs: [filters: [], limit: NONE]
         │       └── estimated rows: 0.00
         └── Exchange(Probe)
-            ├── exchange type: Init-Partition
+            ├── exchange type: Hash(t1.a (#0))
             └── TableScan
                 ├── table: default.default.t1
                 ├── read rows: 0