diff --git a/datafusion/__init__.py b/datafusion/__init__.py index 578e7dc05..367cda382 100644 --- a/datafusion/__init__.py +++ b/datafusion/__init__.py @@ -77,6 +77,7 @@ TryCast, Between, Explain, + SubqueryAlias, Extension, ) @@ -130,6 +131,7 @@ "TryCast", "Between", "Explain", + "SubqueryAlias", "Extension", ] diff --git a/datafusion/tests/test_imports.py b/datafusion/tests/test_imports.py index 907544627..0e3a122bd 100644 --- a/datafusion/tests/test_imports.py +++ b/datafusion/tests/test_imports.py @@ -76,6 +76,7 @@ Case, Cast, TryCast, + SubqueryAlias, Between, Explain, Extension, @@ -145,6 +146,7 @@ def test_class_module_is_datafusion(): Cast, TryCast, Between, + SubqueryAlias, Explain, Extension, ]: diff --git a/src/expr.rs b/src/expr.rs index 1d6c99932..9dea30da8 100644 --- a/src/expr.rs +++ b/src/expr.rs @@ -70,6 +70,7 @@ pub mod scalar_variable; pub mod signature; pub mod sort; pub mod subquery; +pub mod subquery_alias; pub mod table_scan; pub mod union; @@ -277,5 +278,6 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/src/expr/subquery_alias.rs b/src/expr/subquery_alias.rs new file mode 100644 index 000000000..5360bbbc4 --- /dev/null +++ b/src/expr/subquery_alias.rs @@ -0,0 +1,91 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt::{self, Display, Formatter}; + +use datafusion_expr::SubqueryAlias; +use pyo3::prelude::*; + +use crate::{common::df_schema::PyDFSchema, sql::logical::PyLogicalPlan}; + +use super::logical_node::LogicalNode; + +#[pyclass(name = "SubqueryAlias", module = "datafusion.expr", subclass)] +#[derive(Clone)] +pub struct PySubqueryAlias { + subquery_alias: SubqueryAlias, +} + +impl From for SubqueryAlias { + fn from(subquery_alias: PySubqueryAlias) -> Self { + subquery_alias.subquery_alias + } +} + +impl From for PySubqueryAlias { + fn from(subquery_alias: SubqueryAlias) -> PySubqueryAlias { + PySubqueryAlias { subquery_alias } + } +} + +impl Display for PySubqueryAlias { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!( + f, + "SubqueryAlias + Inputs(s): {:?} + Alias: {:?} + Schema: {:?}", + self.subquery_alias.input, self.subquery_alias.alias, self.subquery_alias.schema, + ) + } +} + +#[pymethods] +impl PySubqueryAlias { + /// Retrieves the input `LogicalPlan` to this `Projection` node + fn input(&self) -> PyResult> { + Ok(Self::inputs(self)) + } + + /// Resulting Schema for this `Projection` node instance + fn schema(&self) -> PyResult { + Ok((*self.subquery_alias.schema).clone().into()) + } + + fn alias(&self) -> PyResult { + Ok(self.subquery_alias.alias.clone()) + } + + fn __repr__(&self) -> PyResult { + Ok(format!("SubqueryAlias({})", self)) + } + + fn __name__(&self) -> PyResult { + Ok("SubqueryAlias".to_string()) + } +} + +impl LogicalNode for PySubqueryAlias { + fn inputs(&self) -> Vec { + vec![PyLogicalPlan::from((*self.subquery_alias.input).clone())] + } + + fn to_variant(&self, py: Python) -> PyResult { + Ok(self.clone().into_py(py)) + } +}