diff --git a/docs/source/user_guide/global_index.rst b/docs/source/user_guide/global_index.rst index 956d28dba..f553c67c2 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, as Java +Paimon does. 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..c2ad2a7df 100644 --- a/include/paimon/predicate/full_text_search.h +++ b/include/paimon/predicate/full_text_search.h @@ -83,7 +83,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..817701e10 100644 --- a/src/paimon/common/global_index/offset_global_index_reader.cpp +++ b/src/paimon/common/global_index/offset_global_index_reader.cpp @@ -19,13 +19,28 @@ #include "paimon/common/global_index/offset_global_index_reader.h" +#include #include +#include + +#include "paimon/global_index/bitmap_scored_global_index_result.h" +#include "paimon/utils/roaring_bitmap64.h" namespace paimon { +namespace { +std::shared_ptr EmptyScoredResult() { + return std::make_shared(RoaringBitmap64(), std::vector()); +} +} // namespace + +OffsetGlobalIndexReader::OffsetGlobalIndexReader(std::shared_ptr&& wrapped, + int64_t offset, int64_t to) + : OffsetGlobalIndexReader(std::move(wrapped), offset, to, + /*skip_to=*/std::numeric_limits::min()) {} OffsetGlobalIndexReader::OffsetGlobalIndexReader(std::shared_ptr&& wrapped, - int64_t offset) - : wrapped_(std::move(wrapped)), offset_(offset) {} + int64_t offset, int64_t to, int64_t skip_to) + : wrapped_(std::move(wrapped)), offset_(offset), to_(to), skip_to_(skip_to) {} Result> OffsetGlobalIndexReader::VisitIsNotNull() { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr result, wrapped_->VisitIsNotNull()); @@ -137,11 +152,34 @@ 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. 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); + if (full_text_search) { + if (full_text_search->limit <= 0) { + return Status::Invalid("full-text search requires a positive limit"); + } + // Rows up to `skip_to_` are ranked by a reader over an overlapping row range. + if (skip_to_ >= to_) { + return EmptyScoredResult(); + } + bool skips_rows = skip_to_ >= offset_; + if (full_text_search->include_row_ids) { + local_search = full_text_search->OffsetRange(offset_, to_); + if (skips_rows) { + local_search->include_row_ids->RemoveRange(0, skip_to_ - offset_ + 1); + } + if (local_search->include_row_ids->IsEmpty()) { + return EmptyScoredResult(); + } + } else if (skips_rows) { + // `local_to` may be the largest int64 value, so add it apart from the half-open range. + int64_t local_to = to_ - offset_; + RoaringBitmap64 local_row_ids; + local_row_ids.AddRange(skip_to_ - offset_ + 1, local_to); + local_row_ids.Add(local_to); + local_search = std::make_shared( + full_text_search->field_name, full_text_search->query, full_text_search->limit, + std::move(local_row_ids)); + } } 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..b3f8747fe 100644 --- a/src/paimon/common/global_index/offset_global_index_reader.h +++ b/src/paimon/common/global_index/offset_global_index_reader.h @@ -31,11 +31,22 @@ namespace paimon { /// results. This is used to convert local row IDs into global row IDs. class OffsetGlobalIndexReader : public GlobalIndexReader { public: - /// Constructs an OffsetGlobalIndexReader. + /// Constructs an OffsetGlobalIndexReader whose full-text searches rank all rows it covers. /// @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); + + /// Constructs an OffsetGlobalIndexReader whose full-text searches rank rows in [`offset`, + /// `to`], excluding rows up to `skip_to` (inclusive) assigned to earlier readers. + /// No row is ranked if `skip_to` >= `to`. Storing the inclusive end avoids computing + /// INT64_MAX + 1 when an earlier range covers the largest row id. + /// @param skip_to The last global row id already assigned to an earlier reader. + OffsetGlobalIndexReader(std::shared_ptr&& wrapped, int64_t offset, + int64_t to, int64_t skip_to); Result> VisitIsNotNull() override; Result> VisitIsNull() override; @@ -78,6 +89,8 @@ class OffsetGlobalIndexReader : public GlobalIndexReader { private: std::shared_ptr wrapped_; int64_t offset_; + int64_t to_; + int64_t skip_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..e8c154db3 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,32 +228,124 @@ 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); ASSERT_OK(offset_reader->VisitFullTextSearch(search)); - ASSERT_TRUE(fake_reader->CapturedFullTextSearch()->include_row_ids); - ASSERT_TRUE(fake_reader->CapturedFullTextSearch()->include_row_ids->IsEmpty()); + ASSERT_EQ(fake_reader->CapturedFullTextSearch(), search); } + { + auto search = std::make_shared( + "f0", R"({"match":{"query":"paimon"}})", /*limit=*/7, RoaringBitmap64::From({5l, 20l})); + ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitFullTextSearch(search)); + CheckScoredResult(result, {}, {}); + ASSERT_EQ(fake_reader->InvocationCount(), 2); + } +} + +TEST_F(OffsetGlobalIndexReaderTest, TestVisitFullTextSearchRejectsNonPositiveLimit) { + auto fake_reader = std::make_shared(); + auto offset_reader = std::make_shared(fake_reader, 10, /*to=*/19); + auto covered_reader = std::make_shared(fake_reader, 10, /*to=*/19, + /*skip_to=*/19); + for (int32_t limit : {0, -1}) { + auto search = + std::make_shared("f0", R"({"match":{"query":"paimon"}})", limit); + ASSERT_NOK_WITH_MSG(offset_reader->VisitFullTextSearch(search), + "full-text search requires a positive limit"); + ASSERT_NOK_WITH_MSG(covered_reader->VisitFullTextSearch(search), + "full-text search requires a positive limit"); + search->include_row_ids = RoaringBitmap64::From({5, 20}); + ASSERT_NOK_WITH_MSG(offset_reader->VisitFullTextSearch(search), + "full-text search requires a positive limit"); + } + ASSERT_EQ(fake_reader->InvocationCount(), 0); +} + +TEST_F(OffsetGlobalIndexReaderTest, TestVisitFullTextSearchSkipsRows) { + auto fake_reader = std::make_shared(); + fake_reader->SetScoredResult({3, 9}, {0.5f, 0.9f}); + auto offset_reader = std::make_shared(fake_reader, 10, /*to=*/19, + /*skip_to=*/12); { auto search = std::make_shared("f0", R"({"match":{"query":"paimon"}})", /*limit=*/7); + ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitFullTextSearch(search)); + CheckScoredResult(result, {13, 19}, {0.5f, 0.9f}); + ASSERT_EQ(fake_reader->CapturedFullTextSearch()->include_row_ids, + RoaringBitmap64::From({3l, 4l, 5l, 6l, 7l, 8l, 9l})); + } + { + auto search = + std::make_shared("f0", R"({"match":{"query":"paimon"}})", /*limit=*/7, + RoaringBitmap64::From({5l, 12l, 13l, 19l, 20l})); ASSERT_OK(offset_reader->VisitFullTextSearch(search)); - ASSERT_EQ(fake_reader->CapturedFullTextSearch(), search); + ASSERT_EQ(fake_reader->CapturedFullTextSearch()->include_row_ids, + RoaringBitmap64::From({3l, 9l})); + } + { + auto search = std::make_shared( + "f0", R"({"match":{"query":"paimon"}})", /*limit=*/7, RoaringBitmap64::From({5l, 12l})); + ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitFullTextSearch(search)); + CheckScoredResult(result, {}, {}); + ASSERT_EQ(fake_reader->InvocationCount(), 2); + } +} + +TEST_F(OffsetGlobalIndexReaderTest, TestVisitFullTextSearchEndingAtLargestRowId) { + constexpr int64_t kMax = std::numeric_limits::max(); + { + auto fake_reader = std::make_shared(); + fake_reader->SetScoredResult({7, 9}, {0.5f, 0.9f}); + auto offset_reader = std::make_shared( + fake_reader, kMax - 9, /*to=*/kMax, /*skip_to=*/kMax - 3); + auto search = + std::make_shared("f0", R"({"match":{"query":"paimon"}})", /*limit=*/7); + ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitFullTextSearch(search)); + CheckScoredResult(result, {kMax - 2, kMax}, {0.5f, 0.9f}); + ASSERT_EQ(fake_reader->CapturedFullTextSearch()->include_row_ids, + RoaringBitmap64::From({7l, 8l, 9l})); + + search = + std::make_shared("f0", R"({"match":{"query":"paimon"}})", /*limit=*/7, + RoaringBitmap64::From({kMax - 5, kMax - 2, kMax})); + ASSERT_OK(offset_reader->VisitFullTextSearch(search)); + ASSERT_EQ(fake_reader->CapturedFullTextSearch()->include_row_ids, + RoaringBitmap64::From({7l, 9l})); + } + { + auto fake_reader = std::make_shared(); + auto offset_reader = std::make_shared( + fake_reader, /*offset=*/0, /*to=*/kMax, /*skip_to=*/kMax - 1); + auto search = + std::make_shared("f0", R"({"match":{"query":"paimon"}})", /*limit=*/7); + ASSERT_OK(offset_reader->VisitFullTextSearch(search)); + ASSERT_EQ(fake_reader->CapturedFullTextSearch()->include_row_ids, + RoaringBitmap64::From({kMax})); + } + { + auto fake_reader = std::make_shared(); + auto offset_reader = std::make_shared( + fake_reader, kMax - 9, /*to=*/kMax, /*skip_to=*/kMax); + auto search = + std::make_shared("f0", R"({"match":{"query":"paimon"}})", /*limit=*/7); + ASSERT_OK_AND_ASSIGN(auto result, offset_reader->VisitFullTextSearch(search)); + CheckScoredResult(result, {}, {}); + ASSERT_EQ(fake_reader->InvocationCount(), 0); } } @@ -260,7 +353,7 @@ 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 +362,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 +372,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 +384,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 +396,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 +408,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..b5565adc4 100644 --- a/src/paimon/common/global_index/union_global_index_reader.cpp +++ b/src/paimon/common/global_index/union_global_index_reader.cpp @@ -19,12 +19,62 @@ #include "paimon/common/global_index/union_global_index_reader.h" +#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/utils/roaring_bitmap64.h" namespace paimon { +namespace { +/// Returns whether `left` ranks before `right` as in Java `ScoredGlobalIndexResult#topK`: a higher +/// score first, then a lower row id. Scores are compared as Java `Double#compare` does. +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, as Java +/// `ScoredGlobalIndexResult#merge` followed by `#topK` does. 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) @@ -133,9 +183,33 @@ 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. + ScoredReaderAction action = + [full_text_search](const std::shared_ptr& reader) { + return reader->VisitFullTextSearch(full_text_search); + }; + auto results = ExecuteAllReaders>>(action); + 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)); + } + } + if (scored_results.empty()) { + return std::shared_ptr(); + } + return MergeTopK(scored_results, full_text_search->limit); } Result> UnionGlobalIndexReader::UnionScored( 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..8be8b9637 100644 --- a/src/paimon/common/global_index/union_global_index_reader.h +++ b/src/paimon/common/global_index/union_global_index_reader.h @@ -58,6 +58,9 @@ 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. Result> VisitFullTextSearch( const std::shared_ptr& full_text_search) override; 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..6eab13ef7 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,6 +19,7 @@ #include "paimon/common/global_index/union_global_index_reader.h" +#include #include #include #include @@ -31,6 +32,7 @@ #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 +147,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) { @@ -435,10 +441,90 @@ TEST_F(UnionGlobalIndexReaderTest, TestVisitFullTextSearchUnion) { std::vector> readers = {reader1, reader2, reader3}; UnionGlobalIndexReader union_reader(std::move(readers), 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), 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), 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, TestVisitFullTextSearchOrdersSpecialScoresAsJava) { + 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), 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), 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), 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_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, TestVisitVectorSearchAllNullptr) { auto reader1 = std::make_shared(); auto reader2 = std::make_shared(); 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..580ee1176 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,8 @@ */ #include "paimon/core/global_index/global_index_scan_impl.h" +#include +#include #include #include #include @@ -185,7 +187,13 @@ Result>> GlobalIndexScanImpl::Cre } std::vector> union_readers; union_readers.reserve(range_to_metas.size()); + // Assign each full-text row to one reader before requesting per-reader top-k, following + // Java `DataEvolutionFullTextScan#chooseIndexRanges`. Here ranges of the same index type + // are processed in (from, to) order; each skips rows covered by earlier ranges. + int64_t covered_to = std::numeric_limits::min(); for (const auto& [range, metas] : range_to_metas) { + int64_t skip_to = covered_to; + covered_to = std::max(covered_to, range.to); if (row_range_index && !row_range_index->Intersects(range.from, range.to)) { continue; } @@ -199,8 +207,8 @@ 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, skip_to)); } if (union_readers.empty()) { continue; diff --git a/test/inte/global_index_test.cpp b/test/inte/global_index_test.cpp index e6c1abc59..eaf00185f 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,52 @@ 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; + } + } + Result> ScanGlobalIndexAndData( const std::string& table_path, const std::shared_ptr& predicate, const std::map& options = {}, @@ -2803,6 +2851,75 @@ 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) { + arrow::FieldVector fields = {arrow::field("f0", arrow::utf8()), + arrow::field("f1", arrow::int32())}; + auto tmp_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(tmp_dir); + std::map lucene_options = { + {"lucene-fts.write.omit-term-freq-and-position", "false"}, + {"lucene-fts.write.tmp.directory", tmp_dir->Str()}}; + 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"); + std::vector write_cols = schema->field_names(); + + 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, write_cols, src_array)); + ASSERT_OK(Commit(table_path, commit_msgs)); + ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "lucene-fts", + /*options=*/lucene_options, Range(0, 2))); + ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "lucene-fts", + /*options=*/lucene_options, Range(1, 1))); + ASSERT_OK(WriteIndex(table_path, /*partition_filters=*/{}, "f0", "lucene-fts", + /*options=*/lucene_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); + std::string query = R"({"match":{"query":"document","operator":"And"}})"; + 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}); } #endif @@ -2932,6 +3049,19 @@ 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}); + } } #endif