tdcmeehan opened a new pull request, #262:
URL: https://github.com/apache/spark-connect-rust/pull/262
### What changes were proposed in this pull request?
`#[spark_wasm_udf]` functions whose arguments and result are all `i32`,
`i64`, `f32`, `f64`, `bool`, `String`, `Vec<u8>`, or one `Option` of those are
now called once per Arrow batch instead of once per row.
- **Macro** (`spark-connect-macros`): for such a function it also emits a
`spark_batch_udf_<fn>` export, and the generated `udf::<fn>_udf()` builder
calls `.with_batch_entrypoint(...)`. Other signatures (arrays, nested
`Option`s) are unchanged. The batch export names use a prefix of their own, so
they cannot collide with the per-row `spark_udf_<fn>` exports.
- **Build runtime** (`spark-connect-build`): adds the batch ABI to the
runtime compiled into every module. The host writes a descriptor and the Arrow
buffers of each argument column (validity, values, offsets) into a scratch
buffer (`spark_batch_scratch`). The export evaluates every row and returns a
descriptor of the result's Arrow buffers, or an error message.
- **Packer** (`wasm_packer.py`): adds the `WasmArrowUDF` runner. It casts
each argument to the declared Arrow type, copies the buffers in, calls the
batch export once, and builds the result array from the returned buffers. Each
Python worker process keeps one wasmtime engine and compiles each module once,
keyed by its hash. Each runner gets its own wasmtime instance, which is
recreated after a trap.
- **Client** (`wasm_udf.rs`): adds
`UserDefinedFunction::with_batch_entrypoint`, which sets the eval type to
`SQL_SCALAR_ARROW_UDF` (250). The eval type decides which runner is packed: the
batch entrypoint goes into the UDF only under eval type 250. Eval type 250
without a batch entrypoint, or with zero arguments, is an error on the client.
- Tests in a new `tests/wasm-udf` crate, excluded from the workspace like
the WASM UDF examples, plus docs and CI.
### Why are the changes needed?
Spark's Python worker calls a `SQL_ARROW_BATCHED_UDF` once per row. For a
WASM UDF, each call encodes the row's arguments, enters the module, and decodes
the result, which costs about 100 microseconds per row in the measurements
below. With one call per batch, these UDFs cost about the same as a Python
`arrow_udf` doing the same work with `pyarrow.compute`.
Measured end to end against a local Spark 4.3.0 Connect server (`--master
local[1]`, default 10,000-row Arrow batches), in ns per row:
| UDF | built-in expression | per row (today) | per batch (this PR) | Python
`arrow_udf` (`pyarrow.compute`) |
|---|---:|---:|---:|---:|
| `add_one (i64)` | 1.2 / 1.6 | 100,235 | 123.5 | 131.4 |
| `shout (String)` | 109.2 | 103,467 | 301.9 | 304.2 |
- **Method:** each case runs `range(n, 1 partition).select(f(id)).agg(sum)`,
with `sum(length(...))` for the string cases. The cost per row is the
difference between the run times of two row counts divided by the difference in
rows, so query, worker and module start-up costs cancel out. Each row count
takes the best of 3 runs after one warm-up run.
- **Built-ins:** `id + 1` and `concat(upper(cast(id as string)), '!')`.
- **Runs:** the figures are from the first of two consecutive runs. In the
second run every figure was within 15%, except the built-in `id + 1` (a 1-2 ns
figure), so both of its runs are shown. Second run: 102,251 / 125.9 / 127.3 for
`add_one`, and 101.2 / 97,310 / 303.5 / 294.6 for `shout`.
- **Machine:** Intel Xeon Platinum 8559C, 4 cores / 8 threads. OpenJDK
21.0.12.1, Python 3.12.3, pyspark 4.3.0, pyarrow 26.0.0, wasmtime 50.0.1, rustc
1.99.1.
<details>
<summary>How to reproduce (the bench and the script are not part of this
PR)</summary>
Start a Spark 4.3.0 Connect server whose Python workers have `pyarrow` and
`wasmtime`, for example from a `pyspark==4.3.0` install:
```bash
PYSPARK_PYTHON=<python with pyarrow and wasmtime> \
"$SPARK_HOME/sbin/start-connect-server.sh" --master 'local[1]'
```
Put the three bench files in `examples/wasm-udf-bench/`, add
`"examples/wasm-udf-bench"` to the root `Cargo.toml` `exclude` list, and run
the bench and then the `arrow_udf` script (needs `pyspark[connect]==4.3.0` and
`pyarrow`):
```bash
rustup target add wasm32-unknown-unknown
SPARK_CONNECT_PYTHON=<python with pyspark> \
cargo run --release --manifest-path examples/wasm-udf-bench/Cargo.toml
python arrow_udf_bench.py
```
`examples/wasm-udf-bench/Cargo.toml`:
```toml
[workspace]
[package]
name = "wasm-udf-bench"
version = "0.0.0"
edition = "2021"
publish = false
[dependencies]
spark-connect-macros = { package = "apache-spark-connect-macros", path =
"../../crates/spark-connect-macros" }
[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
spark-connect = { package = "apache-spark-connect", path =
"../../crates/spark-connect", features = ["wasm-udf"] }
[build-dependencies]
spark-connect-build = { package = "apache-spark-connect-build", path =
"../../crates/spark-connect-build" }
spark-connect-macros = { package = "apache-spark-connect-macros", path =
"../../crates/spark-connect-macros" }
```
`examples/wasm-udf-bench/build.rs`:
```rust
fn main() {
// Compiles this crate's `#[spark_wasm_udf]` functions to wasm32 and
embeds
// the module (exposed to the crate as the `WASM_UDFS_MODULE` env var).
spark_connect_build::embed_wasm_udf("src/main.rs");
}
```
`examples/wasm-udf-bench/src/main.rs`:
```rust
//! Rust UDFs per row vs per Arrow batch, end to end through a Spark Connect
server.
//!
//! Each case runs `range(n, 1 partition).select(f(id)).agg(sum)`. The cost
per row is the
//! difference between two row counts divided by the difference in rows, so
query, worker and
//! module start-up costs cancel out. Best of 3 per row count, after one
warm-up run.
//!
//! rustup target add wasm32-unknown-unknown
//! SPARK_CONNECT_PYTHON=<python with pyspark> \
//! cargo run --release --manifest-path
examples/wasm-udf-bench/Cargo.toml
use spark_connect_macros::spark_wasm_udf;
#[spark_wasm_udf]
mod udfs {
pub fn add_one(x: i64) -> i64 {
x + 1
}
pub fn shout(s: String) -> String {
let mut t = s.to_uppercase();
t.push('!');
t
}
pub fn half_if_even(x: Option<i64>) -> Option<i64> {
x.filter(|v| v % 2 == 0).map(|v| v / 2)
}
}
#[cfg(not(target_arch = "wasm32"))]
fn main() -> Result<(), Box<dyn std::error::Error>> {
use spark_connect::functions::{col, expr, length, sum};
use spark_connect::row::Value;
use spark_connect::udf::eval_type;
use spark_connect::{Column, SparkSession, SparkSessionBuilder};
use std::time::Instant;
type BoxErr = Box<dyn std::error::Error>;
let spark: SparkSession = SparkSessionBuilder::default()
.remote("sc://127.0.0.1:15002/")
.get_or_create()?;
// Seconds to collect sum(f(id)) (or sum(length(f(id)))) over n ids in
one partition.
let run = |f: &dyn Fn() -> Result<Column, BoxErr>,
sum_len: bool,
n: i64|
-> Result<(f64, i64), BoxErr> {
let total = if sum_len {
sum(length(col("x")))
} else {
sum(col("x"))
};
let df = spark
.range_full(0, n, 1, Some(1))?
.select([f()?.alias("x")])
.agg(vec![total.expression().clone()]);
let start = Instant::now();
let rows = df.collect()?;
let secs = start.elapsed().as_secs_f64();
match rows[0].get(0) {
Some(Value::Long(v)) => Ok((secs, *v)),
other => Err(format!("unexpected result {other:?}").into()),
}
};
let ns_per_row = |name: &str,
f: &dyn Fn() -> Result<Column, BoxErr>,
sum_len: bool,
expected: &dyn Fn(i64) -> i64,
n1: i64,
n2: i64|
-> Result<(), BoxErr> {
run(f, sum_len, n1)?;
let mut best = [f64::MAX; 2];
for (k, n) in [n1, n2].into_iter().enumerate() {
for _ in 0..3 {
let (secs, v) = run(f, sum_len, n)?;
if v != expected(n) {
return Err(
format!("{name}: sum {v} for {n} rows, expected {}",
expected(n)).into(),
);
}
best[k] = best[k].min(secs);
}
}
let ns = (best[1] - best[0]) / (n2 - n1) as f64 * 1e9;
println!(
"{name:22} {ns:10.1} ns/row ({n1:>10} rows {:7.3} s, {n2:>10}
rows {:7.3} s)",
best[0], best[1]
);
Ok(())
};
let add_one_sum = |n: i64| n * (n - 1) / 2 + n;
let shout_len = |n: i64| (0..n).map(|i| i.to_string().len() as i64 +
1).sum::<i64>();
let half_sum = |n: i64| (0..n).filter(|i| i % 2 == 0).map(|i| i /
2).sum::<i64>();
let id_str = || col("id").cast_str("string");
// Same answers on both paths, including null results.
let n = 25_000;
let (_, per_row) = run(
&|| {
Ok(udf::half_if_even_udf()
.with_eval_type(eval_type::SQL_ARROW_BATCHED_UDF)
.call(vec![col("id")])?)
},
false,
n,
)?;
let (_, batch) = run(&|| Ok(udf::half_if_even(col("id"))?), false, n)?;
assert_eq!((per_row, batch), (half_sum(n), half_sum(n)), "half_if_even");
ns_per_row(
"builtin id + 1",
&|| Ok(expr("id + 1")),
false,
&add_one_sum,
2_000_000,
20_000_000,
)?;
ns_per_row(
"add_one per row",
&|| {
Ok(udf::add_one_udf()
.with_eval_type(eval_type::SQL_ARROW_BATCHED_UDF)
.call(vec![col("id")])?)
},
false,
&add_one_sum,
10_000,
50_000,
)?;
ns_per_row(
"add_one per batch",
&|| Ok(udf::add_one(col("id"))?),
false,
&add_one_sum,
2_000_000,
20_000_000,
)?;
ns_per_row(
"builtin shout",
&|| Ok(expr("concat(upper(cast(id as string)), '!')")),
true,
&shout_len,
1_000_000,
5_000_000,
)?;
ns_per_row(
"shout per row",
&|| {
Ok(udf::shout_udf()
.with_eval_type(eval_type::SQL_ARROW_BATCHED_UDF)
.call(vec![id_str()])?)
},
true,
&shout_len,
10_000,
50_000,
)?;
ns_per_row(
"shout per batch",
&|| Ok(udf::shout(id_str())?),
true,
&shout_len,
1_000_000,
5_000_000,
)?;
Ok(())
}
```
`arrow_udf_bench.py`:
```python
"""Python arrow_udf baselines for the WASM UDF bench, measured the same way.
Each case runs `range(n, 1 partition).select(f(id)).agg(sum)`. The cost per
row is the
difference between two row counts divided by the difference in rows, so
query, worker and
start-up costs cancel out. Best of 3 per row count, after one warm-up run.
python arrow_udf_bench.py # Spark Connect server on
sc://localhost:15002
"""
import time
import pyarrow as pa
import pyarrow.compute as pc
from pyspark.sql import SparkSession
from pyspark.sql.functions import arrow_udf, col, length
from pyspark.sql.functions import sum as sum_
spark = SparkSession.builder.remote("sc://localhost:15002").getOrCreate()
@arrow_udf("long")
def add_one(x: pa.Array) -> pa.Array:
return pc.add(x, 1)
@arrow_udf("string")
def shout(s: pa.Array) -> pa.Array:
return pc.binary_join_element_wise(pc.utf8_upper(s), "!", "")
def run(f, sum_len, n):
df = spark.range(0, n, 1, 1).select(f().alias("x"))
df = df.agg(sum_(length(col("x"))) if sum_len else sum_(col("x")))
start = time.perf_counter()
value = df.collect()[0][0]
return time.perf_counter() - start, value
def ns_per_row(name, f, sum_len, expected, n1, n2):
run(f, sum_len, n1)
best = [float("inf")] * 2
for k, n in enumerate((n1, n2)):
for _ in range(3):
secs, v = run(f, sum_len, n)
assert v == expected(n), f"{name}: sum {v} for {n} rows,
expected {expected(n)}"
best[k] = min(best[k], secs)
ns = (best[1] - best[0]) / (n2 - n1) * 1e9
print(
f"{name:22} {ns:10.1f} ns/row ({n1:>10} rows {best[0]:7.3f} s, "
f"{n2:>10} rows {best[1]:7.3f} s)"
)
def add_one_sum(n):
return n * (n - 1) // 2 + n
def shout_len(n):
return sum(len(str(i)) + 1 for i in range(n))
ns_per_row(
"add_one arrow_udf", lambda: add_one(col("id")), False, add_one_sum,
2_000_000, 20_000_000
)
ns_per_row(
"shout arrow_udf",
lambda: shout(col("id").cast("string")),
True,
shout_len,
1_000_000,
5_000_000,
)
```
</details>
### Does this PR introduce _any_ user-facing change?
Yes, compared with released versions:
- **Batch path:** a `#[spark_wasm_udf]` function with a batch-capable
signature (above) is now sent as a scalar Arrow UDF (`SQL_SCALAR_ARROW_UDF`)
and called once per Arrow batch instead of once per row. Results are identical,
including nulls in and out.
- **Per-row path kept:** functions with arrays or nested `Option`s, and UDFs
built by hand with `wasm_udf::udf(...)`, stay per row. To force a batch-capable
UDF per row, use
`udf::<fn>_udf().with_eval_type(eval_type::SQL_ARROW_BATCHED_UDF)`.
- **New API:** `UserDefinedFunction::with_batch_entrypoint(entrypoint)`.
- **Spark versions:** eval type 250 needs Spark 4.1 or later. The client
already requires Spark 4.3.0 or later.
- **Mistyped inputs:** on the batch path, each argument column is converted
to the declared type with a safe Arrow cast, because Spark passes a column with
its own type, e.g. a `LongType` column into an `i32` argument.
- Values that convert exactly (wider integers in range, whole floats) work.
- An out-of-range integer, or a float with a fractional part into an
integer argument, now raises `ArrowInvalid`, e.g. `Float value 1.500000 was
truncated converting to int32`.
- On the per-row path an out-of-range integer raised `struct.error`, and a
fractional float was truncated with `int()`.
- **New error messages on the batch path:**
- A null in an argument that is not an `Option`: ``WASM UDF
'spark_batch_udf_<fn>': argument `x` is null in row 1 of the batch, but its
type is not an Option; declare it as Option<T> to accept nulls``.
- A module whose batch export takes a different number of columns than the
UDF declares (only possible with a hand-built `udf(...)`): `WASM UDF
'<entrypoint>': the batch has 1 columns, but the function takes 2`.
- Arguments of different lengths: `WASM UDF '<entrypoint>' got arguments
of different lengths: [3, 2]`.
- A batch whose input buffers exceed 2 GiB: `WASM UDF '<entrypoint>' got a
batch of N bytes, more than the 2147483647 a WASM module can take; lower
spark.sql.execution.arrow.maxRecordsPerBatch`.
- On the client:
- `WASM UDF '<name>' has eval type SQL_SCALAR_ARROW_UDF, which needs a
batch entrypoint; set one with with_batch_entrypoint`.
- `WASM UDF '<name>' has no arguments, so it cannot be called once per
batch`.
- A hand-built `udf(...)` given a batch entrypoint and a type the batch
path can't carry: `WASM UDF '<name>' has type 'array:i64', so it cannot be
called once per batch; call it per row with
with_eval_type(SQL_ARROW_BATCHED_UDF)`.
### How was this patch tested?
Added tests:
- `tests/wasm-udf` (new; fixture UDFs compiled with the macro and runtime).
These tests drive the real runners with `pyarrow` arrays, the way Spark's
worker does:
- Only batch-capable signatures get the batch runner, and
`with_eval_type(SQL_ARROW_BATCHED_UDF)` forces the per-row one.
- For every supported type, the batch runner returns the same values as
the per-row runner and as the expected values: nulls in and out, two arguments
of mixed types, a batch of only empty strings, and empty batches. All runners
of a module share one compiled module.
- Sliced input (including bitmap bit offsets), chunked input,
`large_string`/`large_binary`, and wider integer columns are accepted; an
overflow raises.
- One runner is called with 2, then 50,000, then 3 rows. The 50,000-row
batch grows the scratch buffer and the module's memory during the call.
- Errors name their cause, and the next batch succeeds. Covered: a null in
a non-`Option` argument, including one whose parameter is `_` and so is named
by position (`` `#1` ``), both column-count mismatches, arguments of different
lengths (including an empty one), a batch export the module lacks, a batch over
2 GiB, a missing `wasmtime` (which raises the install hint), and a panic in
user code that leaves a `RefCell` borrowed.
- Every runner test also checks that the Python process exits without
`Exception ignored` tracebacks on stderr. Without the exit hooks that drop the
cached wasmtime engine and modules, those objects print these when a process
that used the batch runner exits normally.
- The runner tests, including the clean-exit check, also pass with
wasmtime 17.0.0.
- The fixture includes functions named `scratch` and `batch_scratch`, so a
colliding export name fails the build.
- `tests/wasm-udf/tests/e2e_wasm_udf.rs`: against a live server, covers a
batch UDF with nulls, a `String` UDF, a two-argument UDF, an array UDF that
stays per row, and `spark.udf().register(...)` followed by a SQL call.
- `spark-connect-macros`: unit test of which types the batch path accepts.
- `wasm_udf.rs`: unit tests that the eval type decides whether the spec
carries the batch entrypoint, and that eval type 250 is rejected without a
batch entrypoint, without arguments, or with a type the batch path can't carry
(an array argument, an `Option` of an array, a nested `Option` return).
- CI: the `wasm-udf` job also runs the macro unit tests and the runner
tests. The `connect-e2e` job adds the wasm32 target, `wasmtime` and the
packer's `pyspark`, and runs the e2e test.
Ran locally (Spark 4.3.0 Connect server, Java 21, Python 3.12 with pyspark
4.3.0, pyarrow 26.0.0, wasmtime 50.0.1):
```bash
cargo test -p apache-spark-connect --features wasm-udf
# 703 passed, 27 ignored
cargo test -p apache-spark-connect --features wasm-udf -- --ignored
packer_produces # 1 passed
cargo test -p apache-spark-connect-macros
# 1 passed
cargo test --manifest-path tests/wasm-udf/Cargo.toml -- --ignored
# 4 passed
SPARK_REMOTE=sc://localhost:15002 cargo test --manifest-path
tests/wasm-udf/Cargo.toml --test e2e_wasm_udf # 1 passed
cargo test -p apache-spark-connect --release --lib
# 461 passed
cargo test -p apache-spark-connect --release --no-run
cargo test -p apache-spark-connect --release --features wasm-udf --no-run
cargo fmt --all --check
cargo fmt --manifest-path tests/wasm-udf/Cargo.toml --check
cargo build -p apache-spark-connect-macros -p apache-spark-connect-build
cargo build --manifest-path examples/wasm-udfs/Cargo.toml
cargo build --manifest-path examples/wasm-udf-inline/Cargo.toml
cargo build -p examples --bin wasm_udf_macro --features
examples/wasm-udf-macro
python scripts/check_doc_snippets.py
```
`cargo clippy --workspace --all-targets` reports the same diagnostics as on
master, and none in the changed files.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5.5
--
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]
