This repository has been archived by the owner on Feb 18, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 221
/
union.rs
72 lines (58 loc) · 1.94 KB
/
union.rs
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
use std::sync::Arc;
use arrow2::{
array::{
growable::{Growable, GrowableUnion},
*,
},
datatypes::*,
error::Result,
};
#[test]
fn sparse() -> Result<()> {
let fields = vec![
Field::new("a", DataType::Int32, true),
Field::new("b", DataType::Utf8, true),
];
let data_type = DataType::Union(fields, None, UnionMode::Sparse);
let types = vec![0, 0, 1].into();
let fields = vec![
Arc::new(Int32Array::from(&[Some(1), None, Some(2)])) as Arc<dyn Array>,
Arc::new(Utf8Array::<i32>::from(&[Some("a"), Some("b"), Some("c")])) as Arc<dyn Array>,
];
let array = UnionArray::from_data(data_type, types, fields, None);
for length in 1..2 {
for index in 0..(array.len() - length + 1) {
let mut a = GrowableUnion::new(vec![&array], 10);
a.extend(0, index, length);
let expected = array.slice(index, length);
let result: UnionArray = a.into();
assert_eq!(expected, result);
}
}
Ok(())
}
#[test]
fn dense() -> Result<()> {
let fields = vec![
Field::new("a", DataType::Int32, true),
Field::new("b", DataType::Utf8, true),
];
let data_type = DataType::Union(fields, None, UnionMode::Dense);
let types = vec![0, 0, 1].into();
let fields = vec![
Arc::new(Int32Array::from(&[Some(1), None, Some(2)])) as Arc<dyn Array>,
Arc::new(Utf8Array::<i32>::from(&[Some("c")])) as Arc<dyn Array>,
];
let offsets = Some(vec![0, 1, 0].into());
let array = UnionArray::from_data(data_type, types, fields, offsets);
for length in 1..2 {
for index in 0..(array.len() - length + 1) {
let mut a = GrowableUnion::new(vec![&array], 10);
a.extend(0, index, length);
let expected = array.slice(index, length);
let result: UnionArray = a.into();
assert_eq!(expected, result);
}
}
Ok(())
}