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

datafusion_python/expr/
literal.rs
