feat: add automatic variable registration for Arrow-compatible Python objects in SQL queries

Author: kosiew

docs/source/user-guide/sql.rst

@@ -36,4 +36,29 @@ DataFusion also offers a SQL API, read the full reference `here <https://arrow.a
     df = ctx.sql('SELECT "Attack"+"Defense", "Attack"-"Defense" FROM pokemon')
 
     # collect and convert to pandas DataFrame
-    df.to_pandas()
+    df.to_pandas()
+
+Automatic variable registration
+-------------------------------
+
+You can opt-in to DataFusion automatically registering Arrow-compatible Python
+objects that appear in SQL queries. This removes the need to call
+``register_*`` helpers explicitly when working with in-memory data structures.
+
+.. code-block:: python
+
+    import pyarrow as pa
+    from datafusion import SessionContext
+
+    ctx = SessionContext(auto_register_python_variables=True)
+
+    orders = pa.Table.from_pydict({"item": ["apple", "pear"], "qty": [5, 2]})
+
+    result = ctx.sql("SELECT item, qty FROM orders WHERE qty > 2")
+    print(result.to_pandas())
+
+The feature inspects the call stack for variables whose names match missing
+tables and registers them if they expose Arrow data (including pandas and
+Polars DataFrames). Existing contexts can enable or disable the behavior at
+runtime through the :py:attr:`SessionContext.auto_register_python_variables`
+property.

python/datafusion/context.py

@@ -19,8 +19,10 @@
 
 from __future__ import annotations
 
+import inspect
+import re
 import warnings
-from typing import TYPE_CHECKING, Any, Protocol
+from typing import TYPE_CHECKING, Any, Iterator, Protocol
 
 try:
     from warnings import deprecated  # Python 3.13+
@@ -41,6 +43,8 @@
 from ._internal import SQLOptions as SQLOptionsInternal
 from ._internal import expr as expr_internal
 
+_MISSING_TABLE_PATTERN = re.compile(r"(?i)(?:table|view) '([^']+)' not found")
+
 if TYPE_CHECKING:
     import pathlib
     from collections.abc import Sequence
@@ -483,6 +487,8 @@ def __init__(
         self,
         config: SessionConfig | None = None,
         runtime: RuntimeEnvBuilder | None = None,
+        *,
+        auto_register_python_variables: bool = False,
     ) -> None:
         """Main interface for executing queries with DataFusion.
 
@@ -493,6 +499,9 @@ def __init__(
         Args:
             config: Session configuration options.
             runtime: Runtime configuration options.
+            auto_register_python_variables: Automatically register Arrow-like
+                Python objects referenced in SQL queries when they are available
+                in the caller's scope.
 
         Example usage:
 
@@ -508,6 +517,7 @@ def __init__(
         runtime = runtime.config_internal if runtime is not None else None
 
         self.ctx = SessionContextInternal(config, runtime)
+        self._auto_register_python_variables = auto_register_python_variables
 
     def __repr__(self) -> str:
         """Print a string representation of the Session Context."""
@@ -534,8 +544,18 @@ def enable_url_table(self) -> SessionContext:
         klass = self.__class__
         obj = klass.__new__(klass)
         obj.ctx = self.ctx.enable_url_table()
+        obj._auto_register_python_variables = self._auto_register_python_variables
         return obj
 
+    @property
+    def auto_register_python_variables(self) -> bool:
+        """Toggle automatic registration of Python variables in SQL queries."""
+        return self._auto_register_python_variables
+
+    @auto_register_python_variables.setter
+    def auto_register_python_variables(self, enabled: bool) -> None:
+        self._auto_register_python_variables = bool(enabled)
+
     def register_object_store(
         self, schema: str, store: Any, host: str | None = None
     ) -> None:
@@ -600,9 +620,12 @@ def sql(self, query: str, options: SQLOptions | None = None) -> DataFrame:
         Returns:
             DataFrame representation of the SQL query.
         """
-        if options is None:
-            return DataFrame(self.ctx.sql(query))
-        return DataFrame(self.ctx.sql_with_options(query, options.options_internal))
+        options_internal = None if options is None else options.options_internal
+        return self._sql_with_retry(
+            query,
+            options_internal,
+            self._auto_register_python_variables,
+        )
 
     def sql_with_options(self, query: str, options: SQLOptions) -> DataFrame:
         """Create a :py:class:`~datafusion.dataframe.DataFrame` from SQL query text.
@@ -619,6 +642,138 @@ def sql_with_options(self, query: str, options: SQLOptions) -> DataFrame:
         """
         return self.sql(query, options)
 
+    def _sql_with_retry(
+        self,
+        query: str,
+        options_internal: SQLOptionsInternal | None,
+        allow_retry: bool,
+    ) -> DataFrame:
+        try:
+            if options_internal is None:
+                return DataFrame(self.ctx.sql(query))
+            return DataFrame(self.ctx.sql_with_options(query, options_internal))
+        except Exception as exc:
+            if not allow_retry or not self._handle_missing_table_error(exc):
+                raise
+            return self._sql_with_retry(query, options_internal, allow_retry)
+
+    def _handle_missing_table_error(self, error: Exception) -> bool:
+        missing_tables = self._extract_missing_table_names(error)
+        if not missing_tables:
+            return False
+
+        registered_any = False
+        attempted: set[str] = set()
+        for raw_name in missing_tables:
+            for candidate in self._candidate_table_names(raw_name):
+                if candidate in attempted:
+                    continue
+                attempted.add(candidate)
+
+                value = self._lookup_python_variable(candidate)
+                if value is None:
+                    continue
+                if self._register_python_value(candidate, value):
+                    registered_any = True
+                    break
+        return registered_any
+
+    def _candidate_table_names(self, identifier: str) -> Iterator[str]:
+        cleaned = identifier.strip().strip('"')
+        if not cleaned:
+            return
+
+        seen: set[str] = set()
+        candidates = [cleaned]
+        if "." in cleaned:
+            candidates.append(cleaned.rsplit(".", 1)[-1])
+
+        for candidate in candidates:
+            normalized = candidate.strip()
+            if not normalized or normalized in seen:
+                continue
+            seen.add(normalized)
+            yield normalized
+
+    def _extract_missing_table_names(self, error: Exception) -> set[str]:
+        names: set[str] = set()
+        attribute = getattr(error, "missing_table_names", None)
+        if attribute is not None:
+            if isinstance(attribute, (list, tuple, set, frozenset)):
+                for item in attribute:
+                    if item is None:
+                        continue
+                    for candidate in self._candidate_table_names(str(item)):
+                        names.add(candidate)
+            elif attribute is not None:
+                for candidate in self._candidate_table_names(str(attribute)):
+                    names.add(candidate)
+        if names:
+            return names
+
+        message = str(error)
+        return {match.group(1) for match in _MISSING_TABLE_PATTERN.finditer(message)}
+
+    def _lookup_python_variable(self, name: str) -> Any | None:
+        frame = inspect.currentframe()
+        outer = frame.f_back if frame is not None else None
+        lower_name = name.lower()
+
+        try:
+            while outer is not None:
+                for mapping in (outer.f_locals, outer.f_globals):
+                    if not mapping:
+                        continue
+                    if name in mapping:
+                        value = mapping[name]
+                        if value is not None:
+                            return value
+                        # allow outer scopes to provide a non-``None`` value
+                        continue
+                    for key, value in mapping.items():
+                        if value is None:
+                            continue
+                        if key == name or key.lower() == lower_name:
+                            return value
+                outer = outer.f_back
+        finally:
+            del outer
+            del frame
+        return None
+
+    def _register_python_value(self, table_name: str, value: Any) -> bool:
+        if value is None:
+            return False
+
+        registered = False
+        if isinstance(value, DataFrame):
+            self.register_view(table_name, value)
+            registered = True
+        elif isinstance(value, Table):
+            self.register_table(table_name, value)
+            registered = True
+        else:
+            provider = getattr(value, "__datafusion_table_provider__", None)
+            if callable(provider):
+                self.register_table_provider(table_name, value)
+                registered = True
+            elif hasattr(value, "__arrow_c_stream__") or hasattr(
+                value, "__arrow_c_array__"
+            ):
+                self.from_arrow(value, name=table_name)
+                registered = True
+            else:
+                module_name = getattr(type(value), "__module__", "") or ""
+                class_name = getattr(type(value), "__name__", "") or ""
+                if module_name.startswith("pandas.") and class_name == "DataFrame":
+                    self.from_pandas(value, name=table_name)
+                    registered = True
+                elif module_name.startswith("polars") and class_name == "DataFrame":
+                    self.from_polars(value, name=table_name)
+                    registered = True
+
+        return registered
+
     def create_dataframe(
         self,
         partitions: list[list[pa.RecordBatch]],

python/tests/test_context.py

@@ -255,6 +255,69 @@ def test_from_pylist(ctx):
     assert df.collect()[0].num_rows == 3
 
 
+def test_sql_missing_table_without_auto_register(ctx):
+    arrow_table = pa.Table.from_pydict({"value": [1, 2, 3]})  # noqa: F841
+
+    with pytest.raises(Exception, match="not found") as excinfo:
+        ctx.sql("SELECT * FROM arrow_table").collect()
+
+    missing = getattr(excinfo.value, "missing_table_names", None)
+    assert missing is not None
+    assert "arrow_table" in set(ctx._extract_missing_table_names(excinfo.value))
+
+
+def test_sql_auto_register_arrow_table():
+    ctx = SessionContext(auto_register_python_variables=True)
+    arrow_table = pa.Table.from_pydict({"value": [1, 2, 3]})  # noqa: F841
+
+    result = ctx.sql(
+        "SELECT SUM(value) AS total FROM arrow_table",
+    ).collect()
+
+    assert ctx.table_exist("arrow_table")
+    assert result[0].column(0).to_pylist()[0] == 6
+
+
+def test_sql_auto_register_arrow_outer_scope():
+    ctx = SessionContext()
+    ctx.auto_register_python_variables = True
+    arrow_table = pa.Table.from_pydict({"value": [1, 2, 3, 4]})  # noqa: F841
+
+    def run_query():
+        return ctx.sql(
+            "SELECT COUNT(*) AS total_rows FROM arrow_table",
+        ).collect()
+
+    result = run_query()
+    assert result[0].column(0).to_pylist()[0] == 4
+
+
+def test_sql_auto_register_pandas_dataframe():
+    pd = pytest.importorskip("pandas")
+
+    ctx = SessionContext(auto_register_python_variables=True)
+    pandas_df = pd.DataFrame({"value": [1, 2, 3, 4]})  # noqa: F841
+
+    result = ctx.sql(
+        "SELECT AVG(value) AS avg_value FROM pandas_df",
+    ).collect()
+
+    assert pytest.approx(result[0].column(0).to_pylist()[0]) == 2.5
+
+
+def test_sql_auto_register_polars_dataframe():
+    pl = pytest.importorskip("polars")
+
+    ctx = SessionContext(auto_register_python_variables=True)
+    polars_df = pl.DataFrame({"value": [2, 4, 6]})  # noqa: F841
+
+    result = ctx.sql(
+        "SELECT MIN(value) AS min_value FROM polars_df",
+    ).collect()
+
+    assert result[0].column(0).to_pylist()[0] == 2
+
+
 def test_from_pydict(ctx):
     # create a dataframe from Python dictionary
     data = {"a": [1, 2, 3], "b": [4, 5, 6]}
@@ -484,7 +547,7 @@ def test_table_exist(ctx):
 
 
 def test_table_not_found(ctx):
-    from uuid import uuid4
+    from uuid import uuid4  # noqa: PLC0415
 
     with pytest.raises(KeyError):
         ctx.table(f"not-found-{uuid4()}")