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.
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)?,
```
--
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]
