Function zip

arrow-select/src/zip.rs:99–146

Zip two arrays by some boolean mask. - Where `mask` is `true`, values of `truthy` are taken - Where `mask` is `false` or `NULL`, values of `falsy` are taken # Example: `zip` two arrays ``` # use std::sync::Arc; # use arrow_array::{ArrayRef, BooleanArray, Int32Array}; # use arrow_select::zip::zip; // mask: [true, true, false, NULL, true] let mask = BooleanArray::from(vec![ Some(true), Some(true),

(
    mask: &BooleanArray,
    truthy: &dyn Datum,
    falsy: &dyn Datum,
)

/// assert_eq!(&result, &expected);
/// ```
pub fn zip(
    mask: &BooleanArray,
    truthy: &dyn Datum,
    falsy: &dyn Datum,
) -> Result<ArrayRef, ArrowError> {
    let (truthy_array, truthy_is_scalar) = truthy.get();
    let (falsy_array, falsy_is_scalar) = falsy.get();

    if falsy_is_scalar && truthy_is_scalar {
        let zipper = ScalarZipper::try_new(truthy, falsy)?;
        return zipper.zip_impl.create_output(mask);
    }

    let truthy = truthy_array;
    let falsy = falsy_array;

    if truthy.data_type() != falsy.data_type() {
        return Err(ArrowError::InvalidArgumentError(
            "arguments need to have the same data type".into(),
        ));
    }

    if truthy_is_scalar && truthy.len() != 1 {
        return Err(ArrowError::InvalidArgumentError(
            "scalar arrays must have 1 element".into(),
        ));
    }
    if !truthy_is_scalar && truthy.len() != mask.len() {
        return Err(ArrowError::InvalidArgumentError(
            "all arrays should have the same length".into(),
        ));
    }
    if falsy_is_scalar && falsy.len() != 1 {
        return Err(ArrowError::InvalidArgumentError(
            "scalar arrays must have 1 element".into(),
        ));
    }
    if !falsy_is_scalar && falsy.len() != mask.len() {
        return Err(ArrowError::InvalidArgumentError(
            "all arrays should have the same length".into(),
        ));
    }

    let falsy = falsy.to_data();
    let truthy = truthy.to_data();

    zip_impl(mask, &truthy, truthy_is_scalar, &falsy, falsy_is_scalar)
}

fn zip_impl(
    mask: &BooleanArray,