ensure that Cast expressions first updates groups before it flattens (#…

…3401)

polars/polars-core/src/datatypes.rs

@@ -749,6 +749,21 @@ impl DataType {
             _ => true,
         }
     }
+    pub fn is_signed(&self) -> bool {
+        // allow because it cannot be replaced when object feature is activated
+        #[allow(clippy::match_like_matches_macro)]
+        match self {
+            #[cfg(feature = "dtype-i8")]
+            DataType::Int8 => true,
+            #[cfg(feature = "dtype-i16")]
+            DataType::Int16 => true,
+            DataType::Int32 | DataType::Int64 => true,
+            _ => false,
+        }
+    }
+    pub fn is_unsigned(&self) -> bool {
+        self.is_numeric() && !self.is_signed()
+    }
 
     /// Convert to an Arrow data type.
     pub fn to_arrow(&self) -> ArrowDataType {

polars/polars-lazy/src/logical_plan/optimizer/type_coercion.rs

@@ -39,11 +39,17 @@ fn use_supertype(
 
             // cast literal to right type
             (AExpr::Literal(_), _) => {
-                st = type_right.clone();
+                // never cast signed to unsigned
+                if type_right.is_signed() {
+                    st = type_right.clone();
+                }
             }
             // cast literal to left type
             (_, AExpr::Literal(_)) => {
-                st = type_left.clone();
+                // never cast signed to unsigned
+                if type_left.is_signed() {
+                    st = type_left.clone();
+                }
             }
             // do nothing
             _ => {}

polars/polars-lazy/src/physical_plan/expressions/cast.rs

@@ -71,6 +71,8 @@ impl PhysicalExpr for CastExpr {
         state: &ExecutionState,
     ) -> Result<AggregationContext<'a>> {
         let mut ac = self.input.evaluate_on_groups(df, groups, state)?;
+        // before we flatten, make sure that groups are updated
+        ac.groups();
         let s = ac.flat_naive();
         let s = self.finish(s.as_ref())?;
 

polars/tests/it/lazy/expressions/apply.rs

@@ -82,3 +82,30 @@ fn test_expand_list() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+#[cfg(feature = "unique_counts")]
+fn test_update_groups_in_cast() -> Result<()> {
+    let df = df![
+        "group" =>  ["A" ,"A", "A", "B", "B", "B", "B"],
+        "id"=> [1, 2, 1, 4, 5, 4, 6],
+    ]?;
+
+    // optimized to
+    // col("id").unique_counts().cast(int64) * -1
+    // in aggregation that cast coerces a list and the cast may forget to update groups
+    let out = df
+        .lazy()
+        .groupby_stable([col("group")])
+        .agg([col("id").unique_counts() * lit(-1)])
+        .collect()?;
+
+    let expected = df![
+        "group" =>  ["A" ,"B"],
+        "id"=> [AnyValue::List(Series::new("", [-2i64, -1])), AnyValue::List(Series::new("", [-2i64, -1, -1]))]
+    ]?;
+    dbg!(&out, &expected);
+
+    assert!(out.frame_equal(&expected));
+    Ok(())
+}

py-polars/tests/test_queries.py

@@ -86,3 +86,26 @@ def test_overflow_uint16_agg_mean() -> None:
         .to_dict(False)
         == {"col1": ["A"], "col3": [64.0]}
     )
+
+
+def test_binary_on_list_agg_3345() -> None:
+    df = pl.DataFrame(
+        {
+            "group": ["A", "A", "A", "B", "B", "B", "B"],
+            "id": [1, 2, 1, 4, 5, 4, 6],
+        }
+    )
+
+    assert (
+        df.groupby(["group"], maintain_order=True)
+        .agg(
+            [
+                (
+                    (pl.col("id").unique_counts() / pl.col("id").len()).log()
+                    * -1
+                    * (pl.col("id").unique_counts() / pl.col("id").len())
+                ).sum()
+            ]
+        )
+        .to_dict(False)
+    ) == {"group": ["A", "B"], "id": [0.6365141682948128, 1.0397207708399179]}