FelixYBW commented on PR #59313:
URL: https://github.com/apache/spark/pull/59313#issuecomment-6208486344
@HyukjinKwon I'm comparing the implementation vs. what we have done in
Gluten. Claude reported a bug when it work with a pyarrow UDF. FYI, @viirya
``` python
import pyarrow as pa
from pyspark.sql import SparkSession
from pyspark.sql.functions import arrow_udf, col, assert_true, concat_ws, lit
from pyspark.sql.types import LongType
@arrow_udf(LongType())
def ident(v: pa.Array) -> pa.Array:
return v
spark = (SparkSession.builder.master("local[1]")
.config("spark.driver.extraLibraryPath", "/tmp/nds")
.getOrCreate())
def df(end, batch_size):
return (spark.read.format("native_range")
.option("end", str(end)).option("partitions", "1")
.option("batch_size", str(batch_size)).load()
.select("id")
.withColumn("uid", ident(col("id"))))
# Case A: crash. Small input: the reader is closed while outputs are still
pending.
try:
df(200, 100).collect()
except Exception as e:
print("A:", "IndexOutOfBoundsException" in str(e))
# Case B: wrong results. Large input: the reused vectors are overwritten by
read-ahead.
check = df(2_000_000, 1000).select(
assert_true(col("id") == col("uid"),
concat_ws(" ", lit("MISMATCH id"), col("id"), lit("uid"),
col("uid"))))
try:
check.collect() # use collect(), not count(): count() prunes
the check away
except Exception as e:
print("B:", next(l for l in str(e).splitlines() if "MISMATCH" in l))
# Control: same query on the row path, correct
spark.conf.set("spark.sql.execution.arrow.pythonUDF.columnarInput.enabled",
"false")
print("control mismatches:", sum(r.id != r.uid for r in df(200,
100).collect()))
```
reproduce:
A: True
B: ... MISMATCH id 1383000 uid 0 ... (the exact id depends on how far the
runner read ahead)
control mismatches: 0
There are two separate causes:
1. Vector reuse gives wrong results. The Python runner read about 1,400
input batches ahead before returning the first result. Every batch is imported
into the reader's single set of vectors. So the pass-through columns waiting
for UDF results show the latest batch, not their own: result 0 got paired with
id 1,383,000.
2. Early close crashes the task. When reader.next() returns false,
DataSourceRDD's PartitionIterator.hasNext closes the reader immediately.
NativePartitionReader.close() then calls root.close(), which empties the
vectors that pending output batches still point to. A partition with a single
batch is enough to trigger it.
I'd suggest to use ColumnarBatch types in
[SPARK-57468](https://issues.apache.org/jira/browse/SPARK-57468) to solve the
batch type conversions elegantly.
--
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]
