risingwave_common/array/arrow/
mod.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
// Copyright 2024 RisingWave Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// These mods imports arrow_impl.rs to provide FromArrow, ToArrow traits for corresponding arrow versions,
// and the default From/To implementations.

mod arrow_52;
mod arrow_53;
// These mods import mods above and may override some methods.
mod arrow_deltalake;
mod arrow_iceberg;
mod arrow_udf;

pub use arrow_deltalake::DeltaLakeConvert;
pub use arrow_iceberg::{IcebergArrowConvert, IcebergCreateTableArrowConvert};
pub use arrow_udf::UdfArrowConvert;
pub use reexport::*;
/// For other RisingWave crates, they can directly use arrow re-exported here, without adding
/// `arrow` dependencies in their `Cargo.toml`. And they don't need to care about the version.
mod reexport {
    pub use super::arrow_deltalake::{
        arrow_array as arrow_array_deltalake, arrow_buffer as arrow_buffer_deltalake,
        arrow_cast as arrow_cast_deltalake, arrow_schema as arrow_schema_deltalake,
        FromArrow as DeltaLakeFromArrow, ToArrow as DeltaLakeToArrow,
    };
    pub use super::arrow_iceberg::{
        arrow_array as arrow_array_iceberg, arrow_buffer as arrow_buffer_iceberg,
        arrow_cast as arrow_cast_iceberg, arrow_schema as arrow_schema_iceberg,
        FromArrow as IcebergFromArrow, ToArrow as IcebergToArrow,
    };
    pub use super::arrow_udf::{
        arrow_array as arrow_array_udf, arrow_buffer as arrow_buffer_udf,
        arrow_cast as arrow_cast_udf, arrow_schema as arrow_schema_udf, FromArrow as UdfFromArrow,
        ToArrow as UdfToArrow,
    };
}
use crate::types::Interval;

/// Arrow 52 changed the interval type from `i128` to `arrow_buffer::IntervalMonthDayNano`, so
/// we introduced this trait to customize the conversion in `arrow_impl.rs`.
/// We may delete this after all arrow versions are upgraded.
trait ArrowIntervalTypeTrait {
    fn to_interval(self) -> Interval;
    fn from_interval(value: Interval) -> Self;
}

impl ArrowIntervalTypeTrait for i128 {
    fn to_interval(self) -> Interval {
        // XXX: the arrow-rs decoding is incorrect
        // let (months, days, ns) = arrow_array::types::IntervalMonthDayNanoType::to_parts(value);
        let months = self as i32;
        let days = (self >> 32) as i32;
        let ns = (self >> 64) as i64;
        Interval::from_month_day_usec(months, days, ns / 1000)
    }

    fn from_interval(value: Interval) -> i128 {
        // XXX: the arrow-rs encoding is incorrect
        // arrow_array::types::IntervalMonthDayNanoType::make_value(
        //     self.months(),
        //     self.days(),
        //     // TODO: this may overflow and we need `try_into`
        //     self.usecs() * 1000,
        // )
        let m = value.months() as u128 & u32::MAX as u128;
        let d = (value.days() as u128 & u32::MAX as u128) << 32;
        let n = ((value.usecs() * 1000) as u128 & u64::MAX as u128) << 64;
        (m | d | n) as i128
    }
}