fix lazy schema (#4027)

polars/polars-lazy/src/logical_plan/aexpr.rs

@@ -257,7 +257,8 @@ impl AExpr {
                 use AAggExpr::*;
                 match agg {
                     Max(expr) | Sum(expr) | Min(expr) | First(expr) | Last(expr) => {
-                        arena.get(*expr).to_field(schema, ctxt, arena)
+                        // default context because `col()` would return a list in aggregation context
+                        arena.get(*expr).to_field(schema, Context::Default, arena)
                     }
                     Median(expr) => {
                         let mut field = arena.get(*expr).to_field(schema, ctxt, arena)?;
@@ -272,7 +273,9 @@ impl AExpr {
                         Ok(field)
                     }
                     List(expr) => {
-                        let mut field = arena.get(*expr).to_field(schema, ctxt, arena)?;
+                        // default context because `col()` would return a list in aggregation context
+                        let mut field =
+                            arena.get(*expr).to_field(schema, Context::Default, arena)?;
                         field.coerce(DataType::List(field.data_type().clone().into()));
                         Ok(field)
                     }

py-polars/tests/test_schema.py

@@ -0,0 +1,26 @@
+import polars as pl
+
+
+def test_schema_on_agg() -> None:
+    df = pl.DataFrame({"a": ["x", "x", "y", "n"], "b": [1, 2, 3, 4]})
+
+    assert (
+        df.lazy()
+        .groupby("a")
+        .agg(
+            [
+                pl.col("b").min().alias("min"),
+                pl.col("b").max().alias("max"),
+                pl.col("b").sum().alias("sum"),
+                pl.col("b").first().alias("first"),
+                pl.col("b").last().alias("last"),
+            ]
+        )
+    ).schema == {
+        "a": pl.Utf8,
+        "min": pl.Int64,
+        "max": pl.Int64,
+        "sum": pl.Int64,
+        "first": pl.Int64,
+        "last": pl.Int64,
+    }