handling join column expansion during USING JOIN planning

get rid of shared field and move column expansion logic into plan
builder and optimizer.

Author: houqp

ballista/rust/core/proto/ballista.proto

@@ -576,7 +576,6 @@ message HashJoinExecNode {
   PhysicalPlanNode right = 2;
   repeated JoinOn on = 3;
   JoinType join_type = 4;
-  JoinConstraint join_constraint = 5;
 }
 
 message PhysicalColumn {

ballista/rust/core/src/serde/logical_plan/from_proto.rs

@@ -267,9 +267,7 @@ impl TryInto<LogicalPlan> for &protobuf::LogicalPlanNode {
                     ))
                 })?;
 
-                let builder =
-                    LogicalPlanBuilder::from(&convert_box_required!(join.left)?);
-
+                let builder = LogicalPlanBuilder::from(convert_box_required!(join.left)?);
                 let builder = match join_constraint.into() {
                     JoinConstraint::On => builder.join(
                         &convert_box_required!(join.right)?,

ballista/rust/core/src/serde/physical_plan/from_proto.rs

@@ -350,21 +350,11 @@ impl TryInto<Arc<dyn ExecutionPlan>> for &protobuf::PhysicalPlanNode {
                         ))
                     })?;
 
-                let join_constraint =
-                    protobuf::JoinConstraint::from_i32(hashjoin.join_constraint)
-                        .ok_or_else(|| {
-                            proto_error(format!(
-                        "Received a HashJoinNode message with unknown JoinConstraint {}",
-                        hashjoin.join_constraint,
-                    ))
-                        })?;
-
                 Ok(Arc::new(HashJoinExec::try_new(
                     left,
                     right,
                     on,
                     &join_type.into(),
-                    join_constraint.into(),
                     PartitionMode::CollectLeft,
                 )?))
             }

ballista/rust/core/src/serde/physical_plan/mod.rs

@@ -27,7 +27,7 @@ mod roundtrip_tests {
             compute::kernels::sort::SortOptions,
             datatypes::{DataType, Field, Schema},
         },
-        logical_plan::{JoinConstraint, JoinType, Operator},
+        logical_plan::{JoinType, Operator},
         physical_plan::{
             empty::EmptyExec,
             expressions::{binary, col, lit, InListExpr, NotExpr},
@@ -92,7 +92,6 @@ mod roundtrip_tests {
             Arc::new(EmptyExec::new(false, Arc::new(schema_right))),
             on,
             &JoinType::Inner,
-            JoinConstraint::On,
             PartitionMode::CollectLeft,
         )?))
     }

ballista/rust/core/src/serde/physical_plan/to_proto.rs

@@ -136,7 +136,6 @@ impl TryInto<protobuf::PhysicalPlanNode> for Arc<dyn ExecutionPlan> {
                 })
                 .collect();
             let join_type: protobuf::JoinType = exec.join_type().to_owned().into();
-            let join_constraint: protobuf::JoinConstraint = exec.join_constraint().into();
 
             Ok(protobuf::PhysicalPlanNode {
                 physical_plan_type: Some(PhysicalPlanType::HashJoin(Box::new(
@@ -145,7 +144,6 @@ impl TryInto<protobuf::PhysicalPlanNode> for Arc<dyn ExecutionPlan> {
                         right: Some(Box::new(right)),
                         on,
                         join_type: join_type.into(),
-                        join_constraint: join_constraint.into(),
                     },
                 ))),
             })

datafusion/src/logical_plan/builder.rs

@@ -40,7 +40,6 @@ use crate::logical_plan::{
     columnize_expr, normalize_col, normalize_cols, Column, DFField, DFSchema,
     DFSchemaRef, Partitioning,
 };
-use std::collections::HashSet;
 
 /// Default table name for unnamed table
 pub const UNNAMED_TABLE: &str = "?table?";
@@ -213,7 +212,6 @@ impl LogicalPlanBuilder {
     /// * An invalid expression is used (e.g. a `sort` expression)
     pub fn project(&self, expr: impl IntoIterator<Item = Expr>) -> Result<Self> {
         let input_schema = self.plan.schema();
-        let all_schemas = self.plan.all_schemas();
         let mut projected_expr = vec![];
         for e in expr {
             match e {
@@ -223,10 +221,8 @@ impl LogicalPlanBuilder {
                             .push(Expr::Column(input_schema.field(i).qualified_column()))
                     });
                 }
-                _ => projected_expr.push(columnize_expr(
-                    normalize_col(e, &all_schemas)?,
-                    input_schema,
-                )),
+                _ => projected_expr
+                    .push(columnize_expr(normalize_col(e, &self.plan)?, input_schema)),
             }
         }
 
@@ -243,7 +239,7 @@ impl LogicalPlanBuilder {
 
     /// Apply a filter
     pub fn filter(&self, expr: Expr) -> Result<Self> {
-        let expr = normalize_col(expr, &self.plan.all_schemas())?;
+        let expr = normalize_col(expr, &self.plan)?;
         Ok(Self::from(LogicalPlan::Filter {
             predicate: expr,
             input: Arc::new(self.plan.clone()),
@@ -260,9 +256,8 @@ impl LogicalPlanBuilder {
 
     /// Apply a sort
     pub fn sort(&self, exprs: impl IntoIterator<Item = Expr>) -> Result<Self> {
-        let schemas = self.plan.all_schemas();
         Ok(Self::from(LogicalPlan::Sort {
-            expr: normalize_cols(exprs, &schemas)?,
+            expr: normalize_cols(exprs, &self.plan)?,
             input: Arc::new(self.plan.clone()),
         }))
     }
@@ -288,20 +283,15 @@ impl LogicalPlanBuilder {
 
         let left_keys: Vec<Column> = left_keys
             .into_iter()
-            .map(|c| c.into().normalize(&self.plan.all_schemas()))
+            .map(|c| c.into().normalize(&self.plan))
             .collect::<Result<_>>()?;
         let right_keys: Vec<Column> = right_keys
             .into_iter()
-            .map(|c| c.into().normalize(&right.all_schemas()))
+            .map(|c| c.into().normalize(right))
             .collect::<Result<_>>()?;
         let on: Vec<(_, _)> = left_keys.into_iter().zip(right_keys.into_iter()).collect();
-        let join_schema = build_join_schema(
-            self.plan.schema(),
-            right.schema(),
-            &on,
-            &join_type,
-            &JoinConstraint::On,
-        )?;
+        let join_schema =
+            build_join_schema(self.plan.schema(), right.schema(), &join_type)?;
 
         Ok(Self::from(LogicalPlan::Join {
             left: Arc::new(self.plan.clone()),
@@ -323,21 +313,16 @@ impl LogicalPlanBuilder {
         let left_keys: Vec<Column> = using_keys
             .clone()
             .into_iter()
-            .map(|c| c.into().normalize(&self.plan.all_schemas()))
+            .map(|c| c.into().normalize(&self.plan))
             .collect::<Result<_>>()?;
         let right_keys: Vec<Column> = using_keys
             .into_iter()
-            .map(|c| c.into().normalize(&right.all_schemas()))
+            .map(|c| c.into().normalize(right))
             .collect::<Result<_>>()?;
 
         let on: Vec<(_, _)> = left_keys.into_iter().zip(right_keys.into_iter()).collect();
-        let join_schema = build_join_schema(
-            self.plan.schema(),
-            right.schema(),
-            &on,
-            &join_type,
-            &JoinConstraint::Using,
-        )?;
+        let join_schema =
+            build_join_schema(self.plan.schema(), right.schema(), &join_type)?;
 
         Ok(Self::from(LogicalPlan::Join {
             left: Arc::new(self.plan.clone()),
@@ -390,9 +375,8 @@ impl LogicalPlanBuilder {
         group_expr: impl IntoIterator<Item = Expr>,
         aggr_expr: impl IntoIterator<Item = Expr>,
     ) -> Result<Self> {
-        let schemas = self.plan.all_schemas();
-        let group_expr = normalize_cols(group_expr, &schemas)?;
-        let aggr_expr = normalize_cols(aggr_expr, &schemas)?;
+        let group_expr = normalize_cols(group_expr, &self.plan)?;
+        let aggr_expr = normalize_cols(aggr_expr, &self.plan)?;
         let all_expr = group_expr.iter().chain(aggr_expr.iter());
 
         validate_unique_names("Aggregations", all_expr.clone(), self.plan.schema())?;
@@ -436,53 +420,14 @@ impl LogicalPlanBuilder {
 pub fn build_join_schema(
     left: &DFSchema,
     right: &DFSchema,
-    on: &[(Column, Column)],
     join_type: &JoinType,
-    join_constraint: &JoinConstraint,
 ) -> Result<DFSchema> {
     let fields: Vec<DFField> = match join_type {
         JoinType::Inner | JoinType::Left | JoinType::Full | JoinType::Right => {
-            match join_constraint {
-                JoinConstraint::On => {
-                    let right_fields = right.fields().iter();
-                    let left_fields = left.fields().iter();
-                    // left then right
-                    left_fields.chain(right_fields).cloned().collect()
-                }
-                JoinConstraint::Using => {
-                    // using join requires unique join column in the output schema, so we mark all
-                    // right join keys as duplicate
-                    let duplicate_join_names =
-                        on.iter().map(|on| &on.1.name).collect::<HashSet<_>>();
-
-                    let right_fields = right
-                        .fields()
-                        .iter()
-                        .filter(|f| !duplicate_join_names.contains(f.name()))
-                        .cloned();
-
-                    let left_fields = left.fields().iter().map(|f| {
-                        for key in on.iter() {
-                            // update qualifiers for shared fields
-                            if duplicate_join_names.contains(f.name()) {
-                                let mut hs = HashSet::new();
-                                if let Some(q) = &key.0.relation {
-                                    hs.insert(q.to_string());
-                                }
-                                if let Some(q) = &key.1.relation {
-                                    hs.insert(q.to_string());
-                                }
-                                return f.clone().set_shared_qualifiers(hs);
-                            }
-                        }
-
-                        f.clone()
-                    });
-
-                    // left then right
-                    left_fields.chain(right_fields).collect()
-                }
-            }
+            let right_fields = right.fields().iter();
+            let left_fields = left.fields().iter();
+            // left then right
+            left_fields.chain(right_fields).cloned().collect()
         }
         JoinType::Semi | JoinType::Anti => {
             // Only use the left side for the schema

datafusion/src/logical_plan/dfschema.rs

@@ -57,15 +57,6 @@ impl DFSchema {
                         field.qualified_name()
                     )));
                 }
-            } else if let Some(shared_qualifiers) = field.shared_qualifiers() {
-                for qualifier in shared_qualifiers {
-                    if !qualified_names.insert((qualifier, field.name())) {
-                        return Err(DataFusionError::Plan(format!(
-                            "Schema contains duplicate qualified field name '{}'",
-                            field.qualified_name()
-                        )));
-                    }
-                }
             } else if !unqualified_names.insert(field.name()) {
                 return Err(DataFusionError::Plan(format!(
                     "Schema contains duplicate unqualified field name '{}'",
@@ -170,18 +161,8 @@ impl DFSchema {
                 // current field is qualified and not shared between relations, compare both
                 // qualifer and name.
                 (Some(q), Some(field_q)) => q == field_q && field.name() == name,
-                // field to lookup is qualified.
-                // current field is either unqualified or qualified and shared between relations.
-                (Some(q), None) => {
-                    if let Some(shared_q) = field.shared_qualifiers() {
-                        // current field is a shared qualified field, check for all shared
-                        // relation names.
-                        shared_q.contains(q) && field.name() == name
-                    } else {
-                        // current field is unqualifiied
-                        false
-                    }
-                }
+                // field to lookup is qualified but current field is unqualified.
+                (Some(_), None) => false,
                 // field to lookup is unqualified, no need to compare qualifier
                 _ => field.name() == name,
             })
@@ -222,13 +203,17 @@ impl DFSchema {
         }
     }
 
-    /// Find the field with the given name
-    pub fn field_with_unqualified_name(&self, name: &str) -> Result<&DFField> {
-        let matches: Vec<&DFField> = self
-            .fields
+    /// Find all fields match the given name
+    pub fn fields_with_unqualified_name(&self, name: &str) -> Vec<&DFField> {
+        self.fields
             .iter()
             .filter(|field| field.name() == name)
-            .collect();
+            .collect()
+    }
+
+    /// Find the field with the given name
+    pub fn field_with_unqualified_name(&self, name: &str) -> Result<&DFField> {
+        let matches = self.fields_with_unqualified_name(name);
         match matches.len() {
             0 => Err(DataFusionError::Plan(format!(
                 "No field with unqualified name '{}'. Valid fields are {}.",
@@ -417,9 +402,6 @@ impl Display for DFSchema {
 pub struct DFField {
     /// Optional qualifier (usually a table or relation name)
     qualifier: Option<String>,
-    /// Optional set of qualifiers that all share this same field. This is used for `JOIN USING`
-    /// clause where the join keys are combined into a shared column.
-    shared_qualifiers: Option<HashSet<String>>,
     /// Arrow field definition
     field: Field,
 }
@@ -434,7 +416,6 @@ impl DFField {
     ) -> Self {
         DFField {
             qualifier: qualifier.map(|s| s.to_owned()),
-            shared_qualifiers: None,
             field: Field::new(name, data_type, nullable),
         }
     }
@@ -443,7 +424,6 @@ impl DFField {
     pub fn from(field: Field) -> Self {
         Self {
             qualifier: None,
-            shared_qualifiers: None,
             field,
         }
     }
@@ -452,7 +432,6 @@ impl DFField {
     pub fn from_qualified(qualifier: &str, field: Field) -> Self {
         Self {
             qualifier: Some(qualifier.to_owned()),
-            shared_qualifiers: None,
             field,
         }
     }
@@ -502,11 +481,6 @@ impl DFField {
         self.qualifier.as_ref()
     }
 
-    /// Get the optional qualifier
-    pub fn shared_qualifiers(&self) -> Option<&HashSet<String>> {
-        self.shared_qualifiers.as_ref()
-    }
-
     /// Get the arrow field
     pub fn field(&self) -> &Field {
         &self.field
@@ -517,13 +491,6 @@ impl DFField {
         self.qualifier = None;
         self
     }
-
-    /// Return field with shared qualifiers set and qualifier stripped
-    pub fn set_shared_qualifiers(mut self, shared_qualifiers: HashSet<String>) -> Self {
-        self.qualifier = None;
-        self.shared_qualifiers = Some(shared_qualifiers);
-        self
-    }
 }
 
 #[cfg(test)]