perf: remove per-row String allocations from the Spark url functions - #23884
Conversation
url_encode built a String for every row via byte_serialize(..).collect::<String>() and collected the results into a StringArray, so each row paid for its own allocation. The percent-encoded form is now assembled in a single scratch buffer reused across rows and appended to a pre-sized builder. The now-unused UrlEncode::encode helper, which returned a Result that was never an Err, is removed. -56% at 1024 rows and -48% at 8192 against the benchmark added in apache#23882.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #23884 +/- ##
==========================================
- Coverage 80.98% 80.98% -0.01%
==========================================
Files 1106 1106
Lines 383395 383400 +5
Branches 383395 383400 +5
==========================================
- Hits 310509 310488 -21
- Misses 54565 54575 +10
- Partials 18321 18337 +16 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
…key array directly url_decode ended with .map(|parsed| parsed.into_owned()). Both replace_plus and decode_utf8 borrow their input when there is nothing to rewrite, so that allocated a String for every row even when the value needed no decoding at all. decode now returns a Cow and the array paths append it to a pre-sized builder. To let the Cow survive, spark_handled_url_decode's error handler changes from a closure over Result<Option<String>> to an OnDecodeError enum. Its only caller, try_url_decode, passed a closure that mapped Err to Ok(None), which is exactly OnDecodeError::Null. parse_url built its all-null key array for the 2-argument form by calling append_null() once per row on an uncapacitated builder; new_null_array does it in one allocation. url_decode -29% where nothing needs unescaping, -4% where it does, against the benchmark added in apache#23882.
| // `new_null_array` allocates the null array outright, rather than | ||
| // appending one null per row through a builder. | ||
| let key_array = new_null_array(&DataType::Utf8, args[0].len()); | ||
| let key = key_array.as_string::<i32>(); |
There was a problem hiding this comment.
| // `new_null_array` allocates the null array outright, rather than | |
| // appending one null per row through a builder. | |
| let key_array = new_null_array(&DataType::Utf8, args[0].len()); | |
| let key = key_array.as_string::<i32>(); | |
| let key = StringArray::new_null(args[0].len()); |
There was a problem hiding this comment.
@andygrove
Looks good overall. I left one small suggestion to strengthen the direct Rust unit coverage for the new string builder paths.
| }}; | ||
| } | ||
|
|
||
| match &args[0].data_type() { |
There was a problem hiding this comment.
Could we add direct Rust unit tests for the new LargeUtf8 and Utf8View builder paths in url_encode and url_decode? The SLTs already cover these input types, but the Rust unit tests mostly exercise the Utf8 path. A small test using LargeStringArray and StringViewArray would help guard these new optimized branches more directly.
…nd url_decode Requested in code review: apache#23884 (comment) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
In an effort to keep the code moving I took the liberty of merging up and implementing @kosiew 's suggestion (well having Claude Code do it, but I have reviewed it) |
|
Thanks @andygrove and @kosiew |
…pache#23884) ## Which issue does this PR close? N/A ## Rationale for this change Three allocation problems in the `url` module, all on the per-row path. **`url_encode`** built a `String` for every row: ```rust fn encode(value: &str) -> Result<String> { Ok(byte_serialize(value.as_bytes()).collect::<String>()) } ``` **`url_decode`** ended with `.map(|parsed| parsed.into_owned())`. Both `replace_plus` and `decode_utf8` return a `Cow` that *borrows* when there is nothing to rewrite, so `into_owned` allocated a `String` for every row even when the value contained no percent-escapes and no `+` at all. **`parse_url`** built the all-null `key` array for the two-argument form by calling `append_null()` once per row against a builder with no capacity. ## What changes are included in this PR? `url_encode.rs`: - The three type branches build a pre-sized builder and share an `encode_all!` loop that clears and refills one scratch `String` per batch. - `UrlEncode::encode` is removed; it returned a `Result` whose error arm was never constructed and now has no callers. `url_decode.rs`: - `decode` returns `Cow<'_, str>` instead of `String`. When `replace_plus` borrows, the decode borrows straight from the input; when it has already allocated (the input contained `+`), owning the decoded form costs nothing beyond what was already spent. - The array paths append the `Cow` to a pre-sized builder rather than collecting `Result<StringArray>` from a per-row `String`. - **API change:** `spark_handled_url_decode`'s second parameter changes from `impl Fn(Result<Option<String>>) -> Result<Option<String>>` to a new `OnDecodeError` enum. The `String` in that signature is what forced the allocation. Its only caller is `try_url_decode`, whose closure was `Err(_) => Ok(None)` — exactly `OnDecodeError::Null`. `parse_url.rs`: - The `append_null()`-per-row loop becomes `new_null_array(&DataType::Utf8, len)`. No behaviour change in any of the three. ## Are these changes tested? Existing coverage pins the behaviour: `spark/url/url_encode.slt`, `url_decode.slt`, `try_url_decode.slt`, and `parse_url.slt` assert concrete results including reserved characters, `+` handling, malformed percent-encoding (which must error for `url_decode` and yield NULL for `try_url_decode`), and the two- versus three-argument `parse_url` forms. All 5 `spark/url` sqllogictest files pass, along with the 258 `datafusion-spark` unit tests. The `try_url_decode` unit test is what pins the `OnDecodeError::Null` path, since it drives a malformed input through the refactored signature. Benchmarks are added separately in apache#23882 so the baselines can be measured on `main` before this lands. ### Benchmarks Criterion, `apache/main` @ `f1ab86dad` as baseline. Median of the reported change interval. `url_encode`: | Benchmark | 1024 | 8192 | | --- | --- | --- | | `url_encode/utf8` | −57.1% | −48.2% | | `url_encode/largeutf8` | −56.1% | −49.7% | | `url_encode/utf8view` | −56.0% | −46.5% | `url_decode`, split by whether the input actually needs unescaping: | Benchmark | 1024 | 8192 | | --- | --- | --- | | `url_decode/plain_utf8` | −26.4% | −28.8% | | `url_decode/plain_utf8view` | −29.0% | −27.2% | | `url_decode/escaped_utf8` | −5.4% | −3.6% | | `url_decode/escaped_largeutf8` | −4.7% | −3.2% | The `plain` rows are the ones that benefit: with nothing to unescape the decoded value borrows its input. The `escaped` rows still have to allocate, so they gain only the pre-sized builder and the removed intermediate — a few percent, as expected. `parse_url` has no benchmark; its change removes an `append_null()` loop that runs once per batch rather than affecting a measured per-row path. ## Are there any user-facing changes? No change to SQL behaviour — all three functions return identical results. `spark_handled_url_decode` is public Rust API and its signature changes as described above; `OnDecodeError` is new public API on `datafusion-spark`. --------- Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Which issue does this PR close?
N/A
Rationale for this change
Three allocation problems in the
urlmodule, all on the per-row path.url_encodebuilt aStringfor every row:url_decodeended with.map(|parsed| parsed.into_owned()). Bothreplace_plusanddecode_utf8return aCowthat borrows when there isnothing to rewrite, so
into_ownedallocated aStringfor every row even whenthe value contained no percent-escapes and no
+at all.parse_urlbuilt the all-nullkeyarray for the two-argument form bycalling
append_null()once per row against a builder with no capacity.What changes are included in this PR?
url_encode.rs:encode_all!loop that clears and refills one scratch
Stringper batch.UrlEncode::encodeis removed; it returned aResultwhose error arm wasnever constructed and now has no callers.
url_decode.rs:decodereturnsCow<'_, str>instead ofString. Whenreplace_plusborrows, the decode borrows straight from the input; when it has already
allocated (the input contained
+), owning the decoded form costs nothingbeyond what was already spent.
Cowto a pre-sized builder rather thancollecting
Result<StringArray>from a per-rowString.spark_handled_url_decode's second parameter changes fromimpl Fn(Result<Option<String>>) -> Result<Option<String>>to a newOnDecodeErrorenum. TheStringin that signature is what forced theallocation. Its only caller is
try_url_decode, whose closure wasErr(_) => Ok(None)— exactlyOnDecodeError::Null.parse_url.rs:append_null()-per-row loop becomesnew_null_array(&DataType::Utf8, len).No behaviour change in any of the three.
Are these changes tested?
Existing coverage pins the behaviour:
spark/url/url_encode.slt,url_decode.slt,try_url_decode.slt, andparse_url.sltassert concreteresults including reserved characters,
+handling, malformed percent-encoding(which must error for
url_decodeand yield NULL fortry_url_decode), and thetwo- versus three-argument
parse_urlforms. All 5spark/urlsqllogictestfiles pass, along with the 258
datafusion-sparkunit tests.The
try_url_decodeunit test is what pins theOnDecodeError::Nullpath,since it drives a malformed input through the refactored signature.
Benchmarks are added separately in #23882 so the baselines can be measured on
mainbefore this lands.Benchmarks
Criterion,
apache/main@f1ab86dadas baseline. Median of the reportedchange interval.
url_encode:url_encode/utf8url_encode/largeutf8url_encode/utf8viewurl_decode, split by whether the input actually needs unescaping:url_decode/plain_utf8url_decode/plain_utf8viewurl_decode/escaped_utf8url_decode/escaped_largeutf8The
plainrows are the ones that benefit: with nothing to unescape the decodedvalue borrows its input. The
escapedrows still have to allocate, so they gainonly the pre-sized builder and the removed intermediate — a few percent, as
expected.
parse_urlhas no benchmark; its change removes anappend_null()loop thatruns once per batch rather than affecting a measured per-row path.
Are there any user-facing changes?
No change to SQL behaviour — all three functions return identical results.
spark_handled_url_decodeis public Rust API and its signature changes asdescribed above;
OnDecodeErroris new public API ondatafusion-spark.