Remove mutual exclusivity from aggregate and returns_scalar

pola-rs · Feb 17, 2025 · bf6dda3 · bf6dda3
1 parent 5e51218
commit bf6dda3
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 11 deletions.
diff --git a/crates/polars-expr/src/expressions/apply.rs b/crates/polars-expr/src/expressions/apply.rs
@@ -46,13 +46,6 @@ impl ApplyExpr {
         output_field: Field,
         returns_scalar: bool,
     ) -> Self {
-        #[cfg(debug_assertions)]
-        if matches!(options.collect_groups, ApplyOptions::ElementWise)
-            && options.flags.contains(FunctionFlags::RETURNS_SCALAR)
-        {
-            panic!("expr {:?} is not implemented correctly. 'returns_scalar' and 'elementwise' are mutually exclusive", expr)
-        }
-
         Self {
             inputs,
             function,

diff --git a/crates/polars-plan/src/dsl/dt.rs b/crates/polars-plan/src/dsl/dt.rs
@@ -242,10 +242,11 @@ impl DateLikeNameSpace {
 
     /// Round the Datetime/Date range into buckets.
     pub fn round(self, every: Expr) -> Expr {
+        let returns_scalar = all_return_scalar(&self.0);
         self.0.map_many_private(
             FunctionExpr::TemporalExpr(TemporalFunction::Round),
             &[every],
-            false,
+            returns_scalar,
             None,
         )
     }
@@ -254,10 +255,11 @@ impl DateLikeNameSpace {
     /// This will take leap years/ months into account.
     #[cfg(feature = "offset_by")]
     pub fn offset_by(self, by: Expr) -> Expr {
+        let returns_scalar = all_return_scalar(&self.0);
         self.0.map_many_private(
             FunctionExpr::TemporalExpr(TemporalFunction::OffsetBy),
             &[by],
-            false,
+            returns_scalar,
             None,
         )
     }
@@ -269,20 +271,22 @@ impl DateLikeNameSpace {
         ambiguous: Expr,
         non_existent: NonExistent,
     ) -> Expr {
+        let returns_scalar = all_return_scalar(&self.0);
         self.0.map_many_private(
             FunctionExpr::TemporalExpr(TemporalFunction::ReplaceTimeZone(time_zone, non_existent)),
             &[ambiguous],
-            false,
+            returns_scalar,
             None,
         )
     }
 
     /// Combine an existing Date/Datetime with a Time, creating a new Datetime value.
     pub fn combine(self, time: Expr, tu: TimeUnit) -> Expr {
+        let returns_scalar = all_return_scalar(&self.0) && all_return_scalar(&time);
         self.0.map_many_private(
             FunctionExpr::TemporalExpr(TemporalFunction::Combine(tu)),
             &[time],
-            false,
+            returns_scalar,
             None,
         )
     }