fix take/get relation

pola-rs · Oct 23, 2023 · 6b824ec · 6b824ec
1 parent 8eba908
commit 6b824ec
Show file tree

Hide file tree

Showing 24 changed files with 190 additions and 53 deletions.
diff --git a/crates/polars-lazy/src/physical_plan/expressions/apply.rs b/crates/polars-lazy/src/physical_plan/expressions/apply.rs
@@ -36,6 +36,7 @@ impl ApplyExpr {
         allow_threading: bool,
         input_schema: Option<SchemaRef>,
     ) -> Self {
+        #[cfg(debug_assertions)]
         if matches!(options.collect_groups, ApplyOptions::ElementWise) && options.returns_scalar {
             panic!("expr {} is not implemented correctly. 'returns_scalar' and 'elementwise' are mutually exclusive", expr)
         }

diff --git a/crates/polars-lazy/src/physical_plan/expressions/take.rs b/crates/polars-lazy/src/physical_plan/expressions/take.rs
@@ -12,6 +12,7 @@ pub struct TakeExpr {
     pub(crate) phys_expr: Arc<dyn PhysicalExpr>,
     pub(crate) idx: Arc<dyn PhysicalExpr>,
     pub(crate) expr: Expr,
+    pub(crate) returns_scalar: bool,
 }
 
 impl TakeExpr {
@@ -101,12 +102,23 @@ impl PhysicalExpr for TakeExpr {
                     },
                 };
                 let taken = ac.flat_naive().take(&idx)?;
+
+                let taken = if self.returns_scalar {
+                    taken
+                } else {
+                    taken.as_list().into_series()
+                };
+
                 ac.with_series(taken, true, Some(&self.expr))?;
                 return Ok(ac);
             },
-            AggState::AggregatedList(s) => s.list().unwrap().clone(),
+            AggState::AggregatedList(s) => {
+                polars_ensure!(!self.returns_scalar, ComputeError: "expected single index");
+                s.list().unwrap().clone()
+            },
             // Maybe a literal as well, this needs a different path.
             AggState::NotAggregated(_) => {
+                polars_ensure!(!self.returns_scalar, ComputeError: "expected single index");
                 let s = idx.aggregated();
                 s.list().unwrap().clone()
             },
@@ -144,6 +156,13 @@ impl PhysicalExpr for TakeExpr {
                                 },
                             };
                             let taken = ac.flat_naive().take(&idx.into_inner())?;
+
+                            let taken = if self.returns_scalar {
+                                taken
+                            } else {
+                                taken.as_list().into_series()
+                            };
+
                             ac.with_series(taken, true, Some(&self.expr))?;
                             ac.with_update_groups(UpdateGroups::WithGroupsLen);
                             Ok(ac)

diff --git a/crates/polars-lazy/src/physical_plan/planner/expr.rs b/crates/polars-lazy/src/physical_plan/planner/expr.rs
@@ -192,13 +192,18 @@ pub(crate) fn create_physical_expr(
                 node_to_expr(expression, expr_arena),
             )))
         },
-        Take { expr, idx } => {
+        Take {
+            expr,
+            idx,
+            returns_scalar,
+        } => {
             let phys_expr = create_physical_expr(expr, ctxt, expr_arena, schema, state)?;
             let phys_idx = create_physical_expr(idx, ctxt, expr_arena, schema, state)?;
             Ok(Arc::new(TakeExpr {
                 phys_expr,
                 idx: phys_idx,
                 expr: node_to_expr(expression, expr_arena),
+                returns_scalar,
             }))
         },
         SortBy {

diff --git a/crates/polars-plan/src/dsl/binary.rs b/crates/polars-plan/src/dsl/binary.rs
@@ -9,7 +9,7 @@ impl BinaryNameSpace {
         self.0.map_many_private(
             FunctionExpr::BinaryExpr(BinaryFunction::Contains),
             &[pat],
-            true,
+            false,
             true,
         )
     }
@@ -19,7 +19,7 @@ impl BinaryNameSpace {
         self.0.map_many_private(
             FunctionExpr::BinaryExpr(BinaryFunction::EndsWith),
             &[sub],
-            true,
+            false,
             true,
         )
     }
@@ -29,7 +29,7 @@ impl BinaryNameSpace {
         self.0.map_many_private(
             FunctionExpr::BinaryExpr(BinaryFunction::StartsWith),
             &[sub],
-            true,
+            false,
             true,
         )
     }

diff --git a/crates/polars-plan/src/dsl/dt.rs b/crates/polars-plan/src/dsl/dt.rs
@@ -172,7 +172,7 @@ impl DateLikeNameSpace {
         self.0.map_many_private(
             FunctionExpr::TemporalExpr(TemporalFunction::Truncate(offset)),
             &[every, ambiguous],
-            true,
+            false,
             false,
         )
     }

diff --git a/crates/polars-plan/src/dsl/expr.rs b/crates/polars-plan/src/dsl/expr.rs
@@ -89,6 +89,7 @@ pub enum Expr {
     Take {
         expr: Box<Expr>,
         idx: Box<Expr>,
+        returns_scalar: bool,
     },
     SortBy {
         expr: Box<Expr>,

diff --git a/crates/polars-plan/src/dsl/functions/concat.rs b/crates/polars-plan/src/dsl/functions/concat.rs
@@ -12,7 +12,7 @@ pub fn concat_str<E: AsRef<[Expr]>>(s: E, separator: &str) -> Expr {
         options: FunctionOptions {
             collect_groups: ApplyOptions::ElementWise,
             input_wildcard_expansion: true,
-            returns_scalar: true,
+            returns_scalar: false,
             ..Default::default()
         },
     }

diff --git a/crates/polars-plan/src/dsl/functions/horizontal.rs b/crates/polars-plan/src/dsl/functions/horizontal.rs
@@ -201,7 +201,7 @@ pub fn all_horizontal<E: AsRef<[Expr]>>(exprs: E) -> Expr {
         options: FunctionOptions {
             collect_groups: ApplyOptions::ElementWise,
             input_wildcard_expansion: true,
-            returns_scalar: true,
+            returns_scalar: false,
             cast_to_supertypes: false,
             allow_rename: true,
             ..Default::default()
@@ -220,7 +220,7 @@ pub fn any_horizontal<E: AsRef<[Expr]>>(exprs: E) -> Expr {
         options: FunctionOptions {
             collect_groups: ApplyOptions::ElementWise,
             input_wildcard_expansion: true,
-            returns_scalar: true,
+            returns_scalar: false,
             cast_to_supertypes: false,
             allow_rename: true,
             ..Default::default()
@@ -243,7 +243,7 @@ pub fn max_horizontal<E: AsRef<[Expr]>>(exprs: E) -> Expr {
         options: FunctionOptions {
             collect_groups: ApplyOptions::ElementWise,
             input_wildcard_expansion: true,
-            returns_scalar: true,
+            returns_scalar: false,
             allow_rename: true,
             ..Default::default()
         },
@@ -265,7 +265,7 @@ pub fn min_horizontal<E: AsRef<[Expr]>>(exprs: E) -> Expr {
         options: FunctionOptions {
             collect_groups: ApplyOptions::ElementWise,
             input_wildcard_expansion: true,
-            returns_scalar: true,
+            returns_scalar: false,
             allow_rename: true,
             ..Default::default()
         },
@@ -284,7 +284,7 @@ pub fn sum_horizontal<E: AsRef<[Expr]>>(exprs: E) -> Expr {
         options: FunctionOptions {
             collect_groups: ApplyOptions::ElementWise,
             input_wildcard_expansion: true,
-            returns_scalar: true,
+            returns_scalar: false,
             cast_to_supertypes: false,
             allow_rename: true,
             ..Default::default()

diff --git a/crates/polars-plan/src/dsl/list.rs b/crates/polars-plan/src/dsl/list.rs
@@ -137,7 +137,7 @@ impl ListNameSpace {
         self.0.map_many_private(
             FunctionExpr::ListExpr(ListFunction::Get),
             &[index],
-            true,
+            false,
             false,
         )
     }
@@ -152,7 +152,7 @@ impl ListNameSpace {
         self.0.map_many_private(
             FunctionExpr::ListExpr(ListFunction::Take(null_on_oob)),
             &[index],
-            true,
+            false,
             false,
         )
     }
@@ -216,7 +216,7 @@ impl ListNameSpace {
         self.0.map_many_private(
             FunctionExpr::ListExpr(ListFunction::Slice),
             &[offset, length],
-            true,
+            false,
             false,
         )
     }
@@ -296,7 +296,7 @@ impl ListNameSpace {
             .map_many_private(
                 FunctionExpr::ListExpr(ListFunction::Contains),
                 &[other],
-                true,
+                false,
                 false,
             )
             .with_function_options(|mut options| {
@@ -313,7 +313,7 @@ impl ListNameSpace {
             .map_many_private(
                 FunctionExpr::ListExpr(ListFunction::CountMatches),
                 &[other],
-                true,
+                false,
                 false,
             )
             .with_function_options(|mut options| {

diff --git a/crates/polars-plan/src/dsl/mod.rs b/crates/polars-plan/src/dsl/mod.rs
@@ -444,6 +444,16 @@ impl Expr {
         Expr::Take {
             expr: Box::new(self),
             idx: Box::new(idx.into()),
+            returns_scalar: false,
+        }
+    }
+
+    /// Take the values by a single index.
+    pub fn get<E: Into<Expr>>(self, idx: E) -> Self {
+        Expr::Take {
+            expr: Box::new(self),
+            idx: Box::new(idx.into()),
+            returns_scalar: true,
         }
     }
 
@@ -679,7 +689,7 @@ impl Expr {
         self,
         function_expr: FunctionExpr,
         arguments: &[Expr],
-        auto_explode: bool,
+        returns_scalar: bool,
         cast_to_supertypes: bool,
     ) -> Self {
         let mut input = Vec::with_capacity(arguments.len() + 1);
@@ -691,7 +701,7 @@ impl Expr {
             function: function_expr,
             options: FunctionOptions {
                 collect_groups: ApplyOptions::ElementWise,
-                returns_scalar: auto_explode,
+                returns_scalar,
                 cast_to_supertypes,
                 ..Default::default()
             },
@@ -1060,12 +1070,25 @@ impl Expr {
         let other = other.into();
         let has_literal = has_leaf_literal(&other);
 
+        // lit(true).is_in() returns a scalar.
+        let returns_scalar = all_leaf_literal(&self);
+
         let arguments = &[other];
         // we don't have to apply on groups, so this is faster
         if has_literal {
-            self.map_many_private(BooleanFunction::IsIn.into(), arguments, true, true)
+            self.map_many_private(
+                BooleanFunction::IsIn.into(),
+                arguments,
+                returns_scalar,
+                true,
+            )
         } else {
-            self.apply_many_private(BooleanFunction::IsIn.into(), arguments, true, true)
+            self.apply_many_private(
+                BooleanFunction::IsIn.into(),
+                arguments,
+                returns_scalar,
+                true,
+            )
         }
     }
 

diff --git a/crates/polars-plan/src/dsl/string.rs b/crates/polars-plan/src/dsl/string.rs
@@ -13,7 +13,7 @@ impl StringNameSpace {
                 strict: false,
             }),
             &[pat],
-            true,
+            false,
             true,
         )
     }
@@ -28,7 +28,7 @@ impl StringNameSpace {
                 strict,
             }),
             &[pat],
-            true,
+            false,
             true,
         )
     }
@@ -38,7 +38,7 @@ impl StringNameSpace {
         self.0.map_many_private(
             FunctionExpr::StringExpr(StringFunction::EndsWith),
             &[sub],
-            true,
+            false,
             true,
         )
     }
@@ -48,7 +48,7 @@ impl StringNameSpace {
         self.0.map_many_private(
             FunctionExpr::StringExpr(StringFunction::StartsWith),
             &[sub],
-            true,
+            false,
             true,
         )
     }
@@ -131,7 +131,7 @@ impl StringNameSpace {
         self.0.map_many_private(
             StringFunction::CountMatches(literal).into(),
             &[pat],
-            true,
+            false,
             false,
         )
     }
@@ -142,7 +142,7 @@ impl StringNameSpace {
         self.0.map_many_private(
             StringFunction::Strptime(dtype, options).into(),
             &[ambiguous],
-            true,
+            false,
             false,
         )
     }
@@ -207,6 +207,7 @@ impl StringNameSpace {
             .apply_private(StringFunction::ConcatVertical(delimiter.to_owned()).into())
             .with_function_options(|mut options| {
                 options.returns_scalar = true;
+                options.collect_groups = ApplyOptions::GroupWise;
                 options
             })
     }

diff --git a/crates/polars-plan/src/logical_plan/aexpr/mod.rs b/crates/polars-plan/src/logical_plan/aexpr/mod.rs
@@ -146,6 +146,7 @@ pub enum AExpr {
     Take {
         expr: Node,
         idx: Node,
+        returns_scalar: bool,
     },
     SortBy {
         expr: Node,
@@ -267,7 +268,7 @@ impl AExpr {
             },
             Cast { expr, .. } => container.push(*expr),
             Sort { expr, .. } => container.push(*expr),
-            Take { expr, idx } => {
+            Take { expr, idx, .. } => {
                 container.push(*idx);
                 // latest, so that it is popped first
                 container.push(*expr);
@@ -346,7 +347,7 @@ impl AExpr {
                 *left = inputs[1];
                 return self;
             },
-            Take { expr, idx } => {
+            Take { expr, idx, .. } => {
                 *idx = inputs[0];
                 *expr = inputs[1];
                 return self;