feat: Perform tolerance-based comparison for lists and arrays

Author: MariusMerkleQC

diffly/_conditions.py

@@ -22,12 +22,14 @@ def condition_equal_rows(
     abs_tol_by_column: Mapping[str, float],
     rel_tol_by_column: Mapping[str, float],
     abs_tol_temporal_by_column: Mapping[str, dt.timedelta],
+    max_list_lengths_by_column: Mapping[str, int] | None = None,
 ) -> pl.Expr:
     """Build an expression whether two rows are equal, based on all columns' data
     types."""
     if not columns:
         return pl.lit(True)
 
+    _max_list_lengths = max_list_lengths_by_column or {}
     return pl.all_horizontal(
         [
             condition_equal_columns(
@@ -37,6 +39,7 @@ def condition_equal_rows(
                 abs_tol=abs_tol_by_column[column],
                 rel_tol=rel_tol_by_column[column],
                 abs_tol_temporal=abs_tol_temporal_by_column[column],
+                max_list_length=_max_list_lengths.get(column, 0),
             )
             for column in columns
         ]
@@ -50,6 +53,7 @@ def condition_equal_columns(
     abs_tol: float = ABS_TOL_DEFAULT,
     rel_tol: float = REL_TOL_DEFAULT,
     abs_tol_temporal: dt.timedelta = ABS_TOL_TEMPORAL_DEFAULT,
+    max_list_length: int = 0,
 ) -> pl.Expr:
     """Build an expression whether two columns are equal, depending on the columns' data
     types."""
@@ -61,6 +65,7 @@ def condition_equal_columns(
         abs_tol=abs_tol,
         rel_tol=rel_tol,
         abs_tol_temporal=abs_tol_temporal,
+        max_list_length=max_list_length,
     )
 
 
@@ -95,6 +100,7 @@ def _compare_columns(
     abs_tol: float,
     rel_tol: float,
     abs_tol_temporal: dt.timedelta,
+    max_list_length: int = 0,
 ) -> pl.Expr:
     """Build an expression whether two expressions yield the same value.
 
@@ -133,10 +139,18 @@ def _compare_columns(
         elif isinstance(dtype_left, pl.List | pl.Array) and isinstance(
             dtype_right, pl.List | pl.Array
         ):
-            # As of polars 1.28, there is no way to access another column within
-            # `list.eval`. Hence, we necessarily need to resort to a primitive
-            # comparison in this case.
-            pass
+            result = _compare_sequence_columns(
+                col_left=col_left,
+                col_right=col_right,
+                dtype_left=dtype_left,
+                dtype_right=dtype_right,
+                max_list_length=max_list_length,
+                abs_tol=abs_tol,
+                rel_tol=rel_tol,
+                abs_tol_temporal=abs_tol_temporal,
+            )
+            if result is not None:
+                return result
 
     if (
         isinstance(dtype_left, pl.Enum)
@@ -167,6 +181,77 @@ def _compare_columns(
     )
 
 
+def _compare_sequence_columns(
+    col_left: pl.Expr,
+    col_right: pl.Expr,
+    dtype_left: DataType | DataTypeClass,
+    dtype_right: DataType | DataTypeClass,
+    max_list_length: int,
+    abs_tol: float,
+    rel_tol: float,
+    abs_tol_temporal: dt.timedelta,
+) -> pl.Expr | None:
+    """Compare Array/List columns element-wise with tolerance.
+
+    Returns ``None`` if the comparison cannot be performed element-wise (e.g. List vs
+    List without a known ``max_list_length``), signalling to the caller that it should
+    fall back to primitive comparison.
+    """
+    assert isinstance(dtype_left, pl.List | pl.Array)
+    assert isinstance(dtype_right, pl.List | pl.Array)
+    inner_left = dtype_left.inner
+    inner_right = dtype_right.inner
+
+    def _get_element(col: pl.Expr, dtype: DataType | DataTypeClass, i: int) -> pl.Expr:
+        if isinstance(dtype, pl.Array):
+            return col.arr.get(i)
+        return col.list.get(i, null_on_oob=True)
+
+    n: int | None = None
+    length_check: pl.Expr | None = None
+
+    if isinstance(dtype_left, pl.Array) and isinstance(dtype_right, pl.Array):
+        if dtype_left.shape != dtype_right.shape:
+            return pl.repeat(pl.lit(False), pl.len())
+        n = dtype_left.shape[0]
+    elif isinstance(dtype_left, pl.Array) and isinstance(dtype_right, pl.List):
+        n = dtype_left.shape[0]
+        length_check = col_right.list.len().eq(pl.lit(n))
+    elif isinstance(dtype_left, pl.List) and isinstance(dtype_right, pl.Array):
+        n = dtype_right.shape[0]
+        length_check = col_left.list.len().eq(pl.lit(n))
+    else:
+        # List vs List
+        if max_list_length == 0:
+            return None
+        n = max_list_length
+        length_check = col_left.list.len().eq_missing(col_right.list.len())
+
+    if n == 0:
+        if length_check is not None:
+            return _eq_missing(length_check, col_left, col_right)
+        return _eq_missing(pl.lit(True), col_left, col_right)
+
+    elements_match = pl.all_horizontal(
+        [
+            _compare_columns(
+                col_left=_get_element(col_left, dtype_left, i),
+                col_right=_get_element(col_right, dtype_right, i),
+                dtype_left=inner_left,
+                dtype_right=inner_right,
+                abs_tol=abs_tol,
+                rel_tol=rel_tol,
+                abs_tol_temporal=abs_tol_temporal,
+            )
+            for i in range(n)
+        ]
+    )
+
+    if length_check is not None:
+        return _eq_missing(length_check & elements_match, col_left, col_right)
+    return elements_match
+
+
 def _compare_primitive_columns(
     col_left: pl.Expr,
     col_right: pl.Expr,

diffly/comparison.py

@@ -511,6 +511,7 @@ def equal(self, *, check_dtypes: bool = True) -> bool:
                     abs_tol_by_column=self.abs_tol_by_column,
                     rel_tol_by_column=self.rel_tol_by_column,
                     abs_tol_temporal_by_column=self.abs_tol_temporal_by_column,
+                    max_list_lengths_by_column=self._max_list_lengths,
                 ).all()
             )
             .item()
@@ -708,6 +709,26 @@ def _validate_subset_of_common_columns(self, subset: Iterable[str]) -> list[str]
             raise ValueError(f"{difference} are not common columns.")
         return list(subset)
 
+    @cached_property
+    def _max_list_lengths(self) -> dict[str, int]:
+        list_columns = [
+            col
+            for col in self._other_common_columns
+            if isinstance(self.left_schema[col], pl.List)
+            and isinstance(self.right_schema[col], pl.List)
+        ]
+        if not list_columns:
+            return {}
+
+        exprs = [pl.col(col).list.len().max().alias(col) for col in list_columns]
+        [left_max, right_max] = pl.collect_all(
+            [self.left.select(exprs), self.right.select(exprs)]
+        )
+        return {
+            col: max(int(left_max[col].item() or 0), int(right_max[col].item() or 0))
+            for col in list_columns
+        }
+
     def _condition_equal_rows(self, columns: list[str]) -> pl.Expr:
         return condition_equal_rows(
             columns=columns,
@@ -716,6 +737,7 @@ def _condition_equal_rows(self, columns: list[str]) -> pl.Expr:
             abs_tol_by_column=self.abs_tol_by_column,
             rel_tol_by_column=self.rel_tol_by_column,
             abs_tol_temporal_by_column=self.abs_tol_temporal_by_column,
+            max_list_lengths_by_column=self._max_list_lengths,
         )
 
     def _condition_equal_columns(self, column: str) -> pl.Expr:
@@ -726,6 +748,7 @@ def _condition_equal_columns(self, column: str) -> pl.Expr:
             abs_tol=self.abs_tol_by_column[column],
             rel_tol=self.rel_tol_by_column[column],
             abs_tol_temporal=self.abs_tol_temporal_by_column[column],
+            max_list_length=self._max_list_lengths.get(column, 0),
         )
 
     def _equal_rows(self) -> bool:

multi_array.py

@@ -0,0 +1,6 @@
+# %%
+import polars as pl
+# %%
+df = pl.DataFrame({"a": [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]}, schema={"a": pl.Array(inner=pl.UInt8, shape=(2, 2))})
+# %%
+df.select(pl.col("a").arr.get(1))

tests/test_conditions.py

@@ -81,7 +81,7 @@ def test_condition_equal_columns_different_struct_fields() -> None:
 @pytest.mark.parametrize(
     "rhs_type", [pl.Array(pl.Float64, shape=2), pl.List(pl.Float64)]
 )
-def test_condition_equal_columns_list_array_equal_exact(
+def test_condition_equal_columns_list_array_with_tolerance(
     lhs_type: pl.DataType, rhs_type: pl.DataType
 ) -> None:
     # Arrange
@@ -110,12 +110,58 @@ def test_condition_equal_columns_list_array_equal_exact(
                 dtype_right=rhs.schema["a_right"],
                 abs_tol=0.5,
                 rel_tol=0,
+                max_list_length=2,
             )
         )
         .to_series()
     )
 
-    # Assert
+    # Assert: diff is 0.1, within abs_tol=0.5
+    assert actual.to_list() == [True, True]
+
+
+@pytest.mark.parametrize(
+    "lhs_type", [pl.Array(pl.Float64, shape=2), pl.List(pl.Float64)]
+)
+@pytest.mark.parametrize(
+    "rhs_type", [pl.Array(pl.Float64, shape=2), pl.List(pl.Float64)]
+)
+def test_condition_equal_columns_list_array_exceeds_tolerance(
+    lhs_type: pl.DataType, rhs_type: pl.DataType
+) -> None:
+    # Arrange
+    lhs = pl.DataFrame(
+        {
+            "pk": [1, 2],
+            "a_left": [[1.0, 1.1], [2.0, 2.1]],
+        },
+        schema={"pk": pl.Int64, "a_left": lhs_type},
+    )
+    rhs = pl.DataFrame(
+        {
+            "pk": [1, 2],
+            "a_right": [[1.0, 1.1], [2.0, 2.8]],
+        },
+        schema={"pk": pl.Int64, "a_right": rhs_type},
+    )
+
+    # Act
+    actual = (
+        lhs.join(rhs, on="pk", maintain_order="left")
+        .select(
+            condition_equal_columns(
+                "a",
+                dtype_left=lhs.schema["a_left"],
+                dtype_right=rhs.schema["a_right"],
+                abs_tol=0.5,
+                rel_tol=0,
+                max_list_length=2,
+            )
+        )
+        .to_series()
+    )
+
+    # Assert: diff is 0.7, exceeds abs_tol=0.5
     assert actual.to_list() == [True, False]
 
 
@@ -226,6 +272,97 @@ def test_condition_equal_columns_temporal_tolerance() -> None:
     assert actual.to_list() == [True, False, False, True]
 
 
+def test_condition_equal_columns_list_different_lengths() -> None:
+    lhs = pl.DataFrame(
+        {
+            "pk": [1, 2],
+            "a_left": [[1.0, 2.0], [3.0]],
+        },
+    )
+    rhs = pl.DataFrame(
+        {
+            "pk": [1, 2],
+            "a_right": [[1.0, 2.0], [3.0, 4.0]],
+        },
+    )
+
+    actual = (
+        lhs.join(rhs, on="pk", maintain_order="left")
+        .select(
+            condition_equal_columns(
+                "a",
+                dtype_left=lhs.schema["a_left"],
+                dtype_right=rhs.schema["a_right"],
+                abs_tol=0.5,
+                rel_tol=0,
+                max_list_length=2,
+            )
+        )
+        .to_series()
+    )
+    assert actual.to_list() == [True, False]
+
+
+def test_condition_equal_columns_list_nulls() -> None:
+    lhs = pl.DataFrame(
+        {
+            "pk": [1, 2, 3],
+            "a_left": [[1.0, 2.0], None, None],
+        },
+    )
+    rhs = pl.DataFrame(
+        {
+            "pk": [1, 2, 3],
+            "a_right": [[1.0, 2.0], [3.0], None],
+        },
+    )
+
+    actual = (
+        lhs.join(rhs, on="pk", maintain_order="left")
+        .select(
+            condition_equal_columns(
+                "a",
+                dtype_left=lhs.schema["a_left"],
+                dtype_right=rhs.schema["a_right"],
+                max_list_length=2,
+            )
+        )
+        .to_series()
+    )
+    assert actual.to_list() == [True, False, True]
+
+
+def test_condition_equal_columns_array_vs_list_length_mismatch() -> None:
+    lhs = pl.DataFrame(
+        {
+            "pk": [1, 2],
+            "a_left": [[1.0, 2.0], [3.0, 4.0]],
+        },
+        schema={"pk": pl.Int64, "a_left": pl.Array(pl.Float64, shape=2)},
+    )
+    rhs = pl.DataFrame(
+        {
+            "pk": [1, 2],
+            "a_right": [[1.0, 2.0], [3.0]],
+        },
+    )
+
+    actual = (
+        lhs.join(rhs, on="pk", maintain_order="left")
+        .select(
+            condition_equal_columns(
+                "a",
+                dtype_left=lhs.schema["a_left"],
+                dtype_right=rhs.schema["a_right"],
+                abs_tol=0.5,
+                rel_tol=0,
+            )
+        )
+        .to_series()
+    )
+    assert actual.to_list() == [True, False]
+
+
 @pytest.mark.parametrize(
     ("dtype_left", "dtype_right", "can_compare_dtypes"),
     [