vortex_scalar/arrow/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use arrow_array::*;
7use vortex_dtype::datetime::{TemporalMetadata, TimeUnit, is_temporal_ext_type};
8use vortex_dtype::{DType, PType};
9use vortex_error::{VortexError, vortex_bail, vortex_err};
10
11use crate::Scalar;
12use crate::decimal::DecimalValue;
13
14macro_rules! value_to_arrow_scalar {
15    ($V:expr, $AR:ty) => {
16        Ok(std::sync::Arc::new(
17            $V.map(<$AR>::new_scalar)
18                .unwrap_or_else(|| arrow_array::Scalar::new(<$AR>::new_null(1))),
19        ))
20    };
21}
22
23impl TryFrom<&Scalar> for Arc<dyn Datum> {
24    type Error = VortexError;
25
26    fn try_from(value: &Scalar) -> Result<Arc<dyn Datum>, Self::Error> {
27        match value.dtype() {
28            DType::Null => Ok(Arc::new(NullArray::new(1))),
29            DType::Bool(_) => value_to_arrow_scalar!(value.as_bool().value(), BooleanArray),
30            DType::Primitive(ptype, ..) => {
31                let scalar = value.as_primitive();
32                Ok(match ptype {
33                    PType::U8 => scalar
34                        .typed_value()
35                        .map(|i| Arc::new(UInt8Array::new_scalar(i)) as Arc<dyn Datum>)
36                        .unwrap_or_else(|| Arc::new(UInt8Array::new_null(1))),
37                    PType::U16 => scalar
38                        .typed_value()
39                        .map(|i| Arc::new(UInt16Array::new_scalar(i)) as Arc<dyn Datum>)
40                        .unwrap_or_else(|| Arc::new(UInt16Array::new_null(1))),
41                    PType::U32 => scalar
42                        .typed_value()
43                        .map(|i| Arc::new(UInt32Array::new_scalar(i)) as Arc<dyn Datum>)
44                        .unwrap_or_else(|| Arc::new(UInt32Array::new_null(1))),
45                    PType::U64 => scalar
46                        .typed_value()
47                        .map(|i| Arc::new(UInt64Array::new_scalar(i)) as Arc<dyn Datum>)
48                        .unwrap_or_else(|| Arc::new(UInt64Array::new_null(1))),
49                    PType::I8 => scalar
50                        .typed_value()
51                        .map(|i| Arc::new(Int8Array::new_scalar(i)) as Arc<dyn Datum>)
52                        .unwrap_or_else(|| Arc::new(Int8Array::new_null(1))),
53                    PType::I16 => scalar
54                        .typed_value()
55                        .map(|i| Arc::new(Int16Array::new_scalar(i)) as Arc<dyn Datum>)
56                        .unwrap_or_else(|| Arc::new(Int16Array::new_null(1))),
57                    PType::I32 => scalar
58                        .typed_value()
59                        .map(|i| Arc::new(Int32Array::new_scalar(i)) as Arc<dyn Datum>)
60                        .unwrap_or_else(|| Arc::new(Int32Array::new_null(1))),
61                    PType::I64 => scalar
62                        .typed_value()
63                        .map(|i| Arc::new(Int64Array::new_scalar(i)) as Arc<dyn Datum>)
64                        .unwrap_or_else(|| Arc::new(Int64Array::new_null(1))),
65                    PType::F16 => scalar
66                        .typed_value()
67                        .map(|i| Arc::new(Float16Array::new_scalar(i)) as Arc<dyn Datum>)
68                        .unwrap_or_else(|| Arc::new(Float16Array::new_null(1))),
69                    PType::F32 => scalar
70                        .typed_value()
71                        .map(|i| Arc::new(Float32Array::new_scalar(i)) as Arc<dyn Datum>)
72                        .unwrap_or_else(|| Arc::new(Float32Array::new_null(1))),
73                    PType::F64 => scalar
74                        .typed_value()
75                        .map(|i| Arc::new(Float64Array::new_scalar(i)) as Arc<dyn Datum>)
76                        .unwrap_or_else(|| Arc::new(Float64Array::new_null(1))),
77                })
78            }
79            DType::Decimal(..) => match value.as_decimal().decimal_value() {
80                // TODO(joe): replace with decimal32, etc.
81                Some(DecimalValue::I8(v)) => Ok(Arc::new(Decimal128Array::new_scalar(v as i128))),
82                Some(DecimalValue::I16(v)) => Ok(Arc::new(Decimal128Array::new_scalar(v as i128))),
83                Some(DecimalValue::I32(v)) => Ok(Arc::new(Decimal128Array::new_scalar(v as i128))),
84                Some(DecimalValue::I64(v)) => Ok(Arc::new(Decimal128Array::new_scalar(v as i128))),
85                Some(DecimalValue::I128(v128)) => Ok(Arc::new(Decimal128Array::new_scalar(v128))),
86                Some(DecimalValue::I256(v256)) => {
87                    Ok(Arc::new(Decimal256Array::new_scalar(v256.into())))
88                }
89                None => Ok(Arc::new(arrow_array::Scalar::new(
90                    Decimal128Array::new_null(1),
91                ))),
92            },
93            DType::Utf8(_) => {
94                value_to_arrow_scalar!(value.as_utf8().value(), StringViewArray)
95            }
96            DType::Binary(_) => {
97                value_to_arrow_scalar!(value.as_binary().value(), BinaryViewArray)
98            }
99            DType::Struct(..) => {
100                todo!("struct scalar conversion")
101            }
102            DType::List(..) => {
103                todo!("list scalar conversion")
104            }
105            DType::Extension(ext) => {
106                if is_temporal_ext_type(ext.id()) {
107                    let metadata = TemporalMetadata::try_from(ext.as_ref())?;
108                    let storage_scalar = value.as_extension().storage();
109                    let primitive = storage_scalar
110                        .as_primitive_opt()
111                        .ok_or_else(|| vortex_err!("Expected primitive scalar"))?;
112
113                    return match metadata {
114                        TemporalMetadata::Time(u) => match u {
115                            TimeUnit::Ns => value_to_arrow_scalar!(
116                                primitive.as_::<i64>()?,
117                                Time64NanosecondArray
118                            ),
119                            TimeUnit::Us => value_to_arrow_scalar!(
120                                primitive.as_::<i64>()?,
121                                Time64MicrosecondArray
122                            ),
123                            TimeUnit::Ms => value_to_arrow_scalar!(
124                                primitive.as_::<i32>()?,
125                                Time32MillisecondArray
126                            ),
127                            TimeUnit::S => {
128                                value_to_arrow_scalar!(primitive.as_::<i32>()?, Time32SecondArray)
129                            }
130                            TimeUnit::D => {
131                                vortex_bail!("Unsupported TimeUnit {u} for {}", ext.id())
132                            }
133                        },
134                        TemporalMetadata::Date(u) => match u {
135                            TimeUnit::Ms => {
136                                value_to_arrow_scalar!(primitive.as_::<i64>()?, Date64Array)
137                            }
138                            TimeUnit::D => {
139                                value_to_arrow_scalar!(primitive.as_::<i32>()?, Date32Array)
140                            }
141                            _ => vortex_bail!("Unsupported TimeUnit {u} for {}", ext.id()),
142                        },
143                        TemporalMetadata::Timestamp(u, _) => match u {
144                            TimeUnit::Ns => value_to_arrow_scalar!(
145                                primitive.as_::<i64>()?,
146                                TimestampNanosecondArray
147                            ),
148                            TimeUnit::Us => value_to_arrow_scalar!(
149                                primitive.as_::<i64>()?,
150                                TimestampMicrosecondArray
151                            ),
152                            TimeUnit::Ms => value_to_arrow_scalar!(
153                                primitive.as_::<i64>()?,
154                                TimestampMillisecondArray
155                            ),
156                            TimeUnit::S => value_to_arrow_scalar!(
157                                primitive.as_::<i64>()?,
158                                TimestampSecondArray
159                            ),
160                            TimeUnit::D => {
161                                vortex_bail!("Unsupported TimeUnit {u} for {}", ext.id())
162                            }
163                        },
164                    };
165                }
166
167                todo!("Non temporal extension scalar conversion")
168            }
169        }
170    }
171}
172
173#[cfg(test)]
174mod tests;
vortex_scalar/arrow/mod.rs

vortex_scalar/arrow/
mod.rs
