This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion-python.git
The following commit(s) were added to refs/heads/main by this push:
new 3b74dde Add Python wrapper for LogicalPlan::Sort (#196)
3b74dde is described below
commit 3b74ddecddccadb852a073ebae64bddf43e0011e
Author: Andy Grove <[email protected]>
AuthorDate: Sun Feb 19 11:49:24 2023 -0700
Add Python wrapper for LogicalPlan::Sort (#196)
---
datafusion/tests/test_imports.py | 3 +-
src/expr.rs | 2 +
src/expr/sort.rs | 94 ++++++++++++++++++++++++++++++++++++++++
3 files changed, 98 insertions(+), 1 deletion(-)
diff --git a/datafusion/tests/test_imports.py b/datafusion/tests/test_imports.py
index e5d9585..b7f4ef4 100644
--- a/datafusion/tests/test_imports.py
+++ b/datafusion/tests/test_imports.py
@@ -35,6 +35,7 @@ from datafusion.expr import (
Expr,
Projection,
TableScan,
+ Sort,
)
@@ -55,7 +56,7 @@ def test_class_module_is_datafusion():
]:
assert klass.__module__ == "datafusion"
- for klass in [Expr, Projection, TableScan]:
+ for klass in [Expr, Projection, TableScan, Sort]:
assert klass.__module__ == "datafusion.expr"
for klass in [DFField, DFSchema]:
diff --git a/src/expr.rs b/src/expr.rs
index f3695fe..68534bc 100644
--- a/src/expr.rs
+++ b/src/expr.rs
@@ -26,6 +26,7 @@ use datafusion::scalar::ScalarValue;
pub mod logical_node;
pub mod projection;
+pub mod sort;
pub mod table_scan;
/// A PyExpr that can be used on a DataFrame
@@ -143,5 +144,6 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> {
m.add_class::<PyExpr>()?;
m.add_class::<table_scan::PyTableScan>()?;
m.add_class::<projection::PyProjection>()?;
+ m.add_class::<sort::PySort>()?;
Ok(())
}
diff --git a/src/expr/sort.rs b/src/expr/sort.rs
new file mode 100644
index 0000000..1d0a7f6
--- /dev/null
+++ b/src/expr/sort.rs
@@ -0,0 +1,94 @@
+// 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 datafusion_common::DataFusionError;
+use datafusion_expr::logical_plan::Sort;
+use pyo3::prelude::*;
+use std::fmt::{self, Display, Formatter};
+
+use crate::common::df_schema::PyDFSchema;
+use crate::expr::logical_node::LogicalNode;
+use crate::expr::PyExpr;
+use crate::sql::logical::PyLogicalPlan;
+
+#[pyclass(name = "Sort", module = "datafusion.expr", subclass)]
+#[derive(Clone)]
+pub struct PySort {
+ sort: Sort,
+}
+
+impl From<Sort> for PySort {
+ fn from(sort: Sort) -> PySort {
+ PySort { sort }
+ }
+}
+
+impl TryFrom<PySort> for Sort {
+ type Error = DataFusionError;
+
+ fn try_from(agg: PySort) -> Result<Self, Self::Error> {
+ Ok(agg.sort)
+ }
+}
+
+impl Display for PySort {
+ fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+ write!(
+ f,
+ "Sort
+ \nExpr(s): {:?}
+ \nInput: {:?}
+ \nSchema: {:?}",
+ &self.sort.expr,
+ self.sort.input,
+ self.sort.input.schema()
+ )
+ }
+}
+
+#[pymethods]
+impl PySort {
+ /// Retrieves the sort expressions for this `Sort`
+ fn sort_exprs(&self) -> PyResult<Vec<PyExpr>> {
+ Ok(self
+ .sort
+ .expr
+ .iter()
+ .map(|e| PyExpr::from(e.clone()))
+ .collect())
+ }
+
+ /// Retrieves the input `LogicalPlan` to this `Sort` node
+ fn input(&self) -> PyLogicalPlan {
+ PyLogicalPlan::from((*self.sort.input).clone())
+ }
+
+ /// Resulting Schema for this `Sort` node instance
+ fn schema(&self) -> PyDFSchema {
+ self.sort.input.schema().as_ref().clone().into()
+ }
+
+ fn __repr__(&self) -> PyResult<String> {
+ Ok(format!("Sort({})", self))
+ }
+}
+
+impl LogicalNode for PySort {
+ fn input(&self) -> Vec<PyLogicalPlan> {
+ vec![PyLogicalPlan::from((*self.sort.input).clone())]
+ }
+}
