mbutrovich commented on code in PR #26377:
URL: https://github.com/apache/datafusion/pull/26377#discussion_r4286027151
##########
datafusion/functions/src/strings.rs:
##########
@@ -829,38 +832,78 @@ impl StringViewArrayBuilder {
Self::make_long_view_checked(length, buffer_index, offset,
prefix_bytes)
}
- /// See [`BulkNullStringArrayBuilder::append_byte_map`].
+ /// Fallible variant of [`Self::append_byte_map`].
///
/// # Safety
///
/// The bytes produced by applying `map` to each byte of `src`, in order,
/// must form valid UTF-8.
///
- /// # Panics
+ /// # Errors
///
- /// Panics under the same conditions as [`Self::append_value`]: if
- /// `src.len()`, the in-progress buffer offset, or the number of completed
- /// buffers exceeds `i32::MAX`.
+ /// Returns an error under the same conditions as
+ /// [`Self::try_append_value`].
#[inline]
- pub unsafe fn append_byte_map<F: FnMut(u8) -> u8>(&mut self, src: &[u8],
mut map: F) {
- let length: u32 =
- i32::try_from(src.len()).expect("value length exceeds i32::MAX")
as u32;
- if length <= 12 {
- let mut bytes = [0u8; 12];
- for (d, &b) in bytes[..src.len()].iter_mut().zip(src) {
+ pub unsafe fn try_append_byte_map<F: FnMut(u8) -> u8>(
+ &mut self,
+ src: &[u8],
+ mut map: F,
+ ) -> Result<()> {
+ if src.len() <= MAX_INLINE_LEN {
+ // An inline view holds the length in its low 4 bytes, followed by
+ // the value zero-padded to 12 bytes. Writing the mapped bytes
+ // straight into it is faster than going through `inline_view`.
+ let mut view = (src.len() as u128).to_le_bytes();
+ for (d, &b) in view[4..].iter_mut().zip(src) {
*d = map(b);
}
- self.views.push(make_view(&bytes[..src.len()], 0, 0));
- return;
+ self.views.push(u128::from_le_bytes(view));
+ return Ok(());
}
- self.ensure_long_capacity(length);
+ let length: u32 = i32::try_from(src.len())
+ .map_err(|_| string_view_overflow_error("value length"))?
+ as u32;
+ self.try_ensure_long_capacity(length)?;
let cursor = self.in_progress.len();
- let offset: u32 = i32::try_from(cursor).expect("offset exceeds
i32::MAX") as u32;
+ let offset: u32 = i32::try_from(cursor)
+ .map_err(|_| string_view_overflow_error("offset"))?
+ as u32;
+ let buffer_index: u32 = i32::try_from(self.completed.len())
+ .map_err(|_| string_view_overflow_error("buffer count"))?
+ as u32;
self.in_progress.extend(src.iter().map(|&b| map(b)));
Review Comment:
I checked the codegen for this branch. As I read it, `upper` and `lower`
already get vector code here. I built the `string_expressions` bench binary
(release profile with LTO) for `x86_64-apple-darwin` at the default target CPU.
There, `append_case_converted_view` compiles
`self.in_progress.extend(src.iter().map(..))` to an SSE2 loop that handles 32
bytes per iteration:
```
movdqu xmm3, xmmword ptr [rdx + r10] # load 32 source bytes
movdqu xmm4, xmmword ptr [rdx + r10 + 0x10]
paddb xmm5, xmm0 # xmm0 = 0xbf, so b - 'A'
pminub xmm7, xmm1 # xmm1 = 25
pcmpeqb xmm7, xmm5 # b is in 'A'..='Z'
pand xmm7, xmm2 # xmm2 = 0x20
por xmm7, xmm3 # set the lowercase bit
movdqu xmmword ptr [r9 + r10 - 0x10], xmm7 # store into in_progress
```
The `upper` case is the same loop with `pxor`. The generic builder's new
copy-then-map loop compiles to the same thing. On aarch64, both get a NEON loop
that handles 64 bytes per iteration. As I read Arrow's
[`MutableBuffer::extend_from_iter`](https://github.com/apache/arrow-rs/blob/ef1fa157977633f0ba21aa921f9ef3d5c669235b/arrow-buffer/src/buffer/mutable.rs#L1007-L1019),
its loop runs until capacity is reached and checks `iterator.next()` on every
step, so it has two exits. `Vec::extend` on a `TrustedLen` iterator reserves
once and runs a counted loop. So copy-then-map helps the `MutableBuffer`
builder and wouldn't change the `Vec` one. For `translate`, the map is a table
lookup, which stays a scalar byte-at-a-time lookup (`movzx` loads, no xmm
registers) in both builders on x86. That fits `translate` on `Utf8View` staying
flat in your numbers. If that matches your reading, I think this thread can be
resolved without a change.
--
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]
