KonaeAkira commented on code in PR #23730:
URL: https://github.com/apache/datafusion/pull/23730#discussion_r3624773307
##########
datafusion/common/src/utils/mod.rs:
##########
@@ -411,6 +412,52 @@ pub fn split_vec_min_alloc<T>(vec: &mut Vec<T>, n: usize)
-> Vec<T> {
}
}
+/// Splits a vector of offsets at index `n`, returning a vector with
+/// `[offsets[0], ..., offsets[n]]` and leaving the remaining offsets in
shifted
+/// `offsets`, adjusted to start at 0. This function assumes monotonicity of
the
+/// `offsets` elements, so the `offsets[n]` cut-off value is subtracted from
+/// the remaining offset values with no overflow checks, except those enabled
+/// in debug builds.
+///
+/// Allocates for whichever side is smaller, so the new allocation is
+/// `min(n + 1, offsets.len() - n)`. This matters when the split emits a prefix
+/// under memory pressure, where `n` can be close to `offsets.len()`.
+pub fn take_n_offsets<O>(offsets: &mut Vec<O>, n: usize) -> Vec<O>
+where
+ O: OffsetSizeTrait,
+{
+ let cut_offset = offsets[n];
+ if n.saturating_mul(2) < offsets.len() {
+ let mut prefix = Vec::<O>::with_capacity(n + 1);
+ // SAFETY: copying to a newly allocated vector with sufficient
capacity.
+ // The length of `offsets` is checked to exceed n.
+ unsafe {
+ copy_nonoverlapping(offsets.as_ptr(), prefix.as_mut_ptr(), n);
+ *prefix.as_mut_ptr().add(n) = cut_offset;
+ prefix.set_len(n + 1);
+ }
Review Comment:
```suggestion
let prefix: Vec<O> = offsets[..=n].into();
```
This produces almost the exact same assembly: https://godbolt.org/z/zq7ETM85d
##########
datafusion/common/src/utils/mod.rs:
##########
@@ -411,6 +412,52 @@ pub fn split_vec_min_alloc<T>(vec: &mut Vec<T>, n: usize)
-> Vec<T> {
}
}
+/// Splits a vector of offsets at index `n`, returning a vector with
+/// `[offsets[0], ..., offsets[n]]` and leaving the remaining offsets in
shifted
+/// `offsets`, adjusted to start at 0. This function assumes monotonicity of
the
+/// `offsets` elements, so the `offsets[n]` cut-off value is subtracted from
+/// the remaining offset values with no overflow checks, except those enabled
+/// in debug builds.
+///
+/// Allocates for whichever side is smaller, so the new allocation is
+/// `min(n + 1, offsets.len() - n)`. This matters when the split emits a prefix
+/// under memory pressure, where `n` can be close to `offsets.len()`.
+pub fn take_n_offsets<O>(offsets: &mut Vec<O>, n: usize) -> Vec<O>
+where
+ O: OffsetSizeTrait,
+{
+ let cut_offset = offsets[n];
+ if n.saturating_mul(2) < offsets.len() {
+ let mut prefix = Vec::<O>::with_capacity(n + 1);
+ // SAFETY: copying to a newly allocated vector with sufficient
capacity.
+ // The length of `offsets` is checked to exceed n.
+ unsafe {
+ copy_nonoverlapping(offsets.as_ptr(), prefix.as_mut_ptr(), n);
+ *prefix.as_mut_ptr().add(n) = cut_offset;
+ prefix.set_len(n + 1);
+ }
+ // Shift the remaining offsets in place so that the first offset is 0.
+ let dst = offsets.as_mut_ptr();
+ for (i, &offset) in offsets[n..].iter().enumerate() {
+ // SAFETY: the range of iteration is within `offsets.len()`.
+ // In the overlapping region, the destination element is
overwritten
+ // only after it is read from.
+ unsafe {
+ *dst.add(i) = offset - cut_offset;
+ }
+ }
+ offsets.truncate(offsets.len() - n);
Review Comment:
I'm not convinced that this is faster than using `copy_within`and then
subtracting in-place, which is a `memmove` + a vectorizable subtract.
Have you benchmarked against the safe variant to justify the use of `unsafe`?
```suggestion
offsets.copy_within(n.., 0);
offsets.truncate(offsets.len() - n);
offsets.iter_mut().for_each(|o| *o -= cut_offset);
```
--
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]
