fix: work around PyArrow to sort extension columns

Author: lidavidm

adbc_drivers_validation/compare.py

@@ -247,6 +247,39 @@ def compare_schemas(
         compare_fields(expected_field, actual_field, ())
 
 
+def sort_oblivious(table: pyarrow.Table, sort_keys: typing.Any) -> pyarrow.Table:
+    # sort, ignorant of extension types
+    # we have to cast to the storage type and back
+    cols = []
+    fields = []
+    for field, col in zip(table.schema, table.columns):
+        if field.type.id == 31:
+            col = pyarrow.chunked_array(
+                (c.storage for c in col.chunks),
+                type=col.type.storage_type,
+            )
+            field = pyarrow.field(
+                field.name,
+                field.type.storage_type,
+                field.nullable,
+                field.metadata,
+            )
+        fields.append(field)
+        cols.append(col)
+    new_table = pyarrow.Table.from_arrays(cols, schema=pyarrow.schema(fields))
+    new_table = new_table.sort_by(sort_keys)
+
+    cols = []
+    for field, col in zip(table.schema, new_table.columns):
+        if field.type.id == 31:
+            col = pyarrow.chunked_array(
+                (field.type.wrap_array(c) for c in col.chunks),
+                type=field.type,
+            )
+        cols.append(col)
+    return pyarrow.Table.from_arrays(cols, schema=table.schema)
+
+
 def compare_tables(
     expected: pyarrow.Table,
     actual: pyarrow.Table,
@@ -258,8 +291,8 @@ def compare_tables(
     actual = make_nullable(actual)
 
     if sort_keys := meta.get("sort-keys", None):
-        expected = expected.sort_by(sort_keys)
-        actual = actual.sort_by(sort_keys)
+        expected = sort_oblivious(expected, sort_keys)
+        actual = sort_oblivious(actual, sort_keys)
 
     compare_schemas(expected.schema, actual.schema)
     if expected != actual:

tests/test_compare.py

@@ -128,3 +128,53 @@ def test_compare_schemas_nullability():
 )
 def test_scalar_to_py_smart(value: pyarrow.Scalar, expected: typing.Any) -> None:
     assert compare.scalar_to_py_smart(value) == expected
+
+
+@pytest.mark.parametrize(
+    "table, sort_keys, expected",
+    [
+        pytest.param(
+            pyarrow.Table.from_pylist(
+                [
+                    {"idx": 1, "value": 0},
+                    {"idx": 0, "value": 1},
+                    {"idx": 2, "value": 2},
+                ]
+            ),
+            [("idx", "descending")],
+            pyarrow.Table.from_pylist(
+                [
+                    {"idx": 2, "value": 2},
+                    {"idx": 1, "value": 0},
+                    {"idx": 0, "value": 1},
+                ]
+            ),
+            id="int",
+        ),
+        pytest.param(
+            pyarrow.Table.from_pydict(
+                {
+                    "idx": pyarrow.opaque(pyarrow.int64(), "test", "").wrap_array(
+                        pyarrow.array([1, 0, 2])
+                    ),
+                    "value": pyarrow.array([0, 1, 2]),
+                }
+            ),
+            [("idx", "descending")],
+            pyarrow.Table.from_pydict(
+                {
+                    "idx": pyarrow.opaque(pyarrow.int64(), "test", "").wrap_array(
+                        pyarrow.array([2, 1, 0])
+                    ),
+                    "value": pyarrow.array([2, 0, 1]),
+                }
+            ),
+            id="extension",
+        ),
+    ],
+)
+def test_sort_oblivious(
+    table: pyarrow.Table, sort_keys: list, expected: pyarrow.Table
+) -> None:
+    result = compare.sort_oblivious(table, sort_keys)
+    assert result.equals(expected)