[arrow-ballista] branch master updated: Python: add method to get explain output as a string (#593)

[ag...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

agrove pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-ballista.git


The following commit(s) were added to refs/heads/master by this push:
     new dfc168c6 Python: add method to get explain output as a string (#593)
dfc168c6 is described below

commit dfc168c680376e4f5e88fdce26a8fe3d7e6cc4e6
Author: Andy Grove <an...@gmail.com>
AuthorDate: Wed Jan 18 10:49:35 2023 -0700

    Python: add method to get explain output as a string (#593)
---
 python/ballista/tests/test_dataframe.py | 14 ++++++++++++++
 python/src/context.rs                   |  2 +-
 python/src/dataframe.rs                 | 17 ++++++++++++++++-
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/python/ballista/tests/test_dataframe.py b/python/ballista/tests/test_dataframe.py
index a73d48dc..f134199a 100644
--- a/python/ballista/tests/test_dataframe.py
+++ b/python/ballista/tests/test_dataframe.py
@@ -48,6 +48,20 @@ def struct_df():
     return ctx.create_dataframe([[batch]])
 
 
+def test_explain_string(df):
+    df = df.select(
+        column("a") + column("b"),
+        column("a") - column("b"),
+    )
+
+    # execute and collect the first (and only) batch
+    df.collect()
+
+    explain = df.explain_string()
+    assert "logical_plan" in explain
+    assert "physical_plan" in explain
+
+
 def test_select(df):
     df = df.select(
         column("a") + column("b"),
diff --git a/python/src/context.rs b/python/src/context.rs
index 26c5661a..d479a5f4 100644
--- a/python/src/context.rs
+++ b/python/src/context.rs
@@ -264,4 +264,4 @@ fn convert_table_partition_cols(
         .iter()
         .map(|(name, t)| (name.clone(), t.data_type.clone()))
         .collect()
-}
\ No newline at end of file
+}
diff --git a/python/src/dataframe.rs b/python/src/dataframe.rs
index 645cdebc..295f89dc 100644
--- a/python/src/dataframe.rs
+++ b/python/src/dataframe.rs
@@ -159,7 +159,7 @@ impl PyDataFrame {
         Ok(Self::new(df))
     }
 
-    /// Print the query plan
+    /// Print the explain output to stdout
     #[args(verbose = false, analyze = false)]
     fn explain(&self, py: Python, verbose: bool, analyze: bool) -> PyResult<()> {
         let df = self.df.explain(verbose, analyze)?;
@@ -167,4 +167,19 @@ impl PyDataFrame {
         pretty::print_batches(&batches)
             .map_err(|err| PyArrowException::new_err(err.to_string()))
     }
+
+    /// Get the explain output as a string
+    #[args(verbose = false, analyze = false)]
+    fn explain_string(
+        &self,
+        py: Python,
+        verbose: bool,
+        analyze: bool,
+    ) -> PyResult<String> {
+        let df = self.df.explain(verbose, analyze)?;
+        let batches = wait_for_future(py, df.collect())?;
+        let display = pretty::pretty_format_batches(&batches)
+            .map_err(|err| PyArrowException::new_err(err.to_string()))?;
+        Ok(format!("{}", display))
+    }
 }