refactor: opt-in UDTF session injection via with_session flag

Replaces signature sniffing with an explicit ``with_session=True`` kwarg
on ``TableFunction`` / ``udtf``. Avoids name-based detection footguns
(positional-only ``session`` params, accidental ``**kwargs`` opt-in,
shadowing by unrelated params) and makes author intent visible at
registration. Also documents the feature in the UDTF user guide.

Rust field renamed ``accepts_session`` -> ``inject_session_on_call`` to
match the Python-side opt-in semantics.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: timsaucer

crates/core/src/udtf.rs

@@ -39,10 +39,11 @@ use crate::table::PyTable;
 #[derive(Debug, Clone)]
 pub(crate) struct PythonTableFunctionCallable {
     pub(crate) callable: Arc<Py<PyAny>>,
-    /// Whether the callable's signature accepts a ``session`` keyword
-    /// argument (or ``**kwargs``). When true the calling
-    /// :class:`SessionContext` is threaded through on each invocation.
-    pub(crate) accepts_session: bool,
+    /// When true, the calling :class:`SessionContext` is passed to the
+    /// callable as a ``session`` keyword argument on every invocation.
+    /// Opt-in at registration time via ``with_session=True`` on the
+    /// Python wrapper.
+    pub(crate) inject_session_on_call: bool,
 }
 
 /// Represents a user defined table function
@@ -62,12 +63,12 @@ pub(crate) enum PyTableFunctionInner {
 #[pymethods]
 impl PyTableFunction {
     #[new]
-    #[pyo3(signature=(name, func, session, accepts_session=false))]
+    #[pyo3(signature=(name, func, session, inject_session_on_call=false))]
     pub fn new(
         name: &str,
         func: Bound<'_, PyAny>,
         session: Option<Bound<PyAny>>,
-        accepts_session: bool,
+        inject_session_on_call: bool,
     ) -> PyResult<Self> {
         let inner = if func.hasattr("__datafusion_table_function__")? {
             let py = func.py();
@@ -95,7 +96,7 @@ impl PyTableFunction {
         } else {
             PyTableFunctionInner::PythonFunction(PythonTableFunctionCallable {
                 callable: Arc::new(func.unbind()),
-                accepts_session,
+                inject_session_on_call,
             })
         };
 
@@ -162,7 +163,7 @@ fn call_python_table_function(
     func: &PythonTableFunctionCallable,
     args: TableFunctionArgs,
 ) -> DataFusionResult<Arc<dyn TableProvider>> {
-    let py_session = if func.accepts_session {
+    let py_session = if func.inject_session_on_call {
         Some(py_session_from_session(args.session())?)
     } else {
         None

docs/source/user-guide/common-operations/udf-and-udfa.rst

@@ -431,3 +431,39 @@ that you wish to expose via PyO3, you need to expose it as a ``PyCapsule``.
             PyCapsule::new(py, provider, Some(name))
         }
     }
+
+Accessing the Calling Session
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Pure-Python UDTFs can opt into receiving the calling
+:py:class:`~datafusion.SessionContext` by registering with
+``with_session=True``. The context is passed as a ``session`` keyword
+argument on every invocation. Use it to look up registered tables,
+UDFs, or session configuration from inside the callback.
+
+.. code-block:: python
+
+    from datafusion import SessionContext, Table, udtf
+    from datafusion.context import TableProviderExportable
+    import pyarrow as pa
+    import pyarrow.dataset as ds
+
+    @udtf("list_tables", with_session=True)
+    def list_tables(*, session: SessionContext) -> TableProviderExportable:
+        names = sorted(session.catalog().schema().names())
+        batch = pa.RecordBatch.from_pydict({"name": names})
+        return Table(ds.dataset([batch]))
+
+    ctx = SessionContext()
+    ctx.register_batch("t1", pa.RecordBatch.from_pydict({"x": [1]}))
+    ctx.register_udtf(list_tables)
+    ctx.sql("SELECT * FROM list_tables()").show()
+
+Without ``with_session=True``, the callback receives only the positional
+expression arguments. The flag is opt-in so existing UDTFs keep working
+unchanged.
+
+The injected ``session`` is a fresh :py:class:`~datafusion.SessionContext`
+wrapper backed by the same underlying state as the caller, so registries
+(tables, UDFs, catalogs) are visible. Mutations made through it affect
+the live session.

python/datafusion/user_defined.py

@@ -1054,29 +1054,6 @@ def from_pycapsule(func: WindowUDFExportable) -> WindowUDF:
         )
 
 
-def _callable_accepts_session_kwarg(func: object) -> bool:
-    """Return True if ``func`` accepts a ``session`` keyword argument.
-
-    Used to opt a Python UDTF callback into receiving the calling
-    :class:`SessionContext` at invocation time. ``**kwargs`` callables
-    are treated as accepting it; built-ins and objects without an
-    introspectable signature fall back to ``False``.
-    """
-    import inspect  # noqa: PLC0415
-
-    try:
-        signature = inspect.signature(func)
-    except (TypeError, ValueError):
-        return False
-
-    for parameter in signature.parameters.values():
-        if parameter.name == "session":
-            return True
-        if parameter.kind is inspect.Parameter.VAR_KEYWORD:
-            return True
-    return False
-
-
 def _wrap_session_kwarg_for_udtf(func: Callable[..., Any]) -> Callable[..., Any]:
     """Adapt the raw internal session pyo3 object back to a Python wrapper.
 
@@ -1103,23 +1080,27 @@ class TableFunction:
     """
 
     def __init__(
-        self, name: str, func: Callable[[], any], ctx: SessionContext | None = None
+        self,
+        name: str,
+        func: Callable[..., Any],
+        ctx: SessionContext | None = None,
+        *,
+        with_session: bool = False,
     ) -> None:
         """Instantiate a user-defined table function (UDTF).
 
-        If ``func``'s signature accepts a ``session`` keyword (or
-        ``**kwargs``), the calling :class:`SessionContext` is threaded
-        through to it on each invocation. Use it inside the body to look
-        up registered tables, UDFs, or session configuration. Callables
-        whose signatures do not declare ``session`` are invoked with the
-        positional expression arguments only.
+        Set ``with_session=True`` to have the calling
+        :class:`SessionContext` passed as a ``session`` keyword argument
+        on each invocation. Use it inside the callback to look up
+        registered tables, UDFs, or session configuration. When
+        ``with_session`` is ``False`` (the default), ``func`` is invoked
+        with the positional expression arguments only.
 
         See :py:func:`udtf` for a convenience function and argument
         descriptions.
         """
-        accepts_session = _callable_accepts_session_kwarg(func)
-        registered = _wrap_session_kwarg_for_udtf(func) if accepts_session else func
-        self._udtf = df_internal.TableFunction(name, registered, ctx, accepts_session)
+        registered = _wrap_session_kwarg_for_udtf(func) if with_session else func
+        self._udtf = df_internal.TableFunction(name, registered, ctx, with_session)
 
     def __call__(self, *args: Expr) -> Any:
         """Execute the UDTF and return a table provider."""
@@ -1130,47 +1111,66 @@ def __call__(self, *args: Expr) -> Any:
     @staticmethod
     def udtf(
         name: str,
+        *,
+        with_session: bool = False,
     ) -> Callable[..., Any]: ...
 
     @overload
     @staticmethod
     def udtf(
-        func: Callable[[], Any],
+        func: Callable[..., Any],
         name: str,
+        *,
+        with_session: bool = False,
     ) -> TableFunction: ...
 
     @staticmethod
-    def udtf(*args: Any, **kwargs: Any):
-        """Create a new User-Defined Table Function (UDTF)."""
+    def udtf(*args: Any, with_session: bool = False, **kwargs: Any):
+        """Create a new User-Defined Table Function (UDTF).
+
+        Pass ``with_session=True`` to have the calling
+        :class:`SessionContext` injected as a ``session`` keyword
+        argument on each invocation.
+        """
         if args and callable(args[0]):
             # Case 1: Used as a function, require the first parameter to be callable
-            return TableFunction._create_table_udf(*args, **kwargs)
+            return TableFunction._create_table_udf(
+                *args, with_session=with_session, **kwargs
+            )
         if args and hasattr(args[0], "__datafusion_table_function__"):
             # Case 2: We have a datafusion FFI provided function
             return TableFunction(args[1], args[0])
         # Case 3: Used as a decorator with parameters
-        return TableFunction._create_table_udf_decorator(*args, **kwargs)
+        return TableFunction._create_table_udf_decorator(
+            *args, with_session=with_session, **kwargs
+        )
 
     @staticmethod
     def _create_table_udf(
         func: Callable[..., Any],
         name: str,
+        *,
+        with_session: bool = False,
     ) -> TableFunction:
         """Create a TableFunction instance from function arguments."""
         if not callable(func):
             msg = "`func` must be callable."
             raise TypeError(msg)
 
-        return TableFunction(name, func)
+        return TableFunction(name, func, with_session=with_session)
 
     @staticmethod
     def _create_table_udf_decorator(
         name: str | None = None,
-    ) -> Callable[[Callable[[], WindowEvaluator]], Callable[..., Expr]]:
-        """Create a decorator for a WindowUDF."""
-
-        def decorator(func: Callable[[], WindowEvaluator]) -> Callable[..., Expr]:
-            return TableFunction._create_table_udf(func, name)
+        *,
+        with_session: bool = False,
+    ) -> Callable[[Callable[..., Any]], TableFunction]:
+        """Create a decorator for a TableFunction."""
+
+        def decorator(func: Callable[..., Any]) -> TableFunction:
+            return TableFunction._create_table_udf(
+                func, name, with_session=with_session
+            )
 
         return decorator
 

python/tests/test_udtf.py

@@ -137,11 +137,11 @@ def string_arg_func(prefix: Expr) -> TableProviderExportable:
 
 
 def test_python_table_function_receives_session() -> None:
-    """A UDTF whose signature declares ``session`` gets the calling ctx."""
+    """A UDTF registered ``with_session=True`` gets the calling ctx."""
     ctx = SessionContext()
     captured: list[SessionContext] = []
 
-    @udtf("session_aware_func")
+    @udtf("session_aware_func", with_session=True)
     def session_aware_func(*, session: SessionContext) -> TableProviderExportable:
         captured.append(session)
         batch = pa.RecordBatch.from_pydict({"a": [1, 2, 3]})
@@ -165,7 +165,7 @@ def test_python_table_function_session_used_for_metadata() -> None:
 
     seen_tables: list[set[str]] = []
 
-    @udtf("table_inventory")
+    @udtf("table_inventory", with_session=True)
     def table_inventory(*, session: SessionContext) -> TableProviderExportable:
         # Stash the visible tables to verify the session wired through.
         seen_tables.append(session.catalog().schema().names())
@@ -179,8 +179,8 @@ def table_inventory(*, session: SessionContext) -> TableProviderExportable:
     assert result[0].column(0).to_pylist() == ["base_tbl"]
 
 
-def test_python_table_function_class_callable_session_kwarg() -> None:
-    """Class-based UDTFs whose __call__ accepts ``session`` get it too."""
+def test_python_table_function_class_callable_with_session() -> None:
+    """Class-based UDTFs opt in via ``with_session=True``."""
     ctx = SessionContext()
     captured: list[SessionContext] = []
 
@@ -193,9 +193,25 @@ def __call__(
             batch = pa.RecordBatch.from_pydict({"a": list(range(count))})
             return Table(ds.dataset([batch]))
 
-    ctx.register_udtf(udtf(SessionAware(), "session_class_func"))
+    ctx.register_udtf(udtf(SessionAware(), "session_class_func", with_session=True))
     result = ctx.sql("SELECT * FROM session_class_func(3)").collect()
 
     assert len(captured) == 1
     assert isinstance(captured[0], SessionContext)
     assert result[0].column(0).to_pylist() == [0, 1, 2]
+
+
+def test_python_table_function_without_session_flag_no_injection() -> None:
+    """Default registration (no ``with_session``) calls func positionally."""
+    ctx = SessionContext()
+
+    @udtf("plain_func")
+    def plain_func(n: Expr) -> TableProviderExportable:
+        count = n.to_variant().value_i64()
+        batch = pa.RecordBatch.from_pydict({"a": list(range(count))})
+        return Table(ds.dataset([batch]))
+
+    ctx.register_udtf(plain_func)
+    result = ctx.sql("SELECT * FROM plain_func(4)").collect()
+
+    assert result[0].column(0).to_pylist() == [0, 1, 2, 3]