From 3cae61d32c403af6281a322aaf2e0ff18f9500b7 Mon Sep 17 00:00:00 2001 From: "jinli.zjw" Date: Wed, 23 Sep 2026 14:41:51 +0800 Subject: [PATCH 1/3] refactor(lumina): extract shared index components --- CMakeLists.txt | 2 +- .../lumina/lumina_global_index.cpp | 1024 ----------------- .../lumina/CMakeLists.txt | 9 +- .../lumina/lumina_api_test.cpp | 6 +- src/paimon/indexer/lumina/lumina_dataset.cpp | 105 ++ src/paimon/indexer/lumina/lumina_dataset.h | 80 ++ .../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 | 227 ++++ .../lumina/lumina_global_index.h | 79 +- .../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 + .../indexer/lumina/lumina_index_searcher.cpp | 119 ++ .../indexer/lumina/lumina_index_searcher.h | 77 ++ .../lumina/lumina_memory_pool.h | 0 .../indexer/lumina/lumina_tag_utils.cpp | 516 +++++++++ src/paimon/indexer/lumina/lumina_tag_utils.h | 80 ++ .../lumina/lumina_utils.h | 0 24 files changed, 1585 insertions(+), 1102 deletions(-) delete mode 100644 src/paimon/global_index/lumina/lumina_global_index.cpp rename src/paimon/{global_index => indexer}/lumina/CMakeLists.txt (88%) 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 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 (73%) 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 create mode 100644 src/paimon/indexer/lumina/lumina_index_searcher.cpp create mode 100644 src/paimon/indexer/lumina/lumina_index_searcher.h rename src/paimon/{global_index => indexer}/lumina/lumina_memory_pool.h (100%) 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 85fc0d8cd..f55bfadfa 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -583,7 +583,7 @@ if(PAIMON_ENABLE_LANCE) add_subdirectory(src/paimon/format/lance) 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/src/paimon/global_index/lumina/lumina_global_index.cpp b/src/paimon/global_index/lumina/lumina_global_index.cpp deleted file mode 100644 index 98585f21d..000000000 --- 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 88% rename from src/paimon/global_index/lumina/CMakeLists.txt rename to src/paimon/indexer/lumina/CMakeLists.txt index b0496df65..662358a96 100644 --- a/src/paimon/global_index/lumina/CMakeLists.txt +++ b/src/paimon/indexer/lumina/CMakeLists.txt @@ -15,7 +15,14 @@ # 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_global_index.cpp + lumina_global_index_factory.cpp + lumina_index_accumulator.cpp + lumina_index_options.cpp + lumina_index_searcher.cpp + lumina_tag_utils.cpp) add_paimon_lib(paimon_lumina_index SOURCES 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 2c4954b54..c16aa22b7 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 000000000..81079b619 --- /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 000000000..8e820656d --- /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/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 7ac7327fd..fa8aa0b54 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 0ae11692a..168d31685 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 a2514cc6c..1452b90e0 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 000000000..376bfe144 --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_global_index.cpp @@ -0,0 +1,227 @@ +/* + * 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_index_searcher.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(std::unique_ptr searcher, + LuminaIndexSearcher::Open(lumina_options, index_info, in, lumina_pool)); + return std::make_shared(std::move(searcher)); +} + +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(std::unique_ptr&& searcher) + : searcher_(std::move(searcher)) {} + +Result> LuminaIndexReader::VisitVectorSearch( + const std::shared_ptr& vector_search) { + PAIMON_ASSIGN_OR_RAISE(::lumina::api::LuminaSearcher::SearchResult search_result, + searcher_->Search(vector_search)); + + // 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 73% rename from src/paimon/global_index/lumina/lumina_global_index.h rename to src/paimon/indexer/lumina/lumina_global_index.h index c2c30475a..d15078535 100644 --- a/src/paimon/global_index/lumina/lumina_global_index.h +++ b/src/paimon/indexer/lumina/lumina_global_index.h @@ -22,41 +22,21 @@ #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/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_index_searcher.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 +80,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 +99,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,30 +109,16 @@ 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; - }; - - LuminaIndexReader( - const 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); - - ~LuminaIndexReader() override { - [[maybe_unused]] auto status = searcher_->Close(); - } + using IndexInfo = LuminaIndexInfo; + + explicit LuminaIndexReader(std::unique_ptr&& searcher); + + ~LuminaIndexReader() override = default; /// @note `VisitVectorSearch` is thread-safe (not coroutine-safe) while other `VisitXXX` is not /// thread-safe. @@ -246,13 +200,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_; - std::unique_ptr<::lumina::extensions::SearchWithFilterExtension> searcher_with_filter_; - std::unique_ptr<::lumina::extensions::experimental::SearchWithTagExtension> searcher_with_tag_; + std::unique_ptr searcher_; }; } // namespace paimon::lumina 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 e27837da3..6d26ecd8a 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 2d54950ad..1ac1114a6 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 000000000..d6f3993be --- /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 000000000..a7b6fe69d --- /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 000000000..e26b4e14f --- /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 000000000..ff77f032b --- /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/indexer/lumina/lumina_index_searcher.cpp b/src/paimon/indexer/lumina/lumina_index_searcher.cpp new file mode 100644 index 000000000..78c025556 --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_index_searcher.cpp @@ -0,0 +1,119 @@ +/* + * 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_searcher.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 { + +LuminaIndexSearcher::LuminaIndexSearcher(const LuminaIndexInfo& index_info, + const std::shared_ptr& pool, + std::unique_ptr<::lumina::api::LuminaSearcher>&& searcher) + : index_info_(index_info), pool_(pool), searcher_(std::move(searcher)) {} + +LuminaIndexSearcher::~LuminaIndexSearcher() { + [[maybe_unused]] ::lumina::core::Status status = searcher_->Close(); +} + +Result> LuminaIndexSearcher::Open( + const std::map& lumina_options, const LuminaIndexInfo& index_info, + const std::shared_ptr& input, const std::shared_ptr& pool) { + PAIMON_ASSIGN_OR_RAISE(::lumina::api::SearcherOptions searcher_options, + LuminaIndexOptions::CreateSearcherOptions(lumina_options, index_info)); + ::lumina::core::MemoryResourceConfig memory_resource(pool.get()); + PAIMON_ASSIGN_OR_RAISE_FROM_LUMINA( + ::lumina::api::LuminaSearcher lumina_searcher, + ::lumina::api::LuminaSearcher::Create(searcher_options, memory_resource)); + auto searcher = std::unique_ptr(new LuminaIndexSearcher( + index_info, pool, + std::make_unique<::lumina::api::LuminaSearcher>(std::move(lumina_searcher)))); + + auto file_reader = std::make_unique(input); + PAIMON_RETURN_NOT_OK_FROM_LUMINA( + searcher->searcher_->Open(std::move(file_reader), ::lumina::api::IOOptions())); + if (searcher->searcher_->GetMeta().dim != index_info.dimension) { + return Status::Invalid( + fmt::format("Lumina index dimension {} mismatch expected dimension {}", + searcher->searcher_->GetMeta().dim, index_info.dimension)); + } + + searcher->searcher_with_filter_ = + std::make_unique<::lumina::extensions::SearchWithFilterExtension>(); + PAIMON_RETURN_NOT_OK_FROM_LUMINA(searcher->searcher_->Attach(*searcher->searcher_with_filter_)); + if (index_info.has_tag) { + searcher->searcher_with_tag_ = + std::make_unique<::lumina::extensions::experimental::SearchWithTagExtension>(); + PAIMON_RETURN_NOT_OK_FROM_LUMINA( + searcher->searcher_->Attach(*searcher->searcher_with_tag_)); + } + return searcher; +} + +Result<::lumina::api::LuminaSearcher::SearchResult> LuminaIndexSearcher::Search( + const std::shared_ptr& vector_search) const { + 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_index_searcher.h b/src/paimon/indexer/lumina/lumina_index_searcher.h new file mode 100644 index 000000000..29fa1b748 --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_index_searcher.h @@ -0,0 +1,77 @@ +/* + * 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 { + +/// Owns an opened Lumina searcher and its attached search extensions. +class LuminaIndexSearcher { + public: + /// Create and open a Lumina searcher. + /// + /// @param lumina_options Normalized Lumina options without the `lumina.` prefix. + /// @param index_info Information about the built index. + /// @param input Input stream containing the Lumina index. + /// @param pool Memory pool used by Lumina. + /// @return An opened searcher, or an error Status. + static Result> Open( + const std::map& lumina_options, const LuminaIndexInfo& index_info, + const std::shared_ptr& input, const std::shared_ptr& pool); + + ~LuminaIndexSearcher(); + + LuminaIndexSearcher(const LuminaIndexSearcher&) = delete; + LuminaIndexSearcher& operator=(const LuminaIndexSearcher&) = delete; + + /// Execute a vector search against the opened index. + /// + /// @param vector_search Vector search request. + /// @return Native Lumina search result, or an error Status. + Result<::lumina::api::LuminaSearcher::SearchResult> Search( + const std::shared_ptr& vector_search) const; + + private: + LuminaIndexSearcher(const LuminaIndexInfo& index_info, + const std::shared_ptr& pool, + std::unique_ptr<::lumina::api::LuminaSearcher>&& searcher); + + LuminaIndexInfo index_info_; + std::shared_ptr pool_; + std::unique_ptr<::lumina::extensions::SearchWithFilterExtension> searcher_with_filter_; + std::unique_ptr<::lumina::extensions::experimental::SearchWithTagExtension> searcher_with_tag_; + std::unique_ptr<::lumina::api::LuminaSearcher> searcher_; +}; + +} // 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_tag_utils.cpp b/src/paimon/indexer/lumina/lumina_tag_utils.cpp new file mode 100644 index 000000000..979be3310 --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_tag_utils.cpp @@ -0,0 +1,516 @@ +/* + * 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 auto* 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)); + auto* 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 000000000..ff40cf9eb --- /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 c73c17f1b2a394e862b10417b54d580fad493da3 Mon Sep 17 00:00:00 2001 From: "jinli.zjw" Date: Thu, 24 Sep 2026 10:31:45 +0800 Subject: [PATCH 2/3] refactor(lumina): preserve checkpoint support --- src/paimon/indexer/lumina/CMakeLists.txt | 2 +- .../lumina/lumina_checkpoint_manager.cpp | 55 +++ .../lumina/lumina_checkpoint_manager.h | 45 ++ .../indexer/lumina/lumina_global_index.cpp | 61 ++- .../indexer/lumina/lumina_global_index.h | 27 +- .../lumina/lumina_global_index_test.cpp | 449 +++++++++++++++++- .../lumina/lumina_index_accumulator.cpp | 50 +- .../indexer/lumina/lumina_index_accumulator.h | 23 +- .../indexer/lumina/lumina_index_options.cpp | 6 + .../indexer/lumina/lumina_index_options.h | 2 + 10 files changed, 684 insertions(+), 36 deletions(-) create mode 100644 src/paimon/indexer/lumina/lumina_checkpoint_manager.cpp create mode 100644 src/paimon/indexer/lumina/lumina_checkpoint_manager.h diff --git a/src/paimon/indexer/lumina/CMakeLists.txt b/src/paimon/indexer/lumina/CMakeLists.txt index 662358a96..248a32920 100644 --- a/src/paimon/indexer/lumina/CMakeLists.txt +++ b/src/paimon/indexer/lumina/CMakeLists.txt @@ -16,6 +16,7 @@ if(PAIMON_ENABLE_LUMINA) set(PAIMON_LUMINA_INDEX + lumina_checkpoint_manager.cpp lumina_dataset.cpp lumina_global_index.cpp lumina_global_index_factory.cpp @@ -23,7 +24,6 @@ if(PAIMON_ENABLE_LUMINA) lumina_index_options.cpp lumina_index_searcher.cpp lumina_tag_utils.cpp) - add_paimon_lib(paimon_lumina_index SOURCES ${PAIMON_LUMINA_INDEX} diff --git a/src/paimon/indexer/lumina/lumina_checkpoint_manager.cpp b/src/paimon/indexer/lumina/lumina_checkpoint_manager.cpp new file mode 100644 index 000000000..dc66f30e7 --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_checkpoint_manager.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/indexer/lumina/lumina_checkpoint_manager.h" + +#include + +#include "paimon/indexer/lumina/lumina_file_reader.h" +#include "paimon/indexer/lumina/lumina_file_writer.h" +#include "paimon/indexer/lumina/lumina_utils.h" + +namespace paimon::lumina { + +std::unique_ptr<::lumina::io::FileWriter> LuminaCheckpointManager::CreateCkptFileWriter() { + Result> output = file_manager_->CreateCheckpointOutputStream(); + if (!output.ok()) { + return nullptr; + } + std::shared_ptr shared_output = std::move(output).value(); + return std::make_unique(shared_output); +} + +::lumina::core::Result LuminaCheckpointManager::HasCkptFile() { + Result exists = file_manager_->CheckpointExists(); + if (!exists.ok()) { + return ::lumina::core::Result::Err(PaimonToLuminaStatus(exists.status())); + } + return ::lumina::core::Result::Ok(exists.value()); +} + +std::unique_ptr<::lumina::io::FileReader> LuminaCheckpointManager::GetCkptFileReader() { + Result> input = file_manager_->OpenCheckpointInputStream(); + if (!input.ok()) { + return nullptr; + } + std::shared_ptr shared_input = std::move(input).value(); + return std::make_unique(shared_input); +} + +} // namespace paimon::lumina diff --git a/src/paimon/indexer/lumina/lumina_checkpoint_manager.h b/src/paimon/indexer/lumina/lumina_checkpoint_manager.h new file mode 100644 index 000000000..e52316e21 --- /dev/null +++ b/src/paimon/indexer/lumina/lumina_checkpoint_manager.h @@ -0,0 +1,45 @@ +/* + * 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 "lumina/extensions/experimental/CkptManager.h" +#include "paimon/global_index/io/global_index_checkpoint_file_manager.h" + +namespace paimon::lumina { + +/// Adapts Paimon's task-scoped global index checkpoint storage to Lumina. +class LuminaCheckpointManager final : public ::lumina::extensions::experimental::CkptManager { + public: + explicit LuminaCheckpointManager( + const std::shared_ptr& file_manager) + : file_manager_(file_manager) {} + + std::unique_ptr<::lumina::io::FileWriter> CreateCkptFileWriter() override; + + ::lumina::core::Result HasCkptFile() override; + + std::unique_ptr<::lumina::io::FileReader> GetCkptFileReader() override; + + private: + std::shared_ptr file_manager_; +}; + +} // namespace paimon::lumina diff --git a/src/paimon/indexer/lumina/lumina_global_index.cpp b/src/paimon/indexer/lumina/lumina_global_index.cpp index 376bfe144..5a7bf66c6 100644 --- a/src/paimon/indexer/lumina/lumina_global_index.cpp +++ b/src/paimon/indexer/lumina/lumina_global_index.cpp @@ -22,6 +22,7 @@ #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" +#include "glog/logging.h" #include "lumina/api/LuminaBuilder.h" #include "lumina/api/LuminaSearcher.h" #include "lumina/core/Constants.h" @@ -32,6 +33,7 @@ #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_checkpoint_manager.h" #include "paimon/indexer/lumina/lumina_file_writer.h" #include "paimon/indexer/lumina/lumina_index_options.h" #include "paimon/indexer/lumina/lumina_index_searcher.h" @@ -72,10 +74,19 @@ Result> LuminaGlobalIndex::CreateWriter( PAIMON_ASSIGN_OR_RAISE(uint32_t dimension, LuminaIndexOptions::GetDimension(lumina_options)); PAIMON_ASSIGN_OR_RAISE(::lumina::api::BuilderOptions builder_options, LuminaIndexOptions::CreateBuilderOptions(lumina_options)); + std::shared_ptr checkpoint_file_manager; + if (LuminaIndexOptions::IsCheckpointEnabled(lumina_options)) { + checkpoint_file_manager = + std::dynamic_pointer_cast(file_writer); + if (!checkpoint_file_manager || !checkpoint_file_manager->SupportsCheckpoint()) { + return Status::Invalid("Lumina checkpoint requires a checkpoint-capable file writer"); + } + } 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); + ::lumina::api::IOOptions(), lumina_options, std::move(tag_fields), checkpoint_file_manager, + lumina_pool); } Result LuminaIndexReader::GetIndexInfo( @@ -137,7 +148,9 @@ LuminaIndexWriter::LuminaIndexWriter( 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) + std::vector&& tag_fields, + const std::shared_ptr& checkpoint_file_manager, + const std::shared_ptr& pool) : pool_(pool), field_name_(field_name), arrow_type_(arrow_type), @@ -146,7 +159,8 @@ LuminaIndexWriter::LuminaIndexWriter( builder_options_(std::move(builder_options)), io_options_(std::move(io_options)), lumina_options_(lumina_options), - tag_fields_(std::move(tag_fields)) {} + tag_fields_(std::move(tag_fields)), + checkpoint_file_manager_(checkpoint_file_manager) {} Status LuminaIndexWriter::AddBatch(::ArrowArray* arrow_array, std::vector&& relative_row_ids) { @@ -179,9 +193,34 @@ 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())); + + bool had_checkpoint = false; + if (checkpoint_file_manager_) { + PAIMON_ASSIGN_OR_RAISE(had_checkpoint, checkpoint_file_manager_->CheckpointExists()); + } + + auto build_index = [&]() -> Result> { + std::unique_ptr<::lumina::extensions::experimental::CkptManager> checkpoint_manager; + if (checkpoint_file_manager_) { + checkpoint_manager = + std::make_unique(checkpoint_file_manager_); + } + return accumulator_.Build(builder_options_, dimension_, !tag_fields_.empty(), + std::move(checkpoint_manager), pool_.get()); + }; + + Result> build_result = build_index(); + if (!build_result.ok() && had_checkpoint) { + LOG(WARNING) << "Failed to build Lumina index with checkpoint, discard it and rebuild " + "from scratch: " + << build_result.status().ToString(); + PAIMON_RETURN_NOT_OK(checkpoint_file_manager_->DeleteCheckpoint()); + build_result = build_index(); + } + if (!build_result.ok()) { + return build_result.status(); + } + std::unique_ptr build_context = std::move(build_result).value(); // dump index PAIMON_ASSIGN_OR_RAISE(std::string index_file_name, @@ -189,7 +228,8 @@ Result> LuminaIndexWriter::Finish() { 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_)); + PAIMON_RETURN_NOT_OK_FROM_LUMINA( + build_context->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; @@ -197,6 +237,13 @@ Result> LuminaIndexWriter::Finish() { auto meta_bytes = std::make_shared(options_json, pool_->GetPaimonPool().get()); GlobalIndexIOMeta meta(file_manager_->ToPath(index_file_name), file_size, /*metadata=*/meta_bytes); + if (checkpoint_file_manager_) { + Status status = checkpoint_file_manager_->DeleteCheckpoint(); + if (!status.ok()) { + LOG(WARNING) << "Failed to delete Lumina checkpoints after successful build: " + << status.ToString(); + } + } return std::vector({meta}); } diff --git a/src/paimon/indexer/lumina/lumina_global_index.h b/src/paimon/indexer/lumina/lumina_global_index.h index d15078535..452f5561a 100644 --- a/src/paimon/indexer/lumina/lumina_global_index.h +++ b/src/paimon/indexer/lumina/lumina_global_index.h @@ -18,6 +18,7 @@ #pragma once +#include #include #include #include @@ -29,6 +30,7 @@ #include "lumina/api/Options.h" #include "lumina/extensions/experimental/DatasetWithTag.h" #include "paimon/global_index/global_indexer.h" +#include "paimon/global_index/io/global_index_checkpoint_file_manager.h" #include "paimon/indexer/lumina/lumina_index_accumulator.h" #include "paimon/indexer/lumina/lumina_index_options.h" #include "paimon/indexer/lumina/lumina_index_searcher.h" @@ -53,6 +55,8 @@ namespace paimon::lumina { /// lumina.diskann.build.thread_count:64 /// lumina.diskann.build.ef_construction:1024 /// lumina.diskann.build.neighbor_count:64 +/// lumina.extension.build.ckpt.threshold:10000 +/// lumina.extension.build.ckpt.count:3 /// /// - **Index Reader:** /// No configuration required at load time — settings are stored in the index metadata, @@ -67,8 +71,14 @@ class LuminaGlobalIndex : public GlobalIndexer { explicit LuminaGlobalIndex(const std::map& options) : options_(options) {} + bool SupportsCheckpoint() const override { + return true; + } + Result>> GetExtraFieldNames() const override; + /// With checkpoints enabled, file_writer must implement GlobalIndexCheckpointFileManager and + /// return true from SupportsCheckpoint(). Result> CreateWriter( const std::string& field_name, ::ArrowSchema* arrow_schema, const std::shared_ptr& file_writer, @@ -85,14 +95,14 @@ class LuminaGlobalIndex : public GlobalIndexer { class LuminaIndexWriter : public GlobalIndexWriter { public: - 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); + 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& checkpoint_file_manager, + const std::shared_ptr& pool); Status AddBatch(::ArrowArray* arrow_array, std::vector&& relative_row_ids) override; @@ -109,6 +119,7 @@ class LuminaIndexWriter : public GlobalIndexWriter { ::lumina::api::IOOptions io_options_; std::map lumina_options_; std::vector tag_fields_; + std::shared_ptr checkpoint_file_manager_; LuminaIndexAccumulator accumulator_; }; diff --git a/src/paimon/indexer/lumina/lumina_global_index_test.cpp b/src/paimon/indexer/lumina/lumina_global_index_test.cpp index 1ac1114a6..e734d48c3 100644 --- a/src/paimon/indexer/lumina/lumina_global_index_test.cpp +++ b/src/paimon/indexer/lumina/lumina_global_index_test.cpp @@ -23,15 +23,23 @@ #include "arrow/c/bridge.h" #include "arrow/ipc/api.h" #include "gtest/gtest.h" +#include "lumina/api/Dataset.h" +#include "lumina/api/LuminaBuilder.h" +#include "lumina/core/Constants.h" +#include "lumina/extensions/experimental/BuildCombinedExtensionV0.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/date_time_utils.h" #include "paimon/common/utils/path_util.h" #include "paimon/common/utils/string_utils.h" #include "paimon/core/global_index/global_index_file_manager.h" +#include "paimon/core/index/index_checkpoint_path_factory.h" #include "paimon/core/index/index_path_factory.h" +#include "paimon/core/utils/file_store_path_factory.h" #include "paimon/fs/local/local_file_system.h" #include "paimon/global_index/bitmap_scored_global_index_result.h" #include "paimon/global_index/global_index_result.h" +#include "paimon/indexer/lumina/lumina_checkpoint_manager.h" +#include "paimon/indexer/lumina/lumina_memory_pool.h" #include "paimon/predicate/predicate_builder.h" #include "paimon/testing/utils/testharness.h" namespace paimon::lumina::test { @@ -62,6 +70,63 @@ class LuminaGlobalIndexTest : public ::testing::Test { std::string index_path_; }; + class CountingCheckpointFileManager : public GlobalIndexFileManager { + public: + using GlobalIndexFileManager::GlobalIndexFileManager; + + Result CheckpointExists() const override { + check_checkpoint_count_++; + return GlobalIndexFileManager::CheckpointExists(); + } + + Result> CreateCheckpointOutputStream() const override { + create_checkpoint_count_++; + return GlobalIndexFileManager::CreateCheckpointOutputStream(); + } + + Result> OpenCheckpointInputStream() const override { + open_checkpoint_count_++; + return GlobalIndexFileManager::OpenCheckpointInputStream(); + } + + mutable int32_t check_checkpoint_count_ = 0; + mutable int32_t create_checkpoint_count_ = 0; + mutable int32_t open_checkpoint_count_ = 0; + }; + + class TestLuminaDataset : public ::lumina::api::Dataset { + public: + TestLuminaDataset(uint32_t dimension, const std::vector& vectors, + const std::vector<::lumina::core::vector_id_t>& ids) + : dimension_(dimension), vectors_(vectors), ids_(ids) {} + + uint32_t Dim() const noexcept override { + return dimension_; + } + + uint64_t TotalSize() const noexcept override { + return ids_.size(); + } + + ::lumina::core::Result GetNextBatch( + std::vector& vector_buffer, + std::vector<::lumina::core::vector_id_t>& id_buffer) noexcept override { + if (consumed_) { + return ::lumina::core::Result::Ok(0); + } + vector_buffer = vectors_; + id_buffer = ids_; + consumed_ = true; + return ::lumina::core::Result::Ok(ids_.size()); + } + + private: + uint32_t dimension_; + std::vector vectors_; + std::vector<::lumina::core::vector_id_t> ids_; + bool consumed_ = false; + }; + std::unique_ptr<::ArrowSchema> CreateArrowSchema( const std::shared_ptr& data_type) const { auto c_schema = std::make_unique<::ArrowSchema>(); @@ -69,14 +134,47 @@ class LuminaGlobalIndexTest : public ::testing::Test { return c_schema; } + Result> CreateFileStorePathFactory( + const std::string& table_path) const { + return FileStorePathFactory::Create( + table_path, arrow::schema({}), /*partition_keys=*/{}, /*default_part_value=*/"", + /*identifier=*/"mock", /*data_file_prefix=*/"data-", + /*legacy_partition_name_enabled=*/true, /*external_paths=*/{}, + /*global_index_external_path=*/std::nullopt, + /*index_file_in_data_file_dir=*/false, pool_); + } + + Result> CreateCheckpointPathFactory( + const std::string& table_path, const std::string& index_type, const std::string& field_name, + const Range& range) const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr file_store_path_factory, + CreateFileStorePathFactory(table_path)); + return file_store_path_factory->CreateGlobalIndexCheckpointPathFactory( + index_type, field_name, range, "task-1"); + } + + Result> CreateCountingCheckpointFileManager( + const std::shared_ptr& path_factory, const Range& range) const { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr checkpoint_path_factory, + path_factory->CreateGlobalIndexCheckpointPathFactory("lumina", "f0", range, "task-1")); + return std::make_shared( + fs_, path_factory->CreateGlobalIndexFileFactory(), std::move(checkpoint_path_factory)); + } + Result WriteGlobalIndex(const std::string& index_root, const std::shared_ptr& data_type, const std::map& options, const std::shared_ptr& array, const Range& expected_range) const { auto global_index = std::make_shared(options); - auto path_factory = std::make_shared(index_root); - auto file_writer = std::make_shared(fs_, path_factory); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr path_factory, + CreateFileStorePathFactory(index_root)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr checkpoint_path_factory, + path_factory->CreateGlobalIndexCheckpointPathFactory( + "lumina", "f0", expected_range, "task-1")); + auto file_writer = std::make_shared( + fs_, path_factory->CreateGlobalIndexFileFactory(), std::move(checkpoint_path_factory)); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr global_writer, global_index->CreateWriter("f0", CreateArrowSchema(data_type).get(), @@ -127,7 +225,8 @@ class LuminaGlobalIndexTest : public ::testing::Test { const std::map& options, const GlobalIndexIOMeta& meta) const { auto global_index = std::make_shared(options); auto path_factory = std::make_shared(index_root); - auto file_reader = std::make_shared(fs_, path_factory); + auto file_reader = std::make_shared( + fs_, path_factory, /*checkpoint_path_factory=*/nullptr); return global_index->CreateReader(CreateArrowSchema(data_type).get(), file_reader, {meta}, pool_); } @@ -252,12 +351,344 @@ TEST_F(LuminaGlobalIndexTest, TestWithFilter) { } } +TEST_F(LuminaGlobalIndexTest, TestCheckpointCapabilityOnlyRequiredWhenEnabled) { + ASSERT_TRUE(LuminaGlobalIndex(options_).SupportsCheckpoint()); + auto dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(std::shared_ptr path_factory, + CreateFileStorePathFactory(dir->Str())); + auto plain_manager = + std::make_shared(fs_, path_factory->CreateGlobalIndexFileFactory(), + /*checkpoint_path_factory=*/nullptr); + ASSERT_OK_AND_ASSIGN(std::shared_ptr checkpoint_manager, + CreateCountingCheckpointFileManager(path_factory, Range(0, 3))); + std::map options = options_; + options["other.extension.build.ckpt.count"] = "1"; + ASSERT_OK(LuminaGlobalIndex(options).CreateWriter("f0", CreateArrowSchema(data_type_).get(), + plain_manager, pool_)); + ASSERT_OK(LuminaGlobalIndex(options).CreateWriter("f0", CreateArrowSchema(data_type_).get(), + checkpoint_manager, pool_)); + + for (const std::string& checkpoint_key : + std::vector{"extension.build.ckpt.count", "extension.build.ckpt.threshold"}) { + auto enabled = options; + enabled["lumina." + checkpoint_key] = "1"; + ASSERT_NOK_WITH_MSG(LuminaGlobalIndex(enabled).CreateWriter( + "f0", CreateArrowSchema(data_type_).get(), plain_manager, pool_), + "Lumina checkpoint requires a checkpoint-capable file writer"); + ASSERT_NOK_WITH_MSG(LuminaGlobalIndex(enabled).CreateWriter( + "f0", CreateArrowSchema(data_type_).get(), nullptr, pool_), + "Lumina checkpoint requires a checkpoint-capable file writer"); + ASSERT_OK(LuminaGlobalIndex(enabled).CreateWriter("f0", CreateArrowSchema(data_type_).get(), + checkpoint_manager, pool_)); + } + ASSERT_EQ(checkpoint_manager->check_checkpoint_count_, 0); + ASSERT_EQ(checkpoint_manager->create_checkpoint_count_, 0); +} + +TEST_F(LuminaGlobalIndexTest, TestBuildWithCheckpoint) { + auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(test_root_dir); + std::string test_root = test_root_dir->Str(); + + std::map checkpoint_options = options_; + checkpoint_options["lumina.extension.build.ckpt.threshold"] = "1"; + checkpoint_options["lumina.extension.build.ckpt.count"] = "1"; + + auto global_index = std::make_shared(checkpoint_options); + ASSERT_OK_AND_ASSIGN(std::shared_ptr path_factory, + CreateFileStorePathFactory(test_root)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr file_manager, + CreateCountingCheckpointFileManager(path_factory, Range(0, 3))); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr global_writer, + global_index->CreateWriter("f0", CreateArrowSchema(data_type_).get(), file_manager, pool_)); + + ArrowArray c_array; + ASSERT_TRUE(arrow::ExportArray(*array_, &c_array).ok()); + ASSERT_OK(global_writer->AddBatch(&c_array, {0, 1, 2, 3})); + ASSERT_OK_AND_ASSIGN(std::vector result_metas, global_writer->Finish()); + ASSERT_EQ(result_metas.size(), 1); + ASSERT_GT(file_manager->create_checkpoint_count_, 0); + ASSERT_OK_AND_ASSIGN(bool checkpoint_exists, file_manager->CheckpointExists()); + ASSERT_FALSE(checkpoint_exists); + + ASSERT_OK_AND_ASSIGN( + std::shared_ptr reader, + CreateGlobalIndexReader(test_root, data_type_, checkpoint_options, result_metas[0])); + ASSERT_OK_AND_ASSIGN(std::shared_ptr scored_result, + reader->VisitVectorSearch(std::make_shared( + /*field_name=*/"f0", /*limit=*/4, query_, /*filter=*/nullptr, + /*predicate=*/nullptr, /*distance_type=*/std::nullopt, + /*options=*/checkpoint_options))); + CheckResult(scored_result, {3l, 1l, 2l, 0l}, {0.01f, 2.01f, 2.21f, 4.21f}); +} + +TEST_F(LuminaGlobalIndexTest, TestCheckpointFileManagement) { + auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(test_root_dir); + std::string test_root = test_root_dir->Str(); + std::string prefix = "lumina-global-index-f0-10-20-task-1-"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr checkpoint_path_factory, + CreateCheckpointPathFactory(test_root, "lumina", "f0", Range(10, 20))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr path_factory, + CreateFileStorePathFactory(test_root)); + auto file_manager = std::make_shared( + fs_, path_factory->CreateGlobalIndexFileFactory(), std::move(checkpoint_path_factory)); + + ASSERT_OK_AND_ASSIGN(bool checkpoint_exists, file_manager->CheckpointExists()); + ASSERT_FALSE(checkpoint_exists); + ASSERT_OK_AND_ASSIGN(std::unique_ptr first_output, + file_manager->CreateCheckpointOutputStream()); + ASSERT_OK_AND_ASSIGN(std::string first_path, first_output->GetUri()); + ASSERT_TRUE(StringUtils::StartsWith(PathUtil::GetName(first_path), prefix)); + ASSERT_TRUE(StringUtils::EndsWith(first_path, "-0.index.ckpt")); + ASSERT_OK_AND_ASSIGN(int64_t first_written, first_output->Write("old", 3)); + ASSERT_EQ(first_written, 3); + ASSERT_OK(first_output->Close()); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr second_output, + file_manager->CreateCheckpointOutputStream()); + ASSERT_OK_AND_ASSIGN(std::string second_path, second_output->GetUri()); + ASSERT_TRUE(StringUtils::StartsWith(PathUtil::GetName(second_path), prefix)); + ASSERT_TRUE(StringUtils::EndsWith(second_path, "-1.index.ckpt")); + ASSERT_OK_AND_ASSIGN(int64_t second_written, second_output->Write("latest", 6)); + ASSERT_EQ(second_written, 6); + ASSERT_OK(second_output->Close()); + + std::string checkpoint_dir = PathUtil::JoinPath(test_root, "index/checkpoint"); + ASSERT_OK(fs_->WriteFile(PathUtil::JoinPath(checkpoint_dir, prefix + "9.index.ckpt"), "id-nine", + /*overwrite=*/false)); + ASSERT_OK(fs_->WriteFile(PathUtil::JoinPath(checkpoint_dir, prefix + "10.index.ckpt"), "id-ten", + /*overwrite=*/false)); + std::string unrelated_file = + PathUtil::JoinPath(checkpoint_dir, "lumina-global-index-other-10-20-task-2-99.index.ckpt"); + ASSERT_OK(fs_->WriteFile(unrelated_file, "unrelated", /*overwrite=*/false)); + std::string malformed_file = + PathUtil::JoinPath(checkpoint_dir, prefix + "invalid-99.index.ckpt"); + ASSERT_OK(fs_->WriteFile(malformed_file, "malformed", /*overwrite=*/false)); + + ASSERT_OK_AND_ASSIGN(checkpoint_path_factory, + CreateCheckpointPathFactory(test_root, "lumina", "f0", Range(10, 20))); + file_manager = std::make_shared( + fs_, path_factory->CreateGlobalIndexFileFactory(), std::move(checkpoint_path_factory)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr latest_input, + file_manager->OpenCheckpointInputStream()); + char buffer[6]; + ASSERT_OK_AND_ASSIGN(int64_t read_bytes, latest_input->Read(buffer, sizeof(buffer))); + ASSERT_EQ(read_bytes, 6); + ASSERT_EQ(std::string(buffer, sizeof(buffer)), "id-ten"); + ASSERT_OK(latest_input->Close()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr next_output, + file_manager->CreateCheckpointOutputStream()); + ASSERT_OK_AND_ASSIGN(std::string next_path, next_output->GetUri()); + ASSERT_TRUE(StringUtils::EndsWith(next_path, "-11.index.ckpt")); + ASSERT_OK(next_output->Close()); + + ASSERT_OK(file_manager->DeleteCheckpoint()); + ASSERT_OK_AND_ASSIGN(checkpoint_exists, file_manager->CheckpointExists()); + ASSERT_FALSE(checkpoint_exists); + ASSERT_OK_AND_ASSIGN(bool unrelated_exists, fs_->Exists(unrelated_file)); + ASSERT_TRUE(unrelated_exists); + ASSERT_OK_AND_ASSIGN(bool malformed_exists, fs_->Exists(malformed_file)); + ASSERT_TRUE(malformed_exists); +} + +TEST_F(LuminaGlobalIndexTest, TestResumeFromCheckpoint) { + auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(test_root_dir); + std::string test_root = test_root_dir->Str(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr path_factory, + CreateFileStorePathFactory(test_root)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr file_manager, + CreateCountingCheckpointFileManager(path_factory, Range(0, 3))); + + std::map checkpoint_options = options_; + checkpoint_options["lumina.extension.build.ckpt.threshold"] = "1"; + checkpoint_options["lumina.extension.build.ckpt.count"] = "1"; + std::vector vectors = { + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 1.0f, + 1.0f, 0.0f, 1.0f, 0.0f, 1.0f, 1.0f, 1.0f, 1.0f, + }; + std::vector<::lumina::core::vector_id_t> ids = {0, 1, 2, 3}; + + { + ::lumina::api::BuilderOptions builder_options; + builder_options.Set(::lumina::core::kIndexType, ::lumina::core::kIndexTypeBruteforce) + .Set(::lumina::core::kDimension, static_cast(4)) + .Set(::lumina::core::kDistanceMetric, ::lumina::core::kDistanceL2) + .Set(::lumina::core::kEncodingType, ::lumina::core::kEncodingRawf32) + .Set(::lumina::core::kExtensionCkptThreshold, static_cast(1)) + .Set(::lumina::core::kExtensionCkptCount, static_cast(1)); + LuminaMemoryPool lumina_pool(pool_); + ::lumina::core::MemoryResourceConfig memory_resource(&lumina_pool); + auto builder_result = + ::lumina::api::LuminaBuilder::Create(builder_options, memory_resource); + ASSERT_TRUE(builder_result.IsOk()) << builder_result.GetStatus().Message(); + ::lumina::api::LuminaBuilder builder = std::move(builder_result).TakeValue(); + ::lumina::extensions::experimental::BuildWithCheckpointExtension checkpoint_extension; + ASSERT_TRUE(builder.Attach(checkpoint_extension).IsOk()); + ASSERT_TRUE(checkpoint_extension + .LoadCkptManager(std::make_unique(file_manager)) + .IsOk()); + TestLuminaDataset pretrain_dataset(/*dimension=*/4, vectors, ids); + ASSERT_TRUE(builder.PretrainFrom(pretrain_dataset).IsOk()); + TestLuminaDataset insert_dataset(/*dimension=*/4, vectors, ids); + ASSERT_TRUE(builder.InsertFrom(insert_dataset).IsOk()); + } + ASSERT_GT(file_manager->create_checkpoint_count_, 0); + ASSERT_OK_AND_ASSIGN(bool checkpoint_exists, file_manager->CheckpointExists()); + ASSERT_TRUE(checkpoint_exists); + + auto global_index = std::make_shared(checkpoint_options); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr global_writer, + global_index->CreateWriter("f0", CreateArrowSchema(data_type_).get(), file_manager, pool_)); + ArrowArray c_array; + ASSERT_TRUE(arrow::ExportArray(*array_, &c_array).ok()); + ASSERT_OK(global_writer->AddBatch(&c_array, {0, 1, 2, 3})); + ASSERT_OK_AND_ASSIGN(std::vector result_metas, global_writer->Finish()); + ASSERT_EQ(result_metas.size(), 1); + ASSERT_GT(file_manager->open_checkpoint_count_, 0); + ASSERT_OK_AND_ASSIGN(checkpoint_exists, file_manager->CheckpointExists()); + ASSERT_FALSE(checkpoint_exists); + + ASSERT_OK_AND_ASSIGN( + std::shared_ptr reader, + CreateGlobalIndexReader(test_root, data_type_, checkpoint_options, result_metas[0])); + ASSERT_OK_AND_ASSIGN(std::shared_ptr scored_result, + reader->VisitVectorSearch(std::make_shared( + /*field_name=*/"f0", /*limit=*/4, query_, /*filter=*/nullptr, + /*predicate=*/nullptr, /*distance_type=*/std::nullopt, + /*options=*/checkpoint_options))); + CheckResult(scored_result, {3l, 1l, 2l, 0l}, {0.01f, 2.01f, 2.21f, 4.21f}); +} + +TEST_F(LuminaGlobalIndexTest, TestDiscardInvalidCheckpointAndRebuild) { + auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(test_root_dir); + std::string test_root = test_root_dir->Str(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr path_factory, + CreateFileStorePathFactory(test_root)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr file_manager, + CreateCountingCheckpointFileManager(path_factory, Range(0, 3))); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr invalid_checkpoint, + file_manager->CreateCheckpointOutputStream()); + ASSERT_OK_AND_ASSIGN(int64_t written, invalid_checkpoint->Write("invalid checkpoint", 18)); + ASSERT_EQ(written, 18); + ASSERT_OK(invalid_checkpoint->Close()); + + std::map checkpoint_options = options_; + checkpoint_options["lumina.extension.build.ckpt.threshold"] = "1"; + checkpoint_options["lumina.extension.build.ckpt.count"] = "1"; + auto global_index = std::make_shared(checkpoint_options); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr global_writer, + global_index->CreateWriter("f0", CreateArrowSchema(data_type_).get(), file_manager, pool_)); + ArrowArray c_array; + ASSERT_TRUE(arrow::ExportArray(*array_, &c_array).ok()); + ASSERT_OK(global_writer->AddBatch(&c_array, {0, 1, 2, 3})); + ASSERT_OK_AND_ASSIGN(std::vector result_metas, global_writer->Finish()); + ASSERT_EQ(result_metas.size(), 1); + ASSERT_GT(file_manager->open_checkpoint_count_, 0); + ASSERT_OK_AND_ASSIGN(bool checkpoint_exists, file_manager->CheckpointExists()); + ASSERT_FALSE(checkpoint_exists); + + ASSERT_OK_AND_ASSIGN( + std::shared_ptr reader, + CreateGlobalIndexReader(test_root, data_type_, checkpoint_options, result_metas[0])); + ASSERT_OK_AND_ASSIGN(std::shared_ptr scored_result, + reader->VisitVectorSearch(std::make_shared( + /*field_name=*/"f0", /*limit=*/4, query_, /*filter=*/nullptr, + /*predicate=*/nullptr, /*distance_type=*/std::nullopt, + /*options=*/checkpoint_options))); + CheckResult(scored_result, {3l, 1l, 2l, 0l}, {0.01f, 2.01f, 2.21f, 4.21f}); +} + +TEST_F(LuminaGlobalIndexTest, TestBuildFailureDiscardsCheckpointAndRetries) { + class DeletionCountingCheckpointFileManager : public CountingCheckpointFileManager { + public: + using CountingCheckpointFileManager::CountingCheckpointFileManager; + + Status DeleteCheckpoint() const override { + delete_count_++; + return CountingCheckpointFileManager::DeleteCheckpoint(); + } + + mutable int32_t delete_count_ = 0; + }; + + auto dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + ASSERT_OK_AND_ASSIGN(std::shared_ptr path_factory, + CreateFileStorePathFactory(dir->Str())); + ASSERT_OK_AND_ASSIGN(std::unique_ptr checkpoint_path_factory, + path_factory->CreateGlobalIndexCheckpointPathFactory( + "lumina", "f0", Range(0, 3), "task-1")); + auto file_manager = std::make_shared( + fs_, path_factory->CreateGlobalIndexFileFactory(), std::move(checkpoint_path_factory)); + + ::lumina::api::BuilderOptions builder_options; + builder_options.Set(::lumina::core::kIndexType, ::lumina::core::kIndexTypeBruteforce) + .Set(::lumina::core::kDimension, static_cast(4)) + .Set(::lumina::core::kDistanceMetric, ::lumina::core::kDistanceL2) + .Set(::lumina::core::kEncodingType, ::lumina::core::kEncodingRawf32) + .Set(::lumina::core::kExtensionCkptThreshold, static_cast(1)) + .Set(::lumina::core::kExtensionCkptCount, static_cast(1)); + LuminaMemoryPool lumina_pool(pool_); + ::lumina::core::MemoryResourceConfig memory_resource(&lumina_pool); + auto builder_result = ::lumina::api::LuminaBuilder::Create(builder_options, memory_resource); + ASSERT_TRUE(builder_result.IsOk()) << builder_result.GetStatus().Message(); + { + ::lumina::api::LuminaBuilder builder = std::move(builder_result).TakeValue(); + ::lumina::extensions::experimental::BuildWithCheckpointExtension checkpoint_extension; + ASSERT_TRUE(builder.Attach(checkpoint_extension).IsOk()); + ASSERT_TRUE(checkpoint_extension + .LoadCkptManager(std::make_unique(file_manager)) + .IsOk()); + std::vector vectors = { + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 1.0f, + 1.0f, 0.0f, 1.0f, 0.0f, 1.0f, 1.0f, 1.0f, 1.0f, + }; + std::vector<::lumina::core::vector_id_t> ids = {0, 1, 2, 3}; + TestLuminaDataset pretrain_dataset(/*dimension=*/4, vectors, ids); + ASSERT_TRUE(builder.PretrainFrom(pretrain_dataset).IsOk()); + TestLuminaDataset insert_dataset(/*dimension=*/4, vectors, ids); + ASSERT_TRUE(builder.InsertFrom(insert_dataset).IsOk()); + } + ASSERT_OK_AND_ASSIGN(bool checkpoint_exists, file_manager->CheckpointExists()); + ASSERT_TRUE(checkpoint_exists); + + std::map checkpoint_options = options_; + checkpoint_options["lumina.extension.build.ckpt.threshold"] = "1"; + checkpoint_options["lumina.extension.build.ckpt.count"] = "1"; + ASSERT_OK_AND_ASSIGN( + std::shared_ptr writer, + LuminaGlobalIndex(checkpoint_options) + .CreateWriter("f0", CreateArrowSchema(data_type_).get(), file_manager, pool_)); + // The two-row build fails after loading the checkpoint, and must be retried without it. + file_manager->check_checkpoint_count_ = 0; + ArrowArray c_array; + ASSERT_TRUE(arrow::ExportArray(*array_->Slice(0, 2), &c_array).ok()); + ASSERT_OK(writer->AddBatch(&c_array, {0, 1})); + ASSERT_NOK(writer->Finish()); + ASSERT_GT(file_manager->open_checkpoint_count_, 0); + ASSERT_EQ(file_manager->delete_count_, 1); + ASSERT_GE(file_manager->check_checkpoint_count_, 3); + ASSERT_OK_AND_ASSIGN(checkpoint_exists, file_manager->CheckpointExists()); + ASSERT_FALSE(checkpoint_exists); +} + TEST_F(LuminaGlobalIndexTest, TestWriteAndReadWithTagFilter) { auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); ASSERT_TRUE(test_root_dir); std::string test_root = test_root_dir->Str(); std::map tag_options = options_; + tag_options["lumina.extension.build.ckpt.threshold"] = "1"; + tag_options["lumina.extension.build.ckpt.count"] = "1"; tag_options["lumina.extension.build.tag.tag_schema"] = R"({"key_name":"color","type":"enum","value_type":"string"})"; @@ -276,6 +707,9 @@ TEST_F(LuminaGlobalIndexTest, TestWriteAndReadWithTagFilter) { ASSERT_OK_AND_ASSIGN( GlobalIndexIOMeta meta, WriteGlobalIndex(test_root, tag_data_type, tag_options, tag_array, Range(0, 3))); + std::vector checkpoint_files; + ASSERT_OK(fs_->ListDir(PathUtil::JoinPath(test_root, "index/checkpoint"), &checkpoint_files)); + ASSERT_TRUE(checkpoint_files.empty()); ASSERT_OK_AND_ASSIGN(std::shared_ptr reader, CreateGlobalIndexReader(test_root, data_type_, tag_options, meta)); @@ -787,7 +1221,8 @@ TEST_F(LuminaGlobalIndexTest, TestInvalidInputs) { { auto global_index = std::make_shared(options_); auto path_factory = std::make_shared(index_root); - auto file_reader = std::make_shared(fs_, path_factory); + auto file_reader = std::make_shared( + fs_, path_factory, /*checkpoint_path_factory=*/nullptr); ASSERT_NOK_WITH_MSG(global_index->CreateReader(CreateArrowSchema(data_type_).get(), file_reader, {meta, meta}, pool_), @@ -1005,7 +1440,8 @@ TEST_F(LuminaGlobalIndexTest, TestWriteWithAllNullRows) { auto global_index = std::make_shared(options_); auto path_factory = std::make_shared(test_root); - auto file_writer = std::make_shared(fs_, path_factory); + auto file_writer = std::make_shared( + fs_, path_factory, /*checkpoint_path_factory=*/nullptr); ASSERT_OK_AND_ASSIGN( std::shared_ptr global_writer, @@ -1077,7 +1513,8 @@ TEST_F(LuminaGlobalIndexTest, TestWriteWithNullAcrossMultipleBatches) { auto global_index = std::make_shared(options_); auto path_factory = std::make_shared(test_root); - auto file_writer = std::make_shared(fs_, path_factory); + auto file_writer = std::make_shared( + fs_, path_factory, /*checkpoint_path_factory=*/nullptr); ASSERT_OK_AND_ASSIGN( std::shared_ptr global_writer, diff --git a/src/paimon/indexer/lumina/lumina_index_accumulator.cpp b/src/paimon/indexer/lumina/lumina_index_accumulator.cpp index d6f3993be..e655c5d00 100644 --- a/src/paimon/indexer/lumina/lumina_index_accumulator.cpp +++ b/src/paimon/indexer/lumina/lumina_index_accumulator.cpp @@ -81,34 +81,62 @@ Status LuminaIndexAccumulator::AddBatch(const std::shared_ptr LuminaIndexAccumulator::Build( +Result> LuminaIndexAccumulator::Build( const ::lumina::api::BuilderOptions& builder_options, uint32_t dimension, bool with_tag, + std::unique_ptr<::lumina::extensions::experimental::CkptManager> checkpoint_manager, 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)); + auto context = std::make_unique(std::move(builder)); + + if (checkpoint_manager) { + if (!with_tag) { + context->checkpoint_extension = std::make_unique< + ::lumina::extensions::experimental::BuildWithCheckpointExtension>(); + PAIMON_RETURN_NOT_OK_FROM_LUMINA( + context->builder.Attach(*context->checkpoint_extension)); + PAIMON_RETURN_NOT_OK_FROM_LUMINA( + context->checkpoint_extension->LoadCkptManager(std::move(checkpoint_manager))); + } else { + context->checkpoint_tag_extension = std::make_unique< + ::lumina::extensions::experimental::BuildWithCkptAndTagExtension>(); + PAIMON_RETURN_NOT_OK_FROM_LUMINA( + context->builder.Attach(*context->checkpoint_tag_extension)); + PAIMON_RETURN_NOT_OK_FROM_LUMINA( + context->checkpoint_tag_extension->LoadCkptManager(std::move(checkpoint_manager))); + } + } else if (with_tag) { + context->tag_extension = + std::make_unique<::lumina::extensions::experimental::BuildWithTagExtension>(); + PAIMON_RETURN_NOT_OK_FROM_LUMINA(context->builder.Attach(*context->tag_extension)); + } // 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)); + PAIMON_RETURN_NOT_OK_FROM_LUMINA(context->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)); + PAIMON_RETURN_NOT_OK_FROM_LUMINA(context->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)); + if (context->checkpoint_tag_extension) { + PAIMON_RETURN_NOT_OK_FROM_LUMINA( + context->checkpoint_tag_extension->InsertFromWithTag(insert_data)); + } else { + PAIMON_RETURN_NOT_OK_FROM_LUMINA( + context->tag_extension->InsertFromWithTag(insert_data)); + } } - return std::move(builder); + + std::vector>().swap(arrays_); + std::vector>().swap( + tag_data_vec_); + return context; } } // namespace paimon::lumina diff --git a/src/paimon/indexer/lumina/lumina_index_accumulator.h b/src/paimon/indexer/lumina/lumina_index_accumulator.h index a7b6fe69d..bad2c124b 100644 --- a/src/paimon/indexer/lumina/lumina_index_accumulator.h +++ b/src/paimon/indexer/lumina/lumina_index_accumulator.h @@ -20,11 +20,14 @@ #include #include +#include #include #include "arrow/api.h" #include "lumina/api/LuminaBuilder.h" #include "lumina/api/Options.h" +#include "lumina/extensions/experimental/BuildCombinedExtensionV0.h" +#include "lumina/extensions/experimental/CkptManager.h" #include "lumina/extensions/experimental/DatasetWithTag.h" #include "paimon/indexer/lumina/lumina_memory_pool.h" #include "paimon/indexer/lumina/lumina_tag_utils.h" @@ -33,6 +36,19 @@ namespace paimon::lumina { +/// Keeps a built Lumina index and its attached extensions alive until the index is dumped. +struct LuminaIndexBuildContext { + explicit LuminaIndexBuildContext(::lumina::api::LuminaBuilder&& value) + : builder(std::move(value)) {} + + ::lumina::api::LuminaBuilder builder; + std::unique_ptr<::lumina::extensions::experimental::BuildWithCheckpointExtension> + checkpoint_extension; + std::unique_ptr<::lumina::extensions::experimental::BuildWithTagExtension> tag_extension; + std::unique_ptr<::lumina::extensions::experimental::BuildWithCkptAndTagExtension> + checkpoint_tag_extension; +}; + /// Accumulates the non-null vector segments shared by Lumina Global Index and File Index writers. class LuminaIndexAccumulator { public: @@ -40,9 +56,10 @@ class LuminaIndexAccumulator { 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); + Result> Build( + const ::lumina::api::BuilderOptions& builder_options, uint32_t dimension, bool with_tag, + std::unique_ptr<::lumina::extensions::experimental::CkptManager> checkpoint_manager, + LuminaMemoryPool* pool); int64_t IndexedCount() const { return indexed_count_; diff --git a/src/paimon/indexer/lumina/lumina_index_options.cpp b/src/paimon/indexer/lumina/lumina_index_options.cpp index e26b4e14f..526982fbf 100644 --- a/src/paimon/indexer/lumina/lumina_index_options.cpp +++ b/src/paimon/indexer/lumina/lumina_index_options.cpp @@ -71,6 +71,12 @@ Result<::lumina::api::BuilderOptions> LuminaIndexOptions::CreateBuilderOptions( return builder_options; } +bool LuminaIndexOptions::IsCheckpointEnabled( + const std::map& lumina_options) { + return lumina_options.count(std::string(::lumina::core::kExtensionCkptThreshold)) != 0 || + lumina_options.count(std::string(::lumina::core::kExtensionCkptCount)) != 0; +} + Result<::lumina::api::SearcherOptions> LuminaIndexOptions::CreateSearcherOptions( const std::map& lumina_options, const LuminaIndexInfo& index_info) { std::map options = lumina_options; diff --git a/src/paimon/indexer/lumina/lumina_index_options.h b/src/paimon/indexer/lumina/lumina_index_options.h index ff77f032b..1c3cc355e 100644 --- a/src/paimon/indexer/lumina/lumina_index_options.h +++ b/src/paimon/indexer/lumina/lumina_index_options.h @@ -52,6 +52,8 @@ class LuminaIndexOptions { static Result<::lumina::api::BuilderOptions> CreateBuilderOptions( const std::map& lumina_options); + static bool IsCheckpointEnabled(const std::map& lumina_options); + static Result<::lumina::api::SearcherOptions> CreateSearcherOptions( const std::map& lumina_options, const LuminaIndexInfo& index_info); From 027201f487911a213a89ef8b23ef4615aa290814 Mon Sep 17 00:00:00 2001 From: "jinli.zjw" Date: Thu, 24 Sep 2026 11:23:24 +0800 Subject: [PATCH 3/3] fix(lumina): retain data for checkpoint retries --- src/paimon/indexer/lumina/lumina_index_accumulator.cpp | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/paimon/indexer/lumina/lumina_index_accumulator.cpp b/src/paimon/indexer/lumina/lumina_index_accumulator.cpp index e655c5d00..749a73bd4 100644 --- a/src/paimon/indexer/lumina/lumina_index_accumulator.cpp +++ b/src/paimon/indexer/lumina/lumina_index_accumulator.cpp @@ -120,6 +120,9 @@ Result> LuminaIndexAccumulator::Build( // insert data if (!with_tag) { LuminaDataset insert_data(indexed_count_, dimension, arrays_, array_start_ids_); + if (!context->checkpoint_extension) { + std::vector>().swap(arrays_); + } PAIMON_RETURN_NOT_OK_FROM_LUMINA(context->builder.InsertFrom(insert_data)); } else { LuminaDatasetWithTag insert_data(indexed_count_, dimension, arrays_, array_start_ids_, @@ -128,14 +131,13 @@ Result> LuminaIndexAccumulator::Build( PAIMON_RETURN_NOT_OK_FROM_LUMINA( context->checkpoint_tag_extension->InsertFromWithTag(insert_data)); } else { + std::vector>().swap(arrays_); + std::vector>().swap( + tag_data_vec_); PAIMON_RETURN_NOT_OK_FROM_LUMINA( context->tag_extension->InsertFromWithTag(insert_data)); } } - - std::vector>().swap(arrays_); - std::vector>().swap( - tag_data_vec_); return context; }