feat: optimize broadcast join

databendlabs · May 9, 2023 · 17ae705 · 17ae705
1 parent 6993c44
commit 17ae705
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 2 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(|child| check_partition(child))
+                {
+                    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) {
+            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
@@ -342,6 +342,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(),
@@ -423,8 +428,6 @@ 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 {