fix: Fix struct broadcasting comparisons (#19003)

pola-rs · Sep 29, 2024 · 57df7b0 · 57df7b0
1 parent 570e3a2
commit 57df7b0
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/crates/polars-core/src/chunked_array/comparison/mod.rs b/crates/polars-core/src/chunked_array/comparison/mod.rs
@@ -764,8 +764,10 @@ where
     F: Fn(&Series, &Series) -> BooleanChunked,
     R: Fn(BooleanChunked, BooleanChunked) -> BooleanChunked,
 {
-    if a.len() != b.len() || a.struct_fields().len() != b.struct_fields().len() {
-        // polars_ensure!(a.len() == 1 || b.len() == 1, ShapeMismatch: "length lhs: {}, length rhs: {}", a.len(), b.len());
+    let len_a = a.len();
+    let len_b = b.len();
+    let broadcasts = len_a == 1 || len_b == 1;
+    if (a.len() != b.len() && !broadcasts) || a.struct_fields().len() != b.struct_fields().len() {
         BooleanChunked::full(PlSmallStr::EMPTY, value, a.len())
     } else {
         let (a, b) = align_chunks_binary(a, b);

diff --git a/py-polars/tests/unit/operations/test_comparison.py b/py-polars/tests/unit/operations/test_comparison.py
@@ -394,3 +394,10 @@ def test_nested_binary_literal_super_type_12227() -> None:
 
     result = pl.select((pl.lit(0) + (pl.lit(0) == pl.lit(0)) * pl.lit(0.1)) + pl.lit(0))
     assert result.item() == 0.1
+
+
+def test_struct_broadcasting_comparison() -> None:
+    df = pl.DataFrame({"foo": [{"a": 1}, {"a": 2}, {"a": 1}]})
+    assert df.select(eq=pl.col.foo == pl.col.foo.last()).to_dict(as_series=False) == {
+        "eq": [True, False, True]
+    }