From fd8665cf6bba382c544b79f8198988f8bb4cb15f Mon Sep 17 00:00:00 2001
From: etiennebacher <etienne.bacher@protonmail.com>
Date: Wed, 27 Nov 2024 15:01:44 +0100
Subject: [PATCH] init

---
 py-polars/polars/functions/eager.py           | 25 ++++++++++++++++---
 .../tests/unit/functions/test_functions.py    |  7 ++++++
 py-polars/tests/unit/test_errors.py           | 10 ++++++++
 3 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/py-polars/polars/functions/eager.py b/py-polars/polars/functions/eager.py
index e65f4aa949c0..5de66bab13f3 100644
--- a/py-polars/polars/functions/eager.py
+++ b/py-polars/polars/functions/eager.py
@@ -31,14 +31,14 @@ def concat(
     parallel: bool = True,
 ) -> PolarsType:
     """
-    Combine multiple DataFrames, LazyFrames, or Series into a single object.
+    Combine multiple DataFrames, LazyFrames, Series, or Expr into a single object.
 
     Parameters
     ----------
     items
-        DataFrames, LazyFrames, or Series to concatenate.
+        DataFrames, LazyFrames, Series, or Expr to concatenate.
     how : {'vertical', 'vertical_relaxed', 'diagonal', 'diagonal_relaxed', 'horizontal', 'align'}
-        Series only support the `vertical` strategy.
+        Series and Expr only support the `vertical` strategy.
 
         * vertical: Applies multiple `vstack` operations.
         * vertical_relaxed: Same as `vertical`, but additionally coerces columns to
@@ -127,6 +127,19 @@ def concat(
     │ 2   ┆ 4    ┆ 5    ┆ null │
     │ 3   ┆ null ┆ 6    ┆ 8    │
     └─────┴──────┴──────┴──────┘
+    >>> df = pl.DataFrame({"a": [1, 2], "b": [3, 4]})
+    >>> df.select(pl.concat([pl.col("a"), pl.col("b")]))
+    shape: (4, 1)
+    ┌─────┐
+    │ a   │
+    │ --- │
+    │ i64 │
+    ╞═════╡
+    │ 1   │
+    │ 2   │
+    │ 3   │
+    │ 4   │
+    └─────┘
     """  # noqa: W505
     # unpack/standardise (handles generator input)
     elems = list(items)
@@ -250,7 +263,11 @@ def concat(
             raise ValueError(msg)
 
     elif isinstance(first, pl.Expr):
-        return wrap_expr(plr.concat_expr([e._pyexpr for e in elems], rechunk))
+        if how == "vertical":
+            return wrap_expr(plr.concat_expr([e._pyexpr for e in elems], rechunk))
+        else:
+            msg = "Expr only supports 'vertical' concat strategy"
+            raise ValueError(msg)
     else:
         msg = f"did not expect type: {type(first).__name__!r} in `concat`"
         raise TypeError(msg)
diff --git a/py-polars/tests/unit/functions/test_functions.py b/py-polars/tests/unit/functions/test_functions.py
index 05bd11976fd8..f68bfa995d47 100644
--- a/py-polars/tests/unit/functions/test_functions.py
+++ b/py-polars/tests/unit/functions/test_functions.py
@@ -183,6 +183,13 @@ def test_concat_vertical() -> None:
     assert_frame_equal(result, expected)
 
 
+@pytest.mark.may_fail_auto_streaming
+def test_concat_expr() -> None:
+    dat = pl.DataFrame({"a": [1, 2], "b": [3, 4]})
+    out = dat.select(pl.concat([pl.col("a"), pl.col("b") + 1]))
+    assert out.to_dict(as_series=False) == {"a": [1, 2, 4, 5]}
+
+
 def test_extend_ints() -> None:
     a = pl.DataFrame({"a": [1 for _ in range(1)]}, schema={"a": pl.Int64})
     with pytest.raises(pl.exceptions.SchemaError):
diff --git a/py-polars/tests/unit/test_errors.py b/py-polars/tests/unit/test_errors.py
index be299fb23996..bbb12172fd7d 100644
--- a/py-polars/tests/unit/test_errors.py
+++ b/py-polars/tests/unit/test_errors.py
@@ -283,6 +283,16 @@ def test_series_concat_err(how: ConcatMethod) -> None:
         pl.concat([s, s], how=how)
 
 
+@pytest.mark.parametrize("how", ["horizontal", "diagonal"])
+def test_expr_concat_err(how: ConcatMethod) -> None:
+    e = pl.lit([1, 2, 3])
+    with pytest.raises(
+        ValueError,
+        match="Expr only supports 'vertical' concat strategy",
+    ):
+        pl.concat([e, e], how=how)
+
+
 def test_invalid_sort_by() -> None:
     df = pl.DataFrame(
         {