Add Python wrapper for LogicalPlan::Limit (#193)

datafusion/tests/test_imports.py

@@ -35,6 +35,7 @@
     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]:

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(())

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())]
+    }
+}

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())
     }