This is an automated email from the ASF dual-hosted git repository.

HyukjinKwon pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark-connect-rust.git


The following commit(s) were added to refs/heads/master by this push:
     new f3c67fb  [SPARK-59113] Fix PySpark UDF parity in the Rust drop-in
f3c67fb is described below

commit f3c67fb435eafc31c6ec30d3277813fe2b3d6a63
Author: Hyukjin Kwon <[email protected]>
AuthorDate: Sun Aug 30 05:25:36 2026 +0900

    [SPARK-59113] Fix PySpark UDF parity in the Rust drop-in
    
    A full behavioral UDF sweep against real PySpark 4.2.0 (live Spark Connect
    server) found the UDF subsystem broadly broken. Fix it end to end.
    
    - Scalar / arrow-optimized python UDFs treated a string column arg as a 
literal
      (udf(f)("c") vs udf(f)(F.col("c"))), and @udf / @pandas_udf / @arrow_udf
      decorator forms raised TypeError. Fixed the UDF __call__ arg coercion 
(str->col)
      and the decorator/return-type handling.
    - pandas_udf family: the command pickled a DDL string instead of a concrete
      DataType (the worker's Arrow serializer failed), and the eval-type was not
      inferred from type hints. Now pickles a real DataType and infers the eval 
type
      (scalar 200 / iterator 204 / grouped-agg 202).
    - Grouped/map UDFs (applyInPandas, mapInPandas, mapInArrow, applyInArrow,
      cogroup): the GroupMap / MapPartitions / CoGroupMap relation did not pass 
the
      input column names, so the worker's pandas/Arrow DataFrame was unnamed
      (KeyError). Now passes the child columns as the UDF arguments.
    - spark.udf.register(name, f, T) did not notify the server 
(UNRESOLVED_ROUTINE in
      SQL). Now sends a RegisterFunction command with the cloudpickled 
PythonUDF.
    
    Verified live: pandas_udf scalar/iterator/grouped-agg, applyInPandas,
    mapInPandas, mapInArrow, applyInArrow, cogroup, and register+SQL all return
    correct results. Example battery 32/32; offline suite 236 passed;
    `cargo build --workspace --all-targets` clean; rustfmt clean.
    
    Co-authored-by: Isaac <[email protected]>
---
 crates/pyspark-rs/src/dataframe.rs  | 15 +++++-
 crates/pyspark-rs/src/group.rs      | 23 ++++++++-
 crates/pyspark-rs/src/session.rs    | 30 +++++++++++
 crates/spark-connect/src/group.rs   | 14 ++++++
 python/pyspark/sql/udf.py           | 99 +++++++++++++++++++++++++++++++------
 python/tests/test_dropin_offline.py | 84 ++++++++++++++++++++++++++++++-
 6 files changed, 247 insertions(+), 18 deletions(-)

diff --git a/crates/pyspark-rs/src/dataframe.rs 
b/crates/pyspark-rs/src/dataframe.rs
index 7f0e04c..5840d18 100644
--- a/crates/pyspark-rs/src/dataframe.rs
+++ b/crates/pyspark-rs/src/dataframe.rs
@@ -1381,12 +1381,14 @@ impl PyDataFrame {
         schema: Bound<'_, PyAny>,
         is_barrier: bool,
     ) -> PyResult<PyDataFrame> {
+        let cols = py.detach(|| self.dataframe.columns()).to_pyerr()?;
         let udf = build_map_udf(
             py,
             "mapInPandas",
             &func,
             &schema,
             eval_type::SQL_MAP_PANDAS_ITER_UDF,
+            &cols,
         )?;
         Ok(PyDataFrame::new(
             self.dataframe.map_in_pandas(udf, is_barrier),
@@ -1402,12 +1404,14 @@ impl PyDataFrame {
         schema: Bound<'_, PyAny>,
         is_barrier: bool,
     ) -> PyResult<PyDataFrame> {
+        let cols = py.detach(|| self.dataframe.columns()).to_pyerr()?;
         let udf = build_map_udf(
             py,
             "mapInArrow",
             &func,
             &schema,
             eval_type::SQL_MAP_ARROW_ITER_UDF,
+            &cols,
         )?;
         Ok(PyDataFrame::new(
             self.dataframe.map_in_arrow(udf, is_barrier),
@@ -2202,15 +2206,24 @@ pub(crate) fn build_map_udf(
     func: &Bound<'_, PyAny>,
     schema: &Bound<'_, PyAny>,
     eval_type_id: i32,
+    input_cols: &[String],
 ) -> PyResult<CommonInlineUserDefinedFunctionExpression> {
     let return_type = resolve_datatype(schema)?;
     let tup = PyTuple::new(py, [func.clone(), schema.clone()])?;
     let command = py_cloudpickle(py, tup.as_any())?;
     let payload = PythonUDFPayload::new(return_type, eval_type_id, command, 
py_version(py));
+    // The worker reconstructs the input pandas/Arrow DataFrame from the UDF's 
argument
+    // columns; without them the input columns are unnamed (KeyError in the 
user fn).
+    // Mirror the reference Connect client, which passes the child's columns 
as the
+    // function arguments.
+    let arguments = input_cols
+        .iter()
+        .map(|c| spark_connect::functions::col(c).expression().clone())
+        .collect();
     Ok(CommonInlineUserDefinedFunctionExpression::new(
         name.to_string(),
         true,
-        vec![],
+        arguments,
         payload,
     ))
 }
diff --git a/crates/pyspark-rs/src/group.rs b/crates/pyspark-rs/src/group.rs
index d4270a7..4f129a0 100644
--- a/crates/pyspark-rs/src/group.rs
+++ b/crates/pyspark-rs/src/group.rs
@@ -5,6 +5,7 @@ use spark_connect::group::{CoGroupedData as RustCoGroupedData, 
GroupedData as Ru
 use spark_connect::udf::eval_type;
 
 use crate::dataframe::{build_map_udf, PyDataFrame};
+use crate::errors::ResultExt;
 use crate::functions::to_column;
 
 /// Python wrapper for Spark GroupedData.
@@ -152,12 +153,14 @@ impl PyGroupedData {
         func: Bound<'_, PyAny>,
         schema: Bound<'_, PyAny>,
     ) -> PyResult<PyDataFrame> {
+        let cols = py.detach(|| self.grouped_data.input_columns()).to_pyerr()?;
         let udf = build_map_udf(
             py,
             "applyInPandas",
             &func,
             &schema,
             eval_type::SQL_GROUPED_MAP_PANDAS_UDF,
+            &cols,
         )?;
         Ok(PyDataFrame::new(self.grouped_data.apply_in_pandas(udf)))
     }
@@ -183,12 +186,14 @@ impl PyGroupedData {
         func: Bound<'_, PyAny>,
         schema: Bound<'_, PyAny>,
     ) -> PyResult<PyDataFrame> {
+        let cols = py.detach(|| self.grouped_data.input_columns()).to_pyerr()?;
         let udf = build_map_udf(
             py,
             "applyInArrow",
             &func,
             &schema,
             eval_type::SQL_GROUPED_MAP_ARROW_UDF,
+            &cols,
         )?;
         Ok(PyDataFrame::new(self.grouped_data.apply_in_arrow(udf)))
     }
@@ -213,12 +218,14 @@ impl PyGroupedData {
         output_mode: &str,
         timeout_conf: &str,
     ) -> PyResult<PyDataFrame> {
+        let cols = py.detach(|| self.grouped_data.input_columns()).to_pyerr()?;
         let udf = build_map_udf(
             py,
             "applyInPandasWithState",
             &func,
             &output_struct_type,
             eval_type::SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE,
+            &cols,
         )?;
         let state_schema = 
crate::dataframe::resolve_datatype(&state_struct_type)?;
         Ok(PyDataFrame::new(
@@ -255,12 +262,14 @@ impl PyGroupedData {
             initialState.is_some(),
             true,
         )?;
+        let cols = py.detach(|| self.grouped_data.input_columns()).to_pyerr()?;
         let udf = build_map_udf(
             py,
             "transformWithStateInPandas",
             &func,
             &outputStructType,
             et,
+            &cols,
         )?;
         let out = crate::types::py_to_data_type(&outputStructType)?;
         let etc = 
(!eventTimeColumnName.is_empty()).then_some(eventTimeColumnName);
@@ -298,7 +307,15 @@ impl PyGroupedData {
             initialState.is_some(),
             false,
         )?;
-        let udf = build_map_udf(py, "transformWithState", &func, 
&outputStructType, et)?;
+        let cols = py.detach(|| self.grouped_data.input_columns()).to_pyerr()?;
+        let udf = build_map_udf(
+            py,
+            "transformWithState",
+            &func,
+            &outputStructType,
+            et,
+            &cols,
+        )?;
         let etc = 
(!eventTimeColumnName.is_empty()).then_some(eventTimeColumnName);
         Ok(PyDataFrame::new(self.grouped_data.transform_with_state(
             udf,
@@ -361,12 +378,14 @@ impl PyCoGroupedData {
         func: Bound<'_, PyAny>,
         schema: Bound<'_, PyAny>,
     ) -> PyResult<PyDataFrame> {
+        let cols = py.detach(|| self.inner.input_columns()).to_pyerr()?;
         let udf = build_map_udf(
             py,
             "applyInPandas",
             &func,
             &schema,
             eval_type::SQL_COGROUPED_MAP_PANDAS_UDF,
+            &cols,
         )?;
         Ok(PyDataFrame::new(self.inner.apply_in_pandas(udf)))
     }
@@ -379,12 +398,14 @@ impl PyCoGroupedData {
         func: Bound<'_, PyAny>,
         schema: Bound<'_, PyAny>,
     ) -> PyResult<PyDataFrame> {
+        let cols = py.detach(|| self.inner.input_columns()).to_pyerr()?;
         let udf = build_map_udf(
             py,
             "applyInArrow",
             &func,
             &schema,
             eval_type::SQL_COGROUPED_MAP_ARROW_UDF,
+            &cols,
         )?;
         Ok(PyDataFrame::new(self.inner.apply_in_arrow(udf)))
     }
diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs
index ea91b6b..3445c2a 100644
--- a/crates/pyspark-rs/src/session.rs
+++ b/crates/pyspark-rs/src/session.rs
@@ -5,6 +5,7 @@ use pyo3::types::{PyDict, PyList};
 use spark_connect::row::{Row, Value};
 use spark_connect::session::SparkSession;
 use spark_connect::types::{DataType, StructField};
+use spark_connect::udf::{CommonInlineUserDefinedFunctionExpression, 
PythonUDFPayload};
 
 use crate::catalog::PyCatalog;
 use crate::dataframe::PyDataFrame;
@@ -500,6 +501,35 @@ impl PySparkSession {
         .to_pyerr()
     }
 
+    /// Register a Python UDF by name so it resolves in SQL (used by
+    /// `UDFRegistration.register`). Mirrors the reference 
`client.register_udf`: build a
+    /// `CommonInlineUserDefinedFunction` carrying a `PythonUDF` (cloudpickled 
command +
+    /// output type + eval type) with NO arguments, and send it as the 
`RegisterFunction`
+    /// command so the server's session function registry resolves `name` in 
SQL.
+    #[pyo3(name = "_registerPythonUdf", signature = (name, return_type, 
eval_type, command, python_ver, deterministic=true))]
+    #[allow(non_snake_case)]
+    fn register_python_udf_py(
+        &self,
+        py: Python<'_>,
+        name: &str,
+        return_type: &Bound<'_, PyAny>,
+        eval_type: i32,
+        command: Vec<u8>,
+        python_ver: String,
+        deterministic: bool,
+    ) -> PyResult<()> {
+        let return_data_type = crate::types::py_to_data_type(return_type)?;
+        let payload = PythonUDFPayload::new(return_data_type, eval_type, 
command, python_ver);
+        let udf_expr = CommonInlineUserDefinedFunctionExpression::new(
+            name.to_string(),
+            deterministic,
+            Vec::new(),
+            payload,
+        );
+        py.detach(|| self.session.register_function(udf_expr))
+            .to_pyerr()
+    }
+
     #[pyo3(name = "sessionId")]
     fn session_id_camel(&self) -> String {
         self.session.session_id().to_string()
diff --git a/crates/spark-connect/src/group.rs 
b/crates/spark-connect/src/group.rs
index 1f7ba8c..677da84 100644
--- a/crates/spark-connect/src/group.rs
+++ b/crates/spark-connect/src/group.rs
@@ -111,6 +111,12 @@ impl GroupedData {
         }
     }
 
+    /// Column names of the underlying (pre-grouping) DataFrame. These are 
passed as the
+    /// map/apply UDF's argument columns so the worker's pandas/Arrow input is 
named.
+    pub fn input_columns(&self) -> Result<Vec<String>> {
+        self.dataframe.columns()
+    }
+
     /// Grouping-key expressions for this grouped data.
     fn grouping_expressions(&self) -> Vec<Expression> {
         self.group_cols
@@ -495,6 +501,14 @@ pub struct CoGroupedData {
 }
 
 impl CoGroupedData {
+    /// Column names of both sides (left then right), passed as the cogroup 
UDF's
+    /// argument columns so the worker's two pandas/Arrow inputs are named.
+    pub fn input_columns(&self) -> Result<Vec<String>> {
+        let mut cols = self.left.input_columns()?;
+        cols.extend(self.right.input_columns()?);
+        Ok(cols)
+    }
+
     /// Apply a pandas UDF to each cogroup (`cogroup(...).applyInPandas`).
     ///
     /// `func` is built on the Python side (cloudpickled, eval type
diff --git a/python/pyspark/sql/udf.py b/python/pyspark/sql/udf.py
index 11ce874..e7034e2 100644
--- a/python/pyspark/sql/udf.py
+++ b/python/pyspark/sql/udf.py
@@ -23,10 +23,26 @@ Mirrors pyspark.sql.connect.udf and 
pyspark.sql.functions.udf/pandas_udf.
 import sys
 from typing import Callable, Optional, Any
 
-from pyspark.sql.types import DataType, StringType
+from pyspark.sql.types import DataType, StringType, _parse_datatype_json_string
 from pyspark.serializers import CloudPickleSerializer
 
 
+def _as_concrete_datatype(returnType: Any) -> DataType:
+    """Return a concrete, picklable pyspark ``DataType``.
+
+    A DDL string (e.g. ``"double"``, ``"a int, b string"``) is parsed via the
+    Rust DDL parser and rebuilt into a concrete type through
+    ``_parse_datatype_json_string``. This matters because the cloudpickled UDF
+    command must carry a real ``DataType`` (not a DDL string): the pandas/Arrow
+    worker pairs each result column with its spark type to build the Arrow 
batch,
+    so a bare string breaks it ("not enough values to unpack"). Mirrors how the
+    reference Connect client pickles the *parsed* output type.
+    """
+    if isinstance(returnType, str):
+        return _parse_datatype_json_string(DataType.fromDDL(returnType).json())
+    return returnType
+
+
 class UserDefinedFunction:
     """
     Represents a user-defined function (UDF) that can be called with columns.
@@ -42,7 +58,9 @@ class UserDefinedFunction:
         deterministic: bool = True,
     ):
         self.func = func
-        self.returnType = returnType
+        # Normalize to a concrete, picklable DataType so the cloudpickled 
command
+        # carries a real type (required by the pandas/Arrow worker, not just 
str).
+        self.returnType = _as_concrete_datatype(returnType)
         self.evalType = evalType
         self.name = name or (
             func.__name__ if hasattr(func, "__name__") else "udf"
@@ -50,9 +68,11 @@ class UserDefinedFunction:
         self.deterministic = deterministic
         self.python_ver = f"{sys.version_info.major}.{sys.version_info.minor}"
 
-        # Cloudpickle the function (the command is the pickled (func, 
returnType) tuple)
+        # Cloudpickle the function (the command is the pickled (func, 
returnType) tuple).
+        # Use the normalized concrete DataType, not the raw arg, so the worker 
gets a
+        # real type it can turn into an Arrow schema.
         serializer = CloudPickleSerializer()
-        self.command = serializer.dumps((func, returnType))
+        self.command = serializer.dumps((func, self.returnType))
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
         """
@@ -60,6 +80,12 @@ class UserDefinedFunction:
         Returns a Column representing the UDF call.
         """
         from pyspark import _pyspark
+        from pyspark.sql.functions import col
+
+        # A UDF argument is ColumnOrName: a bare string is a COLUMN NAME (like 
PySpark),
+        # not a string literal. Coerce str -> col(str) before handing to the 
Rust layer
+        # (whose to_column would otherwise treat a str as a literal).
+        cols = [col(a) if isinstance(a, str) else a for a in args]
 
         # pyfunc_make_udf accepts any DataType object or DDL string as the 
return type
         # and builds the proto output type; the pickled command carries (func, 
returnType)
@@ -70,7 +96,7 @@ class UserDefinedFunction:
             self.evalType,
             self.command,
             self.python_ver,
-            *args,
+            *cols,
             deterministic=self.deterministic,
         )
 
@@ -113,6 +139,11 @@ def udf(
     >>> u = F.udf(lambda x: x + 1, IntegerType())
     >>> result = spark.range(3).select(u(F.col('id')).alias('inc_id'))
     """
+    # Decorator form @udf(returnType) / @udf("ddl"): the first positional is 
the
+    # return type (a DataType or DDL string), not the function.
+    if f is not None and not callable(f):
+        returnType, f = f, None
+
     if returnType is None:
         returnType = StringType()
 
@@ -134,7 +165,7 @@ def udf(
 def pandas_udf(
     f: Optional[Callable[..., Any]] = None,
     returnType: Optional[DataType] = None,
-    functionType: str = "scalar",
+    functionType: Optional[str] = None,
 ) -> Any:
     """
     Create a pandas UDF.
@@ -163,6 +194,11 @@ def pandas_udf(
     ...     return s + 1
     >>> result = spark.range(3).select(inc_id(F.col('id')).alias('inc_id'))
     """
+    # Decorator form @pandas_udf(returnType) / @pandas_udf("ddl"): the first
+    # positional is the return type, not the function.
+    if f is not None and not callable(f):
+        returnType, f = f, None
+
     if returnType is None:
         returnType = StringType()
 
@@ -177,10 +213,24 @@ def pandas_udf(
         "cogrouped_map": 206,  # SQL_COGROUPED_MAP_PANDAS_UDF
     }
 
-    evalType = eval_type_map.get(functionType, 200)  # Default to scalar
-
     def _pandas_udf_decorator(func):
-        return UserDefinedFunction(func, returnType, evalType)
+        et = eval_type_map.get(functionType) if functionType is not None else 
None
+        if et is None:
+            # Modern form @pandas_udf("type"): infer the eval type from the 
function's
+            # pandas type hints (Series->Series = scalar, 
Iterator[Series]->Iterator =
+            # scalar_iter, Series->scalar = grouped_agg), mirroring
+            # pyspark.sql.pandas.functions.pandas_udf.
+            try:
+                from inspect import signature
+                from typing import get_type_hints
+                from pyspark.sql.pandas.typehints import infer_eval_type
+
+                et = infer_eval_type(signature(func), get_type_hints(func))
+            except Exception:
+                et = None
+        if et is None:
+            et = 200  # SQL_SCALAR_PANDAS_UDF
+        return UserDefinedFunction(func, returnType, et)
 
     if f is not None:
         # Direct call
@@ -210,6 +260,10 @@ def arrow_udf(
     functionType : str, optional
         "scalar" (default) or "scalar_iter".
     """
+    # Decorator form @arrow_udf(returnType) / @arrow_udf("ddl").
+    if f is not None and not callable(f):
+        returnType, f = f, None
+
     if returnType is None:
         returnType = StringType()
 
@@ -262,11 +316,28 @@ class UDFRegistration:
         >>> spark.udf.register("inc_id", lambda x: x + 1, IntegerType())
         >>> result = spark.sql("SELECT inc_id(id) AS incremented FROM 
range(3)")
         """
-        if returnType is None:
-            returnType = StringType()
-
-        # Create a UDF with the given name
-        udf = UserDefinedFunction(f, returnType, 100, name)
+        # If f is already a UDF (from udf()/pandas_udf()), reuse its packed 
command;
+        # otherwise wrap the plain Python function as a SQL_BATCHED_UDF.
+        if hasattr(f, "command") and hasattr(f, "evalType"):
+            if returnType is not None:
+                from pyspark.errors import PySparkTypeError
+
+                raise PySparkTypeError(
+                    errorClass="CANNOT_SPECIFY_RETURN_TYPE_FOR_UDF",
+                    messageParameters={"arg_name": "f", "return_type": 
str(returnType)},
+                )
+            udf = f
+        else:
+            if returnType is None:
+                returnType = StringType()
+            udf = UserDefinedFunction(f, returnType, 100, name)
+
+        # Send the registration to the Connect server so `name` resolves in SQL
+        # (mirrors the reference client.register_udf: a RegisterFunction 
command
+        # carrying the cloudpickled PythonUDF with no bound arguments).
+        self.spark_session._registerPythonUdf(
+            name, udf.returnType, udf.evalType, udf.command, udf.python_ver, 
udf.deterministic
+        )
         return udf
 
     def registerJavaFunction(
diff --git a/python/tests/test_dropin_offline.py 
b/python/tests/test_dropin_offline.py
index 2283de7..b7402a4 100644
--- a/python/tests/test_dropin_offline.py
+++ b/python/tests/test_dropin_offline.py
@@ -654,13 +654,26 @@ def test_udf_defaults_and_registration():
     # pandas_udf grouped_map/other functionType branches
     for ft in ["scalar", "grouped_map", "grouped_agg"]:
         assert pandas_udf(lambda s: s, T.IntegerType(), functionType=ft) is 
not None
-    # UDFRegistration.register returns a UDF
+    # UDFRegistration.register wires to the session's _registerPythonUdf (so 
the name
+    # resolves in SQL) and returns the UDF.
     from pyspark.sql.udf import UDFRegistration
-    reg = UDFRegistration(object())
+
+    class _StubSession:
+        def __init__(self):
+            self.calls = []
+
+        def _registerPythonUdf(self, name, returnType, evalType, command, 
python_ver, deterministic):
+            self.calls.append((name, evalType, isinstance(command, (bytes, 
bytearray))))
+
+    stub = _StubSession()
+    reg = UDFRegistration(stub)
     r = reg.register("f", lambda x: x + 1, T.IntegerType())
     assert r.name == "f"
     r2 = reg.register("g", lambda x: x)  # returnType default
     assert r2.name == "g"
+    # both were sent to the server with cloudpickled command bytes
+    assert [c[0] for c in stub.calls] == ["f", "g"]
+    assert all(c[2] for c in stub.calls)
 
 
 def test_util_print_exec_and_skip_env(monkeypatch):
@@ -2488,3 +2501,70 @@ def test_generated_fn_accepts_str_column_name():
     assert repr(F.struct("id", "g")) == repr(F.struct(F.col("id"), F.col("g")))
     # literal-position args stay literal (date_format format, regexp pattern)
     assert "yyyy" in repr(F.date_format(F.col("d"), "yyyy"))
+
+
+def test_udf_string_arg_is_column_name_not_literal():
+    # udf(f)("colname") must reference a column (ColumnOrName), not a string 
literal.
+    from pyspark.sql.functions import udf, col
+    from pyspark.sql.types import IntegerType
+    u = udf(lambda x: x + 1, IntegerType())
+    c = u("id")
+    # Same rendered expression whether the arg is a str or an explicit col().
+    assert repr(c) == repr(u(col("id")))
+
+
+def test_udf_and_pandas_udf_decorator_forms():
+    # @udf(returnType) / @pandas_udf(returnType) return a decorator that wraps 
a function.
+    from pyspark.sql.functions import udf, pandas_udf
+    from pyspark.sql.types import IntegerType
+    from pyspark.sql.udf import UserDefinedFunction
+
+    @udf(IntegerType())
+    def inc(x):
+        return x + 1
+    assert isinstance(inc, UserDefinedFunction)
+
+    @pandas_udf("long")
+    def p(s):
+        return s + 1
+    assert isinstance(p, UserDefinedFunction)
+
+
+def test_udf_return_type_normalized_to_concrete_datatype():
+    # A DDL-string returnType must become a concrete, picklable DataType so the
+    # cloudpickled command carries a real type (the pandas/Arrow worker needs 
it).
+    import pickle
+    from pyspark.sql.functions import pandas_udf
+    from pyspark.sql.types import DoubleType
+
+    @pandas_udf("double")
+    def f(s):
+        return s
+    assert isinstance(f.returnType, DoubleType), type(f.returnType)
+    # command is cloudpickle of (func, returnType) -> must not carry a bare 
str.
+    assert isinstance(f.command, (bytes, bytearray))
+    pickle.dumps(f.returnType)  # concrete types pickle
+
+
+def test_pandas_udf_eval_type_inferred_from_type_hints():
+    # @pandas_udf("type") infers eval type from the function's pandas type 
hints:
+    # Series->Series = scalar(200), Iterator->Iterator = scalar_iter(204),
+    # Series->scalar = grouped_agg(202).
+    from typing import Iterator
+    import pandas as pd
+    from pyspark.sql.functions import pandas_udf
+
+    @pandas_udf("double")
+    def scal(s: pd.Series) -> pd.Series:
+        return s
+    assert scal.evalType == 200, scal.evalType
+
+    @pandas_udf("double")
+    def itr(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
+        yield from it
+    assert itr.evalType == 204, itr.evalType
+
+    @pandas_udf("double")
+    def agg(s: pd.Series) -> float:
+        return float(s.mean())
+    assert agg.evalType == 202, agg.evalType


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to