refactor: improve auto-registration logic for Arrow and DataFrame objects in SessionContext

Author: kosiew

python/datafusion/context.py

@@ -797,23 +797,29 @@ def _register_python_object(self, name: str, obj: Any) -> bool:
         if isinstance(obj, DataFrame):
             self.register_view(name, obj)
             registered = True
-        elif (
-            obj.__class__.__module__.startswith("polars.")
-            and obj.__class__.__name__ == "DataFrame"
-        ):
-            self.from_polars(obj, name=name)
-            registered = True
-        elif (
-            obj.__class__.__module__.startswith("pandas.")
-            and obj.__class__.__name__ == "DataFrame"
-        ):
-            self.from_pandas(obj, name=name)
-            registered = True
-        elif isinstance(obj, (pa.Table, pa.RecordBatch, pa.RecordBatchReader)) or (
-            hasattr(obj, "__arrow_c_stream__") or hasattr(obj, "__arrow_c_array__")
-        ):
+        elif isinstance(obj, (pa.Table, pa.RecordBatch, pa.RecordBatchReader)):
             self.from_arrow(obj, name=name)
             registered = True
+        else:
+            exports_arrow_capsule = hasattr(obj, "__arrow_c_stream__") or hasattr(
+                obj, "__arrow_c_array__"
+            )
+
+            if exports_arrow_capsule:
+                self.from_arrow(obj, name=name)
+                registered = True
+            elif (
+                obj.__class__.__module__.startswith("polars.")
+                and obj.__class__.__name__ == "DataFrame"
+            ):
+                self.from_polars(obj, name=name)
+                registered = True
+            elif (
+                obj.__class__.__module__.startswith("pandas.")
+                and obj.__class__.__name__ == "DataFrame"
+            ):
+                self.from_pandas(obj, name=name)
+                registered = True
 
         if registered:
             try:

python/tests/test_context.py

@@ -372,12 +372,23 @@ def test_sql_auto_register_case_insensitive_lookup():
     assert batches[0].column(0).to_pylist()[0] == 5
 
 
-def test_sql_auto_register_pandas_dataframe():
+def test_sql_auto_register_pandas_dataframe(monkeypatch):
     pd = pytest.importorskip("pandas")
 
     ctx = SessionContext(auto_register_python_objects=True)
     pandas_df = pd.DataFrame({"value": [1, 2, 3, 4]})  # noqa: F841
 
+    if not (
+        hasattr(pandas_df, "__arrow_c_stream__")
+        or hasattr(pandas_df, "__arrow_c_array__")
+    ):
+        pytest.skip("pandas does not expose Arrow capsule export")
+
+    def fail_from_pandas(*args, **kwargs):  # noqa: ANN002, ANN003
+        raise AssertionError("from_pandas should not be called during auto-registration")
+
+    monkeypatch.setattr(SessionContext, "from_pandas", fail_from_pandas)
+
     result = ctx.sql(
         "SELECT AVG(value) AS avg_value FROM pandas_df",
     ).collect()