refactor(udf): raise KeyError on UDF/UDAF/UDWF lookup miss

`SessionContext.udf` / `udaf` / `udwf` previously surfaced upstream
`DataFusionError::Plan` as a generic exception whose message ("There
is no UDF named ...") is set by DataFusion and can drift between
releases. Pre-check membership via `udfs()` / `udafs()` / `udwfs()`
and raise `PyKeyError` on miss so callers get the Pythonic
dict-style lookup behavior and tests are no longer coupled to the
upstream wording.

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

Author: timsaucer

crates/core/src/context.rs

@@ -1072,18 +1072,27 @@ impl PySessionContext {
         self.ctx.deregister_udwf(name);
     }
 
-    pub fn udf(&self, name: &str) -> PyDataFusionResult<PyScalarUDF> {
-        let function = (*self.ctx.udf(name)?).clone();
+    pub fn udf(&self, name: &str) -> PyResult<PyScalarUDF> {
+        if !self.ctx.udfs().contains(name) {
+            return Err(PyKeyError::new_err(format!("no UDF named '{name}'")));
+        }
+        let function = (*self.ctx.udf(name).map_err(py_datafusion_err)?).clone();
         Ok(PyScalarUDF { function })
     }
 
-    pub fn udaf(&self, name: &str) -> PyDataFusionResult<PyAggregateUDF> {
-        let function = (*self.ctx.udaf(name)?).clone();
+    pub fn udaf(&self, name: &str) -> PyResult<PyAggregateUDF> {
+        if !self.ctx.udafs().contains(name) {
+            return Err(PyKeyError::new_err(format!("no UDAF named '{name}'")));
+        }
+        let function = (*self.ctx.udaf(name).map_err(py_datafusion_err)?).clone();
         Ok(PyAggregateUDF { function })
     }
 
-    pub fn udwf(&self, name: &str) -> PyDataFusionResult<PyWindowUDF> {
-        let function = (*self.ctx.udwf(name)?).clone();
+    pub fn udwf(&self, name: &str) -> PyResult<PyWindowUDF> {
+        if !self.ctx.udwfs().contains(name) {
+            return Err(PyKeyError::new_err(format!("no UDWF named '{name}'")));
+        }
+        let function = (*self.ctx.udwf(name).map_err(py_datafusion_err)?).clone();
         Ok(PyWindowUDF { function })
     }
 

python/datafusion/context.py

@@ -1323,7 +1323,7 @@ def udf(self, name: str) -> ScalarUDF:
             name: Name of the registered scalar UDF.
 
         Raises:
-            Exception: If no scalar UDF is registered under ``name``.
+            KeyError: If no scalar UDF is registered under ``name``.
 
         Examples:
             Register a UDF, then look it up by name and use it in the
@@ -1371,7 +1371,7 @@ def udaf(self, name: str) -> AggregateUDF:
             name: Name of the registered aggregate UDF.
 
         Raises:
-            Exception: If no aggregate UDF is registered under ``name``.
+            KeyError: If no aggregate UDF is registered under ``name``.
 
         Examples:
             Look up a built-in aggregate by name and use it in
@@ -1404,7 +1404,7 @@ def udwf(self, name: str) -> WindowUDF:
             name: Name of the registered window UDF.
 
         Raises:
-            Exception: If no window UDF is registered under ``name``.
+            KeyError: If no window UDF is registered under ``name``.
 
         Examples:
             Look up a built-in window function by name and use it in

python/tests/test_udf.py

@@ -93,7 +93,7 @@ def test_udf_lookup(ctx, df) -> None:
     result = df_result.collect()[0].column(0)
     assert result == pa.array([False, False, True])
 
-    with pytest.raises(Exception, match="no UDF named"):
+    with pytest.raises(KeyError, match="no UDF named"):
         ctx.udf("does_not_exist")
 
 
@@ -133,7 +133,7 @@ def test_udaf_lookup_builtin(ctx, df) -> None:
     result = df.aggregate([], [sum_fn(column("a")).alias("total")]).collect()
     assert result[0].column(0).to_pylist() == [6]
 
-    with pytest.raises(Exception, match="no UDAF named"):
+    with pytest.raises(KeyError, match="no UDAF named"):
         ctx.udaf("does_not_exist")
 
 
@@ -143,7 +143,7 @@ def test_udwf_lookup_builtin(ctx, df) -> None:
     result = df.select(column("a"), rn().alias("rn")).collect()
     assert result[0].column(1).to_pylist() == [1, 2, 3]
 
-    with pytest.raises(Exception, match="no UDWF named"):
+    with pytest.raises(KeyError, match="no UDWF named"):
         ctx.udwf("does_not_exist")