Ignore tracing of subordinate operator of the traced opeator within t…

…he same driver
facebookincubator · Oct 11, 2024 · 725f6e7 · 725f6e7
1 parent 9cf4ee0
commit 725f6e7
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 4 deletions.
diff --git a/velox/exec/Driver.h b/velox/exec/Driver.h
@@ -277,6 +277,7 @@ struct DriverCtx {
   std::shared_ptr<Task> task;
   Driver* driver{nullptr};
   facebook::velox::process::ThreadDebugInfo threadDebugInfo;
+  std::unordered_set<int32_t> tracedOperatorIds;
 
   DriverCtx(
       std::shared_ptr<Task> _task,

diff --git a/velox/exec/Operator.cpp b/velox/exec/Operator.cpp
@@ -115,10 +115,18 @@ void Operator::maybeSetTracer() {
     return;
   }
 
-  if (queryTraceConfig->queryNodes.count(planNodeId()) == 0) {
+  const auto nodeId = planNodeId();
+  if (queryTraceConfig->queryNodes.count(nodeId) == 0) {
     return;
   }
 
+  auto& tracedOpIds = operatorCtx_->driverCtx()->tracedOperatorIds;
+  if (const auto iter = tracedOpIds.find(operatorId());
+      iter != tracedOpIds.end()) {
+    return;
+  }
+  tracedOpIds.emplace(operatorId());
+
   const auto pipelineId = operatorCtx_->driverCtx()->pipelineId;
   const auto driverId = operatorCtx_->driverCtx()->driverId;
   LOG(INFO) << "Trace data for operator type: " << operatorType()

diff --git a/velox/tool/trace/tests/TableWriterReplayerTest.cpp b/velox/tool/trace/tests/TableWriterReplayerTest.cpp
@@ -159,16 +159,24 @@ class TableWriterReplayerTest : public HiveConnectorTestBase {
           connector::CommitStrategy::kNoCommit) {
     return [=](core::PlanNodeId nodeId,
                core::PlanNodePtr source) -> core::PlanNodePtr {
+      std::shared_ptr<core::AggregationNode> aggNode = nullptr;
+      if (aggregationNode == nullptr) {
+        aggNode = generateAggregationNode(
+            "c0", nodeId, {}, core::AggregationNode::Step::kPartial, source);
+      } else {
+        aggNode = aggregationNode;
+      }
+
       return std::make_shared<core::TableWriteNode>(
           nodeId,
           inputColumns,
           tableColumnNames,
-          aggregationNode,
+          aggNode,
           insertHandle,
           hasPartitioningScheme,
-          TableWriteTraits::outputType(aggregationNode),
+          TableWriteTraits::outputType(aggNode),
           commitStrategy,
-          std::move(source));
+          source);
     };
   }
 
@@ -227,6 +235,31 @@ class TableWriterReplayerTest : public HiveConnectorTestBase {
     }
   }
 
+  static std::shared_ptr<core::AggregationNode> generateAggregationNode(
+      const std::string& name,
+      const core::PlanNodeId nodeId,
+      const std::vector<core::FieldAccessTypedExprPtr>& groupingKeys,
+      AggregationNode::Step step,
+      const PlanNodePtr& source) {
+    core::TypedExprPtr inputField =
+        std::make_shared<const core::FieldAccessTypedExpr>(BIGINT(), name);
+    auto callExpr = std::make_shared<const core::CallTypedExpr>(
+        BIGINT(), std::vector<core::TypedExprPtr>{inputField}, "min");
+    std::vector<std::string> aggregateNames = {"min"};
+    std::vector<core::AggregationNode::Aggregate> aggregates = {
+        core::AggregationNode::Aggregate{
+            callExpr, {{BIGINT()}}, nullptr, {}, {}}};
+    return std::make_shared<core::AggregationNode>(
+        nodeId,
+        step,
+        groupingKeys,
+        std::vector<core::FieldAccessTypedExprPtr>{},
+        aggregateNames,
+        aggregates,
+        false, // ignoreNullKeys
+        source);
+  }
+
   std::string tableWriteNodeId_;
   FileFormat fileFormat_{FileFormat::DWRF};
 };