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/// Check whether the string value at index `i` is null for string or
172/// dictionary-encoded string arrays.
173///
174/// Returns `true` when the value is null or the array type is not a string
175/// type, which corresponds to [`string_array_value_at_index`] returning `None`.
176///
177/// # Panics
178///
179/// If index `i` is out of bounds.
180pub fn is_string_null_at(array: &ArrayRef, i: usize) -> bool {
181    match array.data_type() {
182        DataType::Utf8 => {
183            let array = array.as_string::<i32>();
184            !array.is_valid(i)
185        }
186        DataType::LargeUtf8 => {
187            let array = array.as_string::<i64>();
188            !array.is_valid(i)
189        }
190        DataType::Utf8View => {
191            let array = array.as_string_view();
192            !array.is_valid(i)
193        }
194        DataType::Dictionary(key_type, value_type)
195            if key_type.is_integer() && value_type.is_string() =>
196        {
197            downcast_dictionary_array! {
198                array => array
199                    .key(i)
200                    .is_none_or(|key| is_string_null_at(array.values(), key)),
201                _ => true,
202            }
203        }
204        _ => true,
205    }
206}
207
208/// Get the string value at index `i` for `Utf8`, `LargeUtf8`, or `Utf8View` arrays.
209///
210/// Note: This method does not check for nulls and the value is arbitrary
211/// if [`is_null`](arrow::array::Array::is_null) returns true for the index.
212///
213/// # Panics
214/// 1. if index `i` is out of bounds;
215/// 2. or the array is not a string type.
216pub fn string_array_value(array: &ArrayRef, i: usize) -> &str {
217    match array.data_type() {
218        DataType::Utf8 => array.as_string::<i32>().value(i),
219        DataType::LargeUtf8 => array.as_string::<i64>().value(i),
220        DataType::Utf8View => array.as_string_view().value(i),
221        _ => unreachable!(),
222    }
223}
224
225/// Get the binary value at index `i` for `Binary`, `LargeBinary`, or `BinaryView` arrays.
226///
227/// Note: This method does not check for nulls and the value is arbitrary
228/// if [`is_null`](arrow::array::Array::is_null) returns true for the index.
229///
230/// # Panics
231/// 1. if index `i` is out of bounds;
232/// 2. or the array is not a binary type.
233pub fn binary_array_value(array: &ArrayRef, i: usize) -> &[u8] {
234    match array.data_type() {
235        DataType::Binary => array.as_binary::<i32>().value(i),
236        DataType::LargeBinary => array.as_binary::<i64>().value(i),
237        DataType::BinaryView => array.as_binary_view().value(i),
238        _ => unreachable!(),
239    }
240}
241
242/// Get the integer value (`i64`) at index `i` for any integer array.
243///
244/// Returns `None` when:
245///
246/// - the array type is not an integer type;
247/// - the value is larger than `i64::MAX`;
248/// - the value is null.
249///
250/// # Panics
251///
252/// If index `i` is out of bounds.
253pub fn int_array_value_at_index(array: &ArrayRef, i: usize) -> Option<i64> {
254    match array.data_type() {
255        DataType::Int8 => {
256            let array = array.as_primitive::<Int8Type>();
257            array.is_valid(i).then(|| array.value(i) as i64)
258        }
259        DataType::Int16 => {
260            let array = array.as_primitive::<Int16Type>();
261            array.is_valid(i).then(|| array.value(i) as i64)
262        }
263        DataType::Int32 => {
264            let array = array.as_primitive::<Int32Type>();
265            array.is_valid(i).then(|| array.value(i) as i64)
266        }
267        DataType::Int64 => {
268            let array = array.as_primitive::<Int64Type>();
269            array.is_valid(i).then(|| array.value(i))
270        }
271        DataType::UInt8 => {
272            let array = array.as_primitive::<UInt8Type>();
273            array.is_valid(i).then(|| array.value(i) as i64)
274        }
275        DataType::UInt16 => {
276            let array = array.as_primitive::<UInt16Type>();
277            array.is_valid(i).then(|| array.value(i) as i64)
278        }
279        DataType::UInt32 => {
280            let array = array.as_primitive::<UInt32Type>();
281            array.is_valid(i).then(|| array.value(i) as i64)
282        }
283        DataType::UInt64 => {
284            let array = array.as_primitive::<UInt64Type>();
285            array
286                .is_valid(i)
287                .then(|| {
288                    let i = array.value(i);
289                    if i <= i64::MAX as u64 {
290                        Some(i as i64)
291                    } else {
292                        None
293                    }
294                })
295                .flatten()
296        }
297        _ => None,
298    }
299}
300
301#[cfg(test)]
302mod tests {
303    use std::sync::Arc;
304
305    use arrow::array::StringDictionaryBuilder;
306    use arrow::datatypes::ArrowDictionaryKeyType;
307
308    use super::*;
309
310    fn assert_dictionary_key_type<K: ArrowDictionaryKeyType>() {
311        let mut builder = StringDictionaryBuilder::<K>::new();
312        builder.append("foo").unwrap();
313        builder.append("bar").unwrap();
314        builder.append("foo").unwrap();
315        builder.append_null();
316        let array: ArrayRef = Arc::new(builder.finish());
317
318        assert_eq!(Some("foo"), string_array_value_at_index(&array, 0));
319        assert_eq!(Some("bar"), string_array_value_at_index(&array, 1));
320        assert_eq!(Some("foo"), string_array_value_at_index(&array, 2));
321        assert_eq!(None, string_array_value_at_index(&array, 3));
322    }
323
324    #[test]
325    fn reads_dictionary_encoded_strings_with_all_key_types() {
326        assert_dictionary_key_type::<Int8Type>();
327        assert_dictionary_key_type::<Int16Type>();
328        assert_dictionary_key_type::<Int32Type>();
329        assert_dictionary_key_type::<Int64Type>();
330        assert_dictionary_key_type::<UInt8Type>();
331        assert_dictionary_key_type::<UInt16Type>();
332        assert_dictionary_key_type::<UInt32Type>();
333        assert_dictionary_key_type::<UInt64Type>();
334    }
335}
datatypes/arrow_array.rs

datatypes/
arrow_array.rs
