mbutrovich commented on code in PR #11559:
URL: https://github.com/apache/arrow-rs/pull/11559#discussion_r4361462702
##########
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:
Filed https://github.com/apache/arrow-rs/issues/11862.
--
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]
