refactor[rust]: is_unique/is_duplicated -> FunctionExpr (#4711)

pola-rs · Sep 3, 2022 · 8db19b8 · 8db19b8
1 parent 08afe57
commit 8db19b8
Show file tree

Hide file tree

Showing 12 changed files with 19 additions and 54 deletions.
diff --git a/polars/polars-lazy/src/dsl/expr.rs b/polars/polars-lazy/src/dsl/expr.rs
@@ -311,8 +311,6 @@ pub enum Expr {
         function: FunctionExpr,
         options: FunctionOptions,
     },
-    Duplicated(Box<Expr>),
-    IsUnique(Box<Expr>),
     Explode(Box<Expr>),
     Filter {
         input: Box<Expr>,

diff --git a/polars/polars-lazy/src/dsl/function_expr/dispatch.rs b/polars/polars-lazy/src/dsl/function_expr/dispatch.rs
@@ -21,3 +21,11 @@ pub(super) fn is_not_null(s: &Series) -> Result<Series> {
 pub(super) fn is_not(s: &Series) -> Result<Series> {
     Ok(s.bool()?.not().into_series())
 }
+
+pub(super) fn is_unique(s: &Series) -> Result<Series> {
+    s.is_unique().map(|ca| ca.into_series())
+}
+
+pub(super) fn is_duplicated(s: &Series) -> Result<Series> {
+    s.is_duplicated().map(|ca| ca.into_series())
+}
diff --git a/polars/polars-lazy/src/dsl/function_expr/mod.rs b/polars/polars-lazy/src/dsl/function_expr/mod.rs
@@ -102,6 +102,8 @@ pub enum FunctionExpr {
     IsNull,
     IsNotNull,
     Not,
+    IsUnique,
+    IsDuplicated,
 }
 
 impl Display for FunctionExpr {
@@ -152,6 +154,8 @@ impl Display for FunctionExpr {
             Not => write!(f, "is_not"),
             IsNull => write!(f, "is_null"),
             IsNotNull => write!(f, "is_not_null"),
+            IsUnique => write!(f, "is_unique"),
+            IsDuplicated => write!(f, "is_duplicated"),
         }
     }
 }
@@ -309,6 +313,8 @@ impl From<FunctionExpr> for SpecialEq<Arc<dyn SeriesUdf>> {
             IsNull => map!(dispatch::is_null),
             IsNotNull => map!(dispatch::is_not_null),
             Not => map!(dispatch::is_not),
+            IsUnique => map!(dispatch::is_unique),
+            IsDuplicated => map!(dispatch::is_duplicated),
         }
     }
 }

diff --git a/polars/polars-lazy/src/dsl/function_expr/schema.rs b/polars/polars-lazy/src/dsl/function_expr/schema.rs
@@ -155,7 +155,7 @@ impl FunctionExpr {
             #[cfg(feature = "top_k")]
             TopK { .. } => same_type(),
             Shift(..) | Reverse => same_type(),
-            IsNotNull | IsNull | Not => with_dtype(DataType::Boolean),
+            IsNotNull | IsNull | Not | IsUnique | IsDuplicated => with_dtype(DataType::Boolean),
         }
     }
 }
diff --git a/polars/polars-lazy/src/dsl/mod.rs b/polars/polars-lazy/src/dsl/mod.rs
@@ -1137,13 +1137,13 @@ impl Expr {
     /// Get a mask of duplicated values
     #[allow(clippy::wrong_self_convention)]
     pub fn is_duplicated(self) -> Self {
-        Expr::Duplicated(Box::new(self))
+        self.apply_private(FunctionExpr::IsDuplicated)
     }
 
     /// Get a mask of unique values
     #[allow(clippy::wrong_self_convention)]
     pub fn is_unique(self) -> Self {
-        Expr::IsUnique(Box::new(self))
+        self.apply_private(FunctionExpr::IsUnique)
     }
 
     /// and operation

diff --git a/polars/polars-lazy/src/logical_plan/aexpr.rs b/polars/polars-lazy/src/logical_plan/aexpr.rs
@@ -35,8 +35,6 @@ pub enum AAggExpr {
 // AExpr representation of Nodes which are allocated in an Arena
 #[derive(Clone, Debug)]
 pub enum AExpr {
-    IsUnique(Node),
-    Duplicated(Node),
     Explode(Node),
     Alias(Node, Arc<str>),
     Column(Arc<str>),
@@ -127,10 +125,9 @@ impl AExpr {
             | Window { .. }
             | Count
             | Slice { .. }
-            | Duplicated(_)
             | Take { .. }
             | Nth(_)
-            | IsUnique(_) => true,
+             => true,
             | Alias(_, _)
             | Explode(_)
             | Column(_)
@@ -196,14 +193,6 @@ impl AExpr {
                 let e = arena.get(*function);
                 e.to_field(schema, ctxt, arena)
             }
-            IsUnique(expr) => {
-                let field = arena.get(*expr).to_field(schema, ctxt, arena)?;
-                Ok(Field::new(field.name(), DataType::Boolean))
-            }
-            Duplicated(expr) => {
-                let field = arena.get(*expr).to_field(schema, ctxt, arena)?;
-                Ok(Field::new(field.name(), DataType::Boolean))
-            }
             Explode(expr) => {
                 let field = arena.get(*expr).to_field(schema, ctxt, arena)?;
 

diff --git a/polars/polars-lazy/src/logical_plan/conversion.rs b/polars/polars-lazy/src/logical_plan/conversion.rs
@@ -9,8 +9,6 @@ fn to_aexprs(input: Vec<Expr>, arena: &mut Arena<AExpr>) -> Vec<Node> {
 /// converts expression to AExpr and adds it to the arena, which uses an arena (Vec) for allocation
 pub(crate) fn to_aexpr(expr: Expr, arena: &mut Arena<AExpr>) -> Node {
     let v = match expr {
-        Expr::IsUnique(expr) => AExpr::IsUnique(to_aexpr(*expr, arena)),
-        Expr::Duplicated(expr) => AExpr::Duplicated(to_aexpr(*expr, arena)),
         Expr::Explode(expr) => AExpr::Explode(to_aexpr(*expr, arena)),
         Expr::Alias(e, name) => AExpr::Alias(to_aexpr(*e, arena), name),
         Expr::Literal(value) => AExpr::Literal(value),
@@ -443,8 +441,6 @@ pub(crate) fn node_to_expr(node: Node, expr_arena: &Arena<AExpr>) -> Expr {
     let expr = expr_arena.get(node).clone();
 
     match expr {
-        AExpr::Duplicated(node) => Expr::Duplicated(Box::new(node_to_expr(node, expr_arena))),
-        AExpr::IsUnique(node) => Expr::IsUnique(Box::new(node_to_expr(node, expr_arena))),
         AExpr::Explode(node) => Expr::Explode(Box::new(node_to_expr(node, expr_arena))),
         AExpr::Alias(expr, name) => {
             let exp = node_to_expr(expr, expr_arena);

diff --git a/polars/polars-lazy/src/logical_plan/format.rs b/polars/polars-lazy/src/logical_plan/format.rs
@@ -200,9 +200,7 @@ impl fmt::Debug for Expr {
             } => write!(f, "{:?}.over({:?})", function, partition_by),
             Nth(i) => write!(f, "nth({})", i),
             Count => write!(f, "count()"),
-            IsUnique(expr) => write!(f, "{:?}.unique()", expr),
             Explode(expr) => write!(f, "{:?}.explode()", expr),
-            Duplicated(expr) => write!(f, "{:?}.is_duplicate()", expr),
             Alias(expr, name) => write!(f, "{:?}.alias(\"{}\")", expr, name),
             Column(name) => write!(f, "col(\"{}\")", name),
             Literal(v) => {

diff --git a/polars/polars-lazy/src/logical_plan/iterator.rs b/polars/polars-lazy/src/logical_plan/iterator.rs
@@ -61,8 +61,6 @@ macro_rules! push_expr {
             // as the root columns/ input columns by `_suffix` and `_keep_name` etc.
             AnonymousFunction { input, .. } => input.$iter().rev().for_each(|e| $push(e)),
             Function { input, .. } => input.$iter().rev().for_each(|e| $push(e)),
-            Duplicated(e) => $push(e),
-            IsUnique(e) => $push(e),
             Explode(e) => $push(e),
             Window {
                 function,
@@ -227,8 +225,6 @@ impl AExpr {
             {
                 input.iter().rev().for_each(push)
             }
-            Duplicated(e) => push(e),
-            IsUnique(e) => push(e),
             Explode(e) => push(e),
             Window {
                 function,

diff --git a/polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/mod.rs b/polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/mod.rs
@@ -426,7 +426,7 @@ impl PredicatePushDown {
                 for (_, predicate) in acc_predicates {
                     // unique and duplicated can be caused by joins
                     let matches =
-                        |e: &AExpr| matches!(e, AExpr::IsUnique(_) | AExpr::Duplicated(_));
+                        |e: &AExpr| matches!(e, AExpr::Function{function: FunctionExpr::IsDuplicated | FunctionExpr::IsUnique, ..});
 
                     let checks_nulls =
                         |e: &AExpr| matches!(e, AExpr::Function{function: FunctionExpr::IsNotNull | FunctionExpr::IsNull, ..} ) ||

diff --git a/polars/polars-lazy/src/physical_plan/planner/expr.rs b/polars/polars-lazy/src/physical_plan/planner/expr.rs
@@ -524,32 +524,6 @@ impl PhysicalPlanner {
                     expr: node_to_expr(expression, expr_arena),
                 }))
             }
-            Duplicated(expr) => {
-                let input = self.create_physical_expr(expr, ctxt, expr_arena)?;
-                let function = SpecialEq::new(Arc::new(move |s: &mut [Series]| {
-                    let s = std::mem::take(&mut s[0]);
-                    s.is_duplicated().map(|ca| ca.into_series())
-                }) as Arc<dyn SeriesUdf>);
-                Ok(Arc::new(ApplyExpr::new_minimal(
-                    vec![input],
-                    function,
-                    node_to_expr(expression, expr_arena),
-                    ApplyOptions::ApplyGroups,
-                )))
-            }
-            IsUnique(expr) => {
-                let input = self.create_physical_expr(expr, ctxt, expr_arena)?;
-                let function = SpecialEq::new(Arc::new(move |s: &mut [Series]| {
-                    let s = std::mem::take(&mut s[0]);
-                    s.is_unique().map(|ca| ca.into_series())
-                }) as Arc<dyn SeriesUdf>);
-                Ok(Arc::new(ApplyExpr::new_minimal(
-                    vec![input],
-                    function,
-                    node_to_expr(expression, expr_arena),
-                    ApplyOptions::ApplyGroups,
-                )))
-            }
             Explode(expr) => {
                 let input = self.create_physical_expr(expr, ctxt, expr_arena)?;
                 let function = SpecialEq::new(Arc::new(move |s: &mut [Series]| {

diff --git a/py-polars/tests/test_object.py → py-polars/tests/unit/test_object.py b/py-polars/tests/test_object.py → py-polars/tests/unit/test_object.py