fix: Fix build failure and change lookup result iterator return type (f…

…acebookincubator#12224) Summary: Pull Request resolved: facebookincubator#12224 Reviewed By: zacw7 Differential Revision: D68936223 fbshipit-source-id: 40de5a0f289f937b5fed1bcc79bb87e92a4c2c33
ArnavBalyan · Jan 31, 2025 · 253cec9 · 253cec9
1 parent 95898ce
commit 253cec9
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 8 deletions.
diff --git a/velox/connectors/Connector.h b/velox/connectors/Connector.h
@@ -357,7 +357,7 @@ class IndexSource {
         velox::ContinueFuture& future) = 0;
   };
 
-  virtual std::unique_ptr<LookupResultIterator> lookup(
+  virtual std::shared_ptr<LookupResultIterator> lookup(
       const LookupRequest& request) = 0;
 
   virtual std::unordered_map<std::string, RuntimeCounter> runtimeStats() = 0;

diff --git a/velox/exec/IndexLookupJoin.cpp b/velox/exec/IndexLookupJoin.cpp
@@ -64,7 +64,6 @@ IndexLookupJoin::IndexLookupJoin(
       expressionEvaluator_(connectorQueryCtx_->expressionEvaluator()),
       connector_(connector::getConnector(lookupTableHandle_->connectorId())),
       joinNode_{joinNode} {
-  VELOX_CHECK_EQ(joinNode_->sources()[1], joinNode_->lookupSource());
   duplicateJoinKeyCheck(joinNode_->leftKeys());
   duplicateJoinKeyCheck(joinNode_->rightKeys());
 }
@@ -159,8 +158,8 @@ void IndexLookupJoin::initLookupInput() {
           lookupIndexOpt.has_value(),
           "Lookup condition column {} is not found",
           columnName);
-      // A lookup column can only be used in on lookup condition.
-      VELOX_CHECK(
+      // A lookup column can only be used in one lookup condition.
+      VELOX_CHECK_EQ(
           lookupConditionColumnNames.count(columnName),
           0,
           "Lookup condition column {} from lookup table used in more than one lookup conditions",

diff --git a/velox/exec/IndexLookupJoin.h b/velox/exec/IndexLookupJoin.h
@@ -125,7 +125,7 @@ class IndexLookupJoin : public Operator {
   ContinueFuture lookupFuture_{ContinueFuture::makeEmpty()};
   // Used to fetch lookup results for each input batch, and reset after
   // processing all the outputs from the result.
-  std::unique_ptr<connector::IndexSource::LookupResultIterator>
+  std::shared_ptr<connector::IndexSource::LookupResultIterator>
       lookupResultIter_;
   // Used to store the lookup result fetched from 'lookupResultIter_' for output
   // processing. We might split the output result into multiple output batches

diff --git a/velox/exec/tests/utils/TestIndexStorageConnector.cpp b/velox/exec/tests/utils/TestIndexStorageConnector.cpp
@@ -48,7 +48,7 @@ TestIndexSource::TestIndexSource(
   initOutputProjections();
 }
 
-std::unique_ptr<connector::IndexSource::LookupResultIterator>
+std::shared_ptr<connector::IndexSource::LookupResultIterator>
 TestIndexSource::lookup(const LookupRequest& request) {
   const auto numInputRows = request.input->size();
   auto& hashTable = tableHandle_->indexTable()->table;
@@ -67,7 +67,7 @@ TestIndexSource::lookup(const LookupRequest& request) {
   auto& rows = lookup->rows;
   rows.resize(request.input->size());
   std::iota(rows.begin(), rows.end(), 0);
-  return std::make_unique<ResultIterator>(
+  return std::make_shared<ResultIterator>(
       this->shared_from_this(),
       request,
       std::move(lookup),

diff --git a/velox/exec/tests/utils/TestIndexStorageConnector.h b/velox/exec/tests/utils/TestIndexStorageConnector.h
@@ -114,7 +114,7 @@ class TestIndexSource : public connector::IndexSource,
       connector::ConnectorQueryCtx* connectorQueryCtx,
       folly::Executor* executor);
 
-  std::unique_ptr<LookupResultIterator> lookup(
+  std::shared_ptr<LookupResultIterator> lookup(
       const LookupRequest& request) override;
 
   std::unordered_map<std::string, RuntimeCounter> runtimeStats() override {