Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 16 additions & 19 deletions src/paimon/common/global_index/union_global_index_reader.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -39,92 +39,92 @@ Result<std::shared_ptr<GlobalIndexResult>> UnionGlobalIndexReader::VisitIsNull()

Result<std::shared_ptr<GlobalIndexResult>> UnionGlobalIndexReader::VisitEqual(
const Literal& literal) {
return Union([&literal](const std::shared_ptr<GlobalIndexReader>& reader) {
return Union([literal](const std::shared_ptr<GlobalIndexReader>& reader) {
return reader->VisitEqual(literal);
});
}

Result<std::shared_ptr<GlobalIndexResult>> UnionGlobalIndexReader::VisitNotEqual(
const Literal& literal) {
return Union([&literal](const std::shared_ptr<GlobalIndexReader>& reader) {
return Union([literal](const std::shared_ptr<GlobalIndexReader>& reader) {
return reader->VisitNotEqual(literal);
});
}

Result<std::shared_ptr<GlobalIndexResult>> UnionGlobalIndexReader::VisitLessThan(
const Literal& literal) {
return Union([&literal](const std::shared_ptr<GlobalIndexReader>& reader) {
return Union([literal](const std::shared_ptr<GlobalIndexReader>& reader) {
return reader->VisitLessThan(literal);
});
}

Result<std::shared_ptr<GlobalIndexResult>> UnionGlobalIndexReader::VisitLessOrEqual(
const Literal& literal) {
return Union([&literal](const std::shared_ptr<GlobalIndexReader>& reader) {
return Union([literal](const std::shared_ptr<GlobalIndexReader>& reader) {
return reader->VisitLessOrEqual(literal);
});
}

Result<std::shared_ptr<GlobalIndexResult>> UnionGlobalIndexReader::VisitGreaterThan(
const Literal& literal) {
return Union([&literal](const std::shared_ptr<GlobalIndexReader>& reader) {
return Union([literal](const std::shared_ptr<GlobalIndexReader>& reader) {
return reader->VisitGreaterThan(literal);
});
}

Result<std::shared_ptr<GlobalIndexResult>> UnionGlobalIndexReader::VisitGreaterOrEqual(
const Literal& literal) {
return Union([&literal](const std::shared_ptr<GlobalIndexReader>& reader) {
return Union([literal](const std::shared_ptr<GlobalIndexReader>& reader) {
return reader->VisitGreaterOrEqual(literal);
});
}

Result<std::shared_ptr<GlobalIndexResult>> UnionGlobalIndexReader::VisitIn(
const std::vector<Literal>& literals) {
return Union([&literals](const std::shared_ptr<GlobalIndexReader>& reader) {
return Union([literals](const std::shared_ptr<GlobalIndexReader>& reader) {
return reader->VisitIn(literals);
});
}

Result<std::shared_ptr<GlobalIndexResult>> UnionGlobalIndexReader::VisitNotIn(
const std::vector<Literal>& literals) {
return Union([&literals](const std::shared_ptr<GlobalIndexReader>& reader) {
return Union([literals](const std::shared_ptr<GlobalIndexReader>& reader) {
return reader->VisitNotIn(literals);
});
}

Result<std::shared_ptr<GlobalIndexResult>> UnionGlobalIndexReader::VisitStartsWith(
const Literal& prefix) {
return Union([&prefix](const std::shared_ptr<GlobalIndexReader>& reader) {
return Union([prefix](const std::shared_ptr<GlobalIndexReader>& reader) {
return reader->VisitStartsWith(prefix);
});
}

Result<std::shared_ptr<GlobalIndexResult>> UnionGlobalIndexReader::VisitEndsWith(
const Literal& suffix) {
return Union([&suffix](const std::shared_ptr<GlobalIndexReader>& reader) {
return Union([suffix](const std::shared_ptr<GlobalIndexReader>& reader) {
return reader->VisitEndsWith(suffix);
});
}

Result<std::shared_ptr<GlobalIndexResult>> UnionGlobalIndexReader::VisitContains(
const Literal& literal) {
return Union([&literal](const std::shared_ptr<GlobalIndexReader>& reader) {
return Union([literal](const std::shared_ptr<GlobalIndexReader>& reader) {
return reader->VisitContains(literal);
});
}

Result<std::shared_ptr<GlobalIndexResult>> UnionGlobalIndexReader::VisitLike(
const Literal& literal) {
return Union([&literal](const std::shared_ptr<GlobalIndexReader>& reader) {
return Union([literal](const std::shared_ptr<GlobalIndexReader>& reader) {
return reader->VisitLike(literal);
});
}

Result<std::shared_ptr<ScoredGlobalIndexResult>> UnionGlobalIndexReader::VisitVectorSearch(
const std::shared_ptr<VectorSearch>& vector_search) {
auto results = ExecuteAllReaders<Result<std::shared_ptr<ScoredGlobalIndexResult>>>(
[&vector_search](const std::shared_ptr<GlobalIndexReader>& reader)
[vector_search](const std::shared_ptr<GlobalIndexReader>& reader)
-> Result<std::shared_ptr<ScoredGlobalIndexResult>> {
return reader->VisitVectorSearch(vector_search);
});
Expand Down Expand Up @@ -155,15 +155,13 @@ Result<std::shared_ptr<ScoredGlobalIndexResult>> UnionGlobalIndexReader::VisitVe

Result<std::shared_ptr<GlobalIndexResult>> UnionGlobalIndexReader::VisitFullTextSearch(
const std::shared_ptr<FullTextSearch>& full_text_search) {
return Union([&full_text_search](const std::shared_ptr<GlobalIndexReader>& reader) {
return Union([full_text_search](const std::shared_ptr<GlobalIndexReader>& reader) {
return reader->VisitFullTextSearch(full_text_search);
});
}

Result<std::shared_ptr<GlobalIndexResult>> UnionGlobalIndexReader::Union(ReaderAction action) {
auto results = ExecuteAllReaders<Result<std::shared_ptr<GlobalIndexResult>>>(
[&action](const std::shared_ptr<GlobalIndexReader>& reader)
-> Result<std::shared_ptr<GlobalIndexResult>> { return action(reader); });
auto results = ExecuteAllReaders<Result<std::shared_ptr<GlobalIndexResult>>>(action);

std::shared_ptr<GlobalIndexResult> merged_result = nullptr;
for (auto& result_or_status : results) {
Expand Down Expand Up @@ -207,8 +205,7 @@ std::vector<R> UnionGlobalIndexReader::ExecuteAllReaders(
std::vector<std::future<R>> futures;
futures.reserve(readers_.size());
for (const auto& reader : readers_) {
futures.push_back(
Via(executor_.get(), [&action, reader]() -> R { return action(reader); }));
futures.push_back(Via(executor_.get(), [action, reader]() -> R { return action(reader); }));
}
return CollectAll(futures);
}
Expand Down
63 changes: 63 additions & 0 deletions src/paimon/common/global_index/union_global_index_reader_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
#include <atomic>
#include <map>
#include <memory>
#include <queue>
#include <stdexcept>
#include <string>
#include <vector>

Expand Down Expand Up @@ -54,6 +56,11 @@ class FakeReader : public GlobalIndexReader {
error_message_ = message;
}

void SetThrowException(const std::string& message) {
throw_exception_ = true;
exception_message_ = message;
}

/// Sets a scored result returned by VisitVectorSearch.
void SetScoredResult(const std::vector<int64_t>& row_ids, const std::vector<float>& scores) {
scored_row_ids_ = row_ids;
Expand Down Expand Up @@ -161,6 +168,9 @@ class FakeReader : public GlobalIndexReader {
private:
Result<std::shared_ptr<GlobalIndexResult>> MakeResult() {
invocation_count_++;
if (throw_exception_) {
throw std::runtime_error(exception_message_);
}
if (return_error_) {
return Status::Invalid(error_message_);
}
Expand All @@ -176,14 +186,50 @@ class FakeReader : public GlobalIndexReader {
std::vector<int64_t> default_result_;
bool return_nullptr_ = false;
bool return_error_ = false;
bool throw_exception_ = false;
std::string error_message_;
std::string exception_message_;
std::vector<int64_t> scored_row_ids_;
std::vector<float> scored_scores_;
bool has_scored_result_ = false;
bool thread_safe_ = true;
std::atomic<int32_t> invocation_count_{0};
};

// Runs the first task immediately and defers the rest. This makes it possible to verify that
// queued tasks own their action even if collecting an earlier future throws.
class DeferAfterFirstExecutor : public Executor {
public:
void Add(std::function<void()> func) override {
if (submission_count_++ == 0) {
func();
} else {
pending_tasks_.push(std::move(func));
}
}

void ShutdownNow() override {
std::queue<std::function<void()>> empty;
pending_tasks_.swap(empty);
}

uint32_t GetThreadNum() const override {
return 1;
}

void RunPendingTasks() {
while (!pending_tasks_.empty()) {
std::function<void()> task = std::move(pending_tasks_.front());
pending_tasks_.pop();
task();
}
}

private:
uint32_t submission_count_ = 0;
std::queue<std::function<void()>> pending_tasks_;
};

class UnionGlobalIndexReaderTest : public ::testing::Test {
public:
static void CheckResult(const std::shared_ptr<GlobalIndexResult>& result,
Expand Down Expand Up @@ -341,6 +387,23 @@ TEST_F(UnionGlobalIndexReaderTest, TestErrorPropagationWithExecutor) {
ASSERT_NOK_WITH_MSG(union_reader.VisitIsNotNull(), "Unknown error for reader2");
}

TEST_F(UnionGlobalIndexReaderTest, TestDeferredTaskOwnsActionAfterEarlierFutureThrows) {
auto throwing_reader = std::make_shared<FakeReader>();
auto deferred_reader = std::make_shared<FakeReader>();
throwing_reader->SetThrowException("reader exception");
deferred_reader->SetDefaultResult({2});

std::vector<std::shared_ptr<GlobalIndexReader>> readers = {throwing_reader, deferred_reader};
auto executor = std::make_shared<DeferAfterFirstExecutor>();
UnionGlobalIndexReader union_reader(std::move(readers), executor);
ASSERT_THROW(
{ [[maybe_unused]] auto result = union_reader.VisitIsNotNull(); }, std::runtime_error);
ASSERT_EQ(deferred_reader->InvocationCount(), 0);

executor->RunPendingTasks();
ASSERT_EQ(deferred_reader->InvocationCount(), 1);
}

TEST_F(UnionGlobalIndexReaderTest, TestVisitEqualUnion) {
auto reader1 = std::make_shared<FakeReader>();
auto reader2 = std::make_shared<FakeReader>();
Expand Down
Loading
Loading