Add DataFrame methods for accessing plans

datafusion/tests/test_dataframe.py

@@ -258,6 +258,21 @@ def test_explain(df):
  df.explain()
 
 
+def test_logical_plan(df):
+ plan = df.logical_plan()
+ assert plan is not None
+
+
+def test_optimized_logical_plan(df):
+ plan = df.optimized_logical_plan()
+ assert plan is not None
+
+
+def test_execution_plan(df):
+ plan = df.execution_plan()
+ assert plan is not None
+
+
 def test_repartition(df):
  df.repartition(2)
 

src/dataframe.rs

@@ -15,6 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::logical::PyLogicalPlan;
+use crate::physical_plan::PyExecutionPlan;
 use crate::utils::wait_for_future;
 use crate::{errors::DataFusionError, expression::PyExpr};
 use datafusion::arrow::datatypes::Schema;
@@ -202,6 +204,22 @@ impl PyDataFrame {
  pretty::print_batches(&batches).map_err(|err| PyArrowException::new_err(err.to_string()))
  }
 
+ /// Get the logical plan for this `DataFrame`
+ fn logical_plan(&self) -> PyResult<PyLogicalPlan> {
+ Ok(self.df.as_ref().clone().into_optimized_plan()?.into())
+ }
+
+ /// Get the optimized logical plan for this `DataFrame`
+ fn optimized_logical_plan(&self) -> PyResult<PyLogicalPlan> {
+ Ok(self.df.as_ref().clone().into_optimized_plan()?.into())
+ }
+
+ /// Get the execution plan for this `DataFrame`
+ fn execution_plan(&self, py: Python) -> PyResult<PyExecutionPlan> {
+ let plan = wait_for_future(py, self.df.as_ref().clone().create_physical_plan())?;
+ Ok(plan.into())
+ }
+
  /// Repartition a `DataFrame` based on a logical partitioning scheme.
  fn repartition(&self, num: usize) -> PyResult<Self> {
  let new_df = self

src/lib.rs

@@ -35,6 +35,7 @@ mod expression;
 #[allow(clippy::borrow_deref_ref)]
 mod functions;
 pub mod logical;
+pub mod physical_plan;
 mod pyarrow_filter_expression;
 pub mod store;
 pub mod substrait;
@@ -65,6 +66,7 @@ fn _internal(py: Python, m: &PyModule) -> PyResult<()> {
  m.add_class::<udaf::PyAggregateUDF>()?;
  m.add_class::<config::PyConfig>()?;
  m.add_class::<logical::PyLogicalPlan>()?;
+ m.add_class::<physical_plan::PyExecutionPlan>()?;
 
  // Register the functions as a submodule
  let funcs = PyModule::new(py, "functions")?;

src/logical.rs

@@ -20,7 +20,7 @@ use std::sync::Arc;
 use datafusion_expr::LogicalPlan;
 use pyo3::prelude::*;
 
-#[pyclass(name = "LogicalPlan", module = "substrait", subclass)]
+#[pyclass(name = "LogicalPlan", module = "datafusion", subclass)]
 #[derive(Debug, Clone)]
 pub struct PyLogicalPlan {
  pub(crate) plan: Arc<LogicalPlan>,

src/physical_plan.rs

@@ -0,0 +1,46 @@
+// 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::physical_plan::ExecutionPlan;
+use std::sync::Arc;
+
+use pyo3::prelude::*;
+
+#[pyclass(name = "ExecutionPlan", module = "datafusion", subclass)]
+#[derive(Debug, Clone)]
+pub struct PyExecutionPlan {
+ pub(crate) plan: Arc<dyn ExecutionPlan>,
+}
+
+impl PyExecutionPlan {
+ /// creates a new PyPhysicalPlan
+ pub fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
+ Self { plan }
+ }
+}
+
+impl From<PyExecutionPlan> for Arc<dyn ExecutionPlan> {
+ fn from(plan: PyExecutionPlan) -> Arc<dyn ExecutionPlan> {
+ plan.plan.clone()
+ }
+}
+
+impl From<Arc<dyn ExecutionPlan>> for PyExecutionPlan {
+ fn from(plan: Arc<dyn ExecutionPlan>) -> PyExecutionPlan {
+ PyExecutionPlan { plan: plan.clone() }
+ }
+}