Skip to main content

arrow_select/
zip.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! [`zip`]: Combine values from two arrays based on boolean mask
6
7use 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
27/// Zip two arrays by some boolean mask.
28///
29/// - Where `mask` is `true`, values of `truthy` are taken
30/// - Where `mask` is `false` or `NULL`, values of `falsy` are taken
31///
32/// # Example: `zip` two arrays
33/// ```
34/// # use std::sync::Arc;
35/// # use arrow_array::{ArrayRef, BooleanArray, Int32Array};
36/// # use arrow_select::zip::zip;
37/// // mask: [true, true, false, NULL, true]
38/// let mask = BooleanArray::from(vec![
39///   Some(true), Some(true), Some(false), None, Some(true)
40/// ]);
41/// // truthy array: [1, NULL, 3, 4, 5]
42/// let truthy = Int32Array::from(vec![
43///   Some(1), None, Some(3), Some(4), Some(5)
44/// ]);
45/// // falsy array: [10, 20, 30, 40, 50]
46/// let falsy = Int32Array::from(vec![
47///   Some(10), Some(20), Some(30), Some(40), Some(50)
48/// ]);
49/// // zip with this mask select the first, second and last value from `truthy`
50/// // and the third and fourth value from `falsy`
51/// let result = zip(&mask, &truthy, &falsy).unwrap();
52/// // Expected: [1, NULL, 30, 40, 5]
53/// let expected: ArrayRef = Arc::new(Int32Array::from(vec![
54///   Some(1), None, Some(30), Some(40), Some(5)
55/// ]));
56/// assert_eq!(&result, &expected);
57/// ```
58///
59/// # Example: `zip` and array with a scalar
60///
61/// Use `zip` to replace certain values in an array with a scalar
62///
63/// ```
64/// # use std::sync::Arc;
65/// # use arrow_array::{ArrayRef, BooleanArray, Int32Array};
66/// # use arrow_select::zip::zip;
67/// // mask: [true, true, false, NULL, true]
68/// let mask = BooleanArray::from(vec![
69///   Some(true), Some(true), Some(false), None, Some(true)
70/// ]);
71/// //  array: [1, NULL, 3, 4, 5]
72/// let arr = Int32Array::from(vec![
73///   Some(1), None, Some(3), Some(4), Some(5)
74/// ]);
75/// // scalar: 42
76/// let scalar = Int32Array::new_scalar(42);
77/// // zip the array with the  mask select the first, second and last value from `arr`
78/// // and fill the third and fourth value with the scalar 42
79/// let result = zip(&mask, &arr, &scalar).unwrap();
80/// // Expected: [1, NULL, 42, 42, 5]
81/// let expected: ArrayRef = Arc::new(Int32Array::from(vec![
82///   Some(1), None, Some(42), Some(42), Some(5)
83/// ]));
84/// assert_eq!(&result, &expected);
85/// ```
86pub 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
146/// [`zip`] of a primitive `array` with a `scalar` on the other side, or `None` for the types it
147/// does not handle. `scalar_is_truthy` says which side the scalar is on.
148///
149/// The output values are a select per row between the array's value and the scalar's, and its
150/// null buffer has a closed form, so no row goes through [`MutableArrayData`]. With a null
151/// scalar the values under the scalar's rows are never read, so the array's values buffer is
152/// shared as it is.
153fn 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    // The rows that take the array's value
165    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
199/// An integer lane that [`blend`] selects with.
200trait BlendLane:
201    ArrowNativeType + BitAnd<Output = Self> + BitOr<Output = Self> + Not<Output = Self>
202{
203    /// `mask` (0 or -1) at this width.
204    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/// For each bit of `word`, an `i32` with all bits set if it is set, else clear.
227#[inline]
228fn lane_masks(word: u64) -> [i32; 64] {
229    // Each mask tests a constant bit of one byte of the word. This compiles without a branch per
230    // bit on x86_64 (SSE2) and aarch64.
231    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
243/// The values of `array` where `from_array` is set and the scalar's value elsewhere. Each value
244/// is read as `K` integers `U`, so a 16 or 32 byte value is 2 or 4 lanes of `i64` under one mask.
245/// `None` if either values buffer is not aligned to `U`.
246fn 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    // Each block of 64 rows is appended as a whole, so the output is written once and never
263    // zeroed first.
264    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                // The lane masks go to an array first. Computed inside the select, they let LLVM
276                // turn it into a branch per row around a conditional load, which mispredicts on
277                // every other row of a random mask.
278                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    // Interleave's fixed dispatch and index construction costs are not competitive
312    // for small arrays. For larger arrays, use the run count that determines the
313    // amount of work performed by the MutableArrayData implementation.
314    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            // A scalar is a single row array, so every row reads its first value
336            (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        // `MutableArrayData` concatenates the two dictionaries, which can overflow
357        // the key type even when their merged values fit, so let `interleave`
358        // merge them instead
359        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    // the SlicesIterator slices only the true values. So the gaps left by this iterator we need to
372    // fill with falsy values
373
374    // keep track of how much is filled
375    let mut filled = 0;
376
377    for (start, end) in SlicesIterator::from(&mask_buffer) {
378        // the gap needs to be filled with falsy values
379        if start > filled {
380            if falsy_is_scalar {
381                for _ in filled..start {
382                    // Copy the first item from the 'falsy' array into the output buffer.
383                    mutable.try_extend(1, 0, 1)?;
384                }
385            } else {
386                mutable.try_extend(1, filled, start)?;
387            }
388        }
389        // fill with truthy values
390        if truthy_is_scalar {
391            for _ in start..end {
392                // Copy the first item from the 'truthy' array into the output buffer.
393                mutable.try_extend(0, 0, 1)?;
394            }
395        } else {
396            mutable.try_extend(0, start, end)?;
397        }
398        filled = end;
399    }
400    // the remaining part is falsy
401    if filled < mask.len() {
402        if falsy_is_scalar {
403            for _ in filled..mask.len() {
404                // Copy the first item from the 'falsy' array into the output buffer.
405                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/// Zipper for 2 scalars
417///
418/// Useful for using in `IF <expr> THEN <scalar> ELSE <scalar> END` expressions
419///
420/// # Example
421/// ```
422/// # use std::sync::Arc;
423/// # use arrow_array::{ArrayRef, BooleanArray, Int32Array, Scalar, cast::AsArray, types::Int32Type};
424///
425/// # use arrow_select::zip::ScalarZipper;
426/// let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
427/// let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
428/// let zipper = ScalarZipper::try_new(&scalar_truthy, &scalar_falsy).unwrap();
429///
430/// // Later when we have a boolean mask
431/// let mask = BooleanArray::from(vec![true, false, true, false, true]);
432/// let result = zipper.zip(&mask).unwrap();
433/// let actual = result.as_primitive::<Int32Type>();
434/// let expected = Int32Array::from(vec![Some(42), Some(123), Some(42), Some(123), Some(42)]);
435/// ```
436///
437#[derive(Debug, Clone)]
438pub struct ScalarZipper {
439    zip_impl: Arc<dyn ZipImpl>,
440}
441
442impl ScalarZipper {
443    /// Try to create a new ScalarZipper from two scalar Datum
444    ///
445    /// # Errors
446    /// returns error if:
447    /// - the two Datum have different data types
448    /// - either Datum is not a scalar (or has more than 1 element)
449    ///
450    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    /// Creating output array based on input boolean array and the two scalar values the zipper was created with
518    /// See struct level documentation for examples.
519    pub fn zip(&self, mask: &BooleanArray) -> Result<ArrayRef, ArrowError> {
520        self.zip_impl.create_output(mask)
521    }
522}
523
524/// Impl for creating output array based on a mask
525trait ZipImpl: Debug + Send + Sync {
526    /// Creating output array based on input boolean array
527    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    /// return an output array that has
587    /// `value` in all locations where predicate is true
588    /// `null` otherwise
589    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        // Nulls are treated as false
605        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                // If a value is true we need the TRUTHY and the null buffer will have 1 (meaning not null)
619                // If a value is false we need the FALSY and the null buffer will have 0 (meaning null)
620
621                Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, truthy_val)
622            }
623            (None, Some(falsy_val)) => {
624                // Flipping the boolean buffer as we want the opposite of the TRUE case
625                //
626                // if the condition is true we want null so we need to NOT the value so we get 0 (meaning null)
627                // if the condition is false we want the FALSY value so we need to NOT the value so we get 1 (meaning not null)
628                let predicate = predicate.not();
629
630                Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, falsy_val)
631            }
632            (None, None) => {
633                // All values are null
634                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        // Keep decimal precisions, scales or timestamps timezones
645        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    /// return an output array that has
687    /// `value` in all locations where predicate is true
688    /// `null` otherwise
689    #[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        // Fast path for all nulls
699        if number_of_true == 0 {
700            // All values are null
701            let nulls = NullBuffer::new_null(predicate.len());
702
703            return Ok((
704                // Empty bytes
705                Buffer::from(&[]),
706                // All nulls so all lengths are 0
707                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        // If a value is true we need the TRUTHY and the null buffer will have 1 (meaning not null)
724        // If a value is false we need the FALSY and the null buffer will have 0 (meaning null)
725        let nulls = NullBuffer::new(predicate);
726
727        Ok((bytes, offsets, Some(nulls)))
728    }
729
730    /// Create a [`Buffer`] where `value` slice is repeated `number_of_values` times
731    /// and [`OffsetBuffer`] where there are `number_of_values` lengths, and all equals to `value` length
732    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                // All values are falsy
760                return Self::get_bytes_and_offset_for_all_same_value(predicate.len(), falsy_val);
761            }
762            n if n == predicate.len() => {
763                // All values are truthy
764                return Self::get_bytes_and_offset_for_all_same_value(predicate.len(), truthy_val);
765            }
766
767            _ => {
768                // Fallback
769            }
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        // keep track of how much is filled
788        let mut filled = 0;
789
790        let truthy_len = truthy_val.len();
791        let falsy_len = falsy_val.len();
792
793        // Each run's offsets are bounded by the checked total output size.
794        SlicesIterator::from(predicate).try_for_each(|(start, end)| -> Result<(), ArrowError> {
795            // the gap needs to be filled with falsy values
796            if start > filled {
797                let false_repeat_count = start - filled;
798                // Push false value `repeat_count` times
799                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            // fill with truthy values
813            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        // the remaining part is falsy
827        if filled < predicate.len() {
828            let false_repeat_count = predicate.len() - filled;
829            // Copy the first item from the 'falsy' array into the output buffer.
830            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        // SAFETY: offsets start at zero, are monotonically increasing, and fit in T::Offset.
844        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        // Nulls are treated as false
853        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                    // Flipping the boolean buffer as we want the opposite of the TRUE case
868                    //
869                    // if the condition is true we want null so we need to NOT the value so we get 0 (meaning null)
870                    // if the condition is false we want the FALSE value so we need to NOT the value so we get 1 (meaning not null)
871                    let predicate = predicate.not();
872                    Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, falsy_val)?
873                }
874                (None, None) => {
875                    // All values are null
876                    let nulls = NullBuffer::new_null(result_len);
877
878                    (
879                        // Empty bytes
880                        Buffer::from(&[]),
881                        // All nulls so all lengths are 0
882                        OffsetBuffer::<T::Offset>::new_zeroed(predicate.len()),
883                        Some(nulls),
884                    )
885                }
886            };
887
888        let output = unsafe {
889            // Safety: the values are based on valid inputs
890            // and `try_new` is expensive for strings as it validate that the input is valid utf8
891            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    // Nulls are treated as false
900    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        // Fast path for all nulls
953        if number_of_true == 0 {
954            // All values are null
955            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        // If value is true and we want to handle the TRUTHY case, the null buffer will have 1 (meaning not null)
964        // If value is false and we want to handle the FALSY case, the null buffer will have 0 (meaning null)
965        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                // all values are falsy
982                Ok((vec![falsy_view; result_len].into(), falsy_buffers, None))
983            }
984            n if n == predicate.len() => {
985                // all values are truthy
986                Ok((vec![truthy_view; result_len].into(), truthy_buffers, None))
987            }
988            _ => {
989                let mut buffers: Vec<Buffer> = truthy_buffers.to_vec();
990
991                // Check if we need to keep falsy buffer in case view isn't inline
992                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        // Nulls are treated as false
1062        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                // All values are null
1088                (
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        // Exercise runs crossing 64-bit chunk boundaries and trailing padding.
1122        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        // Exercise a non-zero bit offset, as masks may be sliced.
1129        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        // Concatenating two dictionaries of 100 values overflows Int8 keys, but
1138        // the values are the same in both, so the merged dictionary fits
1139        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        // Exercise arbitrary bit offsets as masks may be sliced
1200        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, // null treated as false even though in the original mask it was true
1389                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), // true in mask but null
1400            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, // null treated as false even though in the original mask it was true
1417                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), // true in mask but null
1428            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, // null treated as false even though in the original mask it was true
1444                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", // true in mask but null
1452            "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, // null treated as false even though in the original mask it was true
1468                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"), // true in mask but null
1479            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        // mask ends with false
1521        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        // Repeating a 64 KiB scalar 32,768 times exceeds i32::MAX.
1538        // The size must be rejected before allocating the output buffer.
1539        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            // 1. Test where left is null and right is non-null
1564            //    and mask is all true
1565            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            // 2. Test where left is null and right is non-null
1573            //    and mask is all false
1574            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            // 3. Test where left is non-null and right is null
1582            //    and mask is all true
1583            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            // 4. Test where left is non-null and right is null
1591            //    and mask is all false
1592            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            // 5. Test where both left and right are not null
1600            //    and mask is all true
1601            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            // 6. Test where both left and right are not null
1609            //    and mask is all false
1610            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            // 7. Test where both left and right are null
1619            //    and mask is random
1620            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        // test with different mask length as well
1642        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            // Non valid UTF8 bytes
1673            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 to ensure that the precision and scale are kept when zipping Decimal128 data
1711    #[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 to ensure that the timezone is kept when zipping TimestampArray data
1729    #[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        // non null Scalar vs non null Scalar
1753        test_zip_output_data_type(&scalar_1, &scalar_2, 10);
1754
1755        // null Scalar vs non-null Scalar (and vice versa)
1756        test_zip_output_data_type(&null_scalar, &scalar_1, 10);
1757        test_zip_output_data_type(&scalar_1, &null_scalar, 10);
1758
1759        // non-null Scalar and array (and vice versa)
1760        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        // Array and null scalar (and vice versa)
1764        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        // Both arrays
1769        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        // Try different masks to test different paths
1777        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        // Both scalars are inlined, but each slice keeps its array's data buffer
1973        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        // A 13 byte falsy value is not inlined, so it must be remapped after the truthy buffers
1989        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        // A 12 byte inlined value whose bytes 8 to 11 are 0xFF must not be read as a buffer index
2005        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    /// Primitive data of `data_type` with random bytes and nulls, sliced at `offset`.
2025    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}