Skip to main content

dbx_core/api/
traits.rs

1//! API 트레이트 정의
2
3use crate::error::DbxResult;
4use arrow::array::{Array, ArrayRef, RecordBatch};
5use arrow::datatypes::DataType;
6
7/// RecordBatch에서 구조체로 변환하는 트레이트
8pub trait FromRow: Sized {
9    fn from_row(batch: &RecordBatch, row_idx: usize) -> DbxResult<Self>;
10}
11
12/// Rust 타입을 Arrow DataType으로 변환하는 트레이트
13pub trait IntoArrowType {
14    fn arrow_type() -> DataType;
15    fn is_nullable() -> bool;
16}
17
18/// Arrow Array에서 Rust 타입으로 변환하는 트레이트
19pub trait FromColumn: Sized {
20    fn from_column(column: &ArrayRef, row_idx: usize) -> DbxResult<Self>;
21}
22
23// 기본 타입 구현
24impl IntoArrowType for i32 {
25    fn arrow_type() -> DataType {
26        DataType::Int32
27    }
28    fn is_nullable() -> bool {
29        false
30    }
31}
32
33impl FromColumn for i32 {
34    fn from_column(column: &ArrayRef, row_idx: usize) -> DbxResult<Self> {
35        use arrow::array::AsArray;
36        Ok(column
37            .as_primitive::<arrow::datatypes::Int32Type>()
38            .value(row_idx))
39    }
40}
41
42impl IntoArrowType for i64 {
43    fn arrow_type() -> DataType {
44        DataType::Int64
45    }
46    fn is_nullable() -> bool {
47        false
48    }
49}
50
51impl FromColumn for i64 {
52    fn from_column(column: &ArrayRef, row_idx: usize) -> DbxResult<Self> {
53        use arrow::array::AsArray;
54        Ok(column
55            .as_primitive::<arrow::datatypes::Int64Type>()
56            .value(row_idx))
57    }
58}
59
60impl IntoArrowType for f64 {
61    fn arrow_type() -> DataType {
62        DataType::Float64
63    }
64    fn is_nullable() -> bool {
65        false
66    }
67}
68
69impl FromColumn for f64 {
70    fn from_column(column: &ArrayRef, row_idx: usize) -> DbxResult<Self> {
71        use arrow::array::AsArray;
72        Ok(column
73            .as_primitive::<arrow::datatypes::Float64Type>()
74            .value(row_idx))
75    }
76}
77
78impl IntoArrowType for String {
79    fn arrow_type() -> DataType {
80        DataType::Utf8
81    }
82    fn is_nullable() -> bool {
83        false
84    }
85}
86
87impl FromColumn for String {
88    fn from_column(column: &ArrayRef, row_idx: usize) -> DbxResult<Self> {
89        use arrow::array::AsArray;
90        Ok(column.as_string::<i32>().value(row_idx).to_string())
91    }
92}
93
94impl IntoArrowType for bool {
95    fn arrow_type() -> DataType {
96        DataType::Boolean
97    }
98    fn is_nullable() -> bool {
99        false
100    }
101}
102
103impl FromColumn for bool {
104    fn from_column(column: &ArrayRef, row_idx: usize) -> DbxResult<Self> {
105        use arrow::array::AsArray;
106        Ok(column.as_boolean().value(row_idx))
107    }
108}
109
110// Option<T> 구현
111impl<T: IntoArrowType> IntoArrowType for Option<T> {
112    fn arrow_type() -> DataType {
113        T::arrow_type()
114    }
115    fn is_nullable() -> bool {
116        true
117    }
118}
119
120impl<T: FromColumn> FromColumn for Option<T> {
121    fn from_column(column: &ArrayRef, row_idx: usize) -> DbxResult<Self> {
122        if column.is_null(row_idx) {
123            Ok(None)
124        } else {
125            Ok(Some(T::from_column(column, row_idx)?))
126        }
127    }
128}