Explain bindings (#264)

* Introduce to_variant trait function to LogicalNode and create Explain LogicalNode bindings

* Cargo fmt

* Add missing classes to list of exports so test_imports will pass

datafusion/__init__.py

@@ -76,6 +76,7 @@
     Cast,
     TryCast,
     Between,
+    Explain,
 )
 
 __version__ = importlib_metadata.version(__name__)
@@ -127,6 +128,7 @@
     "Cast",
     "TryCast",
     "Between",
+    "Explain",
 ]
 
 

src/errors.rs

@@ -82,3 +82,7 @@ pub fn py_runtime_err(e: impl Debug) -> PyErr {
 pub fn py_datafusion_err(e: impl Debug) -> PyErr {
     PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("{e:?}"))
 }
+
+pub fn py_unsupported_variant_err(e: impl Debug) -> PyErr {
+    PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("{e:?}"))
+}

src/expr.rs

@@ -50,6 +50,7 @@ pub mod column;
 pub mod cross_join;
 pub mod empty_relation;
 pub mod exists;
+pub mod explain;
 pub mod filter;
 pub mod grouping_set;
 pub mod in_list;
@@ -260,10 +261,7 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> {
     m.add_class::<cast::PyTryCast>()?;
     m.add_class::<between::PyBetween>()?;
     m.add_class::<indexed_field::PyGetIndexedField>()?;
-    // operators
-    m.add_class::<table_scan::PyTableScan>()?;
-    m.add_class::<projection::PyProjection>()?;
-    m.add_class::<filter::PyFilter>()?;
+    m.add_class::<explain::PyExplain>()?;
     m.add_class::<limit::PyLimit>()?;
     m.add_class::<aggregate::PyAggregate>()?;
     m.add_class::<sort::PySort>()?;
@@ -274,5 +272,8 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> {
     m.add_class::<join::PyJoinConstraint>()?;
     m.add_class::<cross_join::PyCrossJoin>()?;
     m.add_class::<union::PyUnion>()?;
+    m.add_class::<filter::PyFilter>()?;
+    m.add_class::<projection::PyProjection>()?;
+    m.add_class::<table_scan::PyTableScan>()?;
     Ok(())
 }

src/expr/aggregate.rs

@@ -20,8 +20,8 @@ use datafusion_expr::logical_plan::Aggregate;
 use pyo3::prelude::*;
 use std::fmt::{self, Display, Formatter};
 
+use super::logical_node::LogicalNode;
 use crate::common::df_schema::PyDFSchema;
-use crate::expr::logical_node::LogicalNode;
 use crate::expr::PyExpr;
 use crate::sql::logical::PyLogicalPlan;
 
@@ -103,4 +103,8 @@ impl LogicalNode for PyAggregate {
     fn inputs(&self) -> Vec<PyLogicalPlan> {
         vec![PyLogicalPlan::from((*self.aggregate.input).clone())]
     }
+
+    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
+        Ok(self.clone().into_py(py))
+    }
 }

src/expr/analyze.rs

@@ -19,8 +19,8 @@ use datafusion_expr::logical_plan::Analyze;
 use pyo3::prelude::*;
 use std::fmt::{self, Display, Formatter};
 
+use super::logical_node::LogicalNode;
 use crate::common::df_schema::PyDFSchema;
-use crate::expr::logical_node::LogicalNode;
 use crate::sql::logical::PyLogicalPlan;
 
 #[pyclass(name = "Analyze", module = "datafusion.expr", subclass)]
@@ -77,4 +77,8 @@ impl LogicalNode for PyAnalyze {
     fn inputs(&self) -> Vec<PyLogicalPlan> {
         vec![PyLogicalPlan::from((*self.analyze.input).clone())]
     }
+
+    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
+        Ok(self.clone().into_py(py))
+    }
 }

src/expr/cross_join.rs

@@ -19,8 +19,8 @@ use datafusion_expr::logical_plan::CrossJoin;
 use pyo3::prelude::*;
 use std::fmt::{self, Display, Formatter};
 
+use super::logical_node::LogicalNode;
 use crate::common::df_schema::PyDFSchema;
-use crate::expr::logical_node::LogicalNode;
 use crate::sql::logical::PyLogicalPlan;
 
 #[pyclass(name = "CrossJoin", module = "datafusion.expr", subclass)]
@@ -87,4 +87,8 @@ impl LogicalNode for PyCrossJoin {
             PyLogicalPlan::from((*self.cross_join.right).clone()),
         ]
     }
+
+    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
+        Ok(self.clone().into_py(py))
+    }
 }

src/expr/empty_relation.rs

@@ -15,11 +15,13 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::common::df_schema::PyDFSchema;
+use crate::{common::df_schema::PyDFSchema, sql::logical::PyLogicalPlan};
 use datafusion_expr::EmptyRelation;
 use pyo3::prelude::*;
 use std::fmt::{self, Display, Formatter};
 
+use super::logical_node::LogicalNode;
+
 #[pyclass(name = "EmptyRelation", module = "datafusion.expr", subclass)]
 #[derive(Clone)]
 pub struct PyEmptyRelation {
@@ -70,3 +72,14 @@ impl PyEmptyRelation {
         Ok("EmptyRelation".to_string())
     }
 }
+
+impl LogicalNode for PyEmptyRelation {
+    fn inputs(&self) -> Vec<PyLogicalPlan> {
+        // table scans are leaf nodes and do not have inputs
+        vec![]
+    }
+
+    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
+        Ok(self.clone().into_py(py))
+    }
+}

src/expr/explain.rs

@@ -0,0 +1,110 @@
+// 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::{logical_plan::Explain, LogicalPlan};
+use pyo3::prelude::*;
+
+use crate::{common::df_schema::PyDFSchema, errors::py_type_err, sql::logical::PyLogicalPlan};
+
+use super::logical_node::LogicalNode;
+
+#[pyclass(name = "Explain", module = "datafusion.expr", subclass)]
+#[derive(Clone)]
+pub struct PyExplain {
+    explain: Explain,
+}
+
+impl From<PyExplain> for Explain {
+    fn from(explain: PyExplain) -> Self {
+        explain.explain
+    }
+}
+
+impl From<Explain> for PyExplain {
+    fn from(explain: Explain) -> PyExplain {
+        PyExplain { explain }
+    }
+}
+
+impl Display for PyExplain {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        write!(
+            f,
+            "Explain
+            verbose: {:?}
+            plan: {:?}
+            stringified_plans: {:?}
+            schema: {:?}
+            logical_optimization_succeeded: {:?}",
+            &self.explain.verbose,
+            &self.explain.plan,
+            &self.explain.stringified_plans,
+            &self.explain.schema,
+            &self.explain.logical_optimization_succeeded
+        )
+    }
+}
+
+#[pymethods]
+impl PyExplain {
+    fn explain_string(&self) -> PyResult<Vec<String>> {
+        let mut string_plans: Vec<String> = Vec::new();
+        for stringified_plan in &self.explain.stringified_plans {
+            string_plans.push((*stringified_plan.plan).clone());
+        }
+        Ok(string_plans)
+    }
+
+    fn verbose(&self) -> bool {
+        self.explain.verbose
+    }
+
+    fn plan(&self) -> PyResult<PyLogicalPlan> {
+        Ok(PyLogicalPlan::from((*self.explain.plan).clone()))
+    }
+
+    fn schema(&self) -> PyDFSchema {
+        (*self.explain.schema).clone().into()
+    }
+
+    fn logical_optimization_succceeded(&self) -> bool {
+        self.explain.logical_optimization_succeeded
+    }
+}
+
+impl TryFrom<LogicalPlan> for PyExplain {
+    type Error = PyErr;
+
+    fn try_from(logical_plan: LogicalPlan) -> Result<Self, Self::Error> {
+        match logical_plan {
+            LogicalPlan::Explain(explain) => Ok(PyExplain { explain }),
+            _ => Err(py_type_err("unexpected plan")),
+        }
+    }
+}
+
+impl LogicalNode for PyExplain {
+    fn inputs(&self) -> Vec<PyLogicalPlan> {
+        vec![]
+    }
+
+    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
+        Ok(self.clone().into_py(py))
+    }
+}

src/expr/filter.rs

@@ -47,8 +47,8 @@ impl Display for PyFilter {
         write!(
             f,
             "Filter
-            \nPredicate: {:?}
-            \nInput: {:?}",
+            Predicate: {:?}
+            Input: {:?}",
             &self.filter.predicate, &self.filter.input
         )
     }
@@ -80,4 +80,8 @@ impl LogicalNode for PyFilter {
     fn inputs(&self) -> Vec<PyLogicalPlan> {
         vec![PyLogicalPlan::from((*self.filter.input).clone())]
     }
+
+    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
+        Ok(self.clone().into_py(py))
+    }
 }

src/expr/join.rs

@@ -95,14 +95,14 @@ impl Display for PyJoin {
         write!(
             f,
             "Join
-            \nLeft: {:?}
-            \nRight: {:?}
-            \nOn: {:?}
-            \nFilter: {:?}
-            \nJoinType: {:?}
-            \nJoinConstraint: {:?}
-            \nSchema: {:?}
-            \nNullEqualsNull: {:?}",
+            Left: {:?}
+            Right: {:?}
+            On: {:?}
+            Filter: {:?}
+            JoinType: {:?}
+            JoinConstraint: {:?}
+            Schema: {:?}
+            NullEqualsNull: {:?}",
             &self.join.left,
             &self.join.right,
             &self.join.on,
@@ -178,4 +178,8 @@ impl LogicalNode for PyJoin {
             PyLogicalPlan::from((*self.join.right).clone()),
         ]
     }
+
+    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
+        Ok(self.clone().into_py(py))
+    }
 }

src/expr/limit.rs

@@ -46,9 +46,9 @@ impl Display for PyLimit {
         write!(
             f,
             "Limit
-            \nSkip: {}
-            \nFetch: {:?}
-            \nInput: {:?}",
+            Skip: {}
+            Fetch: {:?}
+            Input: {:?}",
             &self.limit.skip, &self.limit.fetch, &self.limit.input
         )
     }
@@ -85,4 +85,8 @@ impl LogicalNode for PyLimit {
     fn inputs(&self) -> Vec<PyLogicalPlan> {
         vec![PyLogicalPlan::from((*self.limit.input).clone())]
     }
+
+    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
+        Ok(self.clone().into_py(py))
+    }
 }

src/expr/logical_node.rs

@@ -15,11 +15,15 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use pyo3::{PyObject, PyResult, Python};
+
 use crate::sql::logical::PyLogicalPlan;
 
 /// Representation of a `LogicalNode` in the in overall `LogicalPlan`
 /// any "node" shares these common traits in common.
 pub trait LogicalNode {
     /// The input plan to the current logical node instance.
     fn inputs(&self) -> Vec<PyLogicalPlan>;
+
+    fn to_variant(&self, py: Python) -> PyResult<PyObject>;
 }

src/expr/projection.rs

@@ -96,4 +96,8 @@ impl LogicalNode for PyProjection {
     fn inputs(&self) -> Vec<PyLogicalPlan> {
         vec![PyLogicalPlan::from((*self.projection.input).clone())]
     }
+
+    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
+        Ok(self.clone().into_py(py))
+    }
 }

src/expr/sort.rs

@@ -91,4 +91,8 @@ impl LogicalNode for PySort {
     fn inputs(&self) -> Vec<PyLogicalPlan> {
         vec![PyLogicalPlan::from((*self.sort.input).clone())]
     }
+
+    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
+        Ok(self.clone().into_py(py))
+    }
 }

src/expr/table_scan.rs

@@ -52,9 +52,9 @@ impl Display for PyTableScan {
         write!(
             f,
             "TableScan\nTable Name: {}
-            \nProjections: {:?}
-            \nProjected Schema: {:?}
-            \nFilters: {:?}",
+            Projections: {:?}
+            Projected Schema: {:?}
+            Filters: {:?}",
             &self.table_scan.table_name,
             &self.py_projections(),
             &self.py_schema(),
@@ -131,4 +131,8 @@ impl LogicalNode for PyTableScan {
         // table scans are leaf nodes and do not have inputs
         vec![]
     }
+
+    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
+        Ok(self.clone().into_py(py))
+    }
 }