1use crate::filter::{SlicesIterator, prep_null_mask_filter};
8use arrow_array::cast::AsArray;
9use arrow_array::types::{
10 BinaryType, BinaryViewType, ByteArrayType, ByteViewType, LargeBinaryType, LargeUtf8Type,
11 StringViewType, Utf8Type,
12};
13use arrow_array::*;
14use arrow_buffer::{
15 ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer, OffsetBuffer, ScalarBuffer,
16 ToByteSlice,
17};
18use arrow_data::transform::MutableArrayData;
19use arrow_data::{ArrayData, ByteView, MAX_INLINE_VIEW_LEN};
20use arrow_schema::{ArrowError, DataType};
21use std::fmt::{Debug, Formatter};
22use std::hash::Hash;
23use std::marker::PhantomData;
24use std::ops::{BitAnd, BitOr, Not};
25use std::sync::{Arc, OnceLock};
26
27pub fn zip(
87 mask: &BooleanArray,
88 truthy: &dyn Datum,
89 falsy: &dyn Datum,
90) -> Result<ArrayRef, ArrowError> {
91 let (truthy_array, truthy_is_scalar) = truthy.get();
92 let (falsy_array, falsy_is_scalar) = falsy.get();
93
94 if falsy_is_scalar && truthy_is_scalar {
95 let zipper = ScalarZipper::try_new(truthy, falsy)?;
96 return zipper.zip_impl.create_output(mask);
97 }
98
99 let truthy = truthy_array;
100 let falsy = falsy_array;
101
102 if truthy.data_type() != falsy.data_type() {
103 return Err(ArrowError::InvalidArgumentError(
104 "arguments need to have the same data type".into(),
105 ));
106 }
107
108 if truthy_is_scalar && truthy.len() != 1 {
109 return Err(ArrowError::InvalidArgumentError(
110 "scalar arrays must have 1 element".into(),
111 ));
112 }
113 if !truthy_is_scalar && truthy.len() != mask.len() {
114 return Err(ArrowError::InvalidArgumentError(
115 "all arrays should have the same length".into(),
116 ));
117 }
118 if falsy_is_scalar && falsy.len() != 1 {
119 return Err(ArrowError::InvalidArgumentError(
120 "scalar arrays must have 1 element".into(),
121 ));
122 }
123 if !falsy_is_scalar && falsy.len() != mask.len() {
124 return Err(ArrowError::InvalidArgumentError(
125 "all arrays should have the same length".into(),
126 ));
127 }
128
129 if truthy_is_scalar != falsy_is_scalar {
130 let (array, scalar) = if truthy_is_scalar {
131 (falsy, truthy)
132 } else {
133 (truthy, falsy)
134 };
135 if let Some(output) = zip_primitive_with_scalar(mask, array, scalar, truthy_is_scalar) {
136 return output;
137 }
138 }
139
140 let falsy = falsy.to_data();
141 let truthy = truthy.to_data();
142
143 zip_impl(mask, &truthy, truthy_is_scalar, &falsy, falsy_is_scalar)
144}
145
146fn zip_primitive_with_scalar(
154 mask: &BooleanArray,
155 array: &dyn Array,
156 scalar: &dyn Array,
157 scalar_is_truthy: bool,
158) -> Option<Result<ArrayRef, ArrowError>> {
159 let width = array.data_type().primitive_width()?;
160 let array = array.to_data();
161 let scalar = scalar.to_data();
162 let len = array.len();
163
164 let mut from_array = maybe_prep_null_mask_filter(mask);
166 if scalar_is_truthy {
167 from_array = !&from_array;
168 }
169
170 let scalar_valid = scalar.is_valid(0);
171 let nulls = match (scalar_valid, array.nulls()) {
172 (true, None) => None,
173 (true, Some(nulls)) => Some(nulls.inner() | &!&from_array),
174 (false, None) => Some(from_array.clone()),
175 (false, Some(nulls)) => Some(nulls.inner() & &from_array),
176 };
177
178 let values = if scalar_valid {
179 match width {
180 1 => blend::<i8, 1>(&from_array, &array, &scalar)?,
181 2 => blend::<i16, 1>(&from_array, &array, &scalar)?,
182 4 => blend::<i32, 1>(&from_array, &array, &scalar)?,
183 8 => blend::<i64, 1>(&from_array, &array, &scalar)?,
184 16 => blend::<i64, 2>(&from_array, &array, &scalar)?,
185 32 => blend::<i64, 4>(&from_array, &array, &scalar)?,
186 _ => return None,
187 }
188 } else {
189 array.buffers()[0].slice_with_length(array.offset() * width, len * width)
190 };
191
192 let data = ArrayData::builder(array.data_type().clone())
193 .len(len)
194 .add_buffer(values)
195 .nulls(nulls.map(NullBuffer::new));
196 Some(data.build().map(make_array))
197}
198
199trait BlendLane:
201 ArrowNativeType + BitAnd<Output = Self> + BitOr<Output = Self> + Not<Output = Self>
202{
203 fn from_mask(mask: i32) -> Self;
205}
206
207macro_rules! blend_lane {
208 ($($t:ty),*) => {
209 $(impl BlendLane for $t {
210 #[inline]
211 fn from_mask(mask: i32) -> Self {
212 mask as $t
213 }
214 })*
215 };
216}
217blend_lane!(i8, i16, i64);
218
219impl BlendLane for i32 {
220 #[inline]
221 fn from_mask(mask: i32) -> Self {
222 mask
223 }
224}
225
226#[inline]
228fn lane_masks(word: u64) -> [i32; 64] {
229 let mut masks = [0; 64];
232 let (mask_bytes, _) = masks.as_chunks_mut::<8>();
233 for (masks, byte) in mask_bytes.iter_mut().zip(word.to_le_bytes()) {
234 let byte = i32::from(byte);
235 for (j, mask) in masks.iter_mut().enumerate() {
236 let bit = 1 << j;
237 *mask = -i32::from(byte & bit == bit);
238 }
239 }
240 masks
241}
242
243fn blend<U: BlendLane, const K: usize>(
247 from_array: &BooleanBuffer,
248 array: &ArrayData,
249 scalar: &ArrayData,
250) -> Option<Buffer> {
251 fn rows<U: ArrowNativeType, const K: usize>(data: &ArrayData) -> Option<&[[U; K]]> {
252 let buffer = &data.buffers()[0];
253 let aligned = buffer.as_ptr().align_offset(align_of::<U>()) == 0;
254 aligned.then(|| {
255 let values = &buffer.typed_data::<U>()[data.offset() * K..][..data.len() * K];
256 values.as_chunks::<K>().0
257 })
258 }
259 let values = rows::<U, K>(array)?;
260 let scalar = rows::<U, K>(scalar)?[0];
261
262 let mut output = MutableBuffer::with_capacity(values.len() * K * size_of::<U>());
265 let (value_blocks, value_rest) = values.as_chunks::<64>();
266 let mut words = from_array.bit_chunks().iter_padded();
267 let fill = [scalar; 64];
268 let mut blended = fill;
269
270 for (values, word) in value_blocks.iter().zip(&mut words) {
271 let block = match word {
272 u64::MAX => values,
273 0 => &fill,
274 _ => {
275 let masks = lane_masks(word);
279 for ((output, values), &mask) in blended.iter_mut().zip(values).zip(&masks) {
280 let lane = U::from_mask(mask);
281 for ((output, &value), &scalar) in output.iter_mut().zip(values).zip(&scalar) {
282 *output = (value & lane) | (scalar & !lane);
283 }
284 }
285 &blended
286 }
287 };
288 output.extend_from_slice(block.as_flattened());
289 }
290 let word = words.next().unwrap_or_default();
291 for (j, value) in value_rest.iter().enumerate() {
292 output.extend_from_slice(if (word >> j) & 1 != 0 { value } else { &scalar });
293 }
294 Some(output.into())
295}
296
297fn count_true_runs(mask: &BooleanBuffer) -> usize {
298 let mut slices = 0;
299 let mut previous = 0;
300 for chunk in mask.bit_chunks().iter_padded() {
301 let starts = chunk & !((chunk << 1) | previous);
302 slices += starts.count_ones() as usize;
303 previous = chunk >> 63;
304 }
305 slices
306}
307
308fn should_use_interleave(mask: &BooleanBuffer) -> bool {
309 const MIN_LEN: usize = 1024;
310
311 mask.len() >= MIN_LEN && count_true_runs(mask) > mask.len() / 8
315}
316
317fn interleave_arrays(
318 mask: &BooleanBuffer,
319 truthy: &ArrayData,
320 truthy_is_scalar: bool,
321 falsy: &ArrayData,
322 falsy_is_scalar: bool,
323) -> Result<ArrayRef, ArrowError> {
324 let truthy = make_array(truthy.clone());
325 let falsy = make_array(falsy.clone());
326 let indices: Vec<_> = mask
327 .iter()
328 .enumerate()
329 .map(|(idx, selected)| {
330 let is_scalar = if selected {
331 truthy_is_scalar
332 } else {
333 falsy_is_scalar
334 };
335 (usize::from(!selected), if is_scalar { 0 } else { idx })
337 })
338 .collect();
339 crate::interleave::interleave(&[truthy.as_ref(), falsy.as_ref()], &indices)
340}
341
342fn zip_impl(
343 mask: &BooleanArray,
344 truthy: &ArrayData,
345 truthy_is_scalar: bool,
346 falsy: &ArrayData,
347 falsy_is_scalar: bool,
348) -> Result<ArrayRef, ArrowError> {
349 let mask_buffer = maybe_prep_null_mask_filter(mask);
350 if !truthy_is_scalar && !falsy_is_scalar && should_use_interleave(&mask_buffer) {
351 return interleave_arrays(&mask_buffer, truthy, false, falsy, false);
352 }
353
354 let mut mutable = match MutableArrayData::try_new(vec![truthy, falsy], false, truthy.len()) {
355 Ok(mutable) => mutable,
356 Err(ArrowError::DictionaryKeyOverflowError) => {
360 return interleave_arrays(
361 &mask_buffer,
362 truthy,
363 truthy_is_scalar,
364 falsy,
365 falsy_is_scalar,
366 );
367 }
368 Err(e) => return Err(e),
369 };
370
371 let mut filled = 0;
376
377 for (start, end) in SlicesIterator::from(&mask_buffer) {
378 if start > filled {
380 if falsy_is_scalar {
381 for _ in filled..start {
382 mutable.try_extend(1, 0, 1)?;
384 }
385 } else {
386 mutable.try_extend(1, filled, start)?;
387 }
388 }
389 if truthy_is_scalar {
391 for _ in start..end {
392 mutable.try_extend(0, 0, 1)?;
394 }
395 } else {
396 mutable.try_extend(0, start, end)?;
397 }
398 filled = end;
399 }
400 if filled < mask.len() {
402 if falsy_is_scalar {
403 for _ in filled..mask.len() {
404 mutable.try_extend(1, 0, 1)?;
406 }
407 } else {
408 mutable.try_extend(1, filled, mask.len())?;
409 }
410 }
411
412 let data = mutable.freeze();
413 Ok(make_array(data))
414}
415
416#[derive(Debug, Clone)]
438pub struct ScalarZipper {
439 zip_impl: Arc<dyn ZipImpl>,
440}
441
442impl ScalarZipper {
443 pub fn try_new(truthy: &dyn Datum, falsy: &dyn Datum) -> Result<Self, ArrowError> {
451 let (truthy, truthy_is_scalar) = truthy.get();
452 let (falsy, falsy_is_scalar) = falsy.get();
453
454 if truthy.data_type() != falsy.data_type() {
455 return Err(ArrowError::InvalidArgumentError(
456 "arguments need to have the same data type".into(),
457 ));
458 }
459
460 if !truthy_is_scalar {
461 return Err(ArrowError::InvalidArgumentError(
462 "only scalar arrays are supported".into(),
463 ));
464 }
465
466 if !falsy_is_scalar {
467 return Err(ArrowError::InvalidArgumentError(
468 "only scalar arrays are supported".into(),
469 ));
470 }
471
472 if truthy.len() != 1 {
473 return Err(ArrowError::InvalidArgumentError(
474 "scalar arrays must have 1 element".into(),
475 ));
476 }
477 if falsy.len() != 1 {
478 return Err(ArrowError::InvalidArgumentError(
479 "scalar arrays must have 1 element".into(),
480 ));
481 }
482
483 macro_rules! primitive_size_helper {
484 ($t:ty) => {
485 Arc::new(PrimitiveScalarImpl::<$t>::new(truthy, falsy)) as Arc<dyn ZipImpl>
486 };
487 }
488
489 let zip_impl = downcast_primitive! {
490 truthy.data_type() => (primitive_size_helper),
491 DataType::Utf8 => {
492 Arc::new(BytesScalarImpl::<Utf8Type>::new(truthy, falsy)) as Arc<dyn ZipImpl>
493 },
494 DataType::LargeUtf8 => {
495 Arc::new(BytesScalarImpl::<LargeUtf8Type>::new(truthy, falsy)) as Arc<dyn ZipImpl>
496 },
497 DataType::Binary => {
498 Arc::new(BytesScalarImpl::<BinaryType>::new(truthy, falsy)) as Arc<dyn ZipImpl>
499 },
500 DataType::LargeBinary => {
501 Arc::new(BytesScalarImpl::<LargeBinaryType>::new(truthy, falsy)) as Arc<dyn ZipImpl>
502 },
503 DataType::Utf8View => {
504 Arc::new(ByteViewScalarImpl::<StringViewType>::new(truthy, falsy)) as Arc<dyn ZipImpl>
505 },
506 DataType::BinaryView => {
507 Arc::new(ByteViewScalarImpl::<BinaryViewType>::new(truthy, falsy)) as Arc<dyn ZipImpl>
508 },
509 _ => {
510 Arc::new(FallbackImpl::new(truthy, falsy)) as Arc<dyn ZipImpl>
511 },
512 };
513
514 Ok(Self { zip_impl })
515 }
516
517 pub fn zip(&self, mask: &BooleanArray) -> Result<ArrayRef, ArrowError> {
520 self.zip_impl.create_output(mask)
521 }
522}
523
524trait ZipImpl: Debug + Send + Sync {
526 fn create_output(&self, input: &BooleanArray) -> Result<ArrayRef, ArrowError>;
528}
529
530#[derive(Debug, PartialEq)]
531struct FallbackImpl {
532 truthy: ArrayData,
533 falsy: ArrayData,
534}
535
536impl FallbackImpl {
537 fn new(left: &dyn Array, right: &dyn Array) -> Self {
538 Self {
539 truthy: left.to_data(),
540 falsy: right.to_data(),
541 }
542 }
543}
544
545impl ZipImpl for FallbackImpl {
546 fn create_output(&self, predicate: &BooleanArray) -> Result<ArrayRef, ArrowError> {
547 zip_impl(predicate, &self.truthy, true, &self.falsy, true)
548 }
549}
550
551struct PrimitiveScalarImpl<T: ArrowPrimitiveType> {
552 data_type: DataType,
553 truthy: Option<T::Native>,
554 falsy: Option<T::Native>,
555}
556
557impl<T: ArrowPrimitiveType> Debug for PrimitiveScalarImpl<T> {
558 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
559 f.debug_struct("PrimitiveScalarImpl")
560 .field("data_type", &self.data_type)
561 .field("truthy", &self.truthy)
562 .field("falsy", &self.falsy)
563 .finish()
564 }
565}
566
567impl<T: ArrowPrimitiveType> PrimitiveScalarImpl<T> {
568 fn new(truthy: &dyn Array, falsy: &dyn Array) -> Self {
569 Self {
570 data_type: truthy.data_type().clone(),
571 truthy: Self::get_value_from_scalar(truthy),
572 falsy: Self::get_value_from_scalar(falsy),
573 }
574 }
575
576 fn get_value_from_scalar(scalar: &dyn Array) -> Option<T::Native> {
577 if scalar.is_null(0) {
578 None
579 } else {
580 let value = scalar.as_primitive::<T>().value(0);
581
582 Some(value)
583 }
584 }
585
586 fn get_scalar_and_null_buffer_for_single_non_nullable(
590 predicate: BooleanBuffer,
591 value: T::Native,
592 ) -> (Vec<T::Native>, Option<NullBuffer>) {
593 let result_len = predicate.len();
594 let nulls = NullBuffer::new(predicate);
595 let scalars = vec![value; result_len];
596
597 (scalars, Some(nulls))
598 }
599}
600
601impl<T: ArrowPrimitiveType> ZipImpl for PrimitiveScalarImpl<T> {
602 fn create_output(&self, predicate: &BooleanArray) -> Result<ArrayRef, ArrowError> {
603 let result_len = predicate.len();
604 let predicate = maybe_prep_null_mask_filter(predicate);
606
607 let (scalars, nulls): (Vec<T::Native>, Option<NullBuffer>) = match (self.truthy, self.falsy)
608 {
609 (Some(truthy_val), Some(falsy_val)) => {
610 let scalars: Vec<T::Native> = predicate
611 .iter()
612 .map(|b| if b { truthy_val } else { falsy_val })
613 .collect();
614
615 (scalars, None)
616 }
617 (Some(truthy_val), None) => {
618 Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, truthy_val)
622 }
623 (None, Some(falsy_val)) => {
624 let predicate = predicate.not();
629
630 Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, falsy_val)
631 }
632 (None, None) => {
633 let nulls = NullBuffer::new_null(result_len);
635 let scalars = vec![T::default_value(); result_len];
636
637 (scalars, Some(nulls))
638 }
639 };
640
641 let scalars = ScalarBuffer::<T::Native>::from(scalars);
642 let output = PrimitiveArray::<T>::try_new(scalars, nulls)?;
643
644 let output = output.with_data_type(self.data_type.clone());
646
647 Ok(Arc::new(output))
648 }
649}
650
651#[derive(PartialEq, Hash)]
652struct BytesScalarImpl<T: ByteArrayType> {
653 truthy: Option<Vec<u8>>,
654 falsy: Option<Vec<u8>>,
655 phantom: PhantomData<T>,
656}
657
658impl<T: ByteArrayType> Debug for BytesScalarImpl<T> {
659 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
660 f.debug_struct("BytesScalarImpl")
661 .field("truthy", &self.truthy)
662 .field("falsy", &self.falsy)
663 .finish()
664 }
665}
666
667impl<T: ByteArrayType> BytesScalarImpl<T> {
668 fn new(truthy_value: &dyn Array, falsy_value: &dyn Array) -> Self {
669 Self {
670 truthy: Self::get_value_from_scalar(truthy_value),
671 falsy: Self::get_value_from_scalar(falsy_value),
672 phantom: PhantomData,
673 }
674 }
675
676 fn get_value_from_scalar(scalar: &dyn Array) -> Option<Vec<u8>> {
677 if scalar.is_null(0) {
678 None
679 } else {
680 let bytes: &[u8] = scalar.as_bytes::<T>().value(0).as_ref();
681
682 Some(bytes.to_vec())
683 }
684 }
685
686 #[expect(clippy::type_complexity)]
690 fn get_scalar_and_null_buffer_for_single_non_nullable(
691 predicate: BooleanBuffer,
692 value: &[u8],
693 ) -> Result<(Buffer, OffsetBuffer<T::Offset>, Option<NullBuffer>), ArrowError> {
694 let value_length = value.len();
695
696 let number_of_true = predicate.count_set_bits();
697
698 if number_of_true == 0 {
700 let nulls = NullBuffer::new_null(predicate.len());
702
703 return Ok((
704 Buffer::from(&[]),
706 OffsetBuffer::<T::Offset>::new_zeroed(predicate.len()),
708 Some(nulls),
709 ));
710 }
711
712 let offsets = OffsetBuffer::<T::Offset>::from_lengths(
713 predicate.iter().map(|b| if b { value_length } else { 0 }),
714 );
715
716 let mut bytes = MutableBuffer::default();
717 bytes
718 .try_repeat_slice_n_times(value, number_of_true)
719 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
720
721 let bytes = Buffer::from(bytes);
722
723 let nulls = NullBuffer::new(predicate);
726
727 Ok((bytes, offsets, Some(nulls)))
728 }
729
730 fn get_bytes_and_offset_for_all_same_value(
733 number_of_values: usize,
734 value: &[u8],
735 ) -> Result<(Buffer, OffsetBuffer<T::Offset>), ArrowError> {
736 let value_length = value.len();
737
738 let offsets =
739 OffsetBuffer::<T::Offset>::from_repeated_length(value_length, number_of_values);
740
741 let mut bytes = MutableBuffer::default();
742 bytes
743 .try_repeat_slice_n_times(value, number_of_values)
744 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
745 let bytes = Buffer::from(bytes);
746
747 Ok((bytes, offsets))
748 }
749
750 fn create_output_on_non_nulls(
751 predicate: &BooleanBuffer,
752 truthy_val: &[u8],
753 falsy_val: &[u8],
754 ) -> Result<(Buffer, OffsetBuffer<<T as ByteArrayType>::Offset>), ArrowError> {
755 let true_count = predicate.count_set_bits();
756
757 match true_count {
758 0 => {
759 return Self::get_bytes_and_offset_for_all_same_value(predicate.len(), falsy_val);
761 }
762 n if n == predicate.len() => {
763 return Self::get_bytes_and_offset_for_all_same_value(predicate.len(), truthy_val);
765 }
766
767 _ => {
768 }
770 }
771
772 let total_number_of_bytes = true_count
773 .checked_mul(truthy_val.len())
774 .and_then(|truthy_bytes| {
775 let falsy_bytes = (predicate.len() - true_count).checked_mul(falsy_val.len())?;
776 truthy_bytes.checked_add(falsy_bytes)
777 })
778 .ok_or_else(|| ArrowError::MemoryError("zip output size overflow".to_string()))?;
779 T::Offset::from_usize(total_number_of_bytes)
780 .ok_or(ArrowError::OffsetOverflowError(total_number_of_bytes))?;
781 let mut mutable = MutableBuffer::try_with_capacity(total_number_of_bytes)
782 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
783 let mut offsets = Vec::<T::Offset>::with_capacity(predicate.len() + 1);
784 offsets.push(T::Offset::usize_as(0));
785 let mut current_offset: usize = 0;
786
787 let mut filled = 0;
789
790 let truthy_len = truthy_val.len();
791 let falsy_len = falsy_val.len();
792
793 SlicesIterator::from(predicate).try_for_each(|(start, end)| -> Result<(), ArrowError> {
795 if start > filled {
797 let false_repeat_count = start - filled;
798 mutable
800 .try_repeat_slice_n_times(falsy_val, false_repeat_count)
801 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
802
803 let start_offset = current_offset;
804 current_offset += falsy_len * false_repeat_count;
805 offsets.extend(
806 (1..=false_repeat_count)
807 .map(|index| T::Offset::usize_as(start_offset + index * falsy_len)),
808 );
809 }
810
811 let true_repeat_count = end - start;
812 mutable
814 .try_repeat_slice_n_times(truthy_val, true_repeat_count)
815 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
816
817 let start_offset = current_offset;
818 current_offset += truthy_len * true_repeat_count;
819 offsets.extend(
820 (1..=true_repeat_count)
821 .map(|index| T::Offset::usize_as(start_offset + index * truthy_len)),
822 );
823 filled = end;
824 Ok(())
825 })?;
826 if filled < predicate.len() {
828 let false_repeat_count = predicate.len() - filled;
829 mutable
831 .try_repeat_slice_n_times(falsy_val, false_repeat_count)
832 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
833
834 let start_offset = current_offset;
835 current_offset += falsy_len * false_repeat_count;
836 offsets.extend(
837 (1..=false_repeat_count)
838 .map(|index| T::Offset::usize_as(start_offset + index * falsy_len)),
839 );
840 }
841
842 debug_assert_eq!(current_offset, total_number_of_bytes);
843 let offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
845 Ok((mutable.into(), offsets))
846 }
847}
848
849impl<T: ByteArrayType> ZipImpl for BytesScalarImpl<T> {
850 fn create_output(&self, predicate: &BooleanArray) -> Result<ArrayRef, ArrowError> {
851 let result_len = predicate.len();
852 let predicate = maybe_prep_null_mask_filter(predicate);
854
855 let (bytes, offsets, nulls): (Buffer, OffsetBuffer<T::Offset>, Option<NullBuffer>) =
856 match (self.truthy.as_deref(), self.falsy.as_deref()) {
857 (Some(truthy_val), Some(falsy_val)) => {
858 let (bytes, offsets) =
859 Self::create_output_on_non_nulls(&predicate, truthy_val, falsy_val)?;
860
861 (bytes, offsets, None)
862 }
863 (Some(truthy_val), None) => {
864 Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, truthy_val)?
865 }
866 (None, Some(falsy_val)) => {
867 let predicate = predicate.not();
872 Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, falsy_val)?
873 }
874 (None, None) => {
875 let nulls = NullBuffer::new_null(result_len);
877
878 (
879 Buffer::from(&[]),
881 OffsetBuffer::<T::Offset>::new_zeroed(predicate.len()),
883 Some(nulls),
884 )
885 }
886 };
887
888 let output = unsafe {
889 GenericByteArray::<T>::new_unchecked(offsets, bytes, nulls)
892 };
893
894 Ok(Arc::new(output))
895 }
896}
897
898fn maybe_prep_null_mask_filter(predicate: &BooleanArray) -> BooleanBuffer {
899 if predicate.null_count() == 0 {
901 predicate.values().clone()
902 } else {
903 let cleaned = prep_null_mask_filter(predicate);
904 let (boolean_buffer, _) = cleaned.into_parts();
905 boolean_buffer
906 }
907}
908
909struct ByteViewScalarImpl<T: ByteViewType> {
910 truthy_view: Option<u128>,
911 truthy_buffers: Arc<[Buffer]>,
912 falsy_view: Option<u128>,
913 falsy_buffers: Arc<[Buffer]>,
914 phantom: PhantomData<T>,
915}
916
917static EMPTY_ARC: OnceLock<Arc<[Buffer]>> = OnceLock::new();
918fn empty_arc_buffers() -> Arc<[Buffer]> {
919 Arc::clone(EMPTY_ARC.get_or_init(|| Arc::new([])))
920}
921
922impl<T: ByteViewType> ByteViewScalarImpl<T> {
923 fn new(truthy: &dyn Array, falsy: &dyn Array) -> Self {
924 let (truthy_view, truthy_buffers) = Self::get_value_from_scalar(truthy);
925 let (falsy_view, falsy_buffers) = Self::get_value_from_scalar(falsy);
926 Self {
927 truthy_view,
928 truthy_buffers,
929 falsy_view,
930 falsy_buffers,
931 phantom: PhantomData,
932 }
933 }
934
935 fn get_value_from_scalar(scalar: &dyn Array) -> (Option<u128>, Arc<[Buffer]>) {
936 if scalar.is_null(0) {
937 (None, empty_arc_buffers())
938 } else {
939 let (views, buffers, _) = scalar.as_byte_view::<T>().clone().into_parts();
940 (views.first().copied(), buffers)
941 }
942 }
943
944 fn get_views_for_single_non_nullable(
945 predicate: BooleanBuffer,
946 value: u128,
947 buffers: Arc<[Buffer]>,
948 ) -> (ScalarBuffer<u128>, Arc<[Buffer]>, Option<NullBuffer>) {
949 let number_of_true = predicate.count_set_bits();
950 let number_of_values = predicate.len();
951
952 if number_of_true == 0 {
954 return (
956 vec![0; number_of_values].into(),
957 empty_arc_buffers(),
958 Some(NullBuffer::new_null(number_of_values)),
959 );
960 }
961 let bytes = vec![value; number_of_values];
962
963 let nulls = NullBuffer::new(predicate);
966 (bytes.into(), buffers, Some(nulls))
967 }
968
969 #[expect(clippy::type_complexity)]
970 fn get_views_for_non_nullable(
971 predicate: BooleanBuffer,
972 result_len: usize,
973 truthy_view: u128,
974 truthy_buffers: Arc<[Buffer]>,
975 falsy_view: u128,
976 falsy_buffers: Arc<[Buffer]>,
977 ) -> Result<(ScalarBuffer<u128>, Arc<[Buffer]>, Option<NullBuffer>), ArrowError> {
978 let true_count = predicate.count_set_bits();
979 match true_count {
980 0 => {
981 Ok((vec![falsy_view; result_len].into(), falsy_buffers, None))
983 }
984 n if n == predicate.len() => {
985 Ok((vec![truthy_view; result_len].into(), truthy_buffers, None))
987 }
988 _ => {
989 let mut buffers: Vec<Buffer> = truthy_buffers.to_vec();
990
991 let view_falsy = if (falsy_view as u32) <= MAX_INLINE_VIEW_LEN {
993 falsy_view
994 } else {
995 let byte_view_falsy = ByteView::from(falsy_view);
996 let new_index_falsy_buffers =
997 buffers.len() as u32 + byte_view_falsy.buffer_index;
998 buffers.extend(falsy_buffers.iter().cloned());
999 let byte_view_falsy =
1000 byte_view_falsy.with_buffer_index(new_index_falsy_buffers);
1001 byte_view_falsy.as_u128()
1002 };
1003
1004 let total_number_of_bytes = predicate
1005 .len()
1006 .checked_mul(std::mem::size_of::<u128>())
1007 .ok_or_else(|| ArrowError::MemoryError("zip output size overflow".into()))?;
1008 let mut mutable = MutableBuffer::try_with_capacity(total_number_of_bytes)
1009 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
1010 let mut filled = 0;
1011
1012 SlicesIterator::from(&predicate).try_for_each(
1013 |(start, end)| -> Result<(), ArrowError> {
1014 if start > filled {
1015 let false_repeat_count = start - filled;
1016 mutable
1017 .try_repeat_slice_n_times(
1018 view_falsy.to_byte_slice(),
1019 false_repeat_count,
1020 )
1021 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
1022 }
1023 let true_repeat_count = end - start;
1024 mutable
1025 .try_repeat_slice_n_times(
1026 truthy_view.to_byte_slice(),
1027 true_repeat_count,
1028 )
1029 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
1030 filled = end;
1031 Ok(())
1032 },
1033 )?;
1034
1035 if filled < predicate.len() {
1036 let false_repeat_count = predicate.len() - filled;
1037 mutable
1038 .try_repeat_slice_n_times(view_falsy.to_byte_slice(), false_repeat_count)
1039 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
1040 }
1041
1042 let bytes = Buffer::from(mutable);
1043 Ok((bytes.into(), buffers.into(), None))
1044 }
1045 }
1046 }
1047}
1048
1049impl<T: ByteViewType> Debug for ByteViewScalarImpl<T> {
1050 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1051 f.debug_struct("ByteViewScalarImpl")
1052 .field("truthy", &self.truthy_view)
1053 .field("falsy", &self.falsy_view)
1054 .finish()
1055 }
1056}
1057
1058impl<T: ByteViewType> ZipImpl for ByteViewScalarImpl<T> {
1059 fn create_output(&self, predicate: &BooleanArray) -> Result<ArrayRef, ArrowError> {
1060 let result_len = predicate.len();
1061 let predicate = maybe_prep_null_mask_filter(predicate);
1063
1064 let (views, buffers, nulls) = match (self.truthy_view, self.falsy_view) {
1065 (Some(truthy), Some(falsy)) => Self::get_views_for_non_nullable(
1066 predicate,
1067 result_len,
1068 truthy,
1069 Arc::clone(&self.truthy_buffers),
1070 falsy,
1071 Arc::clone(&self.falsy_buffers),
1072 )?,
1073 (Some(truthy), None) => Self::get_views_for_single_non_nullable(
1074 predicate,
1075 truthy,
1076 Arc::clone(&self.truthy_buffers),
1077 ),
1078 (None, Some(falsy)) => {
1079 let predicate = predicate.not();
1080 Self::get_views_for_single_non_nullable(
1081 predicate,
1082 falsy,
1083 Arc::clone(&self.falsy_buffers),
1084 )
1085 }
1086 (None, None) => {
1087 (
1089 vec![0; result_len].into(),
1090 empty_arc_buffers(),
1091 Some(NullBuffer::new_null(result_len)),
1092 )
1093 }
1094 };
1095
1096 let result = unsafe { GenericByteViewArray::<T>::new_unchecked(views, buffers, nulls) };
1097 Ok(Arc::new(result))
1098 }
1099}
1100
1101#[cfg(test)]
1102mod test {
1103 use super::*;
1104 use arrow_array::types::{Int8Type, Int32Type};
1105 use arrow_schema::{IntervalUnit, TimeUnit};
1106 use rand::rngs::StdRng;
1107 use rand::{RngExt, SeedableRng};
1108
1109 #[test]
1110 fn test_count_true_runs() {
1111 let assert_runs = |values: &[bool], expected| {
1112 let mask: BooleanBuffer = values.iter().copied().collect();
1113 assert_eq!(count_true_runs(&mask), expected, "mask: {values:?}");
1114 };
1115
1116 assert_runs(&[], 0);
1117 assert_runs(&[false, false, false], 0);
1118 assert_runs(&[true, true, true], 1);
1119 assert_runs(&[true, false, true, true, false, true], 3);
1120
1121 let mut values = vec![false; 130];
1123 values[0] = true;
1124 values[63..66].fill(true);
1125 values[128..].fill(true);
1126 assert_runs(&values, 3);
1127
1128 let mut offset_values = vec![false; 135];
1130 offset_values[3..133].copy_from_slice(&values);
1131 let offset_mask: BooleanBuffer = offset_values.into_iter().collect();
1132 assert_eq!(count_true_runs(&offset_mask.slice(3, 130)), 3);
1133 }
1134
1135 #[test]
1136 fn test_zip_dictionaries_whose_combined_values_overflow_the_key_type() {
1137 let values: Vec<String> = (0..100).map(|i| format!("v{i}")).collect();
1140 let dictionary = |keys: Vec<i8>| -> ArrayRef {
1141 Arc::new(
1142 DictionaryArray::<Int8Type>::try_new(
1143 Int8Array::from(keys),
1144 Arc::new(StringArray::from(values.clone())),
1145 )
1146 .unwrap(),
1147 )
1148 };
1149 let mask = BooleanArray::from(vec![Some(true), Some(false), None, Some(true)]);
1150 let truthy = dictionary(vec![0, 1, 2, 3]);
1151 let falsy = dictionary(vec![10, 20, 30, 40]);
1152 let truthy_scalar = Scalar::new(dictionary(vec![7]));
1153 let falsy_scalar = Scalar::new(dictionary(vec![5]));
1154
1155 let cases: [(&dyn Datum, &dyn Datum, [&str; 4]); 4] = [
1156 (&truthy, &falsy, ["v0", "v20", "v30", "v3"]),
1157 (&truthy, &falsy_scalar, ["v0", "v5", "v5", "v3"]),
1158 (&truthy_scalar, &falsy, ["v7", "v20", "v30", "v7"]),
1159 (&truthy_scalar, &falsy_scalar, ["v7", "v5", "v5", "v7"]),
1160 ];
1161 for (truthy_datum, falsy_datum, expected) in cases {
1162 let actual = zip(&mask, truthy_datum, falsy_datum).unwrap();
1163 assert_eq!(actual.data_type(), truthy.data_type());
1164 let actual: Vec<_> = actual
1165 .as_dictionary::<Int8Type>()
1166 .downcast_dict::<StringArray>()
1167 .unwrap()
1168 .into_iter()
1169 .collect();
1170 assert_eq!(actual, expected.map(Some));
1171 }
1172 }
1173
1174 #[test]
1175 fn test_should_use_interleave() {
1176 let short: BooleanBuffer = (0..64).map(|i| i % 2 == 0).collect();
1177 assert!(!should_use_interleave(&short));
1178
1179 let fragmented: BooleanBuffer = (0..8192).map(|i| i % 2 == 0).collect();
1180 assert!(should_use_interleave(&fragmented));
1181
1182 let long_runs: BooleanBuffer = (0..8192).map(|i| i < 4096).collect();
1183 assert!(!should_use_interleave(&long_runs));
1184
1185 let sparse: BooleanBuffer = (0..8192).map(|i| i % 10 == 0).collect();
1186 assert!(!should_use_interleave(&sparse));
1187
1188 let dense: BooleanBuffer = (0..8192).map(|i| i % 10 != 0).collect();
1189 assert!(!should_use_interleave(&dense));
1190
1191 let fragmented_head: BooleanBuffer = (0..8192).map(|i| i < 256 && i % 2 == 0).collect();
1192 assert!(!should_use_interleave(&fragmented_head));
1193
1194 let fragmented_edges: BooleanBuffer = (0..8192)
1195 .map(|i| !(256..7936).contains(&i) && i % 2 == 0)
1196 .collect();
1197 assert!(!should_use_interleave(&fragmented_edges));
1198
1199 let offset: BooleanBuffer = (0..8195).map(|i| i >= 3 && i % 2 == 1).collect();
1201 assert!(should_use_interleave(&offset.slice(3, 8192)));
1202 }
1203
1204 #[test]
1205 fn test_interleave_arrays() {
1206 let mask = BooleanArray::from(vec![Some(true), None, Some(true), Some(false)]);
1207 let mask = maybe_prep_null_mask_filter(&mask);
1208 let truthy = Int32Array::from(vec![Some(1), None, Some(3), Some(4)]).to_data();
1209 let falsy = Int32Array::from(vec![Some(10), Some(20), None, Some(40)]).to_data();
1210 let expected = Int32Array::from(vec![Some(1), Some(20), Some(3), Some(40)]);
1211
1212 let actual = interleave_arrays(&mask, &truthy, false, &falsy, false).unwrap();
1213 assert_eq!(actual.as_primitive::<Int32Type>(), &expected);
1214 }
1215
1216 #[test]
1217 fn test_zip_fragmented_array_mask() {
1218 let mask: BooleanArray = (0..8192)
1219 .map(|i| match i % 3 {
1220 0 => Some(true),
1221 1 => Some(false),
1222 _ => None,
1223 })
1224 .collect();
1225 let truthy: Int32Array = (0..8192).map(|i| (i % 7 != 0).then_some(i)).collect();
1226 let falsy: Int32Array = (0..8192).map(|i| (i % 11 != 0).then_some(-i)).collect();
1227 let expected: Int32Array = (0..8192)
1228 .map(|i| {
1229 let array = if i % 3 == 0 { &truthy } else { &falsy };
1230 array.is_valid(i).then(|| array.value(i))
1231 })
1232 .collect();
1233
1234 let actual = zip(&mask, &truthy, &falsy).unwrap();
1235 assert_eq!(actual.as_primitive::<Int32Type>(), &expected);
1236 }
1237
1238 #[test]
1239 fn test_zip_kernel_one() {
1240 let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1241 let b = Int32Array::from(vec![None, Some(3), Some(6), Some(7), Some(3)]);
1242 let mask = BooleanArray::from(vec![true, true, false, false, true]);
1243 let out = zip(&mask, &a, &b).unwrap();
1244 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1245 let expected = Int32Array::from(vec![Some(5), None, Some(6), Some(7), Some(1)]);
1246 assert_eq!(actual, &expected);
1247 }
1248
1249 #[test]
1250 fn test_zip_kernel_two() {
1251 let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1252 let b = Int32Array::from(vec![None, Some(3), Some(6), Some(7), Some(3)]);
1253 let mask = BooleanArray::from(vec![false, false, true, true, false]);
1254 let out = zip(&mask, &a, &b).unwrap();
1255 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1256 let expected = Int32Array::from(vec![None, Some(3), Some(7), None, Some(3)]);
1257 assert_eq!(actual, &expected);
1258 }
1259
1260 #[test]
1261 fn test_zip_kernel_scalar_falsy_1() {
1262 let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1263
1264 let fallback = Scalar::new(Int32Array::from_value(42, 1));
1265
1266 let mask = BooleanArray::from(vec![true, true, false, false, true]);
1267 let out = zip(&mask, &a, &fallback).unwrap();
1268 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1269 let expected = Int32Array::from(vec![Some(5), None, Some(42), Some(42), Some(1)]);
1270 assert_eq!(actual, &expected);
1271 }
1272
1273 #[test]
1274 fn test_zip_kernel_scalar_falsy_2() {
1275 let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1276
1277 let fallback = Scalar::new(Int32Array::from_value(42, 1));
1278
1279 let mask = BooleanArray::from(vec![false, false, true, true, false]);
1280 let out = zip(&mask, &a, &fallback).unwrap();
1281 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1282 let expected = Int32Array::from(vec![Some(42), Some(42), Some(7), None, Some(42)]);
1283 assert_eq!(actual, &expected);
1284 }
1285
1286 #[test]
1287 fn test_zip_kernel_scalar_truthy_1() {
1288 let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1289
1290 let fallback = Scalar::new(Int32Array::from_value(42, 1));
1291
1292 let mask = BooleanArray::from(vec![true, true, false, false, true]);
1293 let out = zip(&mask, &fallback, &a).unwrap();
1294 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1295 let expected = Int32Array::from(vec![Some(42), Some(42), Some(7), None, Some(42)]);
1296 assert_eq!(actual, &expected);
1297 }
1298
1299 #[test]
1300 fn test_zip_kernel_scalar_truthy_2() {
1301 let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1302
1303 let fallback = Scalar::new(Int32Array::from_value(42, 1));
1304
1305 let mask = BooleanArray::from(vec![false, false, true, true, false]);
1306 let out = zip(&mask, &fallback, &a).unwrap();
1307 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1308 let expected = Int32Array::from(vec![Some(5), None, Some(42), Some(42), Some(1)]);
1309 assert_eq!(actual, &expected);
1310 }
1311
1312 #[test]
1313 fn test_zip_kernel_scalar_both_mask_ends_with_true() {
1314 let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1315 let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
1316
1317 let mask = BooleanArray::from(vec![true, true, false, false, true]);
1318 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1319 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1320 let expected = Int32Array::from(vec![Some(42), Some(42), Some(123), Some(123), Some(42)]);
1321 assert_eq!(actual, &expected);
1322 }
1323
1324 #[test]
1325 fn test_zip_kernel_scalar_both_mask_ends_with_false() {
1326 let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1327 let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
1328
1329 let mask = BooleanArray::from(vec![true, true, false, true, false, false]);
1330 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1331 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1332 let expected = Int32Array::from(vec![
1333 Some(42),
1334 Some(42),
1335 Some(123),
1336 Some(42),
1337 Some(123),
1338 Some(123),
1339 ]);
1340 assert_eq!(actual, &expected);
1341 }
1342
1343 #[test]
1344 fn test_zip_kernel_primitive_scalar_none_1() {
1345 let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1346 let scalar_falsy = Scalar::new(Int32Array::new_null(1));
1347
1348 let mask = BooleanArray::from(vec![true, true, false, false, true]);
1349 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1350 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1351 let expected = Int32Array::from(vec![Some(42), Some(42), None, None, Some(42)]);
1352 assert_eq!(actual, &expected);
1353 }
1354
1355 #[test]
1356 fn test_zip_kernel_primitive_scalar_none_2() {
1357 let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1358 let scalar_falsy = Scalar::new(Int32Array::new_null(1));
1359
1360 let mask = BooleanArray::from(vec![false, false, true, true, false]);
1361 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1362 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1363 let expected = Int32Array::from(vec![None, None, Some(42), Some(42), None]);
1364 assert_eq!(actual, &expected);
1365 }
1366
1367 #[test]
1368 fn test_zip_kernel_primitive_scalar_both_null() {
1369 let scalar_truthy = Scalar::new(Int32Array::new_null(1));
1370 let scalar_falsy = Scalar::new(Int32Array::new_null(1));
1371
1372 let mask = BooleanArray::from(vec![false, false, true, true, false]);
1373 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1374 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1375 let expected = Int32Array::from(vec![None, None, None, None, None]);
1376 assert_eq!(actual, &expected);
1377 }
1378
1379 #[test]
1380 fn test_zip_primitive_array_with_nulls_is_mask_should_be_treated_as_false() {
1381 let truthy = Int32Array::from_iter_values(vec![1, 2, 3, 4, 5, 6]);
1382 let falsy = Int32Array::from_iter_values(vec![7, 8, 9, 10, 11, 12]);
1383
1384 let mask = {
1385 let booleans = BooleanBuffer::from(vec![true, true, false, true, false, false]);
1386 let nulls = NullBuffer::from(vec![
1387 true, true, true,
1388 false, true, true,
1390 ]);
1391 BooleanArray::new(booleans, Some(nulls))
1392 };
1393 let out = zip(&mask, &truthy, &falsy).unwrap();
1394 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1395 let expected = Int32Array::from(vec![
1396 Some(1),
1397 Some(2),
1398 Some(9),
1399 Some(10), Some(11),
1401 Some(12),
1402 ]);
1403 assert_eq!(actual, &expected);
1404 }
1405
1406 #[test]
1407 fn test_zip_kernel_primitive_scalar_with_boolean_array_mask_with_nulls_should_be_treated_as_false()
1408 {
1409 let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1410 let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
1411
1412 let mask = {
1413 let booleans = BooleanBuffer::from(vec![true, true, false, true, false, false]);
1414 let nulls = NullBuffer::from(vec![
1415 true, true, true,
1416 false, true, true,
1418 ]);
1419 BooleanArray::new(booleans, Some(nulls))
1420 };
1421 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1422 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1423 let expected = Int32Array::from(vec![
1424 Some(42),
1425 Some(42),
1426 Some(123),
1427 Some(123), Some(123),
1429 Some(123),
1430 ]);
1431 assert_eq!(actual, &expected);
1432 }
1433
1434 #[test]
1435 fn test_zip_string_array_with_nulls_is_mask_should_be_treated_as_false() {
1436 let truthy = StringArray::from_iter_values(vec!["1", "2", "3", "4", "5", "6"]);
1437 let falsy = StringArray::from_iter_values(vec!["7", "8", "9", "10", "11", "12"]);
1438
1439 let mask = {
1440 let booleans = BooleanBuffer::from(vec![true, true, false, true, false, false]);
1441 let nulls = NullBuffer::from(vec![
1442 true, true, true,
1443 false, true, true,
1445 ]);
1446 BooleanArray::new(booleans, Some(nulls))
1447 };
1448 let out = zip(&mask, &truthy, &falsy).unwrap();
1449 let actual = out.as_string::<i32>();
1450 let expected = StringArray::from_iter_values(vec![
1451 "1", "2", "9", "10", "11", "12",
1453 ]);
1454 assert_eq!(actual, &expected);
1455 }
1456
1457 #[test]
1458 fn test_zip_kernel_large_string_scalar_with_boolean_array_mask_with_nulls_should_be_treated_as_false()
1459 {
1460 let scalar_truthy = Scalar::new(LargeStringArray::from_iter_values(["test"]));
1461 let scalar_falsy = Scalar::new(LargeStringArray::from_iter_values(["something else"]));
1462
1463 let mask = {
1464 let booleans = BooleanBuffer::from(vec![true, true, false, true, false, false]);
1465 let nulls = NullBuffer::from(vec![
1466 true, true, true,
1467 false, true, true,
1469 ]);
1470 BooleanArray::new(booleans, Some(nulls))
1471 };
1472 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1473 let actual = out.as_any().downcast_ref::<LargeStringArray>().unwrap();
1474 let expected = LargeStringArray::from_iter(vec![
1475 Some("test"),
1476 Some("test"),
1477 Some("something else"),
1478 Some("something else"), Some("something else"),
1480 Some("something else"),
1481 ]);
1482 assert_eq!(actual, &expected);
1483 }
1484
1485 #[test]
1486 fn test_zip_kernel_bytes_scalar_none_1() {
1487 let scalar_truthy = Scalar::new(StringArray::from_iter_values(["hello"]));
1488 let scalar_falsy = Scalar::new(StringArray::new_null(1));
1489
1490 let mask = BooleanArray::from(vec![true, true, false, false, true]);
1491 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1492 let actual = out.as_any().downcast_ref::<StringArray>().unwrap();
1493 let expected = StringArray::from_iter(vec![
1494 Some("hello"),
1495 Some("hello"),
1496 None,
1497 None,
1498 Some("hello"),
1499 ]);
1500 assert_eq!(actual, &expected);
1501 }
1502
1503 #[test]
1504 fn test_zip_kernel_bytes_scalar_none_2() {
1505 let scalar_truthy = Scalar::new(StringArray::new_null(1));
1506 let scalar_falsy = Scalar::new(StringArray::from_iter_values(["hello"]));
1507
1508 let mask = BooleanArray::from(vec![true, true, false, false, true]);
1509 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1510 let actual = out.as_any().downcast_ref::<StringArray>().unwrap();
1511 let expected = StringArray::from_iter(vec![None, None, Some("hello"), Some("hello"), None]);
1512 assert_eq!(actual, &expected);
1513 }
1514
1515 #[test]
1516 fn test_zip_kernel_bytes_scalar_both() {
1517 let scalar_truthy = Scalar::new(StringArray::from_iter_values(["test"]));
1518 let scalar_falsy = Scalar::new(StringArray::from_iter_values(["something else"]));
1519
1520 let mask = BooleanArray::from(vec![true, true, false, true, false, false]);
1522 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1523 let actual = out.as_any().downcast_ref::<StringArray>().unwrap();
1524 let expected = StringArray::from_iter(vec![
1525 Some("test"),
1526 Some("test"),
1527 Some("something else"),
1528 Some("test"),
1529 Some("something else"),
1530 Some("something else"),
1531 ]);
1532 assert_eq!(actual, &expected);
1533 }
1534
1535 #[test]
1536 fn test_zip_scalar_bytes_offset_overflow() {
1537 let value = vec![0_u8; 65_536];
1540 let large = Scalar::new(BinaryArray::from_iter_values([value.as_slice()]));
1541 let empty = Scalar::new(BinaryArray::from_iter_values([b"".as_slice()]));
1542 let mask = BooleanArray::from_iter((0..65_536).map(|i| Some(i % 2 == 0)));
1543
1544 for (truthy, falsy) in [(&large, &empty), (&empty, &large)] {
1545 assert!(matches!(
1546 zip(&mask, truthy, falsy),
1547 Err(ArrowError::OffsetOverflowError(2_147_483_648))
1548 ));
1549 }
1550 }
1551
1552 #[test]
1553 fn test_zip_scalar_bytes_only_taking_one_side() {
1554 let mask_len = 5;
1555 let all_true_mask = BooleanArray::from(vec![true; mask_len]);
1556 let all_false_mask = BooleanArray::from(vec![false; mask_len]);
1557
1558 let null_scalar = Scalar::new(StringArray::new_null(1));
1559 let non_null_scalar_1 = Scalar::new(StringArray::from_iter_values(["test"]));
1560 let non_null_scalar_2 = Scalar::new(StringArray::from_iter_values(["something else"]));
1561
1562 {
1563 let out = zip(&all_true_mask, &null_scalar, &non_null_scalar_1).unwrap();
1566 let actual = out.as_string::<i32>();
1567 let expected = StringArray::from_iter(std::iter::repeat_n(None::<&str>, mask_len));
1568 assert_eq!(actual, &expected);
1569 }
1570
1571 {
1572 let out = zip(&all_false_mask, &null_scalar, &non_null_scalar_1).unwrap();
1575 let actual = out.as_string::<i32>();
1576 let expected = StringArray::from_iter(std::iter::repeat_n(Some("test"), mask_len));
1577 assert_eq!(actual, &expected);
1578 }
1579
1580 {
1581 let out = zip(&all_true_mask, &non_null_scalar_1, &null_scalar).unwrap();
1584 let actual = out.as_string::<i32>();
1585 let expected = StringArray::from_iter(std::iter::repeat_n(Some("test"), mask_len));
1586 assert_eq!(actual, &expected);
1587 }
1588
1589 {
1590 let out = zip(&all_false_mask, &non_null_scalar_1, &null_scalar).unwrap();
1593 let actual = out.as_string::<i32>();
1594 let expected = StringArray::from_iter(std::iter::repeat_n(None::<&str>, mask_len));
1595 assert_eq!(actual, &expected);
1596 }
1597
1598 {
1599 let out = zip(&all_true_mask, &non_null_scalar_1, &non_null_scalar_2).unwrap();
1602 let actual = out.as_string::<i32>();
1603 let expected = StringArray::from_iter(std::iter::repeat_n(Some("test"), mask_len));
1604 assert_eq!(actual, &expected);
1605 }
1606
1607 {
1608 let out = zip(&all_false_mask, &non_null_scalar_1, &non_null_scalar_2).unwrap();
1611 let actual = out.as_string::<i32>();
1612 let expected =
1613 StringArray::from_iter(std::iter::repeat_n(Some("something else"), mask_len));
1614 assert_eq!(actual, &expected);
1615 }
1616
1617 {
1618 let mask = BooleanArray::from(vec![true, false, true, false, true]);
1621 let out = zip(&mask, &null_scalar, &null_scalar).unwrap();
1622 let actual = out.as_string::<i32>();
1623 let expected = StringArray::from_iter(std::iter::repeat_n(None::<&str>, mask_len));
1624 assert_eq!(actual, &expected);
1625 }
1626 }
1627
1628 #[test]
1629 fn test_scalar_zipper() {
1630 let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1631 let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
1632
1633 let mask = BooleanArray::from(vec![false, false, true, true, false]);
1634
1635 let scalar_zipper = ScalarZipper::try_new(&scalar_truthy, &scalar_falsy).unwrap();
1636 let out = scalar_zipper.zip(&mask).unwrap();
1637 let actual = out.as_primitive::<Int32Type>();
1638 let expected = Int32Array::from(vec![Some(123), Some(123), Some(42), Some(42), Some(123)]);
1639 assert_eq!(actual, &expected);
1640
1641 let mask = BooleanArray::from(vec![true, false, true]);
1643 let out = scalar_zipper.zip(&mask).unwrap();
1644 let actual = out.as_primitive::<Int32Type>();
1645 let expected = Int32Array::from(vec![Some(42), Some(123), Some(42)]);
1646 assert_eq!(actual, &expected);
1647 }
1648
1649 #[test]
1650 fn test_zip_kernel_scalar_strings() {
1651 let scalar_truthy = Scalar::new(StringArray::from(vec!["hello"]));
1652 let scalar_falsy = Scalar::new(StringArray::from(vec!["world"]));
1653
1654 let mask = BooleanArray::from(vec![true, false, true, false, true]);
1655 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1656 let actual = out.as_string::<i32>();
1657 let expected = StringArray::from(vec![
1658 Some("hello"),
1659 Some("world"),
1660 Some("hello"),
1661 Some("world"),
1662 Some("hello"),
1663 ]);
1664 assert_eq!(actual, &expected);
1665 }
1666
1667 #[test]
1668 fn test_zip_kernel_scalar_binary() {
1669 let truthy_bytes: &[u8] = b"\xFF\xFE\xFD";
1670 let falsy_bytes: &[u8] = b"world";
1671 let scalar_truthy = Scalar::new(BinaryArray::from_iter_values(
1672 vec![truthy_bytes],
1674 ));
1675 let scalar_falsy = Scalar::new(BinaryArray::from_iter_values(vec![falsy_bytes]));
1676
1677 let mask = BooleanArray::from(vec![true, false, true, false, true]);
1678 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1679 let actual = out.as_binary::<i32>();
1680 let expected = BinaryArray::from(vec![
1681 Some(truthy_bytes),
1682 Some(falsy_bytes),
1683 Some(truthy_bytes),
1684 Some(falsy_bytes),
1685 Some(truthy_bytes),
1686 ]);
1687 assert_eq!(actual, &expected);
1688 }
1689
1690 #[test]
1691 fn test_zip_kernel_scalar_large_binary() {
1692 let truthy_bytes: &[u8] = b"hey";
1693 let falsy_bytes: &[u8] = b"world";
1694 let scalar_truthy = Scalar::new(LargeBinaryArray::from_iter_values(vec![truthy_bytes]));
1695 let scalar_falsy = Scalar::new(LargeBinaryArray::from_iter_values(vec![falsy_bytes]));
1696
1697 let mask = BooleanArray::from(vec![true, false, true, false, true]);
1698 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1699 let actual = out.as_binary::<i64>();
1700 let expected = LargeBinaryArray::from(vec![
1701 Some(truthy_bytes),
1702 Some(falsy_bytes),
1703 Some(truthy_bytes),
1704 Some(falsy_bytes),
1705 Some(truthy_bytes),
1706 ]);
1707 assert_eq!(actual, &expected);
1708 }
1709
1710 #[test]
1712 fn test_zip_decimal_with_custom_precision_and_scale() {
1713 let arr = Decimal128Array::from_iter_values([12345, 456, 7890, -123223423432432])
1714 .with_precision_and_scale(20, 2)
1715 .unwrap();
1716
1717 let arr: ArrayRef = Arc::new(arr);
1718
1719 let scalar_1 = Scalar::new(arr.slice(0, 1));
1720 let scalar_2 = Scalar::new(arr.slice(1, 1));
1721 let null_scalar = Scalar::new(new_null_array(arr.data_type(), 1));
1722 let array_1: ArrayRef = arr.slice(0, 2);
1723 let array_2: ArrayRef = arr.slice(2, 2);
1724
1725 test_zip_output_data_types_for_input(scalar_1, scalar_2, null_scalar, array_1, array_2);
1726 }
1727
1728 #[test]
1730 fn test_zip_timestamp_with_timezone() {
1731 let arr = TimestampSecondArray::from(vec![0, 1000, 2000, 4000])
1732 .with_timezone("+01:00".to_string());
1733
1734 let arr: ArrayRef = Arc::new(arr);
1735
1736 let scalar_1 = Scalar::new(arr.slice(0, 1));
1737 let scalar_2 = Scalar::new(arr.slice(1, 1));
1738 let null_scalar = Scalar::new(new_null_array(arr.data_type(), 1));
1739 let array_1: ArrayRef = arr.slice(0, 2);
1740 let array_2: ArrayRef = arr.slice(2, 2);
1741
1742 test_zip_output_data_types_for_input(scalar_1, scalar_2, null_scalar, array_1, array_2);
1743 }
1744
1745 fn test_zip_output_data_types_for_input(
1746 scalar_1: Scalar<ArrayRef>,
1747 scalar_2: Scalar<ArrayRef>,
1748 null_scalar: Scalar<ArrayRef>,
1749 array_1: ArrayRef,
1750 array_2: ArrayRef,
1751 ) {
1752 test_zip_output_data_type(&scalar_1, &scalar_2, 10);
1754
1755 test_zip_output_data_type(&null_scalar, &scalar_1, 10);
1757 test_zip_output_data_type(&scalar_1, &null_scalar, 10);
1758
1759 test_zip_output_data_type(&array_1.as_ref(), &scalar_1, array_1.len());
1761 test_zip_output_data_type(&scalar_1, &array_1.as_ref(), array_1.len());
1762
1763 test_zip_output_data_type(&array_1.as_ref(), &null_scalar, array_1.len());
1765
1766 test_zip_output_data_type(&null_scalar, &array_1.as_ref(), array_1.len());
1767
1768 test_zip_output_data_type(&array_1.as_ref(), &array_2.as_ref(), array_1.len());
1770 }
1771
1772 fn test_zip_output_data_type(truthy: &dyn Datum, falsy: &dyn Datum, mask_length: usize) {
1773 let expected_data_type = truthy.get().0.data_type().clone();
1774 assert_eq!(&expected_data_type, falsy.get().0.data_type());
1775
1776 let mask_all_true = BooleanArray::from(vec![true; mask_length]);
1778 let mask_all_false = BooleanArray::from(vec![false; mask_length]);
1779 let mask_some_true_and_false =
1780 BooleanArray::from((0..mask_length).map(|i| i % 2 == 0).collect::<Vec<bool>>());
1781
1782 for mask in [&mask_all_true, &mask_all_false, &mask_some_true_and_false] {
1783 let out = zip(mask, truthy, falsy).unwrap();
1784 assert_eq!(out.data_type(), &expected_data_type);
1785 }
1786 }
1787
1788 #[test]
1789 fn zip_scalar_fallback_impl() {
1790 let truthy_list_item_scalar = Some(vec![Some(1), None, Some(3)]);
1791 let truthy_list_array_scalar =
1792 Scalar::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1793 truthy_list_item_scalar.clone(),
1794 ]));
1795 let falsy_list_item_scalar = Some(vec![None, Some(2), Some(4)]);
1796 let falsy_list_array_scalar =
1797 Scalar::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1798 falsy_list_item_scalar.clone(),
1799 ]));
1800 let mask = BooleanArray::from(vec![true, false, true, false, false, true, false]);
1801 let out = zip(&mask, &truthy_list_array_scalar, &falsy_list_array_scalar).unwrap();
1802 let actual = out.as_list::<i32>();
1803
1804 let expected = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1805 truthy_list_item_scalar.clone(),
1806 falsy_list_item_scalar.clone(),
1807 truthy_list_item_scalar.clone(),
1808 falsy_list_item_scalar.clone(),
1809 falsy_list_item_scalar.clone(),
1810 truthy_list_item_scalar.clone(),
1811 falsy_list_item_scalar.clone(),
1812 ]);
1813 assert_eq!(actual, &expected);
1814 }
1815
1816 #[test]
1817 fn test_zip_kernel_scalar_strings_array_view() {
1818 let scalar_truthy = Scalar::new(StringViewArray::from(vec!["hello"]));
1819 let scalar_falsy = Scalar::new(StringViewArray::from(vec!["world"]));
1820
1821 let mask = BooleanArray::from(vec![true, false, true, false]);
1822 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1823 let actual = out.as_string_view();
1824 let expected = StringViewArray::from(vec![
1825 Some("hello"),
1826 Some("world"),
1827 Some("hello"),
1828 Some("world"),
1829 ]);
1830 assert_eq!(actual, &expected);
1831 }
1832
1833 #[test]
1834 fn test_zip_kernel_scalar_binary_array_view() {
1835 let scalar_truthy = Scalar::new(BinaryViewArray::from_iter_values(vec![b"hello"]));
1836 let scalar_falsy = Scalar::new(BinaryViewArray::from_iter_values(vec![b"world"]));
1837
1838 let mask = BooleanArray::from(vec![true, false]);
1839 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1840 let actual = out.as_byte_view();
1841 let expected = BinaryViewArray::from_iter_values(vec![b"hello", b"world"]);
1842 assert_eq!(actual, &expected);
1843 }
1844
1845 #[test]
1846 fn test_zip_kernel_scalar_strings_array_view_with_nulls() {
1847 let scalar_truthy = Scalar::new(StringViewArray::from_iter_values(["hello"]));
1848 let scalar_falsy = Scalar::new(StringViewArray::new_null(1));
1849
1850 let mask = BooleanArray::from(vec![true, true, false, false, true]);
1851 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1852 let actual = out.as_any().downcast_ref::<StringViewArray>().unwrap();
1853 let expected = StringViewArray::from_iter(vec![
1854 Some("hello"),
1855 Some("hello"),
1856 None,
1857 None,
1858 Some("hello"),
1859 ]);
1860 assert_eq!(actual, &expected);
1861 }
1862
1863 #[test]
1864 fn test_zip_kernel_scalar_strings_array_view_all_true_null() {
1865 let scalar_truthy = Scalar::new(StringViewArray::new_null(1));
1866 let scalar_falsy = Scalar::new(StringViewArray::new_null(1));
1867 let mask = BooleanArray::from(vec![true, true]);
1868 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1869 let actual = out.as_any().downcast_ref::<StringViewArray>().unwrap();
1870 let expected = StringViewArray::from_iter(vec![None::<String>, None]);
1871 assert_eq!(actual, &expected);
1872 }
1873
1874 #[test]
1875 fn test_zip_kernel_scalar_strings_array_view_all_false_null() {
1876 let scalar_truthy = Scalar::new(StringViewArray::new_null(1));
1877 let scalar_falsy = Scalar::new(StringViewArray::new_null(1));
1878 let mask = BooleanArray::from(vec![false, false]);
1879 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1880 let actual = out.as_any().downcast_ref::<StringViewArray>().unwrap();
1881 let expected = StringViewArray::from_iter(vec![None::<String>, None]);
1882 assert_eq!(actual, &expected);
1883 }
1884
1885 #[test]
1886 fn test_zip_kernel_scalar_string_array_view_all_true() {
1887 let scalar_truthy = Scalar::new(StringViewArray::from(vec!["hello"]));
1888 let scalar_falsy = Scalar::new(StringViewArray::from(vec!["world"]));
1889
1890 let mask = BooleanArray::from(vec![true, true]);
1891 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1892 let actual = out.as_string_view();
1893 let expected = StringViewArray::from(vec![Some("hello"), Some("hello")]);
1894 assert_eq!(actual, &expected);
1895 }
1896
1897 #[test]
1898 fn test_zip_kernel_scalar_string_array_view_all_false() {
1899 let scalar_truthy = Scalar::new(StringViewArray::from(vec!["hello"]));
1900 let scalar_falsy = Scalar::new(StringViewArray::from(vec!["world"]));
1901
1902 let mask = BooleanArray::from(vec![false, false]);
1903 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1904 let actual = out.as_string_view();
1905 let expected = StringViewArray::from(vec![Some("world"), Some("world")]);
1906 assert_eq!(actual, &expected);
1907 }
1908
1909 #[test]
1910 fn test_zip_kernel_scalar_strings_large_strings() {
1911 let scalar_truthy = Scalar::new(StringViewArray::from(vec!["longer than 12 bytes"]));
1912 let scalar_falsy = Scalar::new(StringViewArray::from(vec!["another longer than 12 bytes"]));
1913
1914 let mask = BooleanArray::from(vec![true, false]);
1915 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1916 let actual = out.as_string_view();
1917 let expected = StringViewArray::from(vec![
1918 Some("longer than 12 bytes"),
1919 Some("another longer than 12 bytes"),
1920 ]);
1921 assert_eq!(actual, &expected);
1922 }
1923
1924 #[test]
1925 fn test_zip_kernel_scalar_strings_array_view_large_short_strings() {
1926 let scalar_truthy = Scalar::new(StringViewArray::from(vec!["hello"]));
1927 let scalar_falsy = Scalar::new(StringViewArray::from(vec!["longer than 12 bytes"]));
1928
1929 let mask = BooleanArray::from(vec![true, false, true, false]);
1930 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1931 let actual = out.as_string_view();
1932 let expected = StringViewArray::from(vec![
1933 Some("hello"),
1934 Some("longer than 12 bytes"),
1935 Some("hello"),
1936 Some("longer than 12 bytes"),
1937 ]);
1938 assert_eq!(actual, &expected);
1939 }
1940 #[test]
1941 fn test_zip_kernel_scalar_strings_array_view_large_all_true() {
1942 let scalar_truthy = Scalar::new(StringViewArray::from(vec!["longer than 12 bytes"]));
1943 let scalar_falsy = Scalar::new(StringViewArray::from(vec!["another longer than 12 bytes"]));
1944
1945 let mask = BooleanArray::from(vec![true, true]);
1946 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1947 let actual = out.as_string_view();
1948 let expected = StringViewArray::from(vec![
1949 Some("longer than 12 bytes"),
1950 Some("longer than 12 bytes"),
1951 ]);
1952 assert_eq!(actual, &expected);
1953 }
1954
1955 #[test]
1956 fn test_zip_kernel_scalar_strings_array_view_large_all_false() {
1957 let scalar_truthy = Scalar::new(StringViewArray::from(vec!["longer than 12 bytes"]));
1958 let scalar_falsy = Scalar::new(StringViewArray::from(vec!["another longer than 12 bytes"]));
1959
1960 let mask = BooleanArray::from(vec![false, false]);
1961 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1962 let actual = out.as_string_view();
1963 let expected = StringViewArray::from(vec![
1964 Some("another longer than 12 bytes"),
1965 Some("another longer than 12 bytes"),
1966 ]);
1967 assert_eq!(actual, &expected);
1968 }
1969
1970 #[test]
1971 fn test_zip_kernel_scalar_strings_array_view_sliced_inline() {
1972 let truthy =
1974 StringViewArray::from(vec!["x", "a string longer than twelve bytes"]).slice(0, 1);
1975 let falsy =
1976 StringViewArray::from(vec!["hello", "another string longer than twelve"]).slice(0, 1);
1977
1978 let mask = BooleanArray::from(vec![true, false, true, false]);
1979 let out = zip(&mask, &Scalar::new(&truthy), &Scalar::new(&falsy)).unwrap();
1980 out.to_data().validate_full().unwrap();
1981 let actual = out.as_string_view();
1982 let expected = StringViewArray::from(vec!["x", "hello", "x", "hello"]);
1983 assert_eq!(actual, &expected);
1984 }
1985
1986 #[test]
1987 fn test_zip_kernel_scalar_strings_array_view_sliced_long_falsy() {
1988 let truthy =
1990 StringViewArray::from(vec!["a string longer than twelve bytes", "x"]).slice(0, 1);
1991 let falsy = StringViewArray::from(vec!["y", "thirteen byte"]).slice(1, 1);
1992
1993 let mask = BooleanArray::from(vec![true, false]);
1994 let out = zip(&mask, &Scalar::new(&truthy), &Scalar::new(&falsy)).unwrap();
1995 out.to_data().validate_full().unwrap();
1996 let actual = out.as_string_view();
1997 let expected =
1998 StringViewArray::from(vec!["a string longer than twelve bytes", "thirteen byte"]);
1999 assert_eq!(actual, &expected);
2000 }
2001
2002 #[test]
2003 fn test_zip_kernel_scalar_binary_array_view_sliced_inline_max_len() {
2004 let truthy = BinaryViewArray::from_iter_values([
2006 b"x".as_slice(),
2007 b"another value longer than twelve",
2008 ])
2009 .slice(0, 1);
2010 let falsy = BinaryViewArray::from_iter_values([
2011 [0xFFu8; 12].as_slice(),
2012 b"a value longer than twelve bytes",
2013 ])
2014 .slice(0, 1);
2015
2016 let mask = BooleanArray::from(vec![true, false]);
2017 let out = zip(&mask, &Scalar::new(&truthy), &Scalar::new(&falsy)).unwrap();
2018 out.to_data().validate_full().unwrap();
2019 let actual = out.as_binary_view();
2020 let expected = BinaryViewArray::from_iter_values([b"x".as_slice(), &[0xFFu8; 12]]);
2021 assert_eq!(actual, &expected);
2022 }
2023
2024 fn random_primitive(
2026 rng: &mut StdRng,
2027 data_type: &DataType,
2028 len: usize,
2029 offset: usize,
2030 null_fraction: f64,
2031 ) -> ArrayRef {
2032 let width = data_type.primitive_width().unwrap();
2033 let total = len + offset;
2034 let mut bytes = MutableBuffer::from_len_zeroed(total * width);
2035 bytes.iter_mut().for_each(|byte| *byte = rng.random());
2036 let nulls = (null_fraction > 0.0)
2037 .then(|| NullBuffer::from_iter((0..total).map(|_| !rng.random_bool(null_fraction))));
2038 let data = ArrayData::builder(data_type.clone())
2039 .len(total)
2040 .add_buffer(bytes.into())
2041 .nulls(nulls)
2042 .build()
2043 .unwrap();
2044 make_array(data).slice(offset, len)
2045 }
2046
2047 #[test]
2048 fn test_zip_primitive_with_scalar_matches_general_path() {
2049 let mut rng = StdRng::seed_from_u64(0x5ca1_ab1e_0dd5_f00d);
2050 let data_types = [
2051 DataType::Int8,
2052 DataType::Int16,
2053 DataType::Float32,
2054 DataType::Int64,
2055 DataType::Timestamp(TimeUnit::Millisecond, Some("+01:00".into())),
2056 DataType::Decimal128(38, 2),
2057 DataType::Interval(IntervalUnit::MonthDayNano),
2058 DataType::Decimal256(76, 3),
2059 ];
2060 for data_type in &data_types {
2061 for len in [0, 1, 63, 64, 65, 130, 1000] {
2062 for mask_shape in 0..5 {
2063 let mask_offset = rng.random_range(0..9);
2064 let mask_bits: Vec<Option<bool>> = (0..len + mask_offset)
2065 .map(|i| match mask_shape {
2066 0 => Some(rng.random_bool(0.5)),
2067 1 => Some(true),
2068 2 => Some(false),
2069 3 => Some(i / 100 % 2 == 0),
2070 _ => rng.random_bool(0.8).then(|| rng.random_bool(0.5)),
2071 })
2072 .collect();
2073 let mask = BooleanArray::from(mask_bits).slice(mask_offset, len);
2074
2075 for array_nulls in [0.0, 0.3] {
2076 let array_offset = rng.random_range(0..5);
2077 let array =
2078 random_primitive(&mut rng, data_type, len, array_offset, array_nulls);
2079 for scalar_null in [false, true] {
2080 let scalar = random_primitive(
2081 &mut rng,
2082 data_type,
2083 1,
2084 1,
2085 if scalar_null { 1.0 } else { 0.0 },
2086 );
2087 for scalar_is_truthy in [false, true] {
2088 let scalar_datum = Scalar::new(Arc::clone(&scalar));
2089 let (truthy, falsy): (&dyn Datum, &dyn Datum) = if scalar_is_truthy
2090 {
2091 (&scalar_datum, &array)
2092 } else {
2093 (&array, &scalar_datum)
2094 };
2095 let actual = zip(&mask, truthy, falsy).unwrap();
2096 let expected = zip_impl(
2097 &mask,
2098 &truthy.get().0.to_data(),
2099 scalar_is_truthy,
2100 &falsy.get().0.to_data(),
2101 !scalar_is_truthy,
2102 )
2103 .unwrap();
2104 assert_eq!(
2105 &actual, &expected,
2106 "{data_type} len {len} mask shape {mask_shape} array nulls \
2107 {array_nulls} scalar null {scalar_null} scalar truthy \
2108 {scalar_is_truthy}"
2109 );
2110 }
2111 }
2112 }
2113 }
2114 }
2115 }
2116 }
2117}