datafusion_python/expr/
literal.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::errors::PyDataFusionError;
19use datafusion::common::ScalarValue;
20use pyo3::{prelude::*, IntoPyObjectExt};
21
22#[pyclass(name = "Literal", module = "datafusion.expr", subclass)]
23#[derive(Clone)]
24pub struct PyLiteral {
25    pub value: ScalarValue,
26}
27
28impl From<PyLiteral> for ScalarValue {
29    fn from(lit: PyLiteral) -> ScalarValue {
30        lit.value
31    }
32}
33
34impl From<ScalarValue> for PyLiteral {
35    fn from(value: ScalarValue) -> PyLiteral {
36        PyLiteral { value }
37    }
38}
39
40macro_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
49#[pymethods]
50impl PyLiteral {
51    /// Get the data type of this literal value
52    fn data_type(&self) -> String {
53        format!("{}", self.value.data_type())
54    }
55
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!(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) => {
121                Ok((*iv, tz.as_ref().map(|s| s.as_ref().to_string())))
122            }
123            other => Err(unexpected_literal_value(other)),
124        }
125    }
126
127    pub fn value_bool(&self) -> PyResult<Option<bool>> {
128        extract_scalar_value!(self, Boolean)
129    }
130
131    pub fn value_string(&self) -> PyResult<Option<String>> {
132        match &self.value {
133            ScalarValue::Utf8(value) => Ok(value.clone()),
134            other => Err(unexpected_literal_value(other)),
135        }
136    }
137
138    pub fn value_interval_day_time(&self) -> PyResult<Option<(i32, i32)>> {
139        match &self.value {
140            ScalarValue::IntervalDayTime(Some(iv)) => Ok(Some((iv.days, iv.milliseconds))),
141            ScalarValue::IntervalDayTime(None) => Ok(None),
142            other => Err(unexpected_literal_value(other)),
143        }
144    }
145
146    #[allow(clippy::wrong_self_convention)]
147    fn into_type<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
148        self.clone().into_bound_py_any(py)
149    }
150
151    fn __repr__(&self) -> PyResult<String> {
152        Ok(format!("{}", self.value))
153    }
154}
155
156fn unexpected_literal_value(value: &ScalarValue) -> PyErr {
157    PyDataFusionError::Common(format!("getValue<T>() - Unexpected value: {value}")).into()
158}
datafusion_python/expr/literal.rs

datafusion_python/expr/
literal.rs
