feat: Implement IndexLookupJoinReplayer (facebookincubator#13560)

Summary:
Pull Request resolved: facebookincubator#13560

Implement IndexLookupJoinReplayer, a tool to be used to replay indexlookupjoin.
bypass-github-export-checks

Reviewed By: xiaoxmeng

Differential Revision: D75668860

fbshipit-source-id: fe7621fce442f2907e964d5406f36cb62dcd145a

Author: juwentus1234

velox/exec/tests/utils/TestIndexStorageConnector.h

@@ -74,6 +74,9 @@ class TestIndexTableHandle : public connector::ConnectorTableHandle {
     obj["name"] = name();
     obj["connectorId"] = connectorId();
     obj["asyncLookup"] = asyncLookup_;
+    // For testing purpose only, we serialize the index table pointer as an
+    // long integer.
+    obj["indexTable"] = reinterpret_cast<int64_t>(indexTable_.get());
     return obj;
   }
 
@@ -82,8 +85,12 @@ class TestIndexTableHandle : public connector::ConnectorTableHandle {
       void* context) {
     // NOTE: this is only for testing purpose so we don't support to serialize
     // the table.
+    auto ptr_value = obj["indexTable"].asInt();
+    auto indexTablePtr = reinterpret_cast<TestIndexTable*>(ptr_value);
     return std::make_shared<TestIndexTableHandle>(
-        obj["connectorId"].getString(), nullptr, obj["asyncLookup"].asBool());
+        obj["connectorId"].getString(),
+        std::shared_ptr<TestIndexTable>(indexTablePtr, [](TestIndexTable*) {}),
+        obj["asyncLookup"].asBool());
   }
 
   static void registerSerDe() {

velox/tool/trace/CMakeLists.txt

@@ -17,6 +17,7 @@ add_library(
   AggregationReplayer.cpp
   FilterProjectReplayer.cpp
   HashJoinReplayer.cpp
+  IndexLookupJoinReplayer.cpp
   OperatorReplayerBase.cpp
   PartitionedOutputReplayer.cpp
   TableScanReplayer.cpp

velox/tool/trace/IndexLookupJoinReplayer.cpp

@@ -0,0 +1,43 @@
+/*
+ * 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/IndexLookupJoinReplayer.h"
+#include "velox/exec/TraceUtil.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 IndexLookupJoinReplayer::createPlanNode(
+    const core::PlanNode* node,
+    const core::PlanNodeId& nodeId,
+    const core::PlanNodePtr& source) const {
+  const auto* indexLookupJoinNode =
+      dynamic_cast<const core::IndexLookupJoinNode*>(node);
+  VELOX_CHECK_NOT_NULL(indexLookupJoinNode);
+  return std::make_shared<core::IndexLookupJoinNode>(
+      nodeId,
+      indexLookupJoinNode->joinType(),
+      indexLookupJoinNode->leftKeys(),
+      indexLookupJoinNode->rightKeys(),
+      indexLookupJoinNode->joinConditions(),
+      source, // Probe side
+      indexLookupJoinNode->lookupSource(), // Index side
+      indexLookupJoinNode->outputType());
+}
+} // namespace facebook::velox::tool::trace

velox/tool/trace/IndexLookupJoinReplayer.h

@@ -0,0 +1,51 @@
+/*
+ * 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"
+
+namespace facebook::velox::tool::trace {
+/// The replayer to replay the traced 'IndexLookupJoin' operator.
+class IndexLookupJoinReplayer : public OperatorReplayerBase {
+ public:
+  IndexLookupJoinReplayer(
+      const std::string& traceDir,
+      const std::string& queryId,
+      const std::string& taskId,
+      const std::string& nodeId,
+      const std::string& operatorType,
+      const std::string& driverIds,
+      uint64_t queryCapacity,
+      folly::Executor* executor)
+      : OperatorReplayerBase(
+            traceDir,
+            queryId,
+            taskId,
+            nodeId,
+            operatorType,
+            driverIds,
+            queryCapacity,
+            executor) {}
+
+ 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/TraceReplayRunner.cpp

@@ -44,6 +44,7 @@
 #include "velox/tool/trace/AggregationReplayer.h"
 #include "velox/tool/trace/FilterProjectReplayer.h"
 #include "velox/tool/trace/HashJoinReplayer.h"
+#include "velox/tool/trace/IndexLookupJoinReplayer.h"
 #include "velox/tool/trace/OperatorReplayerBase.h"
 #include "velox/tool/trace/PartitionedOutputReplayer.h"
 #include "velox/tool/trace/TableScanReplayer.h"
@@ -396,6 +397,16 @@ TraceReplayRunner::createReplayer() const {
         FLAGS_driver_ids,
         queryCapacityBytes,
         cpuExecutor_.get());
+  } else if (traceNodeName == "IndexLookupJoin") {
+    replayer = std::make_unique<tool::trace::IndexLookupJoinReplayer>(
+        FLAGS_root_dir,
+        FLAGS_query_id,
+        FLAGS_task_id,
+        FLAGS_node_id,
+        traceNodeName,
+        FLAGS_driver_ids,
+        queryCapacityBytes,
+        cpuExecutor_.get());
   } else {
     VELOX_UNSUPPORTED("Unsupported operator type: {}", traceNodeName);
   }

velox/tool/trace/tests/CMakeLists.txt

@@ -17,6 +17,7 @@ add_executable(
   AggregationReplayerTest.cpp
   FilterProjectReplayerTest.cpp
   HashJoinReplayerTest.cpp
+  IndexLookupJoinReplayerTest.cpp
   PartitionedOutputReplayerTest.cpp
   TraceFileToolTest.cpp
   TableScanReplayerTest.cpp