datatypes/
arrow_array.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use arrow::array::{ArrayRef, AsArray};
16use arrow::datatypes::{
17    DataType, DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType,
18    DurationSecondType, Int8Type, Int16Type, Int32Type, Int64Type, Time32MillisecondType,
19    Time32SecondType, Time64MicrosecondType, Time64NanosecondType, TimeUnit,
20    TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
21    TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type,
22};
23use arrow::downcast_dictionary_array;
24use arrow_array::Array;
25use common_time::time::Time;
26use common_time::{Duration, Timestamp};
27
28pub type BinaryArray = arrow::array::BinaryArray;
29pub type LargeBinaryArray = arrow::array::LargeBinaryArray;
30pub type MutableBinaryArray = arrow::array::BinaryBuilder;
31pub type BinaryViewArray = arrow::array::BinaryViewArray;
32pub type MutableBinaryViewArray = arrow::array::BinaryViewBuilder;
33pub type StringArray = arrow::array::StringArray;
34pub type MutableStringArray = arrow::array::StringBuilder;
35pub type LargeStringArray = arrow::array::LargeStringArray;
36pub type MutableLargeStringArray = arrow::array::LargeStringBuilder;
37pub type StringViewArray = arrow::array::StringViewArray;
38pub type MutableStringViewArray = arrow::array::StringViewBuilder;
39
40/// Get the [Timestamp] value at index `i` of the timestamp array.
41///
42/// Note: This method does not check for nulls and the value is arbitrary
43/// if [`is_null`](arrow::array::Array::is_null) returns true for the index.
44///
45/// # Panics
46/// 1. if index `i` is out of bounds;
47/// 2. or the array is not timestamp type.
48pub fn timestamp_array_value(array: &ArrayRef, i: usize) -> Timestamp {
49    let DataType::Timestamp(time_unit, _) = &array.data_type() else {
50        unreachable!()
51    };
52    let v = match time_unit {
53        TimeUnit::Second => {
54            let array = array.as_primitive::<TimestampSecondType>();
55            array.value(i)
56        }
57        TimeUnit::Millisecond => {
58            let array = array.as_primitive::<TimestampMillisecondType>();
59            array.value(i)
60        }
61        TimeUnit::Microsecond => {
62            let array = array.as_primitive::<TimestampMicrosecondType>();
63            array.value(i)
64        }
65        TimeUnit::Nanosecond => {
66            let array = array.as_primitive::<TimestampNanosecondType>();
67            array.value(i)
68        }
69    };
70    Timestamp::new(v, time_unit.into())
71}
72
73/// Get the [Time] value at index `i` of the time array.
74///
75/// Note: This method does not check for nulls and the value is arbitrary
76/// if [`is_null`](arrow::array::Array::is_null) returns true for the index.
77///
78/// # Panics
79/// 1. if index `i` is out of bounds;
80/// 2. or the array is not `Time32` or `Time64` type.
81pub fn time_array_value(array: &ArrayRef, i: usize) -> Time {
82    match array.data_type() {
83        DataType::Time32(time_unit) | DataType::Time64(time_unit) => match time_unit {
84            TimeUnit::Second => {
85                let array = array.as_primitive::<Time32SecondType>();
86                Time::new_second(array.value(i) as i64)
87            }
88            TimeUnit::Millisecond => {
89                let array = array.as_primitive::<Time32MillisecondType>();
90                Time::new_millisecond(array.value(i) as i64)
91            }
92            TimeUnit::Microsecond => {
93                let array = array.as_primitive::<Time64MicrosecondType>();
94                Time::new_microsecond(array.value(i))
95            }
96            TimeUnit::Nanosecond => {
97                let array = array.as_primitive::<Time64NanosecondType>();
98                Time::new_nanosecond(array.value(i))
99            }
100        },
101        _ => unreachable!(),
102    }
103}
104
105/// Get the [Duration] value at index `i` of the duration array.
106///
107/// Note: This method does not check for nulls and the value is arbitrary
108/// if [`is_null`](arrow::array::Array::is_null) returns true for the index.
109///
110/// # Panics
111/// 1. if index `i` is out of bounds;
112/// 2. or the array is not duration type.
113pub fn duration_array_value(array: &ArrayRef, i: usize) -> Duration {
114    let DataType::Duration(time_unit) = array.data_type() else {
115        unreachable!();
116    };
117    let v = match time_unit {
118        TimeUnit::Second => {
119            let array = array.as_primitive::<DurationSecondType>();
120            array.value(i)
121        }
122        TimeUnit::Millisecond => {
123            let array = array.as_primitive::<DurationMillisecondType>();
124            array.value(i)
125        }
126        TimeUnit::Microsecond => {
127            let array = array.as_primitive::<DurationMicrosecondType>();
128            array.value(i)
129        }
130        TimeUnit::Nanosecond => {
131            let array = array.as_primitive::<DurationNanosecondType>();
132            array.value(i)
133        }
134    };
135    Duration::new(v, time_unit.into())
136}
137
138/// Get the string value at index `i` for string or dictionary-encoded string arrays.
139///
140/// Returns `None` when the array type is not a string type or the value is null.
141///
142/// # Panics
143///
144/// If index `i` is out of bounds.
145pub fn string_array_value_at_index(array: &ArrayRef, i: usize) -> Option<&str> {
146    match array.data_type() {
147        DataType::Utf8 => {
148            let array = array.as_string::<i32>();
149            array.is_valid(i).then(|| array.value(i))
150        }
151        DataType::LargeUtf8 => {
152            let array = array.as_string::<i64>();
153            array.is_valid(i).then(|| array.value(i))
154        }
155        DataType::Utf8View => {
156            let array = array.as_string_view();
157            array.is_valid(i).then(|| array.value(i))
158        }
159        DataType::Dictionary(key_type, value_type)
160            if key_type.is_integer() && value_type.is_string() =>
161        {
162            downcast_dictionary_array! {
163                array => string_array_value_at_index(array.values(), array.key(i)?),
164                _ => None,
165            }
166        }
167        _ => None,
168    }
169}
170
171/// Get the string value at index `i` for `Utf8`, `LargeUtf8`, or `Utf8View` arrays.
172///
173/// Note: This method does not check for nulls and the value is arbitrary
174/// if [`is_null`](arrow::array::Array::is_null) returns true for the index.
175///
176/// # Panics
177/// 1. if index `i` is out of bounds;
178/// 2. or the array is not a string type.
179pub fn string_array_value(array: &ArrayRef, i: usize) -> &str {
180    match array.data_type() {
181        DataType::Utf8 => array.as_string::<i32>().value(i),
182        DataType::LargeUtf8 => array.as_string::<i64>().value(i),
183        DataType::Utf8View => array.as_string_view().value(i),
184        _ => unreachable!(),
185    }
186}
187
188/// Get the binary value at index `i` for `Binary`, `LargeBinary`, or `BinaryView` arrays.
189///
190/// Note: This method does not check for nulls and the value is arbitrary
191/// if [`is_null`](arrow::array::Array::is_null) returns true for the index.
192///
193/// # Panics
194/// 1. if index `i` is out of bounds;
195/// 2. or the array is not a binary type.
196pub fn binary_array_value(array: &ArrayRef, i: usize) -> &[u8] {
197    match array.data_type() {
198        DataType::Binary => array.as_binary::<i32>().value(i),
199        DataType::LargeBinary => array.as_binary::<i64>().value(i),
200        DataType::BinaryView => array.as_binary_view().value(i),
201        _ => unreachable!(),
202    }
203}
204
205/// Get the integer value (`i64`) at index `i` for any integer array.
206///
207/// Returns `None` when:
208///
209/// - the array type is not an integer type;
210/// - the value is larger than `i64::MAX`;
211/// - the value is null.
212///
213/// # Panics
214///
215/// If index `i` is out of bounds.
216pub fn int_array_value_at_index(array: &ArrayRef, i: usize) -> Option<i64> {
217    match array.data_type() {
218        DataType::Int8 => {
219            let array = array.as_primitive::<Int8Type>();
220            array.is_valid(i).then(|| array.value(i) as i64)
221        }
222        DataType::Int16 => {
223            let array = array.as_primitive::<Int16Type>();
224            array.is_valid(i).then(|| array.value(i) as i64)
225        }
226        DataType::Int32 => {
227            let array = array.as_primitive::<Int32Type>();
228            array.is_valid(i).then(|| array.value(i) as i64)
229        }
230        DataType::Int64 => {
231            let array = array.as_primitive::<Int64Type>();
232            array.is_valid(i).then(|| array.value(i))
233        }
234        DataType::UInt8 => {
235            let array = array.as_primitive::<UInt8Type>();
236            array.is_valid(i).then(|| array.value(i) as i64)
237        }
238        DataType::UInt16 => {
239            let array = array.as_primitive::<UInt16Type>();
240            array.is_valid(i).then(|| array.value(i) as i64)
241        }
242        DataType::UInt32 => {
243            let array = array.as_primitive::<UInt32Type>();
244            array.is_valid(i).then(|| array.value(i) as i64)
245        }
246        DataType::UInt64 => {
247            let array = array.as_primitive::<UInt64Type>();
248            array
249                .is_valid(i)
250                .then(|| {
251                    let i = array.value(i);
252                    if i <= i64::MAX as u64 {
253                        Some(i as i64)
254                    } else {
255                        None
256                    }
257                })
258                .flatten()
259        }
260        _ => None,
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use std::sync::Arc;
267
268    use arrow::array::StringDictionaryBuilder;
269    use arrow::datatypes::ArrowDictionaryKeyType;
270
271    use super::*;
272
273    fn assert_dictionary_key_type<K: ArrowDictionaryKeyType>() {
274        let mut builder = StringDictionaryBuilder::<K>::new();
275        builder.append("foo").unwrap();
276        builder.append("bar").unwrap();
277        builder.append("foo").unwrap();
278        builder.append_null();
279        let array: ArrayRef = Arc::new(builder.finish());
280
281        assert_eq!(Some("foo"), string_array_value_at_index(&array, 0));
282        assert_eq!(Some("bar"), string_array_value_at_index(&array, 1));
283        assert_eq!(Some("foo"), string_array_value_at_index(&array, 2));
284        assert_eq!(None, string_array_value_at_index(&array, 3));
285    }
286
287    #[test]
288    fn reads_dictionary_encoded_strings_with_all_key_types() {
289        assert_dictionary_key_type::<Int8Type>();
290        assert_dictionary_key_type::<Int16Type>();
291        assert_dictionary_key_type::<Int32Type>();
292        assert_dictionary_key_type::<Int64Type>();
293        assert_dictionary_key_type::<UInt8Type>();
294        assert_dictionary_key_type::<UInt16Type>();
295        assert_dictionary_key_type::<UInt32Type>();
296        assert_dictionary_key_type::<UInt64Type>();
297    }
298}
datatypes/arrow_array.rs

datatypes/
arrow_array.rs
