[SPARK-61152][CONNECT][PYTHON] Raise INVALID_RETURN_TYPE_FOR_PANDAS_UDF for a non-struct return type of mapInPandas, applyInPandas and the other map functions, in the Python client as PySpark does

crates/pyspark-rs/src/dataframe.rs

@@ -3922,6 +3922,18 @@ pub(crate) fn build_map_udf(
 ) -> PyResult<CommonInlineUserDefinedFunctionExpression> {
     let name = udf_name(func)?;
     let return_type = resolve_datatype(py, session, schema)?;
+    // As PySpark's `MapPartitions`, `GroupMap` and the other map nodes do when they build the
+    // UDF, the return type is checked for the eval type: a map or group function returns a
+    // struct, so another type raises `INVALID_RETURN_TYPE_FOR_PANDAS_UDF` here, not at the action.
+    py.import("pyspark.sql.udf")?
+        .getattr("UserDefinedFunction")?
+        .call_method1(
+            "_check_return_type",
+            (
+                crate::types::data_type_to_py(py, &return_type)?,
+                eval_type_id,
+            ),
+        )?;
     // The worker reads the pickled type as a `DataType`, so a DDL string is pickled as the
     // type it parses to, as the reference client pickles the type of `_parse_ddl`.
     let pickled_type = if schema.is_instance_of::<pyo3::types::PyString>() {

python/tests/test_dropin_e2e.py

@@ -1684,6 +1684,22 @@ def count(key, frames, state):
         session.conf.unset("spark.sql.shuffle.partitions")
 
 
+def test_map_functions_reject_a_non_struct_return_type_before_the_server(session):
+    # As PySpark's, a map or group function whose return type is not a struct raises on the
+    # client when it is built, so the server never sees the plan.
+    from pyspark.errors import PySparkTypeError
+
+    df = session.range(2)
+    with pytest.raises(PySparkTypeError) as e:
+        df.mapInPandas(lambda frames: frames, "long")
+    assert e.value.getCondition() == "INVALID_RETURN_TYPE_FOR_PANDAS_UDF"
+    with pytest.raises(PySparkTypeError) as e:
+        df.groupBy("id").applyInArrow(lambda table: table, "long")
+    assert e.value.getCondition() == "INVALID_RETURN_TYPE_FOR_ARROW_UDF"
+    assert e.value.getMessageParameters()["return_type"] == "LongType()"
+    assert df.mapInPandas(lambda frames: frames, "id long").collect() == [Row(id=0), Row(id=1)]
+
+
 def test_variant_spatial_and_user_defined_values_round_trip_as_pyspark(session):
     # As PySpark's: createDataFrame sends a VariantVal, a Geometry, a Geography and a value of
     # a user-defined type, and collects them as they are.

python/tests/test_dropin_offline.py

@@ -8036,6 +8036,44 @@ def test_schema_cache_cold_children_keep_independent_caches(mock_connect):
     assert df.columns == ["id"]
 
 
+def test_map_and_group_functions_reject_a_non_struct_return_type_as_pyspark(mock_connect):
+    # As PySpark's MapPartitions, GroupMap, CoGroupMap and ApplyInPandasWithState, which check
+    # the return type of their UDF when they are built: a type that is not a struct raises
+    # INVALID_RETURN_TYPE_FOR_PANDAS_UDF, or INVALID_RETURN_TYPE_FOR_ARROW_UDF for an Arrow
+    # function, at the call, before the server sees the plan.
+    from pyspark.errors import PySparkTypeError
+
+    spark, service = mock_connect
+    df = spark.range(1)
+    grouped = df.groupBy("id")
+    cogrouped = grouped.cogroup(df.groupBy("id"))
+    pandas_error = "INVALID_RETURN_TYPE_FOR_PANDAS_UDF"
+    arrow_error = "INVALID_RETURN_TYPE_FOR_ARROW_UDF"
+    cases = [
+        (lambda: df.mapInPandas(lambda it: it, T.IntegerType()), pandas_error, "IntegerType()"),
+        (lambda: df.mapInPandas(lambda it: it, "int"), pandas_error, "IntegerType()"),
+        (lambda: df.mapInArrow(lambda it: it, "long"), pandas_error, "LongType()"),
+        (lambda: grouped.applyInPandas(lambda pdf: pdf, "int"), pandas_error, "IntegerType()"),
+        (lambda: grouped.applyInPandas(lambda pdf: pdf, T.ArrayType(T.IntegerType())),
+         pandas_error, "ArrayType(IntegerType(), True)"),
+        (lambda: grouped.applyInArrow(lambda table: table, T.LongType()), arrow_error, "LongType()"),
+        (lambda: cogrouped.applyInPandas(lambda left, right: left, "int"),
+         pandas_error, "IntegerType()"),
+        (lambda: cogrouped.applyInArrow(lambda left, right: left, "long"),
+         arrow_error, "LongType()"),
+        (lambda: grouped.applyInPandasWithState(
+            lambda key, pdfs, state: iter([]), T.IntegerType(),
+            T.StructType([T.StructField("c", T.LongType())]), "append", "NoTimeout"),
+         pandas_error, "IntegerType()"),
+    ]
+    for call, condition, return_type in cases:
+        with pytest.raises(PySparkTypeError) as e:
+            call()
+        assert e.value.getCondition() == condition
+        assert e.value.getMessageParameters()["return_type"] == return_type
+    assert service.execute_requests == []
+
+
 @pytest.mark.parametrize("transform", [
     pytest.param(lambda df: df.selectExpr("id AS renamed"), id="select"),
     pytest.param(lambda df: df.zipWithIndex(), id="zipWithIndex"),