Yicong-Huang commented on code in PR #55222:
URL: https://github.com/apache/spark/pull/55222#discussion_r3077722317
##########
python/pyspark/worker.py:
##########
@@ -2833,6 +2797,144 @@ def grouped_func(
# profiling is not supported for UDF
return grouped_func, None, ser, ser
+ if eval_type == PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF:
+ import pyarrow as pa
+
+ assert num_udfs == 1, "One GROUPED_MAP_ARROW UDF expected here."
+ grouped_udf, arg_offsets, return_type, num_udf_args = udfs[0]
+ parsed_offsets = extract_key_value_indexes(arg_offsets)
+
+ arrow_return_type = to_arrow_type(
+ return_type, timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types
+ )
+ if runner_conf.assign_cols_by_name:
+ expected_cols_and_types = {
+ col.name: to_arrow_type(col.dataType, timezone="UTC") for col
in return_type.fields
+ }
+ else:
+ expected_cols_and_types = [
+ (col.name, to_arrow_type(col.dataType, timezone="UTC"))
+ for col in return_type.fields
+ ]
+
Review Comment:
assertion added! they should be 1.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
