diff --git a/src/expr.rs b/src/expr.rs index f2209b90..c55351ef 100644 --- a/src/expr.rs +++ b/src/expr.rs @@ -122,7 +122,7 @@ impl PyExpr { /// Return the specific expression fn to_variant(&self, py: Python) -> PyResult { Python::with_gil(|_| match &self.expr { - Expr::Alias(alias) => Ok(PyAlias::new(&alias.expr, &alias.name).into_py(py)), + Expr::Alias(alias) => Ok(PyAlias::from(alias.clone()).into_py(py)), Expr::Column(col) => Ok(PyColumn::from(col.clone()).into_py(py)), Expr::ScalarVariable(data_type, variables) => { Ok(PyScalarVariable::new(data_type, variables).into_py(py)) diff --git a/src/expr/alias.rs b/src/expr/alias.rs index 2ce65634..3208800a 100644 --- a/src/expr/alias.rs +++ b/src/expr/alias.rs @@ -19,13 +19,24 @@ use crate::expr::PyExpr; use pyo3::prelude::*; use std::fmt::{self, Display, Formatter}; -use datafusion_expr::Expr; +use datafusion_expr::expr::Alias; #[pyclass(name = "Alias", module = "datafusion.expr", subclass)] #[derive(Clone)] pub struct PyAlias { - expr: PyExpr, - alias_name: String, + alias: Alias, +} + +impl From for PyAlias { + fn from(alias: Alias) -> Self { + Self { alias } + } +} + +impl From for Alias { + fn from(py_alias: PyAlias) -> Self { + py_alias.alias + } } impl Display for PyAlias { @@ -35,29 +46,20 @@ impl Display for PyAlias { "Alias \nExpr: `{:?}` \nAlias Name: `{}`", - &self.expr, &self.alias_name + &self.alias.expr, &self.alias.name ) } } -impl PyAlias { - pub fn new(expr: &Expr, alias_name: &String) -> Self { - Self { - expr: expr.clone().into(), - alias_name: alias_name.to_owned(), - } - } -} - #[pymethods] impl PyAlias { /// Retrieve the "name" of the alias fn alias(&self) -> PyResult { - Ok(self.alias_name.clone()) + Ok(self.alias.name.clone()) } fn expr(&self) -> PyResult { - Ok(self.expr.clone()) + Ok((*self.alias.expr.clone()).into()) } /// Get a String representation of this column