arrow_select/
zip.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! [`zip`]: Combine values from two arrays based on boolean mask
19
20use 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    BooleanBuffer, Buffer, MutableBuffer, NullBuffer, OffsetBuffer, OffsetBufferBuilder,
29    ScalarBuffer, ToByteSlice,
30};
31use arrow_data::transform::MutableArrayData;
32use arrow_data::{ArrayData, ByteView};
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
40/// Zip two arrays by some boolean mask.
41///
42/// - Where `mask` is `true`, values of `truthy` are taken
43/// - Where `mask` is `false` or `NULL`, values of `falsy` are taken
44///
45/// # Example: `zip` two arrays
46/// ```
47/// # use std::sync::Arc;
48/// # use arrow_array::{ArrayRef, BooleanArray, Int32Array};
49/// # use arrow_select::zip::zip;
50/// // mask: [true, true, false, NULL, true]
51/// let mask = BooleanArray::from(vec![
52///   Some(true), Some(true), Some(false), None, Some(true)
53/// ]);
54/// // truthy array: [1, NULL, 3, 4, 5]
55/// let truthy = Int32Array::from(vec![
56///   Some(1), None, Some(3), Some(4), Some(5)
57/// ]);
58/// // falsy array: [10, 20, 30, 40, 50]
59/// let falsy = Int32Array::from(vec![
60///   Some(10), Some(20), Some(30), Some(40), Some(50)
61/// ]);
62/// // zip with this mask select the first, second and last value from `truthy`
63/// // and the third and fourth value from `falsy`
64/// let result = zip(&mask, &truthy, &falsy).unwrap();
65/// // Expected: [1, NULL, 30, 40, 5]
66/// let expected: ArrayRef = Arc::new(Int32Array::from(vec![
67///   Some(1), None, Some(30), Some(40), Some(5)
68/// ]));
69/// assert_eq!(&result, &expected);
70/// ```
71///
72/// # Example: `zip` and array with a scalar
73///
74/// Use `zip` to replace certain values in an array with a scalar
75///
76/// ```
77/// # use std::sync::Arc;
78/// # use arrow_array::{ArrayRef, BooleanArray, Int32Array};
79/// # use arrow_select::zip::zip;
80/// // mask: [true, true, false, NULL, true]
81/// let mask = BooleanArray::from(vec![
82///   Some(true), Some(true), Some(false), None, Some(true)
83/// ]);
84/// //  array: [1, NULL, 3, 4, 5]
85/// let arr = Int32Array::from(vec![
86///   Some(1), None, Some(3), Some(4), Some(5)
87/// ]);
88/// // scalar: 42
89/// let scalar = Int32Array::new_scalar(42);
90/// // zip the array with the  mask select the first, second and last value from `arr`
91/// // and fill the third and fourth value with the scalar 42
92/// let result = zip(&mask, &arr, &scalar).unwrap();
93/// // Expected: [1, NULL, 42, 42, 5]
94/// let expected: ArrayRef = Arc::new(Int32Array::from(vec![
95///   Some(1), None, Some(42), Some(42), Some(5)
96/// ]));
97/// assert_eq!(&result, &expected);
98/// ```
99pub 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    // Interleave's fixed dispatch and index construction costs are not competitive
163    // for small arrays. For larger arrays, use the run count that determines the
164    // amount of work performed by the MutableArrayData implementation.
165    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    // the SlicesIterator slices only the true values. So the gaps left by this iterator we need to
198    // fill with falsy values
199
200    // keep track of how much is filled
201    let mut filled = 0;
202
203    for (start, end) in SlicesIterator::from(&mask_buffer) {
204        // the gap needs to be filled with falsy values
205        if start > filled {
206            if falsy_is_scalar {
207                for _ in filled..start {
208                    // Copy the first item from the 'falsy' array into the output buffer.
209                    mutable.try_extend(1, 0, 1)?;
210                }
211            } else {
212                mutable.try_extend(1, filled, start)?;
213            }
214        }
215        // fill with truthy values
216        if truthy_is_scalar {
217            for _ in start..end {
218                // Copy the first item from the 'truthy' array into the output buffer.
219                mutable.try_extend(0, 0, 1)?;
220            }
221        } else {
222            mutable.try_extend(0, start, end)?;
223        }
224        filled = end;
225    }
226    // the remaining part is falsy
227    if filled < mask.len() {
228        if falsy_is_scalar {
229            for _ in filled..mask.len() {
230                // Copy the first item from the 'falsy' array into the output buffer.
231                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/// Zipper for 2 scalars
243///
244/// Useful for using in `IF <expr> THEN <scalar> ELSE <scalar> END` expressions
245///
246/// # Example
247/// ```
248/// # use std::sync::Arc;
249/// # use arrow_array::{ArrayRef, BooleanArray, Int32Array, Scalar, cast::AsArray, types::Int32Type};
250///
251/// # use arrow_select::zip::ScalarZipper;
252/// let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
253/// let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
254/// let zipper = ScalarZipper::try_new(&scalar_truthy, &scalar_falsy).unwrap();
255///
256/// // Later when we have a boolean mask
257/// let mask = BooleanArray::from(vec![true, false, true, false, true]);
258/// let result = zipper.zip(&mask).unwrap();
259/// let actual = result.as_primitive::<Int32Type>();
260/// let expected = Int32Array::from(vec![Some(42), Some(123), Some(42), Some(123), Some(42)]);
261/// ```
262///
263#[derive(Debug, Clone)]
264pub struct ScalarZipper {
265    zip_impl: Arc<dyn ZipImpl>,
266}
267
268impl ScalarZipper {
269    /// Try to create a new ScalarZipper from two scalar Datum
270    ///
271    /// # Errors
272    /// returns error if:
273    /// - the two Datum have different data types
274    /// - either Datum is not a scalar (or has more than 1 element)
275    ///
276    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    /// Creating output array based on input boolean array and the two scalar values the zipper was created with
344    /// See struct level documentation for examples.
345    pub fn zip(&self, mask: &BooleanArray) -> Result<ArrayRef, ArrowError> {
346        self.zip_impl.create_output(mask)
347    }
348}
349
350/// Impl for creating output array based on a mask
351trait ZipImpl: Debug + Send + Sync {
352    /// Creating output array based on input boolean array
353    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    /// return an output array that has
413    /// `value` in all locations where predicate is true
414    /// `null` otherwise
415    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        // Nulls are treated as false
431        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                // If a value is true we need the TRUTHY and the null buffer will have 1 (meaning not null)
445                // If a value is false we need the FALSY and the null buffer will have 0 (meaning null)
446
447                Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, truthy_val)
448            }
449            (None, Some(falsy_val)) => {
450                // Flipping the boolean buffer as we want the opposite of the TRUE case
451                //
452                // if the condition is true we want null so we need to NOT the value so we get 0 (meaning null)
453                // if the condition is false we want the FALSY value so we need to NOT the value so we get 1 (meaning not null)
454                let predicate = predicate.not();
455
456                Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, falsy_val)
457            }
458            (None, None) => {
459                // All values are null
460                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        // Keep decimal precisions, scales or timestamps timezones
471        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    /// return an output array that has
513    /// `value` in all locations where predicate is true
514    /// `null` otherwise
515    #[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        // Fast path for all nulls
525        if number_of_true == 0 {
526            // All values are null
527            let nulls = NullBuffer::new_null(predicate.len());
528
529            return Ok((
530                // Empty bytes
531                Buffer::from(&[]),
532                // All nulls so all lengths are 0
533                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        // If a value is true we need the TRUTHY and the null buffer will have 1 (meaning not null)
550        // If a value is false we need the FALSY and the null buffer will have 0 (meaning null)
551        let nulls = NullBuffer::new(predicate);
552
553        Ok((bytes, offsets, Some(nulls)))
554    }
555
556    /// Create a [`Buffer`] where `value` slice is repeated `number_of_values` times
557    /// and [`OffsetBuffer`] where there are `number_of_values` lengths, and all equals to `value` length
558    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                // All values are falsy
586                return Self::get_bytes_and_offset_for_all_same_value(predicate.len(), falsy_val);
587            }
588            n if n == predicate.len() => {
589                // All values are truthy
590                return Self::get_bytes_and_offset_for_all_same_value(predicate.len(), truthy_val);
591            }
592
593            _ => {
594                // Fallback
595            }
596        }
597
598        let total_number_of_bytes =
599            true_count * truthy_val.len() + (predicate.len() - true_count) * falsy_val.len();
600        let mut mutable = MutableBuffer::with_capacity(total_number_of_bytes);
601        let mut offset_buffer_builder = OffsetBufferBuilder::<T::Offset>::new(predicate.len());
602
603        // keep track of how much is filled
604        let mut filled = 0;
605
606        let truthy_len = truthy_val.len();
607        let falsy_len = falsy_val.len();
608
609        SlicesIterator::from(predicate).try_for_each(|(start, end)| -> Result<(), ArrowError> {
610            // the gap needs to be filled with falsy values
611            if start > filled {
612                let false_repeat_count = start - filled;
613                // Push false value `repeat_count` times
614                mutable
615                    .try_repeat_slice_n_times(falsy_val, false_repeat_count)
616                    .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
617
618                for _ in 0..false_repeat_count {
619                    offset_buffer_builder.push_length(falsy_len)
620                }
621            }
622
623            let true_repeat_count = end - start;
624            // fill with truthy values
625            mutable
626                .try_repeat_slice_n_times(truthy_val, true_repeat_count)
627                .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
628
629            for _ in 0..true_repeat_count {
630                offset_buffer_builder.push_length(truthy_len)
631            }
632            filled = end;
633            Ok(())
634        })?;
635        // the remaining part is falsy
636        if filled < predicate.len() {
637            let false_repeat_count = predicate.len() - filled;
638            // Copy the first item from the 'falsy' array into the output buffer.
639            mutable
640                .try_repeat_slice_n_times(falsy_val, false_repeat_count)
641                .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
642
643            for _ in 0..false_repeat_count {
644                offset_buffer_builder.push_length(falsy_len)
645            }
646        }
647
648        Ok((mutable.into(), offset_buffer_builder.finish()))
649    }
650}
651
652impl<T: ByteArrayType> ZipImpl for BytesScalarImpl<T> {
653    fn create_output(&self, predicate: &BooleanArray) -> Result<ArrayRef, ArrowError> {
654        let result_len = predicate.len();
655        // Nulls are treated as false
656        let predicate = maybe_prep_null_mask_filter(predicate);
657
658        let (bytes, offsets, nulls): (Buffer, OffsetBuffer<T::Offset>, Option<NullBuffer>) =
659            match (self.truthy.as_deref(), self.falsy.as_deref()) {
660                (Some(truthy_val), Some(falsy_val)) => {
661                    let (bytes, offsets) =
662                        Self::create_output_on_non_nulls(&predicate, truthy_val, falsy_val)?;
663
664                    (bytes, offsets, None)
665                }
666                (Some(truthy_val), None) => {
667                    Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, truthy_val)?
668                }
669                (None, Some(falsy_val)) => {
670                    // Flipping the boolean buffer as we want the opposite of the TRUE case
671                    //
672                    // if the condition is true we want null so we need to NOT the value so we get 0 (meaning null)
673                    // if the condition is false we want the FALSE value so we need to NOT the value so we get 1 (meaning not null)
674                    let predicate = predicate.not();
675                    Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, falsy_val)?
676                }
677                (None, None) => {
678                    // All values are null
679                    let nulls = NullBuffer::new_null(result_len);
680
681                    (
682                        // Empty bytes
683                        Buffer::from(&[]),
684                        // All nulls so all lengths are 0
685                        OffsetBuffer::<T::Offset>::new_zeroed(predicate.len()),
686                        Some(nulls),
687                    )
688                }
689            };
690
691        let output = unsafe {
692            // Safety: the values are based on valid inputs
693            // and `try_new` is expensive for strings as it validate that the input is valid utf8
694            GenericByteArray::<T>::new_unchecked(offsets, bytes, nulls)
695        };
696
697        Ok(Arc::new(output))
698    }
699}
700
701fn maybe_prep_null_mask_filter(predicate: &BooleanArray) -> BooleanBuffer {
702    // Nulls are treated as false
703    if predicate.null_count() == 0 {
704        predicate.values().clone()
705    } else {
706        let cleaned = prep_null_mask_filter(predicate);
707        let (boolean_buffer, _) = cleaned.into_parts();
708        boolean_buffer
709    }
710}
711
712struct ByteViewScalarImpl<T: ByteViewType> {
713    truthy_view: Option<u128>,
714    truthy_buffers: Arc<[Buffer]>,
715    falsy_view: Option<u128>,
716    falsy_buffers: Arc<[Buffer]>,
717    phantom: PhantomData<T>,
718}
719
720static EMPTY_ARC: OnceLock<Arc<[Buffer]>> = OnceLock::new();
721fn empty_arc_buffers() -> Arc<[Buffer]> {
722    Arc::clone(EMPTY_ARC.get_or_init(|| Arc::new([])))
723}
724
725impl<T: ByteViewType> ByteViewScalarImpl<T> {
726    fn new(truthy: &dyn Array, falsy: &dyn Array) -> Self {
727        let (truthy_view, truthy_buffers) = Self::get_value_from_scalar(truthy);
728        let (falsy_view, falsy_buffers) = Self::get_value_from_scalar(falsy);
729        Self {
730            truthy_view,
731            truthy_buffers,
732            falsy_view,
733            falsy_buffers,
734            phantom: PhantomData,
735        }
736    }
737
738    fn get_value_from_scalar(scalar: &dyn Array) -> (Option<u128>, Arc<[Buffer]>) {
739        if scalar.is_null(0) {
740            (None, empty_arc_buffers())
741        } else {
742            let (views, buffers, _) = scalar.as_byte_view::<T>().clone().into_parts();
743            (views.first().copied(), buffers)
744        }
745    }
746
747    fn get_views_for_single_non_nullable(
748        predicate: BooleanBuffer,
749        value: u128,
750        buffers: Arc<[Buffer]>,
751    ) -> (ScalarBuffer<u128>, Arc<[Buffer]>, Option<NullBuffer>) {
752        let number_of_true = predicate.count_set_bits();
753        let number_of_values = predicate.len();
754
755        // Fast path for all nulls
756        if number_of_true == 0 {
757            // All values are null
758            return (
759                vec![0; number_of_values].into(),
760                empty_arc_buffers(),
761                Some(NullBuffer::new_null(number_of_values)),
762            );
763        }
764        let bytes = vec![value; number_of_values];
765
766        // If value is true and we want to handle the TRUTHY case, the null buffer will have 1 (meaning not null)
767        // If value is false and we want to handle the FALSY case, the null buffer will have 0 (meaning null)
768        let nulls = NullBuffer::new(predicate);
769        (bytes.into(), buffers, Some(nulls))
770    }
771
772    #[expect(clippy::type_complexity)]
773    fn get_views_for_non_nullable(
774        predicate: BooleanBuffer,
775        result_len: usize,
776        truthy_view: u128,
777        truthy_buffers: Arc<[Buffer]>,
778        falsy_view: u128,
779        falsy_buffers: Arc<[Buffer]>,
780    ) -> Result<(ScalarBuffer<u128>, Arc<[Buffer]>, Option<NullBuffer>), ArrowError> {
781        let true_count = predicate.count_set_bits();
782        match true_count {
783            0 => {
784                // all values are falsy
785                Ok((vec![falsy_view; result_len].into(), falsy_buffers, None))
786            }
787            n if n == predicate.len() => {
788                // all values are truthy
789                Ok((vec![truthy_view; result_len].into(), truthy_buffers, None))
790            }
791            _ => {
792                let true_count = predicate.count_set_bits();
793                let mut buffers: Vec<Buffer> = truthy_buffers.to_vec();
794
795                // If the falsy buffers are empty, we can use the falsy view as it is, because the value
796                // is completely inlined. Otherwise, we have non-inlined values in the buffer, and we need
797                // to recalculate the falsy view
798                let view_falsy = if falsy_buffers.is_empty() {
799                    falsy_view
800                } else {
801                    let byte_view_falsy = ByteView::from(falsy_view);
802                    let new_index_falsy_buffers =
803                        buffers.len() as u32 + byte_view_falsy.buffer_index;
804                    buffers.extend(falsy_buffers.iter().cloned());
805                    let byte_view_falsy =
806                        byte_view_falsy.with_buffer_index(new_index_falsy_buffers);
807                    byte_view_falsy.as_u128()
808                };
809
810                let total_number_of_bytes = true_count * 16 + (predicate.len() - true_count) * 16;
811                let mut mutable = MutableBuffer::new(total_number_of_bytes);
812                let mut filled = 0;
813
814                SlicesIterator::from(&predicate).try_for_each(
815                    |(start, end)| -> Result<(), ArrowError> {
816                        if start > filled {
817                            let false_repeat_count = start - filled;
818                            mutable
819                                .try_repeat_slice_n_times(
820                                    view_falsy.to_byte_slice(),
821                                    false_repeat_count,
822                                )
823                                .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
824                        }
825                        let true_repeat_count = end - start;
826                        mutable
827                            .try_repeat_slice_n_times(
828                                truthy_view.to_byte_slice(),
829                                true_repeat_count,
830                            )
831                            .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
832                        filled = end;
833                        Ok(())
834                    },
835                )?;
836
837                if filled < predicate.len() {
838                    let false_repeat_count = predicate.len() - filled;
839                    mutable
840                        .try_repeat_slice_n_times(view_falsy.to_byte_slice(), false_repeat_count)
841                        .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
842                }
843
844                let bytes = Buffer::from(mutable);
845                Ok((bytes.into(), buffers.into(), None))
846            }
847        }
848    }
849}
850
851impl<T: ByteViewType> Debug for ByteViewScalarImpl<T> {
852    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
853        f.debug_struct("ByteViewScalarImpl")
854            .field("truthy", &self.truthy_view)
855            .field("falsy", &self.falsy_view)
856            .finish()
857    }
858}
859
860impl<T: ByteViewType> ZipImpl for ByteViewScalarImpl<T> {
861    fn create_output(&self, predicate: &BooleanArray) -> Result<ArrayRef, ArrowError> {
862        let result_len = predicate.len();
863        // Nulls are treated as false
864        let predicate = maybe_prep_null_mask_filter(predicate);
865
866        let (views, buffers, nulls) = match (self.truthy_view, self.falsy_view) {
867            (Some(truthy), Some(falsy)) => Self::get_views_for_non_nullable(
868                predicate,
869                result_len,
870                truthy,
871                Arc::clone(&self.truthy_buffers),
872                falsy,
873                Arc::clone(&self.falsy_buffers),
874            )?,
875            (Some(truthy), None) => Self::get_views_for_single_non_nullable(
876                predicate,
877                truthy,
878                Arc::clone(&self.truthy_buffers),
879            ),
880            (None, Some(falsy)) => {
881                let predicate = predicate.not();
882                Self::get_views_for_single_non_nullable(
883                    predicate,
884                    falsy,
885                    Arc::clone(&self.falsy_buffers),
886                )
887            }
888            (None, None) => {
889                // All values are null
890                (
891                    vec![0; result_len].into(),
892                    empty_arc_buffers(),
893                    Some(NullBuffer::new_null(result_len)),
894                )
895            }
896        };
897
898        let result = unsafe { GenericByteViewArray::<T>::new_unchecked(views, buffers, nulls) };
899        Ok(Arc::new(result))
900    }
901}
902
903#[cfg(test)]
904mod test {
905    use super::*;
906    use arrow_array::types::Int32Type;
907
908    #[test]
909    fn test_count_true_runs() {
910        let assert_runs = |values: &[bool], expected| {
911            let mask: BooleanBuffer = values.iter().copied().collect();
912            assert_eq!(count_true_runs(&mask), expected, "mask: {values:?}");
913        };
914
915        assert_runs(&[], 0);
916        assert_runs(&[false, false, false], 0);
917        assert_runs(&[true, true, true], 1);
918        assert_runs(&[true, false, true, true, false, true], 3);
919
920        // Exercise runs crossing 64-bit chunk boundaries and trailing padding.
921        let mut values = vec![false; 130];
922        values[0] = true;
923        values[63..66].fill(true);
924        values[128..].fill(true);
925        assert_runs(&values, 3);
926
927        // Exercise a non-zero bit offset, as masks may be sliced.
928        let mut offset_values = vec![false; 135];
929        offset_values[3..133].copy_from_slice(&values);
930        let offset_mask: BooleanBuffer = offset_values.into_iter().collect();
931        assert_eq!(count_true_runs(&offset_mask.slice(3, 130)), 3);
932    }
933
934    #[test]
935    fn test_should_use_interleave() {
936        let short: BooleanBuffer = (0..64).map(|i| i % 2 == 0).collect();
937        assert!(!should_use_interleave(&short));
938
939        let fragmented: BooleanBuffer = (0..8192).map(|i| i % 2 == 0).collect();
940        assert!(should_use_interleave(&fragmented));
941
942        let long_runs: BooleanBuffer = (0..8192).map(|i| i < 4096).collect();
943        assert!(!should_use_interleave(&long_runs));
944
945        let sparse: BooleanBuffer = (0..8192).map(|i| i % 10 == 0).collect();
946        assert!(!should_use_interleave(&sparse));
947
948        let dense: BooleanBuffer = (0..8192).map(|i| i % 10 != 0).collect();
949        assert!(!should_use_interleave(&dense));
950
951        let fragmented_head: BooleanBuffer = (0..8192).map(|i| i < 256 && i % 2 == 0).collect();
952        assert!(!should_use_interleave(&fragmented_head));
953
954        let fragmented_edges: BooleanBuffer = (0..8192)
955            .map(|i| !(256..7936).contains(&i) && i % 2 == 0)
956            .collect();
957        assert!(!should_use_interleave(&fragmented_edges));
958
959        // Exercise arbitrary bit offsets as masks may be sliced
960        let offset: BooleanBuffer = (0..8195).map(|i| i >= 3 && i % 2 == 1).collect();
961        assert!(should_use_interleave(&offset.slice(3, 8192)));
962    }
963
964    #[test]
965    fn test_interleave_arrays() {
966        let mask = BooleanArray::from(vec![Some(true), None, Some(true), Some(false)]);
967        let mask = maybe_prep_null_mask_filter(&mask);
968        let truthy = Int32Array::from(vec![Some(1), None, Some(3), Some(4)]).to_data();
969        let falsy = Int32Array::from(vec![Some(10), Some(20), None, Some(40)]).to_data();
970        let expected = Int32Array::from(vec![Some(1), Some(20), Some(3), Some(40)]);
971
972        let actual = interleave_arrays(&mask, &truthy, &falsy).unwrap();
973        assert_eq!(actual.as_primitive::<Int32Type>(), &expected);
974    }
975
976    #[test]
977    fn test_zip_fragmented_array_mask() {
978        let mask: BooleanArray = (0..8192)
979            .map(|i| match i % 3 {
980                0 => Some(true),
981                1 => Some(false),
982                _ => None,
983            })
984            .collect();
985        let truthy: Int32Array = (0..8192).map(|i| (i % 7 != 0).then_some(i)).collect();
986        let falsy: Int32Array = (0..8192).map(|i| (i % 11 != 0).then_some(-i)).collect();
987        let expected: Int32Array = (0..8192)
988            .map(|i| {
989                let array = if i % 3 == 0 { &truthy } else { &falsy };
990                array.is_valid(i).then(|| array.value(i))
991            })
992            .collect();
993
994        let actual = zip(&mask, &truthy, &falsy).unwrap();
995        assert_eq!(actual.as_primitive::<Int32Type>(), &expected);
996    }
997
998    #[test]
999    fn test_zip_kernel_one() {
1000        let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1001        let b = Int32Array::from(vec![None, Some(3), Some(6), Some(7), Some(3)]);
1002        let mask = BooleanArray::from(vec![true, true, false, false, true]);
1003        let out = zip(&mask, &a, &b).unwrap();
1004        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1005        let expected = Int32Array::from(vec![Some(5), None, Some(6), Some(7), Some(1)]);
1006        assert_eq!(actual, &expected);
1007    }
1008
1009    #[test]
1010    fn test_zip_kernel_two() {
1011        let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1012        let b = Int32Array::from(vec![None, Some(3), Some(6), Some(7), Some(3)]);
1013        let mask = BooleanArray::from(vec![false, false, true, true, false]);
1014        let out = zip(&mask, &a, &b).unwrap();
1015        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1016        let expected = Int32Array::from(vec![None, Some(3), Some(7), None, Some(3)]);
1017        assert_eq!(actual, &expected);
1018    }
1019
1020    #[test]
1021    fn test_zip_kernel_scalar_falsy_1() {
1022        let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1023
1024        let fallback = Scalar::new(Int32Array::from_value(42, 1));
1025
1026        let mask = BooleanArray::from(vec![true, true, false, false, true]);
1027        let out = zip(&mask, &a, &fallback).unwrap();
1028        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1029        let expected = Int32Array::from(vec![Some(5), None, Some(42), Some(42), Some(1)]);
1030        assert_eq!(actual, &expected);
1031    }
1032
1033    #[test]
1034    fn test_zip_kernel_scalar_falsy_2() {
1035        let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1036
1037        let fallback = Scalar::new(Int32Array::from_value(42, 1));
1038
1039        let mask = BooleanArray::from(vec![false, false, true, true, false]);
1040        let out = zip(&mask, &a, &fallback).unwrap();
1041        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1042        let expected = Int32Array::from(vec![Some(42), Some(42), Some(7), None, Some(42)]);
1043        assert_eq!(actual, &expected);
1044    }
1045
1046    #[test]
1047    fn test_zip_kernel_scalar_truthy_1() {
1048        let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1049
1050        let fallback = Scalar::new(Int32Array::from_value(42, 1));
1051
1052        let mask = BooleanArray::from(vec![true, true, false, false, true]);
1053        let out = zip(&mask, &fallback, &a).unwrap();
1054        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1055        let expected = Int32Array::from(vec![Some(42), Some(42), Some(7), None, Some(42)]);
1056        assert_eq!(actual, &expected);
1057    }
1058
1059    #[test]
1060    fn test_zip_kernel_scalar_truthy_2() {
1061        let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
1062
1063        let fallback = Scalar::new(Int32Array::from_value(42, 1));
1064
1065        let mask = BooleanArray::from(vec![false, false, true, true, false]);
1066        let out = zip(&mask, &fallback, &a).unwrap();
1067        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1068        let expected = Int32Array::from(vec![Some(5), None, Some(42), Some(42), Some(1)]);
1069        assert_eq!(actual, &expected);
1070    }
1071
1072    #[test]
1073    fn test_zip_kernel_scalar_both_mask_ends_with_true() {
1074        let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1075        let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
1076
1077        let mask = BooleanArray::from(vec![true, true, false, false, true]);
1078        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1079        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1080        let expected = Int32Array::from(vec![Some(42), Some(42), Some(123), Some(123), Some(42)]);
1081        assert_eq!(actual, &expected);
1082    }
1083
1084    #[test]
1085    fn test_zip_kernel_scalar_both_mask_ends_with_false() {
1086        let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1087        let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
1088
1089        let mask = BooleanArray::from(vec![true, true, false, true, false, false]);
1090        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1091        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1092        let expected = Int32Array::from(vec![
1093            Some(42),
1094            Some(42),
1095            Some(123),
1096            Some(42),
1097            Some(123),
1098            Some(123),
1099        ]);
1100        assert_eq!(actual, &expected);
1101    }
1102
1103    #[test]
1104    fn test_zip_kernel_primitive_scalar_none_1() {
1105        let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1106        let scalar_falsy = Scalar::new(Int32Array::new_null(1));
1107
1108        let mask = BooleanArray::from(vec![true, true, false, false, true]);
1109        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1110        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1111        let expected = Int32Array::from(vec![Some(42), Some(42), None, None, Some(42)]);
1112        assert_eq!(actual, &expected);
1113    }
1114
1115    #[test]
1116    fn test_zip_kernel_primitive_scalar_none_2() {
1117        let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1118        let scalar_falsy = Scalar::new(Int32Array::new_null(1));
1119
1120        let mask = BooleanArray::from(vec![false, false, true, true, false]);
1121        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1122        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1123        let expected = Int32Array::from(vec![None, None, Some(42), Some(42), None]);
1124        assert_eq!(actual, &expected);
1125    }
1126
1127    #[test]
1128    fn test_zip_kernel_primitive_scalar_both_null() {
1129        let scalar_truthy = Scalar::new(Int32Array::new_null(1));
1130        let scalar_falsy = Scalar::new(Int32Array::new_null(1));
1131
1132        let mask = BooleanArray::from(vec![false, false, true, true, false]);
1133        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1134        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1135        let expected = Int32Array::from(vec![None, None, None, None, None]);
1136        assert_eq!(actual, &expected);
1137    }
1138
1139    #[test]
1140    fn test_zip_primitive_array_with_nulls_is_mask_should_be_treated_as_false() {
1141        let truthy = Int32Array::from_iter_values(vec![1, 2, 3, 4, 5, 6]);
1142        let falsy = Int32Array::from_iter_values(vec![7, 8, 9, 10, 11, 12]);
1143
1144        let mask = {
1145            let booleans = BooleanBuffer::from(vec![true, true, false, true, false, false]);
1146            let nulls = NullBuffer::from(vec![
1147                true, true, true,
1148                false, // null treated as false even though in the original mask it was true
1149                true, true,
1150            ]);
1151            BooleanArray::new(booleans, Some(nulls))
1152        };
1153        let out = zip(&mask, &truthy, &falsy).unwrap();
1154        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1155        let expected = Int32Array::from(vec![
1156            Some(1),
1157            Some(2),
1158            Some(9),
1159            Some(10), // true in mask but null
1160            Some(11),
1161            Some(12),
1162        ]);
1163        assert_eq!(actual, &expected);
1164    }
1165
1166    #[test]
1167    fn test_zip_kernel_primitive_scalar_with_boolean_array_mask_with_nulls_should_be_treated_as_false()
1168     {
1169        let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1170        let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
1171
1172        let mask = {
1173            let booleans = BooleanBuffer::from(vec![true, true, false, true, false, false]);
1174            let nulls = NullBuffer::from(vec![
1175                true, true, true,
1176                false, // null treated as false even though in the original mask it was true
1177                true, true,
1178            ]);
1179            BooleanArray::new(booleans, Some(nulls))
1180        };
1181        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1182        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1183        let expected = Int32Array::from(vec![
1184            Some(42),
1185            Some(42),
1186            Some(123),
1187            Some(123), // true in mask but null
1188            Some(123),
1189            Some(123),
1190        ]);
1191        assert_eq!(actual, &expected);
1192    }
1193
1194    #[test]
1195    fn test_zip_string_array_with_nulls_is_mask_should_be_treated_as_false() {
1196        let truthy = StringArray::from_iter_values(vec!["1", "2", "3", "4", "5", "6"]);
1197        let falsy = StringArray::from_iter_values(vec!["7", "8", "9", "10", "11", "12"]);
1198
1199        let mask = {
1200            let booleans = BooleanBuffer::from(vec![true, true, false, true, false, false]);
1201            let nulls = NullBuffer::from(vec![
1202                true, true, true,
1203                false, // null treated as false even though in the original mask it was true
1204                true, true,
1205            ]);
1206            BooleanArray::new(booleans, Some(nulls))
1207        };
1208        let out = zip(&mask, &truthy, &falsy).unwrap();
1209        let actual = out.as_string::<i32>();
1210        let expected = StringArray::from_iter_values(vec![
1211            "1", "2", "9", "10", // true in mask but null
1212            "11", "12",
1213        ]);
1214        assert_eq!(actual, &expected);
1215    }
1216
1217    #[test]
1218    fn test_zip_kernel_large_string_scalar_with_boolean_array_mask_with_nulls_should_be_treated_as_false()
1219     {
1220        let scalar_truthy = Scalar::new(LargeStringArray::from_iter_values(["test"]));
1221        let scalar_falsy = Scalar::new(LargeStringArray::from_iter_values(["something else"]));
1222
1223        let mask = {
1224            let booleans = BooleanBuffer::from(vec![true, true, false, true, false, false]);
1225            let nulls = NullBuffer::from(vec![
1226                true, true, true,
1227                false, // null treated as false even though in the original mask it was true
1228                true, true,
1229            ]);
1230            BooleanArray::new(booleans, Some(nulls))
1231        };
1232        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1233        let actual = out.as_any().downcast_ref::<LargeStringArray>().unwrap();
1234        let expected = LargeStringArray::from_iter(vec![
1235            Some("test"),
1236            Some("test"),
1237            Some("something else"),
1238            Some("something else"), // true in mask but null
1239            Some("something else"),
1240            Some("something else"),
1241        ]);
1242        assert_eq!(actual, &expected);
1243    }
1244
1245    #[test]
1246    fn test_zip_kernel_bytes_scalar_none_1() {
1247        let scalar_truthy = Scalar::new(StringArray::from_iter_values(["hello"]));
1248        let scalar_falsy = Scalar::new(StringArray::new_null(1));
1249
1250        let mask = BooleanArray::from(vec![true, true, false, false, true]);
1251        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1252        let actual = out.as_any().downcast_ref::<StringArray>().unwrap();
1253        let expected = StringArray::from_iter(vec![
1254            Some("hello"),
1255            Some("hello"),
1256            None,
1257            None,
1258            Some("hello"),
1259        ]);
1260        assert_eq!(actual, &expected);
1261    }
1262
1263    #[test]
1264    fn test_zip_kernel_bytes_scalar_none_2() {
1265        let scalar_truthy = Scalar::new(StringArray::new_null(1));
1266        let scalar_falsy = Scalar::new(StringArray::from_iter_values(["hello"]));
1267
1268        let mask = BooleanArray::from(vec![true, true, false, false, true]);
1269        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1270        let actual = out.as_any().downcast_ref::<StringArray>().unwrap();
1271        let expected = StringArray::from_iter(vec![None, None, Some("hello"), Some("hello"), None]);
1272        assert_eq!(actual, &expected);
1273    }
1274
1275    #[test]
1276    fn test_zip_kernel_bytes_scalar_both() {
1277        let scalar_truthy = Scalar::new(StringArray::from_iter_values(["test"]));
1278        let scalar_falsy = Scalar::new(StringArray::from_iter_values(["something else"]));
1279
1280        // mask ends with false
1281        let mask = BooleanArray::from(vec![true, true, false, true, false, false]);
1282        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1283        let actual = out.as_any().downcast_ref::<StringArray>().unwrap();
1284        let expected = StringArray::from_iter(vec![
1285            Some("test"),
1286            Some("test"),
1287            Some("something else"),
1288            Some("test"),
1289            Some("something else"),
1290            Some("something else"),
1291        ]);
1292        assert_eq!(actual, &expected);
1293    }
1294
1295    #[test]
1296    fn test_zip_scalar_bytes_only_taking_one_side() {
1297        let mask_len = 5;
1298        let all_true_mask = BooleanArray::from(vec![true; mask_len]);
1299        let all_false_mask = BooleanArray::from(vec![false; mask_len]);
1300
1301        let null_scalar = Scalar::new(StringArray::new_null(1));
1302        let non_null_scalar_1 = Scalar::new(StringArray::from_iter_values(["test"]));
1303        let non_null_scalar_2 = Scalar::new(StringArray::from_iter_values(["something else"]));
1304
1305        {
1306            // 1. Test where left is null and right is non-null
1307            //    and mask is all true
1308            let out = zip(&all_true_mask, &null_scalar, &non_null_scalar_1).unwrap();
1309            let actual = out.as_string::<i32>();
1310            let expected = StringArray::from_iter(std::iter::repeat_n(None::<&str>, mask_len));
1311            assert_eq!(actual, &expected);
1312        }
1313
1314        {
1315            // 2. Test where left is null and right is non-null
1316            //    and mask is all false
1317            let out = zip(&all_false_mask, &null_scalar, &non_null_scalar_1).unwrap();
1318            let actual = out.as_string::<i32>();
1319            let expected = StringArray::from_iter(std::iter::repeat_n(Some("test"), mask_len));
1320            assert_eq!(actual, &expected);
1321        }
1322
1323        {
1324            // 3. Test where left is non-null and right is null
1325            //    and mask is all true
1326            let out = zip(&all_true_mask, &non_null_scalar_1, &null_scalar).unwrap();
1327            let actual = out.as_string::<i32>();
1328            let expected = StringArray::from_iter(std::iter::repeat_n(Some("test"), mask_len));
1329            assert_eq!(actual, &expected);
1330        }
1331
1332        {
1333            // 4. Test where left is non-null and right is null
1334            //    and mask is all false
1335            let out = zip(&all_false_mask, &non_null_scalar_1, &null_scalar).unwrap();
1336            let actual = out.as_string::<i32>();
1337            let expected = StringArray::from_iter(std::iter::repeat_n(None::<&str>, mask_len));
1338            assert_eq!(actual, &expected);
1339        }
1340
1341        {
1342            // 5. Test where both left and right are not null
1343            //    and mask is all true
1344            let out = zip(&all_true_mask, &non_null_scalar_1, &non_null_scalar_2).unwrap();
1345            let actual = out.as_string::<i32>();
1346            let expected = StringArray::from_iter(std::iter::repeat_n(Some("test"), mask_len));
1347            assert_eq!(actual, &expected);
1348        }
1349
1350        {
1351            // 6. Test where both left and right are not null
1352            //    and mask is all false
1353            let out = zip(&all_false_mask, &non_null_scalar_1, &non_null_scalar_2).unwrap();
1354            let actual = out.as_string::<i32>();
1355            let expected =
1356                StringArray::from_iter(std::iter::repeat_n(Some("something else"), mask_len));
1357            assert_eq!(actual, &expected);
1358        }
1359
1360        {
1361            // 7. Test where both left and right are null
1362            //    and mask is random
1363            let mask = BooleanArray::from(vec![true, false, true, false, true]);
1364            let out = zip(&mask, &null_scalar, &null_scalar).unwrap();
1365            let actual = out.as_string::<i32>();
1366            let expected = StringArray::from_iter(std::iter::repeat_n(None::<&str>, mask_len));
1367            assert_eq!(actual, &expected);
1368        }
1369    }
1370
1371    #[test]
1372    fn test_scalar_zipper() {
1373        let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1374        let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
1375
1376        let mask = BooleanArray::from(vec![false, false, true, true, false]);
1377
1378        let scalar_zipper = ScalarZipper::try_new(&scalar_truthy, &scalar_falsy).unwrap();
1379        let out = scalar_zipper.zip(&mask).unwrap();
1380        let actual = out.as_primitive::<Int32Type>();
1381        let expected = Int32Array::from(vec![Some(123), Some(123), Some(42), Some(42), Some(123)]);
1382        assert_eq!(actual, &expected);
1383
1384        // test with different mask length as well
1385        let mask = BooleanArray::from(vec![true, false, true]);
1386        let out = scalar_zipper.zip(&mask).unwrap();
1387        let actual = out.as_primitive::<Int32Type>();
1388        let expected = Int32Array::from(vec![Some(42), Some(123), Some(42)]);
1389        assert_eq!(actual, &expected);
1390    }
1391
1392    #[test]
1393    fn test_zip_kernel_scalar_strings() {
1394        let scalar_truthy = Scalar::new(StringArray::from(vec!["hello"]));
1395        let scalar_falsy = Scalar::new(StringArray::from(vec!["world"]));
1396
1397        let mask = BooleanArray::from(vec![true, false, true, false, true]);
1398        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1399        let actual = out.as_string::<i32>();
1400        let expected = StringArray::from(vec![
1401            Some("hello"),
1402            Some("world"),
1403            Some("hello"),
1404            Some("world"),
1405            Some("hello"),
1406        ]);
1407        assert_eq!(actual, &expected);
1408    }
1409
1410    #[test]
1411    fn test_zip_kernel_scalar_binary() {
1412        let truthy_bytes: &[u8] = b"\xFF\xFE\xFD";
1413        let falsy_bytes: &[u8] = b"world";
1414        let scalar_truthy = Scalar::new(BinaryArray::from_iter_values(
1415            // Non valid UTF8 bytes
1416            vec![truthy_bytes],
1417        ));
1418        let scalar_falsy = Scalar::new(BinaryArray::from_iter_values(vec![falsy_bytes]));
1419
1420        let mask = BooleanArray::from(vec![true, false, true, false, true]);
1421        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1422        let actual = out.as_binary::<i32>();
1423        let expected = BinaryArray::from(vec![
1424            Some(truthy_bytes),
1425            Some(falsy_bytes),
1426            Some(truthy_bytes),
1427            Some(falsy_bytes),
1428            Some(truthy_bytes),
1429        ]);
1430        assert_eq!(actual, &expected);
1431    }
1432
1433    #[test]
1434    fn test_zip_kernel_scalar_large_binary() {
1435        let truthy_bytes: &[u8] = b"hey";
1436        let falsy_bytes: &[u8] = b"world";
1437        let scalar_truthy = Scalar::new(LargeBinaryArray::from_iter_values(vec![truthy_bytes]));
1438        let scalar_falsy = Scalar::new(LargeBinaryArray::from_iter_values(vec![falsy_bytes]));
1439
1440        let mask = BooleanArray::from(vec![true, false, true, false, true]);
1441        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1442        let actual = out.as_binary::<i64>();
1443        let expected = LargeBinaryArray::from(vec![
1444            Some(truthy_bytes),
1445            Some(falsy_bytes),
1446            Some(truthy_bytes),
1447            Some(falsy_bytes),
1448            Some(truthy_bytes),
1449        ]);
1450        assert_eq!(actual, &expected);
1451    }
1452
1453    // Test to ensure that the precision and scale are kept when zipping Decimal128 data
1454    #[test]
1455    fn test_zip_decimal_with_custom_precision_and_scale() {
1456        let arr = Decimal128Array::from_iter_values([12345, 456, 7890, -123223423432432])
1457            .with_precision_and_scale(20, 2)
1458            .unwrap();
1459
1460        let arr: ArrayRef = Arc::new(arr);
1461
1462        let scalar_1 = Scalar::new(arr.slice(0, 1));
1463        let scalar_2 = Scalar::new(arr.slice(1, 1));
1464        let null_scalar = Scalar::new(new_null_array(arr.data_type(), 1));
1465        let array_1: ArrayRef = arr.slice(0, 2);
1466        let array_2: ArrayRef = arr.slice(2, 2);
1467
1468        test_zip_output_data_types_for_input(scalar_1, scalar_2, null_scalar, array_1, array_2);
1469    }
1470
1471    // Test to ensure that the timezone is kept when zipping TimestampArray data
1472    #[test]
1473    fn test_zip_timestamp_with_timezone() {
1474        let arr = TimestampSecondArray::from(vec![0, 1000, 2000, 4000])
1475            .with_timezone("+01:00".to_string());
1476
1477        let arr: ArrayRef = Arc::new(arr);
1478
1479        let scalar_1 = Scalar::new(arr.slice(0, 1));
1480        let scalar_2 = Scalar::new(arr.slice(1, 1));
1481        let null_scalar = Scalar::new(new_null_array(arr.data_type(), 1));
1482        let array_1: ArrayRef = arr.slice(0, 2);
1483        let array_2: ArrayRef = arr.slice(2, 2);
1484
1485        test_zip_output_data_types_for_input(scalar_1, scalar_2, null_scalar, array_1, array_2);
1486    }
1487
1488    fn test_zip_output_data_types_for_input(
1489        scalar_1: Scalar<ArrayRef>,
1490        scalar_2: Scalar<ArrayRef>,
1491        null_scalar: Scalar<ArrayRef>,
1492        array_1: ArrayRef,
1493        array_2: ArrayRef,
1494    ) {
1495        // non null Scalar vs non null Scalar
1496        test_zip_output_data_type(&scalar_1, &scalar_2, 10);
1497
1498        // null Scalar vs non-null Scalar (and vice versa)
1499        test_zip_output_data_type(&null_scalar, &scalar_1, 10);
1500        test_zip_output_data_type(&scalar_1, &null_scalar, 10);
1501
1502        // non-null Scalar and array (and vice versa)
1503        test_zip_output_data_type(&array_1.as_ref(), &scalar_1, array_1.len());
1504        test_zip_output_data_type(&scalar_1, &array_1.as_ref(), array_1.len());
1505
1506        // Array and null scalar (and vice versa)
1507        test_zip_output_data_type(&array_1.as_ref(), &null_scalar, array_1.len());
1508
1509        test_zip_output_data_type(&null_scalar, &array_1.as_ref(), array_1.len());
1510
1511        // Both arrays
1512        test_zip_output_data_type(&array_1.as_ref(), &array_2.as_ref(), array_1.len());
1513    }
1514
1515    fn test_zip_output_data_type(truthy: &dyn Datum, falsy: &dyn Datum, mask_length: usize) {
1516        let expected_data_type = truthy.get().0.data_type().clone();
1517        assert_eq!(&expected_data_type, falsy.get().0.data_type());
1518
1519        // Try different masks to test different paths
1520        let mask_all_true = BooleanArray::from(vec![true; mask_length]);
1521        let mask_all_false = BooleanArray::from(vec![false; mask_length]);
1522        let mask_some_true_and_false =
1523            BooleanArray::from((0..mask_length).map(|i| i % 2 == 0).collect::<Vec<bool>>());
1524
1525        for mask in [&mask_all_true, &mask_all_false, &mask_some_true_and_false] {
1526            let out = zip(mask, truthy, falsy).unwrap();
1527            assert_eq!(out.data_type(), &expected_data_type);
1528        }
1529    }
1530
1531    #[test]
1532    fn zip_scalar_fallback_impl() {
1533        let truthy_list_item_scalar = Some(vec![Some(1), None, Some(3)]);
1534        let truthy_list_array_scalar =
1535            Scalar::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1536                truthy_list_item_scalar.clone(),
1537            ]));
1538        let falsy_list_item_scalar = Some(vec![None, Some(2), Some(4)]);
1539        let falsy_list_array_scalar =
1540            Scalar::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1541                falsy_list_item_scalar.clone(),
1542            ]));
1543        let mask = BooleanArray::from(vec![true, false, true, false, false, true, false]);
1544        let out = zip(&mask, &truthy_list_array_scalar, &falsy_list_array_scalar).unwrap();
1545        let actual = out.as_list::<i32>();
1546
1547        let expected = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1548            truthy_list_item_scalar.clone(),
1549            falsy_list_item_scalar.clone(),
1550            truthy_list_item_scalar.clone(),
1551            falsy_list_item_scalar.clone(),
1552            falsy_list_item_scalar.clone(),
1553            truthy_list_item_scalar.clone(),
1554            falsy_list_item_scalar.clone(),
1555        ]);
1556        assert_eq!(actual, &expected);
1557    }
1558
1559    #[test]
1560    fn test_zip_kernel_scalar_strings_array_view() {
1561        let scalar_truthy = Scalar::new(StringViewArray::from(vec!["hello"]));
1562        let scalar_falsy = Scalar::new(StringViewArray::from(vec!["world"]));
1563
1564        let mask = BooleanArray::from(vec![true, false, true, false]);
1565        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1566        let actual = out.as_string_view();
1567        let expected = StringViewArray::from(vec![
1568            Some("hello"),
1569            Some("world"),
1570            Some("hello"),
1571            Some("world"),
1572        ]);
1573        assert_eq!(actual, &expected);
1574    }
1575
1576    #[test]
1577    fn test_zip_kernel_scalar_binary_array_view() {
1578        let scalar_truthy = Scalar::new(BinaryViewArray::from_iter_values(vec![b"hello"]));
1579        let scalar_falsy = Scalar::new(BinaryViewArray::from_iter_values(vec![b"world"]));
1580
1581        let mask = BooleanArray::from(vec![true, false]);
1582        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1583        let actual = out.as_byte_view();
1584        let expected = BinaryViewArray::from_iter_values(vec![b"hello", b"world"]);
1585        assert_eq!(actual, &expected);
1586    }
1587
1588    #[test]
1589    fn test_zip_kernel_scalar_strings_array_view_with_nulls() {
1590        let scalar_truthy = Scalar::new(StringViewArray::from_iter_values(["hello"]));
1591        let scalar_falsy = Scalar::new(StringViewArray::new_null(1));
1592
1593        let mask = BooleanArray::from(vec![true, true, false, false, true]);
1594        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1595        let actual = out.as_any().downcast_ref::<StringViewArray>().unwrap();
1596        let expected = StringViewArray::from_iter(vec![
1597            Some("hello"),
1598            Some("hello"),
1599            None,
1600            None,
1601            Some("hello"),
1602        ]);
1603        assert_eq!(actual, &expected);
1604    }
1605
1606    #[test]
1607    fn test_zip_kernel_scalar_strings_array_view_all_true_null() {
1608        let scalar_truthy = Scalar::new(StringViewArray::new_null(1));
1609        let scalar_falsy = Scalar::new(StringViewArray::new_null(1));
1610        let mask = BooleanArray::from(vec![true, true]);
1611        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1612        let actual = out.as_any().downcast_ref::<StringViewArray>().unwrap();
1613        let expected = StringViewArray::from_iter(vec![None::<String>, None]);
1614        assert_eq!(actual, &expected);
1615    }
1616
1617    #[test]
1618    fn test_zip_kernel_scalar_strings_array_view_all_false_null() {
1619        let scalar_truthy = Scalar::new(StringViewArray::new_null(1));
1620        let scalar_falsy = Scalar::new(StringViewArray::new_null(1));
1621        let mask = BooleanArray::from(vec![false, false]);
1622        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1623        let actual = out.as_any().downcast_ref::<StringViewArray>().unwrap();
1624        let expected = StringViewArray::from_iter(vec![None::<String>, None]);
1625        assert_eq!(actual, &expected);
1626    }
1627
1628    #[test]
1629    fn test_zip_kernel_scalar_string_array_view_all_true() {
1630        let scalar_truthy = Scalar::new(StringViewArray::from(vec!["hello"]));
1631        let scalar_falsy = Scalar::new(StringViewArray::from(vec!["world"]));
1632
1633        let mask = BooleanArray::from(vec![true, true]);
1634        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1635        let actual = out.as_string_view();
1636        let expected = StringViewArray::from(vec![Some("hello"), Some("hello")]);
1637        assert_eq!(actual, &expected);
1638    }
1639
1640    #[test]
1641    fn test_zip_kernel_scalar_string_array_view_all_false() {
1642        let scalar_truthy = Scalar::new(StringViewArray::from(vec!["hello"]));
1643        let scalar_falsy = Scalar::new(StringViewArray::from(vec!["world"]));
1644
1645        let mask = BooleanArray::from(vec![false, false]);
1646        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1647        let actual = out.as_string_view();
1648        let expected = StringViewArray::from(vec![Some("world"), Some("world")]);
1649        assert_eq!(actual, &expected);
1650    }
1651
1652    #[test]
1653    fn test_zip_kernel_scalar_strings_large_strings() {
1654        let scalar_truthy = Scalar::new(StringViewArray::from(vec!["longer than 12 bytes"]));
1655        let scalar_falsy = Scalar::new(StringViewArray::from(vec!["another longer than 12 bytes"]));
1656
1657        let mask = BooleanArray::from(vec![true, false]);
1658        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1659        let actual = out.as_string_view();
1660        let expected = StringViewArray::from(vec![
1661            Some("longer than 12 bytes"),
1662            Some("another longer than 12 bytes"),
1663        ]);
1664        assert_eq!(actual, &expected);
1665    }
1666
1667    #[test]
1668    fn test_zip_kernel_scalar_strings_array_view_large_short_strings() {
1669        let scalar_truthy = Scalar::new(StringViewArray::from(vec!["hello"]));
1670        let scalar_falsy = Scalar::new(StringViewArray::from(vec!["longer than 12 bytes"]));
1671
1672        let mask = BooleanArray::from(vec![true, false, true, false]);
1673        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1674        let actual = out.as_string_view();
1675        let expected = StringViewArray::from(vec![
1676            Some("hello"),
1677            Some("longer than 12 bytes"),
1678            Some("hello"),
1679            Some("longer than 12 bytes"),
1680        ]);
1681        assert_eq!(actual, &expected);
1682    }
1683    #[test]
1684    fn test_zip_kernel_scalar_strings_array_view_large_all_true() {
1685        let scalar_truthy = Scalar::new(StringViewArray::from(vec!["longer than 12 bytes"]));
1686        let scalar_falsy = Scalar::new(StringViewArray::from(vec!["another longer than 12 bytes"]));
1687
1688        let mask = BooleanArray::from(vec![true, true]);
1689        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1690        let actual = out.as_string_view();
1691        let expected = StringViewArray::from(vec![
1692            Some("longer than 12 bytes"),
1693            Some("longer than 12 bytes"),
1694        ]);
1695        assert_eq!(actual, &expected);
1696    }
1697
1698    #[test]
1699    fn test_zip_kernel_scalar_strings_array_view_large_all_false() {
1700        let scalar_truthy = Scalar::new(StringViewArray::from(vec!["longer than 12 bytes"]));
1701        let scalar_falsy = Scalar::new(StringViewArray::from(vec!["another longer than 12 bytes"]));
1702
1703        let mask = BooleanArray::from(vec![false, false]);
1704        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1705        let actual = out.as_string_view();
1706        let expected = StringViewArray::from(vec![
1707            Some("another longer than 12 bytes"),
1708            Some("another longer than 12 bytes"),
1709        ]);
1710        assert_eq!(actual, &expected);
1711    }
1712}
arrow_select/zip.rs

arrow_select/
zip.rs
