don't automatically flatten. If an apply needs a flatten, just add a …

…flatten to the expr
pola-rs · Nov 20, 2021 · 3ea632e · 3ea632e
1 parent 4dc472f
commit 3ea632e
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 25 deletions.
diff --git a/polars/polars-lazy/src/physical_plan/expressions/apply.rs b/polars/polars-lazy/src/physical_plan/expressions/apply.rs
@@ -180,8 +180,6 @@ impl PhysicalAggregation for ApplyExpr {
                     let mut container = [Default::default()];
                     let name = ac.series().name().to_string();
 
-                    let mut all_unit_length = true;
-
                     let mut ca: ListChunked = ac
                         .aggregated()
                         .list()
@@ -190,21 +188,11 @@ impl PhysicalAggregation for ApplyExpr {
                         .map(|opt_s| {
                             opt_s.and_then(|s| {
                                 container[0] = s;
-                                let out = self.function.call_udf(&mut container).ok();
-
-                                if let Some(s) = &out {
-                                    if s.len() != 1 {
-                                        all_unit_length = false;
-                                    }
-                                }
-                                out
+                                self.function.call_udf(&mut container).ok()
                             })
                         })
                         .collect();
                     ca.rename(&name);
-                    if all_unit_length {
-                        return Ok(Some(ca.explode()?));
-                    }
                     Ok(Some(ca.into_series()))
                 }
                 ApplyOptions::ApplyFlat => {
@@ -232,7 +220,6 @@ impl PhysicalAggregation for ApplyExpr {
 
             match self.collect_groups {
                 ApplyOptions::ApplyGroups => {
-                    let mut all_unit_length = true;
                     let mut container = vec![Default::default(); acs.len()];
                     let name = acs[0].series().name().to_string();
 
@@ -259,20 +246,10 @@ impl PhysicalAggregation for ApplyExpr {
                                     Some(s) => container.push(s),
                                 }
                             }
-                            let out = self.function.call_udf(&mut container).ok();
-
-                            if let Some(s) = &out {
-                                if s.len() != 1 {
-                                    all_unit_length = false;
-                                }
-                            }
-                            out
+                            self.function.call_udf(&mut container).ok()
                         })
                         .collect();
                     ca.rename(&name);
-                    if all_unit_length {
-                        return Ok(Some(ca.explode()?));
-                    }
                     Ok(Some(ca.into_series()))
                 }
                 ApplyOptions::ApplyFlat => {

diff --git a/polars/polars-lazy/src/test.rs b/polars/polars-lazy/src/test.rs
@@ -2201,3 +2201,22 @@ fn test_groupby_on_lists() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_single_group_result() -> Result<()> {
+    // the argsort should not auto explode
+    let df = df![
+        "a" => [1, 2],
+        "b" => [1, 1]
+    ]?;
+
+    let out = df
+        .lazy()
+        .select([col("a").arg_sort(false).over([col("a")]).flatten()])
+        .collect()?;
+
+    let a = out.column("a")?.u32()?;
+    assert_eq!(Vec::from(a), &[Some(0), Some(0)]);
+
+    Ok(())
+}