use default context for col upstream col expression type (#4219)

pola-rs · Aug 2, 2022 · af35ce0 · af35ce0
1 parent b2ff9fa
commit af35ce0
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/polars/polars-lazy/src/logical_plan/aexpr.rs b/polars/polars-lazy/src/logical_plan/aexpr.rs
@@ -334,7 +334,8 @@ impl AExpr {
             } => {
                 let fields = input
                     .iter()
-                    .map(|node| arena.get(*node).to_field(schema, ctxt, arena))
+                    // default context because `col()` would return a list in aggregation context
+                    .map(|node| arena.get(*node).to_field(schema, Context::Default, arena))
                     .collect::<Result<Vec<_>>>()?;
                 Ok(output_type.get_field(schema, ctxt, &fields))
             }
@@ -343,7 +344,8 @@ impl AExpr {
             } => {
                 let fields = input
                     .iter()
-                    .map(|node| arena.get(*node).to_field(schema, ctxt, arena))
+                    // default context because `col()` would return a list in aggregation context
+                    .map(|node| arena.get(*node).to_field(schema, Context::Default, arena))
                     .collect::<Result<Vec<_>>>()?;
                 function.get_field(schema, ctxt, &fields)
             }

diff --git a/py-polars/tests/test_struct.py b/py-polars/tests/test_struct.py
@@ -607,3 +607,10 @@ def test_nested_struct_sliced_append() -> None:
         {"_experience": {"aaid": {"id": "B", "namespace": {"code": "bravo"}}}},
         {"_experience": {"aaid": {"id": "D", "namespace": {"code": "delta"}}}},
     ]
+
+
+def test_struct_groupby_field_agg_4216() -> None:
+    df = pl.DataFrame([{"a": {"b": 1}, "c": 0}])
+    assert df.groupby("c").agg(pl.col("a").struct.field("b").count()).to_dict(
+        False
+    ) == {"c": [0], "b": [1]}