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::ExecutionCtx;
18use crate::IntoArray;
19use crate::arrays::BoolArray;
20use crate::arrays::bool::BoolArrayExt;
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::IntoArray;
234    use crate::LEGACY_SESSION;
235    use crate::VortexSessionExecute;
236    use crate::arrays::ConstantArray;
237    use crate::arrays::PrimitiveArray;
238    use crate::arrays::Struct;
239    use crate::arrays::StructArray;
240    use crate::arrays::VarBinView;
241    use crate::arrow::IntoArrowArray;
242    use crate::assert_arrays_eq;
243    use crate::builders::ArrayBuilder;
244    use crate::builders::BufferGrowthStrategy;
245    use crate::builders::VarBinViewBuilder;
246    use crate::builtins::ArrayBuiltins;
247    use crate::columnar::Columnar;
248    use crate::dtype::DType;
249    use crate::dtype::Nullability;
250    use crate::dtype::PType;
251    use crate::expr::lit;
252    use crate::expr::root;
253    use crate::expr::zip_expr;
254    use crate::scalar::Scalar;
255
256    #[test]
257    fn dtype() {
258        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
259        let expr = zip_expr(lit(true), root(), lit(0i32));
260        let result_dtype = expr.return_dtype(&dtype).unwrap();
261        assert_eq!(
262            result_dtype,
263            DType::Primitive(PType::I32, Nullability::NonNullable)
264        );
265    }
266
267    #[test]
268    fn test_display() {
269        let expr = zip_expr(lit(true), root(), lit(0i32));
270        assert_eq!(expr.to_string(), "zip($, 0i32, true)");
271    }
272
273    #[test]
274    fn test_zip_basic() {
275        let mask = Mask::from_iter([true, false, false, true, false]);
276        let if_true = buffer![10, 20, 30, 40, 50].into_array();
277        let if_false = buffer![1, 2, 3, 4, 5].into_array();
278
279        let result = mask.into_array().zip(if_true, if_false).unwrap();
280        let expected = buffer![10, 2, 3, 40, 5].into_array();
281
282        assert_arrays_eq!(result, expected);
283    }
284
285    #[test]
286    fn test_zip_all_true() {
287        let mask = Mask::new_true(4);
288        let if_true = buffer![10, 20, 30, 40].into_array();
289        let if_false =
290            PrimitiveArray::from_option_iter([Some(1), Some(2), Some(3), None]).into_array();
291
292        let result = mask.into_array().zip(if_true, if_false.clone()).unwrap();
293        let expected =
294            PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), Some(40)]).into_array();
295
296        assert_arrays_eq!(result, expected);
297        assert_eq!(result.dtype(), if_false.dtype())
298    }
299
300    #[test]
301    fn test_zip_all_false_widens_nullability() {
302        let mask = Mask::new_false(4);
303        let if_true =
304            PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), None]).into_array();
305        let if_false = buffer![1i32, 2, 3, 4].into_array();
306
307        let result = mask.into_array().zip(if_true.clone(), if_false).unwrap();
308        let expected =
309            PrimitiveArray::from_option_iter([Some(1), Some(2), Some(3), Some(4)]).into_array();
310
311        assert_arrays_eq!(result, expected);
312        assert_eq!(result.dtype(), if_true.dtype());
313    }
314
315    #[test]
316    fn test_zip_impl_all_true_widens_nullability() -> VortexResult<()> {
317        let mask = Mask::new_true(4);
318        let if_true = buffer![10i32, 20, 30, 40].into_array();
319        let if_false =
320            PrimitiveArray::from_option_iter([Some(1), Some(2), Some(3), None]).into_array();
321
322        let result = zip_impl(&if_true, &if_false, &mask)?;
323        assert_arrays_eq!(
324            result,
325            PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30), Some(40)])
326                .into_array()
327        );
328        assert_eq!(result.dtype(), if_false.dtype());
329        Ok(())
330    }
331
332    #[test]
333    fn test_zip_impl_all_false_widens_nullability() -> VortexResult<()> {
334        let mask = Mask::new_false(4);
335        let if_true =
336            PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), None]).into_array();
337        let if_false = buffer![1i32, 2, 3, 4].into_array();
338
339        let result = zip_impl(&if_true, &if_false, &mask)?;
340        assert_arrays_eq!(
341            result,
342            PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]).into_array()
343        );
344        assert_eq!(result.dtype(), if_true.dtype());
345        Ok(())
346    }
347
348    #[test]
349    #[should_panic]
350    fn test_invalid_lengths() {
351        let mask = Mask::new_false(4);
352        let if_true = buffer![10, 20, 30].into_array();
353        let if_false = buffer![1, 2, 3, 4].into_array();
354
355        let _result = mask.into_array().zip(if_true, if_false).unwrap();
356    }
357
358    #[test]
359    fn test_fragmentation() -> VortexResult<()> {
360        let len = 100;
361
362        let const1 = ConstantArray::new(
363            Scalar::utf8("hello_this_is_a_longer_string", Nullability::Nullable),
364            len,
365        )
366        .into_array();
367
368        let const2 = ConstantArray::new(
369            Scalar::utf8("world_this_is_another_string", Nullability::Nullable),
370            len,
371        )
372        .into_array();
373
374        let indices: Vec<usize> = (0..len).step_by(2).collect();
375        let mask = Mask::from_indices(len, indices);
376        let mask_array = mask.into_array();
377
378        let mut ctx = LEGACY_SESSION.create_execution_ctx();
379        let result = mask_array
380            .zip(const1.clone(), const2.clone())?
381            .execute::<Columnar>(&mut ctx)?
382            .into_array();
383
384        insta::assert_snapshot!(result.display_tree(), @r"
385        root: vortex.varbinview(utf8?, len=100) nbytes=1.66 kB (100.00%) [all_valid]
386          metadata: 
387          buffer: buffer_0 host 29 B (align=1) (1.75%)
388          buffer: buffer_1 host 28 B (align=1) (1.69%)
389          buffer: views host 1.60 kB (align=16) (96.56%)
390        ");
391
392        let wrapped1 = StructArray::try_from_iter([("nested", const1)])?.into_array();
393        let wrapped2 = StructArray::try_from_iter([("nested", const2)])?.into_array();
394
395        let wrapped_result = mask_array
396            .zip(wrapped1, wrapped2)?
397            .execute::<ArrayRef>(&mut ctx)?;
398        assert!(wrapped_result.is::<Struct>());
399
400        Ok(())
401    }
402
403    #[test]
404    fn test_varbinview_zip() {
405        let if_true = {
406            let mut builder = VarBinViewBuilder::new(
407                DType::Utf8(Nullability::NonNullable),
408                10,
409                Default::default(),
410                BufferGrowthStrategy::fixed(64 * 1024),
411                0.0,
412            );
413            for _ in 0..100 {
414                builder.append_value("Hello");
415                builder.append_value("Hello this is a long string that won't be inlined.");
416            }
417            builder.finish()
418        };
419
420        let if_false = {
421            let mut builder = VarBinViewBuilder::new(
422                DType::Utf8(Nullability::NonNullable),
423                10,
424                Default::default(),
425                BufferGrowthStrategy::fixed(64 * 1024),
426                0.0,
427            );
428            for _ in 0..100 {
429                builder.append_value("Hello2");
430                builder.append_value("Hello2 this is a long string that won't be inlined.");
431            }
432            builder.finish()
433        };
434
435        let mask = Mask::from_indices(200, (0..100).filter(|i| i % 3 != 0).collect());
436        let mask_array = mask.clone().into_array();
437
438        let mut ctx = LEGACY_SESSION.create_execution_ctx();
439        let zipped = mask_array
440            .zip(if_true.clone(), if_false.clone())
441            .unwrap()
442            .execute::<ArrayRef>(&mut ctx)
443            .unwrap();
444        let zipped = zipped.as_opt::<VarBinView>().unwrap();
445        assert_eq!(zipped.data_buffers().len(), 2);
446
447        let expected = arrow_zip(
448            mask.into_array()
449                .into_arrow_preferred()
450                .unwrap()
451                .as_boolean(),
452            &if_true.into_arrow_preferred().unwrap(),
453            &if_false.into_arrow_preferred().unwrap(),
454        )
455        .unwrap();
456
457        let actual = zipped.array().clone().into_arrow_preferred().unwrap();
458        assert_eq!(actual.as_ref(), expected.as_ref());
459    }
460}
vortex_array/scalar_fn/fns/zip/mod.rs

vortex_array/scalar_fn/fns/zip/
mod.rs
