fix: Expression pl.concat was incorrectly marked as elementwise (#22019)

Author: orlp

Diff for: crates/polars-plan/src/dsl/functions/concat.rs

@@ -97,7 +97,7 @@ pub fn concat_expr<E: AsRef<[IE]>, IE: Into<Expr> + Clone>(
         input: s,
         function: FunctionExpr::ConcatExpr(rechunk),
         options: FunctionOptions {
-            collect_groups: ApplyOptions::ElementWise,
+            collect_groups: ApplyOptions::GroupWise,
             flags: FunctionFlags::default() | FunctionFlags::INPUT_WILDCARD_EXPANSION,
             cast_options: Some(CastingRules::cast_to_supertypes()),
             ..Default::default()

Diff for: crates/polars-stream/src/physical_plan/lower_expr.rs

@@ -539,6 +539,35 @@ fn lower_exprs_with_ctx(
                 transformed_exprs.push(ctx.expr_arena.add(AExpr::Column(out_name)));
             },
 
+            AExpr::Function {
+                input: ref inner_exprs,
+                function: FunctionExpr::ConcatExpr(_rechunk),
+                options: _,
+            } => {
+                // We have to lower each expression separately as they might have different lengths.
+                let mut concat_streams = Vec::new();
+                let out_name = unique_column_name();
+                for inner_expr in inner_exprs {
+                    let (trans_input, trans_expr) =
+                        lower_exprs_with_ctx(input, &[inner_expr.node()], ctx)?;
+                    let select_expr =
+                        ExprIR::new(trans_expr[0], OutputName::Alias(out_name.clone()));
+                    concat_streams.push(build_select_stream_with_ctx(
+                        trans_input,
+                        &[select_expr],
+                        ctx,
+                    )?);
+                }
+
+                let output_schema = ctx.phys_sm[concat_streams[0].node].output_schema.clone();
+                let node_kind = PhysNodeKind::OrderedUnion {
+                    inputs: concat_streams,
+                };
+                let node_key = ctx.phys_sm.insert(PhysNode::new(output_schema, node_kind));
+                input_streams.insert(PhysStream::first(node_key));
+                transformed_exprs.push(ctx.expr_arena.add(AExpr::Column(out_name)));
+            },
+
             AExpr::Function {
                 input: ref inner_exprs,
                 function: FunctionExpr::Boolean(BooleanFunction::IsIn { nulls_equal }),

Diff for: py-polars/tests/unit/functions/test_concat.py

@@ -1,6 +1,7 @@
 import pytest
 
 import polars as pl
+from polars.testing import assert_frame_equal
 
 
 @pytest.mark.slow
@@ -59,3 +60,37 @@ def test_concat_vertically_relaxed() -> None:
         "a": [1.0, 0.2, 1.0, 2.0],
         "b": [None, 0.1, 2.0, 1.0],
     }
+
+
+def test_concat_group_by() -> None:
+    df = pl.DataFrame(
+        {
+            "g": [0, 0, 0, 0, 1, 1, 1, 1],
+            "a": [0, 1, 2, 3, 4, 5, 6, 7],
+            "b": [8, 9, 10, 11, 12, 13, 14, 15],
+        }
+    )
+    out = df.group_by("g").agg(pl.concat([pl.col.a, pl.col.b]))
+
+    assert_frame_equal(
+        out,
+        pl.DataFrame(
+            {
+                "g": [0, 1],
+                "a": [[0, 1, 2, 3, 8, 9, 10, 11], [4, 5, 6, 7, 12, 13, 14, 15]],
+            }
+        ),
+        check_row_order=False,
+    )
+
+
+def test_concat_19877() -> None:
+    df = pl.DataFrame({"a": [1, 2], "b": [3, 4]})
+    out = df.select(pl.concat([pl.col("a"), pl.col("b")]))
+    assert_frame_equal(out, pl.DataFrame({"a": [1, 2, 3, 4]}))
+
+
+def test_concat_zip_series_21980() -> None:
+    df = pl.DataFrame({"x": 1, "y": 2})
+    out = df.select(pl.concat([pl.col.x, pl.col.y]), pl.Series([3, 4]))
+    assert_frame_equal(out, pl.DataFrame({"x": [1, 2], "": [3, 4]}))