Yaoxuan Wu created SPARK-60594:
----------------------------------
Summary: Python UDTF with TABLE(t) PARTITION BY k puts every row
with a NaN key into its own partition
Key: SPARK-60594
URL: https://issues.apache.org/jira/browse/SPARK-60594
Project: Spark
Issue Type: Bug
Components: PySpark
Affects Versions: 4.1.1, 4.2.0, 4.0.0
Reporter: Yaoxuan Wu
{code:python}
from pyspark.sql.functions import udtf
@udtf(returnType="n: int")
class CountRows:
def __init__(self):
self.n = 0
def eval(self, row):
self.n += 1
def terminate(self):
yield (self.n,)
spark.udtf.register("count_rows", CountRows)
df = spark.createDataFrame([(float("nan"), 1), (float("nan"), 2), (1.0, 3),
(1.0, 4)], "k double, v int")
df.createOrReplaceTempView("t")
df.groupBy("k").count().collect() #
(1.0, 2), (NaN, 2)
spark.sql("SELECT * FROM count_rows(TABLE(t) PARTITION BY k)").collect()
# expected: [Row(n=2), Row(n=2)]
# actual: [Row(n=2), Row(n=1), Row(n=1)]
{code}
Spark groups NaN keys together everywhere else (GROUP BY, window PARTITION BY),
and the plan for the table argument does put the two NaN rows next to each
other in one task. The same happens with {{asTable().partitionBy("k")}}, Arrow
UDTFs and {{@arrow_udtf}}, for keys that contain a NaN (arrays, structs), and
for collated strings ({{'a'}} / {{'A'}} with UTF8_LCASE become two partitions).
Cause: the Python worker finds the partition boundaries by comparing adjacent
keys with Python {{!=}}: {{UDTFWithPartitions._check_partition_boundaries}} in
{{python/pyspark/worker.py}} ({{any(k != v for k, v in zip(cur_partitions_args,
prev_partitions_args))}}), and the {{@arrow_udtf}} path
({{col_array[row_idx].as_py() != col_array[row_idx - 1].as_py()}}). In Python
{{nan != nan}} is True. On master this code is in
{{python/pyspark/eval_handlers/_udtf.py}}, unchanged.
Reproduces on 4.0.0, 4.1.1 and 4.2.0.
--
This message was sent by Atlassian Jira
(v10.3.26#10030026)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
