mbutrovich commented on code in PR #11559:
URL: https://github.com/apache/arrow-rs/pull/11559#discussion_r4361268730
##########
arrow-select/src/zip.rs:
##########
@@ -139,12 +139,177 @@ pub fn zip(
));
}
+ if truthy_is_scalar != falsy_is_scalar {
+ let (array, scalar) = if truthy_is_scalar {
+ (falsy, truthy)
+ } else {
+ (truthy, falsy)
+ };
+ if let Some(output) = zip_primitive_with_scalar(mask, array, scalar,
truthy_is_scalar) {
+ return Ok(output);
+ }
+ }
+
let falsy = falsy.to_data();
let truthy = truthy.to_data();
zip_impl(mask, &truthy, truthy_is_scalar, &falsy, falsy_is_scalar)
}
+/// [`zip`] of a primitive `array` with a `scalar` on the other side, or
`None` for the types it
+/// does not handle. `scalar_is_truthy` says which side the scalar is on.
+///
+/// The output values are a select per row between the array's value and the
scalar's, and its
+/// null buffer has a closed form, so no row goes through
[`MutableArrayData`]. With a null
+/// scalar the values under the scalar's rows are never read, so the array's
values buffer is
+/// shared as it is.
+fn zip_primitive_with_scalar(
+ mask: &BooleanArray,
+ array: &dyn Array,
+ scalar: &dyn Array,
+ scalar_is_truthy: bool,
+) -> Option<ArrayRef> {
+ let width = array.data_type().primitive_width()?;
+ let array = array.to_data();
+ let scalar = scalar.to_data();
+ let len = array.len();
+
+ // The rows that take the array's value
+ let mut from_array = maybe_prep_null_mask_filter(mask);
+ if scalar_is_truthy {
+ from_array = !&from_array;
+ }
+
+ let scalar_valid = scalar.is_valid(0);
+ let nulls = match (scalar_valid, array.nulls()) {
+ (true, None) => None,
+ (true, Some(nulls)) => Some(nulls.inner() | &!&from_array),
+ (false, None) => Some(from_array.clone()),
+ (false, Some(nulls)) => Some(nulls.inner() & &from_array),
+ };
+
+ let values = if scalar_valid {
+ match width {
+ 1 => blend::<i8, 1>(&from_array, &array, &scalar)?,
+ 2 => blend::<i16, 1>(&from_array, &array, &scalar)?,
+ 4 => blend::<i32, 1>(&from_array, &array, &scalar)?,
+ 8 => blend::<i64, 1>(&from_array, &array, &scalar)?,
Review Comment:
What happens here for an `Interval(DayTime)` array whose values buffer is
only 4-byte aligned? arrow-rs requires a fixed-width buffer to be aligned only
to `align_of` of its native type
([`new_fixed_width`](https://github.com/apache/arrow-rs/blob/df4dc3ede59bd43b5f71847231212381d068bb72/arrow-data/src/data.rs#L2058-L2063)),
and
[`IntervalDayTime`](https://github.com/apache/arrow-rs/blob/df4dc3ede59bd43b5f71847231212381d068bb72/arrow-buffer/src/interval.rs#L348-L355)
is two `i32`s, so its alignment is 4. The Arrow format only [recommends 8- or
64-byte
alignment](https://github.com/apache/arrow/blob/373d64c02cb23663b63c5853099fca43a46c39fc/docs/source/format/Columnar.rst?plain=1#L264-L273)
in memory.
For such a buffer (I built one by slicing a `Buffer` 4 bytes in),
`blend::<i64, 1>` returns `None`, and the fallback panics in
`Buffer::typed_data`, because `MutableArrayData` [extends `DayTime` as
`i64`](https://github.com/apache/arrow-rs/blob/df4dc3ede59bd43b5f71847231212381d068bb72/arrow-data/src/transform/mod.rs#L259-L263).
That panic also happens on `main` without this PR, so it needs its own fix.
Could you file an issue for it?
For this PR, what do you think about blending `DayTime` as two `i32` lanes?
The aligned case would then never fall back, and as I read it, the alignment
check in `rows` would no longer fail for any type that reaches `blend`. Could
the test also add `Interval(DayTime)` to `data_types`, plus one case with a
4-byte-aligned buffer? `random_primitive` always produces aligned buffers. This
needs `IntervalUnit` in the top-level `arrow_schema` import.
```suggestion
8 if *array.data_type() ==
DataType::Interval(IntervalUnit::DayTime) => {
// arrow-rs aligns a `DayTime` buffer only to
`align_of::<IntervalDayTime>()`, 4
blend::<i32, 2>(&from_array, &array, &scalar)?
}
8 => blend::<i64, 1>(&from_array, &array, &scalar)?,
```
##########
arrow-select/src/zip.rs:
##########
@@ -139,12 +139,177 @@ pub fn zip(
));
}
+ if truthy_is_scalar != falsy_is_scalar {
+ let (array, scalar) = if truthy_is_scalar {
+ (falsy, truthy)
+ } else {
+ (truthy, falsy)
+ };
+ if let Some(output) = zip_primitive_with_scalar(mask, array, scalar,
truthy_is_scalar) {
+ return Ok(output);
+ }
+ }
+
let falsy = falsy.to_data();
let truthy = truthy.to_data();
zip_impl(mask, &truthy, truthy_is_scalar, &falsy, falsy_is_scalar)
}
+/// [`zip`] of a primitive `array` with a `scalar` on the other side, or
`None` for the types it
+/// does not handle. `scalar_is_truthy` says which side the scalar is on.
+///
+/// The output values are a select per row between the array's value and the
scalar's, and its
+/// null buffer has a closed form, so no row goes through
[`MutableArrayData`]. With a null
+/// scalar the values under the scalar's rows are never read, so the array's
values buffer is
+/// shared as it is.
+fn zip_primitive_with_scalar(
+ mask: &BooleanArray,
+ array: &dyn Array,
+ scalar: &dyn Array,
+ scalar_is_truthy: bool,
+) -> Option<ArrayRef> {
+ let width = array.data_type().primitive_width()?;
+ let array = array.to_data();
+ let scalar = scalar.to_data();
+ let len = array.len();
+
+ // The rows that take the array's value
+ let mut from_array = maybe_prep_null_mask_filter(mask);
+ if scalar_is_truthy {
+ from_array = !&from_array;
+ }
+
+ let scalar_valid = scalar.is_valid(0);
+ let nulls = match (scalar_valid, array.nulls()) {
+ (true, None) => None,
+ (true, Some(nulls)) => Some(nulls.inner() | &!&from_array),
+ (false, None) => Some(from_array.clone()),
+ (false, Some(nulls)) => Some(nulls.inner() & &from_array),
+ };
+
+ let values = if scalar_valid {
+ match width {
+ 1 => blend::<i8, 1>(&from_array, &array, &scalar)?,
+ 2 => blend::<i16, 1>(&from_array, &array, &scalar)?,
+ 4 => blend::<i32, 1>(&from_array, &array, &scalar)?,
+ 8 => blend::<i64, 1>(&from_array, &array, &scalar)?,
+ 16 => blend::<i64, 2>(&from_array, &array, &scalar)?,
+ 32 => blend::<i64, 4>(&from_array, &array, &scalar)?,
+ _ => return None,
+ }
+ } else {
+ array.buffers()[0].slice_with_length(array.offset() * width, len *
width)
+ };
+
+ let data = ArrayData::builder(array.data_type().clone())
+ .len(len)
+ .add_buffer(values)
+ .nulls(nulls.map(NullBuffer::new));
+ Some(make_array(
+ data.build().expect("the layout of the input array"),
+ ))
+}
+
+/// An integer lane that [`blend`] selects with.
+trait BlendLane:
+ ArrowNativeType + BitAnd<Output = Self> + BitOr<Output = Self> +
Not<Output = Self>
+{
+ /// `mask` (0 or -1) at this width.
+ fn from_mask(mask: i32) -> Self;
+}
+
+macro_rules! blend_lane {
+ ($($t:ty),*) => {
+ $(impl BlendLane for $t {
+ #[inline]
+ fn from_mask(mask: i32) -> Self {
+ mask as $t
+ }
+ })*
+ };
+}
+blend_lane!(i8, i16, i64);
+
+impl BlendLane for i32 {
+ #[inline]
+ fn from_mask(mask: i32) -> Self {
+ mask
+ }
+}
+
+/// For each bit of `word`, an `i32` with all bits set if it is set, else
clear.
+#[inline]
+fn lane_masks(word: u64) -> [i32; 64] {
+ // Every lane of a byte tests a constant bit of the same value: an AND and
a compare of
+ // 32-bit lanes, which SSE2 and NEON have. Shifting the word by a
different count per lane
+ // has no SIMD form before AVX2, and SSE2 has no 64-bit compare.
Review Comment:
Did you check which instructions this compiles to? On rustc 1.99.1 release,
I don't see an AND and a compare on either target. On aarch64, `lane_masks`
uses the per-lane variable shift that the comment says has no SIMD form, plus
scalar bit extracts:
```
ushl.4s v2, v1, v0 ; shift each lane by its own count
cmlt.4s v2, v2, #0 ; mask = lane < 0
sbfx w11, w8, #5, #1
```
On x86_64 (SSE2), it emulates the per-lane shift with multiplies, plus
scalar `shl`/`sar`:
```
pmuludq xmm4, xmm3 # multiply by per-lane powers of two
psrad xmm8, 31 # mask = sign bit
```
Neither listing has a `pcmpeqd`, `cmeq` or `cmtst`. Both are branch-free,
which is the property that matters here. Should the comment state that property
instead, so that nobody later keeps the "AND and compare" form for a reason
that doesn't hold?
```suggestion
// Each mask tests a constant bit of one byte of the word. This compiles
without a branch per
// bit on x86_64 (SSE2) and aarch64.
```
##########
arrow-select/src/zip.rs:
##########
@@ -139,12 +139,177 @@ pub fn zip(
));
}
+ if truthy_is_scalar != falsy_is_scalar {
+ let (array, scalar) = if truthy_is_scalar {
+ (falsy, truthy)
+ } else {
+ (truthy, falsy)
+ };
+ if let Some(output) = zip_primitive_with_scalar(mask, array, scalar,
truthy_is_scalar) {
+ return Ok(output);
+ }
+ }
+
let falsy = falsy.to_data();
let truthy = truthy.to_data();
zip_impl(mask, &truthy, truthy_is_scalar, &falsy, falsy_is_scalar)
}
+/// [`zip`] of a primitive `array` with a `scalar` on the other side, or
`None` for the types it
+/// does not handle. `scalar_is_truthy` says which side the scalar is on.
+///
+/// The output values are a select per row between the array's value and the
scalar's, and its
+/// null buffer has a closed form, so no row goes through
[`MutableArrayData`]. With a null
+/// scalar the values under the scalar's rows are never read, so the array's
values buffer is
+/// shared as it is.
+fn zip_primitive_with_scalar(
+ mask: &BooleanArray,
+ array: &dyn Array,
+ scalar: &dyn Array,
+ scalar_is_truthy: bool,
+) -> Option<ArrayRef> {
+ let width = array.data_type().primitive_width()?;
+ let array = array.to_data();
+ let scalar = scalar.to_data();
+ let len = array.len();
+
+ // The rows that take the array's value
+ let mut from_array = maybe_prep_null_mask_filter(mask);
+ if scalar_is_truthy {
+ from_array = !&from_array;
+ }
+
+ let scalar_valid = scalar.is_valid(0);
+ let nulls = match (scalar_valid, array.nulls()) {
+ (true, None) => None,
+ (true, Some(nulls)) => Some(nulls.inner() | &!&from_array),
+ (false, None) => Some(from_array.clone()),
+ (false, Some(nulls)) => Some(nulls.inner() & &from_array),
+ };
+
+ let values = if scalar_valid {
+ match width {
+ 1 => blend::<i8, 1>(&from_array, &array, &scalar)?,
+ 2 => blend::<i16, 1>(&from_array, &array, &scalar)?,
+ 4 => blend::<i32, 1>(&from_array, &array, &scalar)?,
+ 8 => blend::<i64, 1>(&from_array, &array, &scalar)?,
+ 16 => blend::<i64, 2>(&from_array, &array, &scalar)?,
+ 32 => blend::<i64, 4>(&from_array, &array, &scalar)?,
+ _ => return None,
+ }
+ } else {
+ array.buffers()[0].slice_with_length(array.offset() * width, len *
width)
+ };
+
+ let data = ArrayData::builder(array.data_type().clone())
+ .len(len)
+ .add_buffer(values)
+ .nulls(nulls.map(NullBuffer::new));
+ Some(make_array(
+ data.build().expect("the layout of the input array"),
+ ))
Review Comment:
Could this return the error instead of panicking? `zip` already returns a
`Result`, and `PrimitiveScalarImpl` in this file propagates its build error
([line
633](https://github.com/apache/arrow-rs/blob/03bef9fd4622d8c78244b265a1653873e038d3de/arrow-select/src/zip.rs#L633)).
If the function returns `Option<Result<ArrayRef, ArrowError>>` (line 171), the
`?` on `blend` keeps working, and the caller at line 149 becomes `return
output;`.
```suggestion
Some(data.build().map(make_array))
```
--
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]
