fix(rust, python): stable list namespace ouput when all elements are … (

#9610)
pola-rs · Jun 29, 2023 · 8823a11 · 8823a11
1 parent f7da720
commit 8823a11
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 8 deletions.
diff --git a/polars/polars-json/src/ndjson/file.rs b/polars/polars-json/src/ndjson/file.rs
@@ -144,6 +144,5 @@ pub fn infer_iter<A: AsRef<str>>(rows: impl Iterator<Item = A>) -> PolarsResult<
     }
 
     let v: Vec<&DataType> = data_types.iter().collect();
-    dbg!(&v);
-    dbg!(Ok(crate::json::infer_schema::coerce_data_type(&v)))
+    Ok(crate::json::infer_schema::coerce_data_type(&v))
 }
diff --git a/polars/polars-ops/src/chunked_array/list/namespace.rs b/polars/polars-ops/src/chunked_array/list/namespace.rs
@@ -145,26 +145,40 @@ pub trait ListNameSpaceImpl: AsList {
         }
     }
 
+    fn same_type(&self, out: ListChunked) -> ListChunked {
+        let ca = self.as_list();
+        let dtype = ca.dtype();
+        if out.dtype() != dtype {
+            out.cast(ca.dtype()).unwrap().list().unwrap().clone()
+        } else {
+            out
+        }
+    }
+
     #[must_use]
     fn lst_sort(&self, options: SortOptions) -> ListChunked {
         let ca = self.as_list();
-        ca.apply_amortized(|s| s.as_ref().sort_with(options))
+        let out = ca.apply_amortized(|s| s.as_ref().sort_with(options));
+        self.same_type(out)
     }
 
     #[must_use]
     fn lst_reverse(&self) -> ListChunked {
         let ca = self.as_list();
-        ca.apply_amortized(|s| s.as_ref().reverse())
+        let out = ca.apply_amortized(|s| s.as_ref().reverse());
+        self.same_type(out)
     }
 
     fn lst_unique(&self) -> PolarsResult<ListChunked> {
         let ca = self.as_list();
-        ca.try_apply_amortized(|s| s.as_ref().unique())
+        let out = ca.try_apply_amortized(|s| s.as_ref().unique())?;
+        Ok(self.same_type(out))
     }
 
     fn lst_unique_stable(&self) -> PolarsResult<ListChunked> {
         let ca = self.as_list();
-        ca.try_apply_amortized(|s| s.as_ref().unique_stable())
+        let out = ca.try_apply_amortized(|s| s.as_ref().unique_stable())?;
+        Ok(self.same_type(out))
     }
 
     fn lst_arg_min(&self) -> IdxCa {
@@ -195,12 +209,14 @@ pub trait ListNameSpaceImpl: AsList {
 
     fn lst_shift(&self, periods: i64) -> ListChunked {
         let ca = self.as_list();
-        ca.apply_amortized(|s| s.as_ref().shift(periods))
+        let out = ca.apply_amortized(|s| s.as_ref().shift(periods));
+        self.same_type(out)
     }
 
     fn lst_slice(&self, offset: i64, length: usize) -> ListChunked {
         let ca = self.as_list();
-        ca.apply_amortized(|s| s.as_ref().slice(offset, length))
+        let out = ca.apply_amortized(|s| s.as_ref().slice(offset, length));
+        self.same_type(out)
     }
 
     fn lst_lengths(&self) -> IdxCa {

diff --git a/py-polars/tests/unit/test_empty.py b/py-polars/tests/unit/test_empty.py
@@ -78,3 +78,12 @@ def test_empty_groupby_apply_err() -> None:
         pl.ComputeError, match=r"cannot groupby \+ apply on empty 'DataFrame'"
     ):
         df.groupby("x").apply(lambda x: x)
+
+
+def test_empty_list_namespace_output_9585() -> None:
+    dtype = pl.List(pl.Utf8)
+    names = ["sort", "unique", "head", "tail", "shift", "reverse"]
+    df = pl.DataFrame([[None]], schema={"A": dtype})
+    assert df.select(
+        [eval(f"pl.col('A').list.{name}().suffix(f'_{name}')") for name in names]
+    ).dtypes == [dtype] * len(names)