[
https://issues.apache.org/jira/browse/SPARK-60594?page=com.atlassian.jira.plugin.system.issuetabpanels:all-tabpanel
]
ASF GitHub Bot updated SPARK-60594:
-----------------------------------
Labels: correctness pull-request-available (was: correctness)
> 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.0.0, 4.2.0, 4.1.1
> Reporter: Yaoxuan Wu
> Priority: Major
> Labels: correctness, pull-request-available
>
> {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]
