dedup using join column in wildcard expansion

Author: houqp

datafusion/src/logical_plan/builder.rs

@@ -17,7 +17,10 @@
 
 //! This module provides a builder for creating LogicalPlans
 
-use std::{collections::HashMap, sync::Arc};
+use std::{
+    collections::{HashMap, HashSet},
+    sync::Arc,
+};
 
 use arrow::{
     datatypes::{Schema, SchemaRef},
@@ -220,10 +223,33 @@ impl LogicalPlanBuilder {
         for e in expr {
             match e {
                 Expr::Wildcard => {
-                    (0..input_schema.fields().len()).for_each(|i| {
-                        projected_expr
-                            .push(Expr::Column(input_schema.field(i).qualified_column()))
-                    });
+                    let columns_to_skip = self
+                        .plan
+                        .using_columns()?
+                        .into_iter()
+                        // For each USING JOIN condition, only expand to one column in projection
+                        .map(|cols| {
+                            let mut cols = cols.into_iter().collect::<Vec<_>>();
+                            // sort join columns to make sure we consistently keep the same
+                            // qualified column
+                            cols.sort();
+                            cols.into_iter().skip(1)
+                        })
+                        .flatten()
+                        .collect::<HashSet<_>>();
+
+                    if columns_to_skip.is_empty() {
+                        input_schema.fields().iter().for_each(|f| {
+                            projected_expr.push(Expr::Column(f.qualified_column()))
+                        })
+                    } else {
+                        input_schema.fields().iter().for_each(|f| {
+                            let col = f.qualified_column();
+                            if !columns_to_skip.contains(&col) {
+                                projected_expr.push(Expr::Column(col))
+                            }
+                        })
+                    }
                 }
                 _ => projected_expr
                     .push(columnize_expr(normalize_col(e, &self.plan)?, input_schema)),
@@ -587,6 +613,27 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn plan_using_join_wildcard_projection() -> Result<()> {
+        let t2 = LogicalPlanBuilder::scan_empty(Some("t2"), &employee_schema(), None)?
+            .build()?;
+
+        let plan = LogicalPlanBuilder::scan_empty(Some("t1"), &employee_schema(), None)?
+            .join_using(&t2, JoinType::Inner, vec!["id"])?
+            .project(vec![Expr::Wildcard])?
+            .build()?;
+
+        // id column should only show up once in projection
+        let expected = "Projection: #t1.id, #t1.first_name, #t1.last_name, #t1.state, #t1.salary, #t2.first_name, #t2.last_name, #t2.state, #t2.salary\
+        \n  Join: Using #t1.id = #t2.id\
+        \n    TableScan: t1 projection=None\
+        \n    TableScan: t2 projection=None";
+
+        assert_eq!(expected, format!("{:?}", plan));
+
+        Ok(())
+    }
+
     #[test]
     fn plan_builder_union_combined_single_union() -> Result<()> {
         let plan = LogicalPlanBuilder::scan_empty(

datafusion/src/logical_plan/expr.rs

@@ -34,7 +34,7 @@ use std::fmt;
 use std::sync::Arc;
 
 /// A named reference to a qualified field in a schema.
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct Column {
     /// relation/table name.
     pub relation: Option<String>,

datafusion/src/logical_plan/plan.rs

@@ -25,8 +25,8 @@ use crate::error::DataFusionError;
 use crate::logical_plan::dfschema::DFSchemaRef;
 use crate::sql::parser::FileType;
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
-use std::collections::HashSet;
 use std::{
+    collections::HashSet,
     fmt::{self, Display},
     sync::Arc,
 };

datafusion/src/sql/planner.rs

@@ -663,13 +663,14 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         projection: &[SelectItem],
     ) -> Result<Vec<Expr>> {
         let input_schema = plan.schema();
+        let using_columns = plan.using_columns()?;
 
         projection
             .iter()
             .map(|expr| self.sql_select_to_rex(expr, input_schema))
             .collect::<Result<Vec<Expr>>>()?
             .iter()
-            .flat_map(|expr| expand_wildcard(expr, input_schema))
+            .flat_map(|expr| expand_wildcard(expr, input_schema, &using_columns))
             .map(|expr| normalize_col(expr, plan))
             .collect::<Result<Vec<Expr>>>()
     }
@@ -2738,6 +2739,19 @@ mod tests {
         quick_test(sql, expected);
     }
 
+    #[test]
+    fn project_wildcard_on_join_with_using() {
+        let sql = "SELECT * \
+            FROM lineitem \
+            JOIN lineitem as lineitem2 \
+            USING (l_item_id)";
+        let expected = "Projection: #lineitem.l_item_id, #lineitem.l_description, #lineitem.price, #lineitem2.l_description, #lineitem2.price\
+        \n  Join: Using #lineitem.l_item_id = #lineitem2.l_item_id\
+        \n    TableScan: lineitem projection=None\
+        \n    TableScan: lineitem2 projection=None";
+        quick_test(sql, expected);
+    }
+
     #[test]
     fn equijoin_explicit_syntax_3_tables() {
         let sql = "SELECT id, order_id, l_description \

datafusion/src/sql/utils.rs

@@ -23,16 +23,50 @@ use crate::{
     error::{DataFusionError, Result},
     logical_plan::{Column, ExpressionVisitor, Recursion},
 };
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 /// Resolves an `Expr::Wildcard` to a collection of `Expr::Column`'s.
-pub(crate) fn expand_wildcard(expr: &Expr, schema: &DFSchema) -> Vec<Expr> {
+pub(crate) fn expand_wildcard(
+    expr: &Expr,
+    schema: &DFSchema,
+    using_columns: &[HashSet<Column>],
+) -> Vec<Expr> {
     match expr {
-        Expr::Wildcard => schema
-            .fields()
-            .iter()
-            .map(|f| Expr::Column(f.qualified_column()))
-            .collect::<Vec<Expr>>(),
+        Expr::Wildcard => {
+            let columns_to_skip = using_columns
+                .iter()
+                // For each USING JOIN condition, only expand to one column in projection
+                .map(|cols| {
+                    let mut cols = cols.iter().collect::<Vec<_>>();
+                    // sort join columns to make sure we consistently keep the same
+                    // qualified column
+                    cols.sort();
+                    cols.into_iter().skip(1)
+                })
+                .flatten()
+                .collect::<HashSet<_>>();
+
+            if columns_to_skip.is_empty() {
+                schema
+                    .fields()
+                    .iter()
+                    .map(|f| Expr::Column(f.qualified_column()))
+                    .collect::<Vec<Expr>>()
+            } else {
+                schema
+                    .fields()
+                    .iter()
+                    .filter_map(|f| {
+                        let col = f.qualified_column();
+                        if !columns_to_skip.contains(&col) {
+                            Some(Expr::Column(col))
+                        } else {
+                            None
+                        }
+                    })
+                    .collect::<Vec<Expr>>()
+            }
+        }
         _ => vec![expr.clone()],
     }
 }