Copilot commented on code in PR #56533:
URL: https://github.com/apache/spark/pull/56533#discussion_r3417447813
##########
python/pyspark/pandas/strings.py:
##########
@@ -1174,14 +1175,25 @@ def findall(self, pat: str, flags: int = 0) ->
"ps.Series":
2 [b, b]
dtype: object
"""
+ num_groups = re.compile(pat, flags=flags).groups
str_dtype = is_str_dtype(self._data.dtype)
+ if num_groups > 1:
+ return_type = ArrayType(
+ ArrayType(StringType(), containsNull=True), containsNull=True
+ )
+ else:
+ return_type = ArrayType(StringType(), containsNull=True)
# type hint does not support to specify array type yet.
- @pandas_udf( # type: ignore[call-overload]
- returnType=ArrayType(StringType(), containsNull=True)
- )
+ @pandas_udf(returnType=return_type) # type: ignore[call-overload]
def pudf(s: pd.Series) -> pd.Series:
ret = s.str.findall(pat, flags)
+ if num_groups > 1:
+ ret = ret.map(
+ lambda matches: [list(match) for match in matches]
+ if isinstance(matches, list)
+ else matches
+ )
if str_dtype:
# ArrayType does not support NaN, so replace with None
ret = ret.replace(np.nan, None)
Review Comment:
`returnType` is always an `ArrayType(...)` (or nested
`ArrayType(ArrayType(...))`) here, but `np.nan` values can still flow through
when `str_dtype` is false (and also via the `else matches` branch). `np.nan` is
a float sentinel and typically cannot be serialized into Spark `ArrayType`,
which can cause Arrow/Spark conversion errors. Consider normalizing `np.nan` to
`None` unconditionally for this UDF (or at minimum, ensuring non-list nullish
values become `None` when `num_groups > 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]
