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 b8ef9bf Add Python wrapper for LogicalPlan::Limit (#193)
b8ef9bf is described below
commit b8ef9bf54f148fe256ae2f67382f6625e6ceff20
Author: Andy Grove <[email protected]>
AuthorDate: Sun Feb 19 12:50:54 2023 -0700
Add Python wrapper for LogicalPlan::Limit (#193)
---
datafusion/tests/test_imports.py | 3 +-
src/expr.rs | 2 +
src/expr/limit.rs | 88 ++++++++++++++++++++++++++++++++++++++++
src/expr/projection.rs | 39 ++++--------------
4 files changed, 101 insertions(+), 31 deletions(-)
diff --git a/datafusion/tests/test_imports.py b/datafusion/tests/test_imports.py
index 5a6c685..ee47b0e 100644
--- a/datafusion/tests/test_imports.py
+++ b/datafusion/tests/test_imports.py
@@ -35,6 +35,7 @@ from datafusion.expr import (
Expr,
Projection,
TableScan,
+ Limit,
Aggregate,
Sort,
)
@@ -57,7 +58,7 @@ def test_class_module_is_datafusion():
]:
assert klass.__module__ == "datafusion"
- for klass in [Expr, Projection, TableScan, Aggregate, Sort]:
+ for klass in [Expr, Projection, TableScan, Aggregate, Sort, Limit]:
assert klass.__module__ == "datafusion.expr"
for klass in [DFField, DFSchema]:
diff --git a/src/expr.rs b/src/expr.rs
index 15359d4..7ef9407 100644
--- a/src/expr.rs
+++ b/src/expr.rs
@@ -25,6 +25,7 @@ use datafusion_expr::{col, lit, Cast, Expr, GetIndexedField};
use datafusion::scalar::ScalarValue;
pub mod aggregate;
+pub mod limit;
pub mod logical_node;
pub mod projection;
pub mod sort;
@@ -145,6 +146,7 @@ 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::<limit::PyLimit>()?;
m.add_class::<aggregate::PyAggregate>()?;
m.add_class::<sort::PySort>()?;
Ok(())
diff --git a/src/expr/limit.rs b/src/expr/limit.rs
new file mode 100644
index 0000000..a50e5b8
--- /dev/null
+++ b/src/expr/limit.rs
@@ -0,0 +1,88 @@
+// 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_expr::logical_plan::Limit;
+use pyo3::prelude::*;
+use std::fmt::{self, Display, Formatter};
+
+use crate::common::df_schema::PyDFSchema;
+use crate::expr::logical_node::LogicalNode;
+use crate::sql::logical::PyLogicalPlan;
+
+#[pyclass(name = "Limit", module = "datafusion.expr", subclass)]
+#[derive(Clone)]
+pub struct PyLimit {
+ limit: Limit,
+}
+
+impl From<Limit> for PyLimit {
+ fn from(limit: Limit) -> PyLimit {
+ PyLimit { limit }
+ }
+}
+
+impl From<PyLimit> for Limit {
+ fn from(limit: PyLimit) -> Self {
+ limit.limit
+ }
+}
+
+impl Display for PyLimit {
+ fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+ write!(
+ f,
+ "Limit
+ \nSkip: {}
+ \nFetch: {:?}
+ \nInput: {:?}",
+ &self.limit.skip, &self.limit.fetch, &self.limit.input
+ )
+ }
+}
+
+#[pymethods]
+impl PyLimit {
+ /// Retrieves the skip value for this `Limit`
+ fn skip(&self) -> usize {
+ self.limit.skip
+ }
+
+ /// Retrieves the fetch value for this `Limit`
+ fn fetch(&self) -> Option<usize> {
+ self.limit.fetch
+ }
+
+ /// Retrieves the input `LogicalPlan` to this `Limit` node
+ fn input(&self) -> PyLogicalPlan {
+ PyLogicalPlan::from((*self.limit.input).clone())
+ }
+
+ /// Resulting Schema for this `Limit` node instance
+ fn schema(&self) -> PyResult<PyDFSchema> {
+ Ok(self.limit.input.schema().as_ref().clone().into())
+ }
+
+ fn __repr__(&self) -> PyResult<String> {
+ Ok(format!("Limit({})", self))
+ }
+}
+
+impl LogicalNode for PyLimit {
+ fn input(&self) -> Vec<PyLogicalPlan> {
+ vec![PyLogicalPlan::from((*self.limit.input).clone())]
+ }
+}
diff --git a/src/expr/projection.rs b/src/expr/projection.rs
index 6d04e59..2d43632 100644
--- a/src/expr/projection.rs
+++ b/src/expr/projection.rs
@@ -15,13 +15,11 @@
// specific language governing permissions and limitations
// under the License.
-use datafusion_common::DataFusionError;
use datafusion_expr::logical_plan::Projection;
use pyo3::prelude::*;
use std::fmt::{self, Display, Formatter};
use crate::common::df_schema::PyDFSchema;
-use crate::errors::py_runtime_err;
use crate::expr::logical_node::LogicalNode;
use crate::expr::PyExpr;
use crate::sql::logical::PyLogicalPlan;
@@ -38,15 +36,9 @@ impl From<Projection> for PyProjection {
}
}
-impl TryFrom<PyProjection> for Projection {
- type Error = DataFusionError;
-
- fn try_from(py_proj: PyProjection) -> Result<Self, Self::Error> {
- Projection::try_new_with_schema(
- py_proj.projection.expr,
- py_proj.projection.input.clone(),
- py_proj.projection.schema,
- )
+impl From<PyProjection> for Projection {
+ fn from(proj: PyProjection) -> Self {
+ proj.projection
}
}
@@ -66,8 +58,7 @@ impl Display for PyProjection {
#[pymethods]
impl PyProjection {
/// Retrieves the expressions for this `Projection`
- #[pyo3(name = "projections")]
- fn py_projections(&self) -> PyResult<Vec<PyExpr>> {
+ fn projections(&self) -> PyResult<Vec<PyExpr>> {
Ok(self
.projection
.expr
@@ -76,25 +67,13 @@ impl PyProjection {
.collect())
}
- // Retrieves the input `LogicalPlan` to this `Projection` node
- #[pyo3(name = "input")]
- fn py_input(&self) -> PyResult<PyLogicalPlan> {
- // DataFusion make a loose guarantee that each Projection should have
an input, however
- // we check for that hear since we are performing explicit index
retrieval
- let inputs = LogicalNode::input(self);
- if !inputs.is_empty() {
- return Ok(inputs[0].clone());
- }
-
- Err(py_runtime_err(format!(
- "Expected `input` field for Projection node: {}",
- self
- )))
+ /// Retrieves the input `LogicalPlan` to this `Projection` node
+ fn input(&self) -> PyLogicalPlan {
+ PyLogicalPlan::from((*self.projection.input).clone())
}
- // Resulting Schema for this `Projection` node instance
- #[pyo3(name = "schema")]
- fn py_schema(&self) -> PyResult<PyDFSchema> {
+ /// Resulting Schema for this `Projection` node instance
+ fn schema(&self) -> PyResult<PyDFSchema> {
Ok((*self.projection.schema).clone().into())
}
