Add DataFrame union_distinct and fix documentation for distinct

datafusion/core/src/dataframe.rs

@@ -208,7 +208,8 @@ impl DataFrame {
         Ok(Arc::new(DataFrame::new(self.session_state.clone(), &plan)))
     }
 
-    /// Calculate the union two [`DataFrame`]s.  The two [`DataFrame`]s must have exactly the same schema
+    /// Calculate the union of two [`DataFrame`]s, preserving duplicate rows.The
+    /// two [`DataFrame`]s must have exactly the same schema
     ///
     /// ```
     /// # use datafusion::prelude::*;
@@ -228,7 +229,30 @@ impl DataFrame {
         Ok(Arc::new(DataFrame::new(self.session_state.clone(), &plan)))
     }
 
-    /// Calculate the union distinct two [`DataFrame`]s.  The two [`DataFrame`]s must have exactly the same schema
+    /// Calculate the distinct union of two [`DataFrame`]s.  The
+    /// two [`DataFrame`]s must have exactly the same schema
+    ///
+    /// ```
+    /// # use datafusion::prelude::*;
+    /// # use datafusion::error::Result;
+    /// # #[tokio::main]
+    /// # async fn main() -> Result<()> {
+    /// let ctx = SessionContext::new();
+    /// let df = ctx.read_csv("tests/example.csv", CsvReadOptions::new()).await?;
+    /// let df = df.union_distinct(df.clone())?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn union_distinct(&self, dataframe: Arc<DataFrame>) -> Result<Arc<DataFrame>> {
+        Ok(Arc::new(DataFrame::new(
+            self.session_state.clone(),
+            &LogicalPlanBuilder::from(self.plan.clone())
+                .union_distinct(dataframe.plan.clone())?
+                .build()?,
+        )))
+    }
+
+    /// Filter out duplicate rows
     ///
     /// ```
     /// # use datafusion::prelude::*;
@@ -237,7 +261,6 @@ impl DataFrame {
     /// # async fn main() -> Result<()> {
     /// let ctx = SessionContext::new();
     /// let df = ctx.read_csv("tests/example.csv", CsvReadOptions::new()).await?;
-    /// let df = df.union(df.clone())?;
     /// let df = df.distinct()?;
     /// # Ok(())
     /// # }

datafusion/core/src/logical_plan/builder.rs

@@ -442,11 +442,16 @@ impl LogicalPlanBuilder {
         })))
     }
 
-    /// Apply a union
+    /// Apply a union, preserving duplicate rows
     pub fn union(&self, plan: LogicalPlan) -> Result<Self> {
         Ok(Self::from(union_with_alias(self.plan.clone(), plan, None)?))
     }
 
+    /// Apply a union, removing duplicate rows
+    pub fn union_distinct(&self, plan: LogicalPlan) -> Result<Self> {
+        self.union(plan)?.distinct()
+    }
+
     /// Apply deduplication: Only distinct (different) values are returned)
     pub fn distinct(&self) -> Result<Self> {
         let projection_expr = expand_wildcard(self.plan.schema(), &self.plan)?;

datafusion/core/src/sql/planner.rs

@@ -28,8 +28,8 @@ use crate::datasource::TableProvider;
 use crate::logical_plan::window_frames::{WindowFrame, WindowFrameUnits};
 use crate::logical_plan::Expr::Alias;
 use crate::logical_plan::{
-    and, col, lit, normalize_col, normalize_col_with_schemas, union_with_alias, Column,
-    CreateCatalog, CreateCatalogSchema, CreateExternalTable as PlanCreateExternalTable,
+    and, col, lit, normalize_col, normalize_col_with_schemas, Column, CreateCatalog,
+    CreateCatalogSchema, CreateExternalTable as PlanCreateExternalTable,
     CreateMemoryTable, CreateView, DFSchema, DFSchemaRef, DropTable, Expr, FileType,
     LogicalPlan, LogicalPlanBuilder, Operator, PlanType, ToDFSchema, ToStringifiedPlan,
 };
@@ -324,13 +324,12 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 let right_plan =
                     self.set_expr_to_plan(*right, None, ctes, outer_query_schema)?;
                 match (op, all) {
-                    (SetOperator::Union, true) => {
-                        union_with_alias(left_plan, right_plan, alias)
-                    }
-                    (SetOperator::Union, false) => {
-                        let union_plan = union_with_alias(left_plan, right_plan, alias)?;
-                        LogicalPlanBuilder::from(union_plan).distinct()?.build()
-                    }
+                    (SetOperator::Union, true) => LogicalPlanBuilder::from(left_plan)
+                        .union(right_plan)?
+                        .build(),
+                    (SetOperator::Union, false) => LogicalPlanBuilder::from(left_plan)
+                        .union_distinct(right_plan)?
+                        .build(),
                     (SetOperator::Intersect, true) => {
                         LogicalPlanBuilder::intersect(left_plan, right_plan, true)
                     }
@@ -3915,6 +3914,18 @@ mod tests {
 
     #[test]
     fn union() {
+        let sql = "SELECT order_id from orders UNION SELECT order_id FROM orders";
+        let expected = "Projection: #order_id\
+        \n  Aggregate: groupBy=[[#order_id]], aggr=[[]]\
+        \n    Union\n      Projection: #orders.order_id\
+        \n        TableScan: orders projection=None\
+        \n      Projection: #orders.order_id\
+        \n        TableScan: orders projection=None";
+        quick_test(sql, expected);
+    }
+
+    #[test]
+    fn union_all() {
         let sql = "SELECT order_id from orders UNION ALL SELECT order_id FROM orders";
         let expected = "Union\
             \n  Projection: #orders.order_id\