8000 Expand Rust return type support for Arrow DataTypes in ScalarValue (#… · chenqin/arrow-datafusion-python@2bf684e · GitHub
[go: up one dir, main page]
Skip to content

Commit 2bf684e

Browse files
jdye64jdye64
authored
Expand Rust return type support for Arrow DataTypes in ScalarValue (apache#287)
1 parent 1ed980a commit 2bf684e

File tree

2 files changed

+105
-16
lines changed

2 files changed

+105
-16
lines changed

‎src/errors.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@ use datafusion::arrow::error::ArrowError;
2323
use datafusion::error::DataFusionError as InnerDataFusionError;
2424
use pyo3::{exceptions::PyException, PyErr};
2525

26+
pub type Result<T> = std::result::Result<T, DataFusionError>;
27+
2628
#[derive(Debug)]
2729
pub enum DataFusionError {
2830
ExecutionError(InnerDataFusionError),

‎src/expr/literal.rs

Lines changed: 103 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
use crate::errors::py_runtime_err;
18+
use crate::errors::{py_runtime_err, DataFusionError};
1919
use datafusion_common::ScalarValue;
2020
use pyo3::prelude::*;
2121

@@ -37,38 +37,125 @@ impl From<ScalarValue> for PyLiteral {
3737
}
3838
}
3939

40+
macro_rules! extract_scalar_value {
41+
($self: expr, $variant: ident) => {
42+
match &$self.value {
43+
ScalarValue::$variant(value) => Ok(*value),
44+
other => Err(unexpected_literal_value(other)),
45+
}
46+
};
47+
}
48+
4049
#[pymethods]
4150
impl PyLiteral {
4251
/// Get the data type of this literal value
4352
fn data_type(&self) -> String {
4453
format!("{}", self.value.get_datatype())
4554
}
4655

47-
fn value_i32(&self) -> PyResult<i32> {
48-
if let ScalarValue::Int32(Some(n)) = &self.value {
49-
Ok(*n)
50-
} else {
51-
Err(py_runtime_err("Cannot access value as i32"))
56+
pub fn value_f32(&self) -> PyResult<Option<f32>> {
57+
extract_scalar_value!(self, Float32)
58+
}
59+
60+
pub fn value_f64(&self) -> PyResult<Option<f64>> {
61+
extract_scalar_value!(self, Float64)
62+
}
63+
64+
pub fn value_decimal128(&mut self) -> PyResult<(Option<i128>, u8, i8)> {
65+
match &self.value {
66+
ScalarValue::Decimal128(value, precision, scale) => Ok((*value, *precision, *scale)),
67+
other => Err(unexpected_literal_value(other)),
68+
}
69+
}
70+
71+
pub fn value_i8(&self) -> PyResult<Option<i8>> {
72+
extract_scalar_value!(self, Int8)
73+
}
74+
75+
pub fn value_i16(&self) -> PyResult<Option<i16>> {
76+
extract_scalar_value!(self, Int16)
77+
}
78+
79+
pub fn value_i32(&self) -> PyResult<Option<i32>> {
80+
extract_scalar_value!(self, Int32)
81+
}
82+
83+
pub fn value_i64(&self) -> PyResult<Option<i64>> {
84+
extract_scalar_value!(self, Int64)
85+
}
86+
87+
pub fn value_u8(&self) -> PyResult<Option<u8>> {
88+
extract_scalar_value!(self, UInt8)
89+
}
90+
91+
pub fn value_u16(&self) -> PyResult<Option<u16>> {
92+
extract_scalar_value!(self, UInt16)
93+
}
94+
95+
pub fn value_u32(&self) -> PyResult<Option<u32>> {
96+
extract_scalar_value!(self, UInt32)
97+
}
98+
99+
pub fn value_u64(&self) -> PyResult<Option<u64>> {
100+
extract_scalar_value!(self, UInt64)
101+
}
102+
103+
pub fn value_date32(&self) -> PyResult<Option<i32>> {
104+
extract_scalar_value!(self, Date32)
105+
}
106+
107+
pub fn value_date64(&self) -> PyResult<Option<i64>> {
108+
extract_scalar_value!( 9E88 self, Date64)
109+
}
110+
111+
pub fn value_time64(&self) -> PyResult<Option<i64>> {
112+
extract_scalar_value!(self, Time64Nanosecond)
113+
}
114+
115+
pub fn value_timestamp(&mut self) -> PyResult<(Option<i64>, Option<String>)> {
116+
match &self.value {
117+
ScalarValue::TimestampNanosecond(iv, tz)
118+
| ScalarValue::TimestampMicrosecond(iv, tz)
119+
| ScalarValue::TimestampMillisecond(iv, tz)
120+
| ScalarValue::TimestampSecond(iv, tz) => Ok((*iv, tz.clone())),
121+
other => Err(unexpected_literal_value(other)),
52122
}
53123
}
54124

55-
fn value_i64(&self) -> PyResult<i64> {
56-
if let ScalarValue::Int64(Some(n)) = &self.value {
57-
Ok(*n)
58-
} else {
59-
Err(py_runtime_err("Cannot access value as i64"))
125+
pub fn value_bool(&self) -> PyResult<Option<bool>> {
126+
extract_scalar_value!(self, Boolean)
127+
}
128+
129+
pub fn value_string(&self) -> PyResult<Option<String>> {
130+
match &self.value {
131+
ScalarValue::Utf8(value) => Ok(value.clone()),
132+
other => Err(unexpected_literal_value(other)),
60133
}
61134
}
62135

63-
fn value_str(&self) -> PyResult<String> {
64-
if let ScalarValue::Utf8(Some(str)) = &self.value {
65-
Ok(str.clone())
66-
} else {
67-
Err(py_runtime_err("Cannot access value as string"))
136+
pub fn value_interval_day_time(&self) -> PyResult<Option<(i32, i32)>> {
137+
match &self.value {
138+
ScalarValue::IntervalDayTime(Some(iv)) => {
139+
let interval = *iv as u64;
140+
let days = (interval >> 32) as i32;
141+
let ms = interval as i32;
142+
Ok(Some((days, ms)))
143+
}
144+
ScalarValue::IntervalDayTime(None) => Ok(None),
145+
other => Err(unexpected_literal_value(other)),
68146
}
69147
}
70148

149+
#[allow(clippy::wrong_self_convention)]
150+
fn into_type(&self, py: Python) -> PyResult<PyObject> {
151+
Ok(self.clone().into_py(py))
152+
}
153+
71154
fn __repr__(&self) -> PyResult<String> {
72155
Ok(format!("{}", self.value))
73156
}
74157
}
158+
159+
fn unexpected_literal_value(value: &ScalarValue) -> PyErr {
160+
DataFusionError::Common(format!("getValue<T>() - Unexpected value: {value}")).into()
161+
}

0 commit comments

Comments
 (0)
0