vortex_array/scalar_fn/fns/zip/

mod.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

mod kernel;

use std::fmt::Display;
use std::fmt::Formatter;
use std::sync::Arc;

pub use kernel::*;
use vortex_error::VortexResult;
use vortex_error::vortex_ensure;
use vortex_error::vortex_err;
use vortex_mask::Mask;
use vortex_mask::MaskValues;
use vortex_session::VortexSession;
use vortex_session::registry::CachedId;

use crate::ArrayRef;
use crate::ExecutionCtx;
use crate::IntoArray;
use crate::arrays::BoolArray;
use crate::arrays::ScalarFnArray;
use crate::arrays::bool::BoolArrayExt;
use crate::builders::ArrayBuilder;
use crate::builders::builder_with_capacity_in;
use crate::builtins::ArrayBuiltins;
use crate::dtype::DType;
use crate::dtype::StructFields;
use crate::expr::BoundExpression;
use crate::expr::bound;
use crate::expr::display::ExprDisplay;
use crate::scalar_fn::Arity;
use crate::scalar_fn::ChildName;
use crate::scalar_fn::EmptyOptions;
use crate::scalar_fn::ExecutionArgs;
use crate::scalar_fn::ScalarFnId;
use crate::scalar_fn::ScalarFnVTable;
use crate::scalar_fn::ScalarFnVTableExt;
use crate::scalar_fn::fns::literal::Literal;
use crate::validity::Validity;

/// An expression that conditionally selects between two arrays based on a boolean mask.
///
/// For each position `i`, `result[i] = if mask[i] then if_true[i] else if_false[i]`.
///
/// Null values in the mask are treated as false (selecting `if_false`). This follows
/// SQL semantics (DuckDB, Trino) where a null condition falls through to the ELSE branch,
/// rather than Arrow's `if_else` which propagates null conditions to the output.
#[derive(Clone)]
pub struct Zip;

impl Zip {
    /// Creates a lazy conditional selection between `if_true` and `if_false`.
    ///
    /// # Errors
    ///
    /// Returns an error if the children have different lengths, the values have incompatible
    /// dtypes, or `mask` is not boolean data.
    pub fn try_new(
        if_true: ArrayRef,
        if_false: ArrayRef,
        mask: ArrayRef,
    ) -> VortexResult<ScalarFnArray> {
        ScalarFnArray::try_new(Zip.bind(EmptyOptions), [if_true, if_false, mask])
    }
}

impl ScalarFnVTable for Zip {
    type Options = EmptyOptions;

    fn id(&self) -> ScalarFnId {
        static ID: CachedId = CachedId::new("vortex.zip");
        *ID
    }

    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
        Ok(Some(vec![]))
    }

    fn deserialize(
        &self,
        _metadata: &[u8],
        _session: &VortexSession,
    ) -> VortexResult<Self::Options> {
        Ok(EmptyOptions)
    }

    fn arity(&self, _options: &Self::Options) -> Arity {
        Arity::Exact(3)
    }

    fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
        match child_idx {
            0 => ChildName::from("if_true"),
            1 => ChildName::from("if_false"),
            2 => ChildName::from("mask"),
            _ => unreachable!("Invalid child index {} for Zip expression", child_idx),
        }
    }

    fn fmt_sql(
        &self,
        _options: &Self::Options,
        expr: &dyn ExprDisplay,
        f: &mut Formatter<'_>,
    ) -> std::fmt::Result {
        write!(f, "zip(")?;
        Display::fmt(expr.display_child(0), f)?;
        write!(f, ", ")?;
        Display::fmt(expr.display_child(1), f)?;
        write!(f, ", ")?;
        Display::fmt(expr.display_child(2), f)?;
        write!(f, ")")
    }

    fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
        vortex_ensure!(
            matches!(arg_dtypes[2], DType::Bool(_)),
            "zip requires mask to be a boolean type, got {}",
            arg_dtypes[2]
        );
        zip_return_dtype(&arg_dtypes[0], &arg_dtypes[1])
    }

    fn execute(
        &self,
        _options: &Self::Options,
        args: &dyn ExecutionArgs,
        ctx: &mut ExecutionCtx,
    ) -> VortexResult<ArrayRef> {
        let if_true = args.get(0)?;
        let if_false = args.get(1)?;
        let mask_array = args.get(2)?;

        let mask = mask_array
            .execute::<BoolArray>(ctx)?
            .to_mask_fill_null_false(ctx);

        let return_dtype = zip_return_dtype(if_true.dtype(), if_false.dtype())?;

        if mask.all_true() {
            return if_true.cast(return_dtype)?.execute(ctx);
        }

        if mask.all_false() {
            return if_false.cast(return_dtype)?.execute(ctx);
        }

        if !if_true.is_canonical() || !if_false.is_canonical() {
            let if_true = if_true.execute::<ArrayRef>(ctx)?;
            let if_false = if_false.execute::<ArrayRef>(ctx)?;
            return mask.into_array().zip(if_true, if_false);
        }

        zip_impl(&if_true, &if_false, &mask, ctx)
    }

    fn simplify(
        &self,
        _options: &Self::Options,
        expr: &BoundExpression,
    ) -> VortexResult<Option<BoundExpression>> {
        let Some(mask_lit) = expr.child(2).as_opt::<Literal>() else {
            return Ok(None);
        };

        if let Some(mask_val) = mask_lit.as_bool().value() {
            let selected = if mask_val {
                expr.child(0).clone()
            } else {
                expr.child(1).clone()
            };
            // The selected branch can be less nullable than the zip result, which unions
            // both branches' nullability, so keep the dtype that binding inferred.
            return Ok(Some(if selected.dtype() == expr.dtype() {
                selected
            } else {
                bound::cast(selected, expr.dtype().clone())
            }));
        }

        Ok(None)
    }

    fn is_strict(&self, _options: &Self::Options) -> bool {
        // A null in an unselected branch does not force a null output.
        false
    }

    fn is_infallible(&self, _options: &Self::Options) -> bool {
        true
    }
}

pub(crate) fn zip_impl(
    if_true: &ArrayRef,
    if_false: &ArrayRef,
    mask: &Mask,
    ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
    assert_eq!(
        if_true.len(),
        if_false.len(),
        "zip requires arrays to have the same size"
    );

    let return_type = zip_return_dtype(if_true.dtype(), if_false.dtype())?;

    let mask_values = match mask {
        Mask::AllTrue(_) | Mask::AllFalse(0) => return if_true.cast(return_type),
        Mask::AllFalse(_) => return if_false.cast(return_type),
        Mask::Values(values) => values,
    };

    // `append_to_builder` requires exact dtype equality, so normalize branch
    // nullability to the output dtype before appending slices into the builder.
    let if_true = if_true.cast(return_type.clone())?;
    let if_false = if_false.cast(return_type.clone())?;

    zip_impl_with_builder(
        &if_true,
        &if_false,
        mask_values.as_ref(),
        builder_with_capacity_in(&return_type, if_true.len(), ctx.allocator()),
        ctx,
    )
}

fn zip_return_dtype(if_true: &DType, if_false: &DType) -> VortexResult<DType> {
    zip_nullability_union(if_true, if_false).ok_or_else(|| {
        vortex_err!(
            "zip requires if_true and if_false to have the same base type, got {} and {}",
            if_true,
            if_false
        )
    })
}

fn zip_nullability_union(lhs: &DType, rhs: &DType) -> Option<DType> {
    let nullability = lhs.nullability() | rhs.nullability();

    match (lhs, rhs) {
        (DType::List(lhs_element, _), DType::List(rhs_element, _)) => Some(DType::List(
            Arc::new(zip_nullability_union(lhs_element, rhs_element)?),
            nullability,
        )),
        (
            DType::FixedSizeList(lhs_element, lhs_size, _),
            DType::FixedSizeList(rhs_element, rhs_size, _),
        ) if lhs_size == rhs_size => Some(DType::FixedSizeList(
            Arc::new(zip_nullability_union(lhs_element, rhs_element)?),
            *lhs_size,
            nullability,
        )),
        (DType::Map(lhs_map, _), DType::Map(rhs_map, _))
            if lhs_map.keys_sorted() == rhs_map.keys_sorted() =>
        {
            DType::map(
                zip_nullability_union(&lhs_map.key_dtype(), &rhs_map.key_dtype())?,
                zip_nullability_union(&lhs_map.value_dtype(), &rhs_map.value_dtype())?,
                lhs_map.keys_sorted(),
                nullability,
            )
            .ok()
        }
        (DType::Struct(lhs_fields, _), DType::Struct(rhs_fields, _))
            if lhs_fields.names() == rhs_fields.names() =>
        {
            let fields = lhs_fields
                .fields()
                .zip(rhs_fields.fields())
                .map(|(lhs, rhs)| zip_nullability_union(&lhs, &rhs))
                .collect::<Option<Vec<_>>>()?;
            Some(DType::Struct(
                StructFields::new(lhs_fields.names().clone(), fields),
                nullability,
            ))
        }
        (DType::Union(lhs_variants, _), DType::Union(rhs_variants, _))
            if lhs_variants == rhs_variants =>
        {
            Some(DType::Union(lhs_variants.clone(), nullability))
        }
        _ if lhs.eq_ignore_nullability(rhs) => Some(lhs.with_nullability(nullability)),
        _ => None,
    }
}

fn zip_impl_with_builder(
    if_true: &ArrayRef,
    if_false: &ArrayRef,
    mask: &MaskValues,
    mut builder: Box<dyn ArrayBuilder>,
    ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
    for (start, end) in mask.slices() {
        if builder.len() < *start {
            if_false
                .slice(builder.len()..*start)?
                .append_to_builder(builder.as_mut(), ctx)?;
        }
        if_true
            .slice(*start..*end)?
            .append_to_builder(builder.as_mut(), ctx)?;
    }
    if builder.len() < if_false.len() {
        if_false
            .slice(builder.len()..if_false.len())?
            .append_to_builder(builder.as_mut(), ctx)?;
    }
    Ok(builder.finish())
}

/// Combine two validities for a row-wise zip: take `if_true`'s validity where `mask` is set and
/// `if_false`'s where it is not.
///
/// That selection is itself a zip over the two boolean validity bitmaps, so it is built as a (lazy)
/// zip array — reusing the zip machinery rather than re-deriving the mask algebra. Trivial cases
/// where both sides' validity already agrees skip the zip. `mask` must already be null-filled so the
/// selection matches the accompanying value selection. Shared by the per-encoding zip kernels (e.g.
/// `Bool`, `Primitive`) that build their result directly.
pub(crate) fn zip_validity(
    if_true: Validity,
    if_false: Validity,
    mask: &Mask,
) -> VortexResult<Validity> {
    match (&if_true, &if_false) {
        (Validity::NonNullable, Validity::NonNullable) => return Ok(Validity::NonNullable),
        (Validity::AllValid, Validity::AllValid) => return Ok(Validity::AllValid),
        (Validity::AllInvalid, Validity::AllInvalid) => return Ok(Validity::AllInvalid),
        _ => {}
    }

    let len = mask.len();
    let validity = mask
        .clone()
        .into_array()
        .zip(if_true.to_array(len), if_false.to_array(len))?;
    Ok(Validity::Array(validity))
}

#[cfg(test)]
mod tests {
    use vortex_buffer::buffer;
    use vortex_error::VortexResult;
    use vortex_mask::Mask;

    use super::zip_impl;
    use crate::ArrayRef;
    use crate::IntoArray;
    use crate::VortexSessionExecute;
    use crate::array_session;
    use crate::arrays::ConstantArray;
    use crate::arrays::PrimitiveArray;
    use crate::arrays::Struct;
    use crate::arrays::StructArray;
    use crate::arrays::VarBinView;
    use crate::arrays::VarBinViewArray;
    use crate::assert_arrays_eq;
    use crate::builders::ArrayBuilder;
    use crate::builders::BufferGrowthStrategy;
    use crate::builders::VarBinViewBuilder;
    use crate::builtins::ArrayBuiltins;
    use crate::columnar::Columnar;
    use crate::dtype::DType;
    use crate::dtype::Nullability;
    use crate::dtype::PType;
    use crate::expr::lit;
    use crate::expr::root;
    use crate::expr::zip_expr;
    use crate::scalar::Scalar;

    #[test]
    fn dtype() {
        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
        let expr = zip_expr(lit(true), root(), lit(0i32));
        let result_dtype = expr.return_dtype(&dtype).unwrap();
        assert_eq!(
            result_dtype,
            DType::Primitive(PType::I32, Nullability::NonNullable)
        );
    }

    #[test]
    fn test_display() {
        let expr = zip_expr(lit(true), root(), lit(0i32));
        assert_eq!(expr.to_string(), "zip($, 0i32, true)");
    }

    #[test]
    fn test_zip_basic() {
        let mut ctx = array_session().create_execution_ctx();
        let mask = Mask::from_iter([true, false, false, true, false]);
        let if_true = buffer![10, 20, 30, 40, 50].into_array();
        let if_false = buffer![1, 2, 3, 4, 5].into_array();

        let result = mask.into_array().zip(if_true, if_false).unwrap();
        let expected = buffer![10, 2, 3, 40, 5].into_array();

        assert_arrays_eq!(result, expected, &mut ctx);
    }

    #[test]
    fn test_zip_all_true() {
        let mut ctx = array_session().create_execution_ctx();
        let mask = Mask::new_true(4);
        let if_true = buffer![10, 20, 30, 40].into_array();
        let if_false =
            PrimitiveArray::from_option_iter([Some(1), Some(2), Some(3), None]).into_array();

        let result = mask.into_array().zip(if_true, if_false.clone()).unwrap();
        let expected =
            PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), Some(40)]).into_array();

        assert_arrays_eq!(result, expected, &mut ctx);
        assert_eq!(result.dtype(), if_false.dtype())
    }

    #[test]
    fn test_zip_all_false_widens_nullability() {
        let mut ctx = array_session().create_execution_ctx();
        let mask = Mask::new_false(4);
        let if_true =
            PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), None]).into_array();
        let if_false = buffer![1i32, 2, 3, 4].into_array();

        let result = mask.into_array().zip(if_true.clone(), if_false).unwrap();
        let expected =
            PrimitiveArray::from_option_iter([Some(1), Some(2), Some(3), Some(4)]).into_array();

        assert_arrays_eq!(result, expected, &mut ctx);
        assert_eq!(result.dtype(), if_true.dtype());
    }

    #[test]
    fn test_zip_impl_all_true_widens_nullability() -> VortexResult<()> {
        let mask = Mask::new_true(4);
        let if_true = buffer![10i32, 20, 30, 40].into_array();
        let if_false =
            PrimitiveArray::from_option_iter([Some(1), Some(2), Some(3), None]).into_array();

        let mut ctx = array_session().create_execution_ctx();
        let result = zip_impl(&if_true, &if_false, &mask, &mut ctx)?;
        assert_arrays_eq!(
            result,
            PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30), Some(40)])
                .into_array(),
            &mut ctx
        );
        assert_eq!(result.dtype(), if_false.dtype());
        Ok(())
    }

    #[test]
    fn test_zip_impl_all_false_widens_nullability() -> VortexResult<()> {
        let mask = Mask::new_false(4);
        let if_true =
            PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), None]).into_array();
        let if_false = buffer![1i32, 2, 3, 4].into_array();

        let mut ctx = array_session().create_execution_ctx();
        let result = zip_impl(&if_true, &if_false, &mask, &mut ctx)?;
        assert_arrays_eq!(
            result,
            PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]).into_array(),
            &mut ctx
        );
        assert_eq!(result.dtype(), if_true.dtype());
        Ok(())
    }

    #[test]
    #[should_panic]
    fn test_invalid_lengths() {
        let mask = Mask::new_false(4);
        let if_true = buffer![10, 20, 30].into_array();
        let if_false = buffer![1, 2, 3, 4].into_array();

        let _result = mask.into_array().zip(if_true, if_false).unwrap();
    }

    #[test]
    fn test_fragmentation() -> VortexResult<()> {
        let len = 100;

        let const1 = ConstantArray::new(
            Scalar::utf8("hello_this_is_a_longer_string", Nullability::Nullable),
            len,
        )
        .into_array();

        let const2 = ConstantArray::new(
            Scalar::utf8("world_this_is_another_string", Nullability::Nullable),
            len,
        )
        .into_array();

        let indices: Vec<usize> = (0..len).step_by(2).collect();
        let mask = Mask::from_indices(len, indices);
        let mask_array = mask.into_array();

        let mut ctx = array_session().create_execution_ctx();
        let result = mask_array
            .zip(const1.clone(), const2.clone())?
            .execute::<Columnar>(&mut ctx)?
            .into_array();

        insta::assert_snapshot!(result.display_tree(), @r"
        root: vortex.varbinview(utf8?, len=100) nbytes=1.66 kB (100.00%) [all_valid]
          metadata: 
          buffer: buffer_0 host 29 B (align=1) (1.75%)
          buffer: buffer_1 host 28 B (align=1) (1.69%)
          buffer: views host 1.60 kB (align=16) (96.56%)
        ");

        let wrapped1 = StructArray::try_from_iter([("nested", const1)])?.into_array();
        let wrapped2 = StructArray::try_from_iter([("nested", const2)])?.into_array();

        let wrapped_result = mask_array
            .zip(wrapped1, wrapped2)?
            .execute::<ArrayRef>(&mut ctx)?;
        assert!(wrapped_result.is::<Struct>());

        Ok(())
    }

    #[test]
    fn test_varbinview_zip() {
        let if_true = {
            let mut builder = VarBinViewBuilder::new_in(
                DType::Utf8(Nullability::NonNullable),
                10,
                Default::default(),
                BufferGrowthStrategy::fixed(64 * 1024),
                0.0,
                vortex_buffer::BufferAllocatorRef::statically_allocated(),
            );
            for _ in 0..100 {
                builder.append_value("Hello");
                builder.append_value("Hello this is a long string that won't be inlined.");
            }
            builder.finish()
        };

        let if_false = {
            let mut builder = VarBinViewBuilder::new_in(
                DType::Utf8(Nullability::NonNullable),
                10,
                Default::default(),
                BufferGrowthStrategy::fixed(64 * 1024),
                0.0,
                vortex_buffer::BufferAllocatorRef::statically_allocated(),
            );
            for _ in 0..100 {
                builder.append_value("Hello2");
                builder.append_value("Hello2 this is a long string that won't be inlined.");
            }
            builder.finish()
        };

        let mask = Mask::from_indices(200, (0..100).filter(|i| i % 3 != 0));
        let mask_array = mask.clone().into_array();

        let mut ctx = array_session().create_execution_ctx();
        let zipped = mask_array
            .zip(if_true, if_false)
            .unwrap()
            .execute::<ArrayRef>(&mut ctx)
            .unwrap();
        let zipped = zipped.as_opt::<VarBinView>().unwrap();
        assert_eq!(zipped.data_buffers().len(), 2);

        let true_value = |i: usize| {
            if i.is_multiple_of(2) {
                "Hello"
            } else {
                "Hello this is a long string that won't be inlined."
            }
        };
        let false_value = |i: usize| {
            if i.is_multiple_of(2) {
                "Hello2"
            } else {
                "Hello2 this is a long string that won't be inlined."
            }
        };
        let expected = VarBinViewArray::from_iter_str((0..200).map(|i| {
            if mask.value(i) {
                true_value(i)
            } else {
                false_value(i)
            }
        }));
        assert_arrays_eq!(zipped.array().clone(), expected, &mut ctx);
    }
}