feat: Add Distinct Logical Node (#89)

Author: gandronchik

datafusion-cli/Cargo.lock

@@ -235,7 +235,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::*;

datafusion/core/src/dataframe.rs

@@ -38,13 +38,12 @@ use arrow::{
 use datafusion_data_access::object_store::ObjectStore;
 use datafusion_physical_expr::coercion_rule::binary_rule::comparison_eq_coercion;
 use std::convert::TryFrom;
-use std::iter;
 use std::{
     collections::{HashMap, HashSet},
     sync::Arc,
 };
 
-use super::{dfschema::ToDFSchema, expr_rewriter::coerce_plan_expr_for_schema};
+use super::{dfschema::ToDFSchema, expr_rewriter::coerce_plan_expr_for_schema, Distinct};
 use super::{exprlist_to_fields, Expr, JoinConstraint, JoinType, LogicalPlan, PlanType};
 use crate::logical_plan::{
     columnize_expr, normalize_col, normalize_cols, rewrite_sort_cols_by_aggs, Column,
@@ -649,18 +648,34 @@ 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> {
+        // unwrap top-level Distincts, to avoid duplication
+        let left_plan = self.plan.clone();
+        let left_plan: LogicalPlan = match left_plan {
+            LogicalPlan::Distinct(Distinct { input }) => (*input).clone(),
+            _ => left_plan,
+        };
+        let right_plan: LogicalPlan = match plan {
+            LogicalPlan::Distinct(Distinct { input }) => (*input).clone(),
+            _ => plan,
+        };
+
+        Ok(Self::from(LogicalPlan::Distinct(Distinct {
+            input: Arc::new(union_with_alias(left_plan, right_plan, None)?),
+        })))
+    }
+
     /// Apply deduplication: Only distinct (different) values are returned)
     pub fn distinct(&self) -> Result<Self> {
-        let projection_expr = expand_wildcard(self.plan.schema(), &self.plan)?;
-        let plan = LogicalPlanBuilder::from(self.plan.clone())
-            .aggregate(projection_expr, iter::empty::<Expr>())?
-            .build()?;
-        Self::from(plan).project(vec![Expr::Wildcard])
+        Ok(Self::from(LogicalPlan::Distinct(Distinct {
+            input: Arc::new(self.plan.clone()),
+        })))
     }
 
     /// Apply a join with on constraint
@@ -1402,6 +1417,57 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn plan_builder_union_distinct_combined_single_union() -> Result<()> {
+        let plan = LogicalPlanBuilder::scan_empty(
+            Some("employee_csv"),
+            &employee_schema(),
+            Some(vec![3, 4]),
+        )?;
+
+        let plan = plan
+            .union_distinct(plan.build()?)?
+            .union_distinct(plan.build()?)?
+            .union_distinct(plan.build()?)?
+            .build()?;
+
+        // output has only one union
+        let expected = "\
+        Distinct:\
+        \n  Union\
+        \n    TableScan: employee_csv projection=Some([3, 4])\
+        \n    TableScan: employee_csv projection=Some([3, 4])\
+        \n    TableScan: employee_csv projection=Some([3, 4])\
+        \n    TableScan: employee_csv projection=Some([3, 4])";
+
+        assert_eq!(expected, format!("{:?}", plan));
+
+        Ok(())
+    }
+
+    #[test]
+    fn plan_builder_simple_distinct() -> Result<()> {
+        let plan = LogicalPlanBuilder::scan_empty(
+            Some("employee_csv"),
+            &employee_schema(),
+            Some(vec![0, 3]),
+        )?
+        .filter(col("state").eq(lit("CO")))?
+        .project(vec![col("id")])?
+        .distinct()?
+        .build()?;
+
+        let expected = "\
+        Distinct:\
+        \n  Projection: #employee_csv.id\
+        \n    Filter: #employee_csv.state = Utf8(\"CO\")\
+        \n      TableScan: employee_csv projection=Some([0, 3])";
+
+        assert_eq!(expected, format!("{:?}", plan));
+
+        Ok(())
+    }
+
     #[test]
     fn projection_non_unique_names() -> Result<()> {
         let plan = LogicalPlanBuilder::scan_empty(

datafusion/core/src/logical_plan/builder.rs

@@ -65,9 +65,9 @@ pub use expr_visitor::{ExprVisitable, ExpressionVisitor, Recursion};
 pub use extension::UserDefinedLogicalNode;
 pub use operators::Operator;
 pub use plan::{
-    CreateCatalogSchema, CreateExternalTable, CreateMemoryTable, CrossJoin, DropTable,
-    EmptyRelation, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType,
-    PlanVisitor, Repartition, StringifiedPlan, Subquery, TableScan, ToStringifiedPlan,
-    Union, Values,
+    CreateCatalogSchema, CreateExternalTable, CreateMemoryTable, CrossJoin, Distinct,
+    DropTable, EmptyRelation, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning,
+    PlanType, PlanVisitor, Repartition, StringifiedPlan, Subquery, TableScan,
+    ToStringifiedPlan, Union, Values,
 };
 pub use registry::FunctionRegistry;

datafusion/core/src/logical_plan/mod.rs

@@ -297,6 +297,13 @@ pub struct Values {
     pub values: Vec<Vec<Expr>>,
 }
 
+/// Removes duplicate rows from the input
+#[derive(Clone)]
+pub struct Distinct {
+    /// The logical plan that is being DISTINCT'd
+    pub input: Arc<LogicalPlan>,
+}
+
 /// Aggregates its input based on a set of grouping and aggregate
 /// expressions (e.g. SUM).
 #[derive(Clone)]
@@ -417,6 +424,8 @@ pub enum LogicalPlan {
     TableUDFs(TableUDFs),
     /// Extension operator defined outside of DataFusion
     Extension(Extension),
+    /// Remove duplicate rows from the input
+    Distinct(Distinct),
 }
 
 impl LogicalPlan {
@@ -431,6 +440,7 @@ impl LogicalPlan {
             LogicalPlan::Projection(Projection { schema, .. }) => schema,
             LogicalPlan::Subquery(Subquery { schema, .. }) => schema,
             LogicalPlan::Filter(Filter { input, .. }) => input.schema(),
+            LogicalPlan::Distinct(Distinct { input }) => input.schema(),
             LogicalPlan::Window(Window { schema, .. }) => schema,
             LogicalPlan::Aggregate(Aggregate { schema, .. }) => schema,
             LogicalPlan::Sort(Sort { input, .. }) => input.schema(),
@@ -504,6 +514,7 @@ impl LogicalPlan {
             | LogicalPlan::Sort(Sort { input, .. })
             | LogicalPlan::CreateMemoryTable(CreateMemoryTable { input, .. })
             | LogicalPlan::Filter(Filter { input, .. }) => input.all_schemas(),
+            LogicalPlan::Distinct(Distinct { input, .. }) => input.all_schemas(),
             LogicalPlan::DropTable(_) => vec![],
         }
     }
@@ -558,7 +569,8 @@ impl LogicalPlan {
             | LogicalPlan::Analyze { .. }
             | LogicalPlan::Explain { .. }
             | LogicalPlan::Subquery(_)
-            | LogicalPlan::Union(_) => {
+            | LogicalPlan::Union(_)
+            | LogicalPlan::Distinct(_) => {
                 vec![]
             }
         }
@@ -585,6 +597,7 @@ impl LogicalPlan {
             LogicalPlan::Limit(Limit { input, .. }) => vec![input],
             LogicalPlan::Extension(extension) => extension.node.inputs(),
             LogicalPlan::Union(Union { inputs, .. }) => inputs.iter().collect(),
+            LogicalPlan::Distinct(Distinct { input }) => vec![input],
             LogicalPlan::Explain(explain) => vec![&explain.plan],
             LogicalPlan::Analyze(analyze) => vec![&analyze.input],
             LogicalPlan::CreateMemoryTable(CreateMemoryTable { input, .. }) => {
@@ -741,6 +754,7 @@ impl LogicalPlan {
                 }
                 true
             }
+            LogicalPlan::Distinct(Distinct { input }) => input.accept(visitor)?,
             LogicalPlan::Limit(Limit { input, .. }) => input.accept(visitor)?,
             LogicalPlan::CreateMemoryTable(CreateMemoryTable { input, .. }) => {
                 input.accept(visitor)?
@@ -1157,6 +1171,9 @@ impl LogicalPlan {
                     }) => {
                         write!(f, "DropTable: {:?} if not exist:={}", name, if_exists)
                     }
+                    LogicalPlan::Distinct(Distinct { .. }) => {
+                        write!(f, "Distinct:")
+                    }
                     LogicalPlan::Explain { .. } => write!(f, "Explain"),
                     LogicalPlan::Analyze { .. } => write!(f, "Analyze"),
                     LogicalPlan::Union(_) => write!(f, "Union"),

datafusion/core/src/logical_plan/plan.rs

@@ -249,6 +249,7 @@ fn optimize(
         | LogicalPlan::CreateCatalogSchema(_)
         | LogicalPlan::DropTable(_)
         | LogicalPlan::Subquery(_)
+        | LogicalPlan::Distinct(_)
         | LogicalPlan::Extension { .. } => {
             // apply the optimization to all inputs of the plan
             let expr = plan.expressions();

datafusion/core/src/optimizer/common_subexpr_eliminate.rs

@@ -505,6 +505,7 @@ fn optimize_plan(
         | LogicalPlan::DropTable(_)
         | LogicalPlan::CrossJoin(_)
         | LogicalPlan::TableUDFs(_)
+        | LogicalPlan::Distinct(_)
         | LogicalPlan::Extension { .. } => {
             let expr = plan.expressions();
             // collect all required columns by this plan

datafusion/core/src/optimizer/projection_push_down.rs

@@ -26,9 +26,9 @@ use crate::optimizer::optimizer::OptimizerConfig;
 
 use crate::logical_plan::builder::build_table_udf_schema;
 use crate::logical_plan::{
-    build_join_schema, Column, CreateMemoryTable, DFSchemaRef, Expr, ExprVisitable,
-    Limit, LogicalPlan, LogicalPlanBuilder, Operator, Partitioning, Recursion,
-    Repartition, Union, Values,
+    build_join_schema, Column, CreateMemoryTable, DFSchemaRef, Distinct, Expr,
+    ExprVisitable, Limit, LogicalPlan, LogicalPlanBuilder, Operator, Partitioning,
+    Recursion, Repartition, Union, Values,
 };
 use crate::prelude::lit;
 use crate::scalar::ScalarValue;
@@ -264,6 +264,9 @@ pub fn from_plan(
                 alias: alias.clone(),
             }))
         }
+        LogicalPlan::Distinct(Distinct { .. }) => Ok(LogicalPlan::Distinct(Distinct {
+            input: Arc::new(inputs[0].clone()),
+        })),
         LogicalPlan::Analyze(a) => {
             assert!(expr.is_empty());
             assert_eq!(inputs.len(), 1);

datafusion/core/src/optimizer/utils.rs

@@ -24,14 +24,15 @@ use super::{
     hash_join::PartitionMode, udaf, union::UnionExec, values::ValuesExec, windows,
 };
 use crate::execution::context::{ExecutionProps, SessionState};
+use crate::logical_plan::builder::expand_wildcard;
 use crate::logical_plan::plan::{
     Aggregate, EmptyRelation, Filter, Join, Projection, Sort, Subquery, TableScan,
     TableUDFs, Window,
 };
 use crate::logical_plan::{
-    unalias, unnormalize_cols, CrossJoin, DFSchema, Expr, LogicalPlan, Operator,
-    Partitioning as LogicalPartitioning, PlanType, Repartition, ToStringifiedPlan, Union,
-    UserDefinedLogicalNode,
+    unalias, unnormalize_cols, CrossJoin, DFSchema, Distinct, Expr, LogicalPlan,
+    Operator, Partitioning as LogicalPartitioning, PlanType, Repartition,
+    ToStringifiedPlan, Union, UserDefinedLogicalNode,
 };
 use crate::logical_plan::{Limit, Values};
 use crate::physical_optimizer::optimizer::PhysicalOptimizerRule;
@@ -616,6 +617,18 @@ impl DefaultPhysicalPlanner {
                             input_exec,
                         )?))
                 }
+                LogicalPlan::Distinct(Distinct {input}) => {
+                    // Convert distinct to groupby with no aggregations
+                    let group_expr = expand_wildcard(input.schema(), input)?;
+                    let aggregate =  LogicalPlan::Aggregate(Aggregate {
+                            input: input.clone(),
+                            group_expr,
+                            aggr_expr: vec![],
+                            schema: input.schema().clone()
+                        }
+                    );
+                    Ok(self.create_initial_plan(&aggregate, session_state).await?)
+                }
                 LogicalPlan::Projection(Projection { input, expr, .. }) => {
                     let input_exec = self.create_initial_plan(input, session_state).await?;
                     let input_schema = input.as_ref().schema();