feat: Add printPlanWithStats API for Task

facebookincubator · Feb 27, 2025 · 2878fc1 · 2878fc1
1 parent b0d7de9
commit 2878fc1
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 1 deletion.
diff --git a/velox/exec/Task.cpp b/velox/exec/Task.cpp
@@ -30,6 +30,7 @@
 #include "velox/exec/NestedLoopJoinBuild.h"
 #include "velox/exec/OperatorUtils.h"
 #include "velox/exec/OutputBufferManager.h"
+#include "velox/exec/PlanNodeStats.h"
 #include "velox/exec/Task.h"
 #include "velox/exec/TraceUtil.h"
 
@@ -2433,6 +2434,14 @@ ContinueFuture Task::taskDeletionFuture() {
   return std::move(future);
 }
 
+std::string Task::printPlanWithStats(bool includeCustomStats) const {
+  if (planFragment_.planNode) {
+    return exec::printPlanWithStats(
+        *planFragment_.planNode, taskStats_, includeCustomStats);
+  }
+  return "";
+}
+
 std::string Task::toString() const {
   std::lock_guard<std::timed_mutex> l(mutex_);
   std::stringstream out;

diff --git a/velox/exec/Task.h b/velox/exec/Task.h
@@ -116,6 +116,8 @@ class Task : public std::enable_shared_from_this<Task> {
     spillDirectoryCallback_ = std::move(spillDirectoryCallback);
   }
 
+  std::string printPlanWithStats(bool includeCustomStats = false) const;
+
   std::string toString() const;
 
   folly::dynamic toJson() const;

diff --git a/velox/exec/tests/PrintPlanWithStatsTest.cpp b/velox/exec/tests/PrintPlanWithStatsTest.cpp
@@ -133,7 +133,7 @@ TEST_F(PrintPlanWithStatsTest, innerJoinWithTableScan) {
   ensureTaskCompletion(task.get());
   compareOutputs(
       ::testing::UnitTest::GetInstance()->current_test_info()->name(),
-      printPlanWithStats(*op, task->taskStats()),
+      task->printPlanWithStats(),
       {{"-- Project\\[4\\]\\[expressions: \\(c0:INTEGER, ROW\\[\"c0\"\\]\\), \\(p1:BIGINT, plus\\(ROW\\[\"c1\"\\],1\\)\\), \\(p2:BIGINT, plus\\(ROW\\[\"c1\"\\],ROW\\[\"u_c1\"\\]\\)\\)\\] -> c0:INTEGER, p1:BIGINT, p2:BIGINT"},
        {"   Output: 2000 rows \\(.+\\), Cpu time: .+, Blocked wall time: .+, Peak memory: .+, Memory allocations: .+, Threads: 1, CPU breakdown: B/I/O/F (.+/.+/.+/.+)"},
        {"  -- HashJoin\\[3\\]\\[INNER c0=u_c0\\] -> c0:INTEGER, c1:BIGINT, u_c1:BIGINT"},