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

datafusion_python/expr/
literal.rs
