Merge pull request #1387 from codeflash-ai/pyarrow-comparator

feat: add PyArrow support to comparator

Author: KRRT7

codeflash/verification/comparator.py

@@ -26,6 +26,7 @@
 HAS_XARRAY = find_spec("xarray") is not None
 HAS_TENSORFLOW = find_spec("tensorflow") is not None
 HAS_NUMBA = find_spec("numba") is not None
+HAS_PYARROW = find_spec("pyarrow") is not None
 
 # Pattern to match pytest temp directories: /tmp/pytest-of-<user>/pytest-<N>/
 # These paths vary between test runs but are logically equivalent
@@ -354,13 +355,57 @@ def comparator(orig: Any, new: Any, superset_obj: bool = False) -> bool:
                 return False
             return (orig != new).nnz == 0
 
+        if HAS_PYARROW:
+            import pyarrow as pa  # type: ignore  # noqa: PGH003
+
+            if isinstance(orig, pa.Table):
+                if orig.schema != new.schema:
+                    return False
+                if orig.num_rows != new.num_rows:
+                    return False
+                return bool(orig.equals(new))
+
+            if isinstance(orig, pa.RecordBatch):
+                if orig.schema != new.schema:
+                    return False
+                if orig.num_rows != new.num_rows:
+                    return False
+                return bool(orig.equals(new))
+
+            if isinstance(orig, pa.ChunkedArray):
+                if orig.type != new.type:
+                    return False
+                if len(orig) != len(new):
+                    return False
+                return bool(orig.equals(new))
+
+            if isinstance(orig, pa.Array):
+                if orig.type != new.type:
+                    return False
+                if len(orig) != len(new):
+                    return False
+                return bool(orig.equals(new))
+
+            if isinstance(orig, pa.Scalar):
+                if orig.type != new.type:
+                    return False
+                # Handle null scalars
+                if not orig.is_valid and not new.is_valid:
+                    return True
+                if not orig.is_valid or not new.is_valid:
+                    return False
+                return bool(orig.equals(new))
+
+            if isinstance(orig, (pa.Schema, pa.Field, pa.DataType)):
+                return bool(orig.equals(new))
+
         if HAS_PANDAS:
             import pandas  # noqa: ICN001
 
             if isinstance(
                 orig, (pandas.DataFrame, pandas.Series, pandas.Index, pandas.Categorical, pandas.arrays.SparseArray)
             ):
-                return orig.equals(new)
+                return bool(orig.equals(new))
 
             if isinstance(orig, (pandas.CategoricalDtype, pandas.Interval, pandas.Period)):
                 return orig == new
@@ -407,10 +452,10 @@ def comparator(orig: Any, new: Any, superset_obj: bool = False) -> bool:
                 return orig == new
 
         if HAS_NUMBA:
-            import numba  # type: ignore  # noqa: PGH003
-            from numba.core.dispatcher import Dispatcher  # type: ignore  # noqa: PGH003
-            from numba.typed import Dict as NumbaDict  # type: ignore  # noqa: PGH003
-            from numba.typed import List as NumbaList  # type: ignore  # noqa: PGH003
+            import numba
+            from numba.core.dispatcher import Dispatcher
+            from numba.typed import Dict as NumbaDict
+            from numba.typed import List as NumbaList
 
             # Handle numba typed List
             if isinstance(orig, NumbaList):

pyproject.toml

@@ -87,6 +87,7 @@ tests = [
     "jax>=0.4.30",
     "numpy>=2.0.2",
     "pandas>=2.3.3",
+    "pyarrow>=15.0.0",
     "pyrsistent>=0.20.0",
     "scipy>=1.13.1",
     "torch>=2.8.0",

tests/test_comparator.py

@@ -802,6 +802,137 @@ def test_pandas():
     assert comparator(filtered1, filtered2)
 
 
+def test_pyarrow():
+    try:
+        import pyarrow as pa
+    except ImportError:
+        pytest.skip()
+
+    # Test PyArrow Table
+    table1 = pa.table({"a": [1, 2, 3], "b": [4, 5, 6]})
+    table2 = pa.table({"a": [1, 2, 3], "b": [4, 5, 6]})
+    table3 = pa.table({"a": [1, 2, 3], "b": [4, 5, 7]})
+    table4 = pa.table({"a": [1, 2, 3, 4], "b": [4, 5, 6, 7]})
+    table5 = pa.table({"a": [1, 2, 3], "c": [4, 5, 6]})  # different column name
+
+    assert comparator(table1, table2)
+    assert not comparator(table1, table3)
+    assert not comparator(table1, table4)
+    assert not comparator(table1, table5)
+
+    # Test PyArrow RecordBatch
+    batch1 = pa.RecordBatch.from_pydict({"x": [1, 2], "y": [3.0, 4.0]})
+    batch2 = pa.RecordBatch.from_pydict({"x": [1, 2], "y": [3.0, 4.0]})
+    batch3 = pa.RecordBatch.from_pydict({"x": [1, 2], "y": [3.0, 5.0]})
+    batch4 = pa.RecordBatch.from_pydict({"x": [1, 2, 3], "y": [3.0, 4.0, 5.0]})
+
+    assert comparator(batch1, batch2)
+    assert not comparator(batch1, batch3)
+    assert not comparator(batch1, batch4)
+
+    # Test PyArrow Array
+    arr1 = pa.array([1, 2, 3])
+    arr2 = pa.array([1, 2, 3])
+    arr3 = pa.array([1, 2, 4])
+    arr4 = pa.array([1, 2, 3, 4])
+    arr5 = pa.array([1.0, 2.0, 3.0])  # different type
+
+    assert comparator(arr1, arr2)
+    assert not comparator(arr1, arr3)
+    assert not comparator(arr1, arr4)
+    assert not comparator(arr1, arr5)
+
+    # Test PyArrow Array with nulls
+    arr_null1 = pa.array([1, None, 3])
+    arr_null2 = pa.array([1, None, 3])
+    arr_null3 = pa.array([1, 2, 3])
+
+    assert comparator(arr_null1, arr_null2)
+    assert not comparator(arr_null1, arr_null3)
+
+    # Test PyArrow ChunkedArray
+    chunked1 = pa.chunked_array([[1, 2], [3, 4]])
+    chunked2 = pa.chunked_array([[1, 2], [3, 4]])
+    chunked3 = pa.chunked_array([[1, 2], [3, 5]])
+    chunked4 = pa.chunked_array([[1, 2, 3], [4, 5]])
+
+    assert comparator(chunked1, chunked2)
+    assert not comparator(chunked1, chunked3)
+    assert not comparator(chunked1, chunked4)
+
+    # Test PyArrow Scalar
+    scalar1 = pa.scalar(42)
+    scalar2 = pa.scalar(42)
+    scalar3 = pa.scalar(43)
+    scalar4 = pa.scalar(42.0)  # different type
+
+    assert comparator(scalar1, scalar2)
+    assert not comparator(scalar1, scalar3)
+    assert not comparator(scalar1, scalar4)
+
+    # Test null scalars
+    null_scalar1 = pa.scalar(None, type=pa.int64())
+    null_scalar2 = pa.scalar(None, type=pa.int64())
+    null_scalar3 = pa.scalar(None, type=pa.float64())
+
+    assert comparator(null_scalar1, null_scalar2)
+    assert not comparator(null_scalar1, null_scalar3)
+
+    # Test PyArrow Schema
+    schema1 = pa.schema([("a", pa.int64()), ("b", pa.float64())])
+    schema2 = pa.schema([("a", pa.int64()), ("b", pa.float64())])
+    schema3 = pa.schema([("a", pa.int64()), ("c", pa.float64())])
+    schema4 = pa.schema([("a", pa.int32()), ("b", pa.float64())])
+
+    assert comparator(schema1, schema2)
+    assert not comparator(schema1, schema3)
+    assert not comparator(schema1, schema4)
+
+    # Test PyArrow Field
+    field1 = pa.field("name", pa.int64())
+    field2 = pa.field("name", pa.int64())
+    field3 = pa.field("other", pa.int64())
+    field4 = pa.field("name", pa.float64())
+
+    assert comparator(field1, field2)
+    assert not comparator(field1, field3)
+    assert not comparator(field1, field4)
+
+    # Test PyArrow DataType
+    type1 = pa.int64()
+    type2 = pa.int64()
+    type3 = pa.int32()
+    type4 = pa.float64()
+
+    assert comparator(type1, type2)
+    assert not comparator(type1, type3)
+    assert not comparator(type1, type4)
+
+    # Test string arrays
+    str_arr1 = pa.array(["hello", "world"])
+    str_arr2 = pa.array(["hello", "world"])
+    str_arr3 = pa.array(["hello", "there"])
+
+    assert comparator(str_arr1, str_arr2)
+    assert not comparator(str_arr1, str_arr3)
+
+    # Test nested types (struct)
+    struct_arr1 = pa.array([{"x": 1, "y": 2}, {"x": 3, "y": 4}])
+    struct_arr2 = pa.array([{"x": 1, "y": 2}, {"x": 3, "y": 4}])
+    struct_arr3 = pa.array([{"x": 1, "y": 2}, {"x": 3, "y": 5}])
+
+    assert comparator(struct_arr1, struct_arr2)
+    assert not comparator(struct_arr1, struct_arr3)
+
+    # Test list arrays
+    list_arr1 = pa.array([[1, 2], [3, 4, 5]])
+    list_arr2 = pa.array([[1, 2], [3, 4, 5]])
+    list_arr3 = pa.array([[1, 2], [3, 4, 6]])
+
+    assert comparator(list_arr1, list_arr2)
+    assert not comparator(list_arr1, list_arr3)
+
+
 def test_pyrsistent():
     try:
         from pyrsistent import PBag, PClass, PRecord, field, pdeque, pmap, pset, pvector  # type: ignore
@@ -2795,7 +2926,6 @@ def test_torch_runtime_error_wrapping():
     class TorchRuntimeError(Exception):
         """Mock TorchRuntimeError for testing."""
 
-
     # Monkey-patch the __module__ to match torch._dynamo.exc
     TorchRuntimeError.__module__ = "torch._dynamo.exc"