uros-b commented on code in PR #56533:
URL: https://github.com/apache/spark/pull/56533#discussion_r3441310199
##########
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:
Seems a bit orthogonal to the current PR, but is this a real concern
nonetheless? @Yicong-Huang @HyukjinKwon PTAL ^^
--
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]
