datafusion_python/
expr.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 std::collections::HashMap;
19use std::convert::{From, Into};
20use std::sync::Arc;
21
22use datafusion::arrow::datatypes::{DataType, Field};
23use datafusion::arrow::pyarrow::PyArrowType;
24use datafusion::functions::core::expr_ext::FieldAccessor;
25use datafusion::logical_expr::expr::{
26    AggregateFunction, AggregateFunctionParams, FieldMetadata, HigherOrderFunction, InList,
27    InSubquery, Lambda, ScalarFunction, SetComparison, WindowFunction,
28};
29use datafusion::logical_expr::utils::exprlist_to_fields;
30use datafusion::logical_expr::{
31    Between, BinaryExpr, Case, Cast, Expr, ExprFuncBuilder, ExprFunctionExt, Like, LogicalPlan,
32    Operator, TryCast, WindowFunctionDefinition, col, lit, lit_with_metadata,
33};
34use datafusion_proto::logical_plan::{from_proto, to_proto};
35use prost::Message;
36use pyo3::IntoPyObjectExt;
37use pyo3::basic::CompareOp;
38use pyo3::exceptions::PyRuntimeError;
39use pyo3::prelude::*;
40use pyo3::types::PyBytes;
41use window::PyWindowFrame;
42
43use self::alias::PyAlias;
44use self::bool_expr::{
45    PyIsFalse, PyIsNotFalse, PyIsNotNull, PyIsNotTrue, PyIsNotUnknown, PyIsNull, PyIsTrue,
46    PyIsUnknown, PyNegative, PyNot,
47};
48use self::like::{PyILike, PyLike, PySimilarTo};
49use self::scalar_variable::PyScalarVariable;
50use crate::codec::PythonLogicalCodec;
51use crate::common::data_type::{DataTypeMap, NullTreatment, PyScalarValue, RexType};
52use crate::context::PySessionContext;
53use crate::errors::{PyDataFusionResult, py_runtime_err, py_type_err, py_unsupported_variant_err};
54use crate::expr::aggregate_expr::PyAggregateFunction;
55use crate::expr::binary_expr::PyBinaryExpr;
56use crate::expr::column::PyColumn;
57use crate::expr::literal::PyLiteral;
58use crate::functions::add_builder_fns_to_window;
59use crate::pyarrow_util::scalar_to_pyarrow;
60use crate::sql::logical::PyLogicalPlan;
61
62pub mod aggregate;
63pub mod aggregate_expr;
64pub mod alias;
65pub mod analyze;
66pub mod between;
67pub mod binary_expr;
68pub mod bool_expr;
69pub mod case;
70pub mod cast;
71pub mod column;
72pub mod conditional_expr;
73pub mod copy_to;
74pub mod create_catalog;
75pub mod create_catalog_schema;
76pub mod create_external_table;
77pub mod create_function;
78pub mod create_index;
79pub mod create_memory_table;
80pub mod create_view;
81pub mod describe_table;
82pub mod distinct;
83pub mod dml;
84pub mod drop_catalog_schema;
85pub mod drop_function;
86pub mod drop_table;
87pub mod drop_view;
88pub mod empty_relation;
89pub mod exists;
90pub mod explain;
91pub mod extension;
92pub mod filter;
93pub mod grouping_set;
94pub mod higher_order_function;
95pub mod in_list;
96pub mod in_subquery;
97pub mod join;
98pub mod lambda;
99pub mod lambda_variable;
100pub mod like;
101pub mod limit;
102pub mod literal;
103pub mod logical_node;
104pub mod placeholder;
105pub mod projection;
106pub mod recursive_query;
107pub mod repartition;
108pub mod scalar_subquery;
109pub mod scalar_variable;
110pub mod set_comparison;
111pub mod signature;
112pub mod sort;
113pub mod sort_expr;
114pub mod statement;
115pub mod subquery;
116pub mod subquery_alias;
117pub mod table_scan;
118pub mod union;
119pub mod unnest;
120pub mod unnest_expr;
121pub mod values;
122pub mod window;
123
124use sort_expr::{PySortExpr, to_sort_expressions};
125
126/// A PyExpr that can be used on a DataFrame
127#[pyclass(
128    from_py_object,
129    frozen,
130    name = "RawExpr",
131    module = "datafusion.expr",
132    subclass
133)]
134#[derive(Debug, Clone)]
135pub struct PyExpr {
136    pub expr: Expr,
137}
138
139impl From<PyExpr> for Expr {
140    fn from(expr: PyExpr) -> Expr {
141        expr.expr
142    }
143}
144
145impl From<Expr> for PyExpr {
146    fn from(expr: Expr) -> PyExpr {
147        PyExpr { expr }
148    }
149}
150
151/// Convert a list of DataFusion Expr to PyExpr
152pub fn py_expr_list(expr: &[Expr]) -> PyResult<Vec<PyExpr>> {
153    Ok(expr.iter().map(|e| PyExpr::from(e.clone())).collect())
154}
155
156#[pymethods]
157impl PyExpr {
158    /// Return the specific expression
159    fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
160        Python::attach(|_| match &self.expr {
161            Expr::Alias(alias) => Ok(PyAlias::from(alias.clone()).into_bound_py_any(py)?),
162            Expr::Column(col) => Ok(PyColumn::from(col.clone()).into_bound_py_any(py)?),
163            Expr::ScalarVariable(field, variables) => {
164                Ok(PyScalarVariable::new(field, variables).into_bound_py_any(py)?)
165            }
166            Expr::Like(value) => Ok(PyLike::from(value.clone()).into_bound_py_any(py)?),
167            Expr::Literal(value, metadata) => Ok(PyLiteral::new_with_metadata(
168                value.clone(),
169                metadata.clone(),
170            )
171            .into_bound_py_any(py)?),
172            Expr::BinaryExpr(expr) => Ok(PyBinaryExpr::from(expr.clone()).into_bound_py_any(py)?),
173            Expr::Not(expr) => Ok(PyNot::new(*expr.clone()).into_bound_py_any(py)?),
174            Expr::IsNotNull(expr) => Ok(PyIsNotNull::new(*expr.clone()).into_bound_py_any(py)?),
175            Expr::IsNull(expr) => Ok(PyIsNull::new(*expr.clone()).into_bound_py_any(py)?),
176            Expr::IsTrue(expr) => Ok(PyIsTrue::new(*expr.clone()).into_bound_py_any(py)?),
177            Expr::IsFalse(expr) => Ok(PyIsFalse::new(*expr.clone()).into_bound_py_any(py)?),
178            Expr::IsUnknown(expr) => Ok(PyIsUnknown::new(*expr.clone()).into_bound_py_any(py)?),
179            Expr::IsNotTrue(expr) => Ok(PyIsNotTrue::new(*expr.clone()).into_bound_py_any(py)?),
180            Expr::IsNotFalse(expr) => Ok(PyIsNotFalse::new(*expr.clone()).into_bound_py_any(py)?),
181            Expr::IsNotUnknown(expr) => {
182                Ok(PyIsNotUnknown::new(*expr.clone()).into_bound_py_any(py)?)
183            }
184            Expr::Negative(expr) => Ok(PyNegative::new(*expr.clone()).into_bound_py_any(py)?),
185            Expr::AggregateFunction(expr) => {
186                Ok(PyAggregateFunction::from(expr.clone()).into_bound_py_any(py)?)
187            }
188            Expr::SimilarTo(value) => Ok(PySimilarTo::from(value.clone()).into_bound_py_any(py)?),
189            Expr::Between(value) => {
190                Ok(between::PyBetween::from(value.clone()).into_bound_py_any(py)?)
191            }
192            Expr::Case(value) => Ok(case::PyCase::from(value.clone()).into_bound_py_any(py)?),
193            Expr::Cast(value) => Ok(cast::PyCast::from(value.clone()).into_bound_py_any(py)?),
194            Expr::TryCast(value) => Ok(cast::PyTryCast::from(value.clone()).into_bound_py_any(py)?),
195            Expr::ScalarFunction(value) => Err(py_unsupported_variant_err(format!(
196                "Converting Expr::ScalarFunction to a Python object is not implemented: {value:?}"
197            ))),
198            Expr::WindowFunction(value) => Err(py_unsupported_variant_err(format!(
199                "Converting Expr::WindowFunction to a Python object is not implemented: {value:?}"
200            ))),
201            Expr::InList(value) => {
202                Ok(in_list::PyInList::from(value.clone()).into_bound_py_any(py)?)
203            }
204            Expr::Exists(value) => Ok(exists::PyExists::from(value.clone()).into_bound_py_any(py)?),
205            Expr::InSubquery(value) => {
206                Ok(in_subquery::PyInSubquery::from(value.clone()).into_bound_py_any(py)?)
207            }
208            Expr::ScalarSubquery(value) => {
209                Ok(scalar_subquery::PyScalarSubquery::from(value.clone()).into_bound_py_any(py)?)
210            }
211            #[allow(deprecated)]
212            Expr::Wildcard { qualifier, options } => Err(py_unsupported_variant_err(format!(
213                "Converting Expr::Wildcard to a Python object is not implemented : {qualifier:?} {options:?}"
214            ))),
215            Expr::GroupingSet(value) => {
216                Ok(grouping_set::PyGroupingSet::from(value.clone()).into_bound_py_any(py)?)
217            }
218            Expr::Placeholder(value) => {
219                Ok(placeholder::PyPlaceholder::from(value.clone()).into_bound_py_any(py)?)
220            }
221            Expr::OuterReferenceColumn(data_type, column) => {
222                Err(py_unsupported_variant_err(format!(
223                    "Converting Expr::OuterReferenceColumn to a Python object is not implemented: {data_type:?} - {column:?}"
224                )))
225            }
226            Expr::Unnest(value) => {
227                Ok(unnest_expr::PyUnnestExpr::from(value.clone()).into_bound_py_any(py)?)
228            }
229            Expr::SetComparison(value) => {
230                Ok(set_comparison::PySetComparison::from(value.clone()).into_bound_py_any(py)?)
231            }
232            Expr::HigherOrderFunction(value) => Ok(
233                higher_order_function::PyHigherOrderFunction::from(value.clone())
234                    .into_bound_py_any(py)?,
235            ),
236            Expr::Lambda(value) => Ok(lambda::PyLambda::from(value.clone()).into_bound_py_any(py)?),
237            Expr::LambdaVariable(value) => {
238                Ok(lambda_variable::PyLambdaVariable::from(value.clone()).into_bound_py_any(py)?)
239            }
240        })
241    }
242
243    /// Returns the name of this expression as it should appear in a schema. This name
244    /// will not include any CAST expressions.
245    fn schema_name(&self) -> PyResult<String> {
246        Ok(format!("{}", self.expr.schema_name()))
247    }
248
249    /// Returns a full and complete string representation of this expression.
250    fn canonical_name(&self) -> PyResult<String> {
251        Ok(format!("{}", self.expr))
252    }
253
254    /// Returns the name of the Expr variant.
255    /// Ex: 'IsNotNull', 'Literal', 'BinaryExpr', etc
256    fn variant_name(&self) -> PyResult<&str> {
257        Ok(self.expr.variant_name())
258    }
259
260    fn __richcmp__(&self, other: PyExpr, op: CompareOp) -> PyExpr {
261        let expr = match op {
262            CompareOp::Lt => self.expr.clone().lt(other.expr),
263            CompareOp::Le => self.expr.clone().lt_eq(other.expr),
264            CompareOp::Eq => self.expr.clone().eq(other.expr),
265            CompareOp::Ne => self.expr.clone().not_eq(other.expr),
266            CompareOp::Gt => self.expr.clone().gt(other.expr),
267            CompareOp::Ge => self.expr.clone().gt_eq(other.expr),
268        };
269        expr.into()
270    }
271
272    fn __repr__(&self) -> PyResult<String> {
273        Ok(format!("Expr({})", self.expr))
274    }
275
276    fn __add__(&self, rhs: PyExpr) -> PyResult<PyExpr> {
277        Ok((self.expr.clone() + rhs.expr).into())
278    }
279
280    fn __sub__(&self, rhs: PyExpr) -> PyResult<PyExpr> {
281        Ok((self.expr.clone() - rhs.expr).into())
282    }
283
284    fn __truediv__(&self, rhs: PyExpr) -> PyResult<PyExpr> {
285        Ok((self.expr.clone() / rhs.expr).into())
286    }
287
288    fn __mul__(&self, rhs: PyExpr) -> PyResult<PyExpr> {
289        Ok((self.expr.clone() * rhs.expr).into())
290    }
291
292    fn __mod__(&self, rhs: PyExpr) -> PyResult<PyExpr> {
293        let expr = self.expr.clone() % rhs.expr;
294        Ok(expr.into())
295    }
296
297    fn __and__(&self, rhs: PyExpr) -> PyResult<PyExpr> {
298        Ok(self.expr.clone().and(rhs.expr).into())
299    }
300
301    fn __or__(&self, rhs: PyExpr) -> PyResult<PyExpr> {
302        Ok(self.expr.clone().or(rhs.expr).into())
303    }
304
305    fn __invert__(&self) -> PyResult<PyExpr> {
306        let expr = !self.expr.clone();
307        Ok(expr.into())
308    }
309
310    fn __getitem__(&self, key: &str) -> PyResult<PyExpr> {
311        Ok(self.expr.clone().field(key).into())
312    }
313
314    #[staticmethod]
315    pub fn literal(value: PyScalarValue) -> PyExpr {
316        lit(value.0).into()
317    }
318
319    #[staticmethod]
320    pub fn literal_with_metadata(
321        value: PyScalarValue,
322        metadata: HashMap<String, String>,
323    ) -> PyExpr {
324        let metadata = FieldMetadata::new(metadata.into_iter().collect());
325        lit_with_metadata(value.0, Some(metadata)).into()
326    }
327
328    #[staticmethod]
329    pub fn column(value: &str) -> PyExpr {
330        col(value).into()
331    }
332
333    /// assign a name to the PyExpr
334    #[pyo3(signature = (name, metadata=None))]
335    pub fn alias(&self, name: &str, metadata: Option<HashMap<String, String>>) -> PyExpr {
336        let metadata = metadata.map(|m| FieldMetadata::new(m.into_iter().collect()));
337        self.expr.clone().alias_with_metadata(name, metadata).into()
338    }
339
340    /// Create a sort PyExpr from an existing PyExpr.
341    #[pyo3(signature = (ascending=true, nulls_first=true))]
342    pub fn sort(&self, ascending: bool, nulls_first: bool) -> PySortExpr {
343        self.expr.clone().sort(ascending, nulls_first).into()
344    }
345
346    pub fn is_null(&self) -> PyExpr {
347        self.expr.clone().is_null().into()
348    }
349
350    pub fn is_not_null(&self) -> PyExpr {
351        self.expr.clone().is_not_null().into()
352    }
353
354    pub fn cast(&self, to: PyArrowType<DataType>) -> PyExpr {
355        // self.expr.cast_to() requires DFSchema to validate that the cast
356        // is supported, omit that for now
357        let expr = Expr::Cast(Cast::new(Box::new(self.expr.clone()), to.0));
358        expr.into()
359    }
360
361    pub fn try_cast(&self, to: PyArrowType<DataType>) -> PyExpr {
362        let expr = Expr::TryCast(TryCast::new(Box::new(self.expr.clone()), to.0));
363        expr.into()
364    }
365
366    #[pyo3(signature = (low, high, negated=false))]
367    pub fn between(&self, low: PyExpr, high: PyExpr, negated: bool) -> PyExpr {
368        let expr = Expr::Between(Between::new(
369            Box::new(self.expr.clone()),
370            negated,
371            Box::new(low.into()),
372            Box::new(high.into()),
373        ));
374        expr.into()
375    }
376
377    /// A Rex (Row Expression) specifies a single row of data. That specification
378    /// could include user defined functions or types. RexType identifies the row
379    /// as one of the possible valid `RexTypes`.
380    pub fn rex_type(&self) -> PyResult<RexType> {
381        Ok(match self.expr {
382            Expr::Alias(..) => RexType::Alias,
383            Expr::Column(..) => RexType::Reference,
384            Expr::ScalarVariable(..) | Expr::Literal(..) => RexType::Literal,
385            Expr::BinaryExpr { .. }
386            | Expr::Not(..)
387            | Expr::IsNotNull(..)
388            | Expr::Negative(..)
389            | Expr::IsNull(..)
390            | Expr::Like { .. }
391            | Expr::SimilarTo { .. }
392            | Expr::Between { .. }
393            | Expr::Case { .. }
394            | Expr::Cast { .. }
395            | Expr::TryCast { .. }
396            | Expr::ScalarFunction { .. }
397            | Expr::AggregateFunction { .. }
398            | Expr::WindowFunction { .. }
399            | Expr::InList { .. }
400            | Expr::Exists { .. }
401            | Expr::InSubquery { .. }
402            | Expr::GroupingSet(..)
403            | Expr::IsTrue(..)
404            | Expr::IsFalse(..)
405            | Expr::IsUnknown(_)
406            | Expr::IsNotTrue(..)
407            | Expr::IsNotFalse(..)
408            | Expr::Placeholder { .. }
409            | Expr::OuterReferenceColumn(_, _)
410            | Expr::Unnest(_)
411            | Expr::IsNotUnknown(_)
412            | Expr::SetComparison(_)
413            | Expr::HigherOrderFunction(..)
414            | Expr::Lambda(..) => RexType::Call,
415            Expr::LambdaVariable(..) => RexType::Reference,
416            Expr::ScalarSubquery(..) => RexType::ScalarSubquery,
417            #[allow(deprecated)]
418            Expr::Wildcard { .. } => {
419                return Err(py_unsupported_variant_err("Expr::Wildcard is unsupported"));
420            }
421        })
422    }
423
424    /// Given the current `Expr` return the DataTypeMap which represents the
425    /// PythonType, Arrow DataType, and SqlType Enum which represents
426    pub fn types(&self) -> PyResult<DataTypeMap> {
427        Self::_types(&self.expr)
428    }
429
430    /// Extracts the Expr value into a Py<PyAny> that can be shared with Python
431    pub fn python_value<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
432        match &self.expr {
433            Expr::Literal(scalar_value, _) => scalar_to_pyarrow(scalar_value, py),
434            _ => Err(py_type_err(format!(
435                "Non Expr::Literal encountered in types: {:?}",
436                &self.expr
437            ))),
438        }
439    }
440
441    /// Row expressions, Rex(s), operate on the concept of operands. Different variants of Expressions, Expr(s),
442    /// store those operands in different datastructures. This function examines the Expr variant and returns
443    /// the operands to the calling logic as a Vec of PyExpr instances.
444    pub fn rex_call_operands(&self) -> PyResult<Vec<PyExpr>> {
445        match &self.expr {
446            // Expr variants that are themselves the operand to return
447            Expr::Column(..)
448            | Expr::ScalarVariable(..)
449            | Expr::Literal(..)
450            | Expr::LambdaVariable(..) => Ok(vec![PyExpr::from(self.expr.clone())]),
451
452            Expr::Alias(alias) => Ok(vec![PyExpr::from(*alias.expr.clone())]),
453
454            // Expr(s) that house the Expr instance to return in their bounded params
455            Expr::Not(expr)
456            | Expr::IsNull(expr)
457            | Expr::IsNotNull(expr)
458            | Expr::IsTrue(expr)
459            | Expr::IsFalse(expr)
460            | Expr::IsUnknown(expr)
461            | Expr::IsNotTrue(expr)
462            | Expr::IsNotFalse(expr)
463            | Expr::IsNotUnknown(expr)
464            | Expr::Negative(expr)
465            | Expr::Cast(Cast { expr, .. })
466            | Expr::TryCast(TryCast { expr, .. })
467            | Expr::InSubquery(InSubquery { expr, .. })
468            | Expr::SetComparison(SetComparison { expr, .. }) => {
469                Ok(vec![PyExpr::from(*expr.clone())])
470            }
471
472            // Expr variants containing a collection of Expr(s) for operands
473            Expr::AggregateFunction(AggregateFunction {
474                params: AggregateFunctionParams { args, .. },
475                ..
476            })
477            | Expr::ScalarFunction(ScalarFunction { args, .. })
478            | Expr::HigherOrderFunction(HigherOrderFunction { args, .. }) => {
479                Ok(args.iter().map(|arg| PyExpr::from(arg.clone())).collect())
480            }
481            Expr::WindowFunction(boxed_window_fn) => {
482                let args = &boxed_window_fn.params.args;
483                Ok(args.iter().map(|arg| PyExpr::from(arg.clone())).collect())
484            }
485            Expr::Lambda(Lambda { body, .. }) => Ok(vec![PyExpr::from(*body.clone())]),
486
487            // Expr(s) that require more specific processing
488            Expr::Case(Case {
489                expr,
490                when_then_expr,
491                else_expr,
492            }) => {
493                let mut operands: Vec<PyExpr> = Vec::new();
494
495                if let Some(e) = expr {
496                    for (when, then) in when_then_expr {
497                        operands.push(PyExpr::from(Expr::BinaryExpr(BinaryExpr::new(
498                            Box::new(*e.clone()),
499                            Operator::Eq,
500                            Box::new(*when.clone()),
501                        ))));
502                        operands.push(PyExpr::from(*then.clone()));
503                    }
504                } else {
505                    for (when, then) in when_then_expr {
506                        operands.push(PyExpr::from(*when.clone()));
507                        operands.push(PyExpr::from(*then.clone()));
508                    }
509                };
510
511                if let Some(e) = else_expr {
512                    operands.push(PyExpr::from(*e.clone()));
513                };
514
515                Ok(operands)
516            }
517            Expr::InList(InList { expr, list, .. }) => {
518                let mut operands: Vec<PyExpr> = vec![PyExpr::from(*expr.clone())];
519                for list_elem in list {
520                    operands.push(PyExpr::from(list_elem.clone()));
521                }
522
523                Ok(operands)
524            }
525            Expr::BinaryExpr(BinaryExpr { left, right, .. }) => Ok(vec![
526                PyExpr::from(*left.clone()),
527                PyExpr::from(*right.clone()),
528            ]),
529            Expr::Like(Like { expr, pattern, .. }) => Ok(vec![
530                PyExpr::from(*expr.clone()),
531                PyExpr::from(*pattern.clone()),
532            ]),
533            Expr::SimilarTo(Like { expr, pattern, .. }) => Ok(vec![
534                PyExpr::from(*expr.clone()),
535                PyExpr::from(*pattern.clone()),
536            ]),
537            Expr::Between(Between {
538                expr,
539                negated: _,
540                low,
541                high,
542            }) => Ok(vec![
543                PyExpr::from(*expr.clone()),
544                PyExpr::from(*low.clone()),
545                PyExpr::from(*high.clone()),
546            ]),
547
548            // Currently un-support/implemented Expr types for Rex Call operations
549            Expr::GroupingSet(..)
550            | Expr::Unnest(_)
551            | Expr::OuterReferenceColumn(_, _)
552            | Expr::ScalarSubquery(..)
553            | Expr::Placeholder { .. }
554            | Expr::Exists { .. } => Err(py_runtime_err(format!(
555                "Unimplemented Expr type: {}",
556                self.expr
557            ))),
558
559            #[allow(deprecated)]
560            Expr::Wildcard { .. } => {
561                Err(py_unsupported_variant_err("Expr::Wildcard is unsupported"))
562            }
563        }
564    }
565
566    /// Extracts the operator associated with a RexType::Call
567    pub fn rex_call_operator(&self) -> PyResult<String> {
568        Ok(match &self.expr {
569            Expr::BinaryExpr(BinaryExpr {
570                left: _,
571                op,
572                right: _,
573            }) => format!("{op}"),
574            Expr::ScalarFunction(ScalarFunction { func, args: _ }) => func.name().to_string(),
575            Expr::HigherOrderFunction(HigherOrderFunction { func, args: _ }) => {
576                func.name().to_string()
577            }
578            Expr::Lambda(..) => "lambda".to_string(),
579            Expr::Cast { .. } => "cast".to_string(),
580            Expr::Between { .. } => "between".to_string(),
581            Expr::Case { .. } => "case".to_string(),
582            Expr::IsNull(..) => "is null".to_string(),
583            Expr::IsNotNull(..) => "is not null".to_string(),
584            Expr::IsTrue(_) => "is true".to_string(),
585            Expr::IsFalse(_) => "is false".to_string(),
586            Expr::IsUnknown(_) => "is unknown".to_string(),
587            Expr::IsNotTrue(_) => "is not true".to_string(),
588            Expr::IsNotFalse(_) => "is not false".to_string(),
589            Expr::IsNotUnknown(_) => "is not unknown".to_string(),
590            Expr::InList { .. } => "in list".to_string(),
591            Expr::Negative(..) => "negative".to_string(),
592            Expr::Not(..) => "not".to_string(),
593            Expr::Like(Like {
594                negated,
595                case_insensitive,
596                ..
597            }) => {
598                let name = if *case_insensitive { "ilike" } else { "like" };
599                if *negated {
600                    format!("not {name}")
601                } else {
602                    name.to_string()
603                }
604            }
605            Expr::SimilarTo(Like { negated, .. }) => {
606                if *negated {
607                    "not similar to".to_string()
608                } else {
609                    "similar to".to_string()
610                }
611            }
612            _ => {
613                return Err(py_type_err(format!(
614                    "Catch all triggered in get_operator_name: {:?}",
615                    &self.expr
616                )));
617            }
618        })
619    }
620
621    pub fn column_name(&self, plan: PyLogicalPlan) -> PyResult<String> {
622        self._column_name(&plan.plan()).map_err(py_runtime_err)
623    }
624
625    // Expression Function Builder functions
626
627    pub fn order_by(&self, order_by: Vec<PySortExpr>) -> PyExprFuncBuilder {
628        self.expr
629            .clone()
630            .order_by(to_sort_expressions(order_by))
631            .into()
632    }
633
634    pub fn filter(&self, filter: PyExpr) -> PyExprFuncBuilder {
635        self.expr.clone().filter(filter.expr.clone()).into()
636    }
637
638    pub fn distinct(&self) -> PyExprFuncBuilder {
639        self.expr.clone().distinct().into()
640    }
641
642    pub fn null_treatment(&self, null_treatment: NullTreatment) -> PyExprFuncBuilder {
643        self.expr
644            .clone()
645            .null_treatment(Some(null_treatment.into()))
646            .into()
647    }
648
649    pub fn partition_by(&self, partition_by: Vec<PyExpr>) -> PyExprFuncBuilder {
650        let partition_by = partition_by.iter().map(|e| e.expr.clone()).collect();
651        self.expr.clone().partition_by(partition_by).into()
652    }
653
654    pub fn window_frame(&self, window_frame: PyWindowFrame) -> PyExprFuncBuilder {
655        self.expr.clone().window_frame(window_frame.into()).into()
656    }
657
658    #[pyo3(signature = (partition_by=None, window_frame=None, order_by=None, null_treatment=None))]
659    pub fn over(
660        &self,
661        partition_by: Option<Vec<PyExpr>>,
662        window_frame: Option<PyWindowFrame>,
663        order_by: Option<Vec<PySortExpr>>,
664        null_treatment: Option<NullTreatment>,
665    ) -> PyDataFusionResult<PyExpr> {
666        match &self.expr {
667            Expr::AggregateFunction(agg_fn) => {
668                let window_fn = Expr::WindowFunction(Box::new(WindowFunction::new(
669                    WindowFunctionDefinition::AggregateUDF(agg_fn.func.clone()),
670                    agg_fn.params.args.clone(),
671                )));
672
673                add_builder_fns_to_window(
674                    window_fn,
675                    partition_by,
676                    window_frame,
677                    order_by,
678                    null_treatment,
679                )
680            }
681            Expr::WindowFunction(_) => add_builder_fns_to_window(
682                self.expr.clone(),
683                partition_by,
684                window_frame,
685                order_by,
686                null_treatment,
687            ),
688            _ => Err(datafusion::error::DataFusionError::Plan(format!(
689                "Using {} with `over` is not allowed. Must use an aggregate or window function.",
690                self.expr.variant_name()
691            ))
692            .into()),
693        }
694    }
695
696    /// Serialize this `Expr` to protobuf bytes.
697    ///
698    /// When `ctx` is supplied, encoding routes through the session's
699    /// installed `LogicalExtensionCodec` so user FFI codecs see the
700    /// encode path. Without `ctx` a default-inner Python codec is
701    /// used; Python scalar UDFs still inline when in-band encoding
702    /// lands, non-Python UDFs fall through to the default codec.
703    #[pyo3(signature = (ctx=None))]
704    pub fn to_bytes<'py>(
705        &'py self,
706        py: Python<'py>,
707        ctx: Option<PySessionContext>,
708    ) -> PyDataFusionResult<Bound<'py, PyBytes>> {
709        let default_codec;
710        let codec: &dyn datafusion_proto::logical_plan::LogicalExtensionCodec = match ctx {
711            Some(ref ctx) => ctx.logical_codec().as_ref(),
712            None => {
713                default_codec = PythonLogicalCodec::default();
714                &default_codec
715            }
716        };
717        let proto = to_proto::serialize_expr(&self.expr, codec)
718            .map_err(|e| PyRuntimeError::new_err(format!("Unable to serialize expr: {e}")))?;
719        let bytes = proto.encode_to_vec();
720        Ok(PyBytes::new(py, &bytes))
721    }
722
723    /// Decode an `Expr` from protobuf bytes against the session's
724    /// function registry and logical codec.
725    #[staticmethod]
726    pub fn from_bytes(
727        ctx: PySessionContext,
728        proto_msg: Bound<'_, PyBytes>,
729    ) -> PyDataFusionResult<Self> {
730        let bytes: &[u8] = proto_msg.extract().map_err(Into::<PyErr>::into)?;
731        let proto_expr =
732            datafusion_proto::protobuf::LogicalExprNode::decode(bytes).map_err(|e| {
733                PyRuntimeError::new_err(format!(
734                    "Unable to decode expression from serialized bytes: {e}"
735                ))
736            })?;
737
738        let codec = ctx.logical_codec();
739        let task_ctx = ctx.ctx.task_ctx();
740        let expr = from_proto::parse_expr(&proto_expr, task_ctx.as_ref(), codec.as_ref())
741            .map_err(|e| PyRuntimeError::new_err(format!("Unable to decode expr: {e}")))?;
742        Ok(Self { expr })
743    }
744}
745
746#[pyclass(
747    from_py_object,
748    frozen,
749    name = "ExprFuncBuilder",
750    module = "datafusion.expr",
751    subclass
752)]
753#[derive(Debug, Clone)]
754pub struct PyExprFuncBuilder {
755    pub builder: ExprFuncBuilder,
756}
757
758impl From<ExprFuncBuilder> for PyExprFuncBuilder {
759    fn from(builder: ExprFuncBuilder) -> Self {
760        Self { builder }
761    }
762}
763
764#[pymethods]
765impl PyExprFuncBuilder {
766    pub fn order_by(&self, order_by: Vec<PySortExpr>) -> PyExprFuncBuilder {
767        self.builder
768            .clone()
769            .order_by(to_sort_expressions(order_by))
770            .into()
771    }
772
773    pub fn filter(&self, filter: PyExpr) -> PyExprFuncBuilder {
774        self.builder.clone().filter(filter.expr.clone()).into()
775    }
776
777    pub fn distinct(&self) -> PyExprFuncBuilder {
778        self.builder.clone().distinct().into()
779    }
780
781    pub fn null_treatment(&self, null_treatment: NullTreatment) -> PyExprFuncBuilder {
782        self.builder
783            .clone()
784            .null_treatment(Some(null_treatment.into()))
785            .into()
786    }
787
788    pub fn partition_by(&self, partition_by: Vec<PyExpr>) -> PyExprFuncBuilder {
789        let partition_by = partition_by.iter().map(|e| e.expr.clone()).collect();
790        self.builder.clone().partition_by(partition_by).into()
791    }
792
793    pub fn window_frame(&self, window_frame: PyWindowFrame) -> PyExprFuncBuilder {
794        self.builder
795            .clone()
796            .window_frame(window_frame.into())
797            .into()
798    }
799
800    pub fn build(&self) -> PyDataFusionResult<PyExpr> {
801        Ok(self.builder.clone().build().map(|expr| expr.into())?)
802    }
803}
804
805impl PyExpr {
806    pub fn _column_name(&self, plan: &LogicalPlan) -> PyDataFusionResult<String> {
807        let field = Self::expr_to_field(&self.expr, plan)?;
808        Ok(field.name().to_owned())
809    }
810
811    /// Create a [Field] representing an [Expr], given an input [LogicalPlan] to resolve against
812    pub fn expr_to_field(expr: &Expr, input_plan: &LogicalPlan) -> PyDataFusionResult<Arc<Field>> {
813        let fields = exprlist_to_fields(std::slice::from_ref(expr), input_plan)?;
814        Ok(fields[0].1.clone())
815    }
816    fn _types(expr: &Expr) -> PyResult<DataTypeMap> {
817        match expr {
818            Expr::BinaryExpr(BinaryExpr {
819                left: _,
820                op,
821                right: _,
822            }) => match op {
823                Operator::Eq
824                | Operator::NotEq
825                | Operator::Lt
826                | Operator::LtEq
827                | Operator::Gt
828                | Operator::GtEq
829                | Operator::And
830                | Operator::Or
831                | Operator::IsDistinctFrom
832                | Operator::IsNotDistinctFrom
833                | Operator::RegexMatch
834                | Operator::RegexIMatch
835                | Operator::RegexNotMatch
836                | Operator::RegexNotIMatch
837                | Operator::LikeMatch
838                | Operator::ILikeMatch
839                | Operator::NotLikeMatch
840                | Operator::NotILikeMatch => DataTypeMap::map_from_arrow_type(&DataType::Boolean),
841                Operator::Plus | Operator::Minus | Operator::Multiply | Operator::Modulo => {
842                    DataTypeMap::map_from_arrow_type(&DataType::Int64)
843                }
844                Operator::Divide => DataTypeMap::map_from_arrow_type(&DataType::Float64),
845                Operator::StringConcat => DataTypeMap::map_from_arrow_type(&DataType::Utf8),
846                Operator::BitwiseShiftLeft
847                | Operator::BitwiseShiftRight
848                | Operator::BitwiseXor
849                | Operator::BitwiseAnd
850                | Operator::BitwiseOr => DataTypeMap::map_from_arrow_type(&DataType::Binary),
851                Operator::AtArrow
852                | Operator::ArrowAt
853                | Operator::Arrow
854                | Operator::LongArrow
855                | Operator::HashArrow
856                | Operator::HashLongArrow
857                | Operator::AtAt
858                | Operator::IntegerDivide
859                | Operator::HashMinus
860                | Operator::AtQuestion
861                | Operator::Question
862                | Operator::QuestionAnd
863                | Operator::QuestionPipe
864                | Operator::Colon => Err(py_type_err(format!("Unsupported expr: ${op}"))),
865            },
866            Expr::Cast(Cast { expr: _, field }) => {
867                DataTypeMap::map_from_arrow_type(field.data_type())
868            }
869            Expr::Literal(scalar_value, _) => DataTypeMap::map_from_scalar_value(scalar_value),
870            _ => Err(py_type_err(format!(
871                "Non Expr::Literal encountered in types: {expr:?}"
872            ))),
873        }
874    }
875}
876
877/// Initializes the `expr` module to match the pattern of `datafusion-expr` https://docs.rs/datafusion-expr/latest/datafusion_expr/
878pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
879    m.add_class::<PyExpr>()?;
880    m.add_class::<PyColumn>()?;
881    m.add_class::<PyLiteral>()?;
882    m.add_class::<PyBinaryExpr>()?;
883    m.add_class::<PyLiteral>()?;
884    m.add_class::<PyAggregateFunction>()?;
885    m.add_class::<PyNot>()?;
886    m.add_class::<PyIsNotNull>()?;
887    m.add_class::<PyIsNull>()?;
888    m.add_class::<PyIsTrue>()?;
889    m.add_class::<PyIsFalse>()?;
890    m.add_class::<PyIsUnknown>()?;
891    m.add_class::<PyIsNotTrue>()?;
892    m.add_class::<PyIsNotFalse>()?;
893    m.add_class::<PyIsNotUnknown>()?;
894    m.add_class::<PyNegative>()?;
895    m.add_class::<PyLike>()?;
896    m.add_class::<PyILike>()?;
897    m.add_class::<PySimilarTo>()?;
898    m.add_class::<PyScalarVariable>()?;
899    m.add_class::<alias::PyAlias>()?;
900    m.add_class::<in_list::PyInList>()?;
901    m.add_class::<exists::PyExists>()?;
902    m.add_class::<subquery::PySubquery>()?;
903    m.add_class::<in_subquery::PyInSubquery>()?;
904    m.add_class::<scalar_subquery::PyScalarSubquery>()?;
905    m.add_class::<placeholder::PyPlaceholder>()?;
906    m.add_class::<grouping_set::PyGroupingSet>()?;
907    m.add_class::<case::PyCase>()?;
908    m.add_class::<conditional_expr::PyCaseBuilder>()?;
909    m.add_class::<cast::PyCast>()?;
910    m.add_class::<cast::PyTryCast>()?;
911    m.add_class::<between::PyBetween>()?;
912    m.add_class::<explain::PyExplain>()?;
913    m.add_class::<limit::PyLimit>()?;
914    m.add_class::<aggregate::PyAggregate>()?;
915    m.add_class::<sort::PySort>()?;
916    m.add_class::<analyze::PyAnalyze>()?;
917    m.add_class::<empty_relation::PyEmptyRelation>()?;
918    m.add_class::<join::PyJoin>()?;
919    m.add_class::<join::PyJoinType>()?;
920    m.add_class::<join::PyJoinConstraint>()?;
921    m.add_class::<union::PyUnion>()?;
922    m.add_class::<unnest::PyUnnest>()?;
923    m.add_class::<unnest_expr::PyUnnestExpr>()?;
924    m.add_class::<higher_order_function::PyHigherOrderFunction>()?;
925    m.add_class::<lambda::PyLambda>()?;
926    m.add_class::<lambda_variable::PyLambdaVariable>()?;
927    m.add_class::<extension::PyExtension>()?;
928    m.add_class::<filter::PyFilter>()?;
929    m.add_class::<projection::PyProjection>()?;
930    m.add_class::<table_scan::PyTableScan>()?;
931    m.add_class::<create_memory_table::PyCreateMemoryTable>()?;
932    m.add_class::<create_view::PyCreateView>()?;
933    m.add_class::<distinct::PyDistinct>()?;
934    m.add_class::<sort_expr::PySortExpr>()?;
935    m.add_class::<subquery_alias::PySubqueryAlias>()?;
936    m.add_class::<drop_table::PyDropTable>()?;
937    m.add_class::<repartition::PyPartitioning>()?;
938    m.add_class::<repartition::PyRepartition>()?;
939    m.add_class::<window::PyWindowExpr>()?;
940    m.add_class::<window::PyWindowFrame>()?;
941    m.add_class::<window::PyWindowFrameBound>()?;
942    m.add_class::<copy_to::PyCopyTo>()?;
943    m.add_class::<copy_to::PyFileType>()?;
944    m.add_class::<create_catalog::PyCreateCatalog>()?;
945    m.add_class::<create_catalog_schema::PyCreateCatalogSchema>()?;
946    m.add_class::<create_external_table::PyCreateExternalTable>()?;
947    m.add_class::<create_function::PyCreateFunction>()?;
948    m.add_class::<create_function::PyOperateFunctionArg>()?;
949    m.add_class::<create_function::PyCreateFunctionBody>()?;
950    m.add_class::<create_index::PyCreateIndex>()?;
951    m.add_class::<describe_table::PyDescribeTable>()?;
952    m.add_class::<dml::PyDmlStatement>()?;
953    m.add_class::<drop_catalog_schema::PyDropCatalogSchema>()?;
954    m.add_class::<drop_function::PyDropFunction>()?;
955    m.add_class::<drop_view::PyDropView>()?;
956    m.add_class::<recursive_query::PyRecursiveQuery>()?;
957
958    m.add_class::<statement::PyTransactionStart>()?;
959    m.add_class::<statement::PyTransactionEnd>()?;
960    m.add_class::<statement::PySetVariable>()?;
961    m.add_class::<statement::PyPrepare>()?;
962    m.add_class::<statement::PyExecute>()?;
963    m.add_class::<statement::PyDeallocate>()?;
964    m.add_class::<values::PyValues>()?;
965    m.add_class::<statement::PyTransactionAccessMode>()?;
966    m.add_class::<statement::PyTransactionConclusion>()?;
967    m.add_class::<statement::PyTransactionIsolationLevel>()?;
968
969    Ok(())
970}
datafusion_python/expr.rs

datafusion_python/
expr.rs
