Implement lazy union

ad-freiburg · Oct 15, 2024 · 52c76bc · 52c76bc
1 parent aa7b723
commit 52c76bc
Show file tree

Hide file tree

Showing 3 changed files with 140 additions and 8 deletions.
diff --git a/src/engine/Union.cpp b/src/engine/Union.cpp
@@ -158,10 +158,20 @@ size_t Union::getCostEstimate() {
          getSizeEstimateBeforeLimit();
 }
 
-ProtoResult Union::computeResult([[maybe_unused]] bool requestLaziness) {
+ProtoResult Union::computeResult(bool requestLaziness) {
   LOG(DEBUG) << "Union result computation..." << std::endl;
-  std::shared_ptr<const Result> subRes1 = _subtrees[0]->getResult();
-  std::shared_ptr<const Result> subRes2 = _subtrees[1]->getResult();
+  std::shared_ptr<const Result> subRes1 =
+      _subtrees[0]->getResult(requestLaziness);
+  std::shared_ptr<const Result> subRes2 =
+      _subtrees[1]->getResult(requestLaziness);
+
+  if (requestLaziness) {
+    auto localVocab = std::make_shared<LocalVocab>();
+    auto generator =
+        computeResultLazily(std::move(subRes1), std::move(subRes2), localVocab);
+    return {std::move(generator), resultSortedOn(), std::move(localVocab)};
+  }
+
   LOG(DEBUG) << "Union subresult computation done." << std::endl;
 
   IdTable idTable =
@@ -179,10 +189,10 @@ void Union::copyChunked(auto beg, auto end, auto target) const {
   size_t total = end - beg;
   for (size_t i = 0; i < total; i += chunkSize) {
     checkCancellation();
-      size_t actualEnd = std::min(i + chunkSize, total);
-      std::copy(beg + i, beg + actualEnd, target + i);
+    size_t actualEnd = std::min(i + chunkSize, total);
+    std::copy(beg + i, beg + actualEnd, target + i);
   }
-};
+}
 
 // _____________________________________________________________________________
 void Union::fillChunked(auto beg, auto end, const auto& value) const {
@@ -229,3 +239,71 @@ IdTable Union::computeUnion(
   }
   return res;
 }
+
+// _____________________________________________________________________________
+template <bool left>
+std::vector<size_t> Union::computePermutation() const {
+  size_t startOfUndefColumns = _subtrees[left ? 0 : 1]->getResultWidth();
+  std::vector<size_t> permutation(_columnOrigins.size());
+  for (size_t targetColIdx = 0; targetColIdx < _columnOrigins.size();
+       ++targetColIdx) {
+    size_t originIndex = _columnOrigins.at(targetColIdx)[left ? 0 : 1];
+    if (originIndex == NO_COLUMN) {
+      originIndex = startOfUndefColumns;
+      startOfUndefColumns++;
+    }
+    permutation[originIndex] = targetColIdx;
+  }
+  return permutation;
+}
+
+// _____________________________________________________________________________
+IdTable Union::transformToCorrectColumnFormat(
+    IdTable idTable, const std::vector<size_t>& permutation) const {
+  while (idTable.numColumns() < getResultWidth()) {
+    idTable.addEmptyColumn();
+    std::ranges::fill(idTable.getColumn(idTable.numColumns() - 1),
+                      Id::makeUndefined());
+  }
+
+  // i + 1 because once everything expect the last column is in the correct
+  // position, the last column is already in the correct position so the last
+  // iteration would always swap the last column with itself.
+  for (size_t i = 0; i + 1 < permutation.size(); ++i) {
+    size_t ind = permutation[i];
+    while (ind < i) {
+      ind = permutation[ind];
+    }
+    idTable.swapColumns(i, ind);
+  }
+  return idTable;
+}
+
+// _____________________________________________________________________________
+cppcoro::generator<IdTable> Union::computeResultLazily(
+    std::shared_ptr<const Result> result1,
+    std::shared_ptr<const Result> result2,
+    std::shared_ptr<LocalVocab> localVocab) const {
+  if (result1->isFullyMaterialized()) {
+    co_yield computeUnion(result1->idTable(),
+                          IdTable{getResultWidth(), allocator()},
+                          _columnOrigins);
+  } else {
+    std::vector<size_t> permutation = computePermutation<true>();
+    for (IdTable& idTable : result1->idTables()) {
+      co_yield transformToCorrectColumnFormat(std::move(idTable), permutation);
+    }
+  }
+  if (result2->isFullyMaterialized()) {
+    co_yield computeUnion(IdTable{getResultWidth(), allocator()},
+                          result2->idTable(), _columnOrigins);
+  } else {
+    std::vector<size_t> permutation = computePermutation<false>();
+    for (IdTable& idTable : result2->idTables()) {
+      co_yield transformToCorrectColumnFormat(std::move(idTable), permutation);
+    }
+  }
+  std::array<const LocalVocab*, 2> vocabs{&result1->localVocab(),
+                                          &result2->localVocab()};
+  *localVocab = LocalVocab::merge(vocabs);
+}
diff --git a/src/engine/Union.h b/src/engine/Union.h
@@ -70,8 +70,21 @@ class Union : public Operation {
   // A similar timeout-checking replacement for `std::fill`.
   void fillChunked(auto beg, auto end, const auto& value) const;
 
-  virtual ProtoResult computeResult(
-      [[maybe_unused]] bool requestLaziness) override;
+  virtual ProtoResult computeResult(bool requestLaziness) override;
 
   VariableToColumnMap computeVariableToColumnMap() const override;
+
+  // Compute the permutation of the `IdTable` being yielded for the left or
+  // right child depending on `left`. This permutation can then be used to swap
+  // the columns without any copy operations.
+  template <bool left>
+  std::vector<size_t> computePermutation() const;
+
+  IdTable transformToCorrectColumnFormat(
+      IdTable idTable, const std::vector<size_t>& permutation) const;
+
+  cppcoro::generator<IdTable> computeResultLazily(
+      std::shared_ptr<const Result> result1,
+      std::shared_ptr<const Result> result2,
+      std::shared_ptr<LocalVocab> localVocab) const;
 };
diff --git a/test/UnionTest.cpp b/test/UnionTest.cpp
@@ -71,3 +71,44 @@ TEST(UnionTest, computeUnionLarge) {
 
   ASSERT_EQ(result, makeIdTableFromVector(expected));
 }
+
+// _____________________________________________________________________________
+TEST(UnionTest, computeUnionLazy) {
+  auto runTest = [](bool nonLazyChilds,
+                    ad_utility::source_location loc =
+                        ad_utility::source_location::current()) {
+    auto l = generateLocationTrace(loc);
+    auto* qec = ad_utility::testing::getQec();
+    IdTable left = makeIdTableFromVector({{V(1)}, {V(2)}, {V(3)}});
+    auto leftT = ad_utility::makeExecutionTree<ValuesForTesting>(
+        qec, std::move(left), Vars{Variable{"?x"}}, false,
+        std::vector<ColumnIndex>{}, LocalVocab{}, std::nullopt, nonLazyChilds);
+
+    IdTable right = makeIdTableFromVector({{V(4), V(5)}, {V(6), V(7)}});
+    auto rightT = ad_utility::makeExecutionTree<ValuesForTesting>(
+        qec, std::move(right), Vars{Variable{"?u"}, Variable{"?x"}}, false,
+        std::vector<ColumnIndex>{}, LocalVocab{}, std::nullopt, nonLazyChilds);
+
+    Union u{ad_utility::testing::getQec(), std::move(leftT), std::move(rightT)};
+    auto resultTable = u.computeResultOnlyForTesting(true);
+    ASSERT_FALSE(resultTable.isFullyMaterialized());
+    auto& result = resultTable.idTables();
+
+    auto U = Id::makeUndefined();
+    auto expected1 = makeIdTableFromVector({{V(1), U}, {V(2), U}, {V(3), U}});
+    auto expected2 = makeIdTableFromVector({{V(5), V(4)}, {V(7), V(6)}});
+
+    auto iterator = result.begin();
+    ASSERT_NE(iterator, result.end());
+    ASSERT_EQ(*iterator, expected1);
+
+    ++iterator;
+    ASSERT_NE(iterator, result.end());
+    ASSERT_EQ(*iterator, expected2);
+
+    ASSERT_EQ(++iterator, result.end());
+  };
+
+  runTest(false);
+  runTest(true);
+}