vortex_array/scalar_fn/fns/zip/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod kernel;
5
6use std::fmt::Formatter;
7
8pub use kernel::*;
9use vortex_error::VortexExpect as _;
10use vortex_error::VortexResult;
11use vortex_error::vortex_ensure;
12use vortex_mask::Mask;
13use vortex_mask::MaskValues;
14use vortex_session::VortexSession;
15
16use crate::ArrayRef;
17use crate::DynArray;
18use crate::ExecutionCtx;
19use crate::IntoArray;
20use crate::arrays::BoolArray;
21use crate::builders::ArrayBuilder;
22use crate::builders::builder_with_capacity;
23use crate::builtins::ArrayBuiltins;
24use crate::dtype::DType;
25use crate::expr::Expression;
26use crate::scalar_fn::Arity;
27use crate::scalar_fn::ChildName;
28use crate::scalar_fn::EmptyOptions;
29use crate::scalar_fn::ExecutionArgs;
30use crate::scalar_fn::ScalarFnId;
31use crate::scalar_fn::ScalarFnVTable;
32use crate::scalar_fn::SimplifyCtx;
33use crate::scalar_fn::fns::literal::Literal;
34
35/// An expression that conditionally selects between two arrays based on a boolean mask.
36///
37/// For each position `i`, `result[i] = if mask[i] then if_true[i] else if_false[i]`.
38///
39/// Null values in the mask are treated as false (selecting `if_false`). This follows
40/// SQL semantics (DuckDB, Trino) where a null condition falls through to the ELSE branch,
41/// rather than Arrow's `if_else` which propagates null conditions to the output.
42#[derive(Clone)]
43pub struct Zip;
44
45impl ScalarFnVTable for Zip {
46    type Options = EmptyOptions;
47
48    fn id(&self) -> ScalarFnId {
49        ScalarFnId::from("vortex.zip")
50    }
51
52    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
53        Ok(Some(vec![]))
54    }
55
56    fn deserialize(
57        &self,
58        _metadata: &[u8],
59        _session: &VortexSession,
60    ) -> VortexResult<Self::Options> {
61        Ok(EmptyOptions)
62    }
63
64    fn arity(&self, _options: &Self::Options) -> Arity {
65        Arity::Exact(3)
66    }
67
68    fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
69        match child_idx {
70            0 => ChildName::from("if_true"),
71            1 => ChildName::from("if_false"),
72            2 => ChildName::from("mask"),
73            _ => unreachable!("Invalid child index {} for Zip expression", child_idx),
74        }
75    }
76
77    fn fmt_sql(
78        &self,
79        _options: &Self::Options,
80        expr: &Expression,
81        f: &mut Formatter<'_>,
82    ) -> std::fmt::Result {
83        write!(f, "zip(")?;
84        expr.child(0).fmt_sql(f)?;
85        write!(f, ", ")?;
86        expr.child(1).fmt_sql(f)?;
87        write!(f, ", ")?;
88        expr.child(2).fmt_sql(f)?;
89        write!(f, ")")
90    }
91
92    fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
93        vortex_ensure!(
94            arg_dtypes[0].eq_ignore_nullability(&arg_dtypes[1]),
95            "zip requires if_true and if_false to have the same base type, got {} and {}",
96            arg_dtypes[0],
97            arg_dtypes[1]
98        );
99        vortex_ensure!(
100            matches!(arg_dtypes[2], DType::Bool(_)),
101            "zip requires mask to be a boolean type, got {}",
102            arg_dtypes[2]
103        );
104        Ok(arg_dtypes[0]
105            .clone()
106            .union_nullability(arg_dtypes[1].nullability()))
107    }
108
109    fn execute(
110        &self,
111        _options: &Self::Options,
112        args: &dyn ExecutionArgs,
113        ctx: &mut ExecutionCtx,
114    ) -> VortexResult<ArrayRef> {
115        let if_true = args.get(0)?;
116        let if_false = args.get(1)?;
117        let mask_array = args.get(2)?;
118
119        let mask = mask_array
120            .execute::<BoolArray>(ctx)?
121            .to_mask_fill_null_false();
122
123        let return_dtype = if_true
124            .dtype()
125            .clone()
126            .union_nullability(if_false.dtype().nullability());
127
128        if mask.all_true() {
129            return if_true.cast(return_dtype)?.execute(ctx);
130        }
131
132        if mask.all_false() {
133            return if_false.cast(return_dtype)?.execute(ctx);
134        }
135
136        if !if_true.is_canonical() || !if_false.is_canonical() {
137            let if_true = if_true.execute::<ArrayRef>(ctx)?;
138            let if_false = if_false.execute::<ArrayRef>(ctx)?;
139            return mask.into_array().zip(if_true, if_false);
140        }
141
142        zip_impl(&if_true, &if_false, &mask)
143    }
144
145    fn simplify(
146        &self,
147        _options: &Self::Options,
148        expr: &Expression,
149        _ctx: &dyn SimplifyCtx,
150    ) -> VortexResult<Option<Expression>> {
151        let Some(mask_lit) = expr.child(2).as_opt::<Literal>() else {
152            return Ok(None);
153        };
154
155        if let Some(mask_val) = mask_lit.as_bool().value() {
156            if mask_val {
157                return Ok(Some(expr.child(0).clone()));
158            } else {
159                return Ok(Some(expr.child(1).clone()));
160            }
161        }
162
163        Ok(None)
164    }
165
166    fn is_null_sensitive(&self, _options: &Self::Options) -> bool {
167        true
168    }
169
170    fn is_fallible(&self, _options: &Self::Options) -> bool {
171        false
172    }
173}
174
175pub(crate) fn zip_impl(
176    if_true: &ArrayRef,
177    if_false: &ArrayRef,
178    mask: &Mask,
179) -> VortexResult<ArrayRef> {
180    assert_eq!(
181        if_true.len(),
182        if_false.len(),
183        "zip requires arrays to have the same size"
184    );
185
186    let return_type = if_true
187        .dtype()
188        .clone()
189        .union_nullability(if_false.dtype().nullability());
190
191    if mask.all_true() {
192        return if_true.cast(return_type);
193    }
194    if mask.all_false() {
195        return if_false.cast(return_type);
196    }
197
198    zip_impl_with_builder(
199        if_true,
200        if_false,
201        mask.values()
202            .vortex_expect("zip_impl_with_builder: mask is not all-true or all-false"),
203        builder_with_capacity(&return_type, if_true.len()),
204    )
205}
206
207fn zip_impl_with_builder(
208    if_true: &ArrayRef,
209    if_false: &ArrayRef,
210    mask: &MaskValues,
211    mut builder: Box<dyn ArrayBuilder>,
212) -> VortexResult<ArrayRef> {
213    for (start, end) in mask.slices() {
214        builder.extend_from_array(&if_false.slice(builder.len()..*start)?);
215        builder.extend_from_array(&if_true.slice(*start..*end)?);
216    }
217    if builder.len() < if_false.len() {
218        builder.extend_from_array(&if_false.slice(builder.len()..if_false.len())?);
219    }
220    Ok(builder.finish())
221}
222
223#[cfg(test)]
224mod tests {
225    use arrow_array::cast::AsArray;
226    use arrow_select::zip::zip as arrow_zip;
227    use vortex_buffer::buffer;
228    use vortex_error::VortexResult;
229    use vortex_mask::Mask;
230
231    use super::zip_impl;
232    use crate::ArrayRef;
233    use crate::DynArray;
234    use crate::IntoArray;
235    use crate::LEGACY_SESSION;
236    use crate::VortexSessionExecute;
237    use crate::arrays::ConstantArray;
238    use crate::arrays::PrimitiveArray;
239    use crate::arrays::Struct;
240    use crate::arrays::StructArray;
241    use crate::arrays::VarBinView;
242    use crate::arrow::IntoArrowArray;
243    use crate::assert_arrays_eq;
244    use crate::builders::ArrayBuilder;
245    use crate::builders::BufferGrowthStrategy;
246    use crate::builders::VarBinViewBuilder;
247    use crate::builtins::ArrayBuiltins;
248    use crate::columnar::Columnar;
249    use crate::dtype::DType;
250    use crate::dtype::Nullability;
251    use crate::dtype::PType;
252    use crate::expr::lit;
253    use crate::expr::root;
254    use crate::expr::zip_expr;
255    use crate::scalar::Scalar;
256
257    #[test]
258    fn dtype() {
259        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
260        let expr = zip_expr(lit(true), root(), lit(0i32));
261        let result_dtype = expr.return_dtype(&dtype).unwrap();
262        assert_eq!(
263            result_dtype,
264            DType::Primitive(PType::I32, Nullability::NonNullable)
265        );
266    }
267
268    #[test]
269    fn test_display() {
270        let expr = zip_expr(lit(true), root(), lit(0i32));
271        assert_eq!(expr.to_string(), "zip($, 0i32, true)");
272    }
273
274    #[test]
275    fn test_zip_basic() {
276        let mask = Mask::from_iter([true, false, false, true, false]);
277        let if_true = buffer![10, 20, 30, 40, 50].into_array();
278        let if_false = buffer![1, 2, 3, 4, 5].into_array();
279
280        let result = mask.into_array().zip(if_true, if_false).unwrap();
281        let expected = buffer![10, 2, 3, 40, 5].into_array();
282
283        assert_arrays_eq!(result, expected);
284    }
285
286    #[test]
287    fn test_zip_all_true() {
288        let mask = Mask::new_true(4);
289        let if_true = buffer![10, 20, 30, 40].into_array();
290        let if_false =
291            PrimitiveArray::from_option_iter([Some(1), Some(2), Some(3), None]).into_array();
292
293        let result = mask.into_array().zip(if_true, if_false.clone()).unwrap();
294        let expected =
295            PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), Some(40)]).into_array();
296
297        assert_arrays_eq!(result, expected);
298        assert_eq!(result.dtype(), if_false.dtype())
299    }
300
301    #[test]
302    fn test_zip_all_false_widens_nullability() {
303        let mask = Mask::new_false(4);
304        let if_true =
305            PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), None]).into_array();
306        let if_false = buffer![1i32, 2, 3, 4].into_array();
307
308        let result = mask.into_array().zip(if_true.clone(), if_false).unwrap();
309        let expected =
310            PrimitiveArray::from_option_iter([Some(1), Some(2), Some(3), Some(4)]).into_array();
311
312        assert_arrays_eq!(result, expected);
313        assert_eq!(result.dtype(), if_true.dtype());
314    }
315
316    #[test]
317    fn test_zip_impl_all_true_widens_nullability() -> VortexResult<()> {
318        let mask = Mask::new_true(4);
319        let if_true = buffer![10i32, 20, 30, 40].into_array();
320        let if_false =
321            PrimitiveArray::from_option_iter([Some(1), Some(2), Some(3), None]).into_array();
322
323        let result = zip_impl(&if_true, &if_false, &mask)?;
324        assert_arrays_eq!(
325            result,
326            PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30), Some(40)])
327                .into_array()
328        );
329        assert_eq!(result.dtype(), if_false.dtype());
330        Ok(())
331    }
332
333    #[test]
334    fn test_zip_impl_all_false_widens_nullability() -> VortexResult<()> {
335        let mask = Mask::new_false(4);
336        let if_true =
337            PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), None]).into_array();
338        let if_false = buffer![1i32, 2, 3, 4].into_array();
339
340        let result = zip_impl(&if_true, &if_false, &mask)?;
341        assert_arrays_eq!(
342            result,
343            PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]).into_array()
344        );
345        assert_eq!(result.dtype(), if_true.dtype());
346        Ok(())
347    }
348
349    #[test]
350    #[should_panic]
351    fn test_invalid_lengths() {
352        let mask = Mask::new_false(4);
353        let if_true = buffer![10, 20, 30].into_array();
354        let if_false = buffer![1, 2, 3, 4].into_array();
355
356        let _result = mask.into_array().zip(if_true, if_false).unwrap();
357    }
358
359    #[test]
360    fn test_fragmentation() -> VortexResult<()> {
361        let len = 100;
362
363        let const1 = ConstantArray::new(
364            Scalar::utf8("hello_this_is_a_longer_string", Nullability::Nullable),
365            len,
366        )
367        .into_array();
368
369        let const2 = ConstantArray::new(
370            Scalar::utf8("world_this_is_another_string", Nullability::Nullable),
371            len,
372        )
373        .into_array();
374
375        let indices: Vec<usize> = (0..len).step_by(2).collect();
376        let mask = Mask::from_indices(len, indices);
377        let mask_array = mask.into_array();
378
379        let mut ctx = LEGACY_SESSION.create_execution_ctx();
380        let result = mask_array
381            .clone()
382            .zip(const1.clone(), const2.clone())?
383            .execute::<Columnar>(&mut ctx)?
384            .into_array();
385
386        insta::assert_snapshot!(result.display_tree(), @r"
387        root: vortex.varbinview(utf8?, len=100) nbytes=1.66 kB (100.00%) [all_valid]
388          metadata: EmptyMetadata
389          buffer: buffer_0 host 29 B (align=1) (1.75%)
390          buffer: buffer_1 host 28 B (align=1) (1.69%)
391          buffer: views host 1.60 kB (align=16) (96.56%)
392        ");
393
394        let wrapped1 = StructArray::try_from_iter([("nested", const1)])?.into_array();
395        let wrapped2 = StructArray::try_from_iter([("nested", const2)])?.into_array();
396
397        let wrapped_result = mask_array
398            .zip(wrapped1, wrapped2)?
399            .execute::<ArrayRef>(&mut ctx)?;
400        assert!(wrapped_result.is::<Struct>());
401
402        Ok(())
403    }
404
405    #[test]
406    fn test_varbinview_zip() {
407        let if_true = {
408            let mut builder = VarBinViewBuilder::new(
409                DType::Utf8(Nullability::NonNullable),
410                10,
411                Default::default(),
412                BufferGrowthStrategy::fixed(64 * 1024),
413                0.0,
414            );
415            for _ in 0..100 {
416                builder.append_value("Hello");
417                builder.append_value("Hello this is a long string that won't be inlined.");
418            }
419            builder.finish()
420        };
421
422        let if_false = {
423            let mut builder = VarBinViewBuilder::new(
424                DType::Utf8(Nullability::NonNullable),
425                10,
426                Default::default(),
427                BufferGrowthStrategy::fixed(64 * 1024),
428                0.0,
429            );
430            for _ in 0..100 {
431                builder.append_value("Hello2");
432                builder.append_value("Hello2 this is a long string that won't be inlined.");
433            }
434            builder.finish()
435        };
436
437        let mask = Mask::from_indices(200, (0..100).filter(|i| i % 3 != 0).collect());
438        let mask_array = mask.clone().into_array();
439
440        let mut ctx = LEGACY_SESSION.create_execution_ctx();
441        let zipped = mask_array
442            .zip(if_true.clone(), if_false.clone())
443            .unwrap()
444            .execute::<ArrayRef>(&mut ctx)
445            .unwrap();
446        let zipped = zipped.as_opt::<VarBinView>().unwrap();
447        assert_eq!(zipped.nbuffers(), 2);
448
449        let expected = arrow_zip(
450            mask.into_array()
451                .into_arrow_preferred()
452                .unwrap()
453                .as_boolean(),
454            &if_true.into_arrow_preferred().unwrap(),
455            &if_false.into_arrow_preferred().unwrap(),
456        )
457        .unwrap();
458
459        let actual = zipped.clone().into_array().into_arrow_preferred().unwrap();
460        assert_eq!(actual.as_ref(), expected.as_ref());
461    }
462}
vortex_array/scalar_fn/fns/zip/mod.rs

vortex_array/scalar_fn/fns/zip/
mod.rs
