[HSTACK] - expose unparser plan_sql on DF object

Author: aditanase

python/datafusion/dataframe.py

@@ -808,6 +808,9 @@ def count(self) -> int:
     def distributed_plan(self):
         return self.df.distributed_plan()
 
+    def plan_sql(self):
+        return self.df.plan_sql()
+
     @deprecated("Use :py:func:`unnest_columns` instead.")
     def unnest_column(self, column: str, preserve_nulls: bool = True) -> DataFrame:
         """See :py:func:`unnest_columns`."""

src/dataframe.rs

@@ -41,13 +41,14 @@ use datafusion::execution::SendableRecordBatchStream;
 use datafusion::parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel};
 use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 use datafusion::prelude::*;
+use datafusion::sql::unparser::plan_to_sql;
 use datafusion_proto::physical_plan::AsExecutionPlan;
 use datafusion_proto::protobuf::PhysicalPlanNode;
 use prost::Message;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::pybacked::PyBackedStr;
-use pyo3::types::{PyBytes, PyCapsule, PyDict, PyTuple, PyTupleMethods};
+use pyo3::types::{PyBytes, PyCapsule, PyDict, PyString, PyTuple, PyTupleMethods};
 use tokio::task::JoinHandle;
 
 use crate::catalog::PyTable;
@@ -713,6 +714,13 @@ impl PyDataFrame {
         let future_plan = DistributedPlan::try_new(self.df.as_ref());
         wait_for_future(py, future_plan).map_err(py_datafusion_err)
     }
+
+    fn plan_sql(&self, py: Python<'_>) -> PyResult<PyObject> {
+        let logical_plan = self.df.logical_plan();
+
+        let sql = plan_to_sql(logical_plan).map_err(py_datafusion_err)?;
+        Ok(PyString::new(py, sql.to_string().as_ref()).into())
+    }
 }
 
 #[pyclass(get_all)]