perf: Optimize scalar fast path & write() encoding for sha2 #20116
Conversation
Jefffrey
left a comment
Jefffrey
left a comment
There was a problem hiding this comment.
We should look into adding a fast path for when values is array but bit length is scalar; I assume that would be another common usecase
| let bytes = data.as_ref(); | ||
| let mut out = Vec::with_capacity(bytes.len() * 2); | ||
| for &b in bytes { | ||
| out.push(HEX_CHARS[(b >> 4) as usize]); |
There was a problem hiding this comment.
it might be a reason for extra LLVM boundaries check, maybe worth to check also
let hi = b >> 4;
let lo = b & 0x0F;
out.push(HEX_CHARS[hi as usize]);
out.push(HEX_CHARS[lo as usize]);
rustc might be smart enough to rewrite by itself
There was a problem hiding this comment.
Thanks! That looks reasonable.
Co-authored-by: Jeffrey Vo <[email protected]>
| let hi = b >> 4; | ||
| let lo = b & 0x0F; | ||
| out.push(HEX_CHARS[hi as usize]); | ||
| out.push(HEX_CHARS[lo as usize]); |
There was a problem hiding this comment.
The hex crate is used in other datafusion-** crates as an optional dependency. It also uses SIMD to be even faster for bigger input.
Consider using it here too.
There was a problem hiding this comment.
Thank you! I will look into this.
There was a problem hiding this comment.
Since sha2 digests are fixed-size (28/32/48/64 bytes), the LUT approach is already quite fast. I don't know if adding the hex will help here? What do you think?
| fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { | ||
| make_scalar_function(sha2_impl, vec![])(&args.args) | ||
| let values = &args.args[0]; | ||
| let bit_lengths = &args.args[1]; |
There was a problem hiding this comment.
Maybe use take_function_args() for consistency with other functions and better error handling ?
Co-authored-by: Martin Grigorov <[email protected]>
There was a problem hiding this comment.
Thanks @kumarUjjawal it is LGTM, I'll wait for other folks to confirm
Ive restarted CI, seems the CI errors are not PR related
Thank you! |
| where | ||
| BinaryArrType: BinaryArrayType<'a>, | ||
| { | ||
| sha2_binary_bitlen_iter(values, std::iter::repeat(Some(bit_length))) |
There was a problem hiding this comment.
I was thinking along the lines of removing the match logic on the hot loop below, if we know the bit length for all values; I think it'll result in more verbose code but could be worth performance. Can look into this in a followup
| ColumnarValue::Scalar(value_scalar), | ||
| ColumnarValue::Scalar(ScalarValue::Int32(Some(bit_length))), | ||
| ) => { | ||
| if value_scalar.is_null() { |
There was a problem hiding this comment.
We should pull all null checks into a single branch at the top, e.g.
match (values, bit_lengths) {
(ColumnarValue::Scalar(s), _) | (_, ColumnarValue::Scalar(s)) if s.is_null() => { // return scalar null }This means we'd only need 4 arms:
- One arm checking if either is null
- One arm for scalar + scalar
- One arm for array + scalar
- Catch all (array + array, scalar + array)
There was a problem hiding this comment.
Should i make the changes in this pr or will these be in the follow up too?
4 participants
Which issue does this PR close?
Rationale for this change
Spark
sha2currently evaluates scalars viamake_scalar_function(sha2_impl, vec![]), which expands scalar inputs to size-1 arrays before execution. This adds avoidable overhead for scalar evaluation / constant folding scenarios.In addition, the existing digest-to-hex formatting uses
write!(&mut s, "{b:02x}")in a loop, which is significantly slower than a LUT-based hex encoder.What changes are included in this PR?
sha2to avoid scalar→array expansion, andwrite!formatting.sha2/scalar/size=1sha2/array_binary_256/size=1024sha2/array_binary_256/size=4096sha2/array_binary_256/size=8192Are these changes tested?
Yes
Are there any user-facing changes?
No