mightsleep commented on code in PR #11559:
URL: https://github.com/apache/arrow-rs/pull/11559#discussion_r4363354673
##########
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:
Thanks for digging into this and for filing #11862. I'd like to keep this PR
about speed. The fast path dispatches on width only, and an unaligned buffer
falls back to zip_impl through the generic alignment check in rows, so the
behavior for it doesn't change here. Handling DayTime alignment in each kernel
with a fast path would spread the same workaround around, while the fix in
MutableArrayData, as #10171 did for MonthDayNano, covers zip, merge and every
other kernel that copies through it at once. I can pick up #11862 once one of
my open PRs lands, unless someone gets to it first.
--
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]
