feat: add deprecation warnings and alias handling for auto_register_python_variables in SessionContext

Author: kosiew

python/datafusion/context.py

@@ -24,7 +24,7 @@
 import re
 import warnings
 from functools import cache
-from typing import TYPE_CHECKING, Any, Iterator, Protocol
+from typing import TYPE_CHECKING, Any, Protocol
 
 try:
     from warnings import deprecated  # Python 3.13+
@@ -64,6 +64,13 @@ def _load_optional_module(module_name: str) -> Any | None:
         return None
 
 
+_AUTO_REGISTER_PYTHON_VARIABLES_DEPRECATED = (
+    "SessionContext.auto_register_python_variables is deprecated; use "
+    "SessionContext.set_python_table_lookup() or the "
+    "'auto_register_python_objects' keyword argument instead."
+)
+
+
 class ArrowStreamExportable(Protocol):
     """Type hint for object exporting Arrow C Stream via Arrow PyCapsule Interface.
 
@@ -503,7 +510,8 @@ def __init__(
         config: SessionConfig | None = None,
         runtime: RuntimeEnvBuilder | None = None,
         *,
-        auto_register_python_objects: bool = True,
+        auto_register_python_objects: bool | None = None,
+        auto_register_python_variables: bool | None = None,
     ) -> None:
         """Main interface for executing queries with DataFusion.
 
@@ -517,6 +525,9 @@ def __init__(
             auto_register_python_objects: Automatically register referenced
                 Python objects (such as pandas or PyArrow data) when ``sql``
                 queries reference them by name.
+            auto_register_python_variables: Deprecated alias for
+                ``auto_register_python_objects``. When provided, it overrides
+                the automatic registration behavior.
 
         Example usage:
 
@@ -532,7 +543,33 @@ def __init__(
             config.config_internal if config is not None else None,
             runtime.config_internal if runtime is not None else None,
         )
-        self._auto_python_table_lookup = auto_register_python_objects
+        if auto_register_python_variables is not None:
+            warnings.warn(
+                _AUTO_REGISTER_PYTHON_VARIABLES_DEPRECATED,
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        if (
+            auto_register_python_objects is not None
+            and auto_register_python_variables is not None
+            and auto_register_python_objects != auto_register_python_variables
+        ):
+            conflict_message = (
+                "auto_register_python_objects and auto_register_python_variables "
+                "were provided with conflicting values."
+            )
+            raise ValueError(conflict_message)
+
+        if auto_register_python_objects is None:
+            if auto_register_python_variables is None:
+                auto_python_table_lookup = True
+            else:
+                auto_python_table_lookup = auto_register_python_variables
+        else:
+            auto_python_table_lookup = auto_register_python_objects
+
+        self._auto_python_table_lookup = auto_python_table_lookup
 
     def __repr__(self) -> str:
         """Print a string representation of the Session Context."""
@@ -579,6 +616,25 @@ def set_python_table_lookup(self, enabled: bool = True) -> SessionContext:
         self._auto_python_table_lookup = enabled
         return self
 
+    @property
+    def auto_register_python_variables(self) -> bool:
+        """Deprecated alias for :py:meth:`set_python_table_lookup`."""
+        warnings.warn(
+            _AUTO_REGISTER_PYTHON_VARIABLES_DEPRECATED,
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return getattr(self, "_auto_python_table_lookup", True)
+
+    @auto_register_python_variables.setter
+    def auto_register_python_variables(self, enabled: bool) -> None:
+        warnings.warn(
+            _AUTO_REGISTER_PYTHON_VARIABLES_DEPRECATED,
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        self.set_python_table_lookup(enabled)
+
     def register_object_store(
         self, schema: str, store: Any, host: str | None = None
     ) -> None:

python/tests/test_context.py

@@ -257,9 +257,10 @@ def test_from_pylist(ctx):
 
 
 def test_sql_missing_table_without_auto_register(ctx):
+    ctx.set_python_table_lookup(False)
     arrow_table = pa.Table.from_pydict({"value": [1, 2, 3]})  # noqa: F841
 
-    with pytest.raises(Exception, match="not found") as excinfo:
+    with pytest.raises(Exception, match="not found|No table named") as excinfo:
         ctx.sql("SELECT * FROM arrow_table").collect()
 
     # Test that our extraction method works correctly
@@ -319,6 +320,39 @@ def test_sql_auto_register_polars_dataframe():
     assert result[0].column(0).to_pylist()[0] == 2
 
 
+def test_session_context_constructor_alias_disables_lookup():
+    with pytest.deprecated_call():
+        ctx = SessionContext(auto_register_python_variables=False)
+
+    arrow_table = pa.Table.from_pydict({"value": [1, 2, 3]})  # noqa: F841
+
+    with pytest.raises(Exception, match="not found|No table named"):
+        ctx.sql("SELECT * FROM arrow_table").collect()
+
+    with pytest.deprecated_call():
+        assert ctx.auto_register_python_variables is False
+
+
+def test_session_context_property_alias_setter_enables_lookup():
+    ctx = SessionContext(auto_register_python_objects=False)
+    arrow_table = pa.Table.from_pydict({"value": [1, 2, 3]})  # noqa: F841
+
+    with pytest.raises(Exception, match="not found|No table named"):
+        ctx.sql("SELECT COUNT(*) FROM arrow_table").collect()
+
+    with pytest.deprecated_call():
+        ctx.auto_register_python_variables = True
+
+    result = ctx.sql(
+        "SELECT SUM(value) AS total FROM arrow_table",
+    ).collect()
+
+    assert result[0].column(0).to_pylist()[0] == 6
+
+    with pytest.deprecated_call():
+        assert ctx.auto_register_python_variables is True
+
+
 def test_from_pydict(ctx):
     # create a dataframe from Python dictionary
     data = {"a": [1, 2, 3], "b": [4, 5, 6]}