From d9266f09f66e1bfe839650c3955861a4a21ce141 Mon Sep 17 00:00:00 2001 From: "jinli.zjw" Date: Wed, 16 Sep 2026 21:29:57 +0800 Subject: [PATCH 1/3] feat(file-index): add file-local search with Lumina vector indexing --- CMakeLists.txt | 2 +- include/paimon/file_index/file_index_format.h | 4 +- include/paimon/file_index/file_index_reader.h | 13 + include/paimon/file_index/file_indexer.h | 12 +- .../file_index/scored_file_index_result.h | 61 + include/paimon/predicate/vector_search.h | 14 +- include/paimon/read_context.h | 39 +- src/paimon/CMakeLists.txt | 7 +- .../common/file_index/file_index_format.cpp | 38 +- .../file_index/file_index_format_test.cpp | 11 +- .../common/file_index/file_index_reader.cpp | 11 + .../file_index/file_index_reader_test.cpp | 12 + .../file_index/scored_file_index_result.cpp | 57 + .../scored_file_index_result_test.cpp | 55 + .../common/io/byte_array_output_stream.cpp | 17 +- .../common/io/byte_array_output_stream.h | 10 +- .../io/byte_array_output_stream_test.cpp | 12 +- src/paimon/common/io/offset_input_stream.cpp | 3 +- src/paimon/common/io/offset_input_stream.h | 4 +- .../complete_index_score_batch_reader.cpp | 22 +- .../complete_index_score_batch_reader.h | 3 + ...complete_index_score_batch_reader_test.cpp | 51 +- ...complete_index_score_file_batch_reader.cpp | 73 ++ .../complete_index_score_file_batch_reader.h | 55 + .../core/append/append_only_writer_test.cpp | 3 +- src/paimon/core/io/data_file_index_writer.cpp | 10 +- .../core/io/data_file_index_writer_test.cpp | 4 +- src/paimon/core/io/file_index_evaluator.cpp | 119 +- src/paimon/core/io/file_index_evaluator.h | 24 +- .../core/io/file_index_evaluator_test.cpp | 78 +- src/paimon/core/io/file_index_options.cpp | 25 +- src/paimon/core/io/file_index_options.h | 7 +- .../core/io/file_index_options_test.cpp | 13 + .../core/operation/abstract_split_read.cpp | 19 +- .../core/operation/abstract_split_read.h | 3 +- .../operation/append_only_file_store_scan.cpp | 6 +- .../operation/data_evolution_split_read.cpp | 17 +- .../operation/data_evolution_split_read.h | 3 +- .../core/operation/internal_read_context.cpp | 42 +- .../core/operation/internal_read_context.h | 15 +- .../operation/internal_read_context_test.cpp | 45 + .../core/operation/merge_file_split_read.cpp | 3 +- .../core/operation/merge_file_split_read.h | 3 +- .../core/operation/raw_file_split_read.cpp | 64 +- .../core/operation/raw_file_split_read.h | 3 +- .../operation/raw_file_split_read_test.cpp | 142 +++ src/paimon/core/operation/read_context.cpp | 43 +- .../core/operation/read_context_test.cpp | 38 + src/paimon/core/table/source/table_read.cpp | 9 + .../lumina/lumina_global_index.cpp | 1024 ----------------- .../lumina/CMakeLists.txt | 12 +- .../lumina/lumina_api_test.cpp | 6 +- src/paimon/indexer/lumina/lumina_dataset.cpp | 105 ++ src/paimon/indexer/lumina/lumina_dataset.h | 80 ++ .../indexer/lumina/lumina_file_index.cpp | 218 ++++ src/paimon/indexer/lumina/lumina_file_index.h | 107 ++ .../lumina/lumina_file_index_factory.cpp | 46 + .../indexer/lumina/lumina_file_index_test.cpp | 198 ++++ .../lumina/lumina_file_io_test.cpp | 4 +- .../lumina/lumina_file_reader.h | 2 +- .../lumina/lumina_file_writer.h | 2 +- .../indexer/lumina/lumina_global_index.cpp | 241 ++++ .../lumina/lumina_global_index.h | 54 +- .../lumina/lumina_global_index_factory.cpp | 4 +- .../lumina/lumina_global_index_factory.h | 0 .../lumina/lumina_global_index_test.cpp | 2 +- .../lumina/lumina_index_accumulator.cpp | 114 ++ .../indexer/lumina/lumina_index_accumulator.h | 58 + .../indexer/lumina/lumina_index_options.cpp | 114 ++ .../indexer/lumina/lumina_index_options.h | 63 + .../lumina/lumina_memory_pool.h | 0 .../indexer/lumina/lumina_search_utils.cpp | 117 ++ .../indexer/lumina/lumina_search_utils.h | 72 ++ .../indexer/lumina/lumina_tag_utils.cpp | 517 +++++++++ src/paimon/indexer/lumina/lumina_tag_utils.h | 80 ++ .../lumina/lumina_utils.h | 0 76 files changed, 3323 insertions(+), 1241 deletions(-) create mode 100644 include/paimon/file_index/scored_file_index_result.h create mode 100644 src/paimon/common/file_index/scored_file_index_result.cpp create mode 100644 src/paimon/common/file_index/scored_file_index_result_test.cpp rename src/paimon/common/{global_index => reader}/complete_index_score_batch_reader.cpp (84%) rename src/paimon/common/{global_index => reader}/complete_index_score_batch_reader.h (95%) rename src/paimon/common/{global_index => reader}/complete_index_score_batch_reader_test.cpp (70%) create mode 100644 src/paimon/common/reader/complete_index_score_file_batch_reader.cpp create mode 100644 src/paimon/common/reader/complete_index_score_file_batch_reader.h delete mode 100644 src/paimon/global_index/lumina/lumina_global_index.cpp rename src/paimon/{global_index => indexer}/lumina/CMakeLists.txt (84%) rename src/paimon/{global_index => indexer}/lumina/lumina_api_test.cpp (98%) create mode 100644 src/paimon/indexer/lumina/lumina_dataset.cpp create mode 100644 src/paimon/indexer/lumina/lumina_dataset.h create mode 100644 src/paimon/indexer/lumina/lumina_file_index.cpp create mode 100644 src/paimon/indexer/lumina/lumina_file_index.h create mode 100644 src/paimon/indexer/lumina/lumina_file_index_factory.cpp create mode 100644 src/paimon/indexer/lumina/lumina_file_index_test.cpp rename src/paimon/{global_index => indexer}/lumina/lumina_file_io_test.cpp (97%) rename src/paimon/{global_index => indexer}/lumina/lumina_file_reader.h (98%) rename src/paimon/{global_index => indexer}/lumina/lumina_file_writer.h (98%) create mode 100644 src/paimon/indexer/lumina/lumina_global_index.cpp rename src/paimon/{global_index => indexer}/lumina/lumina_global_index.h (83%) rename src/paimon/{global_index => indexer}/lumina/lumina_global_index_factory.cpp (90%) rename src/paimon/{global_index => indexer}/lumina/lumina_global_index_factory.h (100%) rename src/paimon/{global_index => indexer}/lumina/lumina_global_index_test.cpp (99%) create mode 100644 src/paimon/indexer/lumina/lumina_index_accumulator.cpp create mode 100644 src/paimon/indexer/lumina/lumina_index_accumulator.h create mode 100644 src/paimon/indexer/lumina/lumina_index_options.cpp create mode 100644 src/paimon/indexer/lumina/lumina_index_options.h rename src/paimon/{global_index => indexer}/lumina/lumina_memory_pool.h (100%) create mode 100644 src/paimon/indexer/lumina/lumina_search_utils.cpp create mode 100644 src/paimon/indexer/lumina/lumina_search_utils.h create mode 100644 src/paimon/indexer/lumina/lumina_tag_utils.cpp create mode 100644 src/paimon/indexer/lumina/lumina_tag_utils.h rename src/paimon/{global_index => indexer}/lumina/lumina_utils.h (100%) diff --git a/CMakeLists.txt b/CMakeLists.txt index f7f6adb4..632fc858 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -555,7 +555,7 @@ if(PAIMON_ENABLE_MOSAIC) add_subdirectory(src/paimon/format/mosaic) endif() if(PAIMON_ENABLE_LUMINA) - add_subdirectory(src/paimon/global_index/lumina) + add_subdirectory(src/paimon/indexer/lumina) endif() add_subdirectory(src/paimon/global_index/lucene) if(PAIMON_ENABLE_TANTIVY) diff --git a/include/paimon/file_index/file_index_format.h b/include/paimon/file_index/file_index_format.h index 3993b624..b3116f13 100644 --- a/include/paimon/file_index/file_index_format.h +++ b/include/paimon/file_index/file_index_format.h @@ -102,10 +102,12 @@ class PAIMON_EXPORT FileIndexFormat { /// stream. /// @param input_stream Input stream containing serialized index data. /// @param pool Memory pool for temporary allocations during reading. + /// @param options Table options used to configure File Index backends. /// @return A unique pointer to a `Reader` on success, or an error if the stream is invalid /// (e.g., wrong magic, unsupported version, or corrupted data). static Result> CreateReader( - const std::shared_ptr& input_stream, const std::shared_ptr& pool); + const std::shared_ptr& input_stream, const std::shared_ptr& pool, + const std::map& options); /// Creates a `Writer` which serializes a complete V1 file index container. /// diff --git a/include/paimon/file_index/file_index_reader.h b/include/paimon/file_index/file_index_reader.h index 2da36fd2..720e3fc8 100644 --- a/include/paimon/file_index/file_index_reader.h +++ b/include/paimon/file_index/file_index_reader.h @@ -23,7 +23,10 @@ #include #include "paimon/file_index/file_index_result.h" +#include "paimon/file_index/scored_file_index_result.h" +#include "paimon/predicate/full_text_search.h" #include "paimon/predicate/function_visitor.h" +#include "paimon/predicate/vector_search.h" #include "paimon/result.h" #include "paimon/visibility.h" @@ -38,6 +41,8 @@ namespace paimon { /// of the indexed data; otherwise, index filtering results may be incorrect. class PAIMON_EXPORT FileIndexReader : public FunctionVisitor> { public: + virtual ~FileIndexReader() = default; + Result> VisitIsNotNull() override; Result> VisitIsNull() override; @@ -66,6 +71,14 @@ class PAIMON_EXPORT FileIndexReader : public FunctionVisitor> VisitContains(const Literal& literal) override; Result> VisitLike(const Literal& literal) override; + + /// Execute vector search and return scored file-local physical row positions. + virtual Result> VisitVectorSearch( + const std::shared_ptr& vector_search); + + /// Execute full-text search and return matching file-local physical row positions. + virtual Result> VisitFullTextSearch( + const std::shared_ptr& full_text_search); }; } // namespace paimon diff --git a/include/paimon/file_index/file_indexer.h b/include/paimon/file_index/file_indexer.h index 72c1fc9c..47af01c0 100644 --- a/include/paimon/file_index/file_indexer.h +++ b/include/paimon/file_index/file_indexer.h @@ -19,6 +19,9 @@ #pragma once #include +#include +#include +#include #include "paimon/file_index/file_index_reader.h" #include "paimon/file_index/file_index_result.h" @@ -35,6 +38,11 @@ class PAIMON_EXPORT FileIndexer { public: virtual ~FileIndexer() = default; + /// Extra columns required while building this index, in addition to the indexed column. + virtual Result>> GetExtraFieldNames() const { + return std::optional>(std::nullopt); + } + /// Create `FileIndexReader` with input stream. /// /// @param arrow_schema ArrowSchema derived from arrow schema or struct type with @@ -51,8 +59,8 @@ class PAIMON_EXPORT FileIndexer { /// Create `FileIndexWriter` for arrow schema. /// - /// @param arrow_schema ArrowSchema derived from arrow schema or struct type with - /// specified indexed field. + /// @param arrow_schema ArrowSchema containing the indexed field first, followed by any + /// fields requested through `GetExtraFieldNames()`. /// @param pool Memory pool for memory allocation. /// @return A `FileIndexWriter` to write index. virtual Result> CreateWriter( diff --git a/include/paimon/file_index/scored_file_index_result.h b/include/paimon/file_index/scored_file_index_result.h new file mode 100644 index 00000000..a81693f6 --- /dev/null +++ b/include/paimon/file_index/scored_file_index_result.h @@ -0,0 +1,61 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +#include "paimon/file_index/file_index_result.h" +#include "paimon/result.h" +#include "paimon/utils/roaring_bitmap32.h" +#include "paimon/visibility.h" + +namespace paimon { + +/// File-local vector search result. Scores correspond to row positions in ascending order. +class PAIMON_EXPORT ScoredFileIndexResult : public FileIndexResult { + public: + static Result> Create(RoaringBitmap32&& row_positions, + std::vector&& scores); + + bool IsEmpty() const { + return row_positions_.IsEmpty(); + } + + Result IsRemain() const override; + + const RoaringBitmap32& GetRowPositions() const { + return row_positions_; + } + + const std::vector& GetScores() const { + return scores_; + } + + std::string ToString() const override; + + private: + ScoredFileIndexResult(RoaringBitmap32&& row_positions, std::vector&& scores); + + RoaringBitmap32 row_positions_; + std::vector scores_; +}; + +} // namespace paimon diff --git a/include/paimon/predicate/vector_search.h b/include/paimon/predicate/vector_search.h index 44852162..f4155632 100644 --- a/include/paimon/predicate/vector_search.h +++ b/include/paimon/predicate/vector_search.h @@ -30,11 +30,12 @@ namespace paimon { /// `VectorSearch` to perform vector similarity search. struct PAIMON_EXPORT VectorSearch { - /// `PreFilter`: A lightweight pre-filtering function applied **before** similarity - /// scoring. It operates solely on **global row ids** and is typically driven by other global - /// index, such as bitmap, or range index. This filter enables early pruning of irrelevant - /// candidates (e.g., "only consider rows with label X"), significantly reducing the search - /// space. Returns true to include the row in vector search process; false to exclude it. + /// `PreFilter`: A lightweight pre-filtering function applied **before** similarity scoring. + /// The row id domain is determined by the reader executing the search: `GlobalIndexReader` + /// supplies global row ids, while `FileIndexReader` supplies physical row ids local to the + /// current data file. This filter enables early pruning of irrelevant candidates (e.g., "only + /// consider rows with label X"), significantly reducing the search space. Returns true to + /// include the row in vector search process; false to exclude it. /// /// @note Must be thread-safe. using PreFilter = std::function; @@ -65,7 +66,8 @@ struct PAIMON_EXPORT VectorSearch { int32_t limit; /// The query vector (must match the dimensionality of the indexed vectors). std::vector query; - /// A pre-filter based on **global row ids**, implemented by leveraging other global index + /// A pre-filter whose row id domain is determined by the reader: global row ids for Global + /// Index and file-local physical row ids for File Index. std::function pre_filter; /// A runtime filtering condition that may involve graph traversal of /// structured attributes. **Using this parameter often yields better diff --git a/include/paimon/read_context.h b/include/paimon/read_context.h index 25cf7ec5..789f61d6 100644 --- a/include/paimon/read_context.h +++ b/include/paimon/read_context.h @@ -35,11 +35,13 @@ namespace paimon { class Executor; +struct FullTextSearch; class FormatTable; class MemoryPool; class Predicate; class FileSystem; class RealtimeContext; +struct VectorSearch; /// `ReadContext` is some configuration for read operations. /// @@ -51,10 +53,13 @@ class PAIMON_EXPORT ReadContext { ReadContext(const std::string& path, const std::string& branch, const std::vector& read_field_names, const std::vector& read_field_ids, - const std::shared_ptr& predicate, bool enable_predicate_filter, - bool enable_prefetch, bool enable_late_materializing, uint32_t prefetch_batch_count, - uint32_t prefetch_max_parallel_num, bool enable_multi_thread_row_to_batch, - uint32_t row_to_batch_thread_number, const std::optional& table_schema, + const std::shared_ptr& predicate, + const std::shared_ptr& vector_search, + const std::shared_ptr& full_text_search, + bool enable_predicate_filter, bool enable_prefetch, bool enable_late_materializing, + uint32_t prefetch_batch_count, uint32_t prefetch_max_parallel_num, + bool enable_multi_thread_row_to_batch, uint32_t row_to_batch_thread_number, + const std::optional& table_schema, const std::shared_ptr& memory_pool, const std::shared_ptr& executor, const std::shared_ptr& specific_file_system, @@ -94,6 +99,18 @@ class PAIMON_EXPORT ReadContext { return predicate_; } + const std::shared_ptr& GetVectorSearch() const { + return vector_search_; + } + + const std::shared_ptr& GetFullTextSearch() const { + return full_text_search_; + } + + bool HasFileIndexSearch() const { + return vector_search_ != nullptr || full_text_search_ != nullptr; + } + bool EnablePredicateFilter() const { return enable_predicate_filter_; } @@ -177,6 +194,8 @@ class PAIMON_EXPORT ReadContext { std::vector read_field_names_; std::vector read_field_ids_; std::shared_ptr predicate_; + std::shared_ptr vector_search_; + std::shared_ptr full_text_search_; bool enable_predicate_filter_; bool enable_prefetch_; bool enable_late_materializing_; @@ -318,6 +337,18 @@ class PAIMON_EXPORT ReadContextBuilder { /// @return Reference to this builder for method chaining. ReadContextBuilder& SetPredicate(const std::shared_ptr& predicate); + /// Configure file-local vector search. The limit is applied independently to every data file. + /// Vector search and full-text search are mutually exclusive. + /// @param vector_search Vector search request, or nullptr to clear it. + /// @return Reference to this builder for method chaining. + ReadContextBuilder& SetVectorSearch(const std::shared_ptr& vector_search); + + /// Configure file-local full-text search. The limit is applied independently to every data + /// file. Full-text search and vector search are mutually exclusive. + /// @param full_text_search Full-text search request, or nullptr to clear it. + /// @return Reference to this builder for method chaining. + ReadContextBuilder& SetFullTextSearch(const std::shared_ptr& full_text_search); + /// Whether to perform precise filtering according to predicates for data read from format /// reader. /// @param enabled Whether to enable precise filtering (default: false) diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index a25397b9..007ef768 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -71,6 +71,7 @@ set(PAIMON_COMMON_SRCS common/file_index/file_index_format.cpp common/file_index/file_index_reader.cpp common/file_index/file_index_result.cpp + common/file_index/scored_file_index_result.cpp common/format/column_stats.cpp common/format/file_format_factory.cpp common/fs/file_system.cpp @@ -78,7 +79,6 @@ set(PAIMON_COMMON_SRCS common/fs/file_system_factory.cpp common/global_index/union_global_index_reader.cpp common/global_index/offset_global_index_reader.cpp - common/global_index/complete_index_score_batch_reader.cpp common/global_index/bitmap_scored_global_index_result.cpp common/global_index/bitmap_global_index_result.cpp common/global_index/global_index_result.cpp @@ -135,6 +135,8 @@ set(PAIMON_COMMON_SRCS common/predicate/predicate_utils.cpp common/predicate/starts_with.cpp common/reader/batch_reader.cpp + common/reader/complete_index_score_batch_reader.cpp + common/reader/complete_index_score_file_batch_reader.cpp common/reader/concat_batch_reader.cpp common/reader/late_materializing_file_batch_reader.cpp common/reader/data_file_reader_factory.cpp @@ -582,6 +584,7 @@ if(PAIMON_BUILD_TESTS) common/file_index/file_indexer_factory_test.cpp common/file_index/file_index_result_test.cpp common/file_index/file_index_reader_test.cpp + common/file_index/scored_file_index_result_test.cpp common/file_index/file_index_format_test.cpp common/file_index/empty/empty_file_index_reader_test.cpp common/file_index/bitmap/bitmap_index_result_test.cpp @@ -597,7 +600,6 @@ if(PAIMON_BUILD_TESTS) common/file_index/rangebitmap/range_bitmap_type_adapter_test.cpp common/file_index/bloomfilter/bloom_filter_file_index_test.cpp common/file_index/bloomfilter/fast_hash_test.cpp - common/global_index/complete_index_score_batch_reader_test.cpp common/global_index/global_index_result_test.cpp common/global_index/global_index_utils_test.cpp common/global_index/offset_global_index_reader_test.cpp @@ -636,6 +638,7 @@ if(PAIMON_BUILD_TESTS) common/predicate/predicate_test.cpp common/predicate/predicate_utils_test.cpp common/predicate/predicate_validator_test.cpp + common/reader/complete_index_score_batch_reader_test.cpp common/reader/concat_batch_reader_test.cpp common/reader/data_file_reader_factory_test.cpp common/reader/late_materializing_file_batch_reader_test.cpp diff --git a/src/paimon/common/file_index/file_index_format.cpp b/src/paimon/common/file_index/file_index_format.cpp index fab5c7a7..0b5bf4cd 100644 --- a/src/paimon/common/file_index/file_index_format.cpp +++ b/src/paimon/common/file_index/file_index_format.cpp @@ -30,6 +30,7 @@ #include "paimon/common/io/data_output_stream.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/math.h" +#include "paimon/core/io/file_index_options.h" #include "paimon/file_index/file_indexer.h" #include "paimon/file_index/file_indexer_factory.h" #include "paimon/io/byte_array_input_stream.h" @@ -170,7 +171,10 @@ class FileIndexFormatReaderImpl : public FileIndexFormat::Reader { std::unordered_map>>; static Result> Create( - const std::shared_ptr& input_stream, const std::shared_ptr& pool) { + const std::shared_ptr& input_stream, const std::shared_ptr& pool, + const std::map& options) { + PAIMON_ASSIGN_OR_RAISE(FileIndexOptions file_index_options, + FileIndexOptions::FromMap(options)); DataInputStream data_input_stream(input_stream); PAIMON_ASSIGN_OR_RAISE(int64_t magic, data_input_stream.ReadValue()); if (magic != FileIndexFormat::MAGIC) { @@ -207,8 +211,8 @@ class FileIndexFormatReaderImpl : public FileIndexFormat::Reader { index_map[index_type] = std::make_pair(offset, length); } } - return std::unique_ptr( - new FileIndexFormatReaderImpl(input_stream, std::move(header), pool)); + return std::unique_ptr(new FileIndexFormatReaderImpl( + input_stream, std::move(header), pool, std::move(file_index_options))); } Result>> ReadColumnIndex( @@ -224,9 +228,10 @@ class FileIndexFormatReaderImpl : public FileIndexFormat::Reader { if (index_iter != header_.end()) { const auto& index_map = index_iter->second; for (const auto& [index_type, offset_and_length] : index_map) { - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr file_index_reader, - GetFileIndexReader(arrow::schema({column_field}), index_type, - offset_and_length)); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr file_index_reader, + GetFileIndexReader(column_name, arrow::schema({column_field}), index_type, + offset_and_length)); if (file_index_reader) { // skip the index not registered res.push_back(std::move(file_index_reader)); @@ -238,19 +243,24 @@ class FileIndexFormatReaderImpl : public FileIndexFormat::Reader { private: FileIndexFormatReaderImpl(const std::shared_ptr& input_stream, HeaderType&& header, - const std::shared_ptr& pool) - : input_stream_(input_stream), pool_(pool), header_(std::move(header)) { + const std::shared_ptr& pool, FileIndexOptions&& options) + : input_stream_(input_stream), + pool_(pool), + header_(std::move(header)), + options_(std::move(options)) { assert(input_stream_); } Result> GetFileIndexReader( - const std::shared_ptr& arrow_schema, const std::string& index_type, - const std::pair& offset_and_length) const { + const std::string& column_name, const std::shared_ptr& arrow_schema, + const std::string& index_type, const std::pair& offset_and_length) const { if (offset_and_length.first == FileIndexFormat::EMPTY_INDEX_FLAG) { return std::make_shared(); } + const std::map& indexer_options = + options_.GetIndexerOptions(column_name, index_type); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr file_indexer, - FileIndexerFactory::Get(index_type, /*options=*/{})); + FileIndexerFactory::Get(index_type, indexer_options)); // assert(file_indexer); if (!file_indexer) { return std::shared_ptr(); @@ -267,6 +277,7 @@ class FileIndexFormatReaderImpl : public FileIndexFormat::Reader { // get header and cache it. // [column_name : [index_type : {offset, length}]] HeaderType header_; + FileIndexOptions options_; }; const int64_t FileIndexFormat::MAGIC = 1493475289347502LL; @@ -274,8 +285,9 @@ const int32_t FileIndexFormat::EMPTY_INDEX_FLAG = -1; const int32_t FileIndexFormat::V_1 = 1; Result> FileIndexFormat::CreateReader( - const std::shared_ptr& input_stream, const std::shared_ptr& pool) { - return FileIndexFormatReaderImpl::Create(input_stream, pool); + const std::shared_ptr& input_stream, const std::shared_ptr& pool, + const std::map& options) { + return FileIndexFormatReaderImpl::Create(input_stream, pool, options); } Result> FileIndexFormat::CreateWriter( diff --git a/src/paimon/common/file_index/file_index_format_test.cpp b/src/paimon/common/file_index/file_index_format_test.cpp index 40989f7e..2e9181dd 100644 --- a/src/paimon/common/file_index/file_index_format_test.cpp +++ b/src/paimon/common/file_index/file_index_format_test.cpp @@ -72,12 +72,13 @@ TEST_F(FileIndexFormatTest, TestWriteAndReadEmptyIndexGoldenBytes) { ASSERT_OK_AND_ASSIGN(auto writer, FileIndexFormat::CreateWriter(output, pool_)); ASSERT_OK(writer->WriteColumnIndexes(indexes)); ASSERT_OK(writer->Close()); - ASSERT_OK_AND_ASSIGN(std::shared_ptr actual, output->Finish(pool_.get())); + ASSERT_OK_AND_ASSIGN(PAIMON_UNIQUE_PTR actual, output->Finish(pool_.get())); ASSERT_EQ(expected, std::vector(actual->data(), actual->data() + actual->size())); auto schema = arrow::schema({arrow::field("c1", arrow::utf8())}); auto input_stream = std::make_shared(actual->data(), actual->size()); - ASSERT_OK_AND_ASSIGN(auto reader, FileIndexFormat::CreateReader(input_stream, pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, + FileIndexFormat::CreateReader(input_stream, pool_, /*options=*/{})); ASSERT_OK_AND_ASSIGN(auto index_file_readers, reader->ReadColumnIndex("c1", CreateArrowSchema(schema).get())); ASSERT_EQ(1, index_file_readers.size()); @@ -113,7 +114,8 @@ TEST_F(FileIndexFormatTest, TestSimple) { 0, 0, 0, 3, 0, 16, 0, 0, 0, 0, 0, 1, 0, 4, 0, 5, 0}; auto input_stream = std::make_shared( reinterpret_cast(index_file_bytes.data()), index_file_bytes.size()); - ASSERT_OK_AND_ASSIGN(auto reader, FileIndexFormat::CreateReader(input_stream, pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, + FileIndexFormat::CreateReader(input_stream, pool_, /*options=*/{})); { ASSERT_OK_AND_ASSIGN(auto index_file_readers, reader->ReadColumnIndex("f1", CreateArrowSchema(schema).get())); @@ -186,7 +188,8 @@ TEST_F(FileIndexFormatTest, TestBitmapIndexWithTimestamp) { ASSERT_OK(fs->ReadFile(file_name, &index_file_bytes)); auto input_stream = std::make_shared(index_file_bytes.data(), index_file_bytes.size()); - ASSERT_OK_AND_ASSIGN(auto reader, FileIndexFormat::CreateReader(input_stream, pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, + FileIndexFormat::CreateReader(input_stream, pool_, /*options=*/{})); auto check_second = [&](const std::string& field_name) { // data: second // 1745542802000lms, 0ns diff --git a/src/paimon/common/file_index/file_index_reader.cpp b/src/paimon/common/file_index/file_index_reader.cpp index 822e4579..9350e23f 100644 --- a/src/paimon/common/file_index/file_index_reader.cpp +++ b/src/paimon/common/file_index/file_index_reader.cpp @@ -21,6 +21,7 @@ #include #include "paimon/predicate/literal.h" +#include "paimon/status.h" namespace paimon { Result> FileIndexReader::VisitIsNotNull() { @@ -100,4 +101,14 @@ Result> FileIndexReader::VisitNotIn( } return file_index_result; } + +Result> FileIndexReader::VisitVectorSearch( + const std::shared_ptr& vector_search) { + return Status::NotImplemented("This file index does not support vector search"); +} + +Result> FileIndexReader::VisitFullTextSearch( + const std::shared_ptr& full_text_search) { + return Status::NotImplemented("This file index does not support full text search"); +} } // namespace paimon diff --git a/src/paimon/common/file_index/file_index_reader_test.cpp b/src/paimon/common/file_index/file_index_reader_test.cpp index b0db5bb7..a7a4c7db 100644 --- a/src/paimon/common/file_index/file_index_reader_test.cpp +++ b/src/paimon/common/file_index/file_index_reader_test.cpp @@ -21,7 +21,10 @@ #include #include "gtest/gtest.h" +#include "paimon/predicate/full_text_search.h" #include "paimon/predicate/literal.h" +#include "paimon/predicate/vector_search.h" +#include "paimon/testing/utils/testharness.h" namespace paimon::test { @@ -63,5 +66,14 @@ TEST(FileIndexReaderTest, TestDefaultIndexReader) { ASSERT_TRUE(reader.VisitIn({lit0, lit1}).value()->IsRemain().value()); ASSERT_TRUE(reader.VisitNotIn({lit0, lit1}).value()->IsRemain().value()); + + std::shared_ptr vector_search = + std::make_shared("embedding", 1, std::vector{1.0f}, nullptr, nullptr, + std::nullopt, std::map{}); + ASSERT_NOK_WITH_MSG(reader.VisitVectorSearch(vector_search), "does not support vector search"); + std::shared_ptr full_text_search = std::make_shared( + "body", std::nullopt, "word", FullTextSearch::SearchType::MATCH_ANY, std::nullopt); + ASSERT_NOK_WITH_MSG(reader.VisitFullTextSearch(full_text_search), + "does not support full text search"); } } // namespace paimon::test diff --git a/src/paimon/common/file_index/scored_file_index_result.cpp b/src/paimon/common/file_index/scored_file_index_result.cpp new file mode 100644 index 00000000..e12d9662 --- /dev/null +++ b/src/paimon/common/file_index/scored_file_index_result.cpp @@ -0,0 +1,57 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/file_index/scored_file_index_result.h" + +#include + +#include "fmt/format.h" +#include "fmt/ranges.h" +#include "paimon/status.h" + +namespace paimon { + +Result> ScoredFileIndexResult::Create( + RoaringBitmap32&& row_positions, std::vector&& scores) { + if (static_cast(row_positions.Cardinality()) != scores.size()) { + return Status::Invalid(fmt::format("Vector search returned {} row positions but {} scores", + row_positions.Cardinality(), scores.size())); + } + return std::shared_ptr( + new ScoredFileIndexResult(std::move(row_positions), std::move(scores))); +} + +ScoredFileIndexResult::ScoredFileIndexResult(RoaringBitmap32&& row_positions, + std::vector&& scores) + : row_positions_(std::move(row_positions)), scores_(std::move(scores)) {} + +Result ScoredFileIndexResult::IsRemain() const { + return !IsEmpty(); +} + +std::string ScoredFileIndexResult::ToString() const { + std::vector formatted_scores; + formatted_scores.reserve(scores_.size()); + for (float score : scores_) { + formatted_scores.push_back(fmt::format("{:.2f}", score)); + } + return fmt::format("row positions: {}, scores: {{{}}}", row_positions_.ToString(), + fmt::join(formatted_scores, ",")); +} + +} // namespace paimon diff --git a/src/paimon/common/file_index/scored_file_index_result_test.cpp b/src/paimon/common/file_index/scored_file_index_result_test.cpp new file mode 100644 index 00000000..8d06086f --- /dev/null +++ b/src/paimon/common/file_index/scored_file_index_result_test.cpp @@ -0,0 +1,55 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/file_index/scored_file_index_result.h" + +#include + +#include "gtest/gtest.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { + +TEST(ScoredFileIndexResultTest, TestCreate) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr result, + ScoredFileIndexResult::Create(RoaringBitmap32::From({2, 5}), {0.25f, 0.75f})); + EXPECT_FALSE(result->IsEmpty()); + EXPECT_EQ(RoaringBitmap32::From({2, 5}), result->GetRowPositions()); + EXPECT_EQ(std::vector({0.25f, 0.75f}), result->GetScores()); + ASSERT_OK_AND_ASSIGN(bool remain, result->IsRemain()); + EXPECT_TRUE(remain); + EXPECT_EQ("row positions: {2,5}, scores: {0.25,0.75}", result->ToString()); + + std::shared_ptr file_index_result = result; + ASSERT_OK_AND_ASSIGN(remain, file_index_result->IsRemain()); + EXPECT_TRUE(remain); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr empty, + ScoredFileIndexResult::Create(RoaringBitmap32(), {})); + EXPECT_TRUE(empty->IsEmpty()); + ASSERT_OK_AND_ASSIGN(remain, empty->IsRemain()); + EXPECT_FALSE(remain); +} + +TEST(ScoredFileIndexResultTest, TestRejectMismatchedScores) { + ASSERT_NOK_WITH_MSG(ScoredFileIndexResult::Create(RoaringBitmap32::From({1, 3}), {0.5f}), + "2 row positions but 1 scores"); +} + +} // namespace paimon::test diff --git a/src/paimon/common/io/byte_array_output_stream.cpp b/src/paimon/common/io/byte_array_output_stream.cpp index bc6d1f1a..be890330 100644 --- a/src/paimon/common/io/byte_array_output_stream.cpp +++ b/src/paimon/common/io/byte_array_output_stream.cpp @@ -22,7 +22,6 @@ #include #include #include -#include #include "paimon/common/memory/memory_segment_utils.h" #include "paimon/common/utils/math.h" @@ -60,21 +59,19 @@ Status ByteArrayOutputStream::Close() { return Status::OK(); } -Result> ByteArrayOutputStream::Finish(MemoryPool* pool) { +Result> ByteArrayOutputStream::Finish(MemoryPool* pool) { assert(pool); - PAIMON_RETURN_NOT_OK(Close()); - if (result_) { - return result_; + if (finished_) { + return Status::Invalid("Byte array output stream has already been finished"); } + PAIMON_RETURN_NOT_OK(Close()); // TODO(jinli.zjw): Support int64_t lengths in MemorySegmentUtils::CopyToBytes and remove this // limit. const int64_t size = output_->CurrentSize(); PAIMON_RETURN_NOT_OK(ValidateValueInRange(size, "byte array output stream size")); - const std::vector& segments = output_->Segments(); - result_ = std::make_shared(static_cast(size), pool); - MemorySegmentUtils::CopyToBytes(segments, /*offset=*/0, result_.get(), - /*bytes_offset=*/0, static_cast(size)); - return result_; + finished_ = true; + return MemorySegmentUtils::CopyToBytes(output_->Segments(), /*offset=*/0, + static_cast(size), pool); } } // namespace paimon diff --git a/src/paimon/common/io/byte_array_output_stream.h b/src/paimon/common/io/byte_array_output_stream.h index 9b87ca42..75292541 100644 --- a/src/paimon/common/io/byte_array_output_stream.h +++ b/src/paimon/common/io/byte_array_output_stream.h @@ -24,16 +24,17 @@ #include "paimon/common/io/memory_segment_output_stream.h" #include "paimon/fs/file_system.h" +#include "paimon/memory/bytes.h" #include "paimon/result.h" #include "paimon/status.h" +#include "paimon/visibility.h" namespace paimon { -class Bytes; class MemoryPool; /// An in-memory output stream backed by segments allocated from a Paimon MemoryPool. -class ByteArrayOutputStream : public OutputStream { +class PAIMON_EXPORT ByteArrayOutputStream : public OutputStream { public: /// Takes ownership of an initialized segmented output stream. explicit ByteArrayOutputStream(std::unique_ptr&& output); @@ -57,13 +58,14 @@ class ByteArrayOutputStream : public OutputStream { Status Close() override; /// Closes the stream and returns its contents as an exactly-sized contiguous byte array. + /// This method can only be called once. /// @note The caller must keep `pool` alive until the returned bytes are destroyed. - Result> Finish(MemoryPool* pool); + Result> Finish(MemoryPool* pool); private: std::unique_ptr output_; - std::shared_ptr result_; bool closed_ = false; + bool finished_ = false; }; } // namespace paimon diff --git a/src/paimon/common/io/byte_array_output_stream_test.cpp b/src/paimon/common/io/byte_array_output_stream_test.cpp index bd185095..c4adb84f 100644 --- a/src/paimon/common/io/byte_array_output_stream_test.cpp +++ b/src/paimon/common/io/byte_array_output_stream_test.cpp @@ -43,14 +43,16 @@ TEST(ByteArrayOutputStreamTest, TestWriteAndFinish) { ASSERT_EQ(6, position); ASSERT_EQ(pool->CurrentUsage(), pool->MaxMemoryUsage()); - ASSERT_OK_AND_ASSIGN(std::shared_ptr result, stream->Finish(pool.get())); + ASSERT_OK_AND_ASSIGN(PAIMON_UNIQUE_PTR result, stream->Finish(pool.get())); ASSERT_EQ("abcdef", std::string(result->data(), result->size())); ASSERT_NOK_WITH_MSG(stream->Write("x", 1), "closed"); - ASSERT_OK_AND_ASSIGN(std::shared_ptr repeated, stream->Finish(pool.get())); - ASSERT_EQ(result, repeated); + ASSERT_NOK_WITH_MSG(stream->Finish(pool.get()), "already been finished"); stream.reset(); + ASSERT_EQ("abcdef", std::string(result->data(), result->size())); ASSERT_EQ(6, pool->CurrentUsage()); + result.reset(); + ASSERT_EQ(0, pool->CurrentUsage()); } TEST(ByteArrayOutputStreamTest, TestWriteValidation) { @@ -62,7 +64,7 @@ TEST(ByteArrayOutputStreamTest, TestWriteValidation) { ASSERT_NOK(stream->Write("", -1)); ASSERT_OK_AND_ASSIGN(int64_t written, stream->Write(nullptr, 0)); ASSERT_EQ(0, written); - ASSERT_OK_AND_ASSIGN(std::shared_ptr result, stream->Finish(pool.get())); + ASSERT_OK_AND_ASSIGN(PAIMON_UNIQUE_PTR result, stream->Finish(pool.get())); ASSERT_EQ(0, result->size()); } @@ -73,7 +75,7 @@ TEST(ByteArrayOutputStreamTest, TestCallerKeepsMemoryPoolAlive) { std::make_shared(std::move(output)); ASSERT_OK_AND_ASSIGN(int64_t written, stream->Write("data", 4)); ASSERT_EQ(4, written); - ASSERT_OK_AND_ASSIGN(std::shared_ptr result, stream->Finish(pool.get())); + ASSERT_OK_AND_ASSIGN(PAIMON_UNIQUE_PTR result, stream->Finish(pool.get())); stream.reset(); ASSERT_GT(pool->CurrentUsage(), 0); diff --git a/src/paimon/common/io/offset_input_stream.cpp b/src/paimon/common/io/offset_input_stream.cpp index 2df88adc..fe7d0967 100644 --- a/src/paimon/common/io/offset_input_stream.cpp +++ b/src/paimon/common/io/offset_input_stream.cpp @@ -118,7 +118,8 @@ void OffsetInputStream::ReadAsync(char* buffer, int64_t size, int64_t offset, } Status OffsetInputStream::Close() { - return wrapped_->Close(); + // TODO(jinli.zjw): Decide whether closing this view should also close the wrapped input stream. + return Status::OK(); } Result OffsetInputStream::GetUri() const { diff --git a/src/paimon/common/io/offset_input_stream.h b/src/paimon/common/io/offset_input_stream.h index 9fb6f4c0..7c2a4db6 100644 --- a/src/paimon/common/io/offset_input_stream.h +++ b/src/paimon/common/io/offset_input_stream.h @@ -30,7 +30,9 @@ #include "paimon/visibility.h" namespace paimon { -/// A `InputStream` wrapping another `InputStream` with offset and length. +/// An `InputStream` wrapping another `InputStream` with offset and length. +/// @note Not safe for concurrent or interleaved use with other views of the wrapped stream: Seek +/// and sequential Read share the wrapped stream's position and may read from the wrong offset. class PAIMON_EXPORT OffsetInputStream : public InputStream { public: static Result> Create( diff --git a/src/paimon/common/global_index/complete_index_score_batch_reader.cpp b/src/paimon/common/reader/complete_index_score_batch_reader.cpp similarity index 84% rename from src/paimon/common/global_index/complete_index_score_batch_reader.cpp rename to src/paimon/common/reader/complete_index_score_batch_reader.cpp index 75eb3610..44cc9e96 100644 --- a/src/paimon/common/global_index/complete_index_score_batch_reader.cpp +++ b/src/paimon/common/reader/complete_index_score_batch_reader.cpp @@ -17,20 +17,18 @@ * under the License. */ -#include "paimon/common/global_index/complete_index_score_batch_reader.h" +#include "paimon/common/reader/complete_index_score_batch_reader.h" #include #include "arrow/api.h" #include "arrow/array/array_base.h" #include "arrow/array/array_nested.h" -#include "arrow/array/util.h" -#include "arrow/c/abi.h" #include "arrow/c/bridge.h" #include "arrow/scalar.h" +#include "fmt/format.h" #include "paimon/common/reader/reader_utils.h" #include "paimon/common/table/special_fields.h" -#include "paimon/common/types/row_kind.h" #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" @@ -64,6 +62,10 @@ Result CompleteIndexScoreBatchReader::NextBatc PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, reader_->NextBatchWithBitmap()); if (BatchReader::IsEofBatch(batch_with_bitmap)) { + if (!scores_.empty() && score_cursor_ != scores_.size()) { + return Status::Invalid(fmt::format("Index score count {} does not match rows read {}", + scores_.size(), score_cursor_)); + } return batch_with_bitmap; } if (scores_.empty()) { @@ -93,6 +95,12 @@ Result CompleteIndexScoreBatchReader::NextBatc auto* typed_builder = checked_cast(index_score_builder.get()); PAIMON_RETURN_NOT_OK_FROM_ARROW(typed_builder->Reserve(struct_array->length())); bool all_not_null = (struct_array->length() == bitmap.Cardinality()); + size_t score_count = all_not_null ? static_cast(struct_array->length()) + : static_cast(bitmap.Cardinality()); + if (score_cursor_ > scores_.size() || score_count > scores_.size() - score_cursor_) { + return Status::Invalid(fmt::format("Index score count {} is smaller than rows read {}", + scores_.size(), score_cursor_ + score_count)); + } for (int64_t i = 0; i < struct_array->length(); i++) { if (all_not_null || bitmap.Contains(i)) { PAIMON_RETURN_NOT_OK_FROM_ARROW(typed_builder->Append(scores_[score_cursor_++])); @@ -112,4 +120,10 @@ Result CompleteIndexScoreBatchReader::NextBatc PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_)); return batch_with_bitmap; } + +void CompleteIndexScoreBatchReader::ResetScoreState() { + score_cursor_ = 0; + index_score_field_idx_ = -1; + field_names_with_score_.clear(); +} } // namespace paimon diff --git a/src/paimon/common/global_index/complete_index_score_batch_reader.h b/src/paimon/common/reader/complete_index_score_batch_reader.h similarity index 95% rename from src/paimon/common/global_index/complete_index_score_batch_reader.h rename to src/paimon/common/reader/complete_index_score_batch_reader.h index 427d854b..cd3d22cc 100644 --- a/src/paimon/common/global_index/complete_index_score_batch_reader.h +++ b/src/paimon/common/reader/complete_index_score_batch_reader.h @@ -48,6 +48,9 @@ class CompleteIndexScoreBatchReader : public BatchReader { Result NextBatchWithBitmap() override; + /// Reset score alignment after the wrapped reader restarts from its first row. + void ResetScoreState(); + void Close() override { reader_->Close(); } diff --git a/src/paimon/common/global_index/complete_index_score_batch_reader_test.cpp b/src/paimon/common/reader/complete_index_score_batch_reader_test.cpp similarity index 70% rename from src/paimon/common/global_index/complete_index_score_batch_reader_test.cpp rename to src/paimon/common/reader/complete_index_score_batch_reader_test.cpp index ec56e083..f8dac396 100644 --- a/src/paimon/common/global_index/complete_index_score_batch_reader_test.cpp +++ b/src/paimon/common/reader/complete_index_score_batch_reader_test.cpp @@ -17,7 +17,7 @@ * under the License. */ -#include "paimon/common/global_index/complete_index_score_batch_reader.h" +#include "paimon/common/reader/complete_index_score_batch_reader.h" #include "arrow/api.h" #include "arrow/array/array_base.h" @@ -25,6 +25,7 @@ #include "arrow/c/bridge.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" +#include "paimon/common/reader/complete_index_score_file_batch_reader.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/mem_utils.h" @@ -160,4 +161,52 @@ TEST_F(CompleteIndexScoreBatchReaderTest, TestReadWithNullScores) { ASSERT_TRUE(expected_array->Equals(*result_array)); } +TEST_F(CompleteIndexScoreBatchReaderTest, TestFileReaderForwardsOperationsAndResetsScores) { + arrow::FieldVector fields = {arrow::field("f0", arrow::utf8()), + arrow::field("_INDEX_SCORE", arrow::float32())}; + auto data = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([ + ["Alice", null], + ["Bob", null] + ])") + .ValueOrDie(); + auto inner_reader = std::make_unique(data, data->type(), /*batch_size=*/1); + MockFileBatchReader* inner = inner_reader.get(); + auto reader = std::make_unique( + std::move(inner_reader), std::vector{1.25f, 2.5f}, GetArrowPool(GetDefaultPool())); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<::ArrowSchema> file_schema, reader->GetFileSchema()); + EXPECT_TRUE(file_schema); + ASSERT_OK_AND_ASSIGN(uint64_t row_count, reader->GetNumberOfRows()); + EXPECT_EQ(2, row_count); + EXPECT_FALSE(reader->SupportPreciseBitmapSelection()); + reader->Warmup(); + EXPECT_EQ(1, inner->GetWarmupCount()); + + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap first, reader->NextBatchWithBitmap()); + ASSERT_OK_AND_ASSIGN(uint64_t file_row_id, reader->GetPreviousBatchFileRowId(0)); + EXPECT_EQ(0, file_row_id); + auto first_array = + arrow::ImportArray(first.first.first.get(), first.first.second.get()).ValueOrDie(); + auto first_struct = std::dynamic_pointer_cast(first_array); + ASSERT_TRUE(first_struct); + auto first_scores = + std::dynamic_pointer_cast(first_struct->GetFieldByName("_INDEX_SCORE")); + ASSERT_TRUE(first_scores); + EXPECT_FLOAT_EQ(1.25f, first_scores->Value(0)); + + ::ArrowSchema read_schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(fields), &read_schema).ok()); + ASSERT_OK(reader->SetReadSchema(&read_schema, /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap restarted, reader->NextBatchWithBitmap()); + auto restarted_array = + arrow::ImportArray(restarted.first.first.get(), restarted.first.second.get()).ValueOrDie(); + auto restarted_struct = std::dynamic_pointer_cast(restarted_array); + ASSERT_TRUE(restarted_struct); + auto restarted_scores = std::dynamic_pointer_cast( + restarted_struct->GetFieldByName("_INDEX_SCORE")); + ASSERT_TRUE(restarted_scores); + EXPECT_FLOAT_EQ(1.25f, restarted_scores->Value(0)); +} + } // namespace paimon::test diff --git a/src/paimon/common/reader/complete_index_score_file_batch_reader.cpp b/src/paimon/common/reader/complete_index_score_file_batch_reader.cpp new file mode 100644 index 00000000..cce919d9 --- /dev/null +++ b/src/paimon/common/reader/complete_index_score_file_batch_reader.cpp @@ -0,0 +1,73 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/common/reader/complete_index_score_file_batch_reader.h" + +#include + +namespace paimon { +CompleteIndexScoreFileBatchReader::CompleteIndexScoreFileBatchReader( + std::unique_ptr&& reader, const std::vector& scores, + const std::shared_ptr& arrow_pool) + : file_reader_(reader.get()), score_reader_(std::move(reader), scores, arrow_pool) {} + +Result CompleteIndexScoreFileBatchReader::NextBatch() { + return score_reader_.NextBatch(); +} + +Result CompleteIndexScoreFileBatchReader::NextBatchWithBitmap() { + return score_reader_.NextBatchWithBitmap(); +} + +Result> CompleteIndexScoreFileBatchReader::GetFileSchema() const { + return file_reader_->GetFileSchema(); +} + +Status CompleteIndexScoreFileBatchReader::SetReadSchema( + ::ArrowSchema* read_schema, const std::shared_ptr& predicate, + const std::optional& selection_bitmap) { + score_reader_.ResetScoreState(); + return file_reader_->SetReadSchema(read_schema, predicate, selection_bitmap); +} + +Result CompleteIndexScoreFileBatchReader::GetPreviousBatchFileRowId( + uint64_t batch_row_id) const { + return file_reader_->GetPreviousBatchFileRowId(batch_row_id); +} + +Result CompleteIndexScoreFileBatchReader::GetNumberOfRows() const { + return file_reader_->GetNumberOfRows(); +} + +bool CompleteIndexScoreFileBatchReader::SupportPreciseBitmapSelection() const { + return file_reader_->SupportPreciseBitmapSelection(); +} + +void CompleteIndexScoreFileBatchReader::Warmup() { + file_reader_->Warmup(); +} + +void CompleteIndexScoreFileBatchReader::Close() { + score_reader_.Close(); +} + +std::shared_ptr CompleteIndexScoreFileBatchReader::GetReaderMetrics() const { + return score_reader_.GetReaderMetrics(); +} +} // namespace paimon diff --git a/src/paimon/common/reader/complete_index_score_file_batch_reader.h b/src/paimon/common/reader/complete_index_score_file_batch_reader.h new file mode 100644 index 00000000..721f4652 --- /dev/null +++ b/src/paimon/common/reader/complete_index_score_file_batch_reader.h @@ -0,0 +1,55 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include +#include + +#include "paimon/common/reader/complete_index_score_batch_reader.h" +#include "paimon/reader/file_batch_reader.h" + +namespace paimon { +/// Adds index scores to a single-file reader while retaining its file-specific operations. +class CompleteIndexScoreFileBatchReader : public FileBatchReader { + public: + CompleteIndexScoreFileBatchReader(std::unique_ptr&& reader, + const std::vector& scores, + const std::shared_ptr& arrow_pool); + + Result NextBatch() override; + Result NextBatchWithBitmap() override; + Result> GetFileSchema() const override; + Status SetReadSchema(::ArrowSchema* read_schema, const std::shared_ptr& predicate, + const std::optional& selection_bitmap) override; + Result GetPreviousBatchFileRowId(uint64_t batch_row_id) const override; + Result GetNumberOfRows() const override; + bool SupportPreciseBitmapSelection() const override; + void Warmup() override; + void Close() override; + std::shared_ptr GetReaderMetrics() const override; + + private: + // Non-owning; score_reader_ owns this FileBatchReader for the wrapper's lifetime. + FileBatchReader* file_reader_; + CompleteIndexScoreBatchReader score_reader_; +}; +} // namespace paimon diff --git a/src/paimon/core/append/append_only_writer_test.cpp b/src/paimon/core/append/append_only_writer_test.cpp index 835717d4..a335a4c6 100644 --- a/src/paimon/core/append/append_only_writer_test.cpp +++ b/src/paimon/core/append/append_only_writer_test.cpp @@ -429,7 +429,8 @@ TEST_F(AppendOnlyWriterTest, TestWritePublishesEmbeddedBitmapIndex) { auto input = std::make_shared(files[0]->embedded_index->data(), files[0]->embedded_index->size()); - ASSERT_OK_AND_ASSIGN(auto index_reader, FileIndexFormat::CreateReader(input, memory_pool_)); + ASSERT_OK_AND_ASSIGN(auto index_reader, + FileIndexFormat::CreateReader(input, memory_pool_, /*options=*/{})); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*schema, &c_schema).ok()); ASSERT_OK_AND_ASSIGN(auto column_readers, index_reader->ReadColumnIndex("f0", &c_schema)); diff --git a/src/paimon/core/io/data_file_index_writer.cpp b/src/paimon/core/io/data_file_index_writer.cpp index 5075c9a6..f95b4bdd 100644 --- a/src/paimon/core/io/data_file_index_writer.cpp +++ b/src/paimon/core/io/data_file_index_writer.cpp @@ -76,6 +76,13 @@ Result> DataFileIndexWriter::Create( return Status::Invalid( fmt::format("File index type '{}' is not registered", definition.index_type)); } + PAIMON_ASSIGN_OR_RAISE(std::optional> extra_field_names, + indexer->GetExtraFieldNames()); + if (extra_field_names) { + return Status::NotImplemented( + fmt::format("File index '{}.{}' requires extra fields, which are not supported", + definition.index_type, definition.column_name)); + } ::ArrowSchema c_schema; ArrowSchemaMarkReleased(&c_schema); ScopeGuard schema_guard([&c_schema]() { ArrowSchemaRelease(&c_schema); }); @@ -157,7 +164,8 @@ Result> DataFileIndexWriter::SerializeContainer() { FileIndexFormat::CreateWriter(output, pool_)); PAIMON_RETURN_NOT_OK(format_writer->WriteColumnIndexes(column_indexes)); PAIMON_RETURN_NOT_OK(format_writer->Close()); - return output->Finish(pool_.get()); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr bytes, output->Finish(pool_.get())); + return bytes; } Result DataFileIndexWriter::Finish(const std::string& data_file_path) { diff --git a/src/paimon/core/io/data_file_index_writer_test.cpp b/src/paimon/core/io/data_file_index_writer_test.cpp index 524fc139..b6f25f3c 100644 --- a/src/paimon/core/io/data_file_index_writer_test.cpp +++ b/src/paimon/core/io/data_file_index_writer_test.cpp @@ -125,7 +125,7 @@ class DataFileIndexWriterTest : public ::testing::Test { Result> CreateReader( const std::shared_ptr& bytes) const { auto input = std::make_shared(bytes->data(), bytes->size()); - return FileIndexFormat::CreateReader(input, pool_); + return FileIndexFormat::CreateReader(input, pool_, /*options=*/{}); } Result>> ReadColumn( @@ -244,7 +244,7 @@ TEST_F(DataFileIndexWriterTest, TestExternalIndexAndAbortCleanup) { ASSERT_OK_AND_ASSIGN(bool exists, file_system_->Exists(index_path)); ASSERT_TRUE(exists); ASSERT_OK_AND_ASSIGN(std::shared_ptr input, file_system_->Open(index_path)); - ASSERT_OK_AND_ASSIGN(auto reader, FileIndexFormat::CreateReader(input, pool_)); + ASSERT_OK_AND_ASSIGN(auto reader, FileIndexFormat::CreateReader(input, pool_, /*options=*/{})); ASSERT_OK_AND_ASSIGN(auto bitmap_readers, ReadColumn(reader.get(), "f0")); ASSERT_EQ(1, bitmap_readers.size()); ASSERT_OK_AND_ASSIGN(auto equal_result, bitmap_readers[0]->VisitEqual(Literal(1))); diff --git a/src/paimon/core/io/file_index_evaluator.cpp b/src/paimon/core/io/file_index_evaluator.cpp index f2943f0b..7ffa3b32 100644 --- a/src/paimon/core/io/file_index_evaluator.cpp +++ b/src/paimon/core/io/file_index_evaluator.cpp @@ -22,28 +22,30 @@ #include #include #include +#include #include +#include #include "arrow/c/bridge.h" #include "arrow/type.h" #include "fmt/format.h" #include "fmt/ranges.h" #include "paimon/common/utils/arrow/status_utils.h" -#include "paimon/common/utils/date_time_utils.h" -#include "paimon/common/utils/field_type_utils.h" #include "paimon/common/utils/string_utils.h" +#include "paimon/core/core_options.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/file_index/file_index_format.h" #include "paimon/file_index/file_index_reader.h" +#include "paimon/file_index/scored_file_index_result.h" #include "paimon/io/byte_array_input_stream.h" -#include "paimon/memory/bytes.h" #include "paimon/predicate/compound_predicate.h" +#include "paimon/predicate/full_text_search.h" #include "paimon/predicate/function.h" #include "paimon/predicate/leaf_predicate.h" -#include "paimon/predicate/literal.h" #include "paimon/predicate/predicate.h" #include "paimon/predicate/predicate_utils.h" +#include "paimon/predicate/vector_search.h" #include "paimon/status.h" namespace paimon { @@ -51,24 +53,26 @@ class MemoryPool; enum class FieldType; Result> FileIndexEvaluator::Evaluate( - const std::shared_ptr& data_schema, const std::shared_ptr& predicate, + const std::shared_ptr& data_schema, const CoreOptions& core_options, + const std::shared_ptr& predicate, const std::shared_ptr& data_file_path_factory, const std::shared_ptr& file_meta, const std::shared_ptr& file_system, const std::shared_ptr& pool) { - return Evaluate(/*only_use_embedding_index=*/false, data_schema, predicate, + return Evaluate(/*only_use_embedding_index=*/false, data_schema, core_options, predicate, data_file_path_factory, file_meta, file_system, pool); } Result> FileIndexEvaluator::Evaluate( - const std::shared_ptr& data_schema, const std::shared_ptr& predicate, - const std::shared_ptr& file_meta, const std::shared_ptr& pool) { - return Evaluate(/*only_use_embedding_index=*/true, data_schema, predicate, + const std::shared_ptr& data_schema, const CoreOptions& core_options, + const std::shared_ptr& predicate, const std::shared_ptr& file_meta, + const std::shared_ptr& pool) { + return Evaluate(/*only_use_embedding_index=*/true, data_schema, core_options, predicate, /*data_file_path_factory=*/nullptr, file_meta, /*file_system=*/nullptr, pool); } Result> FileIndexEvaluator::Evaluate( bool only_use_embedding_index, const std::shared_ptr& data_schema, - const std::shared_ptr& predicate, + const CoreOptions& core_options, const std::shared_ptr& predicate, const std::shared_ptr& data_file_path_factory, const std::shared_ptr& file_meta, const std::shared_ptr& file_system, const std::shared_ptr& pool) { @@ -83,7 +87,7 @@ Result> FileIndexEvaluator::Evaluate( return FileIndexResult::Remain(); } PAIMON_ASSIGN_OR_RAISE(std::unique_ptr format_reader, - FileIndexFormat::CreateReader(input_stream, pool)); + FileIndexFormat::CreateReader(input_stream, pool, core_options.ToMap())); std::set required_field_names; PAIMON_RETURN_NOT_OK(PredicateUtils::GetAllNames(predicate, &required_field_names)); std::map>> @@ -147,6 +151,63 @@ Result> FileIndexEvaluator::ExtractIndexInputStream return std::shared_ptr(); } +namespace { + +Result>> OpenSearchReaders( + const std::shared_ptr& data_schema, const std::string& field_name, + const CoreOptions& core_options, const std::shared_ptr& file_meta, + const std::shared_ptr& input_stream, const std::shared_ptr& pool) { + std::shared_ptr field = data_schema->GetFieldByName(field_name); + if (!field) { + return Status::Invalid( + fmt::format("Search field '{}' does not exist in data schema", field_name)); + } + if (!input_stream) { + return Status::Invalid(fmt::format("Data file '{}' has no File Index for search field '{}'", + file_meta->file_name, field_name)); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr format_reader, + FileIndexFormat::CreateReader(input_stream, pool, core_options.ToMap())); + ::ArrowSchema c_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportSchema(*arrow::schema({std::move(field)}), &c_schema)); + return format_reader->ReadColumnIndex(field_name, &c_schema); +} + +Result> ExecuteVectorSearch( + const std::vector>& readers, + const std::shared_ptr& vector_search) { + if (readers.empty()) { + return Status::NotImplemented( + fmt::format("No File Index reader supports vector search for field '{}'", + vector_search->field_name)); + } + if (readers.size() != 1) { + return Status::Invalid( + fmt::format("Multiple File Index readers exist for vector search field '{}'", + vector_search->field_name)); + } + return readers[0]->VisitVectorSearch(vector_search); +} + +Result> ExecuteFullTextSearch( + const std::vector>& readers, + const std::shared_ptr& full_text_search) { + if (readers.empty()) { + return Status::NotImplemented( + fmt::format("No File Index reader supports full-text search for field '{}'", + full_text_search->field_name)); + } + if (readers.size() != 1) { + return Status::Invalid( + fmt::format("Multiple File Index readers exist for full-text search field '{}'", + full_text_search->field_name)); + } + return readers[0]->VisitFullTextSearch(full_text_search); +} + +} // namespace + Result> FileIndexEvaluator::Evaluate( const std::shared_ptr& predicate, const std::map>>& @@ -223,4 +284,40 @@ Result> FileIndexEvaluator::EvaluateCompoundPre return Status::Invalid("CompoundPredicate only support And/Or function"); } +Result> FileIndexEvaluator::EvaluateVectorSearch( + const std::shared_ptr& data_schema, const CoreOptions& core_options, + const std::shared_ptr& vector_search, + const std::shared_ptr& path_factory, + const std::shared_ptr& file_meta, const std::shared_ptr& file_system, + const std::shared_ptr& pool) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input_stream, + ExtractIndexInputStream(/*only_use_embedding_index=*/false, path_factory, + file_meta, file_system)); + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + OpenSearchReaders(data_schema, vector_search->field_name, core_options, + file_meta, input_stream, pool)); + return ExecuteVectorSearch(readers, vector_search); +} + +Result> FileIndexEvaluator::EvaluateFullTextSearch( + const std::shared_ptr& data_schema, const CoreOptions& core_options, + const std::shared_ptr& full_text_search, + const std::shared_ptr& path_factory, + const std::shared_ptr& file_meta, const std::shared_ptr& file_system, + const std::shared_ptr& pool) { + if (full_text_search->pre_filter) { + return Status::NotImplemented("File full-text search does not support pre-filter yet"); + } + if (full_text_search->with_score) { + return Status::NotImplemented("File full-text search does not support score output yet"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input_stream, + ExtractIndexInputStream(/*only_use_embedding_index=*/false, path_factory, + file_meta, file_system)); + PAIMON_ASSIGN_OR_RAISE(std::vector> readers, + OpenSearchReaders(data_schema, full_text_search->field_name, + core_options, file_meta, input_stream, pool)); + return ExecuteFullTextSearch(readers, full_text_search); +} + } // namespace paimon diff --git a/src/paimon/core/io/file_index_evaluator.h b/src/paimon/core/io/file_index_evaluator.h index bdec895d..34ac1c9e 100644 --- a/src/paimon/core/io/file_index_evaluator.h +++ b/src/paimon/core/io/file_index_evaluator.h @@ -37,11 +37,15 @@ class Schema; namespace paimon { class LeafPredicate; class CompoundPredicate; +class CoreOptions; class DataFilePathFactory; class FileIndexReader; class MemoryPool; class Predicate; +class ScoredFileIndexResult; struct DataFileMeta; +struct FullTextSearch; +struct VectorSearch; class FileIndexEvaluator { public: @@ -50,22 +54,36 @@ class FileIndexEvaluator { // for scan process, will only use embedding file to evaluate static Result> Evaluate( - const std::shared_ptr& data_schema, + const std::shared_ptr& data_schema, const CoreOptions& core_options, const std::shared_ptr& predicate, const std::shared_ptr& file_meta, const std::shared_ptr& pool); // for read process, will use embedding file or extra index file to evaluate static Result> Evaluate( - const std::shared_ptr& data_schema, + const std::shared_ptr& data_schema, const CoreOptions& core_options, const std::shared_ptr& predicate, const std::shared_ptr& data_file_path_factory, const std::shared_ptr& file_meta, const std::shared_ptr& file_system, const std::shared_ptr& pool); + static Result> EvaluateVectorSearch( + const std::shared_ptr& data_schema, const CoreOptions& core_options, + const std::shared_ptr& vector_search, + const std::shared_ptr& path_factory, + const std::shared_ptr& file_meta, + const std::shared_ptr& file_system, const std::shared_ptr& pool); + + static Result> EvaluateFullTextSearch( + const std::shared_ptr& data_schema, const CoreOptions& core_options, + const std::shared_ptr& full_text_search, + const std::shared_ptr& path_factory, + const std::shared_ptr& file_meta, + const std::shared_ptr& file_system, const std::shared_ptr& pool); + private: static Result> Evaluate( bool only_use_embedding_index, const std::shared_ptr& data_schema, - const std::shared_ptr& predicate, + const CoreOptions& core_options, const std::shared_ptr& predicate, const std::shared_ptr& data_file_path_factory, const std::shared_ptr& file_meta, const std::shared_ptr& file_system, const std::shared_ptr& pool); diff --git a/src/paimon/core/io/file_index_evaluator_test.cpp b/src/paimon/core/io/file_index_evaluator_test.cpp index 1824564e..0eb7bc21 100644 --- a/src/paimon/core/io/file_index_evaluator_test.cpp +++ b/src/paimon/core/io/file_index_evaluator_test.cpp @@ -21,7 +21,9 @@ #include #include #include +#include #include +#include #include "arrow/type_fwd.h" #include "gtest/gtest.h" @@ -29,6 +31,7 @@ #include "paimon/common/fs/external_path_provider.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/date_time_utils.h" +#include "paimon/core/core_options.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/core/manifest/file_source.h" @@ -86,8 +89,8 @@ class FileIndexEvaluatorTest : public ::testing::Test { PredicateBuilder::IsNull(/*field_index=*/2, /*field_name=*/"f2", FieldType::INT); ASSERT_OK_AND_ASSIGN( auto file_index_result, - FileIndexEvaluator::Evaluate(data_schema_, predicate, data_file_path_factory, - data_file_meta, fs_, pool_)); + FileIndexEvaluator::Evaluate(data_schema_, core_options_, predicate, + data_file_path_factory, data_file_meta, fs_, pool_)); ASSERT_TRUE(file_index_result->IsRemain().value()); CheckResult(file_index_result, {7}); } @@ -96,8 +99,8 @@ class FileIndexEvaluatorTest : public ::testing::Test { PredicateBuilder::IsNotNull(/*field_index=*/2, /*field_name=*/"f2", FieldType::INT); ASSERT_OK_AND_ASSIGN( auto file_index_result, - FileIndexEvaluator::Evaluate(data_schema_, predicate, data_file_path_factory, - data_file_meta, fs_, pool_)); + FileIndexEvaluator::Evaluate(data_schema_, core_options_, predicate, + data_file_path_factory, data_file_meta, fs_, pool_)); ASSERT_TRUE(file_index_result->IsRemain().value()); CheckResult(file_index_result, {0, 1, 2, 3, 4, 5, 6}); } @@ -107,8 +110,8 @@ class FileIndexEvaluatorTest : public ::testing::Test { Literal(FieldType::STRING, "Alice", 5)); ASSERT_OK_AND_ASSIGN( auto file_index_result, - FileIndexEvaluator::Evaluate(data_schema_, predicate, data_file_path_factory, - data_file_meta, fs_, pool_)); + FileIndexEvaluator::Evaluate(data_schema_, core_options_, predicate, + data_file_path_factory, data_file_meta, fs_, pool_)); ASSERT_TRUE(file_index_result->IsRemain().value()); CheckResult(file_index_result, {0, 7}); } @@ -118,8 +121,8 @@ class FileIndexEvaluatorTest : public ::testing::Test { Literal(FieldType::STRING, "Alice", 5)); ASSERT_OK_AND_ASSIGN( auto file_index_result, - FileIndexEvaluator::Evaluate(data_schema_, predicate, data_file_path_factory, - data_file_meta, fs_, pool_)); + FileIndexEvaluator::Evaluate(data_schema_, core_options_, predicate, + data_file_path_factory, data_file_meta, fs_, pool_)); ASSERT_TRUE(file_index_result->IsRemain().value()); CheckResult(file_index_result, {1, 2, 3, 4, 5, 6}); } @@ -128,8 +131,8 @@ class FileIndexEvaluatorTest : public ::testing::Test { FieldType::INT, Literal(10)); ASSERT_OK_AND_ASSIGN( auto file_index_result, - FileIndexEvaluator::Evaluate(data_schema_, predicate, data_file_path_factory, - data_file_meta, fs_, pool_)); + FileIndexEvaluator::Evaluate(data_schema_, core_options_, predicate, + data_file_path_factory, data_file_meta, fs_, pool_)); ASSERT_TRUE(file_index_result->IsRemain().value()); ASSERT_FALSE(dynamic_cast(file_index_result.get())); } @@ -138,8 +141,8 @@ class FileIndexEvaluatorTest : public ::testing::Test { /*field_index=*/1, /*field_name=*/"f1", FieldType::INT, Literal(10)); ASSERT_OK_AND_ASSIGN( auto file_index_result, - FileIndexEvaluator::Evaluate(data_schema_, predicate, data_file_path_factory, - data_file_meta, fs_, pool_)); + FileIndexEvaluator::Evaluate(data_schema_, core_options_, predicate, + data_file_path_factory, data_file_meta, fs_, pool_)); ASSERT_TRUE(file_index_result->IsRemain().value()); ASSERT_FALSE(dynamic_cast(file_index_result.get())); } @@ -148,8 +151,8 @@ class FileIndexEvaluatorTest : public ::testing::Test { FieldType::INT, Literal(10)); ASSERT_OK_AND_ASSIGN( auto file_index_result, - FileIndexEvaluator::Evaluate(data_schema_, predicate, data_file_path_factory, - data_file_meta, fs_, pool_)); + FileIndexEvaluator::Evaluate(data_schema_, core_options_, predicate, + data_file_path_factory, data_file_meta, fs_, pool_)); ASSERT_TRUE(file_index_result->IsRemain().value()); ASSERT_FALSE(dynamic_cast(file_index_result.get())); } @@ -158,8 +161,8 @@ class FileIndexEvaluatorTest : public ::testing::Test { FieldType::INT, Literal(10)); ASSERT_OK_AND_ASSIGN( auto file_index_result, - FileIndexEvaluator::Evaluate(data_schema_, predicate, data_file_path_factory, - data_file_meta, fs_, pool_)); + FileIndexEvaluator::Evaluate(data_schema_, core_options_, predicate, + data_file_path_factory, data_file_meta, fs_, pool_)); ASSERT_TRUE(file_index_result->IsRemain().value()); ASSERT_FALSE(dynamic_cast(file_index_result.get())); } @@ -170,8 +173,8 @@ class FileIndexEvaluatorTest : public ::testing::Test { Literal(FieldType::STRING, "Lucy", 4)}); ASSERT_OK_AND_ASSIGN( auto file_index_result, - FileIndexEvaluator::Evaluate(data_schema_, predicate, data_file_path_factory, - data_file_meta, fs_, pool_)); + FileIndexEvaluator::Evaluate(data_schema_, core_options_, predicate, + data_file_path_factory, data_file_meta, fs_, pool_)); ASSERT_TRUE(file_index_result->IsRemain().value()); CheckResult(file_index_result, {0, 1, 4, 5, 7}); } @@ -182,8 +185,8 @@ class FileIndexEvaluatorTest : public ::testing::Test { Literal(FieldType::STRING, "Lucy", 4)}); ASSERT_OK_AND_ASSIGN( auto file_index_result, - FileIndexEvaluator::Evaluate(data_schema_, predicate, data_file_path_factory, - data_file_meta, fs_, pool_)); + FileIndexEvaluator::Evaluate(data_schema_, core_options_, predicate, + data_file_path_factory, data_file_meta, fs_, pool_)); ASSERT_TRUE(file_index_result->IsRemain().value()); CheckResult(file_index_result, {2, 3, 6}); } @@ -197,8 +200,8 @@ class FileIndexEvaluatorTest : public ::testing::Test { PredicateBuilder::And({f0_predicate, f1_predicate})); ASSERT_OK_AND_ASSIGN( auto file_index_result, - FileIndexEvaluator::Evaluate(data_schema_, predicate, data_file_path_factory, - data_file_meta, fs_, pool_)); + FileIndexEvaluator::Evaluate(data_schema_, core_options_, predicate, + data_file_path_factory, data_file_meta, fs_, pool_)); ASSERT_TRUE(file_index_result->IsRemain().value()); CheckResult(file_index_result, {7}); } @@ -212,8 +215,8 @@ class FileIndexEvaluatorTest : public ::testing::Test { PredicateBuilder::Or({f0_predicate, f1_predicate})); ASSERT_OK_AND_ASSIGN( auto file_index_result, - FileIndexEvaluator::Evaluate(data_schema_, predicate, data_file_path_factory, - data_file_meta, fs_, pool_)); + FileIndexEvaluator::Evaluate(data_schema_, core_options_, predicate, + data_file_path_factory, data_file_meta, fs_, pool_)); ASSERT_TRUE(file_index_result->IsRemain().value()); CheckResult(file_index_result, {0, 4, 6, 7}); } @@ -224,8 +227,8 @@ class FileIndexEvaluatorTest : public ::testing::Test { Literal(FieldType::STRING, "unknown", 7)); ASSERT_OK_AND_ASSIGN( auto file_index_result, - FileIndexEvaluator::Evaluate(data_schema_, predicate, data_file_path_factory, - data_file_meta, fs_, pool_)); + FileIndexEvaluator::Evaluate(data_schema_, core_options_, predicate, + data_file_path_factory, data_file_meta, fs_, pool_)); ASSERT_FALSE(file_index_result->IsRemain().value()); } { @@ -242,15 +245,15 @@ class FileIndexEvaluatorTest : public ::testing::Test { PredicateBuilder::And({f1_predicate, f2_predicate, f0_predicate})); ASSERT_OK_AND_ASSIGN( auto file_index_result, - FileIndexEvaluator::Evaluate(data_schema_, predicate, data_file_path_factory, - data_file_meta, fs_, pool_)); + FileIndexEvaluator::Evaluate(data_schema_, core_options_, predicate, + data_file_path_factory, data_file_meta, fs_, pool_)); ASSERT_FALSE(file_index_result->IsRemain().value()); } { // test no predicate ASSERT_OK_AND_ASSIGN( auto file_index_result, - FileIndexEvaluator::Evaluate(data_schema_, /*predicate=*/nullptr, + FileIndexEvaluator::Evaluate(data_schema_, core_options_, /*predicate=*/nullptr, data_file_path_factory, data_file_meta, fs_, pool_)); ASSERT_TRUE(file_index_result->IsRemain().value()); ASSERT_FALSE(dynamic_cast(file_index_result.get())); @@ -261,17 +264,18 @@ class FileIndexEvaluatorTest : public ::testing::Test { FieldType::DOUBLE, Literal(14.1)); ASSERT_OK_AND_ASSIGN( auto file_index_result, - FileIndexEvaluator::Evaluate(data_schema_, /*predicate=*/nullptr, + FileIndexEvaluator::Evaluate(data_schema_, core_options_, /*predicate=*/nullptr, data_file_path_factory, data_file_meta, fs_, pool_)); ASSERT_TRUE(file_index_result->IsRemain().value()); ASSERT_FALSE(dynamic_cast(file_index_result.get())); } } - private: + protected: std::shared_ptr pool_; std::shared_ptr fs_; std::shared_ptr data_schema_; + CoreOptions core_options_; }; TEST_F(FileIndexEvaluatorTest, TestEvaluateEmbeddingIndex) { @@ -391,9 +395,10 @@ TEST_F(FileIndexEvaluatorTest, TestTimestampType) { /*field_index=*/6, /*field_name=*/"ts_tz_micro", FieldType::TIMESTAMP, Literal(Timestamp(1745542602001l, 1000))); ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({in_predicate, greater_than})); - ASSERT_OK_AND_ASSIGN(auto file_index_result, FileIndexEvaluator::Evaluate( - data_schema, predicate, data_file_path_factory, - data_file_meta, /*file_system=*/fs_, pool_)); + ASSERT_OK_AND_ASSIGN( + auto file_index_result, + FileIndexEvaluator::Evaluate(data_schema, core_options_, predicate, data_file_path_factory, + data_file_meta, /*file_system=*/fs_, pool_)); CheckResult(file_index_result, {0, 6}); } @@ -416,8 +421,9 @@ TEST_F(FileIndexEvaluatorTest, TestInvalidEvaluate) { auto predicate = PredicateBuilder::IsNull(/*field_index=*/2, /*field_name=*/"f2", FieldType::INT); ASSERT_NOK_WITH_MSG( - FileIndexEvaluator::Evaluate(data_schema_, predicate, /*data_file_path_factory=*/nullptr, - data_file_meta, /*file_system=*/nullptr, pool_), + FileIndexEvaluator::Evaluate(data_schema_, core_options_, predicate, + /*data_file_path_factory=*/nullptr, data_file_meta, + /*file_system=*/nullptr, pool_), "read process for FileIndexEvaluator must have data_file_path_factory and file_system"); } diff --git a/src/paimon/core/io/file_index_options.cpp b/src/paimon/core/io/file_index_options.cpp index a9587363..811fb5d0 100644 --- a/src/paimon/core/io/file_index_options.cpp +++ b/src/paimon/core/io/file_index_options.cpp @@ -85,13 +85,17 @@ Result FileIndexOptions::FromCoreOptions(const CoreOptions& op } std::vector parts = StringUtils::Split(key.substr(kFileIndexPrefixLength), ".", /*ignore_empty=*/false); - if (parts.size() != 3) { + if (parts.size() < 3) { continue; } + std::string option_name = parts[2]; + for (size_t i = 3; i < parts.size(); ++i) { + option_name.append(".").append(parts[i]); + } bool found = false; for (FileIndexDefinition& definition : result.definitions_) { if (definition.index_type == parts[0] && definition.column_name == parts[1]) { - definition.options[parts[2]] = value; + definition.options[option_name] = value; found = true; break; } @@ -106,4 +110,21 @@ Result FileIndexOptions::FromCoreOptions(const CoreOptions& op return result; } +Result FileIndexOptions::FromMap( + const std::map& raw_options) { + PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, CoreOptions::FromMap(raw_options)); + return FromCoreOptions(core_options); +} + +const std::map& FileIndexOptions::GetIndexerOptions( + const std::string& column_name, const std::string& index_type) const { + static const std::map kEmptyOptions; + for (const FileIndexDefinition& definition : definitions_) { + if (definition.column_name == column_name && definition.index_type == index_type) { + return definition.options; + } + } + return kEmptyOptions; +} + } // namespace paimon diff --git a/src/paimon/core/io/file_index_options.h b/src/paimon/core/io/file_index_options.h index 7b7c019b..15cdeb19 100644 --- a/src/paimon/core/io/file_index_options.h +++ b/src/paimon/core/io/file_index_options.h @@ -36,11 +36,13 @@ struct FileIndexDefinition { std::map options; }; -/// Parsed write-side file index configuration. +/// Parsed file index configuration shared by the write and read paths. class FileIndexOptions { public: static Result FromCoreOptions(const CoreOptions& options); + static Result FromMap(const std::map& options); + const std::vector& Definitions() const { return definitions_; } @@ -53,6 +55,9 @@ class FileIndexOptions { return definitions_.empty(); } + const std::map& GetIndexerOptions( + const std::string& column_name, const std::string& index_type) const; + private: FileIndexOptions() = default; diff --git a/src/paimon/core/io/file_index_options_test.cpp b/src/paimon/core/io/file_index_options_test.cpp index 157203f9..fdfbc80d 100644 --- a/src/paimon/core/io/file_index_options_test.cpp +++ b/src/paimon/core/io/file_index_options_test.cpp @@ -55,4 +55,17 @@ TEST(FileIndexOptionsTest, TestNestedMapColumnSyntax) { "nested map columns is not supported"); } +TEST(FileIndexOptionsTest, TestPreserveDottedBackendOptionName) { + ASSERT_OK_AND_ASSIGN( + FileIndexOptions options, + FileIndexOptions::FromMap({{"file-index.lumina.columns", "f1"}, + {"file-index.lumina.f1.index.dimension", "128"}, + {"file-index.lumina.f1.distance.metric", "cosine"}})); + ASSERT_EQ(1, options.Definitions().size()); + EXPECT_EQ("128", options.Definitions()[0].options.at("index.dimension")); + EXPECT_EQ("cosine", options.Definitions()[0].options.at("distance.metric")); + EXPECT_EQ(options.Definitions()[0].options, options.GetIndexerOptions("f1", "lumina")); + EXPECT_TRUE(options.GetIndexerOptions("missing", "lumina").empty()); +} + } // namespace paimon::test diff --git a/src/paimon/core/operation/abstract_split_read.cpp b/src/paimon/core/operation/abstract_split_read.cpp index f1540843..f88910a6 100644 --- a/src/paimon/core/operation/abstract_split_read.cpp +++ b/src/paimon/core/operation/abstract_split_read.cpp @@ -33,6 +33,7 @@ #include "paimon/common/data/shredding/shredding_file_reader.h" #include "paimon/common/data/variant/variant_shredding_read_plan_factory.h" #include "paimon/common/data/variant/variant_type_utils.h" +#include "paimon/common/reader/complete_index_score_file_batch_reader.h" #include "paimon/common/reader/data_file_reader_factory.h" #include "paimon/common/reader/delegating_prefetch_reader.h" #include "paimon/common/reader/late_materializing_reader_builder.h" @@ -229,10 +230,12 @@ Result> AbstractSplitRead::CreateFieldMappingRe } const auto& predicate = field_mapping->non_partition_info.non_partition_filter; auto all_data_schema = DataField::ConvertDataFieldsToArrowSchema(data_schema->Fields()); - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr final_reader, - ApplyIndexAndDvReaderIfNeeded( - std::move(file_reader), file_meta, all_data_schema, read_schema, - predicate, dv_factory, row_ranges, data_file_path_factory)); + std::vector index_scores; + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr final_reader, + ApplyIndexAndDvReaderIfNeeded(std::move(file_reader), file_meta, all_data_schema, + read_schema, predicate, dv_factory, row_ranges, + data_file_path_factory, &index_scores)); if (!final_reader) { // file is skipped by index or dv return std::unique_ptr(); @@ -243,7 +246,13 @@ Result> AbstractSplitRead::CreateFieldMappingRe field_mapping_builder->GetReadFieldCount(), std::move(final_reader), partition, std::move(field_mapping), std::move(skip_map_selected_keys_filter_field_ids), arrow_pool_)); - return mapping_reader; + std::unique_ptr result = std::move(mapping_reader); + if (!index_scores.empty() && context_->GetVectorSearch() && + raw_read_schema_->GetFieldIndex(SpecialFields::IndexScore().Name()) >= 0) { + result = std::make_unique(std::move(result), + index_scores, arrow_pool_); + } + return result; } Result, std::set>> diff --git a/src/paimon/core/operation/abstract_split_read.h b/src/paimon/core/operation/abstract_split_read.h index 1e5f5644..65136c92 100644 --- a/src/paimon/core/operation/abstract_split_read.h +++ b/src/paimon/core/operation/abstract_split_read.h @@ -93,7 +93,8 @@ class AbstractSplitRead : public SplitRead { const std::shared_ptr& read_schema, const std::shared_ptr& predicate, DeletionVector::Factory dv_factory, const std::optional>& row_ranges, - const std::shared_ptr& data_file_path_factory) const = 0; + const std::shared_ptr& data_file_path_factory, + std::vector* index_scores) const = 0; // 1. project write cols to data schema // 2. add partition fields (if write cols not contain) diff --git a/src/paimon/core/operation/append_only_file_store_scan.cpp b/src/paimon/core/operation/append_only_file_store_scan.cpp index 435ece57..2bed6f70 100644 --- a/src/paimon/core/operation/append_only_file_store_scan.cpp +++ b/src/paimon/core/operation/append_only_file_store_scan.cpp @@ -155,9 +155,9 @@ Result AppendOnlyFileStoreScan::TestFileIndex( } assert(data_predicate); auto data_arrow_schema = DataField::ConvertDataFieldsToArrowSchema(data_schema->Fields()); - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr index_result, - FileIndexEvaluator::Evaluate(data_arrow_schema, data_predicate, meta, pool_)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr index_result, + FileIndexEvaluator::Evaluate(data_arrow_schema, core_options_, + data_predicate, meta, pool_)); return index_result->IsRemain(); } diff --git a/src/paimon/core/operation/data_evolution_split_read.cpp b/src/paimon/core/operation/data_evolution_split_read.cpp index e6edbf0a..5bf46f0d 100644 --- a/src/paimon/core/operation/data_evolution_split_read.cpp +++ b/src/paimon/core/operation/data_evolution_split_read.cpp @@ -39,9 +39,9 @@ #include "paimon/common/data/blob_utils.h" #include "paimon/common/data/blob_view_struct.h" #include "paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader.h" -#include "paimon/common/global_index/complete_index_score_batch_reader.h" #include "paimon/common/reader/blob_fallback_batch_reader.h" #include "paimon/common/reader/blob_view_resolving_batch_reader.h" +#include "paimon/common/reader/complete_index_score_batch_reader.h" #include "paimon/common/reader/complete_row_kind_batch_reader.h" #include "paimon/common/reader/concat_batch_reader.h" #include "paimon/common/table/special_fields.h" @@ -517,10 +517,10 @@ Result DataEvolutionSplitRead::SkipByFileIndex( } auto written_schema = DataField::ConvertDataFieldsToArrowSchema(written_fields); - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr index_result, - FileIndexEvaluator::Evaluate(written_schema, data_predicate, data_file_path_factory, - file, options_.GetFileSystem(), pool_)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr index_result, + FileIndexEvaluator::Evaluate(written_schema, options_, + data_predicate, data_file_path_factory, + file, options_.GetFileSystem(), pool_)); PAIMON_ASSIGN_OR_RAISE(bool is_remain, index_result->IsRemain()); if (!is_remain) { return true; @@ -534,13 +534,14 @@ Result> DataEvolutionSplitRead::ApplyIndexAndDv const std::shared_ptr& data_schema, const std::shared_ptr& read_schema, const std::shared_ptr& predicate, DeletionVector::Factory dv_factory, const std::optional>& row_ranges, - const std::shared_ptr& data_file_path_factory) const { + const std::shared_ptr& data_file_path_factory, + std::vector* index_scores) const { std::shared_ptr file_index_result; if (options_.FileIndexReadEnabled()) { PAIMON_ASSIGN_OR_RAISE( file_index_result, - FileIndexEvaluator::Evaluate(data_schema, predicate, data_file_path_factory, file, - options_.GetFileSystem(), pool_)); + FileIndexEvaluator::Evaluate(data_schema, options_, predicate, data_file_path_factory, + file, options_.GetFileSystem(), pool_)); PAIMON_ASSIGN_OR_RAISE(bool is_remain, file_index_result->IsRemain()); if (!is_remain) { return std::unique_ptr(); diff --git a/src/paimon/core/operation/data_evolution_split_read.h b/src/paimon/core/operation/data_evolution_split_read.h index 2568eae4..21133fb4 100644 --- a/src/paimon/core/operation/data_evolution_split_read.h +++ b/src/paimon/core/operation/data_evolution_split_read.h @@ -95,7 +95,8 @@ class DataEvolutionSplitRead : public AbstractSplitRead { const std::shared_ptr& read_schema, const std::shared_ptr& predicate, DeletionVector::Factory dv_factory, const std::optional>& row_ranges, - const std::shared_ptr& data_file_path_factory) const override; + const std::shared_ptr& data_file_path_factory, + std::vector* index_scores) const override; private: /// Files for partial field. diff --git a/src/paimon/core/operation/internal_read_context.cpp b/src/paimon/core/operation/internal_read_context.cpp index c1a07eca..2d7ef6d0 100644 --- a/src/paimon/core/operation/internal_read_context.cpp +++ b/src/paimon/core/operation/internal_read_context.cpp @@ -36,6 +36,8 @@ #include "paimon/core/options/map_storage_layout.h" #include "paimon/core/schema/arrow_schema_validator.h" #include "paimon/core/utils/nested_projection_utils.h" +#include "paimon/predicate/full_text_search.h" +#include "paimon/predicate/vector_search.h" #include "paimon/status.h" namespace paimon { @@ -147,7 +149,7 @@ Result> InternalReadContext::AlignReadFieldWithTab } std::optional InternalReadContext::TryResolveSpecialFieldById( - int32_t field_id, const CoreOptions& core_options) { + int32_t field_id, const CoreOptions& core_options, bool has_file_index_search) { if (field_id == SpecialFields::ValueKind().Id()) { return SpecialFields::ValueKind(); } @@ -164,7 +166,7 @@ std::optional InternalReadContext::TryResolveSpecialFieldById( return std::nullopt; } if (field_id == SpecialFields::IndexScore().Id()) { - if (core_options.DataEvolutionEnabled()) { + if (core_options.DataEvolutionEnabled() || has_file_index_search) { return SpecialFields::IndexScore(); } return std::nullopt; @@ -173,7 +175,7 @@ std::optional InternalReadContext::TryResolveSpecialFieldById( } std::optional InternalReadContext::TryResolveSpecialFieldByName( - const std::string& name, const CoreOptions& core_options) { + const std::string& name, const CoreOptions& core_options, bool has_file_index_search) { if (name == SpecialFields::ValueKind().Name()) { return SpecialFields::ValueKind(); } @@ -190,7 +192,7 @@ std::optional InternalReadContext::TryResolveSpecialFieldByName( return std::nullopt; } if (name == SpecialFields::IndexScore().Name()) { - if (core_options.DataEvolutionEnabled()) { + if (core_options.DataEvolutionEnabled() || has_file_index_search) { return SpecialFields::IndexScore(); } return std::nullopt; @@ -205,9 +207,31 @@ Result> InternalReadContext::Create( CoreOptions::FromMap(options, context->GetSpecificFileSystem(), context->GetFileSystemSchemeToIdentifierMap())); core_options.WithCache(context->GetCache()); + if (context->HasFileIndexSearch()) { + if (!core_options.FileIndexReadEnabled()) { + return Status::Invalid("File Index read must be enabled for file-local search"); + } + if (!table_schema->PrimaryKeys().empty()) { + return Status::NotImplemented( + "File Index search currently supports append-only tables only"); + } + if (core_options.DataEvolutionEnabled()) { + return Status::NotImplemented( + "File Index search does not support data-evolution reads yet"); + } + if (context->GetRealtimeContext()) { + return Status::NotImplemented("File Index search does not support real-time reads yet"); + } + const std::string& search_field = context->GetVectorSearch() + ? context->GetVectorSearch()->field_name + : context->GetFullTextSearch()->field_name; + PAIMON_ASSIGN_OR_RAISE([[maybe_unused]] DataField field, + table_schema->GetField(search_field)); + } // prepare read schema // Priority: projected_arrow_schema > read_field_ids > read_field_names const bool has_projected_read_schema = context->HasReadSchema(); + const bool has_file_index_search = context->HasFileIndexSearch(); std::vector read_data_fields; if (has_projected_read_schema) { // Nested column pruning path: user provided a read C ArrowSchema @@ -218,8 +242,8 @@ Result> InternalReadContext::Create( read_data_fields.reserve(read_schema->num_fields()); // Align special-field validation with read_field_ids/read_field_names branches. for (const auto& read_field : read_schema->fields()) { - if (auto resolved_special_field = - TryResolveSpecialFieldByName(read_field->name(), core_options)) { + if (auto resolved_special_field = TryResolveSpecialFieldByName( + read_field->name(), core_options, has_file_index_search)) { read_data_fields.push_back(*resolved_special_field); continue; } @@ -244,7 +268,8 @@ Result> InternalReadContext::Create( } else if (!context->GetReadFieldIds().empty()) { read_data_fields.reserve(context->GetReadFieldIds().size()); for (const auto& field_id : context->GetReadFieldIds()) { - if (auto resolved_special_field = TryResolveSpecialFieldById(field_id, core_options)) { + if (auto resolved_special_field = + TryResolveSpecialFieldById(field_id, core_options, has_file_index_search)) { read_data_fields.push_back(*resolved_special_field); continue; } @@ -254,7 +279,8 @@ Result> InternalReadContext::Create( } else if (!context->GetReadFieldNames().empty()) { read_data_fields.reserve(context->GetReadFieldNames().size()); for (const auto& name : context->GetReadFieldNames()) { - if (auto resolved_special_field = TryResolveSpecialFieldByName(name, core_options)) { + if (auto resolved_special_field = + TryResolveSpecialFieldByName(name, core_options, has_file_index_search)) { read_data_fields.push_back(*resolved_special_field); continue; } diff --git a/src/paimon/core/operation/internal_read_context.h b/src/paimon/core/operation/internal_read_context.h index ddbbb49f..564819fb 100644 --- a/src/paimon/core/operation/internal_read_context.h +++ b/src/paimon/core/operation/internal_read_context.h @@ -69,6 +69,15 @@ class InternalReadContext { const std::shared_ptr& GetPredicate() const { return read_context_->GetPredicate(); } + const std::shared_ptr& GetVectorSearch() const { + return read_context_->GetVectorSearch(); + } + const std::shared_ptr& GetFullTextSearch() const { + return read_context_->GetFullTextSearch(); + } + bool HasFileIndexSearch() const { + return read_context_->HasFileIndexSearch(); + } bool EnablePredicateFilter() const { return read_context_->EnablePredicateFilter(); } @@ -131,9 +140,11 @@ class InternalReadContext { const std::shared_ptr& arrow_pool); static std::optional TryResolveSpecialFieldById(int32_t field_id, - const CoreOptions& core_options); + const CoreOptions& core_options, + bool has_file_index_search); static std::optional TryResolveSpecialFieldByName(const std::string& name, - const CoreOptions& core_options); + const CoreOptions& core_options, + bool has_file_index_search); static Result> AlignReadFieldWithTableFieldIds( const std::shared_ptr& read_field, const std::shared_ptr& table_field); diff --git a/src/paimon/core/operation/internal_read_context_test.cpp b/src/paimon/core/operation/internal_read_context_test.cpp index 371d2d76..2ba5efbe 100644 --- a/src/paimon/core/operation/internal_read_context_test.cpp +++ b/src/paimon/core/operation/internal_read_context_test.cpp @@ -29,6 +29,7 @@ #include "paimon/data/shredding/map_shared_shredding_schema_utils.h" #include "paimon/defs.h" #include "paimon/fs/local/local_file_system.h" +#include "paimon/predicate/vector_search.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" @@ -152,6 +153,50 @@ TEST(InternalReadContext, TestReadWithRowTrackingAndScoreFields) { } } +TEST(InternalReadContext, TestFileIndexSearchIndexScoreProjection) { + std::string path = paimon::test::GetDataDir() + "/orc/append_09.db/append_09"; + SchemaManager schema_manager(std::make_shared(), path); + ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); + std::shared_ptr vector_search = std::make_shared( + "f3", /*limit=*/2, std::vector{1.0f}, nullptr, nullptr, + VectorSearch::DistanceType::EUCLIDEAN, std::map{}); + + { + ReadContextBuilder builder(path); + builder.SetReadFieldNames({"f0", "_INDEX_SCORE"}).SetVectorSearch(vector_search); + ASSERT_OK_AND_ASSIGN(auto read_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto internal_context, + InternalReadContext::Create(std::move(read_context), table_schema, + table_schema->Options())); + std::shared_ptr expected = DataField::ConvertDataFieldsToArrowSchema( + {DataField(0, arrow::field("f0", arrow::utf8())), SpecialFields::IndexScore()}); + EXPECT_TRUE(internal_context->GetReadSchema()->Equals(expected)); + } + { + ReadContextBuilder builder(path); + builder.SetReadFieldIds({0, SpecialFieldIds::INDEX_SCORE}).SetVectorSearch(vector_search); + ASSERT_OK_AND_ASSIGN(auto read_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto internal_context, + InternalReadContext::Create(std::move(read_context), table_schema, + table_schema->Options())); + EXPECT_EQ(arrow::Type::FLOAT, internal_context->GetReadSchema()->field(1)->type()->id()); + } + { + std::shared_ptr projected = arrow::schema( + {arrow::field("_INDEX_SCORE", arrow::float32()), arrow::field("f0", arrow::utf8())}); + auto c_schema = std::make_unique(); + ASSERT_TRUE(arrow::ExportSchema(*projected, c_schema.get()).ok()); + ReadContextBuilder builder(path); + builder.SetReadSchema(std::move(c_schema)).SetVectorSearch(vector_search); + ASSERT_OK_AND_ASSIGN(auto read_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto internal_context, + InternalReadContext::Create(std::move(read_context), table_schema, + table_schema->Options())); + EXPECT_EQ(std::vector({"_INDEX_SCORE", "f0"}), + internal_context->GetReadSchema()->field_names()); + } +} + TEST(InternalReadContext, TestReadWithValueKindField) { std::string path = paimon::test::GetDataDir() + "/orc/append_09.db/append_09"; ReadContextBuilder context_builder(path); diff --git a/src/paimon/core/operation/merge_file_split_read.cpp b/src/paimon/core/operation/merge_file_split_read.cpp index 7c6b1724..06ec0e62 100644 --- a/src/paimon/core/operation/merge_file_split_read.cpp +++ b/src/paimon/core/operation/merge_file_split_read.cpp @@ -487,7 +487,8 @@ Result> MergeFileSplitRead::ApplyIndexAndDvRead const std::shared_ptr& data_schema, const std::shared_ptr& read_schema, const std::shared_ptr& predicate, DeletionVector::Factory dv_factory, const std::optional>& ranges, - const std::shared_ptr& data_file_path_factory) const { + const std::shared_ptr& data_file_path_factory, + std::vector* index_scores) const { // merge read does not use index std::shared_ptr deletion_vector; if (dv_factory) { diff --git a/src/paimon/core/operation/merge_file_split_read.h b/src/paimon/core/operation/merge_file_split_read.h index 07b5e70b..3643b6f0 100644 --- a/src/paimon/core/operation/merge_file_split_read.h +++ b/src/paimon/core/operation/merge_file_split_read.h @@ -97,7 +97,8 @@ class MergeFileSplitRead : public AbstractSplitRead { const std::shared_ptr& read_schema, const std::shared_ptr& predicate, DeletionVector::Factory dv_factory, const std::optional>& ranges, - const std::shared_ptr& data_file_path_factory) const override; + const std::shared_ptr& data_file_path_factory, + std::vector* index_scores) const override; Result> CreateSortMergeReaderForSection( const std::vector& section, const BinaryRow& partition, diff --git a/src/paimon/core/operation/raw_file_split_read.cpp b/src/paimon/core/operation/raw_file_split_read.cpp index 8439ba19..c03a2970 100644 --- a/src/paimon/core/operation/raw_file_split_read.cpp +++ b/src/paimon/core/operation/raw_file_split_read.cpp @@ -44,6 +44,7 @@ #include "paimon/core/utils/file_store_path_factory.h" #include "paimon/file_index/bitmap_index_result.h" #include "paimon/file_index/file_index_result.h" +#include "paimon/file_index/scored_file_index_result.h" #include "paimon/memory/memory_pool.h" #include "paimon/reader/file_batch_reader.h" #include "paimon/status.h" @@ -194,13 +195,70 @@ Result> RawFileSplitRead::ApplyIndexAndDvReader const std::shared_ptr& data_schema, const std::shared_ptr& read_schema, const std::shared_ptr& predicate, DeletionVector::Factory dv_factory, const std::optional>& ranges, - const std::shared_ptr& data_file_path_factory) const { + const std::shared_ptr& data_file_path_factory, + std::vector* index_scores) const { + if (context_->HasFileIndexSearch()) { + if (predicate) { + return Status::NotImplemented( + "File Index search does not support combining with read predicates yet"); + } + if (ranges) { + return Status::NotImplemented( + "File Index search does not support indexed split row ranges yet"); + } + if (dv_factory) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr deletion_vector, + dv_factory(file->file_name)); + if (deletion_vector && !deletion_vector->IsEmpty()) { + return Status::NotImplemented( + "File Index search does not support deletion vectors yet"); + } + } + + std::optional search_selection; + if (context_->GetVectorSearch()) { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr search_result, + FileIndexEvaluator::EvaluateVectorSearch( + data_schema, options_, context_->GetVectorSearch(), data_file_path_factory, + file, options_.GetFileSystem(), pool_)); + if (search_result->IsEmpty()) { + return std::unique_ptr(); + } + search_selection = search_result->GetRowPositions(); + *index_scores = search_result->GetScores(); + } else { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr search_result, + FileIndexEvaluator::EvaluateFullTextSearch( + data_schema, options_, context_->GetFullTextSearch(), data_file_path_factory, + file, options_.GetFileSystem(), pool_)); + std::shared_ptr bitmap_result = + std::dynamic_pointer_cast(search_result); + if (!bitmap_result) { + return std::unique_ptr(); + } + PAIMON_ASSIGN_OR_RAISE(const RoaringBitmap32* bitmap, bitmap_result->GetBitmap()); + search_selection = *bitmap; + } + + ::ArrowSchema c_read_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*read_schema, &c_read_schema)); + PAIMON_RETURN_NOT_OK( + file_reader->SetReadSchema(&c_read_schema, /*predicate=*/nullptr, search_selection)); + if (!file_reader->SupportPreciseBitmapSelection()) { + return std::make_unique( + std::move(file_reader), std::move(search_selection).value()); + } + return std::move(file_reader); + } + std::shared_ptr file_index_result; if (options_.FileIndexReadEnabled()) { PAIMON_ASSIGN_OR_RAISE( file_index_result, - FileIndexEvaluator::Evaluate(data_schema, predicate, data_file_path_factory, file, - options_.GetFileSystem(), pool_)); + FileIndexEvaluator::Evaluate(data_schema, options_, predicate, data_file_path_factory, + file, options_.GetFileSystem(), pool_)); PAIMON_ASSIGN_OR_RAISE(bool is_remain, file_index_result->IsRemain()); if (!is_remain) { return std::unique_ptr(); diff --git a/src/paimon/core/operation/raw_file_split_read.h b/src/paimon/core/operation/raw_file_split_read.h index 646f24ac..2d0e8d01 100644 --- a/src/paimon/core/operation/raw_file_split_read.h +++ b/src/paimon/core/operation/raw_file_split_read.h @@ -90,7 +90,8 @@ class RawFileSplitRead : public AbstractSplitRead { const std::shared_ptr& read_schema, const std::shared_ptr& predicate, DeletionVector::Factory dv_factory, const std::optional>& ranges, - const std::shared_ptr& data_file_path_factory) const override; + const std::shared_ptr& data_file_path_factory, + std::vector* index_scores) const override; }; } // namespace paimon diff --git a/src/paimon/core/operation/raw_file_split_read_test.cpp b/src/paimon/core/operation/raw_file_split_read_test.cpp index 90528004..5d72e2c2 100644 --- a/src/paimon/core/operation/raw_file_split_read_test.cpp +++ b/src/paimon/core/operation/raw_file_split_read_test.cpp @@ -18,7 +18,9 @@ #include "paimon/core/operation/raw_file_split_read.h" +#include #include +#include #include #include #include @@ -27,6 +29,8 @@ #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" #include "paimon/common/data/binary_row.h" +#include "paimon/common/io/byte_array_output_stream.h" +#include "paimon/common/io/memory_segment_output_stream.h" #include "paimon/common/reader/concat_batch_reader.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/mem_utils.h" @@ -42,9 +46,16 @@ #include "paimon/core/utils/file_store_path_factory.h" #include "paimon/data/timestamp.h" #include "paimon/executor.h" +#include "paimon/factories/factory.h" +#include "paimon/file_index/file_index_format.h" +#include "paimon/file_index/file_index_reader.h" +#include "paimon/file_index/file_indexer.h" +#include "paimon/file_index/file_indexer_factory.h" +#include "paimon/file_index/scored_file_index_result.h" #include "paimon/format/file_format.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/memory/memory_pool.h" +#include "paimon/predicate/vector_search.h" #include "paimon/read_context.h" #include "paimon/status.h" #include "paimon/table/source/data_split.h" @@ -53,6 +64,62 @@ #include "paimon/testing/utils/testharness.h" namespace paimon::test { +namespace { + +constexpr char kRawVectorSearchIndexType[] = "raw-vector-search-test"; + +class RawVectorSearchTestReader final : public FileIndexReader { + public: + Result> VisitVectorSearch( + const std::shared_ptr&) override { + return ScoredFileIndexResult::Create(RoaringBitmap32::From({0, 2}), {0.25f, 0.75f}); + } +}; + +class RawVectorSearchTestIndexer final : public FileIndexer { + public: + Result> CreateReader( + ::ArrowSchema*, int32_t, int32_t, const std::shared_ptr&, + const std::shared_ptr&) const override { + return std::make_shared(); + } + + Result> CreateWriter( + ::ArrowSchema*, const std::shared_ptr&) const override { + return Status::NotImplemented("Raw vector search test index is read-only"); + } +}; + +class RawVectorSearchTestFactory final : public FileIndexerFactory { + public: + const char* Identifier() const override { + return kRawVectorSearchIndexType; + } + + Result> Create( + const std::map&) const override { + return std::make_unique(); + } +}; + +REGISTER_PAIMON_FACTORY(RawVectorSearchTestFactory); + +Result> MakeRawVectorSearchIndex(const std::shared_ptr& pool) { + FileIndexFormat::ColumnIndexes indexes; + indexes["f3"][kRawVectorSearchIndexType] = std::make_shared("index", pool.get()); + auto segment_output = std::make_unique( + MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool); + auto output = std::make_shared(std::move(segment_output)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr writer, + FileIndexFormat::CreateWriter(output, pool)); + PAIMON_RETURN_NOT_OK(writer->WriteColumnIndexes(indexes)); + PAIMON_RETURN_NOT_OK(writer->Close()); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr bytes, output->Finish(pool.get())); + return bytes; +} + +} // namespace + class RawFileSplitReadTest : public ::testing::Test { std::vector> PrepareDataSplits() const { auto meta1 = std::make_shared( @@ -380,6 +447,81 @@ TEST_F(RawFileSplitReadTest, TestCreateReaderWithNonPartitionWithReserveSequence CheckReadResult(read_schema, expected_array); } +TEST_F(RawFileSplitReadTest, TestVectorSearchScoreProjection) { + std::string path = paimon::test::GetDataDir() + + "/orc/multi_partition_append_table.db/multi_partition_append_table"; + std::shared_ptr vector_search = std::make_shared( + "f3", /*limit=*/2, std::vector{1.0f}, nullptr, nullptr, + VectorSearch::DistanceType::EUCLIDEAN, std::map{}); + ReadContextBuilder context_builder(path); + context_builder.SetReadFieldNames({"_INDEX_SCORE", "f0"}).SetVectorSearch(vector_search); + ASSERT_OK_AND_ASSIGN(std::shared_ptr read_context, context_builder.Finish()); + SchemaManager schema_manager(std::make_shared(), path); + ASSERT_OK_AND_ASSIGN(auto table_schema, schema_manager.ReadSchema(0)); + ASSERT_OK_AND_ASSIGN( + auto internal_context, + InternalReadContext::Create(read_context, table_schema, table_schema->Options())); + + auto data_splits = PrepareDataSplits(); + auto data_split = std::dynamic_pointer_cast(data_splits.front()); + ASSERT_TRUE(data_split); + ASSERT_OK_AND_ASSIGN(data_split->DataFiles().front()->embedded_index, + MakeRawVectorSearchIndex(pool_)); + + const CoreOptions& core_options = internal_context->GetCoreOptions(); + std::shared_ptr data_schema = + DataField::ConvertDataFieldsToArrowSchema(table_schema->Fields()); + ASSERT_OK_AND_ASSIGN(std::vector external_paths, + core_options.CreateExternalPaths()); + ASSERT_OK_AND_ASSIGN(std::optional global_index_external_path, + core_options.CreateGlobalIndexExternalPath()); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr path_factory, + FileStorePathFactory::Create( + path, data_schema, table_schema->PartitionKeys(), + core_options.GetPartitionDefaultName(), core_options.GetFileFormat()->Identifier(), + core_options.DataFilePrefix(), core_options.LegacyPartitionNameEnabled(), + external_paths, global_index_external_path, core_options.IndexFileInDataFileDir(), + pool_)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, + CreateDefaultExecutor(/*thread_count=*/2)); + auto split_read = std::make_unique(path_factory, std::move(internal_context), + pool_, executor); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, split_read->CreateReader(data_split)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + ReadResultCollector::CollectResult(std::move(reader))); + + std::shared_ptr expected; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( + arrow::struct_({arrow::field("_VALUE_KIND", arrow::int8()), + arrow::field("_INDEX_SCORE", arrow::float32()), + arrow::field("f0", arrow::utf8())}), + {R"([[0, 0.25, "Bob"], [0, 0.75, "Tony"]])"}, &expected) + .ok()); + ASSERT_TRUE(result->Equals(expected)); + + ReadContextBuilder without_score_builder(path); + without_score_builder.SetReadFieldNames({"f0"}).SetVectorSearch(vector_search); + ASSERT_OK_AND_ASSIGN(std::shared_ptr without_score_context, + without_score_builder.Finish()); + ASSERT_OK_AND_ASSIGN( + auto without_score_internal_context, + InternalReadContext::Create(without_score_context, table_schema, table_schema->Options())); + auto without_score_split_read = std::make_unique( + path_factory, std::move(without_score_internal_context), pool_, executor); + ASSERT_OK_AND_ASSIGN(std::unique_ptr without_score_reader, + without_score_split_read->CreateReader(data_split)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr without_score_result, + ReadResultCollector::CollectResult(std::move(without_score_reader))); + std::shared_ptr without_score_expected; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( + arrow::struct_({arrow::field("_VALUE_KIND", arrow::int8()), + arrow::field("f0", arrow::utf8())}), + {R"([[0, "Bob"], [0, "Tony"]])"}, &without_score_expected) + .ok()); + ASSERT_TRUE(without_score_result->Equals(without_score_expected)); +} + TEST_F(RawFileSplitReadTest, TestEmptyPlan) { std::string path = paimon::test::GetDataDir() + "/orc/multi_partition_append_table.db/" diff --git a/src/paimon/core/operation/read_context.cpp b/src/paimon/core/operation/read_context.cpp index 7545c543..dad8d686 100644 --- a/src/paimon/core/operation/read_context.cpp +++ b/src/paimon/core/operation/read_context.cpp @@ -26,6 +26,8 @@ #include "paimon/core/utils/branch_manager.h" #include "paimon/executor.h" #include "paimon/memory/memory_pool.h" +#include "paimon/predicate/full_text_search.h" +#include "paimon/predicate/vector_search.h" #include "paimon/status.h" #include "paimon/table/format/format_table.h" @@ -35,8 +37,9 @@ class Predicate; ReadContext::ReadContext( const std::string& path, const std::string& branch, const std::vector& read_field_names, const std::vector& read_field_ids, - const std::shared_ptr& predicate, bool enable_predicate_filter, bool enable_prefetch, - bool enable_late_materializing, uint32_t prefetch_batch_count, + const std::shared_ptr& predicate, const std::shared_ptr& vector_search, + const std::shared_ptr& full_text_search, bool enable_predicate_filter, + bool enable_prefetch, bool enable_late_materializing, uint32_t prefetch_batch_count, uint32_t prefetch_max_parallel_num, bool enable_multi_thread_row_to_batch, uint32_t row_to_batch_thread_number, const std::optional& table_schema, const std::shared_ptr& memory_pool, const std::shared_ptr& executor, @@ -51,6 +54,8 @@ ReadContext::ReadContext( read_field_names_(read_field_names), read_field_ids_(read_field_ids), predicate_(predicate), + vector_search_(vector_search), + full_text_search_(full_text_search), enable_predicate_filter_(enable_predicate_filter), enable_prefetch_(enable_prefetch), enable_late_materializing_(enable_late_materializing), @@ -100,6 +105,8 @@ class ReadContextBuilder::Impl { fs_scheme_to_identifier_map_.clear(); options_.clear(); predicate_.reset(); + vector_search_.reset(); + full_text_search_.reset(); enable_predicate_filter_ = false; enable_prefetch_ = false; enable_late_materializing_ = false; @@ -131,6 +138,8 @@ class ReadContextBuilder::Impl { std::map fs_scheme_to_identifier_map_; std::map options_; std::shared_ptr predicate_; + std::shared_ptr vector_search_; + std::shared_ptr full_text_search_; bool enable_predicate_filter_ = false; bool enable_prefetch_ = false; bool enable_late_materializing_ = false; @@ -203,6 +212,18 @@ ReadContextBuilder& ReadContextBuilder::SetPredicate(const std::shared_ptr& vector_search) { + impl_->vector_search_ = vector_search; + return *this; +} + +ReadContextBuilder& ReadContextBuilder::SetFullTextSearch( + const std::shared_ptr& full_text_search) { + impl_->full_text_search_ = full_text_search; + return *this; +} + ReadContextBuilder& ReadContextBuilder::EnablePredicateFilter(bool enabled) { impl_->enable_predicate_filter_ = enabled; return *this; @@ -298,6 +319,9 @@ ReadContextBuilder& ReadContextBuilder::WithCache(const std::shared_ptr& } Result> ReadContextBuilder::Finish() { + if (impl_->vector_search_ && impl_->full_text_search_) { + return Status::Invalid("VectorSearch and FullTextSearch cannot be configured together"); + } if (impl_->built_from_format_table_ && impl_->format_table_ == nullptr) { return Status::Invalid("cannot read with null format table"); } @@ -352,13 +376,14 @@ Result> ReadContextBuilder::Finish() { } auto ctx = std::make_unique( impl_->path_, impl_->branch_, impl_->read_field_names_, impl_->read_field_ids_, - impl_->predicate_, impl_->enable_predicate_filter_, impl_->enable_prefetch_, - impl_->enable_late_materializing_, impl_->prefetch_batch_count_, - impl_->prefetch_max_parallel_num_, impl_->enable_multi_thread_row_to_batch_, - impl_->row_to_batch_thread_number_, impl_->table_schema_, impl_->memory_pool_, - impl_->executor_, impl_->specific_file_system_, impl_->fs_scheme_to_identifier_map_, - impl_->realtime_context_, impl_->options_, impl_->read_ahead_cache_enabled_, - impl_->cache_config_, impl_->cache_, impl_->format_table_, impl_->warmup_level_); + impl_->predicate_, impl_->vector_search_, impl_->full_text_search_, + impl_->enable_predicate_filter_, impl_->enable_prefetch_, impl_->enable_late_materializing_, + impl_->prefetch_batch_count_, impl_->prefetch_max_parallel_num_, + impl_->enable_multi_thread_row_to_batch_, impl_->row_to_batch_thread_number_, + impl_->table_schema_, impl_->memory_pool_, impl_->executor_, impl_->specific_file_system_, + impl_->fs_scheme_to_identifier_map_, impl_->realtime_context_, impl_->options_, + impl_->read_ahead_cache_enabled_, impl_->cache_config_, impl_->cache_, impl_->format_table_, + impl_->warmup_level_); if (impl_->read_schema_ && impl_->read_schema_->release) { ctx->SetReadSchema(std::move(impl_->read_schema_)); } diff --git a/src/paimon/core/operation/read_context_test.cpp b/src/paimon/core/operation/read_context_test.cpp index 572fd5a5..924c6301 100644 --- a/src/paimon/core/operation/read_context_test.cpp +++ b/src/paimon/core/operation/read_context_test.cpp @@ -27,7 +27,9 @@ #include "paimon/defs.h" #include "paimon/executor.h" #include "paimon/memory/memory_pool.h" +#include "paimon/predicate/full_text_search.h" #include "paimon/predicate/predicate_builder.h" +#include "paimon/predicate/vector_search.h" #include "paimon/status.h" #include "paimon/testing/mock/mock_file_system.h" #include "paimon/testing/utils/testharness.h" @@ -43,6 +45,7 @@ TEST(ReadContextTest, TestDefaultValue) { ASSERT_TRUE(ctx->GetReadFieldIds().empty()); ASSERT_TRUE(ctx->GetOptions().empty()); ASSERT_FALSE(ctx->GetPredicate()); + ASSERT_FALSE(ctx->HasFileIndexSearch()); ASSERT_FALSE(ctx->EnablePredicateFilter()); ASSERT_FALSE(ctx->EnablePrefetch()); ASSERT_TRUE(ctx->ReadAheadCacheEnabled()); @@ -217,4 +220,39 @@ TEST(ReadContextTest, TestSetInvalidReadSchemaIgnored) { ASSERT_EQ(ctx->GetReadSchema(), nullptr); } +TEST(ReadContextTest, TestFileIndexSearchConfigurationAndReset) { + std::shared_ptr vector_search = std::make_shared( + "embedding", /*limit=*/3, std::vector{1.0f, 2.0f}, nullptr, nullptr, + VectorSearch::DistanceType::COSINE, std::map{}); + ReadContextBuilder builder("table_root_path"); + ASSERT_EQ(&builder, &builder.SetVectorSearch(vector_search)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr vector_context, builder.Finish()); + EXPECT_EQ(vector_search, vector_context->GetVectorSearch()); + EXPECT_FALSE(vector_context->GetFullTextSearch()); + EXPECT_TRUE(vector_context->HasFileIndexSearch()); + + std::shared_ptr full_text_search = std::make_shared( + "body", /*limit=*/5, "paimon", FullTextSearch::SearchType::MATCH_ANY, std::nullopt); + ASSERT_EQ(&builder, &builder.SetFullTextSearch(full_text_search)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr full_text_context, builder.Finish()); + EXPECT_FALSE(full_text_context->GetVectorSearch()); + EXPECT_EQ(full_text_search, full_text_context->GetFullTextSearch()); + EXPECT_TRUE(full_text_context->HasFileIndexSearch()); + + // Finish resets search configuration just like every other per-read option. + ASSERT_OK_AND_ASSIGN(std::unique_ptr reset_context, builder.Finish()); + EXPECT_FALSE(reset_context->HasFileIndexSearch()); +} + +TEST(ReadContextTest, TestRejectConflictingFileIndexSearch) { + std::shared_ptr vector_search = + std::make_shared("embedding", /*limit=*/3, std::vector{1.0f}, nullptr, + nullptr, std::nullopt, std::map{}); + std::shared_ptr full_text_search = std::make_shared( + "body", std::nullopt, "paimon", FullTextSearch::SearchType::MATCH_ANY, std::nullopt); + ReadContextBuilder conflicting("table_root_path"); + conflicting.SetVectorSearch(vector_search).SetFullTextSearch(full_text_search); + ASSERT_NOK_WITH_MSG(conflicting.Finish(), "cannot be configured together"); +} + } // namespace paimon::test diff --git a/src/paimon/core/table/source/table_read.cpp b/src/paimon/core/table/source/table_read.cpp index 557f2fd5..9b37f625 100644 --- a/src/paimon/core/table/source/table_read.cpp +++ b/src/paimon/core/table/source/table_read.cpp @@ -182,6 +182,9 @@ Result> TableRead::Create(std::unique_ptrGetFormatTable() != nullptr) { + if (context->HasFileIndexSearch()) { + return Status::NotImplemented("File Index search is not supported for format tables"); + } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr format_table, FormatTable::Copy(context->GetFormatTable(), context->GetOptions())); return NewFormatTableRead(format_table, context); @@ -193,6 +196,9 @@ Result> TableRead::Create(std::unique_ptr system_table_path, SystemTableLoader::TryParsePath(context->GetPath())); if (system_table_path) { + if (context->HasFileIndexSearch()) { + return Status::NotImplemented("File Index search is not supported for system tables"); + } PAIMON_ASSIGN_OR_RAISE( std::shared_ptr system_table, SystemTableLoader::LoadFromPath(tmp_core_options.GetFileSystem(), context->GetPath(), @@ -211,6 +217,9 @@ Result> TableRead::Create(std::unique_ptrGetSpecificTableSchema(), /*schema_manager=*/nullptr, &latest_schema)); if (format_table != nullptr) { + if (context->HasFileIndexSearch()) { + return Status::NotImplemented("File Index search is not supported for format tables"); + } return NewFormatTableRead(format_table, context); } diff --git a/src/paimon/global_index/lumina/lumina_global_index.cpp b/src/paimon/global_index/lumina/lumina_global_index.cpp deleted file mode 100644 index 98585f21..00000000 --- a/src/paimon/global_index/lumina/lumina_global_index.cpp +++ /dev/null @@ -1,1024 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "paimon/global_index/lumina/lumina_global_index.h" - -#include -#include -#include -#include -#include - -#include "arrow/c/bridge.h" -#include "arrow/c/helpers.h" -#include "lumina/api/Dataset.h" -#include "lumina/api/LuminaBuilder.h" -#include "lumina/api/LuminaSearcher.h" -#include "lumina/api/OptionsNormalize.h" -#include "lumina/core/Constants.h" -#include "lumina/core/Status.h" -#include "lumina/core/Types.h" -#include "lumina/extensions/experimental/BuildCombinedExtensionV0.h" -#include "paimon/common/global_index/global_index_utils.h" -#include "paimon/common/utils/checked_cast.h" -#include "paimon/common/utils/options_utils.h" -#include "paimon/common/utils/rapidjson_util.h" -#include "paimon/common/utils/string_utils.h" -#include "paimon/global_index/bitmap_scored_global_index_result.h" -#include "paimon/global_index/lumina/lumina_file_reader.h" -#include "paimon/global_index/lumina/lumina_file_writer.h" -#include "paimon/global_index/lumina/lumina_utils.h" -#include "paimon/predicate/compound_predicate.h" -#include "paimon/predicate/leaf_predicate.h" -#include "rapidjson/document.h" -namespace paimon::lumina { -#define CHECK_NOT_NULL(pointer, error_msg) \ - do { \ - if (!(pointer)) { \ - return Status::Invalid(error_msg); \ - } \ - } while (0) - -namespace { -using TagDimensionData = ::lumina::extensions::experimental::TagDimensionData; -using TagFilter = ::lumina::extensions::experimental::TagFilter; -using TagValue = ::lumina::extensions::experimental::TagValue; -using TagValues = ::lumina::extensions::experimental::TagValues; - -Result GetRequiredStringMember(const rapidjson::Value& obj, - const std::string& field_name, - const std::string& tag_label) { - auto iter = obj.FindMember(field_name.c_str()); - if (iter == obj.MemberEnd()) { - return Status::Invalid( - fmt::format("lumina tag_schema {} missing required field: {}", tag_label, field_name)); - } - if (!iter->value.IsString()) { - return Status::Invalid( - fmt::format("lumina tag_schema {} field {} must be string", tag_label, field_name)); - } - return std::string(iter->value.GetString(), iter->value.GetStringLength()); -} - -Result ParseTagField(const rapidjson::Value& obj, const std::string& tag_label) { - if (!obj.IsObject()) { - return Status::Invalid(fmt::format("lumina tag_schema {} must be object", tag_label)); - } - if (obj.MemberCount() != 3) { - return Status::Invalid(fmt::format( - "lumina tag_schema {} must have exactly 3 fields: key_name, type, value_type", - tag_label)); - } - - PAIMON_ASSIGN_OR_RAISE( - std::string key_name, - GetRequiredStringMember(obj, std::string(::lumina::core::kExtensionTagKName), tag_label)); - PAIMON_ASSIGN_OR_RAISE( - std::string type, - GetRequiredStringMember(obj, std::string(::lumina::core::kExtensionTagType), tag_label)); - PAIMON_ASSIGN_OR_RAISE( - std::string value_type, - GetRequiredStringMember(obj, std::string(::lumina::core::kExtensionTagVType), tag_label)); - - if (key_name.empty()) { - return Status::Invalid( - fmt::format("lumina tag_schema {} key_name must not be empty", tag_label)); - } - LuminaTagField::Type parsed_type; - if (type == std::string(::lumina::core::kExtensionTagTypeEnum)) { - parsed_type = LuminaTagField::Type::ENUM; - } else if (type == std::string(::lumina::core::kExtensionTagTypeRange)) { - parsed_type = LuminaTagField::Type::RANGE; - } else { - return Status::Invalid( - fmt::format("lumina tag_schema {} has unsupported type: {}", tag_label, type)); - } - - LuminaTagField::ValueType parsed_value_type; - if (value_type == std::string(::lumina::core::kExtensionTagVTypeInt32)) { - parsed_value_type = LuminaTagField::ValueType::INT32; - } else if (value_type == std::string(::lumina::core::kExtensionTagVTypeInt64)) { - parsed_value_type = LuminaTagField::ValueType::INT64; - } else if (value_type == std::string(::lumina::core::kExtensionTagVTypeFloat)) { - parsed_value_type = LuminaTagField::ValueType::FLOAT; - } else if (value_type == std::string(::lumina::core::kExtensionTagVTypeDouble)) { - parsed_value_type = LuminaTagField::ValueType::DOUBLE; - } else if (value_type == std::string(::lumina::core::kExtensionTagVTypeString)) { - parsed_value_type = LuminaTagField::ValueType::STRING; - } else { - return Status::Invalid(fmt::format("lumina tag_schema {} has unsupported value_type: {}", - tag_label, value_type)); - } - return LuminaTagField{key_name, parsed_type, parsed_value_type}; -} - -Status ValidateTagArrowType(const LuminaTagField& tag_field, - const std::shared_ptr& field_type) { - auto value_type = field_type; - if (auto list_type = std::dynamic_pointer_cast(field_type)) { - value_type = list_type->value_type(); - } - - bool compatible = false; - switch (tag_field.value_type) { - case LuminaTagField::ValueType::INT32: - compatible = value_type->id() == arrow::Type::INT8 || - value_type->id() == arrow::Type::INT16 || - value_type->id() == arrow::Type::INT32; - break; - case LuminaTagField::ValueType::INT64: - compatible = value_type->id() == arrow::Type::INT64; - break; - case LuminaTagField::ValueType::FLOAT: - compatible = value_type->id() == arrow::Type::FLOAT; - break; - case LuminaTagField::ValueType::DOUBLE: - compatible = value_type->id() == arrow::Type::DOUBLE; - break; - case LuminaTagField::ValueType::STRING: - compatible = value_type->id() == arrow::Type::STRING; - break; - } - if (!compatible) { - return Status::Invalid( - fmt::format("lumina tag field {} type {} is not compatible with tag_schema value_type", - tag_field.name, field_type->ToString())); - } - return Status::OK(); -} - -template -void AppendPrimitiveTagValue(const std::shared_ptr& array, int64_t index, - std::vector* values) { - values->push_back( - static_cast(checked_cast(array.get())->Value(index))); -} - -template -Status AppendTagValue(const std::shared_ptr& array, int64_t index, - std::vector* values) { - if (array->IsNull(index)) { - return Status::OK(); - } - - auto validate_array_type = [&](arrow::Type::type expected_type, - const char* value_type_name) -> Status { - if (array->type_id() != expected_type) { - return Status::Invalid(fmt::format("lumina {} tag field has unsupported arrow type {}", - value_type_name, array->type()->ToString())); - } - return Status::OK(); - }; - - if constexpr (std::is_same_v) { - switch (array->type_id()) { - case arrow::Type::INT8: - AppendPrimitiveTagValue(array, index, values); - break; - case arrow::Type::INT16: - AppendPrimitiveTagValue(array, index, values); - break; - case arrow::Type::INT32: - AppendPrimitiveTagValue(array, index, values); - break; - default: - return Status::Invalid( - fmt::format("lumina integer tag field has unsupported arrow type {}", - array->type()->ToString())); - } - } else if constexpr (std::is_same_v) { - PAIMON_RETURN_NOT_OK(validate_array_type(arrow::Type::INT64, "int64")); - AppendPrimitiveTagValue(array, index, values); - } else if constexpr (std::is_same_v) { - PAIMON_RETURN_NOT_OK(validate_array_type(arrow::Type::FLOAT, "float")); - AppendPrimitiveTagValue(array, index, values); - } else if constexpr (std::is_same_v) { - PAIMON_RETURN_NOT_OK(validate_array_type(arrow::Type::DOUBLE, "double")); - AppendPrimitiveTagValue(array, index, values); - } else if constexpr (std::is_same_v) { - PAIMON_RETURN_NOT_OK(validate_array_type(arrow::Type::STRING, "string")); - auto string_array = checked_cast(array.get()); - auto view = string_array->GetView(index); - values->emplace_back(view.data(), view.size()); - } else { - return Status::Invalid("lumina tag field has unsupported value type"); - } - return Status::OK(); -} - -template -Status ExtractTagValues(const std::shared_ptr& field_array, int64_t segment_start, - int64_t segment_len, std::vector>* values) { - values->resize(segment_len); - auto list_array = std::dynamic_pointer_cast(field_array); - if (list_array) { - auto child_values = list_array->values(); - for (int64_t i = 0; i < segment_len; i++) { - int64_t row = segment_start + i; - if (list_array->IsNull(row)) { - continue; - } - auto value_start = list_array->value_offset(row); - auto value_end = list_array->value_offset(row + 1); - auto& row_values = (*values)[i]; - row_values.reserve(value_end - value_start); - for (int64_t value_index = value_start; value_index < value_end; value_index++) { - PAIMON_RETURN_NOT_OK(AppendTagValue(child_values, value_index, &row_values)); - } - } - return Status::OK(); - } - - for (int64_t i = 0; i < segment_len; i++) { - PAIMON_RETURN_NOT_OK(AppendTagValue(field_array, segment_start + i, &(*values)[i])); - } - return Status::OK(); -} - -Result LiteralToTagValue(const Literal& literal) { - if (literal.IsNull()) { - return Status::Invalid("lumina tag predicate does not support null literal"); - } - switch (literal.GetType()) { - case FieldType::TINYINT: - return TagValue(static_cast(literal.GetValue())); - case FieldType::SMALLINT: - return TagValue(static_cast(literal.GetValue())); - case FieldType::INT: - return TagValue(literal.GetValue()); - case FieldType::BIGINT: - return TagValue(literal.GetValue()); - case FieldType::FLOAT: - return TagValue(literal.GetValue()); - case FieldType::DOUBLE: - return TagValue(literal.GetValue()); - case FieldType::STRING: - return TagValue(literal.GetValue()); - default: - return Status::Invalid( - fmt::format("lumina tag predicate does not support literal type {}", - static_cast(literal.GetType()))); - } -} - -Result GetSingleLiteral(const std::vector& literals, - const std::string& function_name) { - if (literals.size() != 1) { - return Status::Invalid( - fmt::format("lumina tag {} predicate requires one literal", function_name)); - } - return &literals[0]; -} - -Result LiteralsToTagValues(const std::vector& literals) { - if (literals.empty()) { - return Status::Invalid("lumina tag predicate IN requires at least one literal"); - } - - switch (literals[0].GetType()) { - case FieldType::TINYINT: - case FieldType::SMALLINT: - case FieldType::INT: { - std::vector values; - values.reserve(literals.size()); - for (const auto& literal : literals) { - PAIMON_ASSIGN_OR_RAISE(TagValue value, LiteralToTagValue(literal)); - auto typed_value = std::get_if(&value); - CHECK_NOT_NULL(typed_value, - "lumina tag predicate IN literals must have the same value type"); - values.push_back(*typed_value); - } - return TagValues(std::move(values)); - } - case FieldType::BIGINT: { - std::vector values; - values.reserve(literals.size()); - for (const auto& literal : literals) { - PAIMON_ASSIGN_OR_RAISE(TagValue value, LiteralToTagValue(literal)); - auto typed_value = std::get_if(&value); - CHECK_NOT_NULL(typed_value, - "lumina tag predicate IN literals must have the same value type"); - values.push_back(*typed_value); - } - return TagValues(std::move(values)); - } - case FieldType::FLOAT: { - std::vector values; - values.reserve(literals.size()); - for (const auto& literal : literals) { - PAIMON_ASSIGN_OR_RAISE(TagValue value, LiteralToTagValue(literal)); - auto typed_value = std::get_if(&value); - CHECK_NOT_NULL(typed_value, - "lumina tag predicate IN literals must have the same value type"); - values.push_back(*typed_value); - } - return TagValues(std::move(values)); - } - case FieldType::DOUBLE: { - std::vector values; - values.reserve(literals.size()); - for (const auto& literal : literals) { - PAIMON_ASSIGN_OR_RAISE(TagValue value, LiteralToTagValue(literal)); - auto typed_value = std::get_if(&value); - CHECK_NOT_NULL(typed_value, - "lumina tag predicate IN literals must have the same value type"); - values.push_back(*typed_value); - } - return TagValues(std::move(values)); - } - case FieldType::STRING: { - std::vector values; - values.reserve(literals.size()); - for (const auto& literal : literals) { - PAIMON_ASSIGN_OR_RAISE(TagValue value, LiteralToTagValue(literal)); - auto typed_value = std::get_if(&value); - CHECK_NOT_NULL(typed_value, - "lumina tag predicate IN literals must have the same value type"); - values.push_back(std::move(*typed_value)); - } - return TagValues(std::move(values)); - } - default: - return Status::Invalid( - fmt::format("lumina tag predicate IN does not support literal type {}", - static_cast(literals[0].GetType()))); - } -} - -} // namespace - -Result> LuminaIndexWriter::ExtractTagDataForSegment( - const std::shared_ptr& struct_array, - const std::vector& tag_fields, int64_t segment_start, int64_t segment_len) { - std::vector tag_dimensions_data; - tag_dimensions_data.reserve(tag_fields.size()); - for (const auto& tag_field : tag_fields) { - auto field_array = struct_array->GetFieldByName(tag_field.name); - CHECK_NOT_NULL(field_array, - fmt::format("lumina tag field {} not in input array", tag_field.name)); - - TagDimensionData tag_dimension_data; - tag_dimension_data.tagkName = tag_field.name; - switch (tag_field.value_type) { - case LuminaTagField::ValueType::INT32: { - std::vector> values; - PAIMON_RETURN_NOT_OK( - ExtractTagValues(field_array, segment_start, segment_len, &values)); - tag_dimension_data.values = std::move(values); - break; - } - case LuminaTagField::ValueType::INT64: { - std::vector> values; - PAIMON_RETURN_NOT_OK( - ExtractTagValues(field_array, segment_start, segment_len, &values)); - tag_dimension_data.values = std::move(values); - break; - } - case LuminaTagField::ValueType::FLOAT: { - std::vector> values; - PAIMON_RETURN_NOT_OK( - ExtractTagValues(field_array, segment_start, segment_len, &values)); - tag_dimension_data.values = std::move(values); - break; - } - case LuminaTagField::ValueType::DOUBLE: { - std::vector> values; - PAIMON_RETURN_NOT_OK( - ExtractTagValues(field_array, segment_start, segment_len, &values)); - tag_dimension_data.values = std::move(values); - break; - } - case LuminaTagField::ValueType::STRING: { - std::vector> values; - PAIMON_RETURN_NOT_OK(ExtractTagValues(field_array, segment_start, - segment_len, &values)); - tag_dimension_data.values = std::move(values); - break; - } - } - tag_dimensions_data.push_back(std::move(tag_dimension_data)); - } - return tag_dimensions_data; -} - -Result> LuminaGlobalIndex::ParseTagSchema( - const std::map& lumina_options) { - auto iter = lumina_options.find(std::string(::lumina::core::kExtensionTagSchema)); - if (iter == lumina_options.end()) { - return std::vector(); - } - - rapidjson::Document document; - document.Parse(iter->second.c_str()); - if (document.HasParseError()) { - return Status::Invalid("lumina tag_schema must be a valid JSON string"); - } - - std::vector tag_fields; - if (document.IsArray()) { - if (document.Empty()) { - return Status::Invalid("lumina tag_schema must contain at least one tag definition"); - } - tag_fields.reserve(document.Size()); - for (rapidjson::SizeType i = 0; i < document.Size(); i++) { - PAIMON_ASSIGN_OR_RAISE(LuminaTagField field, - ParseTagField(document[i], fmt::format("tag[{}]", i))); - tag_fields.push_back(std::move(field)); - } - } else if (document.IsObject()) { - PAIMON_ASSIGN_OR_RAISE(LuminaTagField field, ParseTagField(document, "tag[0]")); - tag_fields.push_back(std::move(field)); - } else { - return Status::Invalid("lumina tag_schema must be an object or array of objects"); - } - - std::unordered_set seen_names; - for (const auto& field : tag_fields) { - if (!seen_names.insert(field.name).second) { - return Status::Invalid( - fmt::format("lumina tag_schema has duplicate key_name: {}", field.name)); - } - } - return tag_fields; -} - -Status LuminaGlobalIndex::ValidateTagFields(const arrow::StructType& struct_type, - const std::vector& tag_fields) { - for (const auto& tag_field : tag_fields) { - auto field = struct_type.GetFieldByName(tag_field.name); - CHECK_NOT_NULL( - field, fmt::format("lumina tag field {} not exist in arrow schema", tag_field.name)); - PAIMON_RETURN_NOT_OK(ValidateTagArrowType(tag_field, field->type())); - } - return Status::OK(); -} - -Result<::lumina::extensions::experimental::TagFilter> LuminaIndexReader::PredicateToTagFilter( - const std::shared_ptr& predicate) { - if (!predicate) { - return Status::Invalid("lumina tag predicate must not be null"); - } - - auto compound_predicate = std::dynamic_pointer_cast(predicate); - if (compound_predicate) { - std::vector<::lumina::extensions::experimental::TagFilter> children; - children.reserve(compound_predicate->Children().size()); - for (const auto& child : compound_predicate->Children()) { - PAIMON_ASSIGN_OR_RAISE(::lumina::extensions::experimental::TagFilter tag_filter, - PredicateToTagFilter(child)); - children.push_back(std::move(tag_filter)); - } - if (children.empty()) { - return Status::Invalid("lumina tag compound predicate must have at least one child"); - } - if (children.size() == 1) { - return std::move(children.front()); - } - switch (compound_predicate->GetFunction().GetType()) { - case Function::Type::AND: - return ::lumina::extensions::experimental::TagFilter::And(std::move(children)); - case Function::Type::OR: - return ::lumina::extensions::experimental::TagFilter::Or(std::move(children)); - default: - return Status::NotImplemented( - fmt::format("lumina tag predicate does not support compound function {}", - compound_predicate->GetFunction().ToString())); - } - } - - auto leaf_predicate = std::dynamic_pointer_cast(predicate); - if (!leaf_predicate) { - return Status::Invalid( - fmt::format("cannot cast predicate {} to CompoundPredicate or LeafPredicate", - predicate->ToString())); - } - - const auto& literals = leaf_predicate->Literals(); - const auto& field_name = leaf_predicate->FieldName(); - switch (leaf_predicate->GetFunction().GetType()) { - case Function::Type::EQUAL: { - PAIMON_ASSIGN_OR_RAISE(const Literal* literal, GetSingleLiteral(literals, "equal")); - PAIMON_ASSIGN_OR_RAISE(TagValue value, LiteralToTagValue(*literal)); - return ::lumina::extensions::experimental::TagFilter::Eq(field_name, std::move(value)); - } - case Function::Type::GREATER_THAN: { - PAIMON_ASSIGN_OR_RAISE(const Literal* literal, - GetSingleLiteral(literals, "greater than")); - PAIMON_ASSIGN_OR_RAISE(TagValue value, LiteralToTagValue(*literal)); - return ::lumina::extensions::experimental::TagFilter::Gt(field_name, std::move(value)); - } - case Function::Type::GREATER_OR_EQUAL: { - PAIMON_ASSIGN_OR_RAISE(const Literal* literal, - GetSingleLiteral(literals, "greater or equal")); - PAIMON_ASSIGN_OR_RAISE(TagValue value, LiteralToTagValue(*literal)); - return ::lumina::extensions::experimental::TagFilter::Gte(field_name, std::move(value)); - } - case Function::Type::LESS_THAN: { - PAIMON_ASSIGN_OR_RAISE(const Literal* literal, GetSingleLiteral(literals, "less than")); - PAIMON_ASSIGN_OR_RAISE(TagValue value, LiteralToTagValue(*literal)); - return ::lumina::extensions::experimental::TagFilter::Lt(field_name, std::move(value)); - } - case Function::Type::LESS_OR_EQUAL: { - PAIMON_ASSIGN_OR_RAISE(const Literal* literal, - GetSingleLiteral(literals, "less or equal")); - PAIMON_ASSIGN_OR_RAISE(TagValue value, LiteralToTagValue(*literal)); - return ::lumina::extensions::experimental::TagFilter::Lte(field_name, std::move(value)); - } - case Function::Type::IN: { - PAIMON_ASSIGN_OR_RAISE(TagValues values, LiteralsToTagValues(literals)); - return ::lumina::extensions::experimental::TagFilter::In(field_name, std::move(values)); - } - default: - return Status::NotImplemented( - fmt::format("lumina tag predicate does not support leaf function {}", - leaf_predicate->GetFunction().ToString())); - } -} - -Result> LuminaGlobalIndex::CreateWriter( - const std::string& field_name, ::ArrowSchema* arrow_schema, - const std::shared_ptr& file_writer, - const std::shared_ptr& pool) const { - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_type, - arrow::ImportType(arrow_schema)); - // check data type - auto struct_type = std::dynamic_pointer_cast(arrow_type); - CHECK_NOT_NULL(struct_type, "arrow schema must be struct type when create LuminaIndexWriter"); - auto index_field = struct_type->GetFieldByName(field_name); - CHECK_NOT_NULL(index_field, - fmt::format("field {} not exist in arrow schema when create LuminaIndexWriter", - field_name)); - auto list_type = std::dynamic_pointer_cast(index_field->type()); - CHECK_NOT_NULL(list_type, "field type must be list[float] when create LuminaIndexWriter"); - if (list_type->value_type()->id() != arrow::Type::type::FLOAT) { - return Status::Invalid("field type must be list[float] when create LuminaIndexWriter"); - } - - // check options - auto lumina_options = - OptionsUtils::FetchOptionsWithPrefix(LuminaDefines::kOptionKeyPrefix, options_); - PAIMON_ASSIGN_OR_RAISE(std::vector tag_fields, ParseTagSchema(lumina_options)); - PAIMON_RETURN_NOT_OK(ValidateTagFields(*struct_type, tag_fields)); - PAIMON_ASSIGN_OR_RAISE(uint32_t dimension, - OptionsUtils::GetValueFromMap( - lumina_options, std::string(::lumina::core::kDimension))); - - PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( - ::lumina::api::BuilderOptions builder_options, - ::lumina::api::NormalizeBuilderOptions(std::unordered_map( - lumina_options.begin(), lumina_options.end()))); - auto lumina_pool = std::make_shared(pool); - return std::make_shared( - field_name, arrow_type, dimension, file_writer, std::move(builder_options), - ::lumina::api::IOOptions(), lumina_options, std::move(tag_fields), lumina_pool); -} - -Result LuminaIndexReader::GetIndexInfo( - const GlobalIndexIOMeta& io_meta) { - auto meta_bytes = io_meta.metadata; - if (!meta_bytes) { - return Status::Invalid("Lumina global index must have meta data"); - } - std::map lumina_write_options; - PAIMON_RETURN_NOT_OK(RapidJsonUtil::FromJsonString( - std::string(meta_bytes->data(), meta_bytes->size()), &lumina_write_options)); - - // check options - PAIMON_ASSIGN_OR_RAISE(uint32_t dimension, - OptionsUtils::GetValueFromMap( - lumina_write_options, std::string(::lumina::core::kDimension))); - PAIMON_ASSIGN_OR_RAISE(std::string index_type, - OptionsUtils::GetValueFromMap( - lumina_write_options, std::string(::lumina::core::kIndexType))); - PAIMON_ASSIGN_OR_RAISE(std::string distance_type_str, - OptionsUtils::GetValueFromMap( - lumina_write_options, std::string(::lumina::core::kDistanceMetric))); - VectorSearch::DistanceType distance_type = VectorSearch::DistanceType::UNKNOWN; - if (distance_type_str == ::lumina::core::kDistanceL2) { - distance_type = VectorSearch::DistanceType::EUCLIDEAN; - } else if (distance_type_str == ::lumina::core::kDistanceCosine) { - distance_type = VectorSearch::DistanceType::COSINE; - } else if (distance_type_str == ::lumina::core::kDistanceInnerProduct) { - distance_type = VectorSearch::DistanceType::INNER_PRODUCT; - } - if (distance_type == VectorSearch::DistanceType::UNKNOWN) { - return Status::Invalid( - fmt::format("invalid distance type {} for lumina", distance_type_str)); - } - bool has_tag = lumina_write_options.find(std::string(::lumina::core::kExtensionTagSchema)) != - lumina_write_options.end(); - return LuminaIndexReader::IndexInfo({dimension, index_type, distance_type, has_tag}); -} - -Result> LuminaGlobalIndex::CreateReader( - ::ArrowSchema* c_arrow_schema, const std::shared_ptr& file_manager, - const std::vector& files, const std::shared_ptr& pool) const { - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_schema, - arrow::ImportSchema(c_arrow_schema)); - if (files.size() != 1) { - return Status::Invalid("lumina index only has one index file per shard"); - } - const auto& io_meta = files[0]; - // check data type - if (arrow_schema->num_fields() != 1) { - return Status::Invalid("LuminaGlobalIndex now only support one field"); - } - auto index_field = arrow_schema->field(0); - auto list_type = std::dynamic_pointer_cast(index_field->type()); - CHECK_NOT_NULL(list_type, "field type must be list[float] when create LuminaIndexReader"); - if (list_type->value_type()->id() != arrow::Type::type::FLOAT) { - return Status::Invalid("field type must be list[float] when create LuminaIndexReader"); - } - - // get index info from meta - PAIMON_ASSIGN_OR_RAISE(LuminaIndexReader::IndexInfo index_info, - LuminaIndexReader::GetIndexInfo(io_meta)); - - auto lumina_pool = std::make_shared(pool); - ::lumina::core::MemoryResourceConfig memory_resource(lumina_pool.get()); - - auto lumina_options = - OptionsUtils::FetchOptionsWithPrefix(LuminaDefines::kOptionKeyPrefix, options_); - lumina_options[std::string(::lumina::core::kDimension)] = std::to_string(index_info.dimension); - lumina_options[std::string(::lumina::core::kIndexType)] = index_info.index_type; - - PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( - ::lumina::api::SearcherOptions searcher_options, - ::lumina::api::NormalizeSearcherOptions(std::unordered_map( - lumina_options.begin(), lumina_options.end()))); - - PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( - ::lumina::api::LuminaSearcher lumina_searcher, - ::lumina::api::LuminaSearcher::Create(searcher_options, memory_resource)); - auto searcher = std::make_unique<::lumina::api::LuminaSearcher>(std::move(lumina_searcher)); - // get input stream and open index - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr in, - file_manager->GetInputStream(io_meta.file_path)); - auto lumina_file_reader = std::make_unique(in); - PAIMON_RETURN_NOT_OK_FROM_LUMINA( - searcher->Open(std::move(lumina_file_reader), ::lumina::api::IOOptions())); - - // check meta - if (searcher->GetMeta().dim != index_info.dimension) { - return Status::Invalid( - fmt::format("lumina index dimension {} mismatch dimension {} in io meta", - searcher->GetMeta().dim, index_info.dimension)); - } - auto searcher_with_filter = std::make_unique<::lumina::extensions::SearchWithFilterExtension>(); - PAIMON_RETURN_NOT_OK_FROM_LUMINA(searcher->Attach(*searcher_with_filter)); - std::unique_ptr<::lumina::extensions::experimental::SearchWithTagExtension> searcher_with_tag; - if (index_info.has_tag) { - searcher_with_tag = - std::make_unique<::lumina::extensions::experimental::SearchWithTagExtension>(); - PAIMON_RETURN_NOT_OK_FROM_LUMINA(searcher->Attach(*searcher_with_tag)); - } - return std::make_shared(index_info, std::move(searcher), - std::move(searcher_with_filter), - std::move(searcher_with_tag), lumina_pool); -} - -Result>> LuminaGlobalIndex::GetExtraFieldNames() const { - auto lumina_options = - OptionsUtils::FetchOptionsWithPrefix(LuminaDefines::kOptionKeyPrefix, options_); - PAIMON_ASSIGN_OR_RAISE(std::vector tag_fields, ParseTagSchema(lumina_options)); - if (tag_fields.empty()) { - return std::optional>(std::nullopt); - } - std::vector field_names; - field_names.reserve(tag_fields.size()); - for (const auto& tag_field : tag_fields) { - field_names.push_back(tag_field.name); - } - return std::optional>(std::move(field_names)); -} - -class LuminaDataset : public ::lumina::api::Dataset { - public: - LuminaDataset(int64_t element_count, uint32_t dimension, - const std::vector>& array_vec, - const std::vector& start_ids) - : element_count_(element_count), - dimension_(dimension), - array_vec_(array_vec), - start_ids_(start_ids) {} - - uint32_t Dim() const noexcept override { - return dimension_; - } - uint64_t TotalSize() const noexcept override { - return element_count_; - } - - ::lumina::core::Result GetNextBatch( - std::vector& vector_buffer, - std::vector<::lumina::core::vector_id_t>& id_buffer) noexcept override { - if (cursor_ >= array_vec_.size()) { - return ::lumina::core::Result::Ok(0); - } - auto& value_array = array_vec_[cursor_]; - int64_t value_array_length = value_array->length(); - int64_t batch_element_count = value_array_length / dimension_; - const float* value_ptr = value_array->raw_values(); - vector_buffer.resize(value_array_length); - memcpy(vector_buffer.data(), value_ptr, sizeof(float) * value_array_length); - id_buffer.resize(batch_element_count); - std::iota(id_buffer.begin(), id_buffer.end(), - static_cast<::lumina::core::vector_id_t>(start_ids_[cursor_])); - - // release the array when copy to vector_buffer - value_array.reset(); - cursor_++; - return ::lumina::core::Result::Ok(static_cast(batch_element_count)); - } - - private: - int64_t element_count_; - uint32_t dimension_; - std::vector> array_vec_; - std::vector start_ids_; - size_t cursor_ = 0; -}; - -class LuminaDatasetWithTag : public ::lumina::extensions::experimental::DatasetWithTag { - public: - LuminaDatasetWithTag(int64_t element_count, uint32_t dimension, - const std::vector>& array_vec, - const std::vector& start_ids, - const std::vector>& tag_data_vec) - : element_count_(element_count), - dimension_(dimension), - array_vec_(array_vec), - start_ids_(start_ids), - tag_data_vec_(tag_data_vec) {} - - uint32_t Dim() const noexcept override { - return dimension_; - } - uint64_t TotalSize() const noexcept override { - return element_count_; - } - - ::lumina::core::Result GetNextBatch( - std::vector& vector_buffer, std::vector<::lumina::core::vector_id_t>& id_buffer, - std::vector& tag_dimensions_data) noexcept override { - if (cursor_ >= array_vec_.size()) { - return ::lumina::core::Result::Ok(0); - } - auto& value_array = array_vec_[cursor_]; - int64_t value_array_length = value_array->length(); - int64_t batch_element_count = value_array_length / dimension_; - const float* value_ptr = value_array->raw_values(); - vector_buffer.resize(value_array_length); - memcpy(vector_buffer.data(), value_ptr, sizeof(float) * value_array_length); - id_buffer.resize(batch_element_count); - std::iota(id_buffer.begin(), id_buffer.end(), - static_cast<::lumina::core::vector_id_t>(start_ids_[cursor_])); - tag_dimensions_data = std::move(tag_data_vec_[cursor_]); - - value_array.reset(); - cursor_++; - return ::lumina::core::Result::Ok(static_cast(batch_element_count)); - } - - private: - int64_t element_count_; - uint32_t dimension_; - std::vector> array_vec_; - std::vector start_ids_; - std::vector> tag_data_vec_; - size_t cursor_ = 0; -}; - -LuminaIndexWriter::LuminaIndexWriter( - const std::string& field_name, const std::shared_ptr& arrow_type, - uint32_t dimension, const std::shared_ptr& file_manager, - ::lumina::api::BuilderOptions&& builder_options, ::lumina::api::IOOptions&& io_options, - const std::map& lumina_options, - std::vector&& tag_fields, const std::shared_ptr& pool) - : pool_(pool), - field_name_(field_name), - arrow_type_(arrow_type), - dimension_(dimension), - file_manager_(file_manager), - builder_options_(std::move(builder_options)), - io_options_(std::move(io_options)), - lumina_options_(lumina_options), - tag_fields_(std::move(tag_fields)) {} - -Status LuminaIndexWriter::AddBatch(::ArrowArray* arrow_array, - std::vector&& relative_row_ids) { - PAIMON_RETURN_NOT_OK( - GlobalIndexUtils::CheckRelativeRowIds(arrow_array, relative_row_ids, count_)); - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, - arrow::ImportArray(arrow_array, arrow_type_)); - if (array->null_count() != 0) { - return Status::Invalid("arrow_array in LuminaIndexWriter is invalid, must not null"); - } - auto struct_array = std::dynamic_pointer_cast(array); - CHECK_NOT_NULL(struct_array, "invalid input array in LuminaIndexWriter, must be struct array"); - auto field_array = struct_array->GetFieldByName(field_name_); - CHECK_NOT_NULL( - field_array, - fmt::format("invalid input array in LuminaIndexWriter, field {} not in input array", - field_name_)); - int64_t field_length = field_array->length(); - auto list_field_array = std::dynamic_pointer_cast(field_array); - CHECK_NOT_NULL(list_field_array, - "invalid input array in LuminaIndexWriter, field array must be list array"); - - // Split into contiguous non-null segments, skipping null rows in the list field. - int64_t segment_start = -1; - for (int64_t i = 0; i <= field_length; i++) { - bool is_null = (i < field_length) && list_field_array->IsNull(i); - bool is_end = (i == field_length); - - if (!is_null && !is_end && segment_start == -1) { - segment_start = i; - } - - if ((is_null || is_end) && segment_start != -1) { - int64_t segment_len = i - segment_start; - // Use value_offset to precisely locate the float range for this segment - auto value_start_offset = list_field_array->value_offset(segment_start); - auto value_end_offset = list_field_array->value_offset(segment_start + segment_len); - int64_t value_length = value_end_offset - value_start_offset; - auto sliced_values = std::dynamic_pointer_cast( - list_field_array->values()->Slice(value_start_offset, value_length)); - CHECK_NOT_NULL(sliced_values, - "invalid sliced value array in LuminaIndexWriter, must be float array"); - if (sliced_values->null_count() != 0) { - return Status::Invalid( - "field value array in LuminaIndexWriter is invalid, must not null"); - } - for (int64_t row = segment_start; row < segment_start + segment_len; row++) { - int64_t vector_length = - list_field_array->value_offset(row + 1) - list_field_array->value_offset(row); - if (vector_length != static_cast(dimension_)) { - return Status::Invalid(fmt::format( - "invalid input array in LuminaIndexWriter, vector at row [{}] has length " - "[{}], expected dimension [{}]", - row, vector_length, dimension_)); - } - } - if (!tag_fields_.empty()) { - PAIMON_ASSIGN_OR_RAISE(std::vector tag_data, - ExtractTagDataForSegment(struct_array, tag_fields_, - segment_start, segment_len)); - tag_data_vec_.push_back(std::move(tag_data)); - } - array_vec_.push_back(std::move(sliced_values)); - array_start_ids_.push_back(count_ + segment_start); - indexed_count_ += segment_len; - segment_start = -1; - } - } - - count_ += array->length(); - return Status::OK(); -} - -Result> LuminaIndexWriter::Finish() { - if (indexed_count_ == 0) { - return std::vector(); - } - ::lumina::core::MemoryResourceConfig memory_resource(pool_.get()); - PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( - ::lumina::api::LuminaBuilder builder, - ::lumina::api::LuminaBuilder::Create(builder_options_, memory_resource)); - // pretrain - LuminaDataset dataset1(indexed_count_, dimension_, array_vec_, array_start_ids_); - PAIMON_RETURN_NOT_OK_FROM_LUMINA(builder.PretrainFrom(dataset1)); - - // insert data - if (tag_fields_.empty()) { - LuminaDataset dataset2(indexed_count_, dimension_, array_vec_, array_start_ids_); - std::vector>().swap(array_vec_); - PAIMON_RETURN_NOT_OK_FROM_LUMINA(builder.InsertFrom(dataset2)); - } else { - ::lumina::extensions::experimental::BuildWithTagExtension tag_extension; - PAIMON_RETURN_NOT_OK_FROM_LUMINA(builder.Attach(tag_extension)); - LuminaDatasetWithTag dataset2(indexed_count_, dimension_, array_vec_, array_start_ids_, - tag_data_vec_); - std::vector>().swap(array_vec_); - std::vector>().swap(tag_data_vec_); - PAIMON_RETURN_NOT_OK_FROM_LUMINA(tag_extension.InsertFromWithTag(dataset2)); - } - - // dump index - PAIMON_ASSIGN_OR_RAISE(std::string index_file_name, - file_manager_->NewFileName(LuminaDefines::kIdentifier)); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr out, - file_manager_->NewOutputStream(index_file_name)); - auto file_writer = std::make_unique(out); - PAIMON_RETURN_NOT_OK_FROM_LUMINA(builder.Dump(std::move(file_writer), io_options_)); - // prepare GlobalIndexIOMeta - PAIMON_ASSIGN_OR_RAISE(int64_t file_size, file_manager_->GetFileSize(index_file_name)); - std::string options_json; - PAIMON_RETURN_NOT_OK(RapidJsonUtil::ToJsonString(lumina_options_, &options_json)); - auto meta_bytes = std::make_shared(options_json, pool_->GetPaimonPool().get()); - GlobalIndexIOMeta meta(file_manager_->ToPath(index_file_name), file_size, - /*metadata=*/meta_bytes); - return std::vector({meta}); -} - -LuminaIndexReader::LuminaIndexReader( - const LuminaIndexReader::IndexInfo& index_info, - std::unique_ptr<::lumina::api::LuminaSearcher>&& searcher, - std::unique_ptr<::lumina::extensions::SearchWithFilterExtension>&& searcher_with_filter, - std::unique_ptr<::lumina::extensions::experimental::SearchWithTagExtension>&& searcher_with_tag, - const std::shared_ptr& pool) - : index_info_(index_info), - pool_(pool), - searcher_(std::move(searcher)), - searcher_with_filter_(std::move(searcher_with_filter)), - searcher_with_tag_(std::move(searcher_with_tag)) {} - -Result> LuminaIndexReader::VisitVectorSearch( - const std::shared_ptr& vector_search) { - if (vector_search->distance_type && - vector_search->distance_type.value() != index_info_.distance_type) { - return Status::Invalid("distance type for index and search not match"); - } - if (vector_search->query.size() != index_info_.dimension) { - return Status::Invalid("dimension for index and search not match"); - } - - auto lumina_options = OptionsUtils::FetchOptionsWithPrefix(LuminaDefines::kOptionKeyPrefix, - vector_search->options); - auto index_type_iter = lumina_options.find(std::string(::lumina::core::kIndexType)); - if (index_type_iter != lumina_options.end() && - index_type_iter->second != index_info_.index_type) { - return Status::Invalid("index type for index and search not match"); - } - - lumina_options[std::string(::lumina::core::kTopK)] = std::to_string(vector_search->limit); - lumina_options[std::string(::lumina::core::kSearchThreadSafeFilter)] = "true"; - PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( - ::lumina::api::SearchOptions search_options, - ::lumina::api::NormalizeSearchOptions(index_info_.index_type, - std::unordered_map( - lumina_options.begin(), lumina_options.end()))); - - ::lumina::api::Query lumina_query(vector_search->query.data(), vector_search->query.size()); - ::lumina::api::LuminaSearcher::SearchResult search_result; - if (vector_search->predicate) { - if (!searcher_with_tag_) { - return Status::Invalid("lumina index was not built with tag"); - } - PAIMON_ASSIGN_OR_RAISE(::lumina::extensions::experimental::TagFilter tag_filter, - PredicateToTagFilter(vector_search->predicate)); - if (!vector_search->pre_filter) { - PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( - search_result, searcher_with_tag_->SearchWithTag(lumina_query, tag_filter, - search_options, *pool_)); - } else { - auto lumina_filter = [filter = vector_search->pre_filter]( - ::lumina::core::vector_id_t id) -> bool { return filter(id); }; - PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( - search_result, - searcher_with_tag_->SearchWithTagAndFilter(lumina_query, tag_filter, lumina_filter, - search_options, *pool_)); - } - } else if (!vector_search->pre_filter) { - PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA(search_result, - searcher_->Search(lumina_query, search_options, *pool_)); - } else { - auto lumina_filter = [filter = vector_search->pre_filter]( - ::lumina::core::vector_id_t id) -> bool { return filter(id); }; - PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( - search_result, searcher_with_filter_->SearchWithFilter(lumina_query, lumina_filter, - search_options, *pool_)); - } - - // prepare BitmapScoredGlobalIndexResult - std::map id_to_score; - for (const auto& [id, score] : search_result.topk) { - id_to_score[id] = score; - } - - RoaringBitmap64 bitmap; - std::vector scores; - scores.reserve(id_to_score.size()); - for (const auto& [id, score] : id_to_score) { - bitmap.Add(id); - scores.push_back(score); - } - return std::make_shared(std::move(bitmap), std::move(scores)); -} - -} // namespace paimon::lumina diff --git a/src/paimon/global_index/lumina/CMakeLists.txt b/src/paimon/indexer/lumina/CMakeLists.txt similarity index 84% rename from src/paimon/global_index/lumina/CMakeLists.txt rename to src/paimon/indexer/lumina/CMakeLists.txt index b0496df6..7cd898a0 100644 --- a/src/paimon/global_index/lumina/CMakeLists.txt +++ b/src/paimon/indexer/lumina/CMakeLists.txt @@ -15,7 +15,16 @@ # limitations under the License. if(PAIMON_ENABLE_LUMINA) - set(PAIMON_LUMINA_INDEX lumina_global_index.cpp lumina_global_index_factory.cpp) + set(PAIMON_LUMINA_INDEX + lumina_dataset.cpp + lumina_file_index.cpp + lumina_file_index_factory.cpp + lumina_global_index.cpp + lumina_global_index_factory.cpp + lumina_index_accumulator.cpp + lumina_index_options.cpp + lumina_search_utils.cpp + lumina_tag_utils.cpp) add_paimon_lib(paimon_lumina_index SOURCES @@ -38,6 +47,7 @@ if(PAIMON_ENABLE_LUMINA) add_paimon_test(lumina_index_test SOURCES lumina_api_test.cpp + lumina_file_index_test.cpp lumina_file_io_test.cpp lumina_global_index_test.cpp EXTRA_INCLUDES diff --git a/src/paimon/global_index/lumina/lumina_api_test.cpp b/src/paimon/indexer/lumina/lumina_api_test.cpp similarity index 98% rename from src/paimon/global_index/lumina/lumina_api_test.cpp rename to src/paimon/indexer/lumina/lumina_api_test.cpp index 2c4954b5..c16aa22b 100644 --- a/src/paimon/global_index/lumina/lumina_api_test.cpp +++ b/src/paimon/indexer/lumina/lumina_api_test.cpp @@ -22,9 +22,9 @@ #include "lumina/core/Types.h" #include "lumina/extensions/SearchWithFilterExtension.h" #include "paimon/fs/local/local_file_system.h" -#include "paimon/global_index/lumina/lumina_file_reader.h" -#include "paimon/global_index/lumina/lumina_file_writer.h" -#include "paimon/global_index/lumina/lumina_memory_pool.h" +#include "paimon/indexer/lumina/lumina_file_reader.h" +#include "paimon/indexer/lumina/lumina_file_writer.h" +#include "paimon/indexer/lumina/lumina_memory_pool.h" #include "paimon/testing/utils/testharness.h" namespace paimon::lumina::test { class LuminaInterfaceTest : public ::testing::Test { diff --git a/src/paimon/indexer/lumina/lumina_dataset.cpp b/src/paimon/indexer/lumina/lumina_dataset.cpp new file mode 100644 index 00000000..81079b61 --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_dataset.cpp @@ -0,0 +1,105 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/indexer/lumina/lumina_dataset.h" + +#include +#include +#include + +namespace paimon::lumina { + +LuminaDataset::LuminaDataset(int64_t element_count, uint32_t dimension, + const std::vector>& arrays, + const std::vector& start_ids) + : element_count_(element_count), + dimension_(dimension), + arrays_(arrays), + start_ids_(start_ids) {} + +uint32_t LuminaDataset::Dim() const noexcept { + return dimension_; +} + +uint64_t LuminaDataset::TotalSize() const noexcept { + return static_cast(element_count_); +} + +::lumina::core::Result LuminaDataset::GetNextBatch( + std::vector& vector_buffer, + std::vector<::lumina::core::vector_id_t>& id_buffer) noexcept { + if (cursor_ >= arrays_.size()) { + return ::lumina::core::Result::Ok(0); + } + std::shared_ptr& values = arrays_[cursor_]; + int64_t value_count = values->length(); + int64_t vector_count = value_count / dimension_; + vector_buffer.resize(static_cast(value_count)); + std::memcpy(vector_buffer.data(), values->raw_values(), + sizeof(float) * static_cast(value_count)); + id_buffer.resize(static_cast(vector_count)); + std::iota(id_buffer.begin(), id_buffer.end(), + static_cast<::lumina::core::vector_id_t>(start_ids_[cursor_])); + + // release the array when copy to vector_buffer + values.reset(); + ++cursor_; + return ::lumina::core::Result::Ok(static_cast(vector_count)); +} + +LuminaDatasetWithTag::LuminaDatasetWithTag( + int64_t element_count, uint32_t dimension, + const std::vector>& arrays, + const std::vector& start_ids, + const std::vector>& tag_data) + : element_count_(element_count), + dimension_(dimension), + arrays_(arrays), + start_ids_(start_ids), + tag_data_(tag_data) {} + +uint32_t LuminaDatasetWithTag::Dim() const noexcept { + return dimension_; +} + +uint64_t LuminaDatasetWithTag::TotalSize() const noexcept { + return static_cast(element_count_); +} + +::lumina::core::Result LuminaDatasetWithTag::GetNextBatch( + std::vector& vector_buffer, std::vector<::lumina::core::vector_id_t>& id_buffer, + std::vector& tag_dimensions_data) noexcept { + if (cursor_ >= arrays_.size()) { + return ::lumina::core::Result::Ok(0); + } + std::shared_ptr& values = arrays_[cursor_]; + int64_t value_count = values->length(); + int64_t vector_count = value_count / dimension_; + vector_buffer.resize(static_cast(value_count)); + std::memcpy(vector_buffer.data(), values->raw_values(), + sizeof(float) * static_cast(value_count)); + id_buffer.resize(static_cast(vector_count)); + std::iota(id_buffer.begin(), id_buffer.end(), + static_cast<::lumina::core::vector_id_t>(start_ids_[cursor_])); + tag_dimensions_data = std::move(tag_data_[cursor_]); + values.reset(); + ++cursor_; + return ::lumina::core::Result::Ok(static_cast(vector_count)); +} + +} // namespace paimon::lumina diff --git a/src/paimon/indexer/lumina/lumina_dataset.h b/src/paimon/indexer/lumina/lumina_dataset.h new file mode 100644 index 00000000..8e820656 --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_dataset.h @@ -0,0 +1,80 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include + +#include "arrow/array.h" +#include "lumina/api/Dataset.h" +#include "lumina/extensions/experimental/DatasetWithTag.h" + +namespace paimon::lumina { + +class LuminaDataset final : public ::lumina::api::Dataset { + public: + LuminaDataset(int64_t element_count, uint32_t dimension, + const std::vector>& arrays, + const std::vector& start_ids); + + uint32_t Dim() const noexcept override; + + uint64_t TotalSize() const noexcept override; + + ::lumina::core::Result GetNextBatch( + std::vector& vector_buffer, + std::vector<::lumina::core::vector_id_t>& id_buffer) noexcept override; + + private: + int64_t element_count_; + uint32_t dimension_; + std::vector> arrays_; + std::vector start_ids_; + size_t cursor_ = 0; +}; + +class LuminaDatasetWithTag final : public ::lumina::extensions::experimental::DatasetWithTag { + public: + using TagDimensionData = ::lumina::extensions::experimental::TagDimensionData; + + LuminaDatasetWithTag(int64_t element_count, uint32_t dimension, + const std::vector>& arrays, + const std::vector& start_ids, + const std::vector>& tag_data); + + uint32_t Dim() const noexcept override; + + uint64_t TotalSize() const noexcept override; + + ::lumina::core::Result GetNextBatch( + std::vector& vector_buffer, std::vector<::lumina::core::vector_id_t>& id_buffer, + std::vector& tag_dimensions_data) noexcept override; + + private: + int64_t element_count_; + uint32_t dimension_; + std::vector> arrays_; + std::vector start_ids_; + std::vector> tag_data_; + size_t cursor_ = 0; +}; + +} // namespace paimon::lumina diff --git a/src/paimon/indexer/lumina/lumina_file_index.cpp b/src/paimon/indexer/lumina/lumina_file_index.cpp new file mode 100644 index 00000000..a80c66bd --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_file_index.cpp @@ -0,0 +1,218 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/indexer/lumina/lumina_file_index.h" + +#include +#include +#include + +#include "arrow/c/bridge.h" +#include "fmt/format.h" +#include "lumina/api/LuminaBuilder.h" +#include "lumina/core/Types.h" +#include "paimon/common/io/byte_array_output_stream.h" +#include "paimon/common/io/memory_segment_output_stream.h" +#include "paimon/common/io/offset_input_stream.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/file_index/scored_file_index_result.h" +#include "paimon/indexer/lumina/lumina_file_writer.h" +#include "paimon/indexer/lumina/lumina_search_utils.h" +#include "paimon/indexer/lumina/lumina_utils.h" +#include "paimon/memory/bytes.h" +#include "paimon/status.h" + +namespace paimon::lumina { +namespace { + +Result> ImportVectorSchema(::ArrowSchema* c_schema, + const std::string& owner) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr schema, + arrow::ImportSchema(c_schema)); + if (schema->num_fields() == 0) { + return Status::Invalid(fmt::format("{} requires at least one field", owner)); + } + std::shared_ptr list_type = + std::dynamic_pointer_cast(schema->field(0)->type()); + if (!list_type || list_type->value_type()->id() != arrow::Type::FLOAT) { + return Status::Invalid(fmt::format("{} field type must be list[float]", owner)); + } + return schema; +} + +} // namespace + +LuminaFileIndexWriter::LuminaFileIndexWriter(std::string field_name, + std::shared_ptr arrow_type, + const LuminaIndexInfo& index_info, + ::lumina::api::BuilderOptions&& builder_options, + std::vector&& tag_fields, + std::shared_ptr pool) + : field_name_(std::move(field_name)), + arrow_type_(std::move(arrow_type)), + index_info_(index_info), + builder_options_(std::move(builder_options)), + tag_fields_(std::move(tag_fields)), + pool_(std::move(pool)) {} + +Status LuminaFileIndexWriter::AddBatch(::ArrowArray* batch) { + if (serialized_) { + return Status::Invalid("Cannot add data after serializing a Lumina File Index"); + } + if (!batch || !batch->release) { + return Status::Invalid("Lumina File Index batch cannot be null or released"); + } + if (batch->length < 0 || + row_count_ > static_cast(RoaringBitmap32::MAX_VALUE) - batch->length) { + return Status::Invalid("Lumina File Index row count exceeds the bitmap32 limit"); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::ImportArray(batch, arrow_type_)); + if (array->null_count() != 0) { + return Status::Invalid("Lumina File Index struct array must not contain null rows"); + } + std::shared_ptr struct_array = + std::dynamic_pointer_cast(array); + if (!struct_array) { + return Status::Invalid("Lumina File Index input must be a struct array"); + } + std::shared_ptr vectors = + std::dynamic_pointer_cast(struct_array->GetFieldByName(field_name_)); + if (!vectors) { + return Status::Invalid("Lumina File Index field must be a list array"); + } + + PAIMON_RETURN_NOT_OK(accumulator_.AddBatch(struct_array, vectors, index_info_.dimension, + tag_fields_, row_count_)); + row_count_ += array->length(); + return Status::OK(); +} + +Result> LuminaFileIndexWriter::SerializedBytes() const { + if (serialized_) { + return Status::Invalid("Lumina File Index has already been serialized"); + } + serialized_ = true; + if (accumulator_.IndexedCount() == 0) { + return PAIMON_UNIQUE_PTR(); + } + + PAIMON_ASSIGN_OR_RAISE(::lumina::api::LuminaBuilder builder, + accumulator_.Build(builder_options_, index_info_.dimension, + !tag_fields_.empty(), pool_.get())); + + auto segment_output = std::make_unique( + MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool_->GetPaimonPool()); + std::shared_ptr output = + std::make_shared(std::move(segment_output)); + auto file_writer = std::make_unique(output); + PAIMON_RETURN_NOT_OK_FROM_LUMINA( + builder.Dump(std::move(file_writer), ::lumina::api::IOOptions())); + return output->Finish(pool_->GetPaimonPool().get()); +} + +LuminaFileIndexReader::LuminaFileIndexReader( + const LuminaIndexInfo& index_info, std::unique_ptr<::lumina::api::LuminaSearcher>&& searcher, + std::unique_ptr<::lumina::extensions::SearchWithFilterExtension>&& searcher_with_filter, + std::unique_ptr<::lumina::extensions::experimental::SearchWithTagExtension>&& searcher_with_tag, + std::shared_ptr pool) + : index_info_(index_info), + pool_(std::move(pool)), + searcher_(std::move(searcher)), + searcher_with_filter_(std::move(searcher_with_filter)), + searcher_with_tag_(std::move(searcher_with_tag)) {} + +LuminaFileIndexReader::~LuminaFileIndexReader() { + [[maybe_unused]] ::lumina::core::Status status = searcher_->Close(); +} + +Result> LuminaFileIndexReader::VisitVectorSearch( + const std::shared_ptr& vector_search) { + PAIMON_ASSIGN_OR_RAISE(::lumina::api::LuminaSearcher::SearchResult search_result, + LuminaSearchUtils::ExecuteVectorSearch( + *searcher_, *searcher_with_filter_, searcher_with_tag_.get(), + vector_search, index_info_, *pool_)); + + std::map rows_with_scores; + for (const auto& [row_id, score] : search_result.topk) { + if (row_id < 0 || row_id > RoaringBitmap32::MAX_VALUE) { + return Status::Invalid( + fmt::format("Lumina returned out-of-range file row id {}", row_id)); + } + if (!rows_with_scores.emplace(static_cast(row_id), score).second) { + return Status::Invalid(fmt::format("Lumina returned duplicate file row id {}", row_id)); + } + } + RoaringBitmap32 row_positions; + std::vector scores; + scores.reserve(rows_with_scores.size()); + for (const auto& [row_id, score] : rows_with_scores) { + row_positions.Add(row_id); + scores.push_back(score); + } + return ScoredFileIndexResult::Create(std::move(row_positions), std::move(scores)); +} + +Result>> LuminaFileIndexer::GetExtraFieldNames() const { + return LuminaTagUtils::GetExtraFieldNames(options_); +} + +Result> LuminaFileIndexer::CreateWriter( + ::ArrowSchema* c_schema, const std::shared_ptr& pool) const { + if (!pool) { + return Status::Invalid("Lumina File Index memory pool cannot be null"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr schema, + ImportVectorSchema(c_schema, "Lumina File Index writer")); + PAIMON_ASSIGN_OR_RAISE(LuminaIndexInfo index_info, LuminaIndexOptions::GetIndexInfo(options_)); + PAIMON_ASSIGN_OR_RAISE(std::vector tag_fields, + LuminaTagUtils::ParseTagSchema(options_)); + std::shared_ptr struct_type = + std::dynamic_pointer_cast(arrow::struct_(schema->fields())); + PAIMON_RETURN_NOT_OK(LuminaTagUtils::ValidateTagFields(*struct_type, tag_fields)); + PAIMON_ASSIGN_OR_RAISE(::lumina::api::BuilderOptions builder_options, + LuminaIndexOptions::CreateBuilderOptions(options_)); + auto lumina_pool = std::make_shared(pool); + return std::make_shared(schema->field(0)->name(), std::move(struct_type), + index_info, std::move(builder_options), + std::move(tag_fields), std::move(lumina_pool)); +} + +Result> LuminaFileIndexer::CreateReader( + ::ArrowSchema* c_schema, int32_t start, int32_t length, + const std::shared_ptr& input_stream, + const std::shared_ptr& pool) const { + if (!input_stream || !pool) { + return Status::Invalid("Lumina File Index reader requires input stream and memory pool"); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr schema, + ImportVectorSchema(c_schema, "Lumina File Index reader")); + (void)schema; + PAIMON_ASSIGN_OR_RAISE(LuminaIndexInfo index_info, LuminaIndexOptions::GetIndexInfo(options_)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr artifact_input, + OffsetInputStream::Create(input_stream, length, start)); + auto lumina_pool = std::make_shared(pool); + PAIMON_ASSIGN_OR_RAISE( + LuminaSearcherWithExtensions opened, + LuminaSearchUtils::OpenSearcher(options_, index_info, artifact_input, lumina_pool.get())); + return std::make_shared( + index_info, std::move(opened.searcher), std::move(opened.extensions.searcher_with_filter), + std::move(opened.extensions.searcher_with_tag), std::move(lumina_pool)); +} + +} // namespace paimon::lumina diff --git a/src/paimon/indexer/lumina/lumina_file_index.h b/src/paimon/indexer/lumina/lumina_file_index.h new file mode 100644 index 00000000..c851ed92 --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_file_index.h @@ -0,0 +1,107 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "lumina/api/LuminaSearcher.h" +#include "lumina/api/Options.h" +#include "lumina/extensions/SearchWithFilterExtension.h" +#include "lumina/extensions/experimental/SearchWithTagExtension.h" +#include "paimon/file_index/file_indexer.h" +#include "paimon/indexer/lumina/lumina_index_accumulator.h" +#include "paimon/indexer/lumina/lumina_index_options.h" +#include "paimon/indexer/lumina/lumina_memory_pool.h" +#include "paimon/indexer/lumina/lumina_tag_utils.h" + +namespace paimon::lumina { + +class LuminaFileIndexer final : public FileIndexer { + public: + explicit LuminaFileIndexer(const std::map& options) + : options_(options) {} + + Result>> GetExtraFieldNames() const override; + + Result> CreateReader( + ::ArrowSchema* arrow_schema, int32_t start, int32_t length, + const std::shared_ptr& input_stream, + const std::shared_ptr& pool) const override; + + Result> CreateWriter( + ::ArrowSchema* arrow_schema, const std::shared_ptr& pool) const override; + + private: + std::map options_; +}; + +class LuminaFileIndexWriter final : public FileIndexWriter { + public: + LuminaFileIndexWriter(std::string field_name, std::shared_ptr arrow_type, + const LuminaIndexInfo& index_info, + ::lumina::api::BuilderOptions&& builder_options, + std::vector&& tag_fields, + std::shared_ptr pool); + + Status AddBatch(::ArrowArray* batch) override; + + Result> SerializedBytes() const override; + + private: + std::string field_name_; + std::shared_ptr arrow_type_; + LuminaIndexInfo index_info_; + ::lumina::api::BuilderOptions builder_options_; + std::vector tag_fields_; + std::shared_ptr pool_; + int64_t row_count_ = 0; + mutable LuminaIndexAccumulator accumulator_; + mutable bool serialized_ = false; +}; + +class LuminaFileIndexReader final : public FileIndexReader { + public: + LuminaFileIndexReader( + const LuminaIndexInfo& index_info, + std::unique_ptr<::lumina::api::LuminaSearcher>&& searcher, + std::unique_ptr<::lumina::extensions::SearchWithFilterExtension>&& searcher_with_filter, + std::unique_ptr<::lumina::extensions::experimental::SearchWithTagExtension>&& + searcher_with_tag, + std::shared_ptr pool); + + ~LuminaFileIndexReader() override; + + Result> VisitVectorSearch( + const std::shared_ptr& vector_search) override; + + private: + LuminaIndexInfo index_info_; + std::shared_ptr pool_; + std::unique_ptr<::lumina::api::LuminaSearcher> searcher_; + std::unique_ptr<::lumina::extensions::SearchWithFilterExtension> searcher_with_filter_; + std::unique_ptr<::lumina::extensions::experimental::SearchWithTagExtension> searcher_with_tag_; +}; + +} // namespace paimon::lumina diff --git a/src/paimon/indexer/lumina/lumina_file_index_factory.cpp b/src/paimon/indexer/lumina/lumina_file_index_factory.cpp new file mode 100644 index 00000000..ad4d1c3d --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_file_index_factory.cpp @@ -0,0 +1,46 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include + +#include "paimon/factories/factory.h" +#include "paimon/file_index/file_indexer_factory.h" +#include "paimon/indexer/lumina/lumina_file_index.h" + +namespace paimon::lumina { +namespace { + +class LuminaFileIndexFactory final : public FileIndexerFactory { + public: + const char* Identifier() const override { + return "lumina"; + } + + Result> Create( + const std::map& options) const override { + return std::make_unique(options); + } +}; + +REGISTER_PAIMON_FACTORY(LuminaFileIndexFactory); + +} // namespace +} // namespace paimon::lumina diff --git a/src/paimon/indexer/lumina/lumina_file_index_test.cpp b/src/paimon/indexer/lumina/lumina_file_index_test.cpp new file mode 100644 index 00000000..1a7b6203 --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_file_index_test.cpp @@ -0,0 +1,198 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/indexer/lumina/lumina_file_index.h" + +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "gtest/gtest.h" +#include "paimon/io/byte_array_input_stream.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/predicate/vector_search.h" +#include "paimon/testing/utils/testharness.h" +#include "paimon/utils/roaring_bitmap32.h" + +namespace paimon::lumina::test { +namespace { + +std::shared_ptr CreateVectors() { + std::shared_ptr values = + std::make_shared(arrow::default_memory_pool()); + arrow::ListBuilder vectors(arrow::default_memory_pool(), values); + EXPECT_TRUE(vectors.Append().ok()); + EXPECT_TRUE(values->AppendValues({0.0f, 0.0f, 0.0f, 0.0f}).ok()); + EXPECT_TRUE(vectors.AppendNull().ok()); + EXPECT_TRUE(vectors.Append().ok()); + EXPECT_TRUE(values->AppendValues({1.0f, 1.0f, 1.0f, 1.0f}).ok()); + std::shared_ptr vector_array; + EXPECT_TRUE(vectors.Finish(&vector_array).ok()); + return arrow::StructArray::Make({vector_array}, + {arrow::field("embedding", vector_array->type())}) + .ValueOrDie(); +} + +std::shared_ptr CreateTaggedVectors() { + std::shared_ptr values = + std::make_shared(arrow::default_memory_pool()); + arrow::ListBuilder vectors(arrow::default_memory_pool(), values); + EXPECT_TRUE(vectors.Append().ok()); + EXPECT_TRUE(values->AppendValues({0.0f, 0.0f, 0.0f, 0.0f}).ok()); + EXPECT_TRUE(vectors.AppendNull().ok()); + EXPECT_TRUE(vectors.Append().ok()); + EXPECT_TRUE(values->AppendValues({1.0f, 1.0f, 1.0f, 1.0f}).ok()); + EXPECT_TRUE(vectors.Append().ok()); + EXPECT_TRUE(values->AppendValues({1.0f, 1.0f, 1.0f, 1.1f}).ok()); + std::shared_ptr vector_array; + EXPECT_TRUE(vectors.Finish(&vector_array).ok()); + + arrow::StringBuilder colors; + EXPECT_TRUE(colors.AppendValues({"red", "red", "blue", "red"}).ok()); + std::shared_ptr color_array; + EXPECT_TRUE(colors.Finish(&color_array).ok()); + return arrow::StructArray::Make({vector_array, color_array}, + {arrow::field("embedding", vector_array->type()), + arrow::field("color", color_array->type())}) + .ValueOrDie(); +} + +} // namespace + +TEST(LuminaFileIndexTest, RoundTripUsesFileLocalRowPositions) { + const std::map options = { + {"index.dimension", "4"}, + {"index.type", "bruteforce"}, + {"distance.metric", "l2"}, + {"encoding.type", "rawf32"}, + }; + std::shared_ptr pool = GetDefaultPool(); + std::shared_ptr batch = CreateVectors(); + std::shared_ptr schema = + arrow::schema({arrow::field("embedding", arrow::list(arrow::float32()))}); + LuminaFileIndexer indexer(options); + + ::ArrowSchema writer_schema; + ASSERT_TRUE(arrow::ExportSchema(*schema, &writer_schema).ok()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr writer, + indexer.CreateWriter(&writer_schema, pool)); + ::ArrowArray c_batch; + ASSERT_TRUE(arrow::ExportArray(*batch, &c_batch).ok()); + ASSERT_OK(writer->AddBatch(&c_batch)); + ASSERT_OK_AND_ASSIGN(PAIMON_UNIQUE_PTR artifact, writer->SerializedBytes()); + ASSERT_TRUE(artifact); + + std::string container = "prefix"; + container.append(artifact->data(), artifact->size()); + std::shared_ptr input = + std::make_shared(container.data(), container.size()); + ::ArrowSchema reader_schema; + ASSERT_TRUE(arrow::ExportSchema(*schema, &reader_schema).ok()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr reader, + indexer.CreateReader(&reader_schema, /*start=*/6, + static_cast(artifact->size()), input, pool)); + + std::shared_ptr search = std::make_shared( + "embedding", /*limit=*/1, std::vector{1.0f, 1.0f, 1.0f, 1.0f}, + /*pre_filter=*/nullptr, /*predicate=*/nullptr, VectorSearch::DistanceType::EUCLIDEAN, + std::map{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + reader->VisitVectorSearch(search)); + EXPECT_EQ(RoaringBitmap32::From({2}), result->GetRowPositions()); + ASSERT_EQ(1, result->GetScores().size()); + EXPECT_FLOAT_EQ(0.0f, result->GetScores()[0]); + + std::shared_ptr predicate = PredicateBuilder::Equal( + /*field_index=*/1, /*field_name=*/"color", FieldType::STRING, + Literal(FieldType::STRING, "red", 3)); + std::shared_ptr tag_search = std::make_shared( + "embedding", /*limit=*/1, std::vector{1.0f, 1.0f, 1.0f, 1.0f}, + /*pre_filter=*/nullptr, predicate, VectorSearch::DistanceType::EUCLIDEAN, + std::map{}); + ASSERT_NOK_WITH_MSG(reader->VisitVectorSearch(tag_search), + "lumina index was not built with tag"); +} + +TEST(LuminaFileIndexTest, RoundTripWithTagPredicate) { + const std::map options = { + {"index.dimension", "4"}, + {"index.type", "bruteforce"}, + {"distance.metric", "l2"}, + {"encoding.type", "rawf32"}, + {"extension.build.tag.tag_schema", + R"([{"key_name":"color","type":"enum","value_type":"string"}])"}, + }; + std::shared_ptr pool = GetDefaultPool(); + std::shared_ptr batch = CreateTaggedVectors(); + std::shared_ptr writer_schema = arrow::schema(batch->type()->fields()); + std::shared_ptr reader_schema = + arrow::schema({arrow::field("embedding", arrow::list(arrow::float32()))}); + LuminaFileIndexer indexer(options); + + ASSERT_OK_AND_ASSIGN(std::optional> extra_field_names, + indexer.GetExtraFieldNames()); + ASSERT_TRUE(extra_field_names); + EXPECT_EQ(std::vector({"color"}), extra_field_names.value()); + + ::ArrowSchema c_writer_schema; + ASSERT_TRUE(arrow::ExportSchema(*writer_schema, &c_writer_schema).ok()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr writer, + indexer.CreateWriter(&c_writer_schema, pool)); + ::ArrowArray c_batch; + ASSERT_TRUE(arrow::ExportArray(*batch, &c_batch).ok()); + ASSERT_OK(writer->AddBatch(&c_batch)); + ASSERT_OK_AND_ASSIGN(PAIMON_UNIQUE_PTR artifact, writer->SerializedBytes()); + ASSERT_TRUE(artifact); + + std::shared_ptr input = std::make_shared( + artifact->data(), static_cast(artifact->size())); + ::ArrowSchema c_reader_schema; + ASSERT_TRUE(arrow::ExportSchema(*reader_schema, &c_reader_schema).ok()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr reader, + indexer.CreateReader(&c_reader_schema, /*start=*/0, + static_cast(artifact->size()), input, pool)); + + std::shared_ptr predicate = PredicateBuilder::Equal( + /*field_index=*/1, /*field_name=*/"color", FieldType::STRING, + Literal(FieldType::STRING, "red", 3)); + std::shared_ptr search = std::make_shared( + "embedding", /*limit=*/1, std::vector{1.0f, 1.0f, 1.0f, 1.0f}, + /*pre_filter=*/nullptr, predicate, VectorSearch::DistanceType::EUCLIDEAN, + std::map{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + reader->VisitVectorSearch(search)); + EXPECT_EQ(RoaringBitmap32::From({3}), result->GetRowPositions()); + ASSERT_EQ(1, result->GetScores().size()); + EXPECT_NEAR(0.01f, result->GetScores()[0], 1e-5f); + + std::shared_ptr filtered_search = std::make_shared( + "embedding", /*limit=*/4, std::vector{1.0f, 1.0f, 1.0f, 1.0f}, + [](int64_t file_row) { return file_row == 0; }, predicate, + VectorSearch::DistanceType::EUCLIDEAN, std::map{}); + ASSERT_OK_AND_ASSIGN(result, reader->VisitVectorSearch(filtered_search)); + EXPECT_EQ(RoaringBitmap32::From({0}), result->GetRowPositions()); + ASSERT_EQ(1, result->GetScores().size()); + EXPECT_FLOAT_EQ(4.0f, result->GetScores()[0]); +} + +} // namespace paimon::lumina::test diff --git a/src/paimon/global_index/lumina/lumina_file_io_test.cpp b/src/paimon/indexer/lumina/lumina_file_io_test.cpp similarity index 97% rename from src/paimon/global_index/lumina/lumina_file_io_test.cpp rename to src/paimon/indexer/lumina/lumina_file_io_test.cpp index 7ac7327f..fa8aa0b5 100644 --- a/src/paimon/global_index/lumina/lumina_file_io_test.cpp +++ b/src/paimon/indexer/lumina/lumina_file_io_test.cpp @@ -17,8 +17,8 @@ */ #include -#include "paimon/global_index/lumina/lumina_file_reader.h" -#include "paimon/global_index/lumina/lumina_file_writer.h" +#include "paimon/indexer/lumina/lumina_file_reader.h" +#include "paimon/indexer/lumina/lumina_file_writer.h" #include "paimon/testing/utils/testharness.h" namespace paimon::lumina::test { class LuminaFileIOTest : public ::testing::Test { diff --git a/src/paimon/global_index/lumina/lumina_file_reader.h b/src/paimon/indexer/lumina/lumina_file_reader.h similarity index 98% rename from src/paimon/global_index/lumina/lumina_file_reader.h rename to src/paimon/indexer/lumina/lumina_file_reader.h index 0ae11692..168d3168 100644 --- a/src/paimon/global_index/lumina/lumina_file_reader.h +++ b/src/paimon/indexer/lumina/lumina_file_reader.h @@ -25,7 +25,7 @@ #include "lumina/io/FileReader.h" #include "paimon/common/utils/math.h" #include "paimon/fs/file_system.h" -#include "paimon/global_index/lumina/lumina_utils.h" +#include "paimon/indexer/lumina/lumina_utils.h" namespace paimon::lumina { class LuminaFileReader : public ::lumina::io::FileReader { public: diff --git a/src/paimon/global_index/lumina/lumina_file_writer.h b/src/paimon/indexer/lumina/lumina_file_writer.h similarity index 98% rename from src/paimon/global_index/lumina/lumina_file_writer.h rename to src/paimon/indexer/lumina/lumina_file_writer.h index a2514cc6..1452b90e 100644 --- a/src/paimon/global_index/lumina/lumina_file_writer.h +++ b/src/paimon/indexer/lumina/lumina_file_writer.h @@ -24,7 +24,7 @@ #include "lumina/io/FileWriter.h" #include "paimon/common/utils/math.h" #include "paimon/fs/file_system.h" -#include "paimon/global_index/lumina/lumina_utils.h" +#include "paimon/indexer/lumina/lumina_utils.h" namespace paimon::lumina { class LuminaFileWriter : public ::lumina::io::FileWriter { public: diff --git a/src/paimon/indexer/lumina/lumina_global_index.cpp b/src/paimon/indexer/lumina/lumina_global_index.cpp new file mode 100644 index 00000000..8d2ed27c --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_global_index.cpp @@ -0,0 +1,241 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/indexer/lumina/lumina_global_index.h" + +#include + +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "lumina/api/LuminaBuilder.h" +#include "lumina/api/LuminaSearcher.h" +#include "lumina/core/Constants.h" +#include "lumina/core/Status.h" +#include "lumina/core/Types.h" +#include "paimon/common/global_index/global_index_utils.h" +#include "paimon/common/utils/options_utils.h" +#include "paimon/common/utils/rapidjson_util.h" +#include "paimon/common/utils/string_utils.h" +#include "paimon/global_index/bitmap_scored_global_index_result.h" +#include "paimon/indexer/lumina/lumina_file_writer.h" +#include "paimon/indexer/lumina/lumina_index_options.h" +#include "paimon/indexer/lumina/lumina_search_utils.h" +#include "paimon/indexer/lumina/lumina_utils.h" +namespace paimon::lumina { +#define CHECK_NOT_NULL(pointer, error_msg) \ + do { \ + if (!(pointer)) { \ + return Status::Invalid(error_msg); \ + } \ + } while (0) + +Result> LuminaGlobalIndex::CreateWriter( + const std::string& field_name, ::ArrowSchema* arrow_schema, + const std::shared_ptr& file_writer, + const std::shared_ptr& pool) const { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_type, + arrow::ImportType(arrow_schema)); + // check data type + auto struct_type = std::dynamic_pointer_cast(arrow_type); + CHECK_NOT_NULL(struct_type, "arrow schema must be struct type when create LuminaIndexWriter"); + auto index_field = struct_type->GetFieldByName(field_name); + CHECK_NOT_NULL(index_field, + fmt::format("field {} not exist in arrow schema when create LuminaIndexWriter", + field_name)); + auto list_type = std::dynamic_pointer_cast(index_field->type()); + CHECK_NOT_NULL(list_type, "field type must be list[float] when create LuminaIndexWriter"); + if (list_type->value_type()->id() != arrow::Type::type::FLOAT) { + return Status::Invalid("field type must be list[float] when create LuminaIndexWriter"); + } + + // check options + auto lumina_options = + OptionsUtils::FetchOptionsWithPrefix(LuminaDefines::kOptionKeyPrefix, options_); + PAIMON_ASSIGN_OR_RAISE(std::vector tag_fields, + LuminaTagUtils::ParseTagSchema(lumina_options)); + PAIMON_RETURN_NOT_OK(LuminaTagUtils::ValidateTagFields(*struct_type, tag_fields)); + PAIMON_ASSIGN_OR_RAISE(uint32_t dimension, LuminaIndexOptions::GetDimension(lumina_options)); + PAIMON_ASSIGN_OR_RAISE(::lumina::api::BuilderOptions builder_options, + LuminaIndexOptions::CreateBuilderOptions(lumina_options)); + auto lumina_pool = std::make_shared(pool); + return std::make_shared( + field_name, arrow_type, dimension, file_writer, std::move(builder_options), + ::lumina::api::IOOptions(), lumina_options, std::move(tag_fields), lumina_pool); +} + +Result LuminaIndexReader::GetIndexInfo( + const GlobalIndexIOMeta& io_meta) { + auto meta_bytes = io_meta.metadata; + if (!meta_bytes) { + return Status::Invalid("Lumina global index must have meta data"); + } + std::map lumina_write_options; + PAIMON_RETURN_NOT_OK(RapidJsonUtil::FromJsonString( + std::string(meta_bytes->data(), meta_bytes->size()), &lumina_write_options)); + + return LuminaIndexOptions::GetIndexInfo(lumina_write_options); +} + +Result> LuminaGlobalIndex::CreateReader( + ::ArrowSchema* c_arrow_schema, const std::shared_ptr& file_manager, + const std::vector& files, const std::shared_ptr& pool) const { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_schema, + arrow::ImportSchema(c_arrow_schema)); + if (files.size() != 1) { + return Status::Invalid("lumina index only has one index file per shard"); + } + const auto& io_meta = files[0]; + // check data type + if (arrow_schema->num_fields() != 1) { + return Status::Invalid("LuminaGlobalIndex now only support one field"); + } + auto index_field = arrow_schema->field(0); + auto list_type = std::dynamic_pointer_cast(index_field->type()); + CHECK_NOT_NULL(list_type, "field type must be list[float] when create LuminaIndexReader"); + if (list_type->value_type()->id() != arrow::Type::type::FLOAT) { + return Status::Invalid("field type must be list[float] when create LuminaIndexReader"); + } + + // get index info from meta + PAIMON_ASSIGN_OR_RAISE(LuminaIndexReader::IndexInfo index_info, + LuminaIndexReader::GetIndexInfo(io_meta)); + + auto lumina_pool = std::make_shared(pool); + auto lumina_options = + OptionsUtils::FetchOptionsWithPrefix(LuminaDefines::kOptionKeyPrefix, options_); + // get input stream and open index + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr in, + file_manager->GetInputStream(io_meta.file_path)); + PAIMON_ASSIGN_OR_RAISE( + LuminaSearcherWithExtensions opened, + LuminaSearchUtils::OpenSearcher(lumina_options, index_info, in, lumina_pool.get())); + return std::make_shared( + index_info, std::move(opened.searcher), std::move(opened.extensions.searcher_with_filter), + std::move(opened.extensions.searcher_with_tag), lumina_pool); +} + +Result>> LuminaGlobalIndex::GetExtraFieldNames() const { + auto lumina_options = + OptionsUtils::FetchOptionsWithPrefix(LuminaDefines::kOptionKeyPrefix, options_); + return LuminaTagUtils::GetExtraFieldNames(lumina_options); +} + +LuminaIndexWriter::LuminaIndexWriter( + const std::string& field_name, const std::shared_ptr& arrow_type, + uint32_t dimension, const std::shared_ptr& file_manager, + ::lumina::api::BuilderOptions&& builder_options, ::lumina::api::IOOptions&& io_options, + const std::map& lumina_options, + std::vector&& tag_fields, const std::shared_ptr& pool) + : pool_(pool), + field_name_(field_name), + arrow_type_(arrow_type), + dimension_(dimension), + file_manager_(file_manager), + builder_options_(std::move(builder_options)), + io_options_(std::move(io_options)), + lumina_options_(lumina_options), + tag_fields_(std::move(tag_fields)) {} + +Status LuminaIndexWriter::AddBatch(::ArrowArray* arrow_array, + std::vector&& relative_row_ids) { + PAIMON_RETURN_NOT_OK( + GlobalIndexUtils::CheckRelativeRowIds(arrow_array, relative_row_ids, count_)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::ImportArray(arrow_array, arrow_type_)); + if (array->null_count() != 0) { + return Status::Invalid("arrow_array in LuminaIndexWriter is invalid, must not null"); + } + auto struct_array = std::dynamic_pointer_cast(array); + CHECK_NOT_NULL(struct_array, "invalid input array in LuminaIndexWriter, must be struct array"); + auto field_array = struct_array->GetFieldByName(field_name_); + CHECK_NOT_NULL( + field_array, + fmt::format("invalid input array in LuminaIndexWriter, field {} not in input array", + field_name_)); + auto list_field_array = std::dynamic_pointer_cast(field_array); + CHECK_NOT_NULL(list_field_array, + "invalid input array in LuminaIndexWriter, field array must be list array"); + + PAIMON_RETURN_NOT_OK( + accumulator_.AddBatch(struct_array, list_field_array, dimension_, tag_fields_, count_)); + + count_ += array->length(); + return Status::OK(); +} + +Result> LuminaIndexWriter::Finish() { + if (accumulator_.IndexedCount() == 0) { + return std::vector(); + } + PAIMON_ASSIGN_OR_RAISE( + ::lumina::api::LuminaBuilder builder, + accumulator_.Build(builder_options_, dimension_, !tag_fields_.empty(), pool_.get())); + + // dump index + PAIMON_ASSIGN_OR_RAISE(std::string index_file_name, + file_manager_->NewFileName(LuminaDefines::kIdentifier)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr out, + file_manager_->NewOutputStream(index_file_name)); + auto file_writer = std::make_unique(out); + PAIMON_RETURN_NOT_OK_FROM_LUMINA(builder.Dump(std::move(file_writer), io_options_)); + // prepare GlobalIndexIOMeta + PAIMON_ASSIGN_OR_RAISE(int64_t file_size, file_manager_->GetFileSize(index_file_name)); + std::string options_json; + PAIMON_RETURN_NOT_OK(RapidJsonUtil::ToJsonString(lumina_options_, &options_json)); + auto meta_bytes = std::make_shared(options_json, pool_->GetPaimonPool().get()); + GlobalIndexIOMeta meta(file_manager_->ToPath(index_file_name), file_size, + /*metadata=*/meta_bytes); + return std::vector({meta}); +} + +LuminaIndexReader::LuminaIndexReader( + const LuminaIndexReader::IndexInfo& index_info, + std::unique_ptr<::lumina::api::LuminaSearcher>&& searcher, + std::unique_ptr<::lumina::extensions::SearchWithFilterExtension>&& searcher_with_filter, + std::unique_ptr<::lumina::extensions::experimental::SearchWithTagExtension>&& searcher_with_tag, + const std::shared_ptr& pool) + : index_info_(index_info), + pool_(pool), + searcher_(std::move(searcher)), + searcher_with_filter_(std::move(searcher_with_filter)), + searcher_with_tag_(std::move(searcher_with_tag)) {} + +Result> LuminaIndexReader::VisitVectorSearch( + const std::shared_ptr& vector_search) { + PAIMON_ASSIGN_OR_RAISE(::lumina::api::LuminaSearcher::SearchResult search_result, + LuminaSearchUtils::ExecuteVectorSearch( + *searcher_, *searcher_with_filter_, searcher_with_tag_.get(), + vector_search, index_info_, *pool_)); + + // prepare BitmapScoredGlobalIndexResult + std::map id_to_score; + for (const auto& [id, score] : search_result.topk) { + id_to_score[id] = score; + } + + RoaringBitmap64 bitmap; + std::vector scores; + scores.reserve(id_to_score.size()); + for (const auto& [id, score] : id_to_score) { + bitmap.Add(id); + scores.push_back(score); + } + return std::make_shared(std::move(bitmap), std::move(scores)); +} + +} // namespace paimon::lumina diff --git a/src/paimon/global_index/lumina/lumina_global_index.h b/src/paimon/indexer/lumina/lumina_global_index.h similarity index 83% rename from src/paimon/global_index/lumina/lumina_global_index.h rename to src/paimon/indexer/lumina/lumina_global_index.h index c2c30475..63e1f3e8 100644 --- a/src/paimon/global_index/lumina/lumina_global_index.h +++ b/src/paimon/indexer/lumina/lumina_global_index.h @@ -22,7 +22,6 @@ #include #include #include -#include #include #include @@ -32,31 +31,14 @@ #include "lumina/extensions/SearchWithFilterExtension.h" #include "lumina/extensions/experimental/DatasetWithTag.h" #include "lumina/extensions/experimental/SearchWithTagExtension.h" -#include "lumina/extensions/experimental/TagFilter.h" -#include "paimon/global_index/bitmap_global_index_result.h" #include "paimon/global_index/global_indexer.h" -#include "paimon/global_index/lumina/lumina_memory_pool.h" -#include "paimon/global_index/lumina/lumina_utils.h" +#include "paimon/indexer/lumina/lumina_index_accumulator.h" +#include "paimon/indexer/lumina/lumina_index_options.h" +#include "paimon/indexer/lumina/lumina_memory_pool.h" +#include "paimon/indexer/lumina/lumina_tag_utils.h" +#include "paimon/indexer/lumina/lumina_utils.h" namespace paimon::lumina { -struct LuminaTagField { - enum class Type { - ENUM, - RANGE, - }; - - enum class ValueType { - INT32, - INT64, - FLOAT, - DOUBLE, - STRING, - }; - - std::string name; - Type type; - ValueType value_type; -}; /// @note When enabling the lumina global index in `paimon-cpp`, all configuration parameters /// specific to Lumina **must be prefixed with `lumina.`**. @@ -100,12 +82,6 @@ class LuminaGlobalIndex : public GlobalIndexer { const std::shared_ptr& pool) const override; private: - static Result> ParseTagSchema( - const std::map& lumina_options); - - static Status ValidateTagFields(const arrow::StructType& struct_type, - const std::vector& tag_fields); - std::map options_; }; @@ -125,13 +101,7 @@ class LuminaIndexWriter : public GlobalIndexWriter { Result> Finish() override; private: - static Result> - ExtractTagDataForSegment(const std::shared_ptr& struct_array, - const std::vector& tag_fields, int64_t segment_start, - int64_t segment_len); - int64_t count_ = 0; - int64_t indexed_count_ = 0; std::shared_ptr pool_; std::string field_name_; std::shared_ptr arrow_type_; @@ -141,19 +111,12 @@ class LuminaIndexWriter : public GlobalIndexWriter { ::lumina::api::IOOptions io_options_; std::map lumina_options_; std::vector tag_fields_; - std::vector> array_vec_; - std::vector array_start_ids_; - std::vector> tag_data_vec_; + LuminaIndexAccumulator accumulator_; }; class LuminaIndexReader : public GlobalIndexReader { public: - struct IndexInfo { - uint32_t dimension; - std::string index_type; - VectorSearch::DistanceType distance_type; - bool has_tag; - }; + using IndexInfo = LuminaIndexInfo; LuminaIndexReader( const IndexInfo& index_info, std::unique_ptr<::lumina::api::LuminaSearcher>&& searcher, @@ -246,9 +209,6 @@ class LuminaIndexReader : public GlobalIndexReader { static Result GetIndexInfo(const GlobalIndexIOMeta& io_meta); private: - static Result<::lumina::extensions::experimental::TagFilter> PredicateToTagFilter( - const std::shared_ptr& predicate); - LuminaIndexReader::IndexInfo index_info_; std::shared_ptr pool_; std::unique_ptr<::lumina::api::LuminaSearcher> searcher_; diff --git a/src/paimon/global_index/lumina/lumina_global_index_factory.cpp b/src/paimon/indexer/lumina/lumina_global_index_factory.cpp similarity index 90% rename from src/paimon/global_index/lumina/lumina_global_index_factory.cpp rename to src/paimon/indexer/lumina/lumina_global_index_factory.cpp index e27837da..6d26ecd8 100644 --- a/src/paimon/global_index/lumina/lumina_global_index_factory.cpp +++ b/src/paimon/indexer/lumina/lumina_global_index_factory.cpp @@ -16,14 +16,14 @@ * limitations under the License. */ -#include "paimon/global_index/lumina/lumina_global_index_factory.h" +#include "paimon/indexer/lumina/lumina_global_index_factory.h" #include #include #include #include -#include "paimon/global_index/lumina/lumina_global_index.h" +#include "paimon/indexer/lumina/lumina_global_index.h" namespace paimon::lumina { const char LuminaGlobalIndexFactory::IDENTIFIER[] = "lumina-global"; diff --git a/src/paimon/global_index/lumina/lumina_global_index_factory.h b/src/paimon/indexer/lumina/lumina_global_index_factory.h similarity index 100% rename from src/paimon/global_index/lumina/lumina_global_index_factory.h rename to src/paimon/indexer/lumina/lumina_global_index_factory.h diff --git a/src/paimon/global_index/lumina/lumina_global_index_test.cpp b/src/paimon/indexer/lumina/lumina_global_index_test.cpp similarity index 99% rename from src/paimon/global_index/lumina/lumina_global_index_test.cpp rename to src/paimon/indexer/lumina/lumina_global_index_test.cpp index 2d54950a..1ac1114a 100644 --- a/src/paimon/global_index/lumina/lumina_global_index_test.cpp +++ b/src/paimon/indexer/lumina/lumina_global_index_test.cpp @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#include "paimon/global_index/lumina/lumina_global_index.h" +#include "paimon/indexer/lumina/lumina_global_index.h" #include #include diff --git a/src/paimon/indexer/lumina/lumina_index_accumulator.cpp b/src/paimon/indexer/lumina/lumina_index_accumulator.cpp new file mode 100644 index 00000000..d6f3993b --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_index_accumulator.cpp @@ -0,0 +1,114 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/indexer/lumina/lumina_index_accumulator.h" + +#include + +#include "fmt/format.h" +#include "lumina/extensions/experimental/BuildCombinedExtensionV0.h" +#include "paimon/indexer/lumina/lumina_dataset.h" +#include "paimon/indexer/lumina/lumina_utils.h" + +namespace paimon::lumina { + +Status LuminaIndexAccumulator::AddBatch(const std::shared_ptr& struct_array, + const std::shared_ptr& vectors, + uint32_t dimension, + const std::vector& tag_fields, + int64_t first_row_id) { + // Split into contiguous non-null segments, skipping null rows in the list field. + int64_t segment_start = -1; + for (int64_t i = 0; i <= vectors->length(); ++i) { + bool is_null = i < vectors->length() && vectors->IsNull(i); + bool is_end = i == vectors->length(); + if (!is_null && !is_end && segment_start < 0) { + segment_start = i; + } + if ((is_null || is_end) && segment_start >= 0) { + int64_t segment_length = i - segment_start; + // Use value_offset to precisely locate the float range for this segment. + int64_t value_start = vectors->value_offset(segment_start); + int64_t value_end = vectors->value_offset(i); + std::shared_ptr values = + std::dynamic_pointer_cast( + vectors->values()->Slice(value_start, value_end - value_start)); + if (!values) { + return Status::Invalid( + "invalid sliced value array in LuminaIndexWriter, must be float array"); + } + if (values->null_count() != 0) { + return Status::Invalid( + "field value array in LuminaIndexWriter is invalid, must not null"); + } + for (int64_t row = segment_start; row < i; ++row) { + int64_t vector_length = vectors->value_offset(row + 1) - vectors->value_offset(row); + if (vector_length != static_cast(dimension)) { + return Status::Invalid(fmt::format( + "invalid input array in LuminaIndexWriter, vector at row [{}] has length " + "[{}], expected dimension [{}]", + row, vector_length, dimension)); + } + } + if (!tag_fields.empty()) { + PAIMON_ASSIGN_OR_RAISE( + std::vector<::lumina::extensions::experimental::TagDimensionData> tag_data, + LuminaTagUtils::ExtractTagDataForSegment(struct_array, tag_fields, + segment_start, segment_length)); + tag_data_vec_.push_back(std::move(tag_data)); + } + arrays_.push_back(std::move(values)); + array_start_ids_.push_back(first_row_id + segment_start); + indexed_count_ += segment_length; + segment_start = -1; + } + } + return Status::OK(); +} + +Result<::lumina::api::LuminaBuilder> LuminaIndexAccumulator::Build( + const ::lumina::api::BuilderOptions& builder_options, uint32_t dimension, bool with_tag, + LuminaMemoryPool* pool) { + ::lumina::core::MemoryResourceConfig memory_resource(pool); + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( + ::lumina::api::LuminaBuilder builder, + ::lumina::api::LuminaBuilder::Create(builder_options, memory_resource)); + + // Pretrain before inserting the accumulated vectors. + LuminaDataset pretrain_data(indexed_count_, dimension, arrays_, array_start_ids_); + PAIMON_RETURN_NOT_OK_FROM_LUMINA(builder.PretrainFrom(pretrain_data)); + + // insert data + if (!with_tag) { + LuminaDataset insert_data(indexed_count_, dimension, arrays_, array_start_ids_); + std::vector>().swap(arrays_); + PAIMON_RETURN_NOT_OK_FROM_LUMINA(builder.InsertFrom(insert_data)); + } else { + ::lumina::extensions::experimental::BuildWithTagExtension tag_extension; + PAIMON_RETURN_NOT_OK_FROM_LUMINA(builder.Attach(tag_extension)); + LuminaDatasetWithTag insert_data(indexed_count_, dimension, arrays_, array_start_ids_, + tag_data_vec_); + std::vector>().swap(arrays_); + std::vector>().swap( + tag_data_vec_); + PAIMON_RETURN_NOT_OK_FROM_LUMINA(tag_extension.InsertFromWithTag(insert_data)); + } + return std::move(builder); +} + +} // namespace paimon::lumina diff --git a/src/paimon/indexer/lumina/lumina_index_accumulator.h b/src/paimon/indexer/lumina/lumina_index_accumulator.h new file mode 100644 index 00000000..a7b6fe69 --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_index_accumulator.h @@ -0,0 +1,58 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +#include "arrow/api.h" +#include "lumina/api/LuminaBuilder.h" +#include "lumina/api/Options.h" +#include "lumina/extensions/experimental/DatasetWithTag.h" +#include "paimon/indexer/lumina/lumina_memory_pool.h" +#include "paimon/indexer/lumina/lumina_tag_utils.h" +#include "paimon/result.h" +#include "paimon/status.h" + +namespace paimon::lumina { + +/// Accumulates the non-null vector segments shared by Lumina Global Index and File Index writers. +class LuminaIndexAccumulator { + public: + Status AddBatch(const std::shared_ptr& struct_array, + const std::shared_ptr& vectors, uint32_t dimension, + const std::vector& tag_fields, int64_t first_row_id); + + Result<::lumina::api::LuminaBuilder> Build(const ::lumina::api::BuilderOptions& builder_options, + uint32_t dimension, bool with_tag, + LuminaMemoryPool* pool); + + int64_t IndexedCount() const { + return indexed_count_; + } + + private: + int64_t indexed_count_ = 0; + std::vector> arrays_; + std::vector array_start_ids_; + std::vector> tag_data_vec_; +}; + +} // namespace paimon::lumina diff --git a/src/paimon/indexer/lumina/lumina_index_options.cpp b/src/paimon/indexer/lumina/lumina_index_options.cpp new file mode 100644 index 00000000..e26b4e14 --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_index_options.cpp @@ -0,0 +1,114 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/indexer/lumina/lumina_index_options.h" + +#include + +#include "fmt/format.h" +#include "lumina/api/OptionsNormalize.h" +#include "lumina/core/Constants.h" +#include "paimon/common/utils/options_utils.h" +#include "paimon/indexer/lumina/lumina_utils.h" + +namespace paimon::lumina { + +Result LuminaIndexOptions::GetDimension( + const std::map& lumina_options) { + return OptionsUtils::GetValueFromMap(lumina_options, + std::string(::lumina::core::kDimension)); +} + +Result LuminaIndexOptions::GetIndexInfo( + const std::map& lumina_options) { + PAIMON_ASSIGN_OR_RAISE(uint32_t dimension, GetDimension(lumina_options)); + PAIMON_ASSIGN_OR_RAISE(std::string index_type, + OptionsUtils::GetValueFromMap( + lumina_options, std::string(::lumina::core::kIndexType))); + PAIMON_ASSIGN_OR_RAISE(std::string distance_type_str, + OptionsUtils::GetValueFromMap( + lumina_options, std::string(::lumina::core::kDistanceMetric))); + + VectorSearch::DistanceType distance_type = VectorSearch::DistanceType::UNKNOWN; + if (distance_type_str == ::lumina::core::kDistanceL2) { + distance_type = VectorSearch::DistanceType::EUCLIDEAN; + } else if (distance_type_str == ::lumina::core::kDistanceCosine) { + distance_type = VectorSearch::DistanceType::COSINE; + } else if (distance_type_str == ::lumina::core::kDistanceInnerProduct) { + distance_type = VectorSearch::DistanceType::INNER_PRODUCT; + } + if (distance_type == VectorSearch::DistanceType::UNKNOWN) { + return Status::Invalid( + fmt::format("invalid distance type {} for lumina", distance_type_str)); + } + + bool has_tag = lumina_options.find(std::string(::lumina::core::kExtensionTagSchema)) != + lumina_options.end(); + return LuminaIndexInfo{dimension, std::move(index_type), distance_type, has_tag}; +} + +Result<::lumina::api::BuilderOptions> LuminaIndexOptions::CreateBuilderOptions( + const std::map& lumina_options) { + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( + ::lumina::api::BuilderOptions builder_options, + ::lumina::api::NormalizeBuilderOptions(std::unordered_map( + lumina_options.begin(), lumina_options.end()))); + return builder_options; +} + +Result<::lumina::api::SearcherOptions> LuminaIndexOptions::CreateSearcherOptions( + const std::map& lumina_options, const LuminaIndexInfo& index_info) { + std::map options = lumina_options; + options[std::string(::lumina::core::kDimension)] = std::to_string(index_info.dimension); + options[std::string(::lumina::core::kIndexType)] = index_info.index_type; + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( + ::lumina::api::SearcherOptions searcher_options, + ::lumina::api::NormalizeSearcherOptions( + std::unordered_map(options.begin(), options.end()))); + return searcher_options; +} + +Result<::lumina::api::SearchOptions> LuminaIndexOptions::CreateSearchOptions( + const VectorSearch& vector_search, const LuminaIndexInfo& index_info) { + if (vector_search.distance_type && + vector_search.distance_type.value() != index_info.distance_type) { + return Status::Invalid("distance type for index and search not match"); + } + if (vector_search.query.size() != index_info.dimension) { + return Status::Invalid("dimension for index and search not match"); + } + + std::map lumina_options = OptionsUtils::FetchOptionsWithPrefix( + LuminaDefines::kOptionKeyPrefix, vector_search.options); + auto index_type_iter = lumina_options.find(std::string(::lumina::core::kIndexType)); + if (index_type_iter != lumina_options.end() && + index_type_iter->second != index_info.index_type) { + return Status::Invalid("index type for index and search not match"); + } + + lumina_options[std::string(::lumina::core::kTopK)] = std::to_string(vector_search.limit); + lumina_options[std::string(::lumina::core::kSearchThreadSafeFilter)] = "true"; + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( + ::lumina::api::SearchOptions search_options, + ::lumina::api::NormalizeSearchOptions(index_info.index_type, + std::unordered_map( + lumina_options.begin(), lumina_options.end()))); + return search_options; +} + +} // namespace paimon::lumina diff --git a/src/paimon/indexer/lumina/lumina_index_options.h b/src/paimon/indexer/lumina/lumina_index_options.h new file mode 100644 index 00000000..ff77f032 --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_index_options.h @@ -0,0 +1,63 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +#include "lumina/api/Options.h" +#include "paimon/predicate/vector_search.h" +#include "paimon/result.h" + +namespace paimon::lumina { + +struct LuminaIndexInfo { + uint32_t dimension; + std::string index_type; + VectorSearch::DistanceType distance_type; + bool has_tag; +}; + +/// Shared parsing and validation for normalized Lumina options. +/// +/// Input maps use native Lumina keys such as `index.dimension`. Callers remain responsible for +/// stripping their own configuration namespace (for example, the `lumina.` Global Index prefix). +class LuminaIndexOptions { + public: + LuminaIndexOptions() = delete; + ~LuminaIndexOptions() = delete; + + static Result GetDimension(const std::map& lumina_options); + + static Result GetIndexInfo( + const std::map& lumina_options); + + static Result<::lumina::api::BuilderOptions> CreateBuilderOptions( + const std::map& lumina_options); + + static Result<::lumina::api::SearcherOptions> CreateSearcherOptions( + const std::map& lumina_options, + const LuminaIndexInfo& index_info); + + static Result<::lumina::api::SearchOptions> CreateSearchOptions( + const VectorSearch& vector_search, const LuminaIndexInfo& index_info); +}; + +} // namespace paimon::lumina diff --git a/src/paimon/global_index/lumina/lumina_memory_pool.h b/src/paimon/indexer/lumina/lumina_memory_pool.h similarity index 100% rename from src/paimon/global_index/lumina/lumina_memory_pool.h rename to src/paimon/indexer/lumina/lumina_memory_pool.h diff --git a/src/paimon/indexer/lumina/lumina_search_utils.cpp b/src/paimon/indexer/lumina/lumina_search_utils.cpp new file mode 100644 index 00000000..765cace9 --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_search_utils.cpp @@ -0,0 +1,117 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/indexer/lumina/lumina_search_utils.h" + +#include + +#include "fmt/format.h" +#include "paimon/indexer/lumina/lumina_file_reader.h" +#include "paimon/indexer/lumina/lumina_tag_utils.h" +#include "paimon/indexer/lumina/lumina_utils.h" +#include "paimon/status.h" + +namespace paimon::lumina { + +Result LuminaSearchUtils::AttachSearchExtensions( + ::lumina::api::LuminaSearcher& searcher, const LuminaIndexInfo& index_info) { + auto searcher_with_filter = std::make_unique<::lumina::extensions::SearchWithFilterExtension>(); + PAIMON_RETURN_NOT_OK_FROM_LUMINA(searcher.Attach(*searcher_with_filter)); + + std::unique_ptr<::lumina::extensions::experimental::SearchWithTagExtension> searcher_with_tag; + if (index_info.has_tag) { + searcher_with_tag = + std::make_unique<::lumina::extensions::experimental::SearchWithTagExtension>(); + PAIMON_RETURN_NOT_OK_FROM_LUMINA(searcher.Attach(*searcher_with_tag)); + } + return LuminaSearchExtensions{std::move(searcher_with_filter), std::move(searcher_with_tag)}; +} + +Result LuminaSearchUtils::OpenSearcher( + const std::map& lumina_options, const LuminaIndexInfo& index_info, + const std::shared_ptr& input, LuminaMemoryPool* pool) { + PAIMON_ASSIGN_OR_RAISE(::lumina::api::SearcherOptions searcher_options, + LuminaIndexOptions::CreateSearcherOptions(lumina_options, index_info)); + ::lumina::core::MemoryResourceConfig memory_resource(pool); + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( + ::lumina::api::LuminaSearcher lumina_searcher, + ::lumina::api::LuminaSearcher::Create(searcher_options, memory_resource)); + auto searcher = std::make_unique<::lumina::api::LuminaSearcher>(std::move(lumina_searcher)); + auto file_reader = std::make_unique(input); + PAIMON_RETURN_NOT_OK_FROM_LUMINA( + searcher->Open(std::move(file_reader), ::lumina::api::IOOptions())); + if (searcher->GetMeta().dim != index_info.dimension) { + return Status::Invalid( + fmt::format("Lumina index dimension {} mismatch expected dimension {}", + searcher->GetMeta().dim, index_info.dimension)); + } + PAIMON_ASSIGN_OR_RAISE(LuminaSearchExtensions extensions, + AttachSearchExtensions(*searcher, index_info)); + return LuminaSearcherWithExtensions{std::move(searcher), std::move(extensions)}; +} + +Result<::lumina::api::LuminaSearcher::SearchResult> LuminaSearchUtils::ExecuteVectorSearch( + ::lumina::api::LuminaSearcher& searcher, + ::lumina::extensions::SearchWithFilterExtension& searcher_with_filter, + ::lumina::extensions::experimental::SearchWithTagExtension* searcher_with_tag, + const std::shared_ptr& vector_search, const LuminaIndexInfo& index_info, + LuminaMemoryPool& pool) { + if (!vector_search) { + return Status::Invalid("Lumina vector search must not be null"); + } + PAIMON_ASSIGN_OR_RAISE(::lumina::api::SearchOptions search_options, + LuminaIndexOptions::CreateSearchOptions(*vector_search, index_info)); + ::lumina::api::Query query(vector_search->query.data(), vector_search->query.size()); + + if (vector_search->predicate) { + if (!searcher_with_tag) { + return Status::Invalid("lumina index was not built with tag"); + } + PAIMON_ASSIGN_OR_RAISE(::lumina::extensions::experimental::TagFilter tag_filter, + LuminaTagUtils::PredicateToTagFilter(vector_search->predicate)); + if (!vector_search->pre_filter) { + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( + ::lumina::api::LuminaSearcher::SearchResult search_result, + searcher_with_tag->SearchWithTag(query, tag_filter, search_options, pool)); + return std::move(search_result); + } + auto filter = [pre_filter = vector_search->pre_filter]( + ::lumina::core::vector_id_t id) -> bool { return pre_filter(id); }; + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( + ::lumina::api::LuminaSearcher::SearchResult search_result, + searcher_with_tag->SearchWithTagAndFilter(query, tag_filter, filter, search_options, + pool)); + return std::move(search_result); + } + + if (!vector_search->pre_filter) { + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( + ::lumina::api::LuminaSearcher::SearchResult search_result, + searcher.Search(query, search_options, pool)); + return std::move(search_result); + } + auto filter = [pre_filter = vector_search->pre_filter](::lumina::core::vector_id_t id) -> bool { + return pre_filter(id); + }; + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( + ::lumina::api::LuminaSearcher::SearchResult search_result, + searcher_with_filter.SearchWithFilter(query, filter, search_options, pool)); + return std::move(search_result); +} + +} // namespace paimon::lumina diff --git a/src/paimon/indexer/lumina/lumina_search_utils.h b/src/paimon/indexer/lumina/lumina_search_utils.h new file mode 100644 index 00000000..41e03c58 --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_search_utils.h @@ -0,0 +1,72 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +#include "lumina/api/LuminaSearcher.h" +#include "lumina/extensions/SearchWithFilterExtension.h" +#include "lumina/extensions/experimental/SearchWithTagExtension.h" +#include "paimon/indexer/lumina/lumina_index_options.h" +#include "paimon/indexer/lumina/lumina_memory_pool.h" +#include "paimon/predicate/vector_search.h" +#include "paimon/result.h" + +namespace paimon { +class InputStream; +} + +namespace paimon::lumina { + +struct LuminaSearchExtensions { + std::unique_ptr<::lumina::extensions::SearchWithFilterExtension> searcher_with_filter; + std::unique_ptr<::lumina::extensions::experimental::SearchWithTagExtension> searcher_with_tag; +}; + +struct LuminaSearcherWithExtensions { + std::unique_ptr<::lumina::api::LuminaSearcher> searcher; + LuminaSearchExtensions extensions; +}; + +/// Shared Lumina vector-search dispatch for Global Index and File Index readers. +class LuminaSearchUtils { + public: + LuminaSearchUtils() = delete; + ~LuminaSearchUtils() = delete; + + static Result OpenSearcher( + const std::map& lumina_options, const LuminaIndexInfo& index_info, + const std::shared_ptr& input, LuminaMemoryPool* pool); + + static Result<::lumina::api::LuminaSearcher::SearchResult> ExecuteVectorSearch( + ::lumina::api::LuminaSearcher& searcher, + ::lumina::extensions::SearchWithFilterExtension& searcher_with_filter, + ::lumina::extensions::experimental::SearchWithTagExtension* searcher_with_tag, + const std::shared_ptr& vector_search, const LuminaIndexInfo& index_info, + LuminaMemoryPool& pool); + + private: + /// Attach the mandatory function filter and the configured optional tag extension. + static Result AttachSearchExtensions( + ::lumina::api::LuminaSearcher& searcher, const LuminaIndexInfo& index_info); +}; + +} // namespace paimon::lumina diff --git a/src/paimon/indexer/lumina/lumina_tag_utils.cpp b/src/paimon/indexer/lumina/lumina_tag_utils.cpp new file mode 100644 index 00000000..aceae5eb --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_tag_utils.cpp @@ -0,0 +1,517 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "paimon/indexer/lumina/lumina_tag_utils.h" + +#include +#include +#include +#include + +#include "fmt/format.h" +#include "lumina/core/Constants.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/predicate/compound_predicate.h" +#include "paimon/predicate/leaf_predicate.h" +#include "rapidjson/document.h" + +namespace paimon::lumina { +namespace { + +using TagDimensionData = ::lumina::extensions::experimental::TagDimensionData; +using TagFilter = ::lumina::extensions::experimental::TagFilter; +using TagValue = ::lumina::extensions::experimental::TagValue; +using TagValues = ::lumina::extensions::experimental::TagValues; + +Result GetRequiredStringMember(const rapidjson::Value& object, + const std::string& field_name, + const std::string& tag_label) { + rapidjson::Value::ConstMemberIterator iter = object.FindMember(field_name.c_str()); + if (iter == object.MemberEnd()) { + return Status::Invalid( + fmt::format("lumina tag_schema {} missing required field: {}", tag_label, field_name)); + } + if (!iter->value.IsString()) { + return Status::Invalid( + fmt::format("lumina tag_schema {} field {} must be string", tag_label, field_name)); + } + return std::string(iter->value.GetString(), iter->value.GetStringLength()); +} + +Result ParseTagField(const rapidjson::Value& object, const std::string& tag_label) { + if (!object.IsObject()) { + return Status::Invalid(fmt::format("lumina tag_schema {} must be object", tag_label)); + } + if (object.MemberCount() != 3) { + return Status::Invalid(fmt::format( + "lumina tag_schema {} must have exactly 3 fields: key_name, type, value_type", + tag_label)); + } + + PAIMON_ASSIGN_OR_RAISE(std::string key_name, + GetRequiredStringMember( + object, std::string(::lumina::core::kExtensionTagKName), tag_label)); + PAIMON_ASSIGN_OR_RAISE( + std::string type, + GetRequiredStringMember(object, std::string(::lumina::core::kExtensionTagType), tag_label)); + PAIMON_ASSIGN_OR_RAISE(std::string value_type, + GetRequiredStringMember( + object, std::string(::lumina::core::kExtensionTagVType), tag_label)); + if (key_name.empty()) { + return Status::Invalid( + fmt::format("lumina tag_schema {} key_name must not be empty", tag_label)); + } + + LuminaTagField::Type parsed_type; + if (type == std::string(::lumina::core::kExtensionTagTypeEnum)) { + parsed_type = LuminaTagField::Type::ENUM; + } else if (type == std::string(::lumina::core::kExtensionTagTypeRange)) { + parsed_type = LuminaTagField::Type::RANGE; + } else { + return Status::Invalid( + fmt::format("lumina tag_schema {} has unsupported type: {}", tag_label, type)); + } + + LuminaTagField::ValueType parsed_value_type; + if (value_type == std::string(::lumina::core::kExtensionTagVTypeInt32)) { + parsed_value_type = LuminaTagField::ValueType::INT32; + } else if (value_type == std::string(::lumina::core::kExtensionTagVTypeInt64)) { + parsed_value_type = LuminaTagField::ValueType::INT64; + } else if (value_type == std::string(::lumina::core::kExtensionTagVTypeFloat)) { + parsed_value_type = LuminaTagField::ValueType::FLOAT; + } else if (value_type == std::string(::lumina::core::kExtensionTagVTypeDouble)) { + parsed_value_type = LuminaTagField::ValueType::DOUBLE; + } else if (value_type == std::string(::lumina::core::kExtensionTagVTypeString)) { + parsed_value_type = LuminaTagField::ValueType::STRING; + } else { + return Status::Invalid(fmt::format("lumina tag_schema {} has unsupported value_type: {}", + tag_label, value_type)); + } + return LuminaTagField{key_name, parsed_type, parsed_value_type}; +} + +Status ValidateTagArrowType(const LuminaTagField& tag_field, + const std::shared_ptr& field_type) { + std::shared_ptr value_type = field_type; + std::shared_ptr list_type = + std::dynamic_pointer_cast(field_type); + if (list_type) { + value_type = list_type->value_type(); + } + + bool compatible = false; + switch (tag_field.value_type) { + case LuminaTagField::ValueType::INT32: + compatible = value_type->id() == arrow::Type::INT8 || + value_type->id() == arrow::Type::INT16 || + value_type->id() == arrow::Type::INT32; + break; + case LuminaTagField::ValueType::INT64: + compatible = value_type->id() == arrow::Type::INT64; + break; + case LuminaTagField::ValueType::FLOAT: + compatible = value_type->id() == arrow::Type::FLOAT; + break; + case LuminaTagField::ValueType::DOUBLE: + compatible = value_type->id() == arrow::Type::DOUBLE; + break; + case LuminaTagField::ValueType::STRING: + compatible = value_type->id() == arrow::Type::STRING; + break; + } + if (!compatible) { + return Status::Invalid( + fmt::format("lumina tag field {} type {} is not compatible with tag_schema value_type", + tag_field.name, field_type->ToString())); + } + return Status::OK(); +} + +template +void AppendPrimitiveTagValue(const std::shared_ptr& array, int64_t index, + std::vector* values) { + values->push_back( + static_cast(checked_cast(array.get())->Value(index))); +} + +template +Status AppendTagValue(const std::shared_ptr& array, int64_t index, + std::vector* values) { + if (array->IsNull(index)) { + return Status::OK(); + } + auto validate_array_type = [&](arrow::Type::type expected_type, + const char* value_type_name) -> Status { + if (array->type_id() != expected_type) { + return Status::Invalid(fmt::format("lumina {} tag field has unsupported arrow type {}", + value_type_name, array->type()->ToString())); + } + return Status::OK(); + }; + + if constexpr (std::is_same_v) { + switch (array->type_id()) { + case arrow::Type::INT8: + AppendPrimitiveTagValue(array, index, values); + break; + case arrow::Type::INT16: + AppendPrimitiveTagValue(array, index, values); + break; + case arrow::Type::INT32: + AppendPrimitiveTagValue(array, index, values); + break; + default: + return Status::Invalid( + fmt::format("lumina integer tag field has unsupported arrow type {}", + array->type()->ToString())); + } + } else if constexpr (std::is_same_v) { + PAIMON_RETURN_NOT_OK(validate_array_type(arrow::Type::INT64, "int64")); + AppendPrimitiveTagValue(array, index, values); + } else if constexpr (std::is_same_v) { + PAIMON_RETURN_NOT_OK(validate_array_type(arrow::Type::FLOAT, "float")); + AppendPrimitiveTagValue(array, index, values); + } else if constexpr (std::is_same_v) { + PAIMON_RETURN_NOT_OK(validate_array_type(arrow::Type::DOUBLE, "double")); + AppendPrimitiveTagValue(array, index, values); + } else if constexpr (std::is_same_v) { + PAIMON_RETURN_NOT_OK(validate_array_type(arrow::Type::STRING, "string")); + const arrow::StringArray* string_array = + checked_cast(array.get()); + std::string_view value = string_array->GetView(index); + values->emplace_back(value.data(), value.size()); + } else { + return Status::Invalid("lumina tag field has unsupported value type"); + } + return Status::OK(); +} + +template +Status ExtractTagValues(const std::shared_ptr& field_array, int64_t segment_start, + int64_t segment_len, std::vector>* values) { + values->resize(segment_len); + std::shared_ptr list_array = + std::dynamic_pointer_cast(field_array); + if (list_array) { + std::shared_ptr child_values = list_array->values(); + for (int64_t i = 0; i < segment_len; ++i) { + int64_t row = segment_start + i; + if (list_array->IsNull(row)) { + continue; + } + int64_t value_start = list_array->value_offset(row); + int64_t value_end = list_array->value_offset(row + 1); + std::vector& row_values = (*values)[i]; + row_values.reserve(value_end - value_start); + for (int64_t value_index = value_start; value_index < value_end; ++value_index) { + PAIMON_RETURN_NOT_OK(AppendTagValue(child_values, value_index, &row_values)); + } + } + return Status::OK(); + } + + for (int64_t i = 0; i < segment_len; ++i) { + PAIMON_RETURN_NOT_OK(AppendTagValue(field_array, segment_start + i, &(*values)[i])); + } + return Status::OK(); +} + +Result LiteralToTagValue(const Literal& literal) { + if (literal.IsNull()) { + return Status::Invalid("lumina tag predicate does not support null literal"); + } + switch (literal.GetType()) { + case FieldType::TINYINT: + return TagValue(static_cast(literal.GetValue())); + case FieldType::SMALLINT: + return TagValue(static_cast(literal.GetValue())); + case FieldType::INT: + return TagValue(literal.GetValue()); + case FieldType::BIGINT: + return TagValue(literal.GetValue()); + case FieldType::FLOAT: + return TagValue(literal.GetValue()); + case FieldType::DOUBLE: + return TagValue(literal.GetValue()); + case FieldType::STRING: + return TagValue(literal.GetValue()); + default: + return Status::Invalid( + fmt::format("lumina tag predicate does not support literal type {}", + static_cast(literal.GetType()))); + } +} + +Result GetSingleLiteral(const std::vector& literals, + const std::string& function_name) { + if (literals.size() != 1) { + return Status::Invalid( + fmt::format("lumina tag {} predicate requires one literal", function_name)); + } + return &literals[0]; +} + +template +Result LiteralsToTypedTagValues(const std::vector& literals) { + std::vector values; + values.reserve(literals.size()); + for (const Literal& literal : literals) { + PAIMON_ASSIGN_OR_RAISE(TagValue value, LiteralToTagValue(literal)); + ValueType* typed_value = std::get_if(&value); + if (!typed_value) { + return Status::Invalid( + "lumina tag predicate IN literals must have the same value type"); + } + values.push_back(std::move(*typed_value)); + } + return TagValues(std::move(values)); +} + +Result LiteralsToTagValues(const std::vector& literals) { + if (literals.empty()) { + return Status::Invalid("lumina tag predicate IN requires at least one literal"); + } + switch (literals[0].GetType()) { + case FieldType::TINYINT: + case FieldType::SMALLINT: + case FieldType::INT: + return LiteralsToTypedTagValues(literals); + case FieldType::BIGINT: + return LiteralsToTypedTagValues(literals); + case FieldType::FLOAT: + return LiteralsToTypedTagValues(literals); + case FieldType::DOUBLE: + return LiteralsToTypedTagValues(literals); + case FieldType::STRING: + return LiteralsToTypedTagValues(literals); + default: + return Status::Invalid( + fmt::format("lumina tag predicate IN does not support literal type {}", + static_cast(literals[0].GetType()))); + } +} + +} // namespace + +Result> LuminaTagUtils::ParseTagSchema( + const std::map& lumina_options) { + auto iter = lumina_options.find(std::string(::lumina::core::kExtensionTagSchema)); + if (iter == lumina_options.end()) { + return std::vector(); + } + + rapidjson::Document document; + document.Parse(iter->second.c_str()); + if (document.HasParseError()) { + return Status::Invalid("lumina tag_schema must be a valid JSON string"); + } + + std::vector tag_fields; + if (document.IsArray()) { + if (document.Empty()) { + return Status::Invalid("lumina tag_schema must contain at least one tag definition"); + } + tag_fields.reserve(document.Size()); + for (rapidjson::SizeType i = 0; i < document.Size(); ++i) { + PAIMON_ASSIGN_OR_RAISE(LuminaTagField field, + ParseTagField(document[i], fmt::format("tag[{}]", i))); + tag_fields.push_back(std::move(field)); + } + } else if (document.IsObject()) { + PAIMON_ASSIGN_OR_RAISE(LuminaTagField field, ParseTagField(document, "tag[0]")); + tag_fields.push_back(std::move(field)); + } else { + return Status::Invalid("lumina tag_schema must be an object or array of objects"); + } + + std::unordered_set seen_names; + for (const LuminaTagField& field : tag_fields) { + if (!seen_names.insert(field.name).second) { + return Status::Invalid( + fmt::format("lumina tag_schema has duplicate key_name: {}", field.name)); + } + } + return tag_fields; +} + +Result>> LuminaTagUtils::GetExtraFieldNames( + const std::map& lumina_options) { + PAIMON_ASSIGN_OR_RAISE(std::vector tag_fields, ParseTagSchema(lumina_options)); + if (tag_fields.empty()) { + return std::optional>(std::nullopt); + } + std::vector field_names; + field_names.reserve(tag_fields.size()); + for (const LuminaTagField& tag_field : tag_fields) { + field_names.push_back(tag_field.name); + } + return std::optional>(std::move(field_names)); +} + +Status LuminaTagUtils::ValidateTagFields(const arrow::StructType& struct_type, + const std::vector& tag_fields) { + for (const LuminaTagField& tag_field : tag_fields) { + std::shared_ptr field = struct_type.GetFieldByName(tag_field.name); + if (!field) { + return Status::Invalid( + fmt::format("lumina tag field {} not exist in arrow schema", tag_field.name)); + } + PAIMON_RETURN_NOT_OK(ValidateTagArrowType(tag_field, field->type())); + } + return Status::OK(); +} + +Result> LuminaTagUtils::ExtractTagDataForSegment( + const std::shared_ptr& struct_array, + const std::vector& tag_fields, int64_t segment_start, int64_t segment_len) { + std::vector tag_dimensions_data; + tag_dimensions_data.reserve(tag_fields.size()); + for (const LuminaTagField& tag_field : tag_fields) { + std::shared_ptr field_array = struct_array->GetFieldByName(tag_field.name); + if (!field_array) { + return Status::Invalid( + fmt::format("lumina tag field {} not in input array", tag_field.name)); + } + + TagDimensionData tag_dimension_data; + tag_dimension_data.tagkName = tag_field.name; + switch (tag_field.value_type) { + case LuminaTagField::ValueType::INT32: { + std::vector> values; + PAIMON_RETURN_NOT_OK( + ExtractTagValues(field_array, segment_start, segment_len, &values)); + tag_dimension_data.values = std::move(values); + break; + } + case LuminaTagField::ValueType::INT64: { + std::vector> values; + PAIMON_RETURN_NOT_OK( + ExtractTagValues(field_array, segment_start, segment_len, &values)); + tag_dimension_data.values = std::move(values); + break; + } + case LuminaTagField::ValueType::FLOAT: { + std::vector> values; + PAIMON_RETURN_NOT_OK( + ExtractTagValues(field_array, segment_start, segment_len, &values)); + tag_dimension_data.values = std::move(values); + break; + } + case LuminaTagField::ValueType::DOUBLE: { + std::vector> values; + PAIMON_RETURN_NOT_OK( + ExtractTagValues(field_array, segment_start, segment_len, &values)); + tag_dimension_data.values = std::move(values); + break; + } + case LuminaTagField::ValueType::STRING: { + std::vector> values; + PAIMON_RETURN_NOT_OK(ExtractTagValues(field_array, segment_start, + segment_len, &values)); + tag_dimension_data.values = std::move(values); + break; + } + } + tag_dimensions_data.push_back(std::move(tag_dimension_data)); + } + return tag_dimensions_data; +} + +Result LuminaTagUtils::PredicateToTagFilter( + const std::shared_ptr& predicate) { + if (!predicate) { + return Status::Invalid("lumina tag predicate must not be null"); + } + + std::shared_ptr compound_predicate = + std::dynamic_pointer_cast(predicate); + if (compound_predicate) { + std::vector children; + children.reserve(compound_predicate->Children().size()); + for (const std::shared_ptr& child : compound_predicate->Children()) { + PAIMON_ASSIGN_OR_RAISE(TagFilter tag_filter, PredicateToTagFilter(child)); + children.push_back(std::move(tag_filter)); + } + if (children.empty()) { + return Status::Invalid("lumina tag compound predicate must have at least one child"); + } + if (children.size() == 1) { + return std::move(children.front()); + } + switch (compound_predicate->GetFunction().GetType()) { + case Function::Type::AND: + return TagFilter::And(std::move(children)); + case Function::Type::OR: + return TagFilter::Or(std::move(children)); + default: + return Status::NotImplemented( + fmt::format("lumina tag predicate does not support compound function {}", + compound_predicate->GetFunction().ToString())); + } + } + + std::shared_ptr leaf_predicate = + std::dynamic_pointer_cast(predicate); + if (!leaf_predicate) { + return Status::Invalid( + fmt::format("cannot cast predicate {} to CompoundPredicate or LeafPredicate", + predicate->ToString())); + } + const std::vector& literals = leaf_predicate->Literals(); + const std::string& field_name = leaf_predicate->FieldName(); + switch (leaf_predicate->GetFunction().GetType()) { + case Function::Type::EQUAL: { + PAIMON_ASSIGN_OR_RAISE(const Literal* literal, GetSingleLiteral(literals, "equal")); + PAIMON_ASSIGN_OR_RAISE(TagValue value, LiteralToTagValue(*literal)); + return TagFilter::Eq(field_name, std::move(value)); + } + case Function::Type::GREATER_THAN: { + PAIMON_ASSIGN_OR_RAISE(const Literal* literal, + GetSingleLiteral(literals, "greater than")); + PAIMON_ASSIGN_OR_RAISE(TagValue value, LiteralToTagValue(*literal)); + return TagFilter::Gt(field_name, std::move(value)); + } + case Function::Type::GREATER_OR_EQUAL: { + PAIMON_ASSIGN_OR_RAISE(const Literal* literal, + GetSingleLiteral(literals, "greater or equal")); + PAIMON_ASSIGN_OR_RAISE(TagValue value, LiteralToTagValue(*literal)); + return TagFilter::Gte(field_name, std::move(value)); + } + case Function::Type::LESS_THAN: { + PAIMON_ASSIGN_OR_RAISE(const Literal* literal, GetSingleLiteral(literals, "less than")); + PAIMON_ASSIGN_OR_RAISE(TagValue value, LiteralToTagValue(*literal)); + return TagFilter::Lt(field_name, std::move(value)); + } + case Function::Type::LESS_OR_EQUAL: { + PAIMON_ASSIGN_OR_RAISE(const Literal* literal, + GetSingleLiteral(literals, "less or equal")); + PAIMON_ASSIGN_OR_RAISE(TagValue value, LiteralToTagValue(*literal)); + return TagFilter::Lte(field_name, std::move(value)); + } + case Function::Type::IN: { + PAIMON_ASSIGN_OR_RAISE(TagValues values, LiteralsToTagValues(literals)); + return TagFilter::In(field_name, std::move(values)); + } + default: + return Status::NotImplemented( + fmt::format("lumina tag predicate does not support leaf function {}", + leaf_predicate->GetFunction().ToString())); + } +} + +} // namespace paimon::lumina diff --git a/src/paimon/indexer/lumina/lumina_tag_utils.h b/src/paimon/indexer/lumina/lumina_tag_utils.h new file mode 100644 index 00000000..ff40cf9e --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_tag_utils.h @@ -0,0 +1,80 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "lumina/extensions/experimental/DatasetWithTag.h" +#include "lumina/extensions/experimental/TagFilter.h" +#include "paimon/predicate/predicate.h" +#include "paimon/result.h" +#include "paimon/status.h" + +namespace paimon::lumina { + +struct LuminaTagField { + enum class Type { + ENUM, + RANGE, + }; + + enum class ValueType { + INT32, + INT64, + FLOAT, + DOUBLE, + STRING, + }; + + std::string name; + Type type; + ValueType value_type; +}; + +/// Shared tag option, Arrow conversion, and predicate conversion helpers for Lumina indexes. +class LuminaTagUtils { + public: + LuminaTagUtils() = delete; + ~LuminaTagUtils() = delete; + + static Result> ParseTagSchema( + const std::map& lumina_options); + + static Result>> GetExtraFieldNames( + const std::map& lumina_options); + + static Status ValidateTagFields(const arrow::StructType& struct_type, + const std::vector& tag_fields); + + static Result> + ExtractTagDataForSegment(const std::shared_ptr& struct_array, + const std::vector& tag_fields, int64_t segment_start, + int64_t segment_len); + + static Result<::lumina::extensions::experimental::TagFilter> PredicateToTagFilter( + const std::shared_ptr& predicate); +}; + +} // namespace paimon::lumina diff --git a/src/paimon/global_index/lumina/lumina_utils.h b/src/paimon/indexer/lumina/lumina_utils.h similarity index 100% rename from src/paimon/global_index/lumina/lumina_utils.h rename to src/paimon/indexer/lumina/lumina_utils.h From a48e00b8b0df24b4dfbf74038fa67f8c76662845 Mon Sep 17 00:00:00 2001 From: "jinli.zjw" Date: Thu, 17 Sep 2026 13:13:12 +0800 Subject: [PATCH 2/3] feat(file-index): support PK table search with deletion vectors --- include/paimon/predicate/full_text_search.h | 5 +- ...complete_index_score_file_batch_reader.cpp | 2 + src/paimon/core/io/file_index_evaluator.cpp | 6 - .../core/operation/internal_read_context.cpp | 8 +- .../operation/internal_read_context_test.cpp | 42 +++++ .../core/operation/raw_file_split_read.cpp | 62 +++++-- .../operation/raw_file_split_read_test.cpp | 168 +++++++++++++++++- .../table/source/key_value_table_read.cpp | 34 ++++ .../core/table/source/key_value_table_read.h | 2 + .../core/table/source/table_read_test.cpp | 31 ++++ .../indexer/lumina/lumina_file_index_test.cpp | 9 +- 11 files changed, 338 insertions(+), 31 deletions(-) diff --git a/include/paimon/predicate/full_text_search.h b/include/paimon/predicate/full_text_search.h index 3e0676c5..59551be8 100644 --- a/include/paimon/predicate/full_text_search.h +++ b/include/paimon/predicate/full_text_search.h @@ -92,8 +92,9 @@ struct PAIMON_EXPORT FullTextSearch { std::string query; /// Type of search to perform. SearchType search_type; - /// A pre-filter based on **global row IDs**, implemented by leveraging another global index. - /// Only rows whose global row ID is present in `pre_filter` will be included during search. + /// A pre-filter whose row ID domain is determined by the reader: global row IDs for + /// GlobalIndexReader and physical file-local row IDs for FileIndexReader. + /// Only rows whose row ID is present in `pre_filter` will be included during search. /// If not set, all rows will be included. std::optional pre_filter; /// Whether to compute and return relevance scores (e.g. BM25). The 4-path matrix: diff --git a/src/paimon/common/reader/complete_index_score_file_batch_reader.cpp b/src/paimon/common/reader/complete_index_score_file_batch_reader.cpp index cce919d9..9bef43d2 100644 --- a/src/paimon/common/reader/complete_index_score_file_batch_reader.cpp +++ b/src/paimon/common/reader/complete_index_score_file_batch_reader.cpp @@ -21,6 +21,8 @@ #include +#include "arrow/c/abi.h" + namespace paimon { CompleteIndexScoreFileBatchReader::CompleteIndexScoreFileBatchReader( std::unique_ptr&& reader, const std::vector& scores, diff --git a/src/paimon/core/io/file_index_evaluator.cpp b/src/paimon/core/io/file_index_evaluator.cpp index 7ffa3b32..74d9d5ee 100644 --- a/src/paimon/core/io/file_index_evaluator.cpp +++ b/src/paimon/core/io/file_index_evaluator.cpp @@ -305,12 +305,6 @@ Result> FileIndexEvaluator::EvaluateFullTextSea const std::shared_ptr& path_factory, const std::shared_ptr& file_meta, const std::shared_ptr& file_system, const std::shared_ptr& pool) { - if (full_text_search->pre_filter) { - return Status::NotImplemented("File full-text search does not support pre-filter yet"); - } - if (full_text_search->with_score) { - return Status::NotImplemented("File full-text search does not support score output yet"); - } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input_stream, ExtractIndexInputStream(/*only_use_embedding_index=*/false, path_factory, file_meta, file_system)); diff --git a/src/paimon/core/operation/internal_read_context.cpp b/src/paimon/core/operation/internal_read_context.cpp index 2d7ef6d0..a75bc6ab 100644 --- a/src/paimon/core/operation/internal_read_context.cpp +++ b/src/paimon/core/operation/internal_read_context.cpp @@ -211,9 +211,13 @@ Result> InternalReadContext::Create( if (!core_options.FileIndexReadEnabled()) { return Status::Invalid("File Index read must be enabled for file-local search"); } - if (!table_schema->PrimaryKeys().empty()) { + if (context->GetFullTextSearch() && context->GetFullTextSearch()->with_score) { return Status::NotImplemented( - "File Index search currently supports append-only tables only"); + "File full-text search does not support score output yet"); + } + if (!table_schema->PrimaryKeys().empty() && !core_options.DeletionVectorsEnabled()) { + return Status::NotImplemented( + "File Index search on primary-key tables requires deletion vectors"); } if (core_options.DataEvolutionEnabled()) { return Status::NotImplemented( diff --git a/src/paimon/core/operation/internal_read_context_test.cpp b/src/paimon/core/operation/internal_read_context_test.cpp index 2ba5efbe..e7e9bae7 100644 --- a/src/paimon/core/operation/internal_read_context_test.cpp +++ b/src/paimon/core/operation/internal_read_context_test.cpp @@ -18,7 +18,12 @@ #include "paimon/core/operation/internal_read_context.h" +#include +#include +#include +#include #include +#include #include "arrow/c/bridge.h" #include "arrow/type.h" @@ -29,6 +34,7 @@ #include "paimon/data/shredding/map_shared_shredding_schema_utils.h" #include "paimon/defs.h" #include "paimon/fs/local/local_file_system.h" +#include "paimon/predicate/full_text_search.h" #include "paimon/predicate/vector_search.h" #include "paimon/status.h" #include "paimon/testing/utils/testharness.h" @@ -197,6 +203,42 @@ TEST(InternalReadContext, TestFileIndexSearchIndexScoreProjection) { } } +TEST(InternalReadContext, TestPrimaryKeyFileIndexSearchRequiresDeletionVectors) { + std::string path = paimon::test::GetDataDir() + + "/orc/pk_table_with_dv_cardinality.db/pk_table_with_dv_cardinality"; + SchemaManager schema_manager(std::make_shared(), path); + ASSERT_OK_AND_ASSIGN(std::shared_ptr table_schema, schema_manager.ReadSchema(0)); + std::shared_ptr search = + std::make_shared("f3", /*limit=*/1, std::vector{1.0f}, nullptr, + nullptr, std::nullopt, std::map{}); + ReadContextBuilder builder(path); + builder.SetReadFieldNames({"f0", "_INDEX_SCORE"}).SetVectorSearch(search); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, builder.Finish()); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr internal, + InternalReadContext::Create(context, table_schema, table_schema->Options())); + EXPECT_TRUE(internal->HasFileIndexSearch()); + + std::map options = table_schema->Options(); + options[Options::DELETION_VECTORS_ENABLED] = "false"; + ASSERT_NOK_WITH_MSG(InternalReadContext::Create(context, table_schema, options), + "primary-key tables requires deletion vectors"); +} + +TEST(InternalReadContext, TestFileFullTextSearchScoreIsNotSupported) { + std::string path = paimon::test::GetDataDir() + "/orc/append_09.db/append_09"; + SchemaManager schema_manager(std::make_shared(), path); + ASSERT_OK_AND_ASSIGN(std::shared_ptr table_schema, schema_manager.ReadSchema(0)); + auto search = std::make_shared( + "f0", /*limit=*/1, "query", FullTextSearch::SearchType::MATCH_ALL, + /*pre_filter=*/std::nullopt, /*with_score=*/true); + ReadContextBuilder builder(path); + builder.SetReadFieldNames({"f0"}).SetFullTextSearch(search); + ASSERT_OK_AND_ASSIGN(std::shared_ptr context, builder.Finish()); + ASSERT_NOK_WITH_MSG(InternalReadContext::Create(context, table_schema, table_schema->Options()), + "does not support score output"); +} + TEST(InternalReadContext, TestReadWithValueKindField) { std::string path = paimon::test::GetDataDir() + "/orc/append_09.db/append_09"; ReadContextBuilder context_builder(path); diff --git a/src/paimon/core/operation/raw_file_split_read.cpp b/src/paimon/core/operation/raw_file_split_read.cpp index c03a2970..e5906f8b 100644 --- a/src/paimon/core/operation/raw_file_split_read.cpp +++ b/src/paimon/core/operation/raw_file_split_read.cpp @@ -46,10 +46,13 @@ #include "paimon/file_index/file_index_result.h" #include "paimon/file_index/scored_file_index_result.h" #include "paimon/memory/memory_pool.h" +#include "paimon/predicate/full_text_search.h" +#include "paimon/predicate/vector_search.h" #include "paimon/reader/file_batch_reader.h" #include "paimon/status.h" #include "paimon/table/source/data_split.h" #include "paimon/utils/roaring_bitmap32.h" +#include "paimon/utils/roaring_bitmap64.h" namespace paimon { class DataFilePathFactory; @@ -206,33 +209,64 @@ Result> RawFileSplitRead::ApplyIndexAndDvReader return Status::NotImplemented( "File Index search does not support indexed split row ranges yet"); } + std::shared_ptr deletion_vector; if (dv_factory) { - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr deletion_vector, - dv_factory(file->file_name)); - if (deletion_vector && !deletion_vector->IsEmpty()) { - return Status::NotImplemented( - "File Index search does not support deletion vectors yet"); - } + PAIMON_ASSIGN_OR_RAISE(deletion_vector, dv_factory(file->file_name)); + } + const RoaringBitmap32* deletion = nullptr; + if (auto* bitmap_dv = dynamic_cast(deletion_vector.get())) { + deletion = bitmap_dv->GetBitmap(); + } else if (deletion_vector && !deletion_vector->IsEmpty()) { + return Status::NotImplemented( + "File Index search only supports bitmap32 deletion vectors"); + } + if (deletion && deletion->Cardinality() == file->row_count) { + return std::unique_ptr(); } std::optional search_selection; if (context_->GetVectorSearch()) { - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr search_result, - FileIndexEvaluator::EvaluateVectorSearch( - data_schema, options_, context_->GetVectorSearch(), data_file_path_factory, - file, options_.GetFileSystem(), pool_)); + std::shared_ptr vector_search = context_->GetVectorSearch(); + if (deletion && !deletion->IsEmpty()) { + // Keep the deletion vector alive for the search callback without copying the + // bitmap. + std::shared_ptr deleted_rows(deletion_vector, deletion); + auto user_filter = vector_search->pre_filter; + vector_search = vector_search->ReplacePreFilter( + [deleted_rows, user_filter, row_count = file->row_count](int64_t row_id) { + return row_id >= 0 && row_id < row_count && + !deleted_rows->Contains(static_cast(row_id)) && + (!user_filter || user_filter(row_id)); + }); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr search_result, + FileIndexEvaluator::EvaluateVectorSearch( + data_schema, options_, vector_search, data_file_path_factory, + file, options_.GetFileSystem(), pool_)); if (search_result->IsEmpty()) { return std::unique_ptr(); } search_selection = search_result->GetRowPositions(); *index_scores = search_result->GetScores(); } else { + std::shared_ptr full_text_search = context_->GetFullTextSearch(); + if (deletion && !deletion->IsEmpty()) { + RoaringBitmap64 eligible; + eligible.AddRange(0, file->row_count); + eligible -= RoaringBitmap64(*deletion); + if (full_text_search->pre_filter) { + eligible &= full_text_search->pre_filter.value(); + } + if (eligible.IsEmpty()) { + return std::unique_ptr(); + } + full_text_search = full_text_search->ReplacePreFilter(eligible); + } PAIMON_ASSIGN_OR_RAISE( std::shared_ptr search_result, - FileIndexEvaluator::EvaluateFullTextSearch( - data_schema, options_, context_->GetFullTextSearch(), data_file_path_factory, - file, options_.GetFileSystem(), pool_)); + FileIndexEvaluator::EvaluateFullTextSearch(data_schema, options_, full_text_search, + data_file_path_factory, file, + options_.GetFileSystem(), pool_)); std::shared_ptr bitmap_result = std::dynamic_pointer_cast(search_result); if (!bitmap_result) { diff --git a/src/paimon/core/operation/raw_file_split_read_test.cpp b/src/paimon/core/operation/raw_file_split_read_test.cpp index 5d72e2c2..a8802b9a 100644 --- a/src/paimon/core/operation/raw_file_split_read_test.cpp +++ b/src/paimon/core/operation/raw_file_split_read_test.cpp @@ -35,6 +35,7 @@ #include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/core/core_options.h" +#include "paimon/core/deletionvectors/bitmap_deletion_vector.h" #include "paimon/core/global_index/indexed_split_impl.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/manifest/file_source.h" @@ -47,6 +48,7 @@ #include "paimon/data/timestamp.h" #include "paimon/executor.h" #include "paimon/factories/factory.h" +#include "paimon/file_index/bitmap_index_result.h" #include "paimon/file_index/file_index_format.h" #include "paimon/file_index/file_index_reader.h" #include "paimon/file_index/file_indexer.h" @@ -55,6 +57,7 @@ #include "paimon/format/file_format.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/memory/memory_pool.h" +#include "paimon/predicate/full_text_search.h" #include "paimon/predicate/vector_search.h" #include "paimon/read_context.h" #include "paimon/status.h" @@ -67,12 +70,25 @@ namespace paimon::test { namespace { constexpr char kRawVectorSearchIndexType[] = "raw-vector-search-test"; +constexpr char kRawFullTextSearchIndexType[] = "raw-full-text-search-test"; class RawVectorSearchTestReader final : public FileIndexReader { public: Result> VisitVectorSearch( - const std::shared_ptr&) override { - return ScoredFileIndexResult::Create(RoaringBitmap32::From({0, 2}), {0.25f, 0.75f}); + const std::shared_ptr& search) override { + RoaringBitmap32 positions; + std::vector scores; + for (const auto& [row_id, score] : + std::vector>{{0, 0.25f}, {2, 0.75f}}) { + if (!search->pre_filter || search->pre_filter(row_id)) { + positions.Add(row_id); + scores.push_back(score); + if (scores.size() == static_cast(search->limit)) { + break; + } + } + } + return ScoredFileIndexResult::Create(std::move(positions), std::move(scores)); } }; @@ -104,9 +120,56 @@ class RawVectorSearchTestFactory final : public FileIndexerFactory { REGISTER_PAIMON_FACTORY(RawVectorSearchTestFactory); -Result> MakeRawVectorSearchIndex(const std::shared_ptr& pool) { +class RawFullTextSearchTestReader final : public FileIndexReader { + public: + Result> VisitFullTextSearch( + const std::shared_ptr& search) override { + RoaringBitmap32 positions; + for (int32_t row_id : {0, 2}) { + if (!search->pre_filter || search->pre_filter->Contains(row_id)) { + positions.Add(row_id); + if (search->limit && positions.Cardinality() == search->limit.value()) { + break; + } + } + } + return std::make_shared( + [bitmap = std::move(positions)]() -> Result { return bitmap; }); + } +}; + +class RawFullTextSearchTestIndexer final : public FileIndexer { + public: + Result> CreateReader( + ::ArrowSchema*, int32_t, int32_t, const std::shared_ptr&, + const std::shared_ptr&) const override { + return std::make_shared(); + } + + Result> CreateWriter( + ::ArrowSchema*, const std::shared_ptr&) const override { + return Status::NotImplemented("Raw full-text search test index is read-only"); + } +}; + +class RawFullTextSearchTestFactory final : public FileIndexerFactory { + public: + const char* Identifier() const override { + return kRawFullTextSearchIndexType; + } + + Result> Create( + const std::map&) const override { + return std::make_unique(); + } +}; + +REGISTER_PAIMON_FACTORY(RawFullTextSearchTestFactory); + +Result> MakeRawSearchIndex(const std::string& index_type, + const std::shared_ptr& pool) { FileIndexFormat::ColumnIndexes indexes; - indexes["f3"][kRawVectorSearchIndexType] = std::make_shared("index", pool.get()); + indexes["f3"][index_type] = std::make_shared("index", pool.get()); auto segment_output = std::make_unique( MemorySegmentOutputStream::DEFAULT_SEGMENT_SIZE, pool); auto output = std::make_shared(std::move(segment_output)); @@ -121,6 +184,7 @@ Result> MakeRawVectorSearchIndex(const std::shared_ptr> PrepareDataSplits() const { auto meta1 = std::make_shared( "data-01b6a930-6564-409b-b8f4-ed1307790d72-0.orc", /*file_size=*/575, /*row_count=*/3, @@ -198,6 +262,52 @@ class RawFileSplitReadTest : public ::testing::Test { return data_splits; } + Result> ReadSearchWithDeletionVector( + const std::shared_ptr& read_context, const std::string& index_type, + const RoaringBitmap32& deleted_rows) const { + SchemaManager schema_manager(std::make_shared(), read_context->GetPath()); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr table_schema, + schema_manager.ReadSchema(0)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr internal_context, + InternalReadContext::Create(read_context, table_schema, table_schema->Options())); + auto split = std::dynamic_pointer_cast(PrepareDataSplits().front()); + if (!split) { + return Status::Invalid("Raw search test split must be a DataSplitImpl"); + } + PAIMON_ASSIGN_OR_RAISE(split->DataFiles().front()->embedded_index, + MakeRawSearchIndex(index_type, pool_)); + + const CoreOptions& options = internal_context->GetCoreOptions(); + std::shared_ptr schema = + DataField::ConvertDataFieldsToArrowSchema(table_schema->Fields()); + PAIMON_ASSIGN_OR_RAISE(std::vector external_paths, + options.CreateExternalPaths()); + PAIMON_ASSIGN_OR_RAISE(std::optional global_index_path, + options.CreateGlobalIndexExternalPath()); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr path_factory, + FileStorePathFactory::Create( + read_context->GetPath(), schema, table_schema->PartitionKeys(), + options.GetPartitionDefaultName(), options.GetFileFormat()->Identifier(), + options.DataFilePrefix(), options.LegacyPartitionNameEnabled(), external_paths, + global_index_path, options.IndexFileInDataFileDir(), pool_)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr executor, + CreateDefaultExecutor(/*thread_count=*/2)); + std::shared_ptr shared_context(std::move(internal_context)); + RawFileSplitRead split_read(path_factory, shared_context, pool_, executor); + auto dv = std::make_shared(deleted_rows); + DeletionVector::Factory dv_factory = + [dv](const std::string&) -> Result> { + return std::shared_ptr(dv); + }; + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr reader, + split_read.CreateReader(split->Partition(), split->Bucket(), split->DataFiles(), + dv_factory, std::nullopt)); + return ReadResultCollector::CollectResult(std::move(reader)); + } + void CheckReadResult(const std::shared_ptr& read_schema, const std::shared_ptr& expected_array) const { std::string path = paimon::test::GetDataDir() + @@ -247,7 +357,7 @@ class RawFileSplitReadTest : public ::testing::Test { ASSERT_TRUE(result_array->Equals(expected_array)); } - private: + protected: std::shared_ptr pool_ = GetDefaultPool(); }; @@ -466,7 +576,7 @@ TEST_F(RawFileSplitReadTest, TestVectorSearchScoreProjection) { auto data_split = std::dynamic_pointer_cast(data_splits.front()); ASSERT_TRUE(data_split); ASSERT_OK_AND_ASSIGN(data_split->DataFiles().front()->embedded_index, - MakeRawVectorSearchIndex(pool_)); + MakeRawSearchIndex(kRawVectorSearchIndexType, pool_)); const CoreOptions& core_options = internal_context->GetCoreOptions(); std::shared_ptr data_schema = @@ -522,6 +632,52 @@ TEST_F(RawFileSplitReadTest, TestVectorSearchScoreProjection) { ASSERT_TRUE(without_score_result->Equals(without_score_expected)); } +TEST_F(RawFileSplitReadTest, TestVectorSearchFiltersDeletionVectorBeforeTopK) { + std::string path = paimon::test::GetDataDir() + + "/orc/multi_partition_append_table.db/multi_partition_append_table"; + auto user_filter = [](int64_t row_id) { return row_id != 1; }; + std::shared_ptr search = std::make_shared( + "f3", /*limit=*/1, std::vector{1.0f}, user_filter, nullptr, + VectorSearch::DistanceType::EUCLIDEAN, std::map{}); + ReadContextBuilder builder(path); + builder.SetReadFieldNames({"_INDEX_SCORE", "f0"}).SetVectorSearch(search); + ASSERT_OK_AND_ASSIGN(std::shared_ptr read_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + ReadSearchWithDeletionVector(read_context, kRawVectorSearchIndexType, + RoaringBitmap32::From({0}))); + std::shared_ptr expected; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( + arrow::struct_({arrow::field("_VALUE_KIND", arrow::int8()), + arrow::field("_INDEX_SCORE", arrow::float32()), + arrow::field("f0", arrow::utf8())}), + {R"([[0, 0.75, "Tony"]])"}, &expected) + .ok()); + EXPECT_TRUE(result->Equals(expected)); +} + +TEST_F(RawFileSplitReadTest, TestFullTextSearchFiltersDeletionVectorBeforeLimit) { + std::string path = paimon::test::GetDataDir() + + "/orc/multi_partition_append_table.db/multi_partition_append_table"; + RoaringBitmap64 user_filter; + user_filter.Add(0); + user_filter.Add(2); + auto search = std::make_shared( + "f3", /*limit=*/1, "match", FullTextSearch::SearchType::MATCH_ALL, user_filter); + ReadContextBuilder builder(path); + builder.SetReadFieldNames({"f0"}).SetFullTextSearch(search); + ASSERT_OK_AND_ASSIGN(std::shared_ptr read_context, builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, + ReadSearchWithDeletionVector(read_context, kRawFullTextSearchIndexType, + RoaringBitmap32::From({0}))); + std::shared_ptr expected; + ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON( + arrow::struct_({arrow::field("_VALUE_KIND", arrow::int8()), + arrow::field("f0", arrow::utf8())}), + {R"([[0, "Tony"]])"}, &expected) + .ok()); + EXPECT_TRUE(result->Equals(expected)); +} + TEST_F(RawFileSplitReadTest, TestEmptyPlan) { std::string path = paimon::test::GetDataDir() + "/orc/multi_partition_append_table.db/" diff --git a/src/paimon/core/table/source/key_value_table_read.cpp b/src/paimon/core/table/source/key_value_table_read.cpp index 9feae79e..904e5bc4 100644 --- a/src/paimon/core/table/source/key_value_table_read.cpp +++ b/src/paimon/core/table/source/key_value_table_read.cpp @@ -193,8 +193,37 @@ void KeyValueTableRead::ForceKeepDelete(bool force_keep_delete) { } } +Status KeyValueTableRead::ValidateFileIndexSearchSplit(const std::shared_ptr& split) const { + if (std::dynamic_pointer_cast(split)) { + return Status::NotImplemented( + "File Index search does not support indexed split row ranges"); + } + if (std::dynamic_pointer_cast(split)) { + return Status::NotImplemented("File Index search does not support real-time splits"); + } + auto data_split = std::dynamic_pointer_cast(split); + if (!data_split) { + return Status::Invalid("File Index search requires a data split"); + } + for (const auto& read : split_reads_) { + if (dynamic_cast(read.get())) { + PAIMON_ASSIGN_OR_RAISE(bool matched, read->Match(split, force_keep_delete_)); + if (matched) { + return Status::OK(); + } + break; + } + } + return Status::NotImplemented( + "File Index search requires a non-streaming, raw-convertible primary-key split " + "without physical delete rows"); +} + Result> KeyValueTableRead::CreateReader( const std::shared_ptr& split) { + if (context_->HasFileIndexSearch()) { + PAIMON_RETURN_NOT_OK(ValidateFileIndexSearchSplit(split)); + } std::shared_ptr realtime_split = std::dynamic_pointer_cast(split); if (realtime_split) { return CreateRealtimeReader(realtime_split, /*release_ticket=*/true); @@ -253,6 +282,11 @@ Result> KeyValueTableRead::CreateReader( Result> KeyValueTableRead::CreateReader( const std::vector>& splits) { + if (context_->HasFileIndexSearch()) { + for (const auto& split : splits) { + PAIMON_RETURN_NOT_OK(ValidateFileIndexSearchSplit(split)); + } + } std::vector> readers; readers.reserve(splits.size()); std::vector> realtime_splits; diff --git a/src/paimon/core/table/source/key_value_table_read.h b/src/paimon/core/table/source/key_value_table_read.h index 43164739..09a0ac6b 100644 --- a/src/paimon/core/table/source/key_value_table_read.h +++ b/src/paimon/core/table/source/key_value_table_read.h @@ -65,6 +65,8 @@ class KeyValueTableRead : public TableRead { Result> CreateRealtimeReader( const std::shared_ptr& realtime_split, bool release_ticket); + Status ValidateFileIndexSearchSplit(const std::shared_ptr& split) const; + std::vector> split_reads_; std::shared_ptr path_factory_; std::shared_ptr context_; diff --git a/src/paimon/core/table/source/table_read_test.cpp b/src/paimon/core/table/source/table_read_test.cpp index f2a8644e..12e6452c 100644 --- a/src/paimon/core/table/source/table_read_test.cpp +++ b/src/paimon/core/table/source/table_read_test.cpp @@ -27,16 +27,20 @@ #include #include "gtest/gtest.h" +#include "paimon/common/data/binary_row.h" #include "paimon/core/core_options.h" #include "paimon/core/operation/abstract_split_read.h" #include "paimon/core/operation/split_read.h" #include "paimon/core/table/source/append_only_table_read.h" +#include "paimon/core/table/source/data_split_impl.h" #include "paimon/core/table/source/key_value_table_read.h" #include "paimon/defs.h" #include "paimon/predicate/literal.h" #include "paimon/predicate/predicate_builder.h" +#include "paimon/predicate/vector_search.h" #include "paimon/read_context.h" #include "paimon/status.h" +#include "paimon/testing/utils/binary_row_generator.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { @@ -125,6 +129,33 @@ TEST(TableReadTest, TestCreateKeyValueTableRead) { ASSERT_TRUE(key_value_table_read); } +TEST(TableReadTest, TestPrimaryKeyFileIndexSearchDoesNotFallBackToMergeReader) { + std::string path = paimon::test::GetDataDir() + + "/orc/pk_table_with_dv_cardinality.db/pk_table_with_dv_cardinality/"; + auto search = + std::make_shared("f3", /*limit=*/1, std::vector{1.0f}, nullptr, + nullptr, std::nullopt, std::map{}); + ReadContextBuilder context_builder(path); + context_builder.SetReadFieldNames({"f0"}).SetVectorSearch(search); + ASSERT_OK_AND_ASSIGN(std::unique_ptr context, context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr reader, TableRead::Create(std::move(context))); + + BinaryRow partition = BinaryRowGenerator::GenerateRow({10}, GetDefaultPool().get()); + DataSplitImpl::Builder builder(partition, /*bucket=*/0, /*bucket_path=*/"", /*data_files=*/{}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr non_raw, + builder.WithSnapshot(1).IsStreaming(false).RawConvertible(false).Build()); + ASSERT_NOK_WITH_MSG(reader->CreateReader(non_raw), "raw-convertible primary-key split"); + + DataSplitImpl::Builder raw_builder(partition, /*bucket=*/0, /*bucket_path=*/"", + /*data_files=*/{}); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr raw, + raw_builder.WithSnapshot(1).IsStreaming(false).RawConvertible(true).Build()); + ASSERT_OK(reader->CreateReader(raw)); + std::vector> splits = {raw, non_raw}; + ASSERT_NOK_WITH_MSG(reader->CreateReader(splits), "raw-convertible primary-key split"); +} + TEST(TableReadTest, TestCreateAppendOnlyTableRead) { std::string path = paimon::test::GetDataDir() + "/orc/append_09.db/append_09"; ReadContextBuilder context_builder(path); diff --git a/src/paimon/indexer/lumina/lumina_file_index_test.cpp b/src/paimon/indexer/lumina/lumina_file_index_test.cpp index 1a7b6203..15f37be9 100644 --- a/src/paimon/indexer/lumina/lumina_file_index_test.cpp +++ b/src/paimon/indexer/lumina/lumina_file_index_test.cpp @@ -122,6 +122,13 @@ TEST(LuminaFileIndexTest, RoundTripUsesFileLocalRowPositions) { ASSERT_EQ(1, result->GetScores().size()); EXPECT_FLOAT_EQ(0.0f, result->GetScores()[0]); + std::shared_ptr filtered_search = + search->ReplacePreFilter([](int64_t file_row) { return file_row == 0; }); + ASSERT_OK_AND_ASSIGN(result, reader->VisitVectorSearch(filtered_search)); + EXPECT_EQ(RoaringBitmap32::From({0}), result->GetRowPositions()); + ASSERT_EQ(1, result->GetScores().size()); + EXPECT_FLOAT_EQ(4.0f, result->GetScores()[0]); + std::shared_ptr predicate = PredicateBuilder::Equal( /*field_index=*/1, /*field_name=*/"color", FieldType::STRING, Literal(FieldType::STRING, "red", 3)); @@ -186,7 +193,7 @@ TEST(LuminaFileIndexTest, RoundTripWithTagPredicate) { EXPECT_NEAR(0.01f, result->GetScores()[0], 1e-5f); std::shared_ptr filtered_search = std::make_shared( - "embedding", /*limit=*/4, std::vector{1.0f, 1.0f, 1.0f, 1.0f}, + "embedding", /*limit=*/1, std::vector{1.0f, 1.0f, 1.0f, 1.0f}, [](int64_t file_row) { return file_row == 0; }, predicate, VectorSearch::DistanceType::EUCLIDEAN, std::map{}); ASSERT_OK_AND_ASSIGN(result, reader->VisitVectorSearch(filtered_search)); From 5eefab65e9d7efd58c484f3da0c33dce5160d5dc Mon Sep 17 00:00:00 2001 From: "jinli.zjw" Date: Fri, 18 Sep 2026 15:54:23 +0800 Subject: [PATCH 3/3] fix(test): resolve file index search CI failures --- src/paimon/common/io/byte_array_output_stream_test.cpp | 2 +- .../reader/complete_index_score_batch_reader_test.cpp | 2 -- src/paimon/core/operation/raw_file_split_read_test.cpp | 10 ++++++++-- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/src/paimon/common/io/byte_array_output_stream_test.cpp b/src/paimon/common/io/byte_array_output_stream_test.cpp index c4adb84f..0fce3a3b 100644 --- a/src/paimon/common/io/byte_array_output_stream_test.cpp +++ b/src/paimon/common/io/byte_array_output_stream_test.cpp @@ -50,7 +50,7 @@ TEST(ByteArrayOutputStreamTest, TestWriteAndFinish) { ASSERT_NOK_WITH_MSG(stream->Finish(pool.get()), "already been finished"); stream.reset(); ASSERT_EQ("abcdef", std::string(result->data(), result->size())); - ASSERT_EQ(6, pool->CurrentUsage()); + ASSERT_EQ(result->size() + sizeof(Bytes), pool->CurrentUsage()); result.reset(); ASSERT_EQ(0, pool->CurrentUsage()); } diff --git a/src/paimon/common/reader/complete_index_score_batch_reader_test.cpp b/src/paimon/common/reader/complete_index_score_batch_reader_test.cpp index f8dac396..97b73f26 100644 --- a/src/paimon/common/reader/complete_index_score_batch_reader_test.cpp +++ b/src/paimon/common/reader/complete_index_score_batch_reader_test.cpp @@ -174,8 +174,6 @@ TEST_F(CompleteIndexScoreBatchReaderTest, TestFileReaderForwardsOperationsAndRes auto reader = std::make_unique( std::move(inner_reader), std::vector{1.25f, 2.5f}, GetArrowPool(GetDefaultPool())); - ASSERT_OK_AND_ASSIGN(std::unique_ptr<::ArrowSchema> file_schema, reader->GetFileSchema()); - EXPECT_TRUE(file_schema); ASSERT_OK_AND_ASSIGN(uint64_t row_count, reader->GetNumberOfRows()); EXPECT_EQ(2, row_count); EXPECT_FALSE(reader->SupportPreciseBitmapSelection()); diff --git a/src/paimon/core/operation/raw_file_split_read_test.cpp b/src/paimon/core/operation/raw_file_split_read_test.cpp index a8802b9a..bb9caa75 100644 --- a/src/paimon/core/operation/raw_file_split_read_test.cpp +++ b/src/paimon/core/operation/raw_file_split_read_test.cpp @@ -26,6 +26,7 @@ #include #include "arrow/api.h" +#include "arrow/c/bridge.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" #include "paimon/common/data/binary_row.h" @@ -34,6 +35,7 @@ #include "paimon/common/reader/concat_batch_reader.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" #include "paimon/core/core_options.h" #include "paimon/core/deletionvectors/bitmap_deletion_vector.h" #include "paimon/core/global_index/indexed_split_impl.h" @@ -95,8 +97,10 @@ class RawVectorSearchTestReader final : public FileIndexReader { class RawVectorSearchTestIndexer final : public FileIndexer { public: Result> CreateReader( - ::ArrowSchema*, int32_t, int32_t, const std::shared_ptr&, + ::ArrowSchema* arrow_schema, int32_t, int32_t, const std::shared_ptr&, const std::shared_ptr&) const override { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW([[maybe_unused]] std::shared_ptr schema, + arrow::ImportSchema(arrow_schema)); return std::make_shared(); } @@ -141,8 +145,10 @@ class RawFullTextSearchTestReader final : public FileIndexReader { class RawFullTextSearchTestIndexer final : public FileIndexer { public: Result> CreateReader( - ::ArrowSchema*, int32_t, int32_t, const std::shared_ptr&, + ::ArrowSchema* arrow_schema, int32_t, int32_t, const std::shared_ptr&, const std::shared_ptr&) const override { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW([[maybe_unused]] std::shared_ptr schema, + arrow::ImportSchema(arrow_schema)); return std::make_shared(); }