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 datafusion::common::ScalarValue;
19use datafusion::logical_expr::expr::FieldMetadata;
20use pyo3::prelude::*;
21use pyo3::IntoPyObjectExt;
22
23use crate::errors::PyDataFusionError;
24
25#[pyclass(name = "Literal", module = "datafusion.expr", subclass, frozen)]
26#[derive(Clone)]
27pub struct PyLiteral {
28    pub value: ScalarValue,
29    pub metadata: Option<FieldMetadata>,
30}
31
32impl PyLiteral {
33    pub fn new_with_metadata(value: ScalarValue, metadata: Option<FieldMetadata>) -> PyLiteral {
34        Self { value, metadata }
35    }
36}
37
38impl From<PyLiteral> for ScalarValue {
39    fn from(lit: PyLiteral) -> ScalarValue {
40        lit.value
41    }
42}
43
44impl From<ScalarValue> for PyLiteral {
45    fn from(value: ScalarValue) -> PyLiteral {
46        PyLiteral {
47            value,
48            metadata: None,
49        }
50    }
51}
52
53macro_rules! extract_scalar_value {
54    ($self: expr, $variant: ident) => {
55        match &$self.value {
56            ScalarValue::$variant(value) => Ok(*value),
57            other => Err(unexpected_literal_value(other)),
58        }
59    };
60}
61
62#[pymethods]
63impl PyLiteral {
64    /// Get the data type of this literal value
65    fn data_type(&self) -> String {
66        format!("{}", self.value.data_type())
67    }
68
69    pub fn value_f32(&self) -> PyResult<Option<f32>> {
70        extract_scalar_value!(self, Float32)
71    }
72
73    pub fn value_f64(&self) -> PyResult<Option<f64>> {
74        extract_scalar_value!(self, Float64)
75    }
76
77    pub fn value_decimal128(&self) -> PyResult<(Option<i128>, u8, i8)> {
78        match &self.value {
79            ScalarValue::Decimal128(value, precision, scale) => Ok((*value, *precision, *scale)),
80            other => Err(unexpected_literal_value(other)),
81        }
82    }
83
84    pub fn value_i8(&self) -> PyResult<Option<i8>> {
85        extract_scalar_value!(self, Int8)
86    }
87
88    pub fn value_i16(&self) -> PyResult<Option<i16>> {
89        extract_scalar_value!(self, Int16)
90    }
91
92    pub fn value_i32(&self) -> PyResult<Option<i32>> {
93        extract_scalar_value!(self, Int32)
94    }
95
96    pub fn value_i64(&self) -> PyResult<Option<i64>> {
97        extract_scalar_value!(self, Int64)
98    }
99
100    pub fn value_u8(&self) -> PyResult<Option<u8>> {
101        extract_scalar_value!(self, UInt8)
102    }
103
104    pub fn value_u16(&self) -> PyResult<Option<u16>> {
105        extract_scalar_value!(self, UInt16)
106    }
107
108    pub fn value_u32(&self) -> PyResult<Option<u32>> {
109        extract_scalar_value!(self, UInt32)
110    }
111
112    pub fn value_u64(&self) -> PyResult<Option<u64>> {
113        extract_scalar_value!(self, UInt64)
114    }
115
116    pub fn value_date32(&self) -> PyResult<Option<i32>> {
117        extract_scalar_value!(self, Date32)
118    }
119
120    pub fn value_date64(&self) -> PyResult<Option<i64>> {
121        extract_scalar_value!(self, Date64)
122    }
123
124    pub fn value_time64(&self) -> PyResult<Option<i64>> {
125        extract_scalar_value!(self, Time64Nanosecond)
126    }
127
128    pub fn value_timestamp(&self) -> PyResult<(Option<i64>, Option<String>)> {
129        match &self.value {
130            ScalarValue::TimestampNanosecond(iv, tz)
131            | ScalarValue::TimestampMicrosecond(iv, tz)
132            | ScalarValue::TimestampMillisecond(iv, tz)
133            | ScalarValue::TimestampSecond(iv, tz) => {
134                Ok((*iv, tz.as_ref().map(|s| s.as_ref().to_string())))
135            }
136            other => Err(unexpected_literal_value(other)),
137        }
138    }
139
140    pub fn value_bool(&self) -> PyResult<Option<bool>> {
141        extract_scalar_value!(self, Boolean)
142    }
143
144    pub fn value_string(&self) -> PyResult<Option<String>> {
145        match &self.value {
146            ScalarValue::Utf8(value) => Ok(value.clone()),
147            other => Err(unexpected_literal_value(other)),
148        }
149    }
150
151    pub fn value_interval_day_time(&self) -> PyResult<Option<(i32, i32)>> {
152        match &self.value {
153            ScalarValue::IntervalDayTime(Some(iv)) => Ok(Some((iv.days, iv.milliseconds))),
154            ScalarValue::IntervalDayTime(None) => Ok(None),
155            other => Err(unexpected_literal_value(other)),
156        }
157    }
158
159    #[allow(clippy::wrong_self_convention)]
160    fn into_type<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
161        self.clone().into_bound_py_any(py)
162    }
163
164    fn __repr__(&self) -> PyResult<String> {
165        Ok(format!("{}", self.value))
166    }
167}
168
169fn unexpected_literal_value(value: &ScalarValue) -> PyErr {
170    PyDataFusionError::Common(format!("getValue<T>() - Unexpected value: {value}")).into()
171}
datafusion_python/expr/literal.rs

datafusion_python/expr/
literal.rs
