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 zip_impl(
149    mask: &BooleanArray,
150    truthy: &ArrayData,
151    truthy_is_scalar: bool,
152    falsy: &ArrayData,
153    falsy_is_scalar: bool,
154) -> Result<ArrayRef, ArrowError> {
155    let mut mutable = MutableArrayData::new(vec![truthy, falsy], false, truthy.len());
156
157    // the SlicesIterator slices only the true values. So the gaps left by this iterator we need to
158    // fill with falsy values
159
160    // keep track of how much is filled
161    let mut filled = 0;
162
163    let mask_buffer = maybe_prep_null_mask_filter(mask);
164    for (start, end) in SlicesIterator::from(&mask_buffer) {
165        // the gap needs to be filled with falsy values
166        if start > filled {
167            if falsy_is_scalar {
168                for _ in filled..start {
169                    // Copy the first item from the 'falsy' array into the output buffer.
170                    mutable.try_extend(1, 0, 1)?;
171                }
172            } else {
173                mutable.try_extend(1, filled, start)?;
174            }
175        }
176        // fill with truthy values
177        if truthy_is_scalar {
178            for _ in start..end {
179                // Copy the first item from the 'truthy' array into the output buffer.
180                mutable.try_extend(0, 0, 1)?;
181            }
182        } else {
183            mutable.try_extend(0, start, end)?;
184        }
185        filled = end;
186    }
187    // the remaining part is falsy
188    if filled < mask.len() {
189        if falsy_is_scalar {
190            for _ in filled..mask.len() {
191                // Copy the first item from the 'falsy' array into the output buffer.
192                mutable.try_extend(1, 0, 1)?;
193            }
194        } else {
195            mutable.try_extend(1, filled, mask.len())?;
196        }
197    }
198
199    let data = mutable.freeze();
200    Ok(make_array(data))
201}
202
203/// Zipper for 2 scalars
204///
205/// Useful for using in `IF <expr> THEN <scalar> ELSE <scalar> END` expressions
206///
207/// # Example
208/// ```
209/// # use std::sync::Arc;
210/// # use arrow_array::{ArrayRef, BooleanArray, Int32Array, Scalar, cast::AsArray, types::Int32Type};
211///
212/// # use arrow_select::zip::ScalarZipper;
213/// let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
214/// let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
215/// let zipper = ScalarZipper::try_new(&scalar_truthy, &scalar_falsy).unwrap();
216///
217/// // Later when we have a boolean mask
218/// let mask = BooleanArray::from(vec![true, false, true, false, true]);
219/// let result = zipper.zip(&mask).unwrap();
220/// let actual = result.as_primitive::<Int32Type>();
221/// let expected = Int32Array::from(vec![Some(42), Some(123), Some(42), Some(123), Some(42)]);
222/// ```
223///
224#[derive(Debug, Clone)]
225pub struct ScalarZipper {
226    zip_impl: Arc<dyn ZipImpl>,
227}
228
229impl ScalarZipper {
230    /// Try to create a new ScalarZipper from two scalar Datum
231    ///
232    /// # Errors
233    /// returns error if:
234    /// - the two Datum have different data types
235    /// - either Datum is not a scalar (or has more than 1 element)
236    ///
237    pub fn try_new(truthy: &dyn Datum, falsy: &dyn Datum) -> Result<Self, ArrowError> {
238        let (truthy, truthy_is_scalar) = truthy.get();
239        let (falsy, falsy_is_scalar) = falsy.get();
240
241        if truthy.data_type() != falsy.data_type() {
242            return Err(ArrowError::InvalidArgumentError(
243                "arguments need to have the same data type".into(),
244            ));
245        }
246
247        if !truthy_is_scalar {
248            return Err(ArrowError::InvalidArgumentError(
249                "only scalar arrays are supported".into(),
250            ));
251        }
252
253        if !falsy_is_scalar {
254            return Err(ArrowError::InvalidArgumentError(
255                "only scalar arrays are supported".into(),
256            ));
257        }
258
259        if truthy.len() != 1 {
260            return Err(ArrowError::InvalidArgumentError(
261                "scalar arrays must have 1 element".into(),
262            ));
263        }
264        if falsy.len() != 1 {
265            return Err(ArrowError::InvalidArgumentError(
266                "scalar arrays must have 1 element".into(),
267            ));
268        }
269
270        macro_rules! primitive_size_helper {
271            ($t:ty) => {
272                Arc::new(PrimitiveScalarImpl::<$t>::new(truthy, falsy)) as Arc<dyn ZipImpl>
273            };
274        }
275
276        let zip_impl = downcast_primitive! {
277            truthy.data_type() => (primitive_size_helper),
278            DataType::Utf8 => {
279                Arc::new(BytesScalarImpl::<Utf8Type>::new(truthy, falsy)) as Arc<dyn ZipImpl>
280            },
281            DataType::LargeUtf8 => {
282                Arc::new(BytesScalarImpl::<LargeUtf8Type>::new(truthy, falsy)) as Arc<dyn ZipImpl>
283            },
284            DataType::Binary => {
285                Arc::new(BytesScalarImpl::<BinaryType>::new(truthy, falsy)) as Arc<dyn ZipImpl>
286            },
287            DataType::LargeBinary => {
288                Arc::new(BytesScalarImpl::<LargeBinaryType>::new(truthy, falsy)) as Arc<dyn ZipImpl>
289            },
290            DataType::Utf8View => {
291                Arc::new(ByteViewScalarImpl::<StringViewType>::new(truthy, falsy)) as Arc<dyn ZipImpl>
292            },
293            DataType::BinaryView => {
294                Arc::new(ByteViewScalarImpl::<BinaryViewType>::new(truthy, falsy)) as Arc<dyn ZipImpl>
295            },
296            _ => {
297                Arc::new(FallbackImpl::new(truthy, falsy)) as Arc<dyn ZipImpl>
298            },
299        };
300
301        Ok(Self { zip_impl })
302    }
303
304    /// Creating output array based on input boolean array and the two scalar values the zipper was created with
305    /// See struct level documentation for examples.
306    pub fn zip(&self, mask: &BooleanArray) -> Result<ArrayRef, ArrowError> {
307        self.zip_impl.create_output(mask)
308    }
309}
310
311/// Impl for creating output array based on a mask
312trait ZipImpl: Debug + Send + Sync {
313    /// Creating output array based on input boolean array
314    fn create_output(&self, input: &BooleanArray) -> Result<ArrayRef, ArrowError>;
315}
316
317#[derive(Debug, PartialEq)]
318struct FallbackImpl {
319    truthy: ArrayData,
320    falsy: ArrayData,
321}
322
323impl FallbackImpl {
324    fn new(left: &dyn Array, right: &dyn Array) -> Self {
325        Self {
326            truthy: left.to_data(),
327            falsy: right.to_data(),
328        }
329    }
330}
331
332impl ZipImpl for FallbackImpl {
333    fn create_output(&self, predicate: &BooleanArray) -> Result<ArrayRef, ArrowError> {
334        zip_impl(predicate, &self.truthy, true, &self.falsy, true)
335    }
336}
337
338struct PrimitiveScalarImpl<T: ArrowPrimitiveType> {
339    data_type: DataType,
340    truthy: Option<T::Native>,
341    falsy: Option<T::Native>,
342}
343
344impl<T: ArrowPrimitiveType> Debug for PrimitiveScalarImpl<T> {
345    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
346        f.debug_struct("PrimitiveScalarImpl")
347            .field("data_type", &self.data_type)
348            .field("truthy", &self.truthy)
349            .field("falsy", &self.falsy)
350            .finish()
351    }
352}
353
354impl<T: ArrowPrimitiveType> PrimitiveScalarImpl<T> {
355    fn new(truthy: &dyn Array, falsy: &dyn Array) -> Self {
356        Self {
357            data_type: truthy.data_type().clone(),
358            truthy: Self::get_value_from_scalar(truthy),
359            falsy: Self::get_value_from_scalar(falsy),
360        }
361    }
362
363    fn get_value_from_scalar(scalar: &dyn Array) -> Option<T::Native> {
364        if scalar.is_null(0) {
365            None
366        } else {
367            let value = scalar.as_primitive::<T>().value(0);
368
369            Some(value)
370        }
371    }
372
373    /// return an output array that has
374    /// `value` in all locations where predicate is true
375    /// `null` otherwise
376    fn get_scalar_and_null_buffer_for_single_non_nullable(
377        predicate: BooleanBuffer,
378        value: T::Native,
379    ) -> (Vec<T::Native>, Option<NullBuffer>) {
380        let result_len = predicate.len();
381        let nulls = NullBuffer::new(predicate);
382        let scalars = vec![value; result_len];
383
384        (scalars, Some(nulls))
385    }
386}
387
388impl<T: ArrowPrimitiveType> ZipImpl for PrimitiveScalarImpl<T> {
389    fn create_output(&self, predicate: &BooleanArray) -> Result<ArrayRef, ArrowError> {
390        let result_len = predicate.len();
391        // Nulls are treated as false
392        let predicate = maybe_prep_null_mask_filter(predicate);
393
394        let (scalars, nulls): (Vec<T::Native>, Option<NullBuffer>) = match (self.truthy, self.falsy)
395        {
396            (Some(truthy_val), Some(falsy_val)) => {
397                let scalars: Vec<T::Native> = predicate
398                    .iter()
399                    .map(|b| if b { truthy_val } else { falsy_val })
400                    .collect();
401
402                (scalars, None)
403            }
404            (Some(truthy_val), None) => {
405                // If a value is true we need the TRUTHY and the null buffer will have 1 (meaning not null)
406                // If a value is false we need the FALSY and the null buffer will have 0 (meaning null)
407
408                Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, truthy_val)
409            }
410            (None, Some(falsy_val)) => {
411                // Flipping the boolean buffer as we want the opposite of the TRUE case
412                //
413                // if the condition is true we want null so we need to NOT the value so we get 0 (meaning null)
414                // if the condition is false we want the FALSY value so we need to NOT the value so we get 1 (meaning not null)
415                let predicate = predicate.not();
416
417                Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, falsy_val)
418            }
419            (None, None) => {
420                // All values are null
421                let nulls = NullBuffer::new_null(result_len);
422                let scalars = vec![T::default_value(); result_len];
423
424                (scalars, Some(nulls))
425            }
426        };
427
428        let scalars = ScalarBuffer::<T::Native>::from(scalars);
429        let output = PrimitiveArray::<T>::try_new(scalars, nulls)?;
430
431        // Keep decimal precisions, scales or timestamps timezones
432        let output = output.with_data_type(self.data_type.clone());
433
434        Ok(Arc::new(output))
435    }
436}
437
438#[derive(PartialEq, Hash)]
439struct BytesScalarImpl<T: ByteArrayType> {
440    truthy: Option<Vec<u8>>,
441    falsy: Option<Vec<u8>>,
442    phantom: PhantomData<T>,
443}
444
445impl<T: ByteArrayType> Debug for BytesScalarImpl<T> {
446    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
447        f.debug_struct("BytesScalarImpl")
448            .field("truthy", &self.truthy)
449            .field("falsy", &self.falsy)
450            .finish()
451    }
452}
453
454impl<T: ByteArrayType> BytesScalarImpl<T> {
455    fn new(truthy_value: &dyn Array, falsy_value: &dyn Array) -> Self {
456        Self {
457            truthy: Self::get_value_from_scalar(truthy_value),
458            falsy: Self::get_value_from_scalar(falsy_value),
459            phantom: PhantomData,
460        }
461    }
462
463    fn get_value_from_scalar(scalar: &dyn Array) -> Option<Vec<u8>> {
464        if scalar.is_null(0) {
465            None
466        } else {
467            let bytes: &[u8] = scalar.as_bytes::<T>().value(0).as_ref();
468
469            Some(bytes.to_vec())
470        }
471    }
472
473    /// return an output array that has
474    /// `value` in all locations where predicate is true
475    /// `null` otherwise
476    #[expect(clippy::type_complexity)]
477    fn get_scalar_and_null_buffer_for_single_non_nullable(
478        predicate: BooleanBuffer,
479        value: &[u8],
480    ) -> Result<(Buffer, OffsetBuffer<T::Offset>, Option<NullBuffer>), ArrowError> {
481        let value_length = value.len();
482
483        let number_of_true = predicate.count_set_bits();
484
485        // Fast path for all nulls
486        if number_of_true == 0 {
487            // All values are null
488            let nulls = NullBuffer::new_null(predicate.len());
489
490            return Ok((
491                // Empty bytes
492                Buffer::from(&[]),
493                // All nulls so all lengths are 0
494                OffsetBuffer::<T::Offset>::new_zeroed(predicate.len()),
495                Some(nulls),
496            ));
497        }
498
499        let offsets = OffsetBuffer::<T::Offset>::from_lengths(
500            predicate.iter().map(|b| if b { value_length } else { 0 }),
501        );
502
503        let mut bytes = MutableBuffer::with_capacity(0);
504        bytes
505            .try_repeat_slice_n_times(value, number_of_true)
506            .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
507
508        let bytes = Buffer::from(bytes);
509
510        // If a value is true we need the TRUTHY and the null buffer will have 1 (meaning not null)
511        // If a value is false we need the FALSY and the null buffer will have 0 (meaning null)
512        let nulls = NullBuffer::new(predicate);
513
514        Ok((bytes, offsets, Some(nulls)))
515    }
516
517    /// Create a [`Buffer`] where `value` slice is repeated `number_of_values` times
518    /// and [`OffsetBuffer`] where there are `number_of_values` lengths, and all equals to `value` length
519    fn get_bytes_and_offset_for_all_same_value(
520        number_of_values: usize,
521        value: &[u8],
522    ) -> Result<(Buffer, OffsetBuffer<T::Offset>), ArrowError> {
523        let value_length = value.len();
524
525        let offsets =
526            OffsetBuffer::<T::Offset>::from_repeated_length(value_length, number_of_values);
527
528        let mut bytes = MutableBuffer::with_capacity(0);
529        bytes
530            .try_repeat_slice_n_times(value, number_of_values)
531            .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
532        let bytes = Buffer::from(bytes);
533
534        Ok((bytes, offsets))
535    }
536
537    fn create_output_on_non_nulls(
538        predicate: &BooleanBuffer,
539        truthy_val: &[u8],
540        falsy_val: &[u8],
541    ) -> Result<(Buffer, OffsetBuffer<<T as ByteArrayType>::Offset>), ArrowError> {
542        let true_count = predicate.count_set_bits();
543
544        match true_count {
545            0 => {
546                // All values are falsy
547                return Self::get_bytes_and_offset_for_all_same_value(predicate.len(), falsy_val);
548            }
549            n if n == predicate.len() => {
550                // All values are truthy
551                return Self::get_bytes_and_offset_for_all_same_value(predicate.len(), truthy_val);
552            }
553
554            _ => {
555                // Fallback
556            }
557        }
558
559        let total_number_of_bytes =
560            true_count * truthy_val.len() + (predicate.len() - true_count) * falsy_val.len();
561        let mut mutable = MutableBuffer::with_capacity(total_number_of_bytes);
562        let mut offset_buffer_builder = OffsetBufferBuilder::<T::Offset>::new(predicate.len());
563
564        // keep track of how much is filled
565        let mut filled = 0;
566
567        let truthy_len = truthy_val.len();
568        let falsy_len = falsy_val.len();
569
570        SlicesIterator::from(predicate).try_for_each(|(start, end)| -> Result<(), ArrowError> {
571            // the gap needs to be filled with falsy values
572            if start > filled {
573                let false_repeat_count = start - filled;
574                // Push false value `repeat_count` times
575                mutable
576                    .try_repeat_slice_n_times(falsy_val, false_repeat_count)
577                    .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
578
579                for _ in 0..false_repeat_count {
580                    offset_buffer_builder.push_length(falsy_len)
581                }
582            }
583
584            let true_repeat_count = end - start;
585            // fill with truthy values
586            mutable
587                .try_repeat_slice_n_times(truthy_val, true_repeat_count)
588                .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
589
590            for _ in 0..true_repeat_count {
591                offset_buffer_builder.push_length(truthy_len)
592            }
593            filled = end;
594            Ok(())
595        })?;
596        // the remaining part is falsy
597        if filled < predicate.len() {
598            let false_repeat_count = predicate.len() - filled;
599            // Copy the first item from the 'falsy' array into the output buffer.
600            mutable
601                .try_repeat_slice_n_times(falsy_val, false_repeat_count)
602                .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
603
604            for _ in 0..false_repeat_count {
605                offset_buffer_builder.push_length(falsy_len)
606            }
607        }
608
609        Ok((mutable.into(), offset_buffer_builder.finish()))
610    }
611}
612
613impl<T: ByteArrayType> ZipImpl for BytesScalarImpl<T> {
614    fn create_output(&self, predicate: &BooleanArray) -> Result<ArrayRef, ArrowError> {
615        let result_len = predicate.len();
616        // Nulls are treated as false
617        let predicate = maybe_prep_null_mask_filter(predicate);
618
619        let (bytes, offsets, nulls): (Buffer, OffsetBuffer<T::Offset>, Option<NullBuffer>) =
620            match (self.truthy.as_deref(), self.falsy.as_deref()) {
621                (Some(truthy_val), Some(falsy_val)) => {
622                    let (bytes, offsets) =
623                        Self::create_output_on_non_nulls(&predicate, truthy_val, falsy_val)?;
624
625                    (bytes, offsets, None)
626                }
627                (Some(truthy_val), None) => {
628                    Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, truthy_val)?
629                }
630                (None, Some(falsy_val)) => {
631                    // Flipping the boolean buffer as we want the opposite of the TRUE case
632                    //
633                    // if the condition is true we want null so we need to NOT the value so we get 0 (meaning null)
634                    // if the condition is false we want the FALSE value so we need to NOT the value so we get 1 (meaning not null)
635                    let predicate = predicate.not();
636                    Self::get_scalar_and_null_buffer_for_single_non_nullable(predicate, falsy_val)?
637                }
638                (None, None) => {
639                    // All values are null
640                    let nulls = NullBuffer::new_null(result_len);
641
642                    (
643                        // Empty bytes
644                        Buffer::from(&[]),
645                        // All nulls so all lengths are 0
646                        OffsetBuffer::<T::Offset>::new_zeroed(predicate.len()),
647                        Some(nulls),
648                    )
649                }
650            };
651
652        let output = unsafe {
653            // Safety: the values are based on valid inputs
654            // and `try_new` is expensive for strings as it validate that the input is valid utf8
655            GenericByteArray::<T>::new_unchecked(offsets, bytes, nulls)
656        };
657
658        Ok(Arc::new(output))
659    }
660}
661
662fn maybe_prep_null_mask_filter(predicate: &BooleanArray) -> BooleanBuffer {
663    // Nulls are treated as false
664    if predicate.null_count() == 0 {
665        predicate.values().clone()
666    } else {
667        let cleaned = prep_null_mask_filter(predicate);
668        let (boolean_buffer, _) = cleaned.into_parts();
669        boolean_buffer
670    }
671}
672
673struct ByteViewScalarImpl<T: ByteViewType> {
674    truthy_view: Option<u128>,
675    truthy_buffers: Arc<[Buffer]>,
676    falsy_view: Option<u128>,
677    falsy_buffers: Arc<[Buffer]>,
678    phantom: PhantomData<T>,
679}
680
681static EMPTY_ARC: OnceLock<Arc<[Buffer]>> = OnceLock::new();
682fn empty_arc_buffers() -> Arc<[Buffer]> {
683    Arc::clone(EMPTY_ARC.get_or_init(|| Arc::new([])))
684}
685
686impl<T: ByteViewType> ByteViewScalarImpl<T> {
687    fn new(truthy: &dyn Array, falsy: &dyn Array) -> Self {
688        let (truthy_view, truthy_buffers) = Self::get_value_from_scalar(truthy);
689        let (falsy_view, falsy_buffers) = Self::get_value_from_scalar(falsy);
690        Self {
691            truthy_view,
692            truthy_buffers,
693            falsy_view,
694            falsy_buffers,
695            phantom: PhantomData,
696        }
697    }
698
699    fn get_value_from_scalar(scalar: &dyn Array) -> (Option<u128>, Arc<[Buffer]>) {
700        if scalar.is_null(0) {
701            (None, empty_arc_buffers())
702        } else {
703            let (views, buffers, _) = scalar.as_byte_view::<T>().clone().into_parts();
704            (views.first().copied(), buffers)
705        }
706    }
707
708    fn get_views_for_single_non_nullable(
709        predicate: BooleanBuffer,
710        value: u128,
711        buffers: Arc<[Buffer]>,
712    ) -> (ScalarBuffer<u128>, Arc<[Buffer]>, Option<NullBuffer>) {
713        let number_of_true = predicate.count_set_bits();
714        let number_of_values = predicate.len();
715
716        // Fast path for all nulls
717        if number_of_true == 0 {
718            // All values are null
719            return (
720                vec![0; number_of_values].into(),
721                empty_arc_buffers(),
722                Some(NullBuffer::new_null(number_of_values)),
723            );
724        }
725        let bytes = vec![value; number_of_values];
726
727        // If value is true and we want to handle the TRUTHY case, the null buffer will have 1 (meaning not null)
728        // If value is false and we want to handle the FALSY case, the null buffer will have 0 (meaning null)
729        let nulls = NullBuffer::new(predicate);
730        (bytes.into(), buffers, Some(nulls))
731    }
732
733    #[expect(clippy::type_complexity)]
734    fn get_views_for_non_nullable(
735        predicate: BooleanBuffer,
736        result_len: usize,
737        truthy_view: u128,
738        truthy_buffers: Arc<[Buffer]>,
739        falsy_view: u128,
740        falsy_buffers: Arc<[Buffer]>,
741    ) -> Result<(ScalarBuffer<u128>, Arc<[Buffer]>, Option<NullBuffer>), ArrowError> {
742        let true_count = predicate.count_set_bits();
743        match true_count {
744            0 => {
745                // all values are falsy
746                Ok((vec![falsy_view; result_len].into(), falsy_buffers, None))
747            }
748            n if n == predicate.len() => {
749                // all values are truthy
750                Ok((vec![truthy_view; result_len].into(), truthy_buffers, None))
751            }
752            _ => {
753                let true_count = predicate.count_set_bits();
754                let mut buffers: Vec<Buffer> = truthy_buffers.to_vec();
755
756                // If the falsy buffers are empty, we can use the falsy view as it is, because the value
757                // is completely inlined. Otherwise, we have non-inlined values in the buffer, and we need
758                // to recalculate the falsy view
759                let view_falsy = if falsy_buffers.is_empty() {
760                    falsy_view
761                } else {
762                    let byte_view_falsy = ByteView::from(falsy_view);
763                    let new_index_falsy_buffers =
764                        buffers.len() as u32 + byte_view_falsy.buffer_index;
765                    buffers.extend(falsy_buffers.iter().cloned());
766                    let byte_view_falsy =
767                        byte_view_falsy.with_buffer_index(new_index_falsy_buffers);
768                    byte_view_falsy.as_u128()
769                };
770
771                let total_number_of_bytes = true_count * 16 + (predicate.len() - true_count) * 16;
772                let mut mutable = MutableBuffer::new(total_number_of_bytes);
773                let mut filled = 0;
774
775                SlicesIterator::from(&predicate).try_for_each(
776                    |(start, end)| -> Result<(), ArrowError> {
777                        if start > filled {
778                            let false_repeat_count = start - filled;
779                            mutable
780                                .try_repeat_slice_n_times(
781                                    view_falsy.to_byte_slice(),
782                                    false_repeat_count,
783                                )
784                                .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
785                        }
786                        let true_repeat_count = end - start;
787                        mutable
788                            .try_repeat_slice_n_times(
789                                truthy_view.to_byte_slice(),
790                                true_repeat_count,
791                            )
792                            .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
793                        filled = end;
794                        Ok(())
795                    },
796                )?;
797
798                if filled < predicate.len() {
799                    let false_repeat_count = predicate.len() - filled;
800                    mutable
801                        .try_repeat_slice_n_times(view_falsy.to_byte_slice(), false_repeat_count)
802                        .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
803                }
804
805                let bytes = Buffer::from(mutable);
806                Ok((bytes.into(), buffers.into(), None))
807            }
808        }
809    }
810}
811
812impl<T: ByteViewType> Debug for ByteViewScalarImpl<T> {
813    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
814        f.debug_struct("ByteViewScalarImpl")
815            .field("truthy", &self.truthy_view)
816            .field("falsy", &self.falsy_view)
817            .finish()
818    }
819}
820
821impl<T: ByteViewType> ZipImpl for ByteViewScalarImpl<T> {
822    fn create_output(&self, predicate: &BooleanArray) -> Result<ArrayRef, ArrowError> {
823        let result_len = predicate.len();
824        // Nulls are treated as false
825        let predicate = maybe_prep_null_mask_filter(predicate);
826
827        let (views, buffers, nulls) = match (self.truthy_view, self.falsy_view) {
828            (Some(truthy), Some(falsy)) => Self::get_views_for_non_nullable(
829                predicate,
830                result_len,
831                truthy,
832                Arc::clone(&self.truthy_buffers),
833                falsy,
834                Arc::clone(&self.falsy_buffers),
835            )?,
836            (Some(truthy), None) => Self::get_views_for_single_non_nullable(
837                predicate,
838                truthy,
839                Arc::clone(&self.truthy_buffers),
840            ),
841            (None, Some(falsy)) => {
842                let predicate = predicate.not();
843                Self::get_views_for_single_non_nullable(
844                    predicate,
845                    falsy,
846                    Arc::clone(&self.falsy_buffers),
847                )
848            }
849            (None, None) => {
850                // All values are null
851                (
852                    vec![0; result_len].into(),
853                    empty_arc_buffers(),
854                    Some(NullBuffer::new_null(result_len)),
855                )
856            }
857        };
858
859        let result = unsafe { GenericByteViewArray::<T>::new_unchecked(views, buffers, nulls) };
860        Ok(Arc::new(result))
861    }
862}
863
864#[cfg(test)]
865mod test {
866    use super::*;
867    use arrow_array::types::Int32Type;
868
869    #[test]
870    fn test_zip_kernel_one() {
871        let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
872        let b = Int32Array::from(vec![None, Some(3), Some(6), Some(7), Some(3)]);
873        let mask = BooleanArray::from(vec![true, true, false, false, true]);
874        let out = zip(&mask, &a, &b).unwrap();
875        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
876        let expected = Int32Array::from(vec![Some(5), None, Some(6), Some(7), Some(1)]);
877        assert_eq!(actual, &expected);
878    }
879
880    #[test]
881    fn test_zip_kernel_two() {
882        let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
883        let b = Int32Array::from(vec![None, Some(3), Some(6), Some(7), Some(3)]);
884        let mask = BooleanArray::from(vec![false, false, true, true, false]);
885        let out = zip(&mask, &a, &b).unwrap();
886        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
887        let expected = Int32Array::from(vec![None, Some(3), Some(7), None, Some(3)]);
888        assert_eq!(actual, &expected);
889    }
890
891    #[test]
892    fn test_zip_kernel_scalar_falsy_1() {
893        let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
894
895        let fallback = Scalar::new(Int32Array::from_value(42, 1));
896
897        let mask = BooleanArray::from(vec![true, true, false, false, true]);
898        let out = zip(&mask, &a, &fallback).unwrap();
899        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
900        let expected = Int32Array::from(vec![Some(5), None, Some(42), Some(42), Some(1)]);
901        assert_eq!(actual, &expected);
902    }
903
904    #[test]
905    fn test_zip_kernel_scalar_falsy_2() {
906        let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
907
908        let fallback = Scalar::new(Int32Array::from_value(42, 1));
909
910        let mask = BooleanArray::from(vec![false, false, true, true, false]);
911        let out = zip(&mask, &a, &fallback).unwrap();
912        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
913        let expected = Int32Array::from(vec![Some(42), Some(42), Some(7), None, Some(42)]);
914        assert_eq!(actual, &expected);
915    }
916
917    #[test]
918    fn test_zip_kernel_scalar_truthy_1() {
919        let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
920
921        let fallback = Scalar::new(Int32Array::from_value(42, 1));
922
923        let mask = BooleanArray::from(vec![true, true, false, false, true]);
924        let out = zip(&mask, &fallback, &a).unwrap();
925        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
926        let expected = Int32Array::from(vec![Some(42), Some(42), Some(7), None, Some(42)]);
927        assert_eq!(actual, &expected);
928    }
929
930    #[test]
931    fn test_zip_kernel_scalar_truthy_2() {
932        let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
933
934        let fallback = Scalar::new(Int32Array::from_value(42, 1));
935
936        let mask = BooleanArray::from(vec![false, false, true, true, false]);
937        let out = zip(&mask, &fallback, &a).unwrap();
938        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
939        let expected = Int32Array::from(vec![Some(5), None, Some(42), Some(42), Some(1)]);
940        assert_eq!(actual, &expected);
941    }
942
943    #[test]
944    fn test_zip_kernel_scalar_both_mask_ends_with_true() {
945        let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
946        let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
947
948        let mask = BooleanArray::from(vec![true, true, false, false, true]);
949        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
950        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
951        let expected = Int32Array::from(vec![Some(42), Some(42), Some(123), Some(123), Some(42)]);
952        assert_eq!(actual, &expected);
953    }
954
955    #[test]
956    fn test_zip_kernel_scalar_both_mask_ends_with_false() {
957        let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
958        let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
959
960        let mask = BooleanArray::from(vec![true, true, false, true, false, false]);
961        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
962        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
963        let expected = Int32Array::from(vec![
964            Some(42),
965            Some(42),
966            Some(123),
967            Some(42),
968            Some(123),
969            Some(123),
970        ]);
971        assert_eq!(actual, &expected);
972    }
973
974    #[test]
975    fn test_zip_kernel_primitive_scalar_none_1() {
976        let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
977        let scalar_falsy = Scalar::new(Int32Array::new_null(1));
978
979        let mask = BooleanArray::from(vec![true, true, false, false, true]);
980        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
981        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
982        let expected = Int32Array::from(vec![Some(42), Some(42), None, None, Some(42)]);
983        assert_eq!(actual, &expected);
984    }
985
986    #[test]
987    fn test_zip_kernel_primitive_scalar_none_2() {
988        let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
989        let scalar_falsy = Scalar::new(Int32Array::new_null(1));
990
991        let mask = BooleanArray::from(vec![false, false, true, true, false]);
992        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
993        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
994        let expected = Int32Array::from(vec![None, None, Some(42), Some(42), None]);
995        assert_eq!(actual, &expected);
996    }
997
998    #[test]
999    fn test_zip_kernel_primitive_scalar_both_null() {
1000        let scalar_truthy = Scalar::new(Int32Array::new_null(1));
1001        let scalar_falsy = Scalar::new(Int32Array::new_null(1));
1002
1003        let mask = BooleanArray::from(vec![false, false, true, true, false]);
1004        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1005        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1006        let expected = Int32Array::from(vec![None, None, None, None, None]);
1007        assert_eq!(actual, &expected);
1008    }
1009
1010    #[test]
1011    fn test_zip_primitive_array_with_nulls_is_mask_should_be_treated_as_false() {
1012        let truthy = Int32Array::from_iter_values(vec![1, 2, 3, 4, 5, 6]);
1013        let falsy = Int32Array::from_iter_values(vec![7, 8, 9, 10, 11, 12]);
1014
1015        let mask = {
1016            let booleans = BooleanBuffer::from(vec![true, true, false, true, false, false]);
1017            let nulls = NullBuffer::from(vec![
1018                true, true, true,
1019                false, // null treated as false even though in the original mask it was true
1020                true, true,
1021            ]);
1022            BooleanArray::new(booleans, Some(nulls))
1023        };
1024        let out = zip(&mask, &truthy, &falsy).unwrap();
1025        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1026        let expected = Int32Array::from(vec![
1027            Some(1),
1028            Some(2),
1029            Some(9),
1030            Some(10), // true in mask but null
1031            Some(11),
1032            Some(12),
1033        ]);
1034        assert_eq!(actual, &expected);
1035    }
1036
1037    #[test]
1038    fn test_zip_kernel_primitive_scalar_with_boolean_array_mask_with_nulls_should_be_treated_as_false()
1039     {
1040        let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1041        let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
1042
1043        let mask = {
1044            let booleans = BooleanBuffer::from(vec![true, true, false, true, false, false]);
1045            let nulls = NullBuffer::from(vec![
1046                true, true, true,
1047                false, // null treated as false even though in the original mask it was true
1048                true, true,
1049            ]);
1050            BooleanArray::new(booleans, Some(nulls))
1051        };
1052        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1053        let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
1054        let expected = Int32Array::from(vec![
1055            Some(42),
1056            Some(42),
1057            Some(123),
1058            Some(123), // true in mask but null
1059            Some(123),
1060            Some(123),
1061        ]);
1062        assert_eq!(actual, &expected);
1063    }
1064
1065    #[test]
1066    fn test_zip_string_array_with_nulls_is_mask_should_be_treated_as_false() {
1067        let truthy = StringArray::from_iter_values(vec!["1", "2", "3", "4", "5", "6"]);
1068        let falsy = StringArray::from_iter_values(vec!["7", "8", "9", "10", "11", "12"]);
1069
1070        let mask = {
1071            let booleans = BooleanBuffer::from(vec![true, true, false, true, false, false]);
1072            let nulls = NullBuffer::from(vec![
1073                true, true, true,
1074                false, // null treated as false even though in the original mask it was true
1075                true, true,
1076            ]);
1077            BooleanArray::new(booleans, Some(nulls))
1078        };
1079        let out = zip(&mask, &truthy, &falsy).unwrap();
1080        let actual = out.as_string::<i32>();
1081        let expected = StringArray::from_iter_values(vec![
1082            "1", "2", "9", "10", // true in mask but null
1083            "11", "12",
1084        ]);
1085        assert_eq!(actual, &expected);
1086    }
1087
1088    #[test]
1089    fn test_zip_kernel_large_string_scalar_with_boolean_array_mask_with_nulls_should_be_treated_as_false()
1090     {
1091        let scalar_truthy = Scalar::new(LargeStringArray::from_iter_values(["test"]));
1092        let scalar_falsy = Scalar::new(LargeStringArray::from_iter_values(["something else"]));
1093
1094        let mask = {
1095            let booleans = BooleanBuffer::from(vec![true, true, false, true, false, false]);
1096            let nulls = NullBuffer::from(vec![
1097                true, true, true,
1098                false, // null treated as false even though in the original mask it was true
1099                true, true,
1100            ]);
1101            BooleanArray::new(booleans, Some(nulls))
1102        };
1103        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1104        let actual = out.as_any().downcast_ref::<LargeStringArray>().unwrap();
1105        let expected = LargeStringArray::from_iter(vec![
1106            Some("test"),
1107            Some("test"),
1108            Some("something else"),
1109            Some("something else"), // true in mask but null
1110            Some("something else"),
1111            Some("something else"),
1112        ]);
1113        assert_eq!(actual, &expected);
1114    }
1115
1116    #[test]
1117    fn test_zip_kernel_bytes_scalar_none_1() {
1118        let scalar_truthy = Scalar::new(StringArray::from_iter_values(["hello"]));
1119        let scalar_falsy = Scalar::new(StringArray::new_null(1));
1120
1121        let mask = BooleanArray::from(vec![true, true, false, false, true]);
1122        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1123        let actual = out.as_any().downcast_ref::<StringArray>().unwrap();
1124        let expected = StringArray::from_iter(vec![
1125            Some("hello"),
1126            Some("hello"),
1127            None,
1128            None,
1129            Some("hello"),
1130        ]);
1131        assert_eq!(actual, &expected);
1132    }
1133
1134    #[test]
1135    fn test_zip_kernel_bytes_scalar_none_2() {
1136        let scalar_truthy = Scalar::new(StringArray::new_null(1));
1137        let scalar_falsy = Scalar::new(StringArray::from_iter_values(["hello"]));
1138
1139        let mask = BooleanArray::from(vec![true, true, false, false, true]);
1140        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1141        let actual = out.as_any().downcast_ref::<StringArray>().unwrap();
1142        let expected = StringArray::from_iter(vec![None, None, Some("hello"), Some("hello"), None]);
1143        assert_eq!(actual, &expected);
1144    }
1145
1146    #[test]
1147    fn test_zip_kernel_bytes_scalar_both() {
1148        let scalar_truthy = Scalar::new(StringArray::from_iter_values(["test"]));
1149        let scalar_falsy = Scalar::new(StringArray::from_iter_values(["something else"]));
1150
1151        // mask ends with false
1152        let mask = BooleanArray::from(vec![true, true, false, true, false, false]);
1153        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1154        let actual = out.as_any().downcast_ref::<StringArray>().unwrap();
1155        let expected = StringArray::from_iter(vec![
1156            Some("test"),
1157            Some("test"),
1158            Some("something else"),
1159            Some("test"),
1160            Some("something else"),
1161            Some("something else"),
1162        ]);
1163        assert_eq!(actual, &expected);
1164    }
1165
1166    #[test]
1167    fn test_zip_scalar_bytes_only_taking_one_side() {
1168        let mask_len = 5;
1169        let all_true_mask = BooleanArray::from(vec![true; mask_len]);
1170        let all_false_mask = BooleanArray::from(vec![false; mask_len]);
1171
1172        let null_scalar = Scalar::new(StringArray::new_null(1));
1173        let non_null_scalar_1 = Scalar::new(StringArray::from_iter_values(["test"]));
1174        let non_null_scalar_2 = Scalar::new(StringArray::from_iter_values(["something else"]));
1175
1176        {
1177            // 1. Test where left is null and right is non-null
1178            //    and mask is all true
1179            let out = zip(&all_true_mask, &null_scalar, &non_null_scalar_1).unwrap();
1180            let actual = out.as_string::<i32>();
1181            let expected = StringArray::from_iter(std::iter::repeat_n(None::<&str>, mask_len));
1182            assert_eq!(actual, &expected);
1183        }
1184
1185        {
1186            // 2. Test where left is null and right is non-null
1187            //    and mask is all false
1188            let out = zip(&all_false_mask, &null_scalar, &non_null_scalar_1).unwrap();
1189            let actual = out.as_string::<i32>();
1190            let expected = StringArray::from_iter(std::iter::repeat_n(Some("test"), mask_len));
1191            assert_eq!(actual, &expected);
1192        }
1193
1194        {
1195            // 3. Test where left is non-null and right is null
1196            //    and mask is all true
1197            let out = zip(&all_true_mask, &non_null_scalar_1, &null_scalar).unwrap();
1198            let actual = out.as_string::<i32>();
1199            let expected = StringArray::from_iter(std::iter::repeat_n(Some("test"), mask_len));
1200            assert_eq!(actual, &expected);
1201        }
1202
1203        {
1204            // 4. Test where left is non-null and right is null
1205            //    and mask is all false
1206            let out = zip(&all_false_mask, &non_null_scalar_1, &null_scalar).unwrap();
1207            let actual = out.as_string::<i32>();
1208            let expected = StringArray::from_iter(std::iter::repeat_n(None::<&str>, mask_len));
1209            assert_eq!(actual, &expected);
1210        }
1211
1212        {
1213            // 5. Test where both left and right are not null
1214            //    and mask is all true
1215            let out = zip(&all_true_mask, &non_null_scalar_1, &non_null_scalar_2).unwrap();
1216            let actual = out.as_string::<i32>();
1217            let expected = StringArray::from_iter(std::iter::repeat_n(Some("test"), mask_len));
1218            assert_eq!(actual, &expected);
1219        }
1220
1221        {
1222            // 6. Test where both left and right are not null
1223            //    and mask is all false
1224            let out = zip(&all_false_mask, &non_null_scalar_1, &non_null_scalar_2).unwrap();
1225            let actual = out.as_string::<i32>();
1226            let expected =
1227                StringArray::from_iter(std::iter::repeat_n(Some("something else"), mask_len));
1228            assert_eq!(actual, &expected);
1229        }
1230
1231        {
1232            // 7. Test where both left and right are null
1233            //    and mask is random
1234            let mask = BooleanArray::from(vec![true, false, true, false, true]);
1235            let out = zip(&mask, &null_scalar, &null_scalar).unwrap();
1236            let actual = out.as_string::<i32>();
1237            let expected = StringArray::from_iter(std::iter::repeat_n(None::<&str>, mask_len));
1238            assert_eq!(actual, &expected);
1239        }
1240    }
1241
1242    #[test]
1243    fn test_scalar_zipper() {
1244        let scalar_truthy = Scalar::new(Int32Array::from_value(42, 1));
1245        let scalar_falsy = Scalar::new(Int32Array::from_value(123, 1));
1246
1247        let mask = BooleanArray::from(vec![false, false, true, true, false]);
1248
1249        let scalar_zipper = ScalarZipper::try_new(&scalar_truthy, &scalar_falsy).unwrap();
1250        let out = scalar_zipper.zip(&mask).unwrap();
1251        let actual = out.as_primitive::<Int32Type>();
1252        let expected = Int32Array::from(vec![Some(123), Some(123), Some(42), Some(42), Some(123)]);
1253        assert_eq!(actual, &expected);
1254
1255        // test with different mask length as well
1256        let mask = BooleanArray::from(vec![true, false, true]);
1257        let out = scalar_zipper.zip(&mask).unwrap();
1258        let actual = out.as_primitive::<Int32Type>();
1259        let expected = Int32Array::from(vec![Some(42), Some(123), Some(42)]);
1260        assert_eq!(actual, &expected);
1261    }
1262
1263    #[test]
1264    fn test_zip_kernel_scalar_strings() {
1265        let scalar_truthy = Scalar::new(StringArray::from(vec!["hello"]));
1266        let scalar_falsy = Scalar::new(StringArray::from(vec!["world"]));
1267
1268        let mask = BooleanArray::from(vec![true, false, true, false, true]);
1269        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1270        let actual = out.as_string::<i32>();
1271        let expected = StringArray::from(vec![
1272            Some("hello"),
1273            Some("world"),
1274            Some("hello"),
1275            Some("world"),
1276            Some("hello"),
1277        ]);
1278        assert_eq!(actual, &expected);
1279    }
1280
1281    #[test]
1282    fn test_zip_kernel_scalar_binary() {
1283        let truthy_bytes: &[u8] = b"\xFF\xFE\xFD";
1284        let falsy_bytes: &[u8] = b"world";
1285        let scalar_truthy = Scalar::new(BinaryArray::from_iter_values(
1286            // Non valid UTF8 bytes
1287            vec![truthy_bytes],
1288        ));
1289        let scalar_falsy = Scalar::new(BinaryArray::from_iter_values(vec![falsy_bytes]));
1290
1291        let mask = BooleanArray::from(vec![true, false, true, false, true]);
1292        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1293        let actual = out.as_binary::<i32>();
1294        let expected = BinaryArray::from(vec![
1295            Some(truthy_bytes),
1296            Some(falsy_bytes),
1297            Some(truthy_bytes),
1298            Some(falsy_bytes),
1299            Some(truthy_bytes),
1300        ]);
1301        assert_eq!(actual, &expected);
1302    }
1303
1304    #[test]
1305    fn test_zip_kernel_scalar_large_binary() {
1306        let truthy_bytes: &[u8] = b"hey";
1307        let falsy_bytes: &[u8] = b"world";
1308        let scalar_truthy = Scalar::new(LargeBinaryArray::from_iter_values(vec![truthy_bytes]));
1309        let scalar_falsy = Scalar::new(LargeBinaryArray::from_iter_values(vec![falsy_bytes]));
1310
1311        let mask = BooleanArray::from(vec![true, false, true, false, true]);
1312        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1313        let actual = out.as_binary::<i64>();
1314        let expected = LargeBinaryArray::from(vec![
1315            Some(truthy_bytes),
1316            Some(falsy_bytes),
1317            Some(truthy_bytes),
1318            Some(falsy_bytes),
1319            Some(truthy_bytes),
1320        ]);
1321        assert_eq!(actual, &expected);
1322    }
1323
1324    // Test to ensure that the precision and scale are kept when zipping Decimal128 data
1325    #[test]
1326    fn test_zip_decimal_with_custom_precision_and_scale() {
1327        let arr = Decimal128Array::from_iter_values([12345, 456, 7890, -123223423432432])
1328            .with_precision_and_scale(20, 2)
1329            .unwrap();
1330
1331        let arr: ArrayRef = Arc::new(arr);
1332
1333        let scalar_1 = Scalar::new(arr.slice(0, 1));
1334        let scalar_2 = Scalar::new(arr.slice(1, 1));
1335        let null_scalar = Scalar::new(new_null_array(arr.data_type(), 1));
1336        let array_1: ArrayRef = arr.slice(0, 2);
1337        let array_2: ArrayRef = arr.slice(2, 2);
1338
1339        test_zip_output_data_types_for_input(scalar_1, scalar_2, null_scalar, array_1, array_2);
1340    }
1341
1342    // Test to ensure that the timezone is kept when zipping TimestampArray data
1343    #[test]
1344    fn test_zip_timestamp_with_timezone() {
1345        let arr = TimestampSecondArray::from(vec![0, 1000, 2000, 4000])
1346            .with_timezone("+01:00".to_string());
1347
1348        let arr: ArrayRef = Arc::new(arr);
1349
1350        let scalar_1 = Scalar::new(arr.slice(0, 1));
1351        let scalar_2 = Scalar::new(arr.slice(1, 1));
1352        let null_scalar = Scalar::new(new_null_array(arr.data_type(), 1));
1353        let array_1: ArrayRef = arr.slice(0, 2);
1354        let array_2: ArrayRef = arr.slice(2, 2);
1355
1356        test_zip_output_data_types_for_input(scalar_1, scalar_2, null_scalar, array_1, array_2);
1357    }
1358
1359    fn test_zip_output_data_types_for_input(
1360        scalar_1: Scalar<ArrayRef>,
1361        scalar_2: Scalar<ArrayRef>,
1362        null_scalar: Scalar<ArrayRef>,
1363        array_1: ArrayRef,
1364        array_2: ArrayRef,
1365    ) {
1366        // non null Scalar vs non null Scalar
1367        test_zip_output_data_type(&scalar_1, &scalar_2, 10);
1368
1369        // null Scalar vs non-null Scalar (and vice versa)
1370        test_zip_output_data_type(&null_scalar, &scalar_1, 10);
1371        test_zip_output_data_type(&scalar_1, &null_scalar, 10);
1372
1373        // non-null Scalar and array (and vice versa)
1374        test_zip_output_data_type(&array_1.as_ref(), &scalar_1, array_1.len());
1375        test_zip_output_data_type(&scalar_1, &array_1.as_ref(), array_1.len());
1376
1377        // Array and null scalar (and vice versa)
1378        test_zip_output_data_type(&array_1.as_ref(), &null_scalar, array_1.len());
1379
1380        test_zip_output_data_type(&null_scalar, &array_1.as_ref(), array_1.len());
1381
1382        // Both arrays
1383        test_zip_output_data_type(&array_1.as_ref(), &array_2.as_ref(), array_1.len());
1384    }
1385
1386    fn test_zip_output_data_type(truthy: &dyn Datum, falsy: &dyn Datum, mask_length: usize) {
1387        let expected_data_type = truthy.get().0.data_type().clone();
1388        assert_eq!(&expected_data_type, falsy.get().0.data_type());
1389
1390        // Try different masks to test different paths
1391        let mask_all_true = BooleanArray::from(vec![true; mask_length]);
1392        let mask_all_false = BooleanArray::from(vec![false; mask_length]);
1393        let mask_some_true_and_false =
1394            BooleanArray::from((0..mask_length).map(|i| i % 2 == 0).collect::<Vec<bool>>());
1395
1396        for mask in [&mask_all_true, &mask_all_false, &mask_some_true_and_false] {
1397            let out = zip(mask, truthy, falsy).unwrap();
1398            assert_eq!(out.data_type(), &expected_data_type);
1399        }
1400    }
1401
1402    #[test]
1403    fn zip_scalar_fallback_impl() {
1404        let truthy_list_item_scalar = Some(vec![Some(1), None, Some(3)]);
1405        let truthy_list_array_scalar =
1406            Scalar::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1407                truthy_list_item_scalar.clone(),
1408            ]));
1409        let falsy_list_item_scalar = Some(vec![None, Some(2), Some(4)]);
1410        let falsy_list_array_scalar =
1411            Scalar::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1412                falsy_list_item_scalar.clone(),
1413            ]));
1414        let mask = BooleanArray::from(vec![true, false, true, false, false, true, false]);
1415        let out = zip(&mask, &truthy_list_array_scalar, &falsy_list_array_scalar).unwrap();
1416        let actual = out.as_list::<i32>();
1417
1418        let expected = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1419            truthy_list_item_scalar.clone(),
1420            falsy_list_item_scalar.clone(),
1421            truthy_list_item_scalar.clone(),
1422            falsy_list_item_scalar.clone(),
1423            falsy_list_item_scalar.clone(),
1424            truthy_list_item_scalar.clone(),
1425            falsy_list_item_scalar.clone(),
1426        ]);
1427        assert_eq!(actual, &expected);
1428    }
1429
1430    #[test]
1431    fn test_zip_kernel_scalar_strings_array_view() {
1432        let scalar_truthy = Scalar::new(StringViewArray::from(vec!["hello"]));
1433        let scalar_falsy = Scalar::new(StringViewArray::from(vec!["world"]));
1434
1435        let mask = BooleanArray::from(vec![true, false, true, false]);
1436        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1437        let actual = out.as_string_view();
1438        let expected = StringViewArray::from(vec![
1439            Some("hello"),
1440            Some("world"),
1441            Some("hello"),
1442            Some("world"),
1443        ]);
1444        assert_eq!(actual, &expected);
1445    }
1446
1447    #[test]
1448    fn test_zip_kernel_scalar_binary_array_view() {
1449        let scalar_truthy = Scalar::new(BinaryViewArray::from_iter_values(vec![b"hello"]));
1450        let scalar_falsy = Scalar::new(BinaryViewArray::from_iter_values(vec![b"world"]));
1451
1452        let mask = BooleanArray::from(vec![true, false]);
1453        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1454        let actual = out.as_byte_view();
1455        let expected = BinaryViewArray::from_iter_values(vec![b"hello", b"world"]);
1456        assert_eq!(actual, &expected);
1457    }
1458
1459    #[test]
1460    fn test_zip_kernel_scalar_strings_array_view_with_nulls() {
1461        let scalar_truthy = Scalar::new(StringViewArray::from_iter_values(["hello"]));
1462        let scalar_falsy = Scalar::new(StringViewArray::new_null(1));
1463
1464        let mask = BooleanArray::from(vec![true, true, false, false, true]);
1465        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1466        let actual = out.as_any().downcast_ref::<StringViewArray>().unwrap();
1467        let expected = StringViewArray::from_iter(vec![
1468            Some("hello"),
1469            Some("hello"),
1470            None,
1471            None,
1472            Some("hello"),
1473        ]);
1474        assert_eq!(actual, &expected);
1475    }
1476
1477    #[test]
1478    fn test_zip_kernel_scalar_strings_array_view_all_true_null() {
1479        let scalar_truthy = Scalar::new(StringViewArray::new_null(1));
1480        let scalar_falsy = Scalar::new(StringViewArray::new_null(1));
1481        let mask = BooleanArray::from(vec![true, true]);
1482        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1483        let actual = out.as_any().downcast_ref::<StringViewArray>().unwrap();
1484        let expected = StringViewArray::from_iter(vec![None::<String>, None]);
1485        assert_eq!(actual, &expected);
1486    }
1487
1488    #[test]
1489    fn test_zip_kernel_scalar_strings_array_view_all_false_null() {
1490        let scalar_truthy = Scalar::new(StringViewArray::new_null(1));
1491        let scalar_falsy = Scalar::new(StringViewArray::new_null(1));
1492        let mask = BooleanArray::from(vec![false, false]);
1493        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1494        let actual = out.as_any().downcast_ref::<StringViewArray>().unwrap();
1495        let expected = StringViewArray::from_iter(vec![None::<String>, None]);
1496        assert_eq!(actual, &expected);
1497    }
1498
1499    #[test]
1500    fn test_zip_kernel_scalar_string_array_view_all_true() {
1501        let scalar_truthy = Scalar::new(StringViewArray::from(vec!["hello"]));
1502        let scalar_falsy = Scalar::new(StringViewArray::from(vec!["world"]));
1503
1504        let mask = BooleanArray::from(vec![true, true]);
1505        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1506        let actual = out.as_string_view();
1507        let expected = StringViewArray::from(vec![Some("hello"), Some("hello")]);
1508        assert_eq!(actual, &expected);
1509    }
1510
1511    #[test]
1512    fn test_zip_kernel_scalar_string_array_view_all_false() {
1513        let scalar_truthy = Scalar::new(StringViewArray::from(vec!["hello"]));
1514        let scalar_falsy = Scalar::new(StringViewArray::from(vec!["world"]));
1515
1516        let mask = BooleanArray::from(vec![false, false]);
1517        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1518        let actual = out.as_string_view();
1519        let expected = StringViewArray::from(vec![Some("world"), Some("world")]);
1520        assert_eq!(actual, &expected);
1521    }
1522
1523    #[test]
1524    fn test_zip_kernel_scalar_strings_large_strings() {
1525        let scalar_truthy = Scalar::new(StringViewArray::from(vec!["longer than 12 bytes"]));
1526        let scalar_falsy = Scalar::new(StringViewArray::from(vec!["another longer than 12 bytes"]));
1527
1528        let mask = BooleanArray::from(vec![true, false]);
1529        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1530        let actual = out.as_string_view();
1531        let expected = StringViewArray::from(vec![
1532            Some("longer than 12 bytes"),
1533            Some("another longer than 12 bytes"),
1534        ]);
1535        assert_eq!(actual, &expected);
1536    }
1537
1538    #[test]
1539    fn test_zip_kernel_scalar_strings_array_view_large_short_strings() {
1540        let scalar_truthy = Scalar::new(StringViewArray::from(vec!["hello"]));
1541        let scalar_falsy = Scalar::new(StringViewArray::from(vec!["longer than 12 bytes"]));
1542
1543        let mask = BooleanArray::from(vec![true, false, true, false]);
1544        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1545        let actual = out.as_string_view();
1546        let expected = StringViewArray::from(vec![
1547            Some("hello"),
1548            Some("longer than 12 bytes"),
1549            Some("hello"),
1550            Some("longer than 12 bytes"),
1551        ]);
1552        assert_eq!(actual, &expected);
1553    }
1554    #[test]
1555    fn test_zip_kernel_scalar_strings_array_view_large_all_true() {
1556        let scalar_truthy = Scalar::new(StringViewArray::from(vec!["longer than 12 bytes"]));
1557        let scalar_falsy = Scalar::new(StringViewArray::from(vec!["another longer than 12 bytes"]));
1558
1559        let mask = BooleanArray::from(vec![true, true]);
1560        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1561        let actual = out.as_string_view();
1562        let expected = StringViewArray::from(vec![
1563            Some("longer than 12 bytes"),
1564            Some("longer than 12 bytes"),
1565        ]);
1566        assert_eq!(actual, &expected);
1567    }
1568
1569    #[test]
1570    fn test_zip_kernel_scalar_strings_array_view_large_all_false() {
1571        let scalar_truthy = Scalar::new(StringViewArray::from(vec!["longer than 12 bytes"]));
1572        let scalar_falsy = Scalar::new(StringViewArray::from(vec!["another longer than 12 bytes"]));
1573
1574        let mask = BooleanArray::from(vec![false, false]);
1575        let out = zip(&mask, &scalar_truthy, &scalar_falsy).unwrap();
1576        let actual = out.as_string_view();
1577        let expected = StringViewArray::from(vec![
1578            Some("another longer than 12 bytes"),
1579            Some("another longer than 12 bytes"),
1580        ]);
1581        assert_eq!(actual, &expected);
1582    }
1583}
arrow_select/zip.rs

arrow_select/
zip.rs
