feat: Add stats spec to TableWrite/TableWriteMerge toString output (#16738)

Summary:
Pull Request resolved: #16738

Add aggregation step and aggregate names to the detailed toString
output of TableWriteNode and TableWriteMergeNode. Example:

The detailed output now looks like:
  - Without stats: -- TableWrite[1][hive, c0, c1, c2] -> rows:BIGINT, ...
  - With stats: -- TableWrite[1][hive, c0, c1, c2, stats[PARTIAL: a0]] -> rows:BIGINT, ...
  - Merge with stats: -- TableWriteMerge[3][stats[INTERMEDIATE: a0]] -> rows:BIGINT, ...

Reviewed By: xiaoxmeng

Differential Revision: D96205345

fbshipit-source-id: 9ee7d331afdcc9acabb8360d3ac71d889534b16d

Author: mbasmanova

velox/core/PlanNode.cpp

@@ -2772,6 +2772,7 @@ PlanNodePtr LocalMergeNode::create(const folly::dynamic& obj, void* context) {
       std::move(sources));
 }
 
+namespace {
 // Validates that grouping keys in 'spec' are present in 'type' and have no
 // duplicates. 'context' is used in error messages (e.g. "written columns",
 // "source output").
@@ -2792,6 +2793,7 @@ void validateGroupingKeys(
         key->name());
   }
 }
+} // namespace
 
 TableWriteNode::TableWriteNode(
     const PlanNodeId& id,
@@ -2847,8 +2849,35 @@ TableWriteNode::TableWriteNode(
       expectedType->toString());
 }
 
+namespace {
+void addStatsSpecDetails(
+    std::stringstream& stream,
+    const std::optional<ColumnStatsSpec>& spec) {
+  if (!spec.has_value()) {
+    return;
+  }
+  stream << "stats[" << AggregationNode::toName(spec->aggregationStep);
+  if (!spec->groupingKeys.empty()) {
+    stream << " [";
+    addFields(stream, spec->groupingKeys);
+    stream << "]";
+  }
+  stream << ": ";
+  for (auto i = 0; i < spec->aggregates.size(); ++i) {
+    appendComma(i, stream);
+    stream << spec->aggregates[i].call->toString();
+  }
+  stream << "]";
+}
+} // namespace
+
 void TableWriteNode::addDetails(std::stringstream& stream) const {
-  stream << insertTableHandle_->connectorInsertTableHandle()->toString();
+  stream << insertTableHandle_->connectorId() << ", "
+         << folly::join(", ", columnNames_);
+  if (columnStatsSpec_.has_value()) {
+    stream << ", ";
+    addStatsSpecDetails(stream, columnStatsSpec_);
+  }
 }
 
 RowTypePtr ColumnStatsSpec::outputType() const {
@@ -2991,7 +3020,9 @@ TableWriteMergeNode::TableWriteMergeNode(
   }
 }
 
-void TableWriteMergeNode::addDetails(std::stringstream& /* stream */) const {}
+void TableWriteMergeNode::addDetails(std::stringstream& stream) const {
+  addStatsSpecDetails(stream, columnStatsSpec_);
+}
 
 folly::dynamic TableWriteMergeNode::serialize() const {
   auto obj = PlanNode::serialize();

velox/exec/tests/PlanNodeToStringTest.cpp

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include "velox/common/testutil/TempDirectoryPath.h"
 #include "velox/connectors/hive/HiveConnector.h"
 #include "velox/exec/WindowFunction.h"
 #include "velox/exec/tests/utils/PlanBuilder.h"
@@ -988,5 +989,72 @@ TEST_F(PlanNodeToStringTest, markDistinct) {
       op->toString(true, false));
 }
 
+TEST_F(PlanNodeToStringTest, tableWrite) {
+  auto outputDir = common::testutil::TempDirectoryPath::create();
+
+  // TableWrite without stats.
+  {
+    auto plan = PlanBuilder()
+                    .values({data_})
+                    .tableWrite(outputDir->getPath())
+                    .planNode();
+    ASSERT_EQ("-- TableWrite[1]\n", plan->toString());
+    ASSERT_EQ(
+        "-- TableWrite[1][test-hive, c0, c1, c2] -> rows:BIGINT, fragments:VARBINARY, commitcontext:VARBINARY\n",
+        plan->toString(true, false));
+  }
+
+  // TableWrite with stats (no grouping keys) and TableWriteMerge.
+  {
+    core::TableWriteNodePtr writeNode;
+    auto plan = PlanBuilder()
+                    .values({data_})
+                    .tableWrite(
+                        outputDir->getPath(),
+                        dwio::common::FileFormat::DWRF,
+                        {"min(c0)"})
+                    .capturePlanNode(writeNode)
+                    .localGather()
+                    .tableWriteMerge()
+                    .planNode();
+
+    ASSERT_EQ("-- TableWrite[1]\n", writeNode->toString());
+    ASSERT_EQ(
+        "-- TableWrite[1][test-hive, c0, c1, c2, stats[PARTIAL: min(ROW[\"c0\"])]] -> rows:BIGINT, fragments:VARBINARY, commitcontext:VARBINARY, a0:SMALLINT\n",
+        writeNode->toString(true, false));
+
+    ASSERT_EQ("-- TableWriteMerge[3]\n", plan->toString());
+    ASSERT_EQ(
+        "-- TableWriteMerge[3][stats[INTERMEDIATE: min(\"a0\")]] -> rows:BIGINT, fragments:VARBINARY, commitcontext:VARBINARY, a0:SMALLINT\n",
+        plan->toString(true, false));
+  }
+
+  // TableWrite with stats and grouping keys (partitioned table).
+  {
+    core::TableWriteNodePtr writeNode;
+    auto plan = PlanBuilder()
+                    .values({data_})
+                    .tableWrite(
+                        outputDir->getPath(),
+                        {"c2"},
+                        dwio::common::FileFormat::DWRF,
+                        {"min(c0)", "max(c1)"})
+                    .capturePlanNode(writeNode)
+                    .localGather()
+                    .tableWriteMerge()
+                    .planNode();
+
+    ASSERT_EQ("-- TableWrite[1]\n", writeNode->toString());
+    ASSERT_EQ(
+        "-- TableWrite[1][test-hive, c0, c1, c2, stats[PARTIAL [c2]: min(ROW[\"c0\"]), max(ROW[\"c1\"])]] -> rows:BIGINT, fragments:VARBINARY, commitcontext:VARBINARY, c2:BIGINT, a0:SMALLINT, a1:INTEGER\n",
+        writeNode->toString(true, false));
+
+    ASSERT_EQ("-- TableWriteMerge[3]\n", plan->toString());
+    ASSERT_EQ(
+        "-- TableWriteMerge[3][stats[INTERMEDIATE [c2]: min(\"a0\"), max(\"a1\")]] -> rows:BIGINT, fragments:VARBINARY, commitcontext:VARBINARY, c2:BIGINT, a0:SMALLINT, a1:INTEGER\n",
+        plan->toString(true, false));
+  }
+}
+
 } // namespace
 } // namespace facebook::velox::exec

velox/exec/tests/PrintPlanWithStatsTest.cpp

@@ -447,7 +447,7 @@ TEST_F(PrintPlanWithStatsTest, tableWriterWithTableScan) {
   compareOutputs(
       ::testing::UnitTest::GetInstance()->current_test_info()->name(),
       printPlanWithStats(*writePlan, task->taskStats()),
-      {{R"(-- TableWrite\[1\]\[.+InsertTableHandle .+)"},
+      {{R"(-- TableWrite\[1\]\[test-hive, c0, c1, c2, c3, c4, c5\] -> rows:BIGINT, fragments:VARBINARY, commitcontext:VARBINARY)"},
        {"   Output: .+, Physical written output: .+, Cpu time: .+, Blocked wall time: .+, Peak memory: .+, Memory allocations: .+, Threads: 1, CPU breakdown: B/I/O/F (.+/.+/.+/.+)"},
        {R"(  -- TableScan\[0\]\[table: hive_table\] -> c0:BIGINT, c1:INTEGER, c2:SMALLINT, c3:REAL, c4:DOUBLE, c5:VARCHAR)"},
        {R"(     Input: 100 rows \(.+\), Output: 100 rows \(.+\), Cpu time: .+, Blocked wall time: .+, Peak memory: .+, Memory allocations: .+, Threads: 1, Splits: 1, CPU breakdown: B/I/O/F (.+/.+/.+/.+))"}});
@@ -456,7 +456,7 @@ TEST_F(PrintPlanWithStatsTest, tableWriterWithTableScan) {
       ::testing::UnitTest::GetInstance()->current_test_info()->name(),
       printPlanWithStats(*writePlan, task->taskStats(), true),
       {
-          {R"(-- TableWrite\[1\]\[.+InsertTableHandle .+)"},
+          {R"(-- TableWrite\[1\]\[test-hive, c0, c1, c2, c3, c4, c5\] -> rows:BIGINT, fragments:VARBINARY, commitcontext:VARBINARY)"},
           {"   Output: .+, Physical written output: .+, Cpu time: .+, Blocked wall time: .+, Peak memory: .+, Memory allocations: .+, Threads: 1, CPU breakdown: B/I/O/F (.+/.+/.+/.+)"},
           {"      dataSourceLazyCpuNanos\\s+sum: .+, count: .+, min: .+, max: .+"},
           {"      dataSourceLazyInputBytes\\s+sum: .+, count: .+, min: .+, max: .+"},