feat: add _normalize_table_provider utility for consistent table registration

Author: kosiew

python/datafusion/catalog.py

@@ -23,6 +23,7 @@
 from typing import TYPE_CHECKING, Protocol
 
 import datafusion._internal as df_internal
+from datafusion.utils import _normalize_table_provider
 
 if TYPE_CHECKING:
     import pyarrow as pa
@@ -137,9 +138,8 @@ def register_table(
         Objects implementing ``__datafusion_table_provider__`` are also supported
         and treated as :class:`TableProvider` instances.
         """
-        if isinstance(table, Table):
-            return self._raw_schema.register_table(name, table.table)
-        return self._raw_schema.register_table(name, table)
+        provider = _normalize_table_provider(table)
+        return self._raw_schema.register_table(name, provider)
 
     def deregister_table(self, name: str) -> None:
         """Deregister a table provider from this schema."""

python/datafusion/context.py

@@ -34,6 +34,7 @@
 from datafusion.expr import Expr, SortExpr, sort_list_to_raw_sort_list
 from datafusion.record_batch import RecordBatchStream
 from datafusion.user_defined import AggregateUDF, ScalarUDF, TableFunction, WindowUDF
+from datafusion.utils import _normalize_table_provider
 
 from ._internal import RuntimeEnvBuilder as RuntimeEnvBuilderInternal
 from ._internal import SessionConfig as SessionConfigInternal
@@ -766,10 +767,8 @@ def register_table(
                 implementing ``__datafusion_table_provider__`` to add to the session
                 context.
         """
-        if isinstance(table, Table):
-            self.ctx.register_table(name, table.table)
-        else:
-            self.ctx.register_table(name, table)
+        provider = _normalize_table_provider(table)
+        self.ctx.register_table(name, provider)
 
     def deregister_table(self, name: str) -> None:
         """Remove a table from the session."""

python/datafusion/utils.py

@@ -0,0 +1,61 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Miscellaneous helper utilities for DataFusion's Python bindings."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from datafusion._internal import EXPECTED_PROVIDER_MSG
+
+if TYPE_CHECKING:  # pragma: no cover - imported for typing only
+    from datafusion import TableProvider
+    from datafusion.catalog import Table
+    from datafusion.context import TableProviderExportable
+
+
+def _normalize_table_provider(
+    table: Table | TableProvider | TableProviderExportable,
+) -> Any:
+    """Return the underlying provider for supported table inputs.
+
+    Args:
+        table: A :class:`~datafusion.catalog.Table`,
+            :class:`~datafusion.table_provider.TableProvider`, or object exporting a
+            DataFusion table provider via ``__datafusion_table_provider__``.
+
+    Returns:
+        The object expected by the Rust bindings for table registration.
+
+    Raises:
+        TypeError: If ``table`` is not a supported table provider input.
+    """
+
+    from datafusion.catalog import Table as _Table
+    from datafusion.table_provider import TableProvider as _TableProvider
+
+    if isinstance(table, _Table):
+        return table.table
+
+    if isinstance(table, _TableProvider):
+        return table._table_provider
+
+    provider_factory = getattr(table, "__datafusion_table_provider__", None)
+    if callable(provider_factory):
+        return table
+
+    raise TypeError(EXPECTED_PROVIDER_MSG)