[SPARK-43543][PYTHON] Fix nested MapType behavior in Pandas UDF

[xinrong-meng]

What changes were proposed in this pull request?

Fix nested MapType behavior in Pandas UDF (and Arrow-optimized Python UDF).

Previously during Arrow-pandas conversion, only the outermost layer is converted to a dictionary; but now nested MapType will be converted to nested dictionaries.

That applies to Spark Connect as well.

Why are the changes needed?

Correctness and consistency (with createDataFrame and toPandas when Arrow is enabled).

Does this PR introduce any user-facing change?

Yes.

Nested MapType type support is corrected in Pandas UDF

>>> schema = StructType([
...      StructField("id", StringType(), True),
...      StructField("attributes", MapType(StringType(), MapType(StringType(), StringType())), True)
... ])
>>> 
>>> data = [
...    ("1", {"personal": {"name": "John", "city": "New York"}}),
... ]
>>> df = spark.createDataFrame(data, schema)
>>> @pandas_udf(StringType())
... def f(s: pd.Series) -> pd.Series:
...    return s.astype(str)
... 
>>> df.select(f(df.attributes)).show(truncate=False)

The results of df.select(f(df.attributes)).show(truncate=False) is corrected

FROM

+------------------------------------------------------+                        
|f(attributes)                                         |
+------------------------------------------------------+
|{'personal': [('name', 'John'), ('city', 'New York')]}|
+------------------------------------------------------+

TO

>>> df.select(f(df.attributes)).show(truncate=False)
+--------------------------------------------------+
|f(attributes)                                     |
+--------------------------------------------------+
|{'personal': {'name': 'John', 'city': 'New York'}}|
+--------------------------------------------------+

Another more obvious example:

>>> @pandas_udf(StringType())
... def extract_name(s:pd.Series) -> pd.Series:
...     return s.apply(lambda x: x['personal']['name'])
...
>>> df.select(extract_name(df.attributes)).show(truncate=False)

df.select(extract_name(df.attributes)).show(truncate=False) is corrected

FROM

org.apache.spark.api.python.PythonException: Traceback (most recent call last):
...
TypeError: list indices must be integers or slices, not str

TO

+------------------------+
|extract_name(attributes)|
+------------------------+
|John                    |
+------------------------+

How was this patch tested?

Unit tests.

[xinrong-meng]

Inherit _create_array of ArrowStreamPandasSerializer. After the change, it is consistent with createDataFrame from a pandas DataFrame when Arrow is enabled.

[dongjoon-hyun]

Could you file a JIRA issue officially and make this IDed TODO like TODO(SPARK-XXX)?

[xinrong-meng]

Certainly, done!

[xinrong-meng]

@ueshin @HyukjinKwon @zhengruifeng would you please review?

[xinrong-meng]

Merged to master, thank you!

[xinrong-meng]

Please free to leave comments if any, I'll adjust them in follow-ups.