1use crate::filter::{SlicesIterator, prep_null_mask_filter};
21use arrow_array::cast::AsArray;
22use arrow_array::types::{
23 BinaryType, BinaryViewType, ByteArrayType, ByteViewType, LargeBinaryType, LargeUtf8Type,
24 StringViewType, Utf8Type,
25};
26use arrow_array::*;
27use arrow_buffer::{
28 ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer, OffsetBuffer, ScalarBuffer,
29 ToByteSlice,
30};
31use arrow_data::transform::MutableArrayData;
32use arrow_data::{ArrayData, ByteView, MAX_INLINE_VIEW_LEN};
33use arrow_schema::{ArrowError, DataType};
34use std::fmt::{Debug, Formatter};
35use std::hash::Hash;
36use std::marker::PhantomData;
37use std::ops::Not;
38use std::sync::{Arc, OnceLock};
39
40pub fn zip(
100 mask: &BooleanArray,
101 truthy: &dyn Datum,
102 falsy: &dyn Datum,
103) -> Result<ArrayRef, ArrowError> {
104 let (truthy_array, truthy_is_scalar) = truthy.get();
105 let (falsy_array, falsy_is_scalar) = falsy.get();
106
107 if falsy_is_scalar && truthy_is_scalar {
108 let zipper = ScalarZipper::try_new(truthy, falsy)?;
109 return zipper.zip_impl.create_output(mask);
110 }
111
112 let truthy = truthy_array;
113 let falsy = falsy_array;
114
115 if truthy.data_type() != falsy.data_type() {
116 return Err(ArrowError::InvalidArgumentError(
117 "arguments need to have the same data type".into(),
118 ));
119 }
120
121 if truthy_is_scalar && truthy.len() != 1 {
122 return Err(ArrowError::InvalidArgumentError(
123 "scalar arrays must have 1 element".into(),
124 ));
125 }
126 if !truthy_is_scalar && truthy.len() != mask.len() {
127 return Err(ArrowError::InvalidArgumentError(
128 "all arrays should have the same length".into(),
129 ));
130 }
131 if falsy_is_scalar && falsy.len() != 1 {
132 return Err(ArrowError::InvalidArgumentError(
133 "scalar arrays must have 1 element".into(),
134 ));
135 }
136 if !falsy_is_scalar && falsy.len() != mask.len() {
137 return Err(ArrowError::InvalidArgumentError(
138 "all arrays should have the same length".into(),
139 ));
140 }
141
142 let falsy = falsy.to_data();
143 let truthy = truthy.to_data();
144
145 zip_impl(mask, &truthy, truthy_is_scalar, &falsy, falsy_is_scalar)
146}
147
148fn count_true_runs(mask: &BooleanBuffer) -> usize {
149 let mut slices = 0;
150 let mut previous = 0;
151 for chunk in mask.bit_chunks().iter_padded() {
152 let starts = chunk & !((chunk << 1) | previous);
153 slices += starts.count_ones() as usize;
154 previous = chunk >> 63;
155 }
156 slices
157}
158
159fn should_use_interleave(mask: &BooleanBuffer) -> bool {
160 const MIN_LEN: usize = 1024;
161
162 mask.len() >= MIN_LEN && count_true_runs(mask) > mask.len() / 8
166}
167
168fn interleave_arrays(
169 mask: &BooleanBuffer,
170 truthy: &ArrayData,
171 falsy: &ArrayData,
172) -> Result<ArrayRef, ArrowError> {
173 let truthy = make_array(truthy.clone());
174 let falsy = make_array(falsy.clone());
175 let indices: Vec<_> = mask
176 .iter()
177 .enumerate()
178 .map(|(idx, selected)| (usize::from(!selected), idx))
179 .collect();
180 crate::interleave::interleave(&[truthy.as_ref(), falsy.as_ref()], &indices)
181}
182
183fn zip_impl(
184 mask: &BooleanArray,
185 truthy: &ArrayData,
186 truthy_is_scalar: bool,
187 falsy: &ArrayData,
188 falsy_is_scalar: bool,
189) -> Result<ArrayRef, ArrowError> {
190 let mask_buffer = maybe_prep_null_mask_filter(mask);
191 if !truthy_is_scalar && !falsy_is_scalar && should_use_interleave(&mask_buffer) {
192 return interleave_arrays(&mask_buffer, truthy, falsy);
193 }
194
195 let mut mutable = MutableArrayData::new(vec![truthy, falsy], false, truthy.len());
196
197 let mut filled = 0;
202
203 for (start, end) in SlicesIterator::from(&mask_buffer) {
204 if start > filled {
206 if falsy_is_scalar {
207 for _ in filled..start {
208 mutable.try_extend(1, 0, 1)?;
210 }
211 } else {
212 mutable.try_extend(1, filled, start)?;
213 }
214 }
215 if truthy_is_scalar {
217 for _ in start..end {
218 mutable.try_extend(0, 0, 1)?;
220 }
221 } else {
222 mutable.try_extend(0, start, end)?;
223 }
224 filled = end;
225 }
226 if filled < mask.len() {
228 if falsy_is_scalar {
229 for _ in filled..mask.len() {
230 mutable.try_extend(1, 0, 1)?;
232 }
233 } else {
234 mutable.try_extend(1, filled, mask.len())?;
235 }
236 }
237
238 let data = mutable.freeze();
239 Ok(make_array(data))
240}
241
242#[derive(Debug, Clone)]
264pub struct ScalarZipper {
265 zip_impl: Arc<dyn ZipImpl>,
266}
267
268impl ScalarZipper {
269 pub fn try_new(truthy: &dyn Datum, falsy: &dyn Datum) -> Result<Self, ArrowError> {
277 let (truthy, truthy_is_scalar) = truthy.get();
278 let (falsy, falsy_is_scalar) = falsy.get();
279
280 if truthy.data_type() != falsy.data_type() {
281 return Err(ArrowError::InvalidArgumentError(
282 "arguments need to have the same data type".into(),
283 ));
284 }
285
286 if !truthy_is_scalar {
287 return Err(ArrowError::InvalidArgumentError(
288 "only scalar arrays are supported".into(),
289 ));
290 }
291
292 if !falsy_is_scalar {
293 return Err(ArrowError::InvalidArgumentError(
294 "only scalar arrays are supported".into(),
295 ));
296 }
297
298 if truthy.len() != 1 {
299 return Err(ArrowError::InvalidArgumentError(
300 "scalar arrays must have 1 element".into(),
301 ));
302 }
303 if falsy.len() != 1 {
304 return Err(ArrowError::InvalidArgumentError(
305 "scalar arrays must have 1 element".into(),
306 ));
307 }
308
309 macro_rules! primitive_size_helper {
310 ($t:ty) => {
311 Arc::new(PrimitiveScalarImpl::<$t>::new(truthy, falsy)) as Arc<dyn ZipImpl>
312 };
313 }
314
315 let zip_impl = downcast_primitive! {
316 truthy.data_type() => (primitive_size_helper),
317 DataType::Utf8 => {
318 Arc::new(BytesScalarImpl::<Utf8Type>::new(truthy, falsy)) as Arc<dyn ZipImpl>
319 },
320 DataType::LargeUtf8 => {
321 Arc::new(BytesScalarImpl::<LargeUtf8Type>::new(truthy, falsy)) as Arc<dyn ZipImpl>
322 },
323 DataType::Binary => {
324 Arc::new(BytesScalarImpl::<BinaryType>::new(truthy, falsy)) as Arc<dyn ZipImpl>
325 },
326 DataType::LargeBinary => {
327 Arc::new(BytesScalarImpl::<LargeBinaryType>::new(truthy, falsy)) as Arc<dyn ZipImpl>
328 },
329 DataType::Utf8View => {
330 Arc::new(ByteViewScalarImpl::<StringViewType>::new(truthy, falsy)) as Arc<dyn ZipImpl>
331 },
332 DataType::BinaryView => {
333 Arc::new(ByteViewScalarImpl::<BinaryViewType>::new(truthy, falsy)) as Arc<dyn ZipImpl>
334 },
335 _ => {
336 Arc::new(FallbackImpl::new(truthy, falsy)) as Arc<dyn ZipImpl>
337 },
338 };
339
340 Ok(Self { zip_impl })
341 }
342
343 pub fn zip(&self, mask: &BooleanArray) -> Result<ArrayRef, ArrowError> {
346 self.zip_impl.create_output(mask)
347 }
348}
349
350trait ZipImpl: Debug + Send + Sync {
352 fn create_output(&self, input: &BooleanArray) -> Result<ArrayRef, ArrowError>;
354}
355
356#[derive(Debug, PartialEq)]
357struct FallbackImpl {
358 truthy: ArrayData,
359 falsy: ArrayData,
360}
361
362impl FallbackImpl {
363 fn new(left: &dyn Array, right: &dyn Array) -> Self {
364 Self {
365 truthy: left.to_data(),
366 falsy: right.to_data(),
367 }
368 }
369}
370
371impl ZipImpl for FallbackImpl {
372 fn create_output(&self, predicate: &BooleanArray) -> Result<ArrayRef, ArrowError> {
373 zip_impl(predicate, &self.truthy, true, &self.falsy, true)
374 }
375}
376
377struct PrimitiveScalarImpl<T: ArrowPrimitiveType> {
378 data_type: DataType,
379 truthy: Option<T::Native>,
380 falsy: Option<T::Native>,
381}
382
383impl<T: ArrowPrimitiveType> Debug for PrimitiveScalarImpl<T> {
384 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
385 f.debug_struct("PrimitiveScalarImpl")
386 .field("data_type", &self.data_type)
387 .field("truthy", &self.truthy)
388 .field("falsy", &self.falsy)
389 .finish()
390 }
391}
392
393impl<T: ArrowPrimitiveType> PrimitiveScalarImpl<T> {
394 fn new(truthy: &dyn Array, falsy: &dyn Array) -> Self {
395 Self {
396 data_type: truthy.data_type().clone(),
397 truthy: Self::get_value_from_scalar(truthy),
398 falsy: Self::get_value_from_scalar(falsy),
399 }
400 }
401
402 fn get_value_from_scalar(scalar: &dyn Array) -> Option<T::Native> {
403 if scalar.is_null(0) {
404 None
405 } else {
406 let value = scalar.as_primitive::<T>().value(0);
407
408 Some(value)
409 }
410 }
411
412 fn get_scalar_and_null_buffer_for_single_non_nullable(
416 predicate: BooleanBuffer,
417 value: T::Native,
418 ) -> (Vec<T::Native>, Option<NullBuffer>) {
419 let result_len = predicate.len();
420 let nulls = NullBuffer::new(predicate);
421 let scalars = vec![value; result_len];
422
423 (scalars, Some(nulls))
424 }
425}
426
427impl<T: ArrowPrimitiveType> ZipImpl for PrimitiveScalarImpl<T> {
428 fn create_output(&self, predicate: &BooleanArray) -> Result<ArrayRef, ArrowError> {
429 let result_len = predicate.len();
430 let predicate = maybe_prep_null_mask_filter(predicate);
432
433 let (scalars, nulls): (Vec<T::Native>, Option<NullBuffer>) = match (self.truthy, self.falsy)
434 {
435 (Some(truthy_val), Some(falsy_val)) => {
436 let scalars: Vec<T::Native> = predicate
437 .iter()
438 .map(|b| if b { truthy_val } else { falsy_val })
439 .collect();
440
441 (scalars, None)
442 }
443 (Some(truthy_val), None) => {
444 Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, truthy_val)
448 }
449 (None, Some(falsy_val)) => {
450 let predicate = predicate.not();
455
456 Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, falsy_val)
457 }
458 (None, None) => {
459 let nulls = NullBuffer::new_null(result_len);
461 let scalars = vec![T::default_value(); result_len];
462
463 (scalars, Some(nulls))
464 }
465 };
466
467 let scalars = ScalarBuffer::<T::Native>::from(scalars);
468 let output = PrimitiveArray::<T>::try_new(scalars, nulls)?;
469
470 let output = output.with_data_type(self.data_type.clone());
472
473 Ok(Arc::new(output))
474 }
475}
476
477#[derive(PartialEq, Hash)]
478struct BytesScalarImpl<T: ByteArrayType> {
479 truthy: Option<Vec<u8>>,
480 falsy: Option<Vec<u8>>,
481 phantom: PhantomData<T>,
482}
483
484impl<T: ByteArrayType> Debug for BytesScalarImpl<T> {
485 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
486 f.debug_struct("BytesScalarImpl")
487 .field("truthy", &self.truthy)
488 .field("falsy", &self.falsy)
489 .finish()
490 }
491}
492
493impl<T: ByteArrayType> BytesScalarImpl<T> {
494 fn new(truthy_value: &dyn Array, falsy_value: &dyn Array) -> Self {
495 Self {
496 truthy: Self::get_value_from_scalar(truthy_value),
497 falsy: Self::get_value_from_scalar(falsy_value),
498 phantom: PhantomData,
499 }
500 }
501
502 fn get_value_from_scalar(scalar: &dyn Array) -> Option<Vec<u8>> {
503 if scalar.is_null(0) {
504 None
505 } else {
506 let bytes: &[u8] = scalar.as_bytes::<T>().value(0).as_ref();
507
508 Some(bytes.to_vec())
509 }
510 }
511
512 #[expect(clippy::type_complexity)]
516 fn get_scalar_and_null_buffer_for_single_non_nullable(
517 predicate: BooleanBuffer,
518 value: &[u8],
519 ) -> Result<(Buffer, OffsetBuffer<T::Offset>, Option<NullBuffer>), ArrowError> {
520 let value_length = value.len();
521
522 let number_of_true = predicate.count_set_bits();
523
524 if number_of_true == 0 {
526 let nulls = NullBuffer::new_null(predicate.len());
528
529 return Ok((
530 Buffer::from(&[]),
532 OffsetBuffer::<T::Offset>::new_zeroed(predicate.len()),
534 Some(nulls),
535 ));
536 }
537
538 let offsets = OffsetBuffer::<T::Offset>::from_lengths(
539 predicate.iter().map(|b| if b { value_length } else { 0 }),
540 );
541
542 let mut bytes = MutableBuffer::with_capacity(0);
543 bytes
544 .try_repeat_slice_n_times(value, number_of_true)
545 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
546
547 let bytes = Buffer::from(bytes);
548
549 let nulls = NullBuffer::new(predicate);
552
553 Ok((bytes, offsets, Some(nulls)))
554 }
555
556 fn get_bytes_and_offset_for_all_same_value(
559 number_of_values: usize,
560 value: &[u8],
561 ) -> Result<(Buffer, OffsetBuffer<T::Offset>), ArrowError> {
562 let value_length = value.len();
563
564 let offsets =
565 OffsetBuffer::<T::Offset>::from_repeated_length(value_length, number_of_values);
566
567 let mut bytes = MutableBuffer::with_capacity(0);
568 bytes
569 .try_repeat_slice_n_times(value, number_of_values)
570 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
571 let bytes = Buffer::from(bytes);
572
573 Ok((bytes, offsets))
574 }
575
576 fn create_output_on_non_nulls(
577 predicate: &BooleanBuffer,
578 truthy_val: &[u8],
579 falsy_val: &[u8],
580 ) -> Result<(Buffer, OffsetBuffer<<T as ByteArrayType>::Offset>), ArrowError> {
581 let true_count = predicate.count_set_bits();
582
583 match true_count {
584 0 => {
585 return Self::get_bytes_and_offset_for_all_same_value(predicate.len(), falsy_val);
587 }
588 n if n == predicate.len() => {
589 return Self::get_bytes_and_offset_for_all_same_value(predicate.len(), truthy_val);
591 }
592
593 _ => {
594 }
596 }
597
598 let total_number_of_bytes = true_count
599 .checked_mul(truthy_val.len())
600 .and_then(|truthy_bytes| {
601 let falsy_bytes = (predicate.len() - true_count).checked_mul(falsy_val.len())?;
602 truthy_bytes.checked_add(falsy_bytes)
603 })
604 .ok_or_else(|| ArrowError::MemoryError("zip output size overflow".to_string()))?;
605 T::Offset::from_usize(total_number_of_bytes)
606 .ok_or(ArrowError::OffsetOverflowError(total_number_of_bytes))?;
607 let mut mutable = MutableBuffer::with_capacity(total_number_of_bytes);
608 let mut offsets = Vec::<T::Offset>::with_capacity(predicate.len() + 1);
609 offsets.push(T::Offset::usize_as(0));
610 let mut current_offset: usize = 0;
611
612 let mut filled = 0;
614
615 let truthy_len = truthy_val.len();
616 let falsy_len = falsy_val.len();
617
618 SlicesIterator::from(predicate).try_for_each(|(start, end)| -> Result<(), ArrowError> {
620 if start > filled {
622 let false_repeat_count = start - filled;
623 mutable
625 .try_repeat_slice_n_times(falsy_val, false_repeat_count)
626 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
627
628 let start_offset = current_offset;
629 current_offset += falsy_len * false_repeat_count;
630 offsets.extend(
631 (1..=false_repeat_count)
632 .map(|index| T::Offset::usize_as(start_offset + index * falsy_len)),
633 );
634 }
635
636 let true_repeat_count = end - start;
637 mutable
639 .try_repeat_slice_n_times(truthy_val, true_repeat_count)
640 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
641
642 let start_offset = current_offset;
643 current_offset += truthy_len * true_repeat_count;
644 offsets.extend(
645 (1..=true_repeat_count)
646 .map(|index| T::Offset::usize_as(start_offset + index * truthy_len)),
647 );
648 filled = end;
649 Ok(())
650 })?;
651 if filled < predicate.len() {
653 let false_repeat_count = predicate.len() - filled;
654 mutable
656 .try_repeat_slice_n_times(falsy_val, false_repeat_count)
657 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
658
659 let start_offset = current_offset;
660 current_offset += falsy_len * false_repeat_count;
661 offsets.extend(
662 (1..=false_repeat_count)
663 .map(|index| T::Offset::usize_as(start_offset + index * falsy_len)),
664 );
665 }
666
667 debug_assert_eq!(current_offset, total_number_of_bytes);
668 let offsets = unsafe { OffsetBuffer::new_unchecked(offsets.into()) };
670 Ok((mutable.into(), offsets))
671 }
672}
673
674impl<T: ByteArrayType> ZipImpl for BytesScalarImpl<T> {
675 fn create_output(&self, predicate: &BooleanArray) -> Result<ArrayRef, ArrowError> {
676 let result_len = predicate.len();
677 let predicate = maybe_prep_null_mask_filter(predicate);
679
680 let (bytes, offsets, nulls): (Buffer, OffsetBuffer<T::Offset>, Option<NullBuffer>) =
681 match (self.truthy.as_deref(), self.falsy.as_deref()) {
682 (Some(truthy_val), Some(falsy_val)) => {
683 let (bytes, offsets) =
684 Self::create_output_on_non_nulls(&predicate, truthy_val, falsy_val)?;
685
686 (bytes, offsets, None)
687 }
688 (Some(truthy_val), None) => {
689 Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, truthy_val)?
690 }
691 (None, Some(falsy_val)) => {
692 let predicate = predicate.not();
697 Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, falsy_val)?
698 }
699 (None, None) => {
700 let nulls = NullBuffer::new_null(result_len);
702
703 (
704 Buffer::from(&[]),
706 OffsetBuffer::<T::Offset>::new_zeroed(predicate.len()),
708 Some(nulls),
709 )
710 }
711 };
712
713 let output = unsafe {
714 GenericByteArray::<T>::new_unchecked(offsets, bytes, nulls)
717 };
718
719 Ok(Arc::new(output))
720 }
721}
722
723fn maybe_prep_null_mask_filter(predicate: &BooleanArray) -> BooleanBuffer {
724 if predicate.null_count() == 0 {
726 predicate.values().clone()
727 } else {
728 let cleaned = prep_null_mask_filter(predicate);
729 let (boolean_buffer, _) = cleaned.into_parts();
730 boolean_buffer
731 }
732}
733
734struct ByteViewScalarImpl<T: ByteViewType> {
735 truthy_view: Option<u128>,
736 truthy_buffers: Arc<[Buffer]>,
737 falsy_view: Option<u128>,
738 falsy_buffers: Arc<[Buffer]>,
739 phantom: PhantomData<T>,
740}
741
742static EMPTY_ARC: OnceLock<Arc<[Buffer]>> = OnceLock::new();
743fn empty_arc_buffers() -> Arc<[Buffer]> {
744 Arc::clone(EMPTY_ARC.get_or_init(|| Arc::new([])))
745}
746
747impl<T: ByteViewType> ByteViewScalarImpl<T> {
748 fn new(truthy: &dyn Array, falsy: &dyn Array) -> Self {
749 let (truthy_view, truthy_buffers) = Self::get_value_from_scalar(truthy);
750 let (falsy_view, falsy_buffers) = Self::get_value_from_scalar(falsy);
751 Self {
752 truthy_view,
753 truthy_buffers,
754 falsy_view,
755 falsy_buffers,
756 phantom: PhantomData,
757 }
758 }
759
760 fn get_value_from_scalar(scalar: &dyn Array) -> (Option<u128>, Arc<[Buffer]>) {
761 if scalar.is_null(0) {
762 (None, empty_arc_buffers())
763 } else {
764 let (views, buffers, _) = scalar.as_byte_view::<T>().clone().into_parts();
765 (views.first().copied(), buffers)
766 }
767 }
768
769 fn get_views_for_single_non_nullable(
770 predicate: BooleanBuffer,
771 value: u128,
772 buffers: Arc<[Buffer]>,
773 ) -> (ScalarBuffer<u128>, Arc<[Buffer]>, Option<NullBuffer>) {
774 let number_of_true = predicate.count_set_bits();
775 let number_of_values = predicate.len();
776
777 if number_of_true == 0 {
779 return (
781 vec![0; number_of_values].into(),
782 empty_arc_buffers(),
783 Some(NullBuffer::new_null(number_of_values)),
784 );
785 }
786 let bytes = vec![value; number_of_values];
787
788 let nulls = NullBuffer::new(predicate);
791 (bytes.into(), buffers, Some(nulls))
792 }
793
794 #[expect(clippy::type_complexity)]
795 fn get_views_for_non_nullable(
796 predicate: BooleanBuffer,
797 result_len: usize,
798 truthy_view: u128,
799 truthy_buffers: Arc<[Buffer]>,
800 falsy_view: u128,
801 falsy_buffers: Arc<[Buffer]>,
802 ) -> Result<(ScalarBuffer<u128>, Arc<[Buffer]>, Option<NullBuffer>), ArrowError> {
803 let true_count = predicate.count_set_bits();
804 match true_count {
805 0 => {
806 Ok((vec![falsy_view; result_len].into(), falsy_buffers, None))
808 }
809 n if n == predicate.len() => {
810 Ok((vec![truthy_view; result_len].into(), truthy_buffers, None))
812 }
813 _ => {
814 let true_count = predicate.count_set_bits();
815 let mut buffers: Vec<Buffer> = truthy_buffers.to_vec();
816
817 let view_falsy = if (falsy_view as u32) <= MAX_INLINE_VIEW_LEN {
819 falsy_view
820 } else {
821 let byte_view_falsy = ByteView::from(falsy_view);
822 let new_index_falsy_buffers =
823 buffers.len() as u32 + byte_view_falsy.buffer_index;
824 buffers.extend(falsy_buffers.iter().cloned());
825 let byte_view_falsy =
826 byte_view_falsy.with_buffer_index(new_index_falsy_buffers);
827 byte_view_falsy.as_u128()
828 };
829
830 let total_number_of_bytes = true_count * 16 + (predicate.len() - true_count) * 16;
831 let mut mutable = MutableBuffer::new(total_number_of_bytes);
832 let mut filled = 0;
833
834 SlicesIterator::from(&predicate).try_for_each(
835 |(start, end)| -> Result<(), ArrowError> {
836 if start > filled {
837 let false_repeat_count = start - filled;
838 mutable
839 .try_repeat_slice_n_times(
840 view_falsy.to_byte_slice(),
841 false_repeat_count,
842 )
843 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
844 }
845 let true_repeat_count = end - start;
846 mutable
847 .try_repeat_slice_n_times(
848 truthy_view.to_byte_slice(),
849 true_repeat_count,
850 )
851 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
852 filled = end;
853 Ok(())
854 },
855 )?;
856
857 if filled < predicate.len() {
858 let false_repeat_count = predicate.len() - filled;
859 mutable
860 .try_repeat_slice_n_times(view_falsy.to_byte_slice(), false_repeat_count)
861 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
862 }
863
864 let bytes = Buffer::from(mutable);
865 Ok((bytes.into(), buffers.into(), None))
866 }
867 }
868 }
869}
870
871impl<T: ByteViewType> Debug for ByteViewScalarImpl<T> {
872 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
873 f.debug_struct("ByteViewScalarImpl")
874 .field("truthy", &self.truthy_view)
875 .field("falsy", &self.falsy_view)
876 .finish()
877 }
878}
879
880impl<T: ByteViewType> ZipImpl for ByteViewScalarImpl<T> {
881 fn create_output(&self, predicate: &BooleanArray) -> Result<ArrayRef, ArrowError> {
882 let result_len = predicate.len();
883 let predicate = maybe_prep_null_mask_filter(predicate);
885
886 let (views, buffers, nulls) = match (self.truthy_view, self.falsy_view) {
887 (Some(truthy), Some(falsy)) => Self::get_views_for_non_nullable(
888 predicate,
889 result_len,
890 truthy,
891 Arc::clone(&self.truthy_buffers),
892 falsy,
893 Arc::clone(&self.falsy_buffers),
894 )?,
895 (Some(truthy), None) => Self::get_views_for_single_non_nullable(
896 predicate,
897 truthy,
898 Arc::clone(&self.truthy_buffers),
899 ),
900 (None, Some(falsy)) => {
901 let predicate = predicate.not();
902 Self::get_views_for_single_non_nullable(
903 predicate,
904 falsy,
905 Arc::clone(&self.falsy_buffers),
906 )
907 }
908 (None, None) => {
909 (
911 vec![0; result_len].into(),
912 empty_arc_buffers(),
913 Some(NullBuffer::new_null(result_len)),
914 )
915 }
916 };
917
918 let result = unsafe { GenericByteViewArray::<T>::new_unchecked(views, buffers, nulls) };
919 Ok(Arc::new(result))
920 }
921}
922
923#[cfg(test)]
924mod test {
925 use super::*;
926 use arrow_array::types::Int32Type;
927
928 #[test]
929 fn test_count_true_runs() {
930 let assert_runs = |values: &[bool], expected| {
931 let mask: BooleanBuffer = values.iter().copied().collect();
932 assert_eq!(count_true_runs(&mask), expected, "mask: {values:?}");
933 };
934
935 assert_runs(&[], 0);
936 assert_runs(&[false, false, false], 0);
937 assert_runs(&[true, true, true], 1);
938 assert_runs(&[true, false, true, true, false, true], 3);
939
940 let mut values = vec![false; 130];
942 values[0] = true;
943 values[63..66].fill(true);
944 values[128..].fill(true);
945 assert_runs(&values, 3);
946
947 let mut offset_values = vec![false; 135];
949 offset_values[3..133].copy_from_slice(&values);
950 let offset_mask: BooleanBuffer = offset_values.into_iter().collect();
951 assert_eq!(count_true_runs(&offset_mask.slice(3, 130)), 3);
952 }
953
954 #[test]
955 fn test_should_use_interleave() {
956 let short: BooleanBuffer = (0..64).map(|i| i % 2 == 0).collect();
957 assert!(!should_use_interleave(&short));
958
959 let fragmented: BooleanBuffer = (0..8192).map(|i| i % 2 == 0).collect();
960 assert!(should_use_interleave(&fragmented));
961
962 let long_runs: BooleanBuffer = (0..8192).map(|i| i < 4096).collect();
963 assert!(!should_use_interleave(&long_runs));
964
965 let sparse: BooleanBuffer = (0..8192).map(|i| i % 10 == 0).collect();
966 assert!(!should_use_interleave(&sparse));
967
968 let dense: BooleanBuffer = (0..8192).map(|i| i % 10 != 0).collect();
969 assert!(!should_use_interleave(&dense));
970
971 let fragmented_head: BooleanBuffer = (0..8192).map(|i| i < 256 && i % 2 == 0).collect();
972 assert!(!should_use_interleave(&fragmented_head));
973
974 let fragmented_edges: BooleanBuffer = (0..8192)
975 .map(|i| !(256..7936).contains(&i) && i % 2 == 0)
976 .collect();
977 assert!(!should_use_interleave(&fragmented_edges));
978
979 let offset: BooleanBuffer = (0..8195).map(|i| i >= 3 && i % 2 == 1).collect();
981 assert!(should_use_interleave(&offset.slice(3, 8192)));
982 }
983
984 #[test]
985 fn test_interleave_arrays() {
986 let mask = BooleanArray::from(vec![Some(true), None, Some(true), Some(false)]);
987 let mask = maybe_prep_null_mask_filter(&mask);
988 let truthy = Int32Array::from(vec![Some(1), None, Some(3), Some(4)]).to_data();
989 let falsy = Int32Array::from(vec![Some(10), Some(20), None, Some(40)]).to_data();
990 let expected = Int32Array::from(vec![Some(1), Some(20), Some(3), Some(40)]);
991
992 let actual = interleave_arrays(&mask, &truthy, &falsy).unwrap();
993 assert_eq!(actual.as_primitive::<Int32Type>(), &expected);
994 }
995
996 #[test]
997 fn test_zip_fragmented_array_mask() {
998 let mask: BooleanArray = (0..8192)
999 .map(|i| match i % 3 {
1000 0 => Some(true),
1001 1 => Some(false),
1002 _ => None,
1003 })
1004 .collect();
1005 let truthy: Int32Array = (0..8192).map(|i| (i % 7 != 0).then_some(i)).collect();
1006 let falsy: Int32Array = (0..8192).map(|i| (i % 11 != 0).then_some(-i)).collect();
1007 let expected: Int32Array = (0..8192)
1008 .map(|i| {
1009 let array = if i % 3 == 0 { &truthy } else { &falsy };
1010 array.is_valid(i).then(|| array.value(i))
1011 })
1012 .collect();
1013
1014 let actual = zip(&mask, &truthy, &falsy).unwrap();
1015 assert_eq!(actual.as_primitive::<Int32Type>(), &expected);
1016 }
1017
1018 #[test]
1019 fn test_zip_kernel_one() {
1020 let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1021 let b = Int32Array::from(vec![None, Some(3), Some(6), Some(7), Some(3)]);
1022 let mask = BooleanArray::from(vec![true, true, false, false, true]);
1023 let out = zip(&mask, &a, &b).unwrap();
1024 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1025 let expected = Int32Array::from(vec![Some(5), None, Some(6), Some(7), Some(1)]);
1026 assert_eq!(actual, &expected);
1027 }
1028
1029 #[test]
1030 fn test_zip_kernel_two() {
1031 let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1032 let b = Int32Array::from(vec![None, Some(3), Some(6), Some(7), Some(3)]);
1033 let mask = BooleanArray::from(vec![false, false, true, true, false]);
1034 let out = zip(&mask, &a, &b).unwrap();
1035 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1036 let expected = Int32Array::from(vec![None, Some(3), Some(7), None, Some(3)]);
1037 assert_eq!(actual, &expected);
1038 }
1039
1040 #[test]
1041 fn test_zip_kernel_scalar_falsy_1() {
1042 let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1043
1044 let fallback = Scalar::new(Int32Array::from_value(42, 1));
1045
1046 let mask = BooleanArray::from(vec![true, true, false, false, true]);
1047 let out = zip(&mask, &a, &fallback).unwrap();
1048 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1049 let expected = Int32Array::from(vec![Some(5), None, Some(42), Some(42), Some(1)]);
1050 assert_eq!(actual, &expected);
1051 }
1052
1053 #[test]
1054 fn test_zip_kernel_scalar_falsy_2() {
1055 let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1056
1057 let fallback = Scalar::new(Int32Array::from_value(42, 1));
1058
1059 let mask = BooleanArray::from(vec![false, false, true, true, false]);
1060 let out = zip(&mask, &a, &fallback).unwrap();
1061 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1062 let expected = Int32Array::from(vec![Some(42), Some(42), Some(7), None, Some(42)]);
1063 assert_eq!(actual, &expected);
1064 }
1065
1066 #[test]
1067 fn test_zip_kernel_scalar_truthy_1() {
1068 let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1069
1070 let fallback = Scalar::new(Int32Array::from_value(42, 1));
1071
1072 let mask = BooleanArray::from(vec![true, true, false, false, true]);
1073 let out = zip(&mask, &fallback, &a).unwrap();
1074 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1075 let expected = Int32Array::from(vec![Some(42), Some(42), Some(7), None, Some(42)]);
1076 assert_eq!(actual, &expected);
1077 }
1078
1079 #[test]
1080 fn test_zip_kernel_scalar_truthy_2() {
1081 let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1082
1083 let fallback = Scalar::new(Int32Array::from_value(42, 1));
1084
1085 let mask = BooleanArray::from(vec![false, false, true, true, false]);
1086 let out = zip(&mask, &fallback, &a).unwrap();
1087 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1088 let expected = Int32Array::from(vec![Some(5), None, Some(42), Some(42), Some(1)]);
1089 assert_eq!(actual, &expected);
1090 }
1091
1092 #[test]
1093 fn test_zip_kernel_scalar_both_mask_ends_with_true() {
1094 let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1095 let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
1096
1097 let mask = BooleanArray::from(vec![true, true, false, false, true]);
1098 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1099 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1100 let expected = Int32Array::from(vec![Some(42), Some(42), Some(123), Some(123), Some(42)]);
1101 assert_eq!(actual, &expected);
1102 }
1103
1104 #[test]
1105 fn test_zip_kernel_scalar_both_mask_ends_with_false() {
1106 let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1107 let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
1108
1109 let mask = BooleanArray::from(vec![true, true, false, true, false, false]);
1110 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1111 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1112 let expected = Int32Array::from(vec![
1113 Some(42),
1114 Some(42),
1115 Some(123),
1116 Some(42),
1117 Some(123),
1118 Some(123),
1119 ]);
1120 assert_eq!(actual, &expected);
1121 }
1122
1123 #[test]
1124 fn test_zip_kernel_primitive_scalar_none_1() {
1125 let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1126 let scalar_falsy = Scalar::new(Int32Array::new_null(1));
1127
1128 let mask = BooleanArray::from(vec![true, true, false, false, true]);
1129 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1130 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1131 let expected = Int32Array::from(vec![Some(42), Some(42), None, None, Some(42)]);
1132 assert_eq!(actual, &expected);
1133 }
1134
1135 #[test]
1136 fn test_zip_kernel_primitive_scalar_none_2() {
1137 let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1138 let scalar_falsy = Scalar::new(Int32Array::new_null(1));
1139
1140 let mask = BooleanArray::from(vec![false, false, true, true, false]);
1141 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1142 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1143 let expected = Int32Array::from(vec![None, None, Some(42), Some(42), None]);
1144 assert_eq!(actual, &expected);
1145 }
1146
1147 #[test]
1148 fn test_zip_kernel_primitive_scalar_both_null() {
1149 let scalar_truthy = Scalar::new(Int32Array::new_null(1));
1150 let scalar_falsy = Scalar::new(Int32Array::new_null(1));
1151
1152 let mask = BooleanArray::from(vec![false, false, true, true, false]);
1153 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1154 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1155 let expected = Int32Array::from(vec![None, None, None, None, None]);
1156 assert_eq!(actual, &expected);
1157 }
1158
1159 #[test]
1160 fn test_zip_primitive_array_with_nulls_is_mask_should_be_treated_as_false() {
1161 let truthy = Int32Array::from_iter_values(vec![1, 2, 3, 4, 5, 6]);
1162 let falsy = Int32Array::from_iter_values(vec![7, 8, 9, 10, 11, 12]);
1163
1164 let mask = {
1165 let booleans = BooleanBuffer::from(vec![true, true, false, true, false, false]);
1166 let nulls = NullBuffer::from(vec![
1167 true, true, true,
1168 false, true, true,
1170 ]);
1171 BooleanArray::new(booleans, Some(nulls))
1172 };
1173 let out = zip(&mask, &truthy, &falsy).unwrap();
1174 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1175 let expected = Int32Array::from(vec![
1176 Some(1),
1177 Some(2),
1178 Some(9),
1179 Some(10), Some(11),
1181 Some(12),
1182 ]);
1183 assert_eq!(actual, &expected);
1184 }
1185
1186 #[test]
1187 fn test_zip_kernel_primitive_scalar_with_boolean_array_mask_with_nulls_should_be_treated_as_false()
1188 {
1189 let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1190 let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
1191
1192 let mask = {
1193 let booleans = BooleanBuffer::from(vec![true, true, false, true, false, false]);
1194 let nulls = NullBuffer::from(vec![
1195 true, true, true,
1196 false, true, true,
1198 ]);
1199 BooleanArray::new(booleans, Some(nulls))
1200 };
1201 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1202 let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1203 let expected = Int32Array::from(vec![
1204 Some(42),
1205 Some(42),
1206 Some(123),
1207 Some(123), Some(123),
1209 Some(123),
1210 ]);
1211 assert_eq!(actual, &expected);
1212 }
1213
1214 #[test]
1215 fn test_zip_string_array_with_nulls_is_mask_should_be_treated_as_false() {
1216 let truthy = StringArray::from_iter_values(vec!["1", "2", "3", "4", "5", "6"]);
1217 let falsy = StringArray::from_iter_values(vec!["7", "8", "9", "10", "11", "12"]);
1218
1219 let mask = {
1220 let booleans = BooleanBuffer::from(vec![true, true, false, true, false, false]);
1221 let nulls = NullBuffer::from(vec![
1222 true, true, true,
1223 false, true, true,
1225 ]);
1226 BooleanArray::new(booleans, Some(nulls))
1227 };
1228 let out = zip(&mask, &truthy, &falsy).unwrap();
1229 let actual = out.as_string::<i32>();
1230 let expected = StringArray::from_iter_values(vec![
1231 "1", "2", "9", "10", "11", "12",
1233 ]);
1234 assert_eq!(actual, &expected);
1235 }
1236
1237 #[test]
1238 fn test_zip_kernel_large_string_scalar_with_boolean_array_mask_with_nulls_should_be_treated_as_false()
1239 {
1240 let scalar_truthy = Scalar::new(LargeStringArray::from_iter_values(["test"]));
1241 let scalar_falsy = Scalar::new(LargeStringArray::from_iter_values(["something else"]));
1242
1243 let mask = {
1244 let booleans = BooleanBuffer::from(vec![true, true, false, true, false, false]);
1245 let nulls = NullBuffer::from(vec![
1246 true, true, true,
1247 false, true, true,
1249 ]);
1250 BooleanArray::new(booleans, Some(nulls))
1251 };
1252 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1253 let actual = out.as_any().downcast_ref::<LargeStringArray>().unwrap();
1254 let expected = LargeStringArray::from_iter(vec![
1255 Some("test"),
1256 Some("test"),
1257 Some("something else"),
1258 Some("something else"), Some("something else"),
1260 Some("something else"),
1261 ]);
1262 assert_eq!(actual, &expected);
1263 }
1264
1265 #[test]
1266 fn test_zip_kernel_bytes_scalar_none_1() {
1267 let scalar_truthy = Scalar::new(StringArray::from_iter_values(["hello"]));
1268 let scalar_falsy = Scalar::new(StringArray::new_null(1));
1269
1270 let mask = BooleanArray::from(vec![true, true, false, false, true]);
1271 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1272 let actual = out.as_any().downcast_ref::<StringArray>().unwrap();
1273 let expected = StringArray::from_iter(vec![
1274 Some("hello"),
1275 Some("hello"),
1276 None,
1277 None,
1278 Some("hello"),
1279 ]);
1280 assert_eq!(actual, &expected);
1281 }
1282
1283 #[test]
1284 fn test_zip_kernel_bytes_scalar_none_2() {
1285 let scalar_truthy = Scalar::new(StringArray::new_null(1));
1286 let scalar_falsy = Scalar::new(StringArray::from_iter_values(["hello"]));
1287
1288 let mask = BooleanArray::from(vec![true, true, false, false, true]);
1289 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1290 let actual = out.as_any().downcast_ref::<StringArray>().unwrap();
1291 let expected = StringArray::from_iter(vec![None, None, Some("hello"), Some("hello"), None]);
1292 assert_eq!(actual, &expected);
1293 }
1294
1295 #[test]
1296 fn test_zip_kernel_bytes_scalar_both() {
1297 let scalar_truthy = Scalar::new(StringArray::from_iter_values(["test"]));
1298 let scalar_falsy = Scalar::new(StringArray::from_iter_values(["something else"]));
1299
1300 let mask = BooleanArray::from(vec![true, true, false, true, false, false]);
1302 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1303 let actual = out.as_any().downcast_ref::<StringArray>().unwrap();
1304 let expected = StringArray::from_iter(vec![
1305 Some("test"),
1306 Some("test"),
1307 Some("something else"),
1308 Some("test"),
1309 Some("something else"),
1310 Some("something else"),
1311 ]);
1312 assert_eq!(actual, &expected);
1313 }
1314
1315 #[test]
1316 fn test_zip_scalar_bytes_offset_overflow() {
1317 let value = vec![0_u8; 65_536];
1320 let large = Scalar::new(BinaryArray::from_iter_values([value.as_slice()]));
1321 let empty = Scalar::new(BinaryArray::from_iter_values([b"".as_slice()]));
1322 let mask = BooleanArray::from_iter((0..65_536).map(|i| Some(i % 2 == 0)));
1323
1324 for (truthy, falsy) in [(&large, &empty), (&empty, &large)] {
1325 assert!(matches!(
1326 zip(&mask, truthy, falsy),
1327 Err(ArrowError::OffsetOverflowError(2_147_483_648))
1328 ));
1329 }
1330 }
1331
1332 #[test]
1333 fn test_zip_scalar_bytes_only_taking_one_side() {
1334 let mask_len = 5;
1335 let all_true_mask = BooleanArray::from(vec![true; mask_len]);
1336 let all_false_mask = BooleanArray::from(vec![false; mask_len]);
1337
1338 let null_scalar = Scalar::new(StringArray::new_null(1));
1339 let non_null_scalar_1 = Scalar::new(StringArray::from_iter_values(["test"]));
1340 let non_null_scalar_2 = Scalar::new(StringArray::from_iter_values(["something else"]));
1341
1342 {
1343 let out = zip(&all_true_mask, &null_scalar, &non_null_scalar_1).unwrap();
1346 let actual = out.as_string::<i32>();
1347 let expected = StringArray::from_iter(std::iter::repeat_n(None::<&str>, mask_len));
1348 assert_eq!(actual, &expected);
1349 }
1350
1351 {
1352 let out = zip(&all_false_mask, &null_scalar, &non_null_scalar_1).unwrap();
1355 let actual = out.as_string::<i32>();
1356 let expected = StringArray::from_iter(std::iter::repeat_n(Some("test"), mask_len));
1357 assert_eq!(actual, &expected);
1358 }
1359
1360 {
1361 let out = zip(&all_true_mask, &non_null_scalar_1, &null_scalar).unwrap();
1364 let actual = out.as_string::<i32>();
1365 let expected = StringArray::from_iter(std::iter::repeat_n(Some("test"), mask_len));
1366 assert_eq!(actual, &expected);
1367 }
1368
1369 {
1370 let out = zip(&all_false_mask, &non_null_scalar_1, &null_scalar).unwrap();
1373 let actual = out.as_string::<i32>();
1374 let expected = StringArray::from_iter(std::iter::repeat_n(None::<&str>, mask_len));
1375 assert_eq!(actual, &expected);
1376 }
1377
1378 {
1379 let out = zip(&all_true_mask, &non_null_scalar_1, &non_null_scalar_2).unwrap();
1382 let actual = out.as_string::<i32>();
1383 let expected = StringArray::from_iter(std::iter::repeat_n(Some("test"), mask_len));
1384 assert_eq!(actual, &expected);
1385 }
1386
1387 {
1388 let out = zip(&all_false_mask, &non_null_scalar_1, &non_null_scalar_2).unwrap();
1391 let actual = out.as_string::<i32>();
1392 let expected =
1393 StringArray::from_iter(std::iter::repeat_n(Some("something else"), mask_len));
1394 assert_eq!(actual, &expected);
1395 }
1396
1397 {
1398 let mask = BooleanArray::from(vec![true, false, true, false, true]);
1401 let out = zip(&mask, &null_scalar, &null_scalar).unwrap();
1402 let actual = out.as_string::<i32>();
1403 let expected = StringArray::from_iter(std::iter::repeat_n(None::<&str>, mask_len));
1404 assert_eq!(actual, &expected);
1405 }
1406 }
1407
1408 #[test]
1409 fn test_scalar_zipper() {
1410 let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1411 let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
1412
1413 let mask = BooleanArray::from(vec![false, false, true, true, false]);
1414
1415 let scalar_zipper = ScalarZipper::try_new(&scalar_truthy, &scalar_falsy).unwrap();
1416 let out = scalar_zipper.zip(&mask).unwrap();
1417 let actual = out.as_primitive::<Int32Type>();
1418 let expected = Int32Array::from(vec![Some(123), Some(123), Some(42), Some(42), Some(123)]);
1419 assert_eq!(actual, &expected);
1420
1421 let mask = BooleanArray::from(vec![true, false, true]);
1423 let out = scalar_zipper.zip(&mask).unwrap();
1424 let actual = out.as_primitive::<Int32Type>();
1425 let expected = Int32Array::from(vec![Some(42), Some(123), Some(42)]);
1426 assert_eq!(actual, &expected);
1427 }
1428
1429 #[test]
1430 fn test_zip_kernel_scalar_strings() {
1431 let scalar_truthy = Scalar::new(StringArray::from(vec!["hello"]));
1432 let scalar_falsy = Scalar::new(StringArray::from(vec!["world"]));
1433
1434 let mask = BooleanArray::from(vec![true, false, true, false, true]);
1435 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1436 let actual = out.as_string::<i32>();
1437 let expected = StringArray::from(vec![
1438 Some("hello"),
1439 Some("world"),
1440 Some("hello"),
1441 Some("world"),
1442 Some("hello"),
1443 ]);
1444 assert_eq!(actual, &expected);
1445 }
1446
1447 #[test]
1448 fn test_zip_kernel_scalar_binary() {
1449 let truthy_bytes: &[u8] = b"\xFF\xFE\xFD";
1450 let falsy_bytes: &[u8] = b"world";
1451 let scalar_truthy = Scalar::new(BinaryArray::from_iter_values(
1452 vec![truthy_bytes],
1454 ));
1455 let scalar_falsy = Scalar::new(BinaryArray::from_iter_values(vec![falsy_bytes]));
1456
1457 let mask = BooleanArray::from(vec![true, false, true, false, true]);
1458 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1459 let actual = out.as_binary::<i32>();
1460 let expected = BinaryArray::from(vec![
1461 Some(truthy_bytes),
1462 Some(falsy_bytes),
1463 Some(truthy_bytes),
1464 Some(falsy_bytes),
1465 Some(truthy_bytes),
1466 ]);
1467 assert_eq!(actual, &expected);
1468 }
1469
1470 #[test]
1471 fn test_zip_kernel_scalar_large_binary() {
1472 let truthy_bytes: &[u8] = b"hey";
1473 let falsy_bytes: &[u8] = b"world";
1474 let scalar_truthy = Scalar::new(LargeBinaryArray::from_iter_values(vec![truthy_bytes]));
1475 let scalar_falsy = Scalar::new(LargeBinaryArray::from_iter_values(vec![falsy_bytes]));
1476
1477 let mask = BooleanArray::from(vec![true, false, true, false, true]);
1478 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1479 let actual = out.as_binary::<i64>();
1480 let expected = LargeBinaryArray::from(vec![
1481 Some(truthy_bytes),
1482 Some(falsy_bytes),
1483 Some(truthy_bytes),
1484 Some(falsy_bytes),
1485 Some(truthy_bytes),
1486 ]);
1487 assert_eq!(actual, &expected);
1488 }
1489
1490 #[test]
1492 fn test_zip_decimal_with_custom_precision_and_scale() {
1493 let arr = Decimal128Array::from_iter_values([12345, 456, 7890, -123223423432432])
1494 .with_precision_and_scale(20, 2)
1495 .unwrap();
1496
1497 let arr: ArrayRef = Arc::new(arr);
1498
1499 let scalar_1 = Scalar::new(arr.slice(0, 1));
1500 let scalar_2 = Scalar::new(arr.slice(1, 1));
1501 let null_scalar = Scalar::new(new_null_array(arr.data_type(), 1));
1502 let array_1: ArrayRef = arr.slice(0, 2);
1503 let array_2: ArrayRef = arr.slice(2, 2);
1504
1505 test_zip_output_data_types_for_input(scalar_1, scalar_2, null_scalar, array_1, array_2);
1506 }
1507
1508 #[test]
1510 fn test_zip_timestamp_with_timezone() {
1511 let arr = TimestampSecondArray::from(vec![0, 1000, 2000, 4000])
1512 .with_timezone("+01:00".to_string());
1513
1514 let arr: ArrayRef = Arc::new(arr);
1515
1516 let scalar_1 = Scalar::new(arr.slice(0, 1));
1517 let scalar_2 = Scalar::new(arr.slice(1, 1));
1518 let null_scalar = Scalar::new(new_null_array(arr.data_type(), 1));
1519 let array_1: ArrayRef = arr.slice(0, 2);
1520 let array_2: ArrayRef = arr.slice(2, 2);
1521
1522 test_zip_output_data_types_for_input(scalar_1, scalar_2, null_scalar, array_1, array_2);
1523 }
1524
1525 fn test_zip_output_data_types_for_input(
1526 scalar_1: Scalar<ArrayRef>,
1527 scalar_2: Scalar<ArrayRef>,
1528 null_scalar: Scalar<ArrayRef>,
1529 array_1: ArrayRef,
1530 array_2: ArrayRef,
1531 ) {
1532 test_zip_output_data_type(&scalar_1, &scalar_2, 10);
1534
1535 test_zip_output_data_type(&null_scalar, &scalar_1, 10);
1537 test_zip_output_data_type(&scalar_1, &null_scalar, 10);
1538
1539 test_zip_output_data_type(&array_1.as_ref(), &scalar_1, array_1.len());
1541 test_zip_output_data_type(&scalar_1, &array_1.as_ref(), array_1.len());
1542
1543 test_zip_output_data_type(&array_1.as_ref(), &null_scalar, array_1.len());
1545
1546 test_zip_output_data_type(&null_scalar, &array_1.as_ref(), array_1.len());
1547
1548 test_zip_output_data_type(&array_1.as_ref(), &array_2.as_ref(), array_1.len());
1550 }
1551
1552 fn test_zip_output_data_type(truthy: &dyn Datum, falsy: &dyn Datum, mask_length: usize) {
1553 let expected_data_type = truthy.get().0.data_type().clone();
1554 assert_eq!(&expected_data_type, falsy.get().0.data_type());
1555
1556 let mask_all_true = BooleanArray::from(vec![true; mask_length]);
1558 let mask_all_false = BooleanArray::from(vec![false; mask_length]);
1559 let mask_some_true_and_false =
1560 BooleanArray::from((0..mask_length).map(|i| i % 2 == 0).collect::<Vec<bool>>());
1561
1562 for mask in [&mask_all_true, &mask_all_false, &mask_some_true_and_false] {
1563 let out = zip(mask, truthy, falsy).unwrap();
1564 assert_eq!(out.data_type(), &expected_data_type);
1565 }
1566 }
1567
1568 #[test]
1569 fn zip_scalar_fallback_impl() {
1570 let truthy_list_item_scalar = Some(vec![Some(1), None, Some(3)]);
1571 let truthy_list_array_scalar =
1572 Scalar::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1573 truthy_list_item_scalar.clone(),
1574 ]));
1575 let falsy_list_item_scalar = Some(vec![None, Some(2), Some(4)]);
1576 let falsy_list_array_scalar =
1577 Scalar::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1578 falsy_list_item_scalar.clone(),
1579 ]));
1580 let mask = BooleanArray::from(vec![true, false, true, false, false, true, false]);
1581 let out = zip(&mask, &truthy_list_array_scalar, &falsy_list_array_scalar).unwrap();
1582 let actual = out.as_list::<i32>();
1583
1584 let expected = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1585 truthy_list_item_scalar.clone(),
1586 falsy_list_item_scalar.clone(),
1587 truthy_list_item_scalar.clone(),
1588 falsy_list_item_scalar.clone(),
1589 falsy_list_item_scalar.clone(),
1590 truthy_list_item_scalar.clone(),
1591 falsy_list_item_scalar.clone(),
1592 ]);
1593 assert_eq!(actual, &expected);
1594 }
1595
1596 #[test]
1597 fn test_zip_kernel_scalar_strings_array_view() {
1598 let scalar_truthy = Scalar::new(StringViewArray::from(vec!["hello"]));
1599 let scalar_falsy = Scalar::new(StringViewArray::from(vec!["world"]));
1600
1601 let mask = BooleanArray::from(vec![true, false, true, false]);
1602 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1603 let actual = out.as_string_view();
1604 let expected = StringViewArray::from(vec![
1605 Some("hello"),
1606 Some("world"),
1607 Some("hello"),
1608 Some("world"),
1609 ]);
1610 assert_eq!(actual, &expected);
1611 }
1612
1613 #[test]
1614 fn test_zip_kernel_scalar_binary_array_view() {
1615 let scalar_truthy = Scalar::new(BinaryViewArray::from_iter_values(vec![b"hello"]));
1616 let scalar_falsy = Scalar::new(BinaryViewArray::from_iter_values(vec![b"world"]));
1617
1618 let mask = BooleanArray::from(vec![true, false]);
1619 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1620 let actual = out.as_byte_view();
1621 let expected = BinaryViewArray::from_iter_values(vec![b"hello", b"world"]);
1622 assert_eq!(actual, &expected);
1623 }
1624
1625 #[test]
1626 fn test_zip_kernel_scalar_strings_array_view_with_nulls() {
1627 let scalar_truthy = Scalar::new(StringViewArray::from_iter_values(["hello"]));
1628 let scalar_falsy = Scalar::new(StringViewArray::new_null(1));
1629
1630 let mask = BooleanArray::from(vec![true, true, false, false, true]);
1631 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1632 let actual = out.as_any().downcast_ref::<StringViewArray>().unwrap();
1633 let expected = StringViewArray::from_iter(vec![
1634 Some("hello"),
1635 Some("hello"),
1636 None,
1637 None,
1638 Some("hello"),
1639 ]);
1640 assert_eq!(actual, &expected);
1641 }
1642
1643 #[test]
1644 fn test_zip_kernel_scalar_strings_array_view_all_true_null() {
1645 let scalar_truthy = Scalar::new(StringViewArray::new_null(1));
1646 let scalar_falsy = Scalar::new(StringViewArray::new_null(1));
1647 let mask = BooleanArray::from(vec![true, true]);
1648 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1649 let actual = out.as_any().downcast_ref::<StringViewArray>().unwrap();
1650 let expected = StringViewArray::from_iter(vec![None::<String>, None]);
1651 assert_eq!(actual, &expected);
1652 }
1653
1654 #[test]
1655 fn test_zip_kernel_scalar_strings_array_view_all_false_null() {
1656 let scalar_truthy = Scalar::new(StringViewArray::new_null(1));
1657 let scalar_falsy = Scalar::new(StringViewArray::new_null(1));
1658 let mask = BooleanArray::from(vec![false, false]);
1659 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1660 let actual = out.as_any().downcast_ref::<StringViewArray>().unwrap();
1661 let expected = StringViewArray::from_iter(vec![None::<String>, None]);
1662 assert_eq!(actual, &expected);
1663 }
1664
1665 #[test]
1666 fn test_zip_kernel_scalar_string_array_view_all_true() {
1667 let scalar_truthy = Scalar::new(StringViewArray::from(vec!["hello"]));
1668 let scalar_falsy = Scalar::new(StringViewArray::from(vec!["world"]));
1669
1670 let mask = BooleanArray::from(vec![true, true]);
1671 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1672 let actual = out.as_string_view();
1673 let expected = StringViewArray::from(vec![Some("hello"), Some("hello")]);
1674 assert_eq!(actual, &expected);
1675 }
1676
1677 #[test]
1678 fn test_zip_kernel_scalar_string_array_view_all_false() {
1679 let scalar_truthy = Scalar::new(StringViewArray::from(vec!["hello"]));
1680 let scalar_falsy = Scalar::new(StringViewArray::from(vec!["world"]));
1681
1682 let mask = BooleanArray::from(vec![false, false]);
1683 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1684 let actual = out.as_string_view();
1685 let expected = StringViewArray::from(vec![Some("world"), Some("world")]);
1686 assert_eq!(actual, &expected);
1687 }
1688
1689 #[test]
1690 fn test_zip_kernel_scalar_strings_large_strings() {
1691 let scalar_truthy = Scalar::new(StringViewArray::from(vec!["longer than 12 bytes"]));
1692 let scalar_falsy = Scalar::new(StringViewArray::from(vec!["another longer than 12 bytes"]));
1693
1694 let mask = BooleanArray::from(vec![true, false]);
1695 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1696 let actual = out.as_string_view();
1697 let expected = StringViewArray::from(vec![
1698 Some("longer than 12 bytes"),
1699 Some("another longer than 12 bytes"),
1700 ]);
1701 assert_eq!(actual, &expected);
1702 }
1703
1704 #[test]
1705 fn test_zip_kernel_scalar_strings_array_view_large_short_strings() {
1706 let scalar_truthy = Scalar::new(StringViewArray::from(vec!["hello"]));
1707 let scalar_falsy = Scalar::new(StringViewArray::from(vec!["longer than 12 bytes"]));
1708
1709 let mask = BooleanArray::from(vec![true, false, true, false]);
1710 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1711 let actual = out.as_string_view();
1712 let expected = StringViewArray::from(vec![
1713 Some("hello"),
1714 Some("longer than 12 bytes"),
1715 Some("hello"),
1716 Some("longer than 12 bytes"),
1717 ]);
1718 assert_eq!(actual, &expected);
1719 }
1720 #[test]
1721 fn test_zip_kernel_scalar_strings_array_view_large_all_true() {
1722 let scalar_truthy = Scalar::new(StringViewArray::from(vec!["longer than 12 bytes"]));
1723 let scalar_falsy = Scalar::new(StringViewArray::from(vec!["another longer than 12 bytes"]));
1724
1725 let mask = BooleanArray::from(vec![true, true]);
1726 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1727 let actual = out.as_string_view();
1728 let expected = StringViewArray::from(vec![
1729 Some("longer than 12 bytes"),
1730 Some("longer than 12 bytes"),
1731 ]);
1732 assert_eq!(actual, &expected);
1733 }
1734
1735 #[test]
1736 fn test_zip_kernel_scalar_strings_array_view_large_all_false() {
1737 let scalar_truthy = Scalar::new(StringViewArray::from(vec!["longer than 12 bytes"]));
1738 let scalar_falsy = Scalar::new(StringViewArray::from(vec!["another longer than 12 bytes"]));
1739
1740 let mask = BooleanArray::from(vec![false, false]);
1741 let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1742 let actual = out.as_string_view();
1743 let expected = StringViewArray::from(vec![
1744 Some("another longer than 12 bytes"),
1745 Some("another longer than 12 bytes"),
1746 ]);
1747 assert_eq!(actual, &expected);
1748 }
1749
1750 #[test]
1751 fn test_zip_kernel_scalar_strings_array_view_sliced_inline() {
1752 let truthy =
1754 StringViewArray::from(vec!["x", "a string longer than twelve bytes"]).slice(0, 1);
1755 let falsy =
1756 StringViewArray::from(vec!["hello", "another string longer than twelve"]).slice(0, 1);
1757
1758 let mask = BooleanArray::from(vec![true, false, true, false]);
1759 let out = zip(&mask, &Scalar::new(&truthy), &Scalar::new(&falsy)).unwrap();
1760 out.to_data().validate_full().unwrap();
1761 let actual = out.as_string_view();
1762 let expected = StringViewArray::from(vec!["x", "hello", "x", "hello"]);
1763 assert_eq!(actual, &expected);
1764 }
1765
1766 #[test]
1767 fn test_zip_kernel_scalar_strings_array_view_sliced_long_falsy() {
1768 let truthy =
1770 StringViewArray::from(vec!["a string longer than twelve bytes", "x"]).slice(0, 1);
1771 let falsy = StringViewArray::from(vec!["y", "thirteen byte"]).slice(1, 1);
1772
1773 let mask = BooleanArray::from(vec![true, false]);
1774 let out = zip(&mask, &Scalar::new(&truthy), &Scalar::new(&falsy)).unwrap();
1775 out.to_data().validate_full().unwrap();
1776 let actual = out.as_string_view();
1777 let expected =
1778 StringViewArray::from(vec!["a string longer than twelve bytes", "thirteen byte"]);
1779 assert_eq!(actual, &expected);
1780 }
1781
1782 #[test]
1783 fn test_zip_kernel_scalar_binary_array_view_sliced_inline_max_len() {
1784 let truthy = BinaryViewArray::from_iter_values([
1786 b"x".as_slice(),
1787 b"another value longer than twelve",
1788 ])
1789 .slice(0, 1);
1790 let falsy = BinaryViewArray::from_iter_values([
1791 [0xFFu8; 12].as_slice(),
1792 b"a value longer than twelve bytes",
1793 ])
1794 .slice(0, 1);
1795
1796 let mask = BooleanArray::from(vec![true, false]);
1797 let out = zip(&mask, &Scalar::new(&truthy), &Scalar::new(&falsy)).unwrap();
1798 out.to_data().validate_full().unwrap();
1799 let actual = out.as_binary_view();
1800 let expected = BinaryViewArray::from_iter_values([b"x".as_slice(), &[0xFFu8; 12]]);
1801 assert_eq!(actual, &expected);
1802 }
1803}