Hyukjin Kwon created SPARK-61141:
------------------------------------
Summary: Run Rust WASM UDFs once per Arrow batch in the Spark
Connect Rust client
Key: SPARK-61141
URL: https://issues.apache.org/jira/browse/SPARK-61141
Project: Spark
Issue Type: Improvement
Components: Connect
Affects Versions: connect-rust-4.3.0
Reporter: Hyukjin Kwon
Spark's Python worker calls a SQL_ARROW_BATCHED_UDF once per row. For a Rust
WASM UDF of the Rust client (#[spark_wasm_udf]), each call encodes the row's
arguments, enters the WASM module and decodes the result, which costs about 100
microseconds per row.
A #[spark_wasm_udf] function whose arguments and result are all i32, i64, f32,
f64, bool, String, Vec<u8>, or one Option of those, can instead be sent as a
scalar Arrow UDF (SQL_SCALAR_ARROW_UDF) and called once per Arrow batch:
- the macro also emits a batch export (spark_batch_udf_<fn>), and its builder
calls UserDefinedFunction::with_batch_entrypoint
- the runtime of the WASM build reads the Arrow buffers of the argument columns
and writes those of the result
- the packer ships a WasmArrowUDF runner that copies the buffers in, calls the
batch export once, and builds the result array
Measured against a local Spark 4.3.0 Connect server, this makes such UDFs cost
about as much as a Python arrow_udf doing the same work with pyarrow.compute:
add_one(i64) goes from about 100,000 ns to about 124 ns per row. Signatures
with arrays or nested Options stay per row, and
with_eval_type(SQL_ARROW_BATCHED_UDF) keeps a UDF per row.
--
This message was sent by Atlassian Jira
(v10.3.26#10030026)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
