fix: TestIndexSource::initOutputProjections (facebookincubator#13855)

mbasmanova · facebook-github-bot · commit 8ee33c1b82a4 · 2025-06-23T21:42:58.000-07:00
Summary: Pull Request resolved: facebookincubator#13855 TestIndexSource didn't support using column assignments in TableScan plan node to project columns under different names. Also, introduce TestIndexConnectorFactory::registerConnector helper method to reduce boiler plate. Reviewed By: xiaoxmeng Differential Revision: D77173098 fbshipit-source-id: f6cabde2f173506fe4382e83c2360b16256a9270
diff --git a/velox/exec/tests/IndexLookupJoinTest.cpp b/velox/exec/tests/IndexLookupJoinTest.cpp
@@ -80,25 +80,14 @@ class IndexLookupJoinTest : public IndexLookupJoinTestBase,
     connector::hive::HiveColumnHandle::registerSerDe();
     Type::registerSerDe();
     core::ITypedExpr::registerSerDe();
-    connector::registerConnectorFactory(
-        std::make_shared<TestIndexConnectorFactory>());
-    std::shared_ptr<connector::Connector> connector =
-        connector::getConnectorFactory(kTestIndexConnectorName)
-            ->newConnector(
-                kTestIndexConnectorName,
-                {},
-                nullptr,
-                connectorCpuExecutor_.get());
-    connector::registerConnector(connector);
+    TestIndexConnectorFactory::registerConnector(connectorCpuExecutor_.get());
 
     keyType_ = ROW({"u0", "u1", "u2"}, {BIGINT(), BIGINT(), BIGINT()});
     valueType_ = ROW({"u3", "u4", "u5"}, {BIGINT(), BIGINT(), VARCHAR()});
     tableType_ = concat(keyType_, valueType_);
     probeType_ = ROW(
         {"t0", "t1", "t2", "t3", "t4", "t5"},
         {BIGINT(), BIGINT(), BIGINT(), BIGINT(), ARRAY(BIGINT()), VARCHAR()});
-
-    TestIndexTableHandle::registerSerDe();
   }
 
   void TearDown() override {
@@ -115,13 +104,25 @@ class IndexLookupJoinTest : public IndexLookupJoinTestBase,
 
   // Makes index table handle with the specified index table and async lookup
   // flag.
-  std::shared_ptr<TestIndexTableHandle> makeIndexTableHandle(
+  static std::shared_ptr<TestIndexTableHandle> makeIndexTableHandle(
       const std::shared_ptr<TestIndexTable>& indexTable,
       bool asyncLookup) {
     return std::make_shared<TestIndexTableHandle>(
         kTestIndexConnectorName, indexTable, asyncLookup);
   }
 
+  static std::
+      unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
+      makeIndexColumnHandles(const std::vector<std::string>& names) {
+    std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
+        handles;
+    for (const auto& name : names) {
+      handles.emplace(name, std::make_shared<TestIndexColumnHandle>(name));
+    }
+
+    return handles;
+  }
+
   const std::unique_ptr<folly::CPUThreadPoolExecutor> connectorCpuExecutor_{
       std::make_unique<folly::CPUThreadPoolExecutor>(128)};
 };
@@ -173,8 +174,7 @@ TEST_P(IndexLookupJoinTest, joinCondition) {
 TEST_P(IndexLookupJoinTest, planNodeAndSerde) {
   TestIndexTableHandle::registerSerDe();
 
-  auto indexConnectorHandle = std::make_shared<TestIndexTableHandle>(
-      kTestIndexConnectorName, nullptr, true);
+  auto indexConnectorHandle = makeIndexTableHandle(nullptr, true);
 
   auto left = makeRowVector(
       {"t0", "t1", "t2", "t3", "t4"},
@@ -201,15 +201,15 @@ TEST_P(IndexLookupJoinTest, planNodeAndSerde) {
   auto planBuilder = PlanBuilder();
   auto nonIndexTableScan = std::dynamic_pointer_cast<const core::TableScanNode>(
       PlanBuilder::TableScanBuilder(planBuilder)
-          .outputType(std::dynamic_pointer_cast<const RowType>(right->type()))
+          .outputType(asRowType(right->type()))
           .endTableScan()
           .planNode());
   VELOX_CHECK_NOT_NULL(nonIndexTableScan);
 
   auto indexTableScan = std::dynamic_pointer_cast<const core::TableScanNode>(
       PlanBuilder::TableScanBuilder(planBuilder)
           .tableHandle(indexConnectorHandle)
-          .outputType(std::dynamic_pointer_cast<const RowType>(right->type()))
+          .outputType(asRowType(right->type()))
           .endTableScan()
           .planNode());
   VELOX_CHECK_NOT_NULL(indexTableScan);
@@ -735,13 +735,11 @@ TEST_P(IndexLookupJoinTest, equalJoin) {
     const auto indexTableHandle =
         makeIndexTableHandle(indexTable, GetParam().asyncLookup);
     auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
-    std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
-        columnHandles;
     const auto indexScanNode = makeIndexScanNode(
         planNodeIdGenerator,
         indexTableHandle,
         makeScanOutputType(testData.scanOutputColumns),
-        columnHandles);
+        makeIndexColumnHandles(testData.scanOutputColumns));
 
     auto plan = makeLookupPlan(
         planNodeIdGenerator,
@@ -1193,13 +1191,11 @@ TEST_P(IndexLookupJoinTest, betweenJoinCondition) {
     const auto indexTableHandle =
         makeIndexTableHandle(indexTable, GetParam().asyncLookup);
     auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
-    std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
-        columnHandles;
     const auto indexScanNode = makeIndexScanNode(
         planNodeIdGenerator,
         indexTableHandle,
         makeScanOutputType(testData.lookupOutputColumns),
-        columnHandles);
+        makeIndexColumnHandles(testData.lookupOutputColumns));
 
     auto plan = makeLookupPlan(
         planNodeIdGenerator,
@@ -1517,13 +1513,11 @@ TEST_P(IndexLookupJoinTest, inJoinCondition) {
     const auto indexTableHandle =
         makeIndexTableHandle(indexTable, GetParam().asyncLookup);
     auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
-    std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
-        columnHandles;
     const auto indexScanNode = makeIndexScanNode(
         planNodeIdGenerator,
         indexTableHandle,
         makeScanOutputType(testData.lookupOutputColumns),
-        columnHandles);
+        makeIndexColumnHandles(testData.lookupOutputColumns));
 
     auto plan = makeLookupPlan(
         planNodeIdGenerator,
@@ -1568,13 +1562,11 @@ DEBUG_ONLY_TEST_P(IndexLookupJoinTest, connectorError) {
   const auto indexTableHandle =
       makeIndexTableHandle(indexTable, GetParam().asyncLookup);
   auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
-  std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
-      columnHandles;
   const auto indexScanNode = makeIndexScanNode(
       planNodeIdGenerator,
       indexTableHandle,
       makeScanOutputType({"u0", "u1", "u2", "u5"}),
-      columnHandles);
+      makeIndexColumnHandles({"u0", "u1", "u2", "u5"}));
 
   auto plan = makeLookupPlan(
       planNodeIdGenerator,
@@ -1637,13 +1629,11 @@ DEBUG_ONLY_TEST_P(IndexLookupJoinTest, prefetch) {
   const auto indexTableHandle =
       makeIndexTableHandle(indexTable, GetParam().asyncLookup);
   auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
-  std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
-      columnHandles;
   const auto indexScanNode = makeIndexScanNode(
       planNodeIdGenerator,
       indexTableHandle,
       makeScanOutputType({"u0", "u1", "u2", "u3", "u5"}),
-      columnHandles);
+      makeIndexColumnHandles({"u0", "u1", "u2", "u3", "u5"}));
 
   auto plan = makeLookupPlan(
       planNodeIdGenerator,
@@ -1739,13 +1729,12 @@ TEST_P(IndexLookupJoinTest, outputBatchSizeWithInnerJoin) {
     const auto indexTableHandle =
         makeIndexTableHandle(indexTable, GetParam().asyncLookup);
     auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
-    std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
-        columnHandles;
+
     const auto indexScanNode = makeIndexScanNode(
         planNodeIdGenerator,
         indexTableHandle,
         makeScanOutputType({"u0", "u1", "u2", "u5"}),
-        columnHandles);
+        makeIndexColumnHandles({"u0", "u1", "u2", "u5"}));
 
     auto plan = makeLookupPlan(
         planNodeIdGenerator,
@@ -1847,13 +1836,11 @@ TEST_P(IndexLookupJoinTest, outputBatchSizeWithLeftJoin) {
     const auto indexTableHandle =
         makeIndexTableHandle(indexTable, GetParam().asyncLookup);
     auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
-    std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
-        columnHandles;
     const auto indexScanNode = makeIndexScanNode(
         planNodeIdGenerator,
         indexTableHandle,
         makeScanOutputType({"u0", "u1", "u2", "u5"}),
-        columnHandles);
+        makeIndexColumnHandles({"u0", "u1", "u2", "u5"}));
 
     auto plan = makeLookupPlan(
         planNodeIdGenerator,
@@ -1919,13 +1906,11 @@ DEBUG_ONLY_TEST_P(IndexLookupJoinTest, runtimeStats) {
   const auto indexTableHandle =
       makeIndexTableHandle(indexTable, GetParam().asyncLookup);
   auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
-  std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
-      columnHandles;
   const auto indexScanNode = makeIndexScanNode(
       planNodeIdGenerator,
       indexTableHandle,
       makeScanOutputType({"u0", "u1", "u2", "u3", "u5"}),
-      columnHandles);
+      makeIndexColumnHandles({"u0", "u1", "u2", "u3", "u5"}));
 
   auto plan = makeLookupPlan(
       planNodeIdGenerator,
@@ -2009,13 +1994,11 @@ TEST_P(IndexLookupJoinTest, barrier) {
   const auto indexTableHandle =
       makeIndexTableHandle(indexTable, GetParam().asyncLookup);
   auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
-  std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
-      columnHandles;
   const auto indexScanNode = makeIndexScanNode(
       planNodeIdGenerator,
       indexTableHandle,
       makeScanOutputType({"u0", "u1", "u2", "u3", "u5"}),
-      columnHandles);
+      makeIndexColumnHandles({"u0", "u1", "u2", "u3", "u5"}));
 
   auto plan = makeLookupPlan(
       planNodeIdGenerator,
@@ -2084,13 +2067,11 @@ TEST_P(IndexLookupJoinTest, joinFuzzer) {
   std::random_device rd;
   std::mt19937 g(rd());
   std::shuffle(scanOutput.begin(), scanOutput.end(), g);
-  std::unordered_map<std::string, std::shared_ptr<connector::ColumnHandle>>
-      columnHandles;
   const auto indexScanNode = makeIndexScanNode(
       planNodeIdGenerator,
       indexTableHandle,
       makeScanOutputType(scanOutput),
-      columnHandles);
+      makeIndexColumnHandles(scanOutput));
 
   auto plan = makeLookupPlan(
       planNodeIdGenerator,
diff --git a/velox/exec/tests/utils/IndexLookupJoinTestBase.cpp b/velox/exec/tests/utils/IndexLookupJoinTestBase.cpp
@@ -283,7 +283,7 @@ IndexLookupJoinTestBase::makeIndexScanNode(
     const std::shared_ptr<facebook::velox::connector::ConnectorTableHandle>
         indexTableHandle,
     const facebook::velox::RowTypePtr& outputType,
-    std::unordered_map<
+    const std::unordered_map<
         std::string,
         std::shared_ptr<facebook::velox::connector::ColumnHandle>>&
         assignments) {
diff --git a/velox/exec/tests/utils/IndexLookupJoinTestBase.h b/velox/exec/tests/utils/IndexLookupJoinTestBase.h
@@ -125,7 +125,7 @@ class IndexLookupJoinTestBase
       const std::shared_ptr<facebook::velox::connector::ConnectorTableHandle>
           indexTableHandle,
       const facebook::velox::RowTypePtr& outputType,
-      std::unordered_map<
+      const std::unordered_map<
           std::string,
           std::shared_ptr<facebook::velox::connector::ColumnHandle>>&
           assignments);
diff --git a/velox/exec/tests/utils/TestIndexStorageConnector.cpp b/velox/exec/tests/utils/TestIndexStorageConnector.cpp
@@ -176,6 +176,9 @@ TestIndexSource::TestIndexSource(
     size_t numEqualJoinKeys,
     const core::TypedExprPtr& joinConditionExpr,
     const std::shared_ptr<TestIndexTableHandle>& tableHandle,
+    const std::unordered_map<
+        std::string,
+        std::shared_ptr<TestIndexColumnHandle>>& columnHandles,
     connector::ConnectorQueryCtx* connectorQueryCtx,
     folly::Executor* executor)
     : tableHandle_(tableHandle),
@@ -192,7 +195,9 @@ TestIndexSource::TestIndexSource(
               : nullptr),
       pool_(connectorQueryCtx_->memoryPool()->shared_from_this()),
       executor_(executor) {
-  VELOX_CHECK_NOT_NULL(executor_);
+  if (tableHandle_->asyncLookup()) {
+    VELOX_CHECK_NOT_NULL(executor_);
+  }
   VELOX_CHECK_LE(outputType_->size(), valueType_->size() + keyType_->size());
   VELOX_CHECK_LE(numEqualJoinKeys_, keyType_->size());
   for (int i = 0; i < numEqualJoinKeys_; ++i) {
@@ -202,7 +207,7 @@ TestIndexSource::TestIndexSource(
         keyType_->toString(),
         inputType_->toString());
   }
-  initOutputProjections();
+  initOutputProjections(columnHandles);
   initConditionProjections();
 }
 
@@ -260,20 +265,25 @@ void TestIndexSource::initConditionProjections() {
   conditionInputType_ = ROW(std::move(names), std::move(types));
 }
 
-void TestIndexSource::initOutputProjections() {
+void TestIndexSource::initOutputProjections(
+    const std::unordered_map<
+        std::string,
+        std::shared_ptr<TestIndexColumnHandle>>& columnHandles) {
   VELOX_CHECK(lookupOutputProjections_.empty());
+
   lookupOutputProjections_.reserve(outputType_->size());
   for (auto outputChannel = 0; outputChannel < outputType_->size();
        ++outputChannel) {
     const auto outputName = outputType_->nameOf(outputChannel);
-    if (valueType_->containsChild(outputName)) {
-      const auto tableValueChannel = valueType_->getChildIdx(outputName);
+    const auto columnName = columnHandles.at(outputName)->name();
+    if (valueType_->containsChild(columnName)) {
+      const auto tableValueChannel = valueType_->getChildIdx(columnName);
       // The hash table layout is: index columns, value columns.
       lookupOutputProjections_.emplace_back(
           keyType_->size() + tableValueChannel, outputChannel);
       continue;
     }
-    const auto tableKeyChannel = keyType_->getChildIdx(outputName);
+    const auto tableKeyChannel = keyType_->getChildIdx(columnName);
     lookupOutputProjections_.emplace_back(tableKeyChannel, outputChannel);
   }
   VELOX_CHECK_EQ(lookupOutputProjections_.size(), outputType_->size());
@@ -547,12 +557,24 @@ std::shared_ptr<connector::IndexSource> TestIndexConnector::createIndexSource(
   const auto& indexTable = testIndexTableHandle->indexTable();
   auto joinConditionExpr =
       toJoinConditionExpr(joinConditions, indexTable, inputType, columnHandles);
+
+  std::unordered_map<std::string, std::shared_ptr<TestIndexColumnHandle>>
+      testColumnHandles;
+  for (const auto& [name, handle] : columnHandles) {
+    auto testColumnHandle =
+        std::dynamic_pointer_cast<TestIndexColumnHandle>(handle);
+    VELOX_CHECK_NOT_NULL(testColumnHandle);
+
+    testColumnHandles.emplace(name, testColumnHandle);
+  }
+
   return std::make_shared<TestIndexSource>(
       inputType,
       outputType,
       numJoinKeys,
-      std::move(joinConditionExpr),
-      std::move(testIndexTableHandle),
+      joinConditionExpr,
+      testIndexTableHandle,
+      testColumnHandles,
       connectorQueryCtx,
       executor_);
 }
diff --git a/velox/exec/tests/utils/TestIndexStorageConnector.h b/velox/exec/tests/utils/TestIndexStorageConnector.h
diff --git a/velox/tool/trace/tests/IndexLookupJoinReplayerTest.cpp b/velox/tool/trace/tests/IndexLookupJoinReplayerTest.cpp
