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, Time32MillisecondType, Time32SecondType, Time64MicrosecondType,
19    Time64NanosecondType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType,
20    TimestampNanosecondType, TimestampSecondType,
21};
22use common_time::time::Time;
23use common_time::{Duration, Timestamp};
24
25pub type BinaryArray = arrow::array::BinaryArray;
26pub type MutableBinaryArray = arrow::array::BinaryBuilder;
27pub type StringArray = arrow::array::StringArray;
28pub type MutableStringArray = arrow::array::StringBuilder;
29pub type LargeStringArray = arrow::array::LargeStringArray;
30pub type MutableLargeStringArray = arrow::array::LargeStringBuilder;
31
32/// Get the [Timestamp] value at index `i` of the timestamp array.
33///
34/// Note: This method does not check for nulls and the value is arbitrary
35/// if [`is_null`](arrow::array::Array::is_null) returns true for the index.
36///
37/// # Panics
38/// 1. if index `i` is out of bounds;
39/// 2. or the array is not timestamp type.
40pub fn timestamp_array_value(array: &ArrayRef, i: usize) -> Timestamp {
41    let DataType::Timestamp(time_unit, _) = &array.data_type() else {
42        unreachable!()
43    };
44    let v = match time_unit {
45        TimeUnit::Second => {
46            let array = array.as_primitive::<TimestampSecondType>();
47            array.value(i)
48        }
49        TimeUnit::Millisecond => {
50            let array = array.as_primitive::<TimestampMillisecondType>();
51            array.value(i)
52        }
53        TimeUnit::Microsecond => {
54            let array = array.as_primitive::<TimestampMicrosecondType>();
55            array.value(i)
56        }
57        TimeUnit::Nanosecond => {
58            let array = array.as_primitive::<TimestampNanosecondType>();
59            array.value(i)
60        }
61    };
62    Timestamp::new(v, time_unit.into())
63}
64
65/// Get the [Time] value at index `i` of the time array.
66///
67/// Note: This method does not check for nulls and the value is arbitrary
68/// if [`is_null`](arrow::array::Array::is_null) returns true for the index.
69///
70/// # Panics
71/// 1. if index `i` is out of bounds;
72/// 2. or the array is not `Time32` or `Time64` type.
73pub fn time_array_value(array: &ArrayRef, i: usize) -> Time {
74    match array.data_type() {
75        DataType::Time32(time_unit) | DataType::Time64(time_unit) => match time_unit {
76            TimeUnit::Second => {
77                let array = array.as_primitive::<Time32SecondType>();
78                Time::new_second(array.value(i) as i64)
79            }
80            TimeUnit::Millisecond => {
81                let array = array.as_primitive::<Time32MillisecondType>();
82                Time::new_millisecond(array.value(i) as i64)
83            }
84            TimeUnit::Microsecond => {
85                let array = array.as_primitive::<Time64MicrosecondType>();
86                Time::new_microsecond(array.value(i))
87            }
88            TimeUnit::Nanosecond => {
89                let array = array.as_primitive::<Time64NanosecondType>();
90                Time::new_nanosecond(array.value(i))
91            }
92        },
93        _ => unreachable!(),
94    }
95}
96
97/// Get the [Duration] value at index `i` of the duration array.
98///
99/// Note: This method does not check for nulls and the value is arbitrary
100/// if [`is_null`](arrow::array::Array::is_null) returns true for the index.
101///
102/// # Panics
103/// 1. if index `i` is out of bounds;
104/// 2. or the array is not duration type.
105pub fn duration_array_value(array: &ArrayRef, i: usize) -> Duration {
106    let DataType::Duration(time_unit) = array.data_type() else {
107        unreachable!();
108    };
109    let v = match time_unit {
110        TimeUnit::Second => {
111            let array = array.as_primitive::<DurationSecondType>();
112            array.value(i)
113        }
114        TimeUnit::Millisecond => {
115            let array = array.as_primitive::<DurationMillisecondType>();
116            array.value(i)
117        }
118        TimeUnit::Microsecond => {
119            let array = array.as_primitive::<DurationMicrosecondType>();
120            array.value(i)
121        }
122        TimeUnit::Nanosecond => {
123            let array = array.as_primitive::<DurationNanosecondType>();
124            array.value(i)
125        }
126    };
127    Duration::new(v, time_unit.into())
128}
datatypes/arrow_array.rs

datatypes/
arrow_array.rs
