From 93b99df434d93cbbfc57551dd5f99a586e186084 Mon Sep 17 00:00:00 2001 From: Jeremy Dyer Date: Sun, 12 Mar 2023 15:32:40 -0400 Subject: [PATCH 1/3] Subquey alias bindings --- datafusion/__init__.py | 2 + datafusion/tests/test_dataframe.py | 1 + datafusion/tests/test_imports.py | 2 + src/expr.rs | 2 + src/expr/subquery_alias.rs | 91 ++++++++++++++++++++++++++++++ 5 files changed, 98 insertions(+) create mode 100644 src/expr/subquery_alias.rs diff --git a/datafusion/__init__.py b/datafusion/__init__.py index 4956d9d30..412027232 100644 --- a/datafusion/__init__.py +++ b/datafusion/__init__.py @@ -77,6 +77,7 @@ TryCast, Between, Explain, + SubqueryAlias, ) __version__ = importlib_metadata.version(__name__) @@ -129,6 +130,7 @@ "TryCast", "Between", "Explain", + "SubqueryAlias", ] diff --git a/datafusion/tests/test_dataframe.py b/datafusion/tests/test_dataframe.py index b87b1f14a..42d90b339 100644 --- a/datafusion/tests/test_dataframe.py +++ b/datafusion/tests/test_dataframe.py @@ -352,6 +352,7 @@ def test_optimized_logical_plan(aggregate_df): expected = "Projection: test.c1, SUM(test.c2)" + breakpoint() assert expected == plan.display() expected = ( diff --git a/datafusion/tests/test_imports.py b/datafusion/tests/test_imports.py index 77309e2bc..30eb4eac8 100644 --- a/datafusion/tests/test_imports.py +++ b/datafusion/tests/test_imports.py @@ -77,6 +77,7 @@ Cast, TryCast, Between, + SubqueryAlias, ) @@ -143,6 +144,7 @@ def test_class_module_is_datafusion(): Cast, TryCast, Between, + SubqueryAlias, ]: assert klass.__module__ == "datafusion.expr" diff --git a/src/expr.rs b/src/expr.rs index 33f36f599..608aff282 100644 --- a/src/expr.rs +++ b/src/expr.rs @@ -69,6 +69,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; @@ -275,5 +276,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..8466e8968 --- /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 + \nInputs(s): {:?} + \nAlias: {:?} + \nSchema: {:?}", + 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!("Projection({})", self)) + } + + fn __name__(&self) -> PyResult { + Ok("Projection".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)) + } +} From b14852b0926d03920f84c4e8049d24e21faf6897 Mon Sep 17 00:00:00 2001 From: Jeremy Dyer Date: Sun, 12 Mar 2023 15:35:15 -0400 Subject: [PATCH 2/3] newline formatting for subquery alias Display function --- datafusion/tests/test_dataframe.py | 1 - 1 file changed, 1 deletion(-) diff --git a/datafusion/tests/test_dataframe.py b/datafusion/tests/test_dataframe.py index 42d90b339..b87b1f14a 100644 --- a/datafusion/tests/test_dataframe.py +++ b/datafusion/tests/test_dataframe.py @@ -352,7 +352,6 @@ def test_optimized_logical_plan(aggregate_df): expected = "Projection: test.c1, SUM(test.c2)" - breakpoint() assert expected == plan.display() expected = ( From b65e9732906321391c901cb2b16eef62d1b37e41 Mon Sep 17 00:00:00 2001 From: Jeremy Dyer Date: Mon, 13 Mar 2023 11:03:14 -0400 Subject: [PATCH 3/3] reorder imports --- datafusion/tests/test_imports.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/tests/test_imports.py b/datafusion/tests/test_imports.py index 30eb4eac8..3615ea4b9 100644 --- a/datafusion/tests/test_imports.py +++ b/datafusion/tests/test_imports.py @@ -76,8 +76,8 @@ Case, Cast, TryCast, - Between, SubqueryAlias, + Between, )