diff --git a/.claude/skills/optimize-comet-expression/SKILL.md b/.claude/skills/optimize-comet-expression/SKILL.md new file mode 100644 index 00000000000..a7d8af89a56 --- /dev/null +++ b/.claude/skills/optimize-comet-expression/SKILL.md @@ -0,0 +1,111 @@ +--- +name: optimize-comet-expression +description: Use when optimizing the performance of an existing native scalar expression in the datafusion-comet-spark-expr crate (native/spark-expr/) — casts, string/JSON/array/math kernels that run per-row or per-batch. Covers benchmarking, keeping output bit-identical, and the no-regression gate. Not for adding new expressions (use implement-comet-expression) or wiring upstream functions (use wire-datafusion-function). +argument-hint: +--- + + + +Optimize the native `$ARGUMENTS` scalar expression in `native/spark-expr/`. + +**The methodology, benchmark shapes, technique catalog, correctness traps, and PR conventions +live in one place for both humans and agents: +`docs/source/contributor-guide/optimizing_expressions.md`. Read it first — it is the source of +truth. This skill only adds the agent execution loop and the discipline that keeps a change from +being a no-op or a regression.** + +## Execution loop + +1. **Read** `docs/source/contributor-guide/optimizing_expressions.md`. +2. **Locate** the implementation under `native/spark-expr/src/` and identify the per-row cost + (allocation, kernel dispatch per element, UTF-8 decoding, regex compilation, bitmap reads). +3. **Baseline first.** Add or extend a criterion benchmark in `native/spark-expr/benches/` + (register it in `native/spark-expr/Cargo.toml`), covering the shapes from the guide: + no-null / sparse-null / dense-null, short / long, valid / invalid, ASCII / non-ASCII. Run it + on `main` before touching any source: + ```sh + cd native && cargo bench --bench -- --save-baseline main + ``` + The baseline must be **built from unmodified source**. `cargo bench` recompiles from whatever + is on disk when it runs, so if your optimization is already in the working tree (or a slow + build is still compiling when you start editing), the "baseline" measures the optimized code + and every speedup looks like zero. Only the benchmark file and `Cargo.toml` bench registration + may be present. If you have already edited the expression, `git stash push -- `, + capture the baseline, then `git stash pop`. Confirm the run finished before editing. +4. **Optimize**, preserving exact semantics. Pick a technique from the catalog in the guide. +5. **Prove correctness.** Run the existing unit tests for the function; they must pass unchanged. + Output must be bit-identical to `main` (values, null buffer, errors). There is no differential + fuzz harness in this repo — the unit tests are the gate. If coverage is thin, add tests (or run + `audit-comet-expression`) before claiming correctness. +6. **Re-measure.** `cargo bench --bench -- --baseline main`. Criterion's "change" compares + two separate process runs, so a small (±a few %) flag on a shape can be cross-run system noise + (thermal, background load), not a real effect. Before trusting any flagged regression, take a + **second independent sample** (`--baseline main` again) — real changes reproduce, noise does + not. A shape whose code path you did not touch cannot truly regress: if it flags, it is noise, + and the fix is another sample, not abandoning the change. +7. **Apply the no-regression gate** (below) before writing any PR. +8. **Finish:** `make format`, build, `cargo clippy --all-targets --workspace -- -D warnings`. + PR title ``perf: optimize `$ARGUMENTS` (Nx faster)``, paste the criterion output for every + shape. +9. **Record the performance audit.** Add a dated `Performance (tuned ...)` line under the + expression's heading on the relevant page in + `docs/source/contributor-guide/expression-audits/` (naming the technique, speedup, PR, and + benchmark file) so contributors can see what has already been tuned. Check this page before + starting, too — it tells you whether the expression was already optimized. + +## The gate: what blocks a submission + +- **A shape got meaningfully and reproducibly slower.** Do NOT submit. Even a 90%-faster-no-nulls + win does not justify a 30%-slower-dense-nulls loss. Either gate the fast path behind a per-batch + runtime check that picks the right path, or abandon the change. A Comet PR was closed for exactly + this. "Reproducibly" matters: confirm a flagged regression on a second sample first (step 6) — + a ~2% flag on a code path you did not touch is noise, not a blocker. +- **No meaningful speedup on any shape.** There is nothing to submit. A change inside criterion's + noise threshold is not an improvement. +- **Output differs from `main`** on any input, including null placement or error behavior. +- **You only benchmarked one shape.** You have not shown the absence of a regression. + +## Red flags — STOP + +- "It's obviously faster, I don't need to benchmark all shapes" → dense nulls / long values / the + error path are where fast paths regress. Benchmark them. +- "The values are right, I'll assume nulls are fine" → null-buffer bugs are the most common + correctness failure. Diff the null buffer. +- "I'll apply the fallible conversion to every slot" → a garbage value under a null must not raise + an error Spark does not. Use `try_unary` or skip null slots. +- "One benchmark improved, ship it" → re-read the gate. One improvement plus one regression is not + a win. + +## Rationalization table + +| Excuse | Reality | +| -------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| "Measuring the baseline is overhead, I'll benchmark once at the end" | Without a `main` baseline you cannot report a change %, and you cannot tell a win from noise. Baseline first. | +| "Dense-null regression is an edge case" | Real columns have dense nulls. A regression there is a regression. Gate the path or drop it. | +| "Existing tests are enough proof, no need to check the null buffer" | Tests may not assert null placement. Confirm bit-identical output explicitly. | +| "This trades a bit of the null case for a big win elsewhere" | That is a trade-off, not an optimization. Only submit a strict improvement (or a correctly-gated per-batch path). | +| "I'll edit the code, then run the baseline" | `cargo bench` compiles from disk. A baseline built with your change already applied measures the optimized code and hides the real speedup. Baseline on unmodified source; stash the edit if needed. | +| "One shape flagged a 2% regression, abandon it" | Criterion compares separate runs; small flags on an untouched code path are cross-run noise. Take a second sample before deciding — real regressions reproduce. | + +## Related skills + +- `audit-comet-expression` — shore up test coverage before optimizing when the gate is thin. +- `implement-comet-expression` — for a brand-new expression, not an existing one. +- `wire-datafusion-function` — for wiring an existing upstream function. diff --git a/docs/source/contributor-guide/expression-audits/collection_funcs.md b/docs/source/contributor-guide/expression-audits/collection_funcs.md index a2dd7bd511a..94faddcf4af 100644 --- a/docs/source/contributor-guide/expression-audits/collection_funcs.md +++ b/docs/source/contributor-guide/expression-audits/collection_funcs.md @@ -48,5 +48,6 @@ - Spark 4.0.1 (audited 2026-05-27): byte-for-byte identical to 3.5.8. - Spark 4.1.1 (audited 2026-05-27): byte-for-byte identical to 3.5.8. - Known limitation: `Size` over `MapType` falls back to Spark (https://github.com/apache/datafusion-comet/issues/4472). +- Performance (tuned 2026-07-10, PR #4877): compute list row sizes from the offset buffer instead of allocating a sliced `ArrayRef` per row via `list_array.value(i)`, removing one heap allocation per row. ~94% faster. Benchmark: `benches/array_size.rs`. [Spark Expression Support]: ../../user-guide/latest/expressions.md diff --git a/docs/source/contributor-guide/expression-audits/conversion_funcs.md b/docs/source/contributor-guide/expression-audits/conversion_funcs.md index 45b6a8f03b1..b725e3b00a7 100644 --- a/docs/source/contributor-guide/expression-audits/conversion_funcs.md +++ b/docs/source/contributor-guide/expression-audits/conversion_funcs.md @@ -35,5 +35,6 @@ - `spark.sql.legacy.castComplexTypesToString.enabled=true` is not honoured by Comet (https://github.com/apache/datafusion-comet/issues/4492). - `CAST( AS DECIMAL)` rounding may differ from Spark (`Incompatible`, gated by `spark.comet.expression.Cast.allowIncompatible`, tracked at https://github.com/apache/datafusion-comet/issues/1371). - Spark registers the type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `decimal`, `double`, `float`, `int`, `smallint`, `string`, `timestamp`, `tinyint`) as cast aliases. Each lowers to the same `Cast` node, so Comet handles it via the `cast` implementation with the same compatibility profile. +- Performance (tuned 2026-07-14, PR #4920): narrowing integer casts (`spark_cast_int_to_int`) map the values buffer in a single pass with Arrow `unary`/`try_unary` and carry the null buffer over untouched, replacing an element-by-element `Option`/`Result` iterator-collect. Up to 100x faster on narrowing casts. Benchmark: `benches/cast_numeric.rs`. [Spark Expression Support]: ../../user-guide/latest/expressions.md diff --git a/docs/source/contributor-guide/expression-audits/index.md b/docs/source/contributor-guide/expression-audits/index.md index 18b124cba10..5ec596ca51a 100644 --- a/docs/source/contributor-guide/expression-audits/index.md +++ b/docs/source/contributor-guide/expression-audits/index.md @@ -21,7 +21,12 @@ This section records per-version audit notes for Spark expressions that have been audited for Comet compatibility. These are findings from auditing the Spark implementation across versions (3.4.3, 3.5.8, 4.0.1), not a statement of support status. -For the authoritative list of which expressions Comet supports, see the user guide [Spark Expression Support](../../user-guide/latest/expressions.md). To audit an expression, use the `audit-comet-expression` skill, which appends its findings to the relevant category page below. +Two kinds of audit are recorded per expression: + +- **Correctness audits** compare the Spark implementation across versions and note divergences and test-coverage gaps. Use the `audit-comet-expression` skill, which appends its findings to the relevant category page below. +- **Performance audits** record when a native expression has been performance-tuned, so contributors can see at a glance which expressions have already been optimized. Recorded as a dated `Performance (tuned ...)` line under the expression, they name the technique, the measured speedup, the linking PR, and the benchmark file. Use the `optimize-comet-expression` skill (see [Optimizing Scalar Expressions](../optimizing_expressions.md)), which appends this line after a change is benchmarked and merged. + +For the authoritative list of which expressions Comet supports, see the user guide [Spark Expression Support](../../user-guide/latest/expressions.md). ```{toctree} :maxdepth: 1 diff --git a/docs/source/contributor-guide/expression-audits/json_funcs.md b/docs/source/contributor-guide/expression-audits/json_funcs.md index 689142e2a33..9518cc6f70a 100644 --- a/docs/source/contributor-guide/expression-audits/json_funcs.md +++ b/docs/source/contributor-guide/expression-audits/json_funcs.md @@ -42,5 +42,6 @@ ## to_json - Partial native support; options and map/array inputs fall back. +- Performance (tuned 2026-07-13, PR #4902): `escape_string` returns `Cow` (zero-alloc borrow when nothing needs escaping) and bulk-copies unescaped byte runs instead of pushing char-by-char per value. 2x faster. Benchmark: `benches/to_json.rs`. [Spark Expression Support]: ../../user-guide/latest/expressions.md diff --git a/docs/source/contributor-guide/expression-audits/math_funcs.md b/docs/source/contributor-guide/expression-audits/math_funcs.md index 0edb9fac098..f7705d82559 100644 --- a/docs/source/contributor-guide/expression-audits/math_funcs.md +++ b/docs/source/contributor-guide/expression-audits/math_funcs.md @@ -264,6 +264,7 @@ ## unhex - Spark 3.4.3, 3.5.8, 4.0.1, 4.1.1 (audited 2026-05-27): `Unhex(child, failOnError)`. Spark 4.x widens input to `StringTypeWithCollation` and wraps the inner call in try/catch; Comet `CometUnhex` forwards `failOnError` to native `spark_unhex` but does not gate on collation. +- Performance (tuned 2026-07-11, PR #4876): compile-time 256-entry hex lookup table plus a preallocated `BinaryBuilder`, cutting per-byte branching and repeated buffer reallocations. Up to 31% faster on long strings. Benchmark: `benches/unhex.rs`. ## width_bucket diff --git a/docs/source/contributor-guide/expression-audits/url_funcs.md b/docs/source/contributor-guide/expression-audits/url_funcs.md index c262762ef60..2a61c742c86 100644 --- a/docs/source/contributor-guide/expression-audits/url_funcs.md +++ b/docs/source/contributor-guide/expression-audits/url_funcs.md @@ -21,6 +21,10 @@ > Audit notes for expressions in this category that have been audited. Absence of an entry means the expression has not been audited yet, not that it is unsupported. See the user guide [Spark Expression Support] for current support status. +## parse_url + +- Performance (tuned 2026-07-11, PR #4893): cache the compiled query regex in a thread-local single-slot cache keyed by the query key, eliminating per-row `Regex::new()` recompilation in the QUERY-with-key path (the key is constant across a batch in practice). Up to 50x faster. Benchmark: `benches/parse_url.rs`. + ## try_url_decode - 4.0.1, 2026-05-05 diff --git a/docs/source/contributor-guide/index.md b/docs/source/contributor-guide/index.md index 5236d03b9dd..f293961bdad 100644 --- a/docs/source/contributor-guide/index.md +++ b/docs/source/contributor-guide/index.md @@ -79,6 +79,7 @@ Iceberg Spark Tests Debugging Guide Benchmarking Guide +Optimizing Scalar Expressions Profiling Tracing ``` diff --git a/docs/source/contributor-guide/optimizing_expressions.md b/docs/source/contributor-guide/optimizing_expressions.md new file mode 100644 index 00000000000..a000f91398a --- /dev/null +++ b/docs/source/contributor-guide/optimizing_expressions.md @@ -0,0 +1,205 @@ + + +# Optimizing Scalar Expressions + +This guide describes how to optimize the native scalar expression implementations in the +`datafusion-comet-spark-expr` crate (`native/spark-expr/`). These are the per-row and per-batch +kernels that run inside DataFusion for Spark-compatible functions such as casts, string +functions, JSON encoders, and array functions. + +The workflow is deliberate: **measure first, keep the output bit-identical, and prove the win +with a benchmark that covers the shapes where an optimization is most likely to backfire.** Every +step below exists because skipping it has produced either a no-op PR or a regression. + +## When to optimize + +Good candidates share these traits: + +- The function is called once per row or once per element, so allocation and branching costs + multiply by batch size. +- The hot path allocates (a `String`, a `Vec`, a sliced `ArrayRef`, a compiled `Regex`) on every + row, or dispatches an Arrow compute kernel per element. +- The function already has correctness tests, so behavior is pinned. + +Do not optimize speculatively. If you cannot write a benchmark that shows a meaningful +improvement, there is nothing to submit. + +## The workflow + +1. **Read the current implementation** in `native/spark-expr/src/`. Identify the per-row cost: + allocation, kernel dispatch, UTF-8 decoding, regex compilation, bitmap reads. +2. **Write or extend a criterion benchmark** in `native/spark-expr/benches/` (see below). Run it + on `main` to capture a baseline before changing any code. +3. **Apply the optimization**, preserving exact semantics (see [Correctness](#correctness-is-non-negotiable)). +4. **Run the existing unit tests** for the function. They must pass unchanged. Output must be + bit-identical to `main`, including null placement and error behavior. +5. **Re-run the benchmark** against the baseline. Confirm a meaningful speedup on at least one + shape **and no meaningful regression on any shape**. +6. **Submit** with the criterion output pasted into the PR description. + +## Writing a benchmark + +Each optimized expression should have a criterion benchmark under `native/spark-expr/benches/` +registered in `native/spark-expr/Cargo.toml`: + +```toml +[[bench]] +name = "unhex" +harness = false +``` + +Model new benchmarks on an existing one such as `benches/unhex.rs` or `benches/array_size.rs`. +A benchmark builds representative Arrow arrays (typically 8192 rows, the default batch size), +wraps the call in `black_box`, and benches several input **shapes**: + +```rust +fn criterion_benchmark(c: &mut Criterion) { + let size = 8192; + // ... build arrays for each shape ... + let mut bench = |name: &str, arr: &ArrayRef| { + let args = vec![ColumnarValue::Array(Arc::clone(arr))]; + c.bench_function(name, |b| { + b.iter(|| black_box(spark_unhex(black_box(&args)).unwrap())) + }); + }; + bench("spark_unhex: all valid", &all_valid); + bench("spark_unhex: with nulls", &with_nulls); + bench("spark_unhex: long strings", &long_valid); + bench("spark_unhex: invalid inputs", &invalid); + bench("spark_unhex: mixed hex column", &mixed); +} +``` + +### Cover the shapes that break optimizations + +The most common way a "faster" change is actually a regression is that it wins on the shape you +looked at and loses on one you did not. Always include: + +- **No nulls vs. sparse nulls vs. dense nulls.** Fast paths that assume no nulls, or that batch + contiguous non-null runs, can be slower than the base loop when nulls are dense. +- **Short vs. long values.** Allocation savings shrink and copy costs grow as values get longer. +- **Valid vs. invalid inputs.** Error and fallback paths have different costs. +- **ASCII vs. non-ASCII** for string functions with an ASCII fast path. + +### Running benchmarks + +```sh +cd native +# capture the baseline on main +git checkout main +cargo bench --bench unhex -- --save-baseline main +# switch to your branch and compare +git checkout my-branch +cargo bench --bench unhex -- --baseline main +``` + +Criterion reports the percentage change and confidence interval per shape. A change inside the +noise threshold is not an improvement. + +## Correctness is non-negotiable + +The output of the optimized function must be **bit-identical** to `main` for every input: +same values, same null buffer, same errors. Comet does not ship a differential fuzz harness in +this repo, so the existing unit tests are your correctness gate. If coverage is thin, add tests +(or run `audit-comet-expression`) before optimizing. + +Recurring correctness traps: + +- **Null slots must not raise errors.** In ANSI/overflow paths, only apply the fallible + conversion to non-null slots. Arrow's `try_unary` does this for you: a garbage value sitting + under a null cannot raise a spurious overflow. Hand-rolled loops that read the value before + checking the null bit will report errors Spark does not. +- **Preserve the exact eval-mode semantics.** For example, Legacy narrowing int casts keep the + low-order bits (a wrapping `as` conversion), while ANSI raises on overflow. These are different + kernels, not a single path with a flag. +- **Carry the null buffer through untouched** when the values transform is infallible and + one-to-one. `unary`/`try_unary` do this; a rebuild-from-iterator does not, and can drop or + shift nulls. +- **Zero-copy buffer reuse must respect offsets and slicing.** Reusing an input array's buffers + is only valid when the layout genuinely matches the output. + +## Proven techniques + +These are the patterns that have produced real speedups in merged and in-flight Comet PRs. Reach +for the lightest one that fits. + +| Technique | What it replaces | Example | +| ------------------------------------------------------------------ | ----------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| **Vectorized Arrow kernels** (`unary`, `try_unary`, `binary`) | `iter().map().collect()` over `Option`/`Result` | `spark_cast_int_to_int` (up to 100x): map the values buffer in one pass, carry the null buffer over | +| **`unary_opt` for null-producing conversions** | A builder loop that appends null on overflow, or `unary` + sentinel + `null_if_overflow` (two passes) | int/float-to-decimal casts (28-50%): map overflow to null in one vectorized pass; for the ANSI throw-on-overflow variant, gate a rare element-wise rescan on an O(1) null-count check | +| **Zero-copy borrow with `Cow`** | Allocating a new `String`/`Vec` per row when most rows are unchanged | `to_json` `escape_string` (2x): borrow when nothing needs escaping, bulk-copy unescaped byte runs | +| **Zero-copy buffer reuse** | Copying every value into a fresh builder | `cast_binary_to_string` default path (up to 7000x): reuse the binary array's buffers instead of copying | +| **Preallocate builders to known size** | Repeated buffer growth/reallocation | `spark_unhex`: preallocate `BinaryBuilder` to the known output length | +| **Compile-time lookup tables** | Per-element range matches / branching | `spark_unhex`: 256-entry hex table instead of per-digit range match | +| **Cache compiled regex** (thread-local, keyed by the constant arg) | `Regex::new()` per row | `parse_url` QUERY-with-key (50x): the key is constant across a batch | +| **Read from the offset buffer directly** | `list_array.value(i)` allocating a sliced `ArrayRef` per row | `spark_size`: compute list lengths from offsets, zero allocation | +| **Typed scans over flat values buffers + hash probe** | A per-element Arrow `eq`/compute kernel that allocates per call | `spark_arrays_overlap` (up to 18x): scan buffers directly, hash probe for large lists | +| **ASCII / byte-offset fast path** | `chars().count()` and per-char UTF-8 decoding | `substring` (up to 10x), `spark_lpad` (2x): slice by byte offset when input is ASCII | +| **`memcpy` from a precomputed buffer** | Char-by-char `push` into a scratch `String` | `spark_lpad`: pad from a precomputed repeating pad buffer, write directly into the builder | + +Cross-cutting principles behind the table: + +- **Hoist work that is constant across the batch** out of the per-row loop (regex compilation, + format setup, capacity that depends only on the whole column). +- **Avoid per-row heap allocation.** Reuse a scratch buffer, or write directly into the output + builder. +- **Prefer bulk operations** (bulk-copy a run, one `extend` per contiguous run) over per-element + operations, but verify the dense-null case does not regress. +- **Watch capacity hints.** A `rows * rows` capacity hint (instead of `rows`) silently + over-allocates; fix these when you find them. + +## The no-regression rule + +**Do not submit an optimization if any benchmark shape is meaningfully slower**, even if another +shape is dramatically faster. A change that is 90% faster with no nulls but 30% slower with dense +nulls is a trade-off, not a win, and should either be gated behind a runtime check that picks the +right path per batch or not submitted at all. A Comet optimization PR was closed for exactly this +reason (a 30% dense-null regression alongside a 90% no-null win). + +## Record the performance audit + +Once a tuning change is merged, record it in the per-expression +[Expression Audits](expression-audits/index.md) so contributors can see which expressions have +already been optimized and avoid re-treading the same ground. Add a dated `Performance (tuned ...)` +line under the expression's heading on the relevant category page (create the heading if the +expression has no audit entry yet): + +```markdown +## unhex + +- Performance (tuned 2026-07-11, PR #4876): compile-time 256-entry hex lookup table plus + preallocated `BinaryBuilder`; up to 31% faster on long strings. Benchmark: `benches/unhex.rs`. +``` + +Name the technique, the measured speedup, the PR, and the benchmark file. Keep it to one line +per tuning pass so the history stays scannable. + +## PR conventions + +Optimization PRs are small and follow a consistent shape: + +- Title: ``perf: optimize `function_name` (Nx faster)`` with the headline speedup. +- Rationale: "Optimize existing expression." +- Changes: one or two sentences naming the specific technique and what per-row cost it removes. +- Testing: "Existing tests." plus the pasted criterion output for every shape. +- Keep the diff focused on one expression. Do not bundle unrelated changes. + +Follow the standard [Pre-PR checklist](development.md): `make format`, build, and +`cargo clippy --all-targets --workspace -- -D warnings`.