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

crates/polars-core/src/frame/horizontal.rs

@@ -76,7 +76,16 @@ impl DataFrame {
 }
 /// Concat [`DataFrame`]s horizontally.
 /// Concat horizontally and extend with null values if lengths don't match
-pub fn concat_df_horizontal(dfs: &[DataFrame], check_duplicates: bool) -> PolarsResult<DataFrame> {
+pub fn concat_df_horizontal<T>(
+    dfs: &[DataFrame],
+    check_duplicates: bool,
+    strict: T,
+) -> PolarsResult<DataFrame>
+where
+    T: Into<Option<bool>>,
+{
+    // Set this to true for v2.0.0 milestone
+    let strict = strict.into().unwrap_or(false);
     let output_height = dfs
         .iter()
         .map(|df| df.height())
@@ -87,6 +96,11 @@ pub fn concat_df_horizontal(dfs: &[DataFrame], check_duplicates: bool) -> Polars
 
     // if not all equal length, extend the DataFrame with nulls
     let dfs = if !dfs.iter().all(|df| df.height() == output_height) {
+        if strict {
+            return Err(
+                polars_err!(ShapeMismatch: "cannot concat dataframes with different heights in 'strict' mode"),
+            );
+        }
         owned_df = dfs
             .iter()
             .cloned()

crates/polars-lazy/src/dsl/functions.rs

@@ -74,6 +74,7 @@ pub fn concat_lf_horizontal<L: AsRef<[LazyFrame]>>(
 
     let options = HConcatOptions {
         parallel: args.parallel,
+        strict: args.strict,
     };
     let lp = DslPlan::HConcat {
         inputs: lfs.iter().map(|lf| lf.logical_plan.clone()).collect(),

crates/polars-mem-engine/src/executors/hconcat.rs

@@ -59,6 +59,6 @@ impl Executor for HConcatExec {
         };
 
         // Invariant of IR. Schema is already checked to contain no duplicates.
-        concat_df_horizontal(&dfs, false)
+        concat_df_horizontal(&dfs, false, self.options.strict)
     }
 }

crates/polars-plan/src/plans/options.rs

@@ -58,6 +58,7 @@ pub struct UnionOptions {
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 pub struct HConcatOptions {
     pub parallel: bool,
+    pub strict: bool,
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Default, Hash)]
@@ -363,6 +364,7 @@ pub struct UnionArgs {
     pub rechunk: bool,
     pub to_supertypes: bool,
     pub diagonal: bool,
+    pub strict: bool,
     // If it is a union from a scan over multiple files.
     pub from_partitioned_ds: bool,
 }
@@ -375,6 +377,8 @@ impl Default for UnionArgs {
             to_supertypes: false,
             diagonal: false,
             from_partitioned_ds: false,
+            // By default, strict should be true in v2.0.0
+            strict: false,
         }
     }
 }

crates/polars-python/src/functions/eager.rs

@@ -78,7 +78,7 @@ pub fn concat_df_diagonal(dfs: &Bound<'_, PyAny>) -> PyResult<PyDataFrame> {
 }
 
 #[pyfunction]
-pub fn concat_df_horizontal(dfs: &Bound<'_, PyAny>) -> PyResult<PyDataFrame> {
+pub fn concat_df_horizontal(dfs: &Bound<'_, PyAny>, strict: bool) -> PyResult<PyDataFrame> {
     let iter = dfs.iter()?;
 
     let dfs = iter
@@ -88,6 +88,6 @@ pub fn concat_df_horizontal(dfs: &Bound<'_, PyAny>) -> PyResult<PyDataFrame> {
         })
         .collect::<PyResult<Vec<_>>>()?;
 
-    let df = functions::concat_df_horizontal(&dfs, true).map_err(PyPolarsErr::from)?;
+    let df = functions::concat_df_horizontal(&dfs, true, strict).map_err(PyPolarsErr::from)?;
     Ok(df.into())
 }

crates/polars-python/src/functions/lazy.rs

@@ -325,7 +325,11 @@ pub fn concat_lf_diagonal(
 }
 
 #[pyfunction]
-pub fn concat_lf_horizontal(lfs: &Bound<'_, PyAny>, parallel: bool) -> PyResult<PyLazyFrame> {
+pub fn concat_lf_horizontal(
+    lfs: &Bound<'_, PyAny>,
+    parallel: bool,
+    strict: bool,
+) -> PyResult<PyLazyFrame> {
     let iter = lfs.iter()?;
 
     let lfs = iter
@@ -339,6 +343,7 @@ pub fn concat_lf_horizontal(lfs: &Bound<'_, PyAny>, parallel: bool) -> PyResult<
         rechunk: false, // No need to rechunk with horizontal concatenation
         parallel,
         to_supertypes: false,
+        strict,
         ..Default::default()
     };
     let lf = dsl::functions::concat_lf_horizontal(lfs, args).map_err(PyPolarsErr::from)?;

crates/polars-stream/src/nodes/zip.rs

@@ -275,7 +275,7 @@ impl ComputeNode for ZipNode {
                 for input_head in &mut self.input_heads {
                     out.push(input_head.take(common_size));
                 }
-                let out_df = concat_df_horizontal(&out, false)?;
+                let out_df = concat_df_horizontal(&out, false, None)?;
                 out.clear();
 
                 let morsel = Morsel::new(out_df, self.out_seq, source_token.clone());
@@ -320,7 +320,7 @@ impl ComputeNode for ZipNode {
                 for input_head in &mut self.input_heads {
                     out.push(input_head.consume_broadcast());
                 }
-                let out_df = concat_df_horizontal(&out, false)?;
+                let out_df = concat_df_horizontal(&out, false, None)?;
                 out.clear();
 
                 let morsel = Morsel::new(out_df, self.out_seq, source_token.clone());

docs/source/src/rust/user-guide/transformations/concatenation.rs

@@ -30,7 +30,8 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             "r2"=> &[7, 8],
             "r3"=> &[9, 10],
     )?;
-    let df_horizontal_concat = polars::functions::concat_df_horizontal(&[df_h1, df_h2], true)?;
+    let df_horizontal_concat =
+        polars::functions::concat_df_horizontal(&[df_h1, df_h2], true, None)?;
     println!("{}", &df_horizontal_concat);
     // --8<-- [end:horizontal]
     //
@@ -43,7 +44,8 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             "r1"=> &[5, 6, 7],
             "r2"=> &[8, 9, 10],
     )?;
-    let df_horizontal_concat = polars::functions::concat_df_horizontal(&[df_h1, df_h2], true)?;
+    let df_horizontal_concat =
+        polars::functions::concat_df_horizontal(&[df_h1, df_h2], true, None)?;
     println!("{}", &df_horizontal_concat);
     // --8<-- [end:horizontal_different_lengths]
 

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 when how=`horizontal`
 
     Examples
     --------
@@ -205,7 +208,7 @@ def concat(
                 )
             ).collect(no_optimization=True)
         elif how == "horizontal":
-            out = wrap_df(plr.concat_df_horizontal(elems))
+            out = wrap_df(plr.concat_df_horizontal(elems, strict=strict))
         else:
             allowed = ", ".join(repr(m) for m in get_args(ConcatMethod))
             msg = f"DataFrame `how` must be one of {{{allowed}}}, got {how!r}"
@@ -235,6 +238,7 @@ def concat(
                 plr.concat_lf_horizontal(
                     elems,
                     parallel=parallel,
+                    strict=strict,
                 )
             )
         else:

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(pl.exceptions.ShapeError):
+        pl.concat([a, b], how="horizontal", strict=True)
+
+    with pytest.raises(pl.exceptions.ShapeError):
+        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]},