Repository navigation
[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 - #341
Closed
HyukjinKwon wants to merge 1 commit into
Closed
HyukjinKwon wants to merge 1 commit into
HyukjinKwon wants to merge 1 commit into
Conversation
…DF for a non-struct return type of mapInPandas, applyInPandas and the other map functions, in the Python client as PySpark does mapInPandas, mapInArrow, applyInPandas, applyInArrow, applyInPandasWithState and the cogrouped applyInPandas and applyInArrow sent a plan whose UDF returns a type that is not a struct, and the server failed with a ClassCastException or an IllegalArgumentException. As PySpark, the UDF build now checks its return type with UserDefinedFunction._check_return_type, which raises INVALID_RETURN_TYPE_FOR_PANDAS_UDF or INVALID_RETURN_TYPE_FOR_ARROW_UDF on the client, at the call. Co-authored-by: Isaac <no-reply@databricks.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
SPARK-61152: Raise
INVALID_RETURN_TYPE_FOR_PANDAS_UDFfor a non-struct return type ofmapInPandas,applyInPandasand the other map functions, in the Python client as PySpark does.PySpark builds the UDF of
mapInPandas,mapInArrow,applyInPandas,applyInArrow,applyInPandasWithStateand the cogroupedapplyInPandasandapplyInArrowthroughUserDefinedFunction, which checks the return type for the eval type. A return type that is not a struct raisesPySparkTypeErrorINVALID_RETURN_TYPE_FOR_PANDAS_UDF, orINVALID_RETURN_TYPE_FOR_ARROW_UDFfor the Arrow functions, at the call, and the plan is never sent.The drop-in already has this check verbatim in
pyspark/sql/udf.py.build_map_udfnow callsUserDefinedFunction._check_return_typewith the resolved return type and the eval type, right after the schema is resolved, as the official client does. The check is a no-op for the transform-with-state eval types, which build withbuild_map_udftoo.Why are the changes needed?
df.mapInPandas(f, "int"),df.mapInPandas(f, IntegerType()),df.mapInArrow(f, "long"),df.groupBy("id").applyInPandas(g, "int")anddf.groupBy("id").applyInArrow(g, LongType())were sent to the server. The server failed withUnknownException(ClassCastException), orIllegalArgumentException: The returnType of the udf must be a struct. PySpark raisesPySparkTypeErrorINVALID_RETURN_TYPE_FOR_PANDAS_UDF(for examplePandas UDF should return StructType for SQL_MAP_PANDAS_ITER_UDF or SQL_MAP_ARROW_ITER_UDF, got IntegerType().) on the client.Does this PR introduce any user-facing change?
Yes. A map or group function with a non-struct return type raises the PySpark
PySparkTypeErroron the client when the call builds the plan, with the same error class, message andreturn_typeas PySpark, instead of a server error.How was this patch tested?
test_map_and_group_functions_reject_a_non_struct_return_type_as_pyspark(offline, mock server): the nine call forms (mapInPandas with a type and a DDL string, mapInArrow, applyInPandas with a DDL string and an ArrayType, applyInArrow, the two cogrouped applies, applyInPandasWithState) raise the expected condition andreturn_typeat the call, and no request is sent. The same calls give the same condition, message and return type with official PySpark 4.3.0.test_map_functions_reject_a_non_struct_return_type_before_the_server(live): mapInPandas and applyInArrow with a non-struct type raise on the client, and a struct schema still runs.wasm-udf, and the offline Python suites (2202 passed, 8 skipped), and the live tests of the map, grouped, cogrouped and stateful functions (22 passed) against a local Spark 4.3.0 server; rustfmt passes, and clippy reports no diagnostics in the changed lines.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5.5
This pull request and its description were written by Isaac.