Add PyDataFrame.explain (#36)

* Add PyDataFrame.explain

* fix indent

* print results in explain rather than returning the data

datafusion/tests/test_dataframe.py

@@ -179,3 +179,11 @@ def test_struct_select(struct_df):
 
     assert result.column(0) == pa.array([5, 7, 9])
     assert result.column(1) == pa.array([-3, -3, -3])
+
+
+def test_explain(df):
+    df = df.select(
+        column("a") + column("b"),
+        column("a") - column("b"),
+    )
+    df.explain()

src/dataframe.rs

@@ -127,4 +127,12 @@ impl PyDataFrame {
             .join(right.df, join_type, &join_keys.0, &join_keys.1)?;
         Ok(Self::new(df))
     }
+
+    /// Print the query plan
+    #[args(verbose = false, analyze = false)]
+    fn explain(&self, py: Python, verbose: bool, analyze: bool) -> PyResult<()> {
+        let df = self.df.explain(verbose, analyze)?;
+        let batches = wait_for_future(py, df.collect())?;
+        Ok(pretty::print_batches(&batches)?)
+    }
 }