1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
use crate::array::*;
use crate::compute::SlicesIterator;
use crate::error::{ArrowError, Result};
pub fn zip(
mask: &BooleanArray,
truthy: &dyn Array,
falsy: &dyn Array,
) -> Result<ArrayRef> {
if truthy.data_type() != falsy.data_type() {
return Err(ArrowError::InvalidArgumentError(
"arguments need to have the same data type".into(),
));
}
if truthy.len() != falsy.len() || falsy.len() != mask.len() {
return Err(ArrowError::InvalidArgumentError(
"all arrays should have the same length".into(),
));
}
let falsy = falsy.data();
let truthy = truthy.data();
let mut mutable = MutableArrayData::new(vec![&*truthy, &*falsy], false, truthy.len());
let mut filled = 0;
SlicesIterator::new(mask).for_each(|(start, end)| {
if start > filled {
mutable.extend(1, filled, start);
}
mutable.extend(0, start, end);
filled = end;
});
if filled < truthy.len() {
mutable.extend(1, filled, truthy.len());
}
let data = mutable.freeze();
Ok(make_array(data))
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test_zip_kernel() {
let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);
let b = Int32Array::from(vec![None, Some(3), Some(6), Some(7), Some(3)]);
let mask = BooleanArray::from(vec![true, true, false, false, true]);
let out = zip(&mask, &a, &b).unwrap();
let actual = out.as_any().downcast_ref::<Int32Array>().unwrap();
let expected = Int32Array::from(vec![Some(5), None, Some(6), Some(7), Some(1)]);
assert_eq!(actual, &expected);
}
}