datafusion_python/expr/
projection.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::fmt::{self, Display, Formatter};
19
20use datafusion::logical_expr::Expr;
21use datafusion::logical_expr::logical_plan::Projection;
22use pyo3::IntoPyObjectExt;
23use pyo3::prelude::*;
24
25use crate::common::df_schema::PyDFSchema;
26use crate::expr::PyExpr;
27use crate::expr::logical_node::LogicalNode;
28use crate::sql::logical::PyLogicalPlan;
29
30#[pyclass(
31    from_py_object,
32    frozen,
33    name = "Projection",
34    module = "datafusion.expr",
35    subclass
36)]
37#[derive(Clone)]
38pub struct PyProjection {
39    pub projection: Projection,
40}
41
42impl PyProjection {
43    pub fn new(projection: Projection) -> Self {
44        Self { projection }
45    }
46}
47
48impl From<Projection> for PyProjection {
49    fn from(projection: Projection) -> PyProjection {
50        PyProjection { projection }
51    }
52}
53
54impl From<PyProjection> for Projection {
55    fn from(proj: PyProjection) -> Self {
56        proj.projection
57    }
58}
59
60impl Display for PyProjection {
61    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
62        write!(
63            f,
64            "Projection
65            \nExpr(s): {:?}
66            \nInput: {:?}
67            \nProjected Schema: {:?}",
68            &self.projection.expr, &self.projection.input, &self.projection.schema,
69        )
70    }
71}
72
73#[pymethods]
74impl PyProjection {
75    /// Retrieves the expressions for this `Projection`
76    fn projections(&self) -> PyResult<Vec<PyExpr>> {
77        Ok(self
78            .projection
79            .expr
80            .iter()
81            .map(|e| PyExpr::from(e.clone()))
82            .collect())
83    }
84
85    /// Retrieves the input `LogicalPlan` to this `Projection` node
86    fn input(&self) -> PyResult<Vec<PyLogicalPlan>> {
87        Ok(Self::inputs(self))
88    }
89
90    /// Resulting Schema for this `Projection` node instance
91    fn schema(&self) -> PyResult<PyDFSchema> {
92        Ok((*self.projection.schema).clone().into())
93    }
94
95    fn __repr__(&self) -> PyResult<String> {
96        Ok(format!("Projection({self})"))
97    }
98
99    fn __name__(&self) -> PyResult<String> {
100        Ok("Projection".to_string())
101    }
102}
103
104impl PyProjection {
105    /// Projection: Gets the names of the fields that should be projected
106    pub fn projected_expressions(local_expr: &PyExpr) -> Vec<PyExpr> {
107        let mut projs: Vec<PyExpr> = Vec::new();
108        match &local_expr.expr {
109            Expr::Alias(alias) => {
110                let py_expr: PyExpr = PyExpr::from(*alias.expr.clone());
111                projs.extend_from_slice(Self::projected_expressions(&py_expr).as_slice());
112            }
113            _ => projs.push(local_expr.clone()),
114        }
115        projs
116    }
117}
118
119impl LogicalNode for PyProjection {
120    fn inputs(&self) -> Vec<PyLogicalPlan> {
121        vec![PyLogicalPlan::from((*self.projection.input).clone())]
122    }
123
124    fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
125        self.clone().into_bound_py_any(py)
126    }
127}
datafusion_python/expr/projection.rs

datafusion_python/expr/
projection.rs
