feat: SessionContext UDF lookup helpers

Expose `udf(name)` / `udaf(name)` / `udwf(name)` lookups symmetric with
the existing `register_udf` / `register_udaf` / `register_udwf` setters,
plus `udfs()` / `udafs()` / `udwfs()` for enumerating registered
function names. Looked-up functions come back as the same
`ScalarUDF` / `AggregateUDF` / `WindowUDF` wrappers users already get
from registration, so they can be called as expressions or re-registered
into a different session.

Returns Vec<String> from the list helpers (sorted) rather than the raw
HashSet upstream returns, so calling code gets a stable ordering.

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

Author: timsaucer

crates/core/src/context.rs

@@ -35,7 +35,6 @@ use datafusion::datasource::listing::{
     ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl,
 };
 use datafusion::datasource::{MemTable, TableProvider};
-use datafusion::execution::TaskContextProvider;
 use datafusion::execution::context::{
     DataFilePaths, SQLOptions, SessionConfig, SessionContext, TaskContext,
 };
@@ -44,6 +43,7 @@ use datafusion::execution::memory_pool::{FairSpillPool, GreedyMemoryPool, Unboun
 use datafusion::execution::options::{ArrowReadOptions, ReadOptions};
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::execution::session_state::SessionStateBuilder;
+use datafusion::execution::{FunctionRegistry, TaskContextProvider};
 use datafusion::prelude::{
     AvroReadOptions, CsvReadOptions, DataFrame, JsonReadOptions, ParquetReadOptions,
 };
@@ -1072,6 +1072,39 @@ impl PySessionContext {
         self.ctx.deregister_udwf(name);
     }
 
+    pub fn udf(&self, name: &str) -> PyDataFusionResult<PyScalarUDF> {
+        let function = (*self.ctx.udf(name)?).clone();
+        Ok(PyScalarUDF { function })
+    }
+
+    pub fn udaf(&self, name: &str) -> PyDataFusionResult<PyAggregateUDF> {
+        let function = (*self.ctx.udaf(name)?).clone();
+        Ok(PyAggregateUDF { function })
+    }
+
+    pub fn udwf(&self, name: &str) -> PyDataFusionResult<PyWindowUDF> {
+        let function = (*self.ctx.udwf(name)?).clone();
+        Ok(PyWindowUDF { function })
+    }
+
+    pub fn udfs(&self) -> Vec<String> {
+        let mut names: Vec<String> = self.ctx.udfs().into_iter().collect();
+        names.sort();
+        names
+    }
+
+    pub fn udafs(&self) -> Vec<String> {
+        let mut names: Vec<String> = self.ctx.udafs().into_iter().collect();
+        names.sort();
+        names
+    }
+
+    pub fn udwfs(&self) -> Vec<String> {
+        let mut names: Vec<String> = self.ctx.udwfs().into_iter().collect();
+        names.sort();
+        names
+    }
+
     #[pyo3(signature = (name="datafusion"))]
     pub fn catalog(&self, py: Python, name: &str) -> PyResult<Py<PyAny>> {
         let catalog = self.ctx.catalog(name).ok_or(PyKeyError::new_err(format!(

python/datafusion/context.py

@@ -1310,6 +1310,65 @@ def deregister_udwf(self, name: str) -> None:
         """
         self.ctx.deregister_udwf(name)
 
+    def udf(self, name: str) -> ScalarUDF:
+        """Look up a registered scalar UDF by name.
+
+        Args:
+            name: Name of the registered scalar UDF.
+
+        Raises:
+            Exception: If no scalar UDF is registered under ``name``.
+        """
+        from datafusion.user_defined import ScalarUDF as _ScalarUDF  # noqa: PLC0415
+
+        wrapper = _ScalarUDF.__new__(_ScalarUDF)
+        wrapper._udf = self.ctx.udf(name)
+        return wrapper
+
+    def udaf(self, name: str) -> AggregateUDF:
+        """Look up a registered aggregate UDF by name.
+
+        Args:
+            name: Name of the registered aggregate UDF.
+
+        Raises:
+            Exception: If no aggregate UDF is registered under ``name``.
+        """
+        from datafusion.user_defined import (  # noqa: PLC0415
+            AggregateUDF as _AggregateUDF,
+        )
+
+        wrapper = _AggregateUDF.__new__(_AggregateUDF)
+        wrapper._udaf = self.ctx.udaf(name)
+        return wrapper
+
+    def udwf(self, name: str) -> WindowUDF:
+        """Look up a registered window UDF by name.
+
+        Args:
+            name: Name of the registered window UDF.
+
+        Raises:
+            Exception: If no window UDF is registered under ``name``.
+        """
+        from datafusion.user_defined import WindowUDF as _WindowUDF  # noqa: PLC0415
+
+        wrapper = _WindowUDF.__new__(_WindowUDF)
+        wrapper._udwf = self.ctx.udwf(name)
+        return wrapper
+
+    def udfs(self) -> list[str]:
+        """Return the sorted names of all registered scalar UDFs."""
+        return self.ctx.udfs()
+
+    def udafs(self) -> list[str]:
+        """Return the sorted names of all registered aggregate UDFs."""
+        return self.ctx.udafs()
+
+    def udwfs(self) -> list[str]:
+        """Return the sorted names of all registered window UDFs."""
+        return self.ctx.udwfs()
+
     def catalog(self, name: str = "datafusion") -> Catalog:
         """Retrieve a catalog by name."""
         return Catalog(self.ctx.catalog(name))

python/tests/test_udf.py

@@ -76,6 +76,27 @@ def test_register_udf(ctx, df) -> None:
     assert result == pa.array([False, False, True])
 
 
+def test_udf_lookup(ctx, df) -> None:
+    is_null = udf(
+        lambda x: x.is_null(),
+        [pa.float64()],
+        pa.bool_(),
+        volatility="immutable",
+        name="lookup_is_null",
+    )
+    ctx.register_udf(is_null)
+
+    assert "lookup_is_null" in ctx.udfs()
+
+    looked_up = ctx.udf("lookup_is_null")
+    df_result = df.select(looked_up(column("b")))
+    result = df_result.collect()[0].column(0)
+    assert result == pa.array([False, False, True])
+
+    with pytest.raises(Exception, match="no UDF named"):
+        ctx.udf("does_not_exist")
+
+
 class OverThresholdUDF:
     def __init__(self, threshold: int = 0) -> None:
         self.threshold = threshold