fix: Fix all-null list aggregations returning Null dtype

pola-rs · Jan 30, 2025 · 7f04089 · 7f04089
1 parent 0bc6cb8
commit 7f04089
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 27 deletions.
diff --git a/crates/polars-ops/src/chunked_array/list/min_max.rs b/crates/polars-ops/src/chunked_array/list/min_max.rs
@@ -89,15 +89,22 @@ pub(super) fn list_min_function(ca: &ListChunked) -> PolarsResult<Series> {
                     unsafe { out.into_series().from_physical_unchecked(dt) }
                 })
             },
-            _ => Ok(ca
-                .try_apply_amortized(|s| {
-                    let s = s.as_ref();
-                    let sc = s.min_reduce()?;
-                    Ok(sc.into_series(s.name().clone()))
-                })?
-                .explode()
-                .unwrap()
-                .into_series()),
+            dt => {
+                let out = ca
+                    .try_apply_amortized(|s| {
+                        let s = s.as_ref();
+                        let sc = s.min_reduce()?;
+                        Ok(sc.into_series(s.name().clone()))
+                    })?
+                    .explode()
+                    .unwrap()
+                    .into_series();
+
+                match out.dtype() {
+                    DataType::Null => out.cast(dt),
+                    _ => Ok(out),
+                }
+            },
         }
     }
 
@@ -199,15 +206,22 @@ pub(super) fn list_max_function(ca: &ListChunked) -> PolarsResult<Series> {
                     unsafe { out.into_series().from_physical_unchecked(dt) }
                 })
             },
-            _ => Ok(ca
-                .try_apply_amortized(|s| {
-                    let s = s.as_ref();
-                    let sc = s.max_reduce()?;
-                    Ok(sc.into_series(s.name().clone()))
-                })?
-                .explode()
-                .unwrap()
-                .into_series()),
+            dt => {
+                let out = ca
+                    .try_apply_amortized(|s| {
+                        let s = s.as_ref();
+                        let sc = s.max_reduce()?;
+                        Ok(sc.into_series(s.name().clone()))
+                    })?
+                    .explode()
+                    .unwrap()
+                    .into_series();
+
+                match out.dtype() {
+                    DataType::Null => out.cast(dt),
+                    _ => Ok(out),
+                }
+            },
         }
     }
 

diff --git a/crates/polars-ops/src/chunked_array/list/sum_mean.rs b/crates/polars-ops/src/chunked_array/list/sum_mean.rs
@@ -106,15 +106,22 @@ pub(super) fn sum_with_nulls(ca: &ListChunked, inner_dtype: &DataType) -> Polars
             out.into_series()
         },
         // slowest sum_as_series path
-        _ => ca
-            .try_apply_amortized(|s| {
-                s.as_ref()
-                    .sum_reduce()
-                    .map(|sc| sc.into_series(PlSmallStr::EMPTY))
-            })?
-            .explode()
-            .unwrap()
-            .into_series(),
+        dt => {
+            let s = ca
+                .try_apply_amortized(|s| {
+                    s.as_ref()
+                        .sum_reduce()
+                        .map(|sc| sc.into_series(PlSmallStr::EMPTY))
+                })?
+                .explode()
+                .unwrap()
+                .into_series();
+
+            match s.dtype() {
+                DataType::Null => s.cast(dt)?,
+                _ => s,
+            }
+        },
     };
     out.rename(ca.name().clone());
     Ok(out)

diff --git a/py-polars/tests/unit/operations/namespaces/list/test_list.py b/py-polars/tests/unit/operations/namespaces/list/test_list.py
@@ -2,6 +2,7 @@
 
 import re
 from datetime import date, datetime
+from typing import TYPE_CHECKING
 
 import numpy as np
 import pytest
@@ -14,6 +15,9 @@
 )
 from polars.testing import assert_frame_equal, assert_series_equal
 
+if TYPE_CHECKING:
+    from polars._typing import PolarsDataType
+
 
 def test_list_arr_get() -> None:
     a = pl.Series("a", [[1, 2, 3], [4, 5], [6, 7, 8, 9]])
@@ -979,3 +983,13 @@ def test_list_eval_element_schema_19345() -> None:
         ),
         pl.DataFrame({"a": [[1]]}),
     )
+
+
+@pytest.mark.parametrize(
+    "inner_dtype",
+    [pl.Int8, pl.Float64, pl.String, pl.Duration],
+)
+def test_list_agg_all_null(inner_dtype: PolarsDataType) -> None:
+    s = pl.Series([None, None], dtype=pl.List(inner_dtype))
+    assert s.list.min().dtype == inner_dtype
+    assert s.list.max().dtype == inner_dtype