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