mightsleep commented on code in PR #11559:
URL: https://github.com/apache/arrow-rs/pull/11559#discussion_r4363377028
##########
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:
my bad, thanks for checking. I had verified this loop for branches, not for
how the masks are computed, and LLVM turns the compare into a shift into the
sign bit, which SSE2 emulates with pmuludq. Took your wording.
--
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]
