test: add test for schema registration with pyarrow dataset

Author: kosiew

python/datafusion/utils.py

@@ -18,10 +18,26 @@
 
 from __future__ import annotations
 
+from importlib import import_module, util
 from typing import TYPE_CHECKING, Any
 
 from datafusion._internal import EXPECTED_PROVIDER_MSG
 
+_PYARROW_DATASET_TYPES: tuple[type[Any], ...]
+_dataset_spec = util.find_spec("pyarrow.dataset")
+if _dataset_spec is None:  # pragma: no cover - optional dependency at runtime
+    _PYARROW_DATASET_TYPES = ()
+else:  # pragma: no cover - exercised in environments with pyarrow installed
+    _dataset_module = import_module("pyarrow.dataset")
+    dataset_base = getattr(_dataset_module, "Dataset", None)
+    dataset_types: set[type[Any]] = set()
+    if isinstance(dataset_base, type):
+        dataset_types.add(dataset_base)
+        for value in vars(_dataset_module).values():
+            if isinstance(value, type) and issubclass(value, dataset_base):
+                dataset_types.add(value)
+    _PYARROW_DATASET_TYPES = tuple(dataset_types)
+
 if TYPE_CHECKING:  # pragma: no cover - imported for typing only
     from datafusion import TableProvider
     from datafusion.catalog import Table
@@ -54,6 +70,9 @@ def _normalize_table_provider(
     if isinstance(table, _TableProvider):
         return table._table_provider
 
+    if _PYARROW_DATASET_TYPES and isinstance(table, _PYARROW_DATASET_TYPES):
+        return table
+
     provider_factory = getattr(table, "__datafusion_table_provider__", None)
     if callable(provider_factory):
         return table

python/tests/test_catalog.py

@@ -164,6 +164,28 @@ def test_python_table_provider(ctx: SessionContext):
     assert schema.table_names() == {"table4"}
 
 
+def test_schema_register_table_with_pyarrow_dataset(ctx: SessionContext):
+    schema = ctx.catalog().schema()
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([1, 2, 3]), pa.array([4, 5, 6])],
+        names=["a", "b"],
+    )
+    dataset = ds.dataset([batch])
+    table_name = "pa_dataset"
+
+    try:
+        schema.register_table(table_name, dataset)
+        assert table_name in schema.table_names()
+
+        result = ctx.sql(f"SELECT a, b FROM {table_name}").collect()
+
+        assert len(result) == 1
+        assert result[0].column(0) == pa.array([1, 2, 3])
+        assert result[0].column(1) == pa.array([4, 5, 6])
+    finally:
+        schema.deregister_table(table_name)
+
+
 def test_schema_register_table_with_dataframe_errors(ctx: SessionContext):
     schema = ctx.catalog().schema()
     df = ctx.from_pydict({"a": [1]})