Yaoxuan Wu created SPARK-60826:
----------------------------------
Summary: Grouped-agg pandas_udf / arrow_udf in an aggregation
without GROUP BY returns no row for an empty input
Key: SPARK-60826
URL: https://issues.apache.org/jira/browse/SPARK-60826
Project: Spark
Issue Type: Bug
Components: PySpark, SQL
Affects Versions: 4.1.1, 4.2.0, 4.0.0
Reporter: Yaoxuan Wu
{code:python}
import pandas as pd
from pyspark.sql import SparkSession, functions as F
from pyspark.sql.functions import pandas_udf
spark = SparkSession.builder.getOrCreate()
@pandas_udf("long")
def cnt(v: pd.Series) -> int:
return int(v.notna().sum())
df = spark.range(3).filter("id > 5") # empty input
df.agg(F.count("id")).show() # one row: 0
df.agg(cnt("id")).show() # expected one row: 0; actual: no row
df.createOrReplaceTempView("t")
spark.udf.register("cnt", cnt)
spark.sql("SELECT cnt(id) FROM t").show() # no row
{code}
An aggregation without grouping keys returns exactly one row, also for an empty
input; Spark's built-in aggregates do, but a grouped-agg pandas UDF (and, in
4.1+, a grouped-agg {{arrow_udf}}) returns no row, in the DataFrame API and in
SQL. Since a pandas aggregate cannot be mixed with built-in aggregates in one
Aggregate, there is no way to get the row back in the same query.
Cause: {{ArrowAggregatePythonExec.doExecute}} ({{AggregateInPandasExec}} in
4.0) skips empty partitions:
{code:scala}
// Map grouped rows to ArrowPythonRunner results, Only execute if partition is
not empty
inputRDD.mapPartitionsInternal { iter => if (iter.isEmpty) iter else {
...
val groupedItr = if (groupingExpressions.isEmpty) {
// Use an empty unsafe row as a place holder for the grouping key
Iterator((new UnsafeRow(), newIter))
{code}
With no grouping expressions the operator requires {{AllTuples}}, so the whole
input is one partition and one group; when that partition is empty it is
returned as is, before the {{groupingExpressions.isEmpty}} branch that creates
the single group. The skip was added by SPARK-28128 ("Pandas Grouped UDFs
should skip over empty partitions"), which is right for grouped aggregation but
not for a global one; {{HashAggregateExec}} / {{SortAggregateExec}} emit the
result of the initial buffer in that case.
Reproduces on 4.0.0, 4.1.1, 4.2.0 and master.
--
This message was sent by Atlassian Jira
(v10.3.26#10030026)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
