feat: Add strict parameter to pl.concat(how='horizontal')

py-polars/polars/functions/eager.py

@@ -29,6 +29,7 @@ def concat(
     how: ConcatMethod = "vertical",
     rechunk: bool = False,
     parallel: bool = True,
+    strict: bool = False,
 ) -> PolarsType:
     """
     Combine multiple DataFrames, LazyFrames, or Series into a single object.
@@ -58,6 +59,8 @@ def concat(
     parallel
         Only relevant for LazyFrames. This determines if the concatenated
         lazy computations may be executed in parallel.
+    strict
+        If True, reject concatenating DataFrames that are not the same height
 
     Examples
     --------
@@ -205,6 +208,12 @@ def concat(
                 )
             ).collect(no_optimization=True)
         elif how == "horizontal":
+            if strict:
+                for e in elems[1:]:
+                    if first.shape[0] != e.shape[0]:
+                        msg = f"Number of rows need to be equal ({first.shape[0]} != {e.shape[0]}) when 'strict' is True"
+                        raise ValueError(msg)
+
             out = wrap_df(plr.concat_df_horizontal(elems))
         else:
             allowed = ", ".join(repr(m) for m in get_args(ConcatMethod))
@@ -231,6 +240,14 @@ def concat(
                 )
             )
         elif how == "horizontal":
+            if strict:
+                nrows = first.select(F.len()).collect()[0, 0]
+                for e in elems[1:]:
+                    nrows2 = e.select(F.len()).collect()[0, 0]
+                    if nrows != nrows2:
+                        msg = f"Number of rows need to be equal ({nrows} != {nrows2}) when 'strict' is True"
+                        raise ValueError(msg)
+
             return wrap_ldf(
                 plr.concat_lf_horizontal(
                     elems,

py-polars/tests/unit/functions/test_concat.py

@@ -22,6 +22,25 @@ def test_concat_lf_stack_overflow() -> None:
     assert bar.collect().shape == (1001, 1)
 
 
+def test_concat_horizontally_strict() -> None:
+    a = pl.DataFrame({"a": [0, 1], "b": [1, 2]})
+    b = pl.DataFrame({"c": [11], "d": [42]})
+
+    with pytest.raises(ValueError):
+        pl.concat([a, b], how="horizontal", strict=True)
+
+    with pytest.raises(ValueError):
+        pl.concat([a.lazy(), b.lazy()], how="horizontal", strict=True)
+
+    out = pl.concat([a, b], how="horizontal", strict=False)
+    assert out.to_dict(as_series=False) == {
+        "a": [0, 1],
+        "b": [1, 2],
+        "c": [11, None],
+        "d": [42, None],
+    }
+
+
 def test_concat_vertically_relaxed() -> None:
     a = pl.DataFrame(
         data={"a": [1, 2, 3], "b": [True, False, None]},