diff --git a/docs/source/user_guide/global_index.rst b/docs/source/user_guide/global_index.rst index 956d28dba..eed71866c 100644 --- a/docs/source/user_guide/global_index.rst +++ b/docs/source/user_guide/global_index.rst @@ -168,9 +168,12 @@ index shard, the engine ranks matching rows by BM25 score and returns at most `` the highest scores. ``include_row_ids`` holds global row ids and restricts the rows that are ranked; an empty set matches no rows. -For both ``full-text`` and ``lucene-fts``, searches across multiple shards currently return the -union of the shard candidates, which can exceed ``limit``. Applying a final global top-k is tracked -in `issue #402 `_. +For both ``full-text`` and ``lucene-fts``, a search across multiple index shards keeps the +``limit`` highest-scoring rows returned by the shards, breaking ties by ascending row id. Which rows +tied at a shard's cutoff score the shard returns is up to its engine. A row in overlapping index +ranges is ranked only by the range that starts first, or ends first if both start at the same row. +A shard ignores the ``include_row_ids`` outside the rows it ranks and is not searched if none +remain. The result iterates rows by ascending row id, not by score. **Configuration:** diff --git a/include/paimon/global_index/global_index_reader.h b/include/paimon/global_index/global_index_reader.h index 6155759bc..94d9ecef8 100644 --- a/include/paimon/global_index/global_index_reader.h +++ b/include/paimon/global_index/global_index_reader.h @@ -45,9 +45,12 @@ class PAIMON_EXPORT GlobalIndexReader : public FunctionVisitor> VisitVectorSearch( const std::shared_ptr& vector_search) = 0; - /// VisitFullTextSearch performs full text search and returns row ids with relevance scores. - /// Each index shard returns at most `full_text_search->limit` rows. Union readers currently - /// retain all shard candidates and can return more than that limit. + /// VisitFullTextSearch performs full text search and returns at most `full_text_search->limit` + /// row ids with the highest relevance scores, iterated by ascending row id rather than score. + /// A search over multiple index shards keeps the `limit` highest-scoring rows returned by the + /// shards, breaking ties by ascending row id; which rows tied at a shard's cutoff score the + /// shard returns is up to its index engine. A row in overlapping index ranges is ranked only by + /// the range that starts first, or ends first if both start at the same row. /// @warning `VisitFullTextSearch` may return error status when it is incorrectly invoked (e.g., /// BTreeGlobalIndexReader call `VisitFullTextSearch`) or the query is invalid. virtual Result> VisitFullTextSearch( diff --git a/include/paimon/predicate/full_text_search.h b/include/paimon/predicate/full_text_search.h index 1adb08759..46e8024db 100644 --- a/include/paimon/predicate/full_text_search.h +++ b/include/paimon/predicate/full_text_search.h @@ -17,6 +17,7 @@ */ #pragma once +#include #include #include #include @@ -49,9 +50,23 @@ struct PAIMON_EXPORT FullTextSearch { } const RoaringBitmap64& global_row_ids = include_row_ids.value(); RoaringBitmap64 local_row_ids; - for (auto iter = global_row_ids.EqualOrLarger(from), end = global_row_ids.End(); - iter != end && *iter <= to; ++iter) { - local_row_ids.Add(*iter - from); + auto iter = global_row_ids.EqualOrLarger(from); + auto end = global_row_ids.End(); + if (iter != end && *iter <= to) { + // Shift one run of consecutive row ids as a range, so that converting a range of rows + // to rank does not visit each row. + int64_t first = *global_row_ids.Begin(); + int64_t last = first + (global_row_ids.Cardinality() - 1); + if (last == RoaringBitmap64::MAX_VALUE || + global_row_ids.EqualOrLarger(last + 1) == end) { + int64_t local_to = std::min(last, to) - from; + local_row_ids.AddRange(*iter - from, local_to); + local_row_ids.Add(local_to); + } else { + for (; iter != end && *iter <= to; ++iter) { + local_row_ids.Add(*iter - from); + } + } } return std::make_shared(field_name, query, limit, std::move(local_row_ids)); } @@ -83,7 +98,8 @@ struct PAIMON_EXPORT FullTextSearch { /// and rejects `boost` queries and a `fuzziness` other than 0. An invalid query is reported as /// an error `Status`. std::string query; - /// Maximum number of highest-scoring rows requested from each index shard; must be positive. + /// Maximum number of highest-scoring rows to return; must be positive. Each index shard is + /// asked for this many rows before the shard results are merged. int32_t limit; /// The row ids to search within, aligned with Java `includeRowIds`. The row id domain is /// determined by the reader: global row ids for GlobalIndexReader and physical file-local row diff --git a/src/paimon/common/global_index/offset_global_index_reader.cpp b/src/paimon/common/global_index/offset_global_index_reader.cpp index 70b6afc69..b66a29b4d 100644 --- a/src/paimon/common/global_index/offset_global_index_reader.cpp +++ b/src/paimon/common/global_index/offset_global_index_reader.cpp @@ -24,8 +24,8 @@ namespace paimon { OffsetGlobalIndexReader::OffsetGlobalIndexReader(std::shared_ptr&& wrapped, - int64_t offset) - : wrapped_(std::move(wrapped)), offset_(offset) {} + int64_t offset, int64_t to) + : wrapped_(std::move(wrapped)), offset_(offset), to_(to) {} Result> OffsetGlobalIndexReader::VisitIsNotNull() { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, wrapped_->VisitIsNotNull()); @@ -137,11 +137,10 @@ Result> OffsetGlobalIndexReader::VisitV Result> OffsetGlobalIndexReader::VisitFullTextSearch( const std::shared_ptr& full_text_search) { - // Convert global include row ids to local row ids. The shard end is unknown here, so row ids - // past the shard are kept; they match no local row. + // Convert global include row ids to local row ids, dropping those outside [offset_, to_]. std::shared_ptr local_search = full_text_search; if (full_text_search && full_text_search->include_row_ids) { - local_search = full_text_search->OffsetRange(offset_, RoaringBitmap64::MAX_VALUE); + local_search = full_text_search->OffsetRange(offset_, to_); } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, wrapped_->VisitFullTextSearch(local_search)); diff --git a/src/paimon/common/global_index/offset_global_index_reader.h b/src/paimon/common/global_index/offset_global_index_reader.h index 082294665..de86fdee3 100644 --- a/src/paimon/common/global_index/offset_global_index_reader.h +++ b/src/paimon/common/global_index/offset_global_index_reader.h @@ -34,8 +34,11 @@ class OffsetGlobalIndexReader : public GlobalIndexReader { /// Constructs an OffsetGlobalIndexReader. /// @param wrapped The inner reader to delegate queries to. Its results are expected to /// contain local row ids. - /// @param offset The offset to add to each row id in the results. - OffsetGlobalIndexReader(std::shared_ptr&& wrapped, int64_t offset); + /// @param offset The first global row id covered by the wrapped reader, added to each row id + /// in the results. + /// @param to The last global row id (inclusive) covered by the wrapped reader. + OffsetGlobalIndexReader(std::shared_ptr&& wrapped, int64_t offset, + int64_t to); Result> VisitIsNotNull() override; Result> VisitIsNull() override; @@ -78,6 +81,7 @@ class OffsetGlobalIndexReader : public GlobalIndexReader { private: std::shared_ptr wrapped_; int64_t offset_; + int64_t to_; }; } // namespace paimon diff --git a/src/paimon/common/global_index/offset_global_index_reader_test.cpp b/src/paimon/common/global_index/offset_global_index_reader_test.cpp index 39fec0b11..83d6079bc 100644 --- a/src/paimon/common/global_index/offset_global_index_reader_test.cpp +++ b/src/paimon/common/global_index/offset_global_index_reader_test.cpp @@ -19,6 +19,7 @@ #include "paimon/common/global_index/offset_global_index_reader.h" +#include #include #include @@ -65,7 +66,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitEqualWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetDefaultResult({0, 1, 3}); - auto offset_reader = std::make_shared(fake_reader, 100); + auto offset_reader = std::make_shared(fake_reader, 100, /*to=*/199); Literal literal_5(5); ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitEqual(literal_5)); @@ -77,7 +78,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitIsNotNullWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetDefaultResult({0, 2, 4, 6}); - auto offset_reader = std::make_shared(fake_reader, 50); + auto offset_reader = std::make_shared(fake_reader, 50, /*to=*/149); ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitIsNotNull()); // row ids {0, 2, 4, 6} + offset 50 -> {50, 52, 54, 56} @@ -88,7 +89,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitIsNullWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetDefaultResult({1, 3}); - auto offset_reader = std::make_shared(fake_reader, 200); + auto offset_reader = std::make_shared(fake_reader, 200, /*to=*/299); ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitIsNull()); // row ids {1, 3} + offset 200 -> {201, 203} @@ -99,7 +100,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitLessThanWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetDefaultResult({0, 1}); - auto offset_reader = std::make_shared(fake_reader, 10); + auto offset_reader = std::make_shared(fake_reader, 10, /*to=*/109); Literal literal_5(5); ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitLessThan(literal_5)); @@ -111,7 +112,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitGreaterOrEqualWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetDefaultResult({3, 4, 5}); - auto offset_reader = std::make_shared(fake_reader, 1000); + auto offset_reader = std::make_shared(fake_reader, 1000, /*to=*/1099); Literal literal_10(10); ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitGreaterOrEqual(literal_10)); @@ -123,7 +124,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitInWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetDefaultResult({0, 2}); - auto offset_reader = std::make_shared(fake_reader, 5); + auto offset_reader = std::make_shared(fake_reader, 5, /*to=*/104); std::vector literals = {Literal(1), Literal(3)}; ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitIn(literals)); @@ -135,7 +136,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitNotInWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetDefaultResult({1, 3, 5}); - auto offset_reader = std::make_shared(fake_reader, 20); + auto offset_reader = std::make_shared(fake_reader, 20, /*to=*/119); std::vector literals = {Literal(2)}; ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitNotIn(literals)); @@ -147,7 +148,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitNotEqualWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetDefaultResult({0, 2, 4}); - auto offset_reader = std::make_shared(fake_reader, 7); + auto offset_reader = std::make_shared(fake_reader, 7, /*to=*/106); Literal literal_1(1); ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitNotEqual(literal_1)); @@ -159,7 +160,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitLessOrEqualWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetDefaultResult({0, 1, 2}); - auto offset_reader = std::make_shared(fake_reader, 30); + auto offset_reader = std::make_shared(fake_reader, 30, /*to=*/129); Literal literal_3(3); ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitLessOrEqual(literal_3)); @@ -171,7 +172,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitGreaterThanWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetDefaultResult({4, 5}); - auto offset_reader = std::make_shared(fake_reader, 15); + auto offset_reader = std::make_shared(fake_reader, 15, /*to=*/114); Literal literal_3(3); ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitGreaterThan(literal_3)); @@ -183,7 +184,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestZeroOffset) { auto fake_reader = std::make_shared(); fake_reader->SetDefaultResult({0, 1, 2}); - auto offset_reader = std::make_shared(fake_reader, 0); + auto offset_reader = std::make_shared(fake_reader, 0, /*to=*/99); ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitIsNotNull()); // row ids {0, 1, 2} + offset 0 -> {0, 1, 2} @@ -194,7 +195,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestEmptyResultWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetDefaultResult({}); - auto offset_reader = std::make_shared(fake_reader, 100); + auto offset_reader = std::make_shared(fake_reader, 100, /*to=*/199); ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitIsNotNull()); // empty result + offset 100 -> still empty @@ -203,14 +204,14 @@ TEST_F(OffsetGlobalIndexReaderTest, TestEmptyResultWithOffset) { TEST_F(OffsetGlobalIndexReaderTest, TestIsThreadSafeDelegated) { auto fake_reader = std::make_shared(); - auto offset_reader = std::make_shared(fake_reader, 100); + auto offset_reader = std::make_shared(fake_reader, 100, /*to=*/199); // MockGlobalIndexReader returns true for IsThreadSafe ASSERT_TRUE(offset_reader->IsThreadSafe()); } TEST_F(OffsetGlobalIndexReaderTest, TestGetIndexTypeDelegated) { auto fake_reader = std::make_shared(); - auto offset_reader = std::make_shared(fake_reader, 100); + auto offset_reader = std::make_shared(fake_reader, 100, /*to=*/199); ASSERT_EQ(offset_reader->GetIndexType(), "fake"); } @@ -218,7 +219,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitFullTextSearchWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetScoredResult({0, 3, 5}, {0.5f, 0.9f, 0.1f}); - auto offset_reader = std::make_shared(fake_reader, 10); + auto offset_reader = std::make_shared(fake_reader, 10, /*to=*/109); ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitFullTextSearch(nullptr)); // row ids {0, 3, 5} + offset 10 -> {10, 13, 15} @@ -227,23 +228,23 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitFullTextSearchWithOffset) { TEST_F(OffsetGlobalIndexReaderTest, TestVisitFullTextSearchConvertsIncludeRowIds) { auto fake_reader = std::make_shared(); - auto offset_reader = std::make_shared(fake_reader, 10); + auto offset_reader = std::make_shared(fake_reader, 10, /*to=*/19); { auto search = std::make_shared("f0", R"({"match":{"query":"paimon"}})", /*limit=*/7, - RoaringBitmap64::From({5l, 10l, 13l, 15l})); + RoaringBitmap64::From({5l, 10l, 13l, 19l, 20l, 150l})); ASSERT_OK(offset_reader->VisitFullTextSearch(search)); auto captured = fake_reader->CapturedFullTextSearch(); ASSERT_TRUE(captured); ASSERT_EQ(captured->field_name, "f0"); ASSERT_EQ(captured->query, R"({"match":{"query":"paimon"}})"); ASSERT_EQ(captured->limit, 7); - ASSERT_EQ(captured->include_row_ids, RoaringBitmap64::From({0l, 3l, 5l})); - ASSERT_EQ(search->include_row_ids, RoaringBitmap64::From({5l, 10l, 13l, 15l})); + ASSERT_EQ(captured->include_row_ids, RoaringBitmap64::From({0l, 3l, 9l})); + ASSERT_EQ(search->include_row_ids, RoaringBitmap64::From({5l, 10l, 13l, 19l, 20l, 150l})); } { - auto search = std::make_shared("f0", R"({"match":{"query":"paimon"}})", - /*limit=*/7, RoaringBitmap64::From({5l})); + auto search = std::make_shared( + "f0", R"({"match":{"query":"paimon"}})", /*limit=*/7, RoaringBitmap64::From({5l, 20l})); ASSERT_OK(offset_reader->VisitFullTextSearch(search)); ASSERT_TRUE(fake_reader->CapturedFullTextSearch()->include_row_ids); ASSERT_TRUE(fake_reader->CapturedFullTextSearch()->include_row_ids->IsEmpty()); @@ -254,13 +255,25 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitFullTextSearchConvertsIncludeRowIds ASSERT_OK(offset_reader->VisitFullTextSearch(search)); ASSERT_EQ(fake_reader->CapturedFullTextSearch(), search); } + { + constexpr int64_t kMax = std::numeric_limits::max(); + auto last_reader = std::make_shared(); + auto last_offset_reader = + std::make_shared(last_reader, kMax - 9, /*to=*/kMax); + auto search = + std::make_shared("f0", R"({"match":{"query":"paimon"}})", /*limit=*/7, + RoaringBitmap64::From({kMax - 5, kMax - 2, kMax})); + ASSERT_OK(last_offset_reader->VisitFullTextSearch(search)); + ASSERT_EQ(last_reader->CapturedFullTextSearch()->include_row_ids, + RoaringBitmap64::From({4l, 7l, 9l})); + } } TEST_F(OffsetGlobalIndexReaderTest, TestVisitVectorSearchWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetScoredResult({0, 2, 5}, {0.9f, 0.7f, 0.3f}); - auto offset_reader = std::make_shared(fake_reader, 100); + auto offset_reader = std::make_shared(fake_reader, 100, /*to=*/199); ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitVectorSearch(nullptr)); // row ids {0, 2, 5} + offset 100 -> {100, 102, 105}, scores unchanged @@ -269,7 +282,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitVectorSearchWithOffset) { TEST_F(OffsetGlobalIndexReaderTest, TestVisitVectorSearchNotSupported) { auto fake_reader = std::make_shared(); - auto offset_reader = std::make_shared(fake_reader, 10); + auto offset_reader = std::make_shared(fake_reader, 10, /*to=*/109); // Propagate the wrapped reader's unsupported-search error. fake_reader->SetReturnError("vector search not supported"); ASSERT_NOK_WITH_MSG(offset_reader->VisitVectorSearch(nullptr), "vector search not supported"); @@ -279,7 +292,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitStartsWithWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetDefaultResult({0, 1, 4}); - auto offset_reader = std::make_shared(fake_reader, 50); + auto offset_reader = std::make_shared(fake_reader, 50, /*to=*/149); Literal prefix(FieldType::STRING, "abc", 3); ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitStartsWith(prefix)); @@ -291,7 +304,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitEndsWithWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetDefaultResult({2, 3}); - auto offset_reader = std::make_shared(fake_reader, 25); + auto offset_reader = std::make_shared(fake_reader, 25, /*to=*/124); Literal suffix(FieldType::STRING, "xyz", 3); ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitEndsWith(suffix)); @@ -303,7 +316,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitContainsWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetDefaultResult({1, 5, 6}); - auto offset_reader = std::make_shared(fake_reader, 300); + auto offset_reader = std::make_shared(fake_reader, 300, /*to=*/399); Literal literal(FieldType::STRING, "foo", 3); ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitContains(literal)); @@ -315,7 +328,7 @@ TEST_F(OffsetGlobalIndexReaderTest, TestVisitLikeWithOffset) { auto fake_reader = std::make_shared(); fake_reader->SetDefaultResult({0, 3, 7}); - auto offset_reader = std::make_shared(fake_reader, 40); + auto offset_reader = std::make_shared(fake_reader, 40, /*to=*/139); Literal pattern(FieldType::STRING, "ab%", 3); ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitLike(pattern)); diff --git a/src/paimon/common/global_index/union_global_index_reader.cpp b/src/paimon/common/global_index/union_global_index_reader.cpp index 27966f6b5..bb11cae39 100644 --- a/src/paimon/common/global_index/union_global_index_reader.cpp +++ b/src/paimon/common/global_index/union_global_index_reader.cpp @@ -19,16 +19,98 @@ #include "paimon/common/global_index/union_global_index_reader.h" +#include +#include #include +#include #include #include "paimon/common/executor/future.h" +#include "paimon/common/utils/fields_comparator.h" +#include "paimon/global_index/bitmap_scored_global_index_result.h" +#include "paimon/predicate/full_text_search.h" +#include "paimon/utils/roaring_bitmap64.h" namespace paimon { +namespace { +std::shared_ptr EmptyScoredResult() { + return std::make_shared(RoaringBitmap64(), std::vector()); +} + +RoaringBitmap64 RowIdsOf(const Range& range) { + // `range.to` may be the largest int64 value, so add it apart from the half-open range. + RoaringBitmap64 row_ids; + row_ids.AddRange(range.from, range.to); + row_ids.Add(range.to); + return row_ids; +} + +/// Returns the row ids of `row_ids` within `range`. Trimming a copy costs time in proportion to +/// the containers of `row_ids`, not to the rows or the span of `range`. +RoaringBitmap64 RowIdsWithin(const RoaringBitmap64& row_ids, const Range& range) { + RoaringBitmap64 result = row_ids; + result.RemoveRange(0, range.from); + if (range.to < RoaringBitmap64::MAX_VALUE) { + // RemoveRange leaves out its end, so the largest int64 value is removed on its own. + result.RemoveRange(range.to + 1, RoaringBitmap64::MAX_VALUE); + result -= RoaringBitmap64::From({RoaringBitmap64::MAX_VALUE}); + } + return result; +} + +/// Returns whether `left` ranks before `right`: a higher score first, then a lower row id. Scores +/// are ordered as -infinity < -0.0 < +0.0 < +infinity < NaN. +bool RanksBefore(const std::pair& left, const std::pair& right) { + int32_t score_order = FieldsComparator::CompareFloatingPoint(left.second, right.second); + if (score_order != 0) { + return score_order > 0; + } + return left.first < right.first; +} + +/// Merges scored results and keeps the `limit` rows that rank first. A row id in more than one +/// result keeps its first score. +Result> MergeTopK( + const std::vector>& results, int32_t limit) { + std::unordered_set merged_row_ids; + std::vector> rows; + for (const auto& result : results) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr iter, + result->CreateScoredIterator()); + while (iter->HasNext()) { + std::pair row = iter->NextWithScore(); + if (merged_row_ids.insert(row.first).second) { + rows.push_back(row); + } + } + } + auto top_k = static_cast(limit); + if (rows.size() > top_k) { + std::nth_element(rows.begin(), rows.begin() + top_k, rows.end(), RanksBefore); + rows.resize(top_k); + } + std::sort(rows.begin(), rows.end()); + RoaringBitmap64 bitmap; + std::vector scores; + scores.reserve(rows.size()); + for (const auto& [row_id, score] : rows) { + bitmap.Add(row_id); + scores.push_back(score); + } + return std::make_shared(std::move(bitmap), std::move(scores)); +} +} // namespace + UnionGlobalIndexReader::UnionGlobalIndexReader( - std::vector>&& readers, - const std::shared_ptr& executor) - : readers_(std::move(readers)), executor_(executor) {} + std::vector>&& readers, std::vector&& row_ranges, + std::vector>&& search_ranges, const std::shared_ptr& executor) + : readers_(std::move(readers)), + executor_(executor), + row_ranges_(std::move(row_ranges)), + search_ranges_(std::move(search_ranges)) { + assert(row_ranges_.size() == search_ranges_.size()); + assert(search_ranges_.empty() || search_ranges_.size() == readers_.size()); +} Result> UnionGlobalIndexReader::VisitIsNotNull() { return Union( @@ -133,9 +215,75 @@ Result> UnionGlobalIndexReader::VisitVe Result> UnionGlobalIndexReader::VisitFullTextSearch( const std::shared_ptr& full_text_search) { - return UnionScored([full_text_search](const std::shared_ptr& reader) { - return reader->VisitFullTextSearch(full_text_search); - }); + if (!full_text_search) { + return Status::Invalid("VisitFullTextSearch: null FullTextSearch pointer"); + } + if (full_text_search->limit <= 0) { + return Status::Invalid("full-text search requires a positive limit"); + } + // If the readers rank disjoint rows, the top `limit` rows of each reader include all of its + // rows scoring above the `limit`-th highest score of all readers, so merging them loses none + // of those rows. Rows tied at that score may already have been cut by a reader's engine. + auto reader_searches = std::make_shared>>(); + reader_searches->reserve(readers_.size()); + for (size_t index = 0; index < readers_.size(); ++index) { + reader_searches->push_back(FullTextSearchForReader(index, full_text_search)); + } + auto results = ExecuteAllReadersWithIndex>>( + [reader_searches](size_t index, const std::shared_ptr& reader) + -> Result> { + const std::shared_ptr& reader_search = (*reader_searches)[index]; + if (!reader_search) { + return EmptyScoredResult(); + } + return reader->VisitFullTextSearch(reader_search); + }); + std::vector> scored_results; + scored_results.reserve(results.size()); + for (auto& result_or_status : results) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, + std::move(result_or_status)); + if (result != nullptr) { + scored_results.push_back(std::move(result)); + } + } + // No reader answered the search, which differs from a search that matches no rows. + if (scored_results.empty()) { + return std::shared_ptr(); + } + return MergeTopK(scored_results, full_text_search->limit); +} + +std::shared_ptr UnionGlobalIndexReader::FullTextSearchForReader( + size_t index, const std::shared_ptr& full_text_search) const { + if (search_ranges_.empty()) { + return full_text_search; + } + const std::optional& search_range = search_ranges_[index]; + if (!search_range) { + return nullptr; + } + bool ranks_all_rows = search_range.value() == row_ranges_[index]; + if (!full_text_search->include_row_ids) { + if (ranks_all_rows) { + return full_text_search; + } + return std::make_shared(full_text_search->field_name, + full_text_search->query, full_text_search->limit, + RowIdsOf(search_range.value())); + } + const RoaringBitmap64& include_row_ids = full_text_search->include_row_ids.value(); + auto first = include_row_ids.EqualOrLarger(search_range->from); + if (first == include_row_ids.End() || *first > search_range->to) { + return nullptr; + } + // The reader drops include row ids outside its rows, so they need no clipping here. + if (ranks_all_rows) { + return full_text_search; + } + return std::make_shared(full_text_search->field_name, full_text_search->query, + full_text_search->limit, + RowIdsWithin(include_row_ids, search_range.value())); } Result> UnionGlobalIndexReader::UnionScored( @@ -198,11 +346,20 @@ bool UnionGlobalIndexReader::IsThreadSafe() const { template std::vector UnionGlobalIndexReader::ExecuteAllReaders( const std::function&)>& action) { + return ExecuteAllReadersWithIndex( + [action](size_t, const std::shared_ptr& reader) -> R { + return action(reader); + }); +} + +template +std::vector UnionGlobalIndexReader::ExecuteAllReadersWithIndex( + const std::function&)>& action) { if (executor_ == nullptr || readers_.size() == 1) { std::vector results; results.reserve(readers_.size()); - for (const auto& reader : readers_) { - results.push_back(action(reader)); + for (size_t index = 0; index < readers_.size(); ++index) { + results.push_back(action(index, readers_[index])); } return results; } @@ -210,8 +367,10 @@ std::vector UnionGlobalIndexReader::ExecuteAllReaders( // Parallel: submit all tasks to executor, then collect results in submission order std::vector> futures; futures.reserve(readers_.size()); - for (const auto& reader : readers_) { - futures.push_back(Via(executor_.get(), [action, reader]() -> R { return action(reader); })); + for (size_t index = 0; index < readers_.size(); ++index) { + futures.push_back(Via(executor_.get(), [action, index, reader = readers_[index]]() -> R { + return action(index, reader); + })); } return CollectAll(futures); } diff --git a/src/paimon/common/global_index/union_global_index_reader.h b/src/paimon/common/global_index/union_global_index_reader.h index 249928608..811e24620 100644 --- a/src/paimon/common/global_index/union_global_index_reader.h +++ b/src/paimon/common/global_index/union_global_index_reader.h @@ -21,11 +21,13 @@ #include #include +#include #include #include #include "paimon/executor.h" #include "paimon/global_index/global_index_reader.h" +#include "paimon/utils/range.h" namespace paimon { /// A GlobalIndexReader that combines results from multiple readers by performing a union @@ -35,7 +37,14 @@ namespace paimon { /// collected in submission order. Otherwise, sub-readers are evaluated sequentially. class UnionGlobalIndexReader : public GlobalIndexReader { public: + /// Constructs a UnionGlobalIndexReader. + /// @param row_ranges The global row ids covered by each reader, in the order of `readers`. + /// @param search_ranges The global row ids each reader ranks in a full-text search, within its + /// row range, in the order of `readers`; std::nullopt if it ranks none. + /// If both are empty, each reader ranks all of its rows. UnionGlobalIndexReader(std::vector>&& readers, + std::vector&& row_ranges, + std::vector>&& search_ranges, const std::shared_ptr& executor); Result> VisitIsNotNull() override; @@ -58,6 +67,11 @@ class UnionGlobalIndexReader : public GlobalIndexReader { Result> VisitVectorSearch( const std::shared_ptr& vector_search) override; + /// Merges the scored results of all readers and keeps the `full_text_search->limit` rows with + /// the highest scores among them, breaking ties by ascending row id. A row id returned by more + /// than one reader keeps the score from the first of those readers. Returns nullptr if no + /// reader returns a result, e.g. if the readers do not support full-text search; a reader that + /// ranks none of the rows to search is not visited and counts as returning an empty result. Result> VisitFullTextSearch( const std::shared_ptr& full_text_search) override; @@ -81,13 +95,26 @@ class UnionGlobalIndexReader : public GlobalIndexReader { /// Executes the given action on all readers and merges the scored results with Union. Result> UnionScored(ScoredReaderAction action); + /// Returns the full-text search to run on the reader at `index`, or nullptr if the reader + /// ranks none of the rows to search. + std::shared_ptr FullTextSearchForReader( + size_t index, const std::shared_ptr& full_text_search) const; + /// Executes the given action on all readers (parallel or sequential) and collects results. template std::vector ExecuteAllReaders( const std::function&)>& action); + /// Executes the given action on each reader and its index (parallel or sequential) and + /// collects results in reader order. + template + std::vector ExecuteAllReadersWithIndex( + const std::function&)>& action); + std::vector> readers_; std::shared_ptr executor_; + std::vector row_ranges_; + std::vector> search_ranges_; }; } // namespace paimon diff --git a/src/paimon/common/global_index/union_global_index_reader_test.cpp b/src/paimon/common/global_index/union_global_index_reader_test.cpp index 5399073b4..8b83e134f 100644 --- a/src/paimon/common/global_index/union_global_index_reader_test.cpp +++ b/src/paimon/common/global_index/union_global_index_reader_test.cpp @@ -19,18 +19,22 @@ #include "paimon/common/global_index/union_global_index_reader.h" +#include #include #include +#include #include #include #include #include #include "gtest/gtest.h" +#include "paimon/common/global_index/offset_global_index_reader.h" #include "paimon/common/global_index/sorted_file_global_index_reader.h" #include "paimon/executor.h" #include "paimon/global_index/bitmap_global_index_result.h" #include "paimon/global_index/bitmap_scored_global_index_result.h" +#include "paimon/predicate/full_text_search.h" #include "paimon/predicate/literal.h" #include "paimon/testing/mock/mock_global_index_reader.h" #include "paimon/testing/utils/testharness.h" @@ -145,6 +149,10 @@ class UnionGlobalIndexReaderTest : public ::testing::Test { << ", expected=" << RoaringBitmap64::From(expected_row_ids).ToString(); ASSERT_EQ(typed_result->GetScores(), expected_scores); } + + static std::shared_ptr MakeFullTextSearch(int32_t limit) { + return std::make_shared("f0", R"({"match":{"query":"apple"}})", limit); + } }; TEST_F(UnionGlobalIndexReaderTest, TestSingleReaderUnion) { @@ -153,7 +161,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestSingleReaderUnion) { auto executor = std::make_shared(); std::vector> readers = {reader}; - UnionGlobalIndexReader union_reader(std::move(readers), executor); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, executor); ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitIsNotNull()); CheckResult(result, {1, 2, 3}); @@ -170,7 +179,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestMultipleReadersUnionSequential) { reader3->SetDefaultResult({5}); std::vector> readers = {reader1, reader2, reader3}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitIsNotNull()); // {1,2} OR {3,4} OR {5} -> {1,2,3,4,5} @@ -187,7 +197,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestMultipleReadersUnionOverlappingIds) { reader2->SetDefaultResult({2, 3, 4}); std::vector> readers = {reader1, reader2}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitIsNotNull()); // {1,2,3} OR {2,3,4} -> {1,2,3,4} @@ -204,7 +215,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestMultipleReadersUnionWithExecutor) { std::vector> readers = {reader1, reader2, reader3}; std::shared_ptr executor = CreateDefaultExecutor(); - UnionGlobalIndexReader union_reader(std::move(readers), executor); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, executor); ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitIsNotNull()); CheckResult(result, {10, 20, 30}); @@ -215,7 +227,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestMultipleReadersUnionWithExecutor) { TEST_F(UnionGlobalIndexReaderTest, TestEmptyReaderList) { std::vector> readers; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); // No readers means no results to merge -> nullptr ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitIsNotNull()); @@ -229,7 +242,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestAllReadersReturnNullptr) { reader2->SetReturnNullptr(); std::vector> readers = {reader1, reader2}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitIsNotNull()); // All readers return nullptr -> merged result is nullptr @@ -245,7 +259,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestPartialReadersReturnNullptr) { reader3->SetReturnNullptr(); std::vector> readers = {reader1, reader2, reader3}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitIsNotNull()); // A missing result from any reader makes the whole union unsupported. @@ -259,7 +274,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestPartialReadersReturnNullptrWithExecutor) reader2->SetReturnNullptr(); std::vector> readers = {reader1, reader2}; - UnionGlobalIndexReader union_reader(std::move(readers), CreateDefaultExecutor()); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, CreateDefaultExecutor()); ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitIsNotNull()); ASSERT_FALSE(result); @@ -292,7 +308,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestErrorPropagationSequential) { reader2->SetReturnError("Unknown error for reader2"); std::vector> readers = {reader1, reader2}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); ASSERT_NOK_WITH_MSG(union_reader.VisitIsNotNull(), "Unknown error for reader2"); } @@ -305,7 +322,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestErrorPropagationWithExecutor) { std::vector> readers = {reader1, reader2}; std::shared_ptr executor = CreateDefaultExecutor(); - UnionGlobalIndexReader union_reader(std::move(readers), executor); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, executor); ASSERT_NOK_WITH_MSG(union_reader.VisitIsNotNull(), "Unknown error for reader2"); } @@ -318,7 +336,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestDeferredTaskOwnsActionAfterEarlierFutureT std::vector> readers = {throwing_reader, deferred_reader}; auto executor = std::make_shared(); - UnionGlobalIndexReader union_reader(std::move(readers), executor); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, executor); ASSERT_THROW( { [[maybe_unused]] auto result = union_reader.VisitIsNotNull(); }, std::runtime_error); ASSERT_EQ(deferred_reader->InvocationCount(), 0); @@ -334,7 +353,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestVisitEqualUnion) { reader2->SetDefaultResult({2}); std::vector> readers = {reader1, reader2}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); Literal literal_42(42); ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitEqual(literal_42)); @@ -348,7 +368,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestVisitNotEqualUnion) { reader2->SetDefaultResult({2}); std::vector> readers = {reader1, reader2}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); Literal literal_42(42); ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitNotEqual(literal_42)); @@ -362,7 +383,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestVisitRangeQueriesUnion) { reader2->SetDefaultResult({2}); std::vector> readers = {reader1, reader2}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); Literal literal_5(5); ASSERT_OK_AND_ASSIGN(auto lt, union_reader.VisitLessThan(literal_5)); @@ -382,7 +404,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestVisitInUnion) { reader2->SetDefaultResult({2, 4}); std::vector> readers = {reader1, reader2}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); std::vector literals = {Literal(10), Literal(20)}; ASSERT_OK_AND_ASSIGN(auto in_result, union_reader.VisitIn(literals)); @@ -399,7 +422,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestVisitStringQueriesUnion) { reader2->SetDefaultResult({2}); std::vector> readers = {reader1, reader2}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); Literal literal_str(FieldType::STRING, "abc", 3); ASSERT_OK_AND_ASSIGN(auto starts, union_reader.VisitStartsWith(literal_str)); @@ -419,7 +443,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestVisitIsNullUnion) { reader2->SetDefaultResult({300}); std::vector> readers = {reader1, reader2}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitIsNull()); CheckResult(result, {100, 200, 300}); @@ -433,12 +458,202 @@ TEST_F(UnionGlobalIndexReaderTest, TestVisitFullTextSearchUnion) { reader3->SetScoredResult({2, 6}, {0.8f, 0.6f}); std::vector> readers = {reader1, reader2, reader3}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); - ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitFullTextSearch(nullptr)); + ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitFullTextSearch(MakeFullTextSearch(10))); CheckScoredResult(result, {1, 2, 5, 6}, {0.9f, 0.8f, 0.5f, 0.6f}); } +TEST_F(UnionGlobalIndexReaderTest, TestVisitFullTextSearchKeepsTopKAcrossReaders) { + auto reader1 = std::make_shared(); + auto reader2 = std::make_shared(); + reader1->SetScoredResult({100, 110, 120}, {0.9f, 0.5f, 0.1f}); + reader2->SetScoredResult({0, 10, 20}, {0.5f, 0.8f, 0.3f}); + + std::vector> readers = {reader1, reader2}; + std::shared_ptr executor = CreateDefaultExecutor(); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, executor); + + ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitFullTextSearch(MakeFullTextSearch(3))); + CheckScoredResult(result, {0, 10, 100}, {0.5f, 0.8f, 0.9f}); +} + +TEST_F(UnionGlobalIndexReaderTest, TestVisitFullTextSearchDuplicateRowIdKeepsFirstScore) { + auto reader1 = std::make_shared(); + auto reader2 = std::make_shared(); + reader1->SetScoredResult({1, 2}, {0.6f, 0.3f}); + reader2->SetScoredResult({2, 3}, {0.9f, 0.5f}); + + std::vector> readers = {reader1, reader2}; + std::shared_ptr executor = CreateDefaultExecutor(); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, executor); + + ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitFullTextSearch(MakeFullTextSearch(10))); + CheckScoredResult(result, {1, 2, 3}, {0.6f, 0.3f, 0.5f}); + ASSERT_OK_AND_ASSIGN(result, union_reader.VisitFullTextSearch(MakeFullTextSearch(2))); + CheckScoredResult(result, {1, 3}, {0.6f, 0.5f}); +} + +TEST_F(UnionGlobalIndexReaderTest, TestVisitFullTextSearchOrdersSpecialScores) { + auto reader1 = std::make_shared(); + auto reader2 = std::make_shared(); + reader1->SetScoredResult({1, 2, 3}, {-0.0f, std::numeric_limits::quiet_NaN(), 1.0f}); + reader2->SetScoredResult({4, 5}, {0.0f, std::numeric_limits::infinity()}); + + std::vector> readers = {reader1, reader2}; + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); + + std::vector>> limit_and_row_ids = { + {1, {2}}, {2, {2, 5}}, {3, {2, 3, 5}}, {4, {2, 3, 4, 5}}}; + for (const auto& [limit, expected_row_ids] : limit_and_row_ids) { + ASSERT_OK_AND_ASSIGN(auto result, + union_reader.VisitFullTextSearch(MakeFullTextSearch(limit))); + auto typed_result = std::dynamic_pointer_cast(result); + ASSERT_TRUE(typed_result); + ASSERT_OK_AND_ASSIGN(const RoaringBitmap64* bitmap, typed_result->GetBitmap()); + ASSERT_EQ(*bitmap, RoaringBitmap64::From(expected_row_ids)) << "limit=" << limit; + } +} + +TEST_F(UnionGlobalIndexReaderTest, TestVisitFullTextSearchAllNullptr) { + auto reader1 = std::make_shared(); + auto reader2 = std::make_shared(); + + std::vector> readers = {reader1, reader2}; + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); + + ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitFullTextSearch(MakeFullTextSearch(10))); + ASSERT_FALSE(result); +} + +TEST_F(UnionGlobalIndexReaderTest, TestVisitFullTextSearchErrors) { + auto reader1 = std::make_shared(); + auto reader2 = std::make_shared(); + reader1->SetScoredResult({1}, {0.9f}); + reader2->SetReturnError("full-text search failure"); + + std::vector> readers = {reader1, reader2}; + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); + + ASSERT_NOK_WITH_MSG(union_reader.VisitFullTextSearch(nullptr), "null FullTextSearch pointer"); + ASSERT_NOK_WITH_MSG(union_reader.VisitFullTextSearch(MakeFullTextSearch(0)), + "full-text search requires a positive limit"); + ASSERT_NOK_WITH_MSG(union_reader.VisitFullTextSearch(MakeFullTextSearch(-1)), + "full-text search requires a positive limit"); + ASSERT_EQ(reader1->InvocationCount(), 0); + ASSERT_EQ(reader2->InvocationCount(), 0); + ASSERT_NOK_WITH_MSG(union_reader.VisitFullTextSearch(MakeFullTextSearch(10)), + "full-text search failure"); +} + +TEST_F(UnionGlobalIndexReaderTest, TestVisitFullTextSearchWithinSearchRanges) { + auto reader1 = std::make_shared(); + auto reader2 = std::make_shared(); + auto reader3 = std::make_shared(); + reader1->SetScoredResult({0, 1, 2}, {0.2f, 0.5f, 0.9f}); + reader2->SetScoredResult({1}, {1.0f}); + reader3->SetScoredResult({3}, {0.7f}); + + std::vector> readers = {reader1, reader2, reader3}; + std::shared_ptr executor = CreateDefaultExecutor(); + UnionGlobalIndexReader union_reader(std::move(readers), + /*row_ranges=*/{Range(0, 2), Range(1, 1), Range(1, 3)}, + /*search_ranges=*/{Range(0, 2), std::nullopt, Range(3, 3)}, + executor); + + auto search = MakeFullTextSearch(2); + ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitFullTextSearch(search)); + CheckScoredResult(result, {2, 3}, {0.9f, 0.7f}); + ASSERT_EQ(reader1->CapturedFullTextSearch(), search); + ASSERT_EQ(reader2->InvocationCount(), 0); + ASSERT_EQ(reader3->CapturedFullTextSearch()->include_row_ids, RoaringBitmap64::From({3l})); + + constexpr int64_t kMax = std::numeric_limits::max(); + search->include_row_ids = RoaringBitmap64::From({1, 2, 3, 5, kMax}); + ASSERT_OK(union_reader.VisitFullTextSearch(search)); + ASSERT_EQ(reader1->CapturedFullTextSearch(), search); + ASSERT_EQ(reader2->InvocationCount(), 0); + ASSERT_EQ(reader3->CapturedFullTextSearch()->include_row_ids, RoaringBitmap64::From({3l})); + ASSERT_EQ(search->include_row_ids, RoaringBitmap64::From({1, 2, 3, 5, kMax})); + + search->include_row_ids = RoaringBitmap64::From({5}); + ASSERT_OK_AND_ASSIGN(result, union_reader.VisitFullTextSearch(search)); + CheckScoredResult(result, {}, {}); + ASSERT_EQ(reader1->InvocationCount(), 2); + ASSERT_EQ(reader2->InvocationCount(), 0); + ASSERT_EQ(reader3->InvocationCount(), 2); +} + +TEST_F(UnionGlobalIndexReaderTest, TestVisitFullTextSearchSparseIncludeInLargeRanges) { + constexpr int64_t kMax = std::numeric_limits::max(); + auto reader1 = std::make_shared(); + auto reader2 = std::make_shared(); + std::vector> readers = {reader1, reader2}; + UnionGlobalIndexReader union_reader( + std::move(readers), /*row_ranges=*/{Range(0, kMax / 2), Range(10, kMax)}, + /*search_ranges=*/{Range(0, kMax / 2), Range(kMax / 2 + 1, kMax)}, /*executor=*/nullptr); + + auto search = MakeFullTextSearch(10); + search->include_row_ids = RoaringBitmap64(); + ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitFullTextSearch(search)); + CheckScoredResult(result, {}, {}); + ASSERT_EQ(reader1->InvocationCount(), 0); + ASSERT_EQ(reader2->InvocationCount(), 0); + + search->include_row_ids = RoaringBitmap64::From({20, kMax / 2 + 1, kMax}); + ASSERT_OK(union_reader.VisitFullTextSearch(search)); + ASSERT_EQ(reader1->CapturedFullTextSearch(), search); + ASSERT_EQ(reader2->CapturedFullTextSearch()->include_row_ids, + RoaringBitmap64::From({kMax / 2 + 1, kMax})); + + search->include_row_ids = RoaringBitmap64::From({20l}); + ASSERT_OK(union_reader.VisitFullTextSearch(search)); + ASSERT_EQ(reader1->InvocationCount(), 2); + ASSERT_EQ(reader2->InvocationCount(), 1); +} + +TEST_F(UnionGlobalIndexReaderTest, TestVisitFullTextSearchOverlappingRangesThroughOffsetReaders) { + constexpr int64_t kMax = std::numeric_limits::max(); + const int64_t to = int64_t{1} << 32; + std::vector> wrapped_readers; + std::vector> readers; + std::vector row_ranges = {Range(0, 99), Range(50, to), Range(kMax - 9, kMax - 3), + Range(kMax - 5, kMax)}; + std::vector> search_ranges = { + Range(0, 99), Range(100, to), Range(kMax - 9, kMax - 3), Range(kMax - 2, kMax)}; + for (const Range& row_range : row_ranges) { + wrapped_readers.push_back(std::make_shared()); + readers.push_back(std::make_shared(wrapped_readers.back(), + row_range.from, row_range.to)); + } + UnionGlobalIndexReader union_reader(std::move(readers), std::move(row_ranges), + std::move(search_ranges), /*executor=*/nullptr); + + auto search = MakeFullTextSearch(1); + ASSERT_OK(union_reader.VisitFullTextSearch(search)); + ASSERT_EQ(wrapped_readers[0]->CapturedFullTextSearch(), search); + RoaringBitmap64 local_row_ids; + local_row_ids.AddRange(50, to - 49); + ASSERT_EQ(wrapped_readers[1]->CapturedFullTextSearch()->include_row_ids, local_row_ids); + ASSERT_EQ(wrapped_readers[2]->CapturedFullTextSearch(), search); + ASSERT_EQ(wrapped_readers[3]->CapturedFullTextSearch()->include_row_ids, + RoaringBitmap64::From({3l, 4l, 5l})); + + RoaringBitmap64 include_row_ids; + include_row_ids.AddRange(0, to + 1); + search->include_row_ids = include_row_ids; + ASSERT_OK(union_reader.VisitFullTextSearch(search)); + ASSERT_EQ(wrapped_readers[1]->CapturedFullTextSearch()->include_row_ids, local_row_ids); + ASSERT_EQ(wrapped_readers[2]->InvocationCount(), 1); + ASSERT_EQ(wrapped_readers[3]->InvocationCount(), 1); +} + TEST_F(UnionGlobalIndexReaderTest, TestVisitVectorSearchAllNullptr) { auto reader1 = std::make_shared(); auto reader2 = std::make_shared(); @@ -447,7 +662,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestVisitVectorSearchAllNullptr) { reader2->SetDefaultResult({2}); std::vector> readers = {reader1, reader2}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitVectorSearch(nullptr)); ASSERT_FALSE(result); @@ -458,7 +674,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestVisitVectorSearchSingleReader) { reader->SetScoredResult({1, 3, 5}, {0.9f, 0.7f, 0.5f}); std::vector> readers = {reader}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitVectorSearch(nullptr)); CheckScoredResult(result, {1, 3, 5}, {0.9f, 0.7f, 0.5f}); @@ -471,7 +688,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestVisitVectorSearchMultipleReadersUnion) { reader2->SetScoredResult({2, 4}, {0.8f, 0.6f}); std::vector> readers = {reader1, reader2}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitVectorSearch(nullptr)); // {1,3} OR {2,4} -> {1,2,3,4}, scores merged in row id order @@ -488,7 +706,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestVisitVectorSearchPartialNullptr) { reader3->SetScoredResult({5, 6}, {0.5f, 0.4f}); std::vector> readers = {reader1, reader2, reader3}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); ASSERT_OK_AND_ASSIGN(auto result, union_reader.VisitVectorSearch(nullptr)); // reader2 nullptr is skipped, {1,2} OR {5,6} -> {1,2,5,6} @@ -502,7 +721,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestVisitVectorSearchErrorPropagation) { reader2->SetReturnError("vector search failure"); std::vector> readers = {reader1, reader2}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); ASSERT_NOK_WITH_MSG(union_reader.VisitVectorSearch(nullptr), "vector search failure"); } @@ -512,7 +732,8 @@ TEST_F(UnionGlobalIndexReaderTest, TestIsThreadSafeReturnsTrueWhenAllReadersAreS auto reader2 = std::make_shared(); std::vector> readers = {reader1, reader2}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); ASSERT_TRUE(union_reader.IsThreadSafe()); } @@ -523,14 +744,16 @@ TEST_F(UnionGlobalIndexReaderTest, TestIsThreadSafeReturnsFalseWhenAnyReaderIsNo reader2->SetThreadSafe(false); std::vector> readers = {reader1, reader2}; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); ASSERT_FALSE(union_reader.IsThreadSafe()); } TEST_F(UnionGlobalIndexReaderTest, TestGetIndexTypeReturnsUnion) { std::vector> readers; - UnionGlobalIndexReader union_reader(std::move(readers), nullptr); + UnionGlobalIndexReader union_reader(std::move(readers), /*row_ranges=*/{}, + /*search_ranges=*/{}, nullptr); ASSERT_EQ(union_reader.GetIndexType(), "union"); } diff --git a/src/paimon/common/predicate/full_text_search_test.cpp b/src/paimon/common/predicate/full_text_search_test.cpp index 2e14dbd38..d5489aed3 100644 --- a/src/paimon/common/predicate/full_text_search_test.cpp +++ b/src/paimon/common/predicate/full_text_search_test.cpp @@ -60,4 +60,28 @@ TEST(FullTextSearchTest, TestOffsetRangeWithIncludeRowIds) { RoaringBitmap64::From({3l, 7l})); } +TEST(FullTextSearchTest, TestOffsetRangeWithIncludeRowIdRun) { + auto row_ids_of = [](int64_t from, int64_t to) { + RoaringBitmap64 row_ids; + row_ids.AddRange(from, to); + row_ids.Add(to); + return row_ids; + }; + FullTextSearch search("f0", R"({"match":{"query":"paimon"}})", /*limit=*/5, + row_ids_of(100, 199)); + ASSERT_EQ(search.OffsetRange(150, 299)->include_row_ids, row_ids_of(0, 49)); + ASSERT_EQ(search.OffsetRange(0, 149)->include_row_ids, row_ids_of(100, 149)); + ASSERT_EQ(search.OffsetRange(120, 130)->include_row_ids, row_ids_of(0, 10)); + ASSERT_EQ(search.OffsetRange(200, 299)->include_row_ids, RoaringBitmap64()); + ASSERT_EQ(search.include_row_ids, row_ids_of(100, 199)); + + constexpr int64_t kMax = RoaringBitmap64::MAX_VALUE; + FullTextSearch last_search("f0", R"({"match":{"query":"paimon"}})", /*limit=*/5, + row_ids_of(kMax - 9, kMax)); + ASSERT_EQ(last_search.OffsetRange(kMax - 5, kMax)->include_row_ids, row_ids_of(0, 5)); + FullTextSearch max_search("f0", R"({"match":{"query":"paimon"}})", /*limit=*/5, + RoaringBitmap64::From({kMax})); + ASSERT_EQ(max_search.OffsetRange(kMax - 5, kMax)->include_row_ids, RoaringBitmap64::From({5l})); +} + } // namespace paimon::test diff --git a/src/paimon/core/global_index/global_index_scan_impl.cpp b/src/paimon/core/global_index/global_index_scan_impl.cpp index c38c50371..eabc10e05 100644 --- a/src/paimon/core/global_index/global_index_scan_impl.cpp +++ b/src/paimon/core/global_index/global_index_scan_impl.cpp @@ -18,6 +18,9 @@ */ #include "paimon/core/global_index/global_index_scan_impl.h" +#include +#include +#include #include #include #include @@ -184,8 +187,21 @@ Result>> GlobalIndexScanImpl::Cre continue; } std::vector> union_readers; + std::vector row_ranges; + std::vector> search_ranges; union_readers.reserve(range_to_metas.size()); + row_ranges.reserve(range_to_metas.size()); + search_ranges.reserve(range_to_metas.size()); + // Assign each full-text row to one reader before requesting per-reader top-k: a row is + // ranked by the first range in (from, to) order that covers it, which leaves each range + // only its rows after the largest `to` of earlier ones. + int64_t covered_to = std::numeric_limits::min(); for (const auto& [range, metas] : range_to_metas) { + std::optional search_range; + if (covered_to < range.to) { + search_range = Range(std::max(range.from, covered_to + 1), range.to); + } + covered_to = std::max(covered_to, range.to); if (row_range_index && !row_range_index->Intersects(range.from, range.to)) { continue; } @@ -199,14 +215,16 @@ Result>> GlobalIndexScanImpl::Cre PAIMON_ASSIGN_OR_RAISE( std::shared_ptr index_reader, indexer->CreateReader(&c_arrow_schema, index_file_manager_, index_io_metas, pool_)); - union_readers.push_back( - std::make_shared(std::move(index_reader), range.from)); + union_readers.push_back(std::make_shared( + std::move(index_reader), range.from, range.to)); + row_ranges.push_back(range); + search_ranges.push_back(search_range); } if (union_readers.empty()) { continue; } - readers.push_back( - std::make_shared(std::move(union_readers), executor_)); + readers.push_back(std::make_shared( + std::move(union_readers), std::move(row_ranges), std::move(search_ranges), executor_)); } return readers; } diff --git a/test/inte/global_index_test.cpp b/test/inte/global_index_test.cpp index e6c1abc59..191d4d1a1 100644 --- a/test/inte/global_index_test.cpp +++ b/test/inte/global_index_test.cpp @@ -16,6 +16,7 @@ * limitations under the License. */ +#include #include #include @@ -74,6 +75,7 @@ class CountingGlobalIndexFileManager : public GlobalIndexFileManager { } // namespace +using ScoredRows = std::vector>; // string: FileFormat, bool: UseSpecificFileSystem using ParamType = std::tuple; /// This is a sdk end-to-end test for global index. @@ -185,6 +187,102 @@ class GlobalIndexTest : public ::testing::Test, public ::testing::WithParamInter return Commit(table_path, {index_commit_msg}); } + static Result SearchFullText( + const std::shared_ptr& reader, const std::string& query, int32_t limit, + const std::optional& include_row_ids) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, + reader->VisitFullTextSearch(std::make_shared( + "f0", query, limit, include_row_ids))); + if (!result) { + return Status::Invalid("full-text search returned no result"); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr iter, + result->CreateScoredIterator()); + ScoredRows rows; + while (iter->HasNext()) { + rows.push_back(iter->NextWithScore()); + } + return rows; + } + + static void CheckFullTextTopK(const std::shared_ptr& reader, + const std::string& query, + const std::optional& include_row_ids, + const std::vector& expected_row_ids) { + ASSERT_OK_AND_ASSIGN(ScoredRows all_rows, + SearchFullText(reader, query, /*limit=*/100, include_row_ids)); + std::vector row_ids; + for (const auto& [row_id, score] : all_rows) { + row_ids.push_back(row_id); + } + ASSERT_EQ(row_ids, expected_row_ids); + ScoredRows ranked_rows = all_rows; + std::sort( + ranked_rows.begin(), ranked_rows.end(), + [](const std::pair& left, const std::pair& right) { + return left.second != right.second ? left.second > right.second + : left.first < right.first; + }); + for (size_t limit = 1; limit < ranked_rows.size(); ++limit) { + ScoredRows expected_rows(ranked_rows.begin(), ranked_rows.begin() + limit); + std::sort(expected_rows.begin(), expected_rows.end()); + ASSERT_OK_AND_ASSIGN( + ScoredRows rows, + SearchFullText(reader, query, static_cast(limit), include_row_ids)); + ASSERT_EQ(rows, expected_rows) << "limit=" << limit; + } + } + + void CheckSearchOverlappingIndexRanges(const std::string& index_type, + const std::map& index_options, + const std::string& query) { + arrow::FieldVector fields = {arrow::field("f0", arrow::utf8()), + arrow::field("f1", arrow::int32())}; + auto schema = arrow::schema(fields); + std::map options = {{Options::FILE_FORMAT, file_format_}, + {Options::FILE_SYSTEM, "local"}, + {Options::ROW_TRACKING_ENABLED, "true"}, + {Options::DATA_EVOLUTION_ENABLED, "true"}}; + CreateTable(/*partition_keys=*/{}, schema, options); + std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar"); + + auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([ +["This is an test document.", 0], +["This is an new document document document.", 1], +["Document document document document test.", 2], +["Document document document document document.", 3] + ])") + .ValueOrDie(); + ASSERT_OK_AND_ASSIGN(auto commit_msgs, + WriteArray(table_path, schema->field_names(), src_array)); + ASSERT_OK(Commit(table_path, commit_msgs)); + ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", index_type, index_options, + Range(0, 2))); + ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", index_type, index_options, + Range(1, 1))); + ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", index_type, index_options, + Range(1, 3))); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr global_index_scan, + GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt, + /*partitions=*/std::nullopt, /*options=*/{}, + fs_, /*executor=*/nullptr, pool_)); + ASSERT_OK_AND_ASSIGN( + auto readers, global_index_scan->CreateReaders("f0", /*row_range_index=*/std::nullopt)); + ASSERT_EQ(readers.size(), 1u); + CheckFullTextTopK(readers[0], query, std::nullopt, {0, 1, 2, 3}); + ASSERT_OK_AND_ASSIGN(ScoredRows top_rows, + SearchFullText(readers[0], query, /*limit=*/1, std::nullopt)); + ASSERT_EQ(top_rows.size(), 1u); + ASSERT_EQ(top_rows[0].first, 3); + CheckFullTextTopK(readers[0], query, RoaringBitmap64::From({1, 2, 3}), {1, 2, 3}); + + ASSERT_OK_AND_ASSIGN(RowRangeIndex row_range_index, RowRangeIndex::Create({Range(3, 3)})); + ASSERT_OK_AND_ASSIGN(readers, global_index_scan->CreateReaders("f0", row_range_index)); + ASSERT_EQ(readers.size(), 1u); + CheckFullTextTopK(readers[0], query, std::nullopt, {3}); + } + Result> ScanGlobalIndexAndData( const std::string& table_path, const std::shared_ptr& predicate, const std::map& options = {}, @@ -2803,6 +2901,28 @@ TEST_P(GlobalIndexTest, TestWriteCommitScanReadLuceneIndexWithPartition) { ASSERT_OK_AND_ASSIGN(RowRangeIndex row_range_index, RowRangeIndex::Create({Range(2, 3)})); scan_and_check_result({{"f1", "20"}}, row_range_index, filter, "row ids: {2}"); } + { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr global_index_scan, + GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt, + /*partitions=*/std::nullopt, lucene_options, fs_, + /*executor=*/nullptr, pool_)); + ASSERT_OK_AND_ASSIGN( + auto readers, global_index_scan->CreateReaders("f0", /*row_range_index=*/std::nullopt)); + ASSERT_EQ(readers.size(), 1u); + std::string query = R"({"match":{"query":"document","operator":"And"}})"; + CheckFullTextTopK(readers[0], query, std::nullopt, {0, 1, 2}); + CheckFullTextTopK(readers[0], query, RoaringBitmap64::From({1, 2, 3}), {1, 2}); + } +} + +TEST_P(GlobalIndexTest, TestLuceneSearchOverlappingIndexRanges) { + auto tmp_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(tmp_dir); + CheckSearchOverlappingIndexRanges("lucene-fts", + {{"lucene-fts.write.omit-term-freq-and-position", "false"}, + {"lucene-fts.write.tmp.directory", tmp_dir->Str()}}, + R"({"match":{"query":"document","operator":"And"}})"); } #endif @@ -2932,6 +3052,24 @@ TEST_P(GlobalIndexTest, TestWriteCommitScanReadFullTextIndexWithPartition) { ASSERT_OK_AND_ASSIGN(RowRangeIndex row_range_index, RowRangeIndex::Create({Range(2, 3)})); scan_and_check_result({{"f1", "20"}}, row_range_index, std::nullopt, "{2}"); } + { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr global_index_scan, + GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt, + /*partitions=*/std::nullopt, /*options=*/{}, fs_, + /*executor=*/nullptr, pool_)); + ASSERT_OK_AND_ASSIGN( + auto readers, global_index_scan->CreateReaders("f0", /*row_range_index=*/std::nullopt)); + ASSERT_EQ(readers.size(), 1u); + std::string query = R"({"match":{"query":"document"}})"; + CheckFullTextTopK(readers[0], query, std::nullopt, {0, 1, 2}); + CheckFullTextTopK(readers[0], query, RoaringBitmap64::From({1, 2, 3}), {1, 2}); + } +} + +TEST_P(GlobalIndexTest, TestFullTextSearchOverlappingIndexRanges) { + CheckSearchOverlappingIndexRanges("full-text", {{"full-text.tokenizer", "default"}}, + R"({"match":{"query":"document"}})"); } #endif