Add hash join replayer

velox/exec/HashBuild.cpp

@@ -314,6 +314,7 @@ void HashBuild::removeInputRowsForAntiJoinFilter() {
 void HashBuild::addInput(RowVectorPtr input) {
   checkRunning();
   ensureInputFits(input);
+  traceInput(input);
 
   TestValue::adjust("facebook::velox::exec::HashBuild::addInput", this);
 

velox/exec/Operator.cpp

@@ -121,6 +121,12 @@ void Operator::maybeSetTracer() {
     return;
   }
 
+  VELOX_CHECK(
+      trace::QueryTraceTraits::supportedOperatorTypes.count(
+          operatorCtx_->operatorType()) != 0,
+      "{} do not support tracing",
+      operatorCtx_->operatorType());
+
   auto& tracedOpMap = operatorCtx_->driverCtx()->tracedOperatorMap;
   if (const auto iter = tracedOpMap.find(operatorId());
       iter != tracedOpMap.end()) {

velox/exec/QueryTraceTraits.h

@@ -31,5 +31,14 @@ struct QueryTraceTraits {
   static inline const std::string kQueryMetaFileName = "query_meta.json";
   static inline const std::string kDataSummaryFileName = "data_summary.json";
   static inline const std::string kDataFileName = "trace.data";
+
+  static inline const std::unordered_set<std::string> supportedOperatorTypes{
+      "TableWriter",
+      "TableScan",
+      "Aggregation",
+      "PartialAggregation",
+      "PartitionedOutput",
+      "HashBuild",
+      "HashProbe"};
 };
 } // namespace facebook::velox::exec::trace

velox/tool/trace/CMakeLists.txt

@@ -17,7 +17,8 @@ velox_add_library(
   AggregationReplayer.cpp
   OperatorReplayerBase.cpp
   PartitionedOutputReplayer.cpp
-  TableWriterReplayer.cpp)
+  TableWriterReplayer.cpp
+  HashJoinReplayer.cpp)
 
 velox_link_libraries(
   velox_query_trace_replayer_base

velox/tool/trace/HashJoinReplayer.cpp

@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "velox/tool/trace/HashJoinReplayer.h"
+
+#include <exec/QueryTraceUtil.h>
+
+#include "velox/exec/QueryDataReader.h"
+#include "velox/exec/tests/utils/PlanBuilder.h"
+
+using namespace facebook::velox;
+using namespace facebook::velox::exec;
+using namespace facebook::velox::exec::test;
+
+namespace facebook::velox::tool::trace {
+core::PlanNodePtr HashJoinReplayer::createPlanNode(
+    const core::PlanNode* node,
+    const core::PlanNodeId& nodeId,
+    const core::PlanNodePtr& source) const {
+  const auto* hashJoinNode = dynamic_cast<const core::HashJoinNode*>(node);
+  return std::make_shared<core::HashJoinNode>(
+      nodeId,
+      hashJoinNode->joinType(),
+      hashJoinNode->isNullAware(),
+      hashJoinNode->leftKeys(),
+      hashJoinNode->rightKeys(),
+      hashJoinNode->filter(),
+      source,
+      PlanBuilder(planNodeIdGenerator_)
+          .traceScan(
+              nodeDir_, exec::trace::getDataType(planFragment_, nodeId_), 1),
+      hashJoinNode->outputType());
+
+  return nullptr;
+}
+} // namespace facebook::velox::tool::trace

velox/tool/trace/HashJoinReplayer.h

@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "velox/core/PlanNode.h"
+#include "velox/tool/trace/OperatorReplayerBase.h"
+
+#include <parse/PlanNodeIdGenerator.h>
+
+namespace facebook::velox::tool::trace {
+/// The replayer to replay the traced 'HashAggregation' operator.
+class HashJoinReplayer : public OperatorReplayerBase {
+ public:
+  HashJoinReplayer(
+      const std::string& rootDir,
+      const std::string& taskId,
+      const std::string& nodeId,
+      const int32_t pipelineId,
+      const std::string& operatorType)
+      : OperatorReplayerBase(
+            rootDir,
+            taskId,
+            nodeId,
+            pipelineId,
+            operatorType) {}
+
+ private:
+  core::PlanNodePtr createPlanNode(
+      const core::PlanNode* node,
+      const core::PlanNodeId& nodeId,
+      const core::PlanNodePtr& source) const override;
+};
+} // namespace facebook::velox::tool::trace

velox/tool/trace/OperatorReplayerBase.cpp

@@ -68,7 +68,7 @@ core::PlanNodePtr OperatorReplayerBase::createPlan() const {
   const auto* replayNode = core::PlanNode::findFirstNode(
       planFragment_.get(),
       [this](const core::PlanNode* node) { return node->id() == nodeId_; });
-  return exec::test::PlanBuilder()
+  return exec::test::PlanBuilder(planNodeIdGenerator_)
       .traceScan(nodeDir_, exec::trace::getDataType(planFragment_, nodeId_))
       .addNode(replayNodeFactory(replayNode))
       .planNode();

velox/tool/trace/OperatorReplayerBase.h

@@ -59,6 +59,9 @@ class OperatorReplayerBase {
   const std::string taskDir_;
   const std::string nodeDir_;
 
+  const auto planNodeIdGenerator_ =
+      std::make_shared<core::PlanNodeIdGenerator>();
+
   std::unordered_map<std::string, std::string> queryConfigs_;
   std::unordered_map<std::string, std::unordered_map<std::string, std::string>>
       connectorConfigs_;