feat(python): Support printPlanWithStats API (#12428)

Summary:
Pull Request resolved: #12428

Add Python API to print the plan executed along with task stats.

Reviewed By: kostasxmeta

Differential Revision: D70029859

fbshipit-source-id: 0d30a19be42d732f5bb5ac7b15ffd716c2137e7b

Author: pedroerp

velox/py/runner/PyLocalRunner.cpp

@@ -21,6 +21,7 @@
 #include "velox/core/PlanNode.h"
 #include "velox/dwio/common/Options.h"
 #include "velox/dwio/dwrf/writer/Writer.h"
+#include "velox/exec/PlanNodeStats.h"
 #include "velox/py/vector/PyVector.h"
 
 namespace facebook::velox::py {
@@ -111,6 +112,11 @@ py::iterator PyLocalRunner::execute() {
   return py::make_iterator(pyIterator_->begin(), pyIterator_->end());
 }
 
+std::string PyLocalRunner::printPlanWithStats() const {
+  return exec::printPlanWithStats(
+      *planNode_, cursor_->task()->taskStats(), true);
+}
+
 void drainAllTasks() {
   auto& executor = folly::QueuedImmediateExecutor::instance();
   std::lock_guard<std::mutex> guard(taskRegistryLock());

velox/py/runner/PyLocalRunner.h

@@ -54,6 +54,10 @@ class PyLocalRunner {
   /// Execute the task and returns an iterable to the output vectors.
   pybind11::iterator execute();
 
+  /// Prints a descriptive debug message containing plan and execution stats.
+  /// If the task hasn't finished, will print the plan with the current stats.
+  std::string printPlanWithStats() const;
+
  private:
   friend class PyTaskIterator;
 

velox/py/runner/runner.cpp

@@ -43,6 +43,13 @@ PYBIND11_MODULE(runner, m) {
         return velox::py::PyLocalRunner{planNode, rootPool, executor};
       }))
       .def("execute", &velox::py::PyLocalRunner::execute)
+      .def(
+          "print_plan_with_stats",
+          &velox::py::PyLocalRunner::printPlanWithStats,
+          py::doc(R"(
+         Prints a descriptive debug message containing plan and execution
+         stats. Needs to called after the task has finished.
+          )"))
       .def(
           "add_file_split",
           &velox::py::PyLocalRunner::addFileSplit,

velox/py/runner/runner.pyi

@@ -25,6 +25,7 @@ class LocalRunner:
     def __init__(self, PlanNode) -> None: ...
     def execute(self) -> Iterator[Vector]: ...
     def add_file_split(self, plan_id: str, file_path: str) -> None: ...
+    def print_plan_with_stats(self) -> str: ...
 
 def register_hive(str) -> None: ...
 def register_tpch(str) -> None: ...

velox/py/tests/test_runner.py

@@ -145,6 +145,7 @@ def test_write_read_file(self):
             iterator = runner.execute()
             output = next(iterator)
             self.assertRaises(StopIteration, next, iterator)
+            self.assertNotEqual(runner.print_plan_with_stats(), "")
 
             output_file_from_table_writer = self.extract_file(output)
             self.assertEqual(output_file, output_file_from_table_writer)