forked from jorgecarleitao/arrow2
-
Notifications
You must be signed in to change notification settings - Fork 0
/
length.rs
88 lines (80 loc) · 2.04 KB
/
length.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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
use arrow2::array::*;
use arrow2::compute::length::*;
use arrow2::datatypes::*;
fn length_test_string<O: Offset>() {
vec![
(
vec![Some("hello"), Some(" "), None],
vec![Some(5usize), Some(1), None],
),
(vec![Some("💖")], vec![Some(4)]),
]
.into_iter()
.for_each(|(input, expected)| {
let array = Utf8Array::<O>::from(&input);
let result = length(&array).unwrap();
let data_type = if O::IS_LARGE {
DataType::Int64
} else {
DataType::Int32
};
let expected = expected
.into_iter()
.map(|x| x.map(|x| O::from_usize(x).unwrap()))
.collect::<PrimitiveArray<O>>()
.to(data_type);
assert_eq!(expected, result.as_ref());
})
}
#[test]
fn large_utf8() {
length_test_string::<i64>()
}
#[test]
fn utf8() {
length_test_string::<i32>()
}
#[test]
fn consistency() {
use arrow2::datatypes::DataType::*;
let datatypes = vec![
Null,
Boolean,
UInt8,
UInt16,
UInt32,
UInt64,
Int8,
Int16,
Int32,
Int64,
Float32,
Float64,
Timestamp(TimeUnit::Second, None),
Timestamp(TimeUnit::Millisecond, None),
Timestamp(TimeUnit::Microsecond, None),
Timestamp(TimeUnit::Nanosecond, None),
Time64(TimeUnit::Microsecond),
Time64(TimeUnit::Nanosecond),
Date32,
Time32(TimeUnit::Second),
Time32(TimeUnit::Millisecond),
Date64,
Utf8,
LargeUtf8,
Binary,
LargeBinary,
Duration(TimeUnit::Second),
Duration(TimeUnit::Millisecond),
Duration(TimeUnit::Microsecond),
Duration(TimeUnit::Nanosecond),
];
datatypes.into_iter().for_each(|d1| {
let array = new_null_array(d1.clone(), 10);
if can_length(&d1) {
assert!(length(array.as_ref()).is_ok());
} else {
assert!(length(array.as_ref()).is_err());
}
});
}