diff --git a/include/paimon/global_index/global_index_io_meta.h b/include/paimon/global_index/global_index_io_meta.h index e9da91cba..522762668 100644 --- a/include/paimon/global_index/global_index_io_meta.h +++ b/include/paimon/global_index/global_index_io_meta.h @@ -16,11 +16,11 @@ #pragma once +#include #include #include #include "paimon/memory/bytes.h" -#include "paimon/utils/range.h" namespace paimon { /// Metadata describing a single file entry in a global index. diff --git a/include/paimon/global_index/global_indexer.h b/include/paimon/global_index/global_indexer.h index 5e529fbe8..a5e881366 100644 --- a/include/paimon/global_index/global_indexer.h +++ b/include/paimon/global_index/global_indexer.h @@ -17,6 +17,7 @@ #pragma once #include +#include #include #include @@ -36,6 +37,11 @@ class PAIMON_EXPORT GlobalIndexer { public: virtual ~GlobalIndexer() = default; + /// Returns additional table fields required during index construction. + virtual Result>> GetExtraFieldNames() const { + return std::optional>(std::nullopt); + } + /// Creates a writer for building a global index on a specific field. /// /// @param field_name Name of the field to be indexed. diff --git a/src/paimon/core/global_index/global_index_scan_impl.cpp b/src/paimon/core/global_index/global_index_scan_impl.cpp index 1c24d2f74..393d524de 100644 --- a/src/paimon/core/global_index/global_index_scan_impl.cpp +++ b/src/paimon/core/global_index/global_index_scan_impl.cpp @@ -20,6 +20,7 @@ #include #include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" #include "paimon/common/global_index/offset_global_index_reader.h" #include "paimon/common/global_index/union_global_index_reader.h" #include "paimon/common/utils/scope_guard.h" @@ -185,7 +186,6 @@ Result>> GlobalIndexScanImpl::Cre if (row_range_index && !row_range_index->Intersects(range.from, range.to)) { continue; } - // TODO(xinyu.lxy): c_arrow_schema may contains additional associated fields. auto arrow_field = DataField::ConvertDataFieldToArrowField(field); auto arrow_schema = arrow::schema({arrow_field}); diff --git a/src/paimon/core/global_index/global_index_write_task.cpp b/src/paimon/core/global_index/global_index_write_task.cpp index d44840957..00c4f73e9 100644 --- a/src/paimon/core/global_index/global_index_write_task.cpp +++ b/src/paimon/core/global_index/global_index_write_task.cpp @@ -16,7 +16,10 @@ #include "paimon/global_index/global_index_write_task.h" +#include + #include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/arrow/status_utils.h" @@ -35,6 +38,17 @@ #include "paimon/table/source/table_read.h" namespace paimon { namespace { +Result> CreateGlobalIndexer(const std::string& index_type, + const CoreOptions& core_options) { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr indexer, + GlobalIndexerFactory::Get(index_type, core_options.ToMap())); + if (!indexer) { + return Status::Invalid( + fmt::format("Unknown index type {}, may not registered", index_type)); + } + return indexer; +} + Result> CreateGlobalIndexFileManager( const std::string& table_path, const std::shared_ptr& table_schema, const CoreOptions& core_options, const std::shared_ptr& pool) { @@ -58,25 +72,81 @@ Result> CreateGlobalIndexFileManager( } Result> CreateGlobalIndexWriter( - const std::string& index_type, const DataField& field, + const GlobalIndexer& indexer, const DataField& field, + const std::vector& extra_fields, const std::shared_ptr& index_file_manager, - const CoreOptions& core_options, const std::shared_ptr& pool) { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr indexer, - GlobalIndexerFactory::Get(index_type, core_options.ToMap())); - if (!indexer) { - return Status::Invalid( - fmt::format("Unknown index type {}, may not registered", index_type)); + const std::shared_ptr& pool) { + arrow::FieldVector arrow_fields; + arrow_fields.reserve(extra_fields.size() + 1); + arrow_fields.push_back(DataField::ConvertDataFieldToArrowField(field)); + for (const auto& extra_field : extra_fields) { + arrow_fields.push_back(DataField::ConvertDataFieldToArrowField(extra_field)); } - // TODO(xinyu.lxy): may add additional fields to read for index write - auto arrow_field = DataField::ConvertDataFieldToArrowField(field); - auto arrow_schema = arrow::schema({arrow_field}); + auto arrow_schema = arrow::schema(arrow_fields); ArrowSchema c_arrow_schema; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*arrow_schema, &c_arrow_schema)); - return indexer->CreateWriter(field.Name(), &c_arrow_schema, index_file_manager, pool); + ScopeGuard guard([&]() { ArrowSchemaRelease(&c_arrow_schema); }); + return indexer.CreateWriter(field.Name(), &c_arrow_schema, index_file_manager, pool); +} + +Result> GetExtraFields(const TableSchema& table_schema, + const std::string& field_name, + const std::vector& extra_field_names) { + std::vector extra_fields; + extra_fields.reserve(extra_field_names.size()); + std::set dedup_field_names; + for (const auto& extra_field_name : extra_field_names) { + if (extra_field_name == field_name) { + return Status::Invalid(fmt::format( + "global index extra field {} must not be the indexed field", extra_field_name)); + } + if (!dedup_field_names.insert(extra_field_name).second) { + return Status::Invalid(fmt::format("global index extra field {} must not be duplicated", + extra_field_name)); + } + PAIMON_ASSIGN_OR_RAISE(DataField extra_field, table_schema.GetField(extra_field_name)); + extra_fields.push_back(extra_field); + } + return extra_fields; +} + +std::vector BuildReadFieldNames(const std::string& field_name, + const std::vector& extra_fields) { + std::vector read_field_names; + read_field_names.reserve(extra_fields.size() + 2); + read_field_names.push_back(field_name); + for (const auto& extra_field : extra_fields) { + read_field_names.push_back(extra_field.Name()); + } + read_field_names.push_back(SpecialFields::RowId().Name()); + return read_field_names; +} + +std::vector BuildWriterFieldNames(const std::string& field_name, + const std::vector& extra_fields) { + std::vector writer_field_names; + writer_field_names.reserve(extra_fields.size() + 1); + writer_field_names.push_back(field_name); + for (const auto& extra_field : extra_fields) { + writer_field_names.push_back(extra_field.Name()); + } + return writer_field_names; +} + +std::optional> GetExtraFieldIds(const std::vector& extra_fields) { + if (extra_fields.empty()) { + return std::nullopt; + } + std::vector extra_field_ids; + extra_field_ids.reserve(extra_fields.size()); + for (const auto& extra_field : extra_fields) { + extra_field_ids.push_back(extra_field.Id()); + } + return extra_field_ids; } Result> CreateBatchReader( - const std::string& table_path, const std::string& field_name, + const std::string& table_path, const std::vector& read_field_names, const std::shared_ptr& indexed_split, const CoreOptions& core_options, const std::shared_ptr& pool) { ReadContextBuilder read_context_builder(table_path); @@ -84,7 +154,7 @@ Result> CreateBatchReader( .WithFileSystem(core_options.GetFileSystem()) .EnablePrefetch(true) .WithMemoryPool(pool) - .SetReadFieldNames({field_name, SpecialFields::RowId().Name()}); + .SetReadFieldNames(read_field_names); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, read_context_builder.Finish()); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr table_read, @@ -92,9 +162,10 @@ Result> CreateBatchReader( return table_read->CreateReader(indexed_split); } -Result> BuildIndex(const std::string& field_name, const Range& range, - BatchReader* batch_reader, - GlobalIndexWriter* global_index_writer) { +Result> BuildIndex( + const std::string& field_name, const Range& range, + const std::vector& writer_field_names, BatchReader* batch_reader, + GlobalIndexWriter* global_index_writer) { while (true) { PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch read_batch, batch_reader->NextBatch()); if (BatchReader::IsEofBatch(read_batch)) { @@ -108,11 +179,6 @@ Result> BuildIndex(const std::string& field_name, return Status::Invalid( "array read from batch reader is not a struct array in GlobalIndexWriteTask"); } - auto indexed_array = struct_array->GetFieldByName(field_name); - if (!indexed_array) { - return Status::Invalid(fmt::format( - "read array does not contain {} field in GlobalIndexWriteTask", field_name)); - } auto row_id_array = struct_array->GetFieldByName(SpecialFields::RowId().Name()); auto typed_row_id_array = std::dynamic_pointer_cast(row_id_array); if (!typed_row_id_array) { @@ -131,8 +197,20 @@ Result> BuildIndex(const std::string& field_name, } relative_row_ids.push_back(row_id - range.from); } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr new_array, - arrow::StructArray::Make({indexed_array}, {field_name})); + std::vector> writer_arrays; + writer_arrays.reserve(writer_field_names.size()); + for (const auto& writer_field_name : writer_field_names) { + auto writer_array = struct_array->GetFieldByName(writer_field_name); + if (!writer_array) { + return Status::Invalid( + fmt::format("read array does not contain {} field in GlobalIndexWriteTask", + writer_field_name)); + } + writer_arrays.push_back(writer_array); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr new_array, + arrow::StructArray::Make(writer_arrays, writer_field_names)); ::ArrowArray c_new_array; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*new_array, &c_new_array)); PAIMON_RETURN_NOT_OK( @@ -144,7 +222,8 @@ Result> BuildIndex(const std::string& field_name, Result> ToCommitMessage( const std::string& index_type, int32_t field_id, const Range& range, const std::vector& global_index_io_metas, const BinaryRow& partition, - int32_t bucket, const std::shared_ptr& file_manager) { + int32_t bucket, const std::shared_ptr& file_manager, + const std::optional>& extra_field_ids) { std::vector> index_file_metas; index_file_metas.reserve(global_index_io_metas.size()); bool is_external_path = file_manager->IsExternalPath(); @@ -157,8 +236,7 @@ Result> ToCommitMessage( index_file_metas.push_back(std::make_shared( index_type, PathUtil::GetName(io_meta.file_path), io_meta.file_size, range.Count(), /*dv_ranges=*/std::nullopt, external_path, - GlobalIndexMeta(range.from, range.to, field_id, - /*extra_field_ids=*/std::nullopt, io_meta.metadata))); + GlobalIndexMeta(range.from, range.to, field_id, extra_field_ids, io_meta.metadata))); } DataIncrement data_increment(std::move(index_file_metas)); return std::make_shared(partition, bucket, @@ -199,6 +277,17 @@ Result> GlobalIndexWriteTask::WriteIndex( } PAIMON_ASSIGN_OR_RAISE(CoreOptions core_options, CoreOptions::FromMap(final_options, file_system)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr indexer, + CreateGlobalIndexer(index_type, core_options)); + PAIMON_ASSIGN_OR_RAISE(std::optional> extra_field_names, + indexer->GetExtraFieldNames()); + PAIMON_ASSIGN_OR_RAISE(DataField field, table_schema->GetField(field_name)); + PAIMON_ASSIGN_OR_RAISE(std::vector extra_fields, + GetExtraFields(*table_schema, field_name, + extra_field_names.value_or(std::vector()))); + std::optional> extra_field_ids = GetExtraFieldIds(extra_fields); + std::vector writer_field_names = BuildWriterFieldNames(field_name, extra_fields); + std::vector read_field_names = BuildReadFieldNames(field_name, extra_fields); // create index file manager PAIMON_ASSIGN_OR_RAISE( @@ -208,13 +297,12 @@ Result> GlobalIndexWriteTask::WriteIndex( // create batch reader PAIMON_ASSIGN_OR_RAISE( std::unique_ptr batch_reader, - CreateBatchReader(table_path, field_name, indexed_split, core_options, pool)); + CreateBatchReader(table_path, read_field_names, indexed_split, core_options, pool)); // create global index writer - PAIMON_ASSIGN_OR_RAISE(DataField field, table_schema->GetField(field_name)); PAIMON_ASSIGN_OR_RAISE( std::shared_ptr global_index_writer, - CreateGlobalIndexWriter(index_type, field, index_file_manager, core_options, pool)); + CreateGlobalIndexWriter(*indexer, field, extra_fields, index_file_manager, pool)); ScopeGuard guard([&]() { global_index_writer.reset(); @@ -222,13 +310,14 @@ Result> GlobalIndexWriteTask::WriteIndex( }); // read from data split and write to index writer - PAIMON_ASSIGN_OR_RAISE( - std::vector global_index_io_metas, - BuildIndex(field_name, range, batch_reader.get(), global_index_writer.get())); + PAIMON_ASSIGN_OR_RAISE(std::vector global_index_io_metas, + BuildIndex(field_name, range, writer_field_names, batch_reader.get(), + global_index_writer.get())); // generate commit message return ToCommitMessage(index_type, field.Id(), range, global_index_io_metas, - data_split->Partition(), data_split->Bucket(), index_file_manager); + data_split->Partition(), data_split->Bucket(), index_file_manager, + extra_field_ids); } } // namespace paimon diff --git a/src/paimon/global_index/lumina/lumina_global_index.cpp b/src/paimon/global_index/lumina/lumina_global_index.cpp index 8e4af3e49..00bac8747 100644 --- a/src/paimon/global_index/lumina/lumina_global_index.cpp +++ b/src/paimon/global_index/lumina/lumina_global_index.cpp @@ -16,6 +16,10 @@ #include "paimon/global_index/lumina/lumina_global_index.h" +#include +#include +#include +#include #include #include "arrow/c/bridge.h" @@ -27,6 +31,7 @@ #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/options_utils.h" #include "paimon/common/utils/rapidjson_util.h" @@ -35,6 +40,9 @@ #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 { \ @@ -43,6 +51,458 @@ namespace paimon::lumina { } \ } 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)); + } + if (key_name.size() > ::lumina::core::kMaxTagKNameLength) { + return Status::Invalid(fmt::format("lumina tag_schema {} key_name exceeds max length {}", + tag_label, ::lumina::core::kMaxTagKNameLength)); + } + + 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::kExtensionTagVTypeInt64)) { + parsed_value_type = LuminaTagField::ValueType::INT64; + } 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)); + } + if (parsed_type == LuminaTagField::Type::RANGE && + parsed_value_type == LuminaTagField::ValueType::STRING) { + return Status::Invalid(fmt::format( + "lumina tag_schema {} range type does not support string value_type", tag_label)); + } + 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(); + } + + // Lumina currently downcasts range tag values to 32-bit precision internally. + // Reject Arrow int64/double inputs to avoid silent precision loss. + bool compatible = false; + switch (tag_field.value_type) { + case LuminaTagField::ValueType::INT64: + compatible = value_type->id() == arrow::Type::INT8 || + value_type->id() == arrow::Type::INT16 || + value_type->id() == arrow::Type::INT32; + break; + case LuminaTagField::ValueType::DOUBLE: + compatible = value_type->id() == arrow::Type::FLOAT; + 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(static_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(); + } + + 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) { + switch (array->type_id()) { + case arrow::Type::FLOAT: + AppendPrimitiveTagValue(array, index, values); + break; + default: + return Status::Invalid( + fmt::format("lumina floating tag field has unsupported arrow type {}", + array->type()->ToString())); + } + } else if constexpr (std::is_same_v) { + if (array->type_id() != arrow::Type::STRING) { + return Status::Invalid( + fmt::format("lumina string tag field has unsupported arrow type {}", + array->type()->ToString())); + } + auto string_array = static_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(static_cast(literal.GetValue())); + case FieldType::FLOAT: + return TagValue(static_cast(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::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::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::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::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, @@ -65,6 +525,8 @@ Result> LuminaGlobalIndex::CreateWriter( // 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))); @@ -76,7 +538,7 @@ Result> LuminaGlobalIndex::CreateWriter( 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, lumina_pool); + ::lumina::api::IOOptions(), lumina_options, std::move(tag_fields), lumina_pool); } Result LuminaIndexReader::GetIndexInfo( @@ -111,7 +573,9 @@ Result LuminaIndexReader::GetIndexInfo( return Status::Invalid( fmt::format("invalid distance type {} for lumina", distance_type_str)); } - return LuminaIndexReader::IndexInfo({dimension, index_type, distance_type}); + 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( @@ -170,8 +634,30 @@ Result> LuminaGlobalIndex::CreateReader( } 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), lumina_pool); + 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 { @@ -199,18 +685,18 @@ class LuminaDataset : public ::lumina::api::Dataset { } auto& value_array = array_vec_[cursor_]; int64_t value_array_length = value_array->length(); - int64_t element_count = value_array_length / dimension_; + 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(element_count); + 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(element_count)); + return ::lumina::core::Result::Ok(static_cast(batch_element_count)); } private: @@ -221,14 +707,62 @@ class LuminaDataset : public ::lumina::api::Dataset { 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, - const std::shared_ptr& pool) +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), @@ -236,7 +770,8 @@ LuminaIndexWriter::LuminaIndexWriter(const std::string& field_name, file_manager_(file_manager), builder_options_(std::move(builder_options)), io_options_(std::move(io_options)), - lumina_options_(lumina_options) {} + lumina_options_(lumina_options), + tag_fields_(std::move(tag_fields)) {} Status LuminaIndexWriter::AddBatch(::ArrowArray* arrow_array, std::vector&& relative_row_ids) { @@ -289,6 +824,12 @@ Status LuminaIndexWriter::AddBatch(::ArrowArray* arrow_array, "multiplied dimension [{}] must match length of field value array [{}]", segment_len, dimension_, sliced_values->length())); } + 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; @@ -313,9 +854,19 @@ Result> LuminaIndexWriter::Finish() { PAIMON_RETURN_NOT_OK_FROM_LUMINA(builder.PretrainFrom(dataset1)); // insert data - LuminaDataset dataset2(indexed_count_, dimension_, array_vec_, array_start_ids_); - std::vector>().swap(array_vec_); - PAIMON_RETURN_NOT_OK_FROM_LUMINA(builder.InsertFrom(dataset2)); + 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, @@ -338,17 +889,16 @@ 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_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->predicate) { - return Status::NotImplemented("lumina index not support predicate in VisitVectorSearch"); - } 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"); @@ -375,7 +925,25 @@ Result> LuminaIndexReader::VisitVectorS ::lumina::api::Query lumina_query(vector_search->query.data(), vector_search->query.size()); ::lumina::api::LuminaSearcher::SearchResult search_result; - if (!vector_search->pre_filter) { + 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 { diff --git a/src/paimon/global_index/lumina/lumina_global_index.h b/src/paimon/global_index/lumina/lumina_global_index.h index e2179ca3d..108d73f49 100644 --- a/src/paimon/global_index/lumina/lumina_global_index.h +++ b/src/paimon/global_index/lumina/lumina_global_index.h @@ -18,20 +18,42 @@ #include #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" namespace paimon::lumina { +struct LuminaTagField { + enum class Type { + ENUM, + RANGE, + }; + + enum class ValueType { + INT64, + 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.`**. /// More options refer to OptionsReference.md in third_party/lumina. @@ -61,6 +83,8 @@ class LuminaGlobalIndex : public GlobalIndexer { explicit LuminaGlobalIndex(const std::map& options) : options_(options) {} + Result>> GetExtraFieldNames() const override; + Result> CreateWriter( const std::string& field_name, ::ArrowSchema* arrow_schema, const std::shared_ptr& file_writer, @@ -72,6 +96,12 @@ 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_; }; @@ -83,6 +113,7 @@ class LuminaIndexWriter : public GlobalIndexWriter { ::lumina::api::BuilderOptions&& builder_options, ::lumina::api::IOOptions&& io_options, const std::map& lumina_options, + std::vector&& tag_fields, const std::shared_ptr& pool); Status AddBatch(::ArrowArray* arrow_array, std::vector&& relative_row_ids) override; @@ -90,6 +121,11 @@ 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_; @@ -100,8 +136,10 @@ class LuminaIndexWriter : public GlobalIndexWriter { ::lumina::api::BuilderOptions builder_options_; ::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_; }; class LuminaIndexReader : public GlobalIndexReader { @@ -110,11 +148,14 @@ class LuminaIndexReader : public GlobalIndexReader { 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 { @@ -201,9 +242,13 @@ 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_; }; } // namespace paimon::lumina diff --git a/src/paimon/global_index/lumina/lumina_global_index_test.cpp b/src/paimon/global_index/lumina/lumina_global_index_test.cpp index e5aeca71a..ddd853770 100644 --- a/src/paimon/global_index/lumina/lumina_global_index_test.cpp +++ b/src/paimon/global_index/lumina/lumina_global_index_test.cpp @@ -250,6 +250,411 @@ TEST_F(LuminaGlobalIndexTest, TestWithFilter) { } } +TEST_F(LuminaGlobalIndexTest, TestWriteAndReadWithTagFilter) { + auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(test_root_dir); + std::string test_root = test_root_dir->Str(); + + std::map tag_options = options_; + tag_options["lumina.extension.build.tag.tag_schema"] = + R"([{"key_name":"color","type":"enum","value_type":"string"}])"; + + std::shared_ptr tag_data_type = arrow::struct_( + {arrow::field("f0", arrow::list(arrow::float32())), arrow::field("color", arrow::utf8())}); + std::shared_ptr tag_array = + arrow::ipc::internal::json::ArrayFromJSON(tag_data_type, + R"([ + [[0.0, 0.0, 0.0, 0.0], "cold"], + [[0.0, 1.0, 0.0, 1.0], "warm"], + [[1.0, 0.0, 1.0, 0.0], "cold"], + [[1.0, 1.0, 1.0, 1.0], "warm"] + ])") + .ValueOrDie(); + + ASSERT_OK_AND_ASSIGN( + GlobalIndexIOMeta meta, + WriteGlobalIndex(test_root, tag_data_type, tag_options, tag_array, Range(0, 3))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr reader, + CreateGlobalIndexReader(test_root, data_type_, tag_options, meta)); + + std::shared_ptr predicate = PredicateBuilder::Equal( + /*field_index=*/1, /*field_name=*/"color", FieldType::STRING, + Literal(FieldType::STRING, "warm", 4)); + { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr scored_result, + reader->VisitVectorSearch(std::make_shared( + /*field_name=*/"f0", /*limit=*/4, query_, /*filter=*/nullptr, predicate, + /*distance_type=*/std::nullopt, /*options=*/tag_options))); + CheckResult(scored_result, {3l, 1l}, {0.01f, 2.01f}); + } + { + auto pre_filter = [](int64_t id) -> bool { return id < 3; }; + ASSERT_OK_AND_ASSIGN(std::shared_ptr filtered_scored_result, + reader->VisitVectorSearch(std::make_shared( + /*field_name=*/"f0", /*limit=*/4, query_, pre_filter, predicate, + /*distance_type=*/std::nullopt, /*options=*/tag_options))); + CheckResult(filtered_scored_result, {1l}, {2.01f}); + } +} + +TEST_F(LuminaGlobalIndexTest, TestWriteAndReadWithMixedTagPredicates) { + auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(test_root_dir); + std::string test_root = test_root_dir->Str(); + + std::map tag_options = options_; + tag_options["lumina.extension.build.tag.tag_schema"] = + R"([{"key_name":"color","type":"enum","value_type":"string"},)" + R"({"key_name":"price","type":"range","value_type":"double"}])"; + + std::shared_ptr tag_data_type = arrow::struct_( + {arrow::field("f0", arrow::list(arrow::float32())), arrow::field("color", arrow::utf8()), + arrow::field("price", arrow::float32())}); + std::shared_ptr tag_array = + arrow::ipc::internal::json::ArrayFromJSON(tag_data_type, + R"([ + [[0.0, 0.0, 0.0, 0.0], "cold", 5.0], + [[0.0, 1.0, 0.0, 1.0], "warm", 10.0], + [[1.0, 0.0, 1.0, 0.0], "warm", 20.0], + [[1.0, 1.0, 1.0, 1.0], "warm", 30.0] + ])") + .ValueOrDie(); + + ASSERT_OK_AND_ASSIGN( + GlobalIndexIOMeta meta, + WriteGlobalIndex(test_root, tag_data_type, tag_options, tag_array, Range(0, 3))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr reader, + CreateGlobalIndexReader(test_root, data_type_, tag_options, meta)); + + std::shared_ptr color_predicate = PredicateBuilder::Equal( + /*field_index=*/1, /*field_name=*/"color", FieldType::STRING, + Literal(FieldType::STRING, "warm", 4)); + std::shared_ptr low_price_predicate = PredicateBuilder::LessOrEqual( + /*field_index=*/2, /*field_name=*/"price", FieldType::FLOAT, Literal(10.0f)); + std::shared_ptr high_price_predicate = PredicateBuilder::GreaterOrEqual( + /*field_index=*/2, /*field_name=*/"price", FieldType::FLOAT, Literal(30.0f)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr price_predicate, + PredicateBuilder::Or({low_price_predicate, high_price_predicate})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr predicate, + PredicateBuilder::And({color_predicate, price_predicate})); + + ASSERT_OK_AND_ASSIGN( + std::shared_ptr scored_result, + reader->VisitVectorSearch(std::make_shared( + /*field_name=*/"f0", /*limit=*/4, query_, /*filter=*/nullptr, predicate, + /*distance_type=*/std::nullopt, /*options=*/tag_options))); + CheckResult(scored_result, {3l, 1l}, {0.01f, 2.01f}); +} + +TEST_F(LuminaGlobalIndexTest, TestWriteAndReadWithIntegerListTagFilter) { + auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(test_root_dir); + std::string test_root = test_root_dir->Str(); + + std::map tag_options = options_; + tag_options["lumina.extension.build.tag.tag_schema"] = + R"([{"key_name":"category_ids","type":"enum","value_type":"int64"}])"; + + std::shared_ptr tag_data_type = + arrow::struct_({arrow::field("f0", arrow::list(arrow::float32())), + arrow::field("category_ids", arrow::list(arrow::int32()))}); + std::shared_ptr tag_array = + arrow::ipc::internal::json::ArrayFromJSON(tag_data_type, + R"([ + [[0.0, 0.0, 0.0, 0.0], [1, 2]], + [[0.0, 1.0, 0.0, 1.0], [3, 8]], + [[1.0, 0.0, 1.0, 0.0], [4, 5]], + [[1.0, 1.0, 1.0, 1.0], [6, 9]] + ])") + .ValueOrDie(); + + ASSERT_OK_AND_ASSIGN( + GlobalIndexIOMeta meta, + WriteGlobalIndex(test_root, tag_data_type, tag_options, tag_array, Range(0, 3))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr reader, + CreateGlobalIndexReader(test_root, data_type_, tag_options, meta)); + + std::shared_ptr predicate = PredicateBuilder::In( + /*field_index=*/1, /*field_name=*/"category_ids", FieldType::INT, {Literal(8), Literal(9)}); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr scored_result, + reader->VisitVectorSearch(std::make_shared( + /*field_name=*/"f0", /*limit=*/4, query_, /*filter=*/nullptr, predicate, + /*distance_type=*/std::nullopt, /*options=*/tag_options))); + CheckResult(scored_result, {3l, 1l}, {0.01f, 2.01f}); +} + +TEST_F(LuminaGlobalIndexTest, TestWriteAndReadWithCompatibleTagArrowTypes) { + auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(test_root_dir); + std::string test_root = test_root_dir->Str(); + + std::map tag_options = options_; + tag_options["lumina.extension.build.tag.tag_schema"] = + R"([{"key_name":"tag_i8","type":"enum","value_type":"int64"},)" + R"({"key_name":"tag_i16","type":"enum","value_type":"int64"},)" + R"({"key_name":"tag_i32","type":"enum","value_type":"int64"},)" + R"({"key_name":"tag_f32","type":"range","value_type":"double"}])"; + + std::shared_ptr tag_data_type = arrow::struct_( + {arrow::field("f0", arrow::list(arrow::float32())), arrow::field("tag_i8", arrow::int8()), + arrow::field("tag_i16", arrow::int16()), arrow::field("tag_i32", arrow::int32()), + arrow::field("tag_f32", arrow::float32())}); + std::shared_ptr tag_array = + arrow::ipc::internal::json::ArrayFromJSON(tag_data_type, + R"([ + [[0.0, 0.0, 0.0, 0.0], 1, 10, 100, 1.5], + [[0.0, 1.0, 0.0, 1.0], 2, 20, 200, 2.5], + [[1.0, 0.0, 1.0, 0.0], 3, 30, 300, 3.5], + [[1.0, 1.0, 1.0, 1.0], 4, 40, 400, 4.5] + ])") + .ValueOrDie(); + + ASSERT_OK_AND_ASSIGN( + GlobalIndexIOMeta meta, + WriteGlobalIndex(test_root, tag_data_type, tag_options, tag_array, Range(0, 3))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr reader, + CreateGlobalIndexReader(test_root, data_type_, tag_options, meta)); + + auto search_and_check = [&](const std::shared_ptr& predicate, + const std::vector& expected_ids, + const std::vector& expected_scores) { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr scored_result, + reader->VisitVectorSearch(std::make_shared( + /*field_name=*/"f0", /*limit=*/4, query_, /*filter=*/nullptr, predicate, + /*distance_type=*/std::nullopt, /*options=*/tag_options))); + CheckResult(scored_result, expected_ids, expected_scores); + }; + + search_and_check(PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"tag_i8", + FieldType::TINYINT, Literal(static_cast(2))), + {1l}, {2.01f}); + search_and_check( + PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"tag_i16", FieldType::SMALLINT, + Literal(static_cast(30))), + {2l}, {2.21f}); + search_and_check(PredicateBuilder::Equal(/*field_index=*/3, /*field_name=*/"tag_i32", + FieldType::INT, Literal(400)), + {3l}, {0.01f}); + search_and_check(PredicateBuilder::GreaterThan(/*field_index=*/4, /*field_name=*/"tag_f32", + FieldType::FLOAT, Literal(4.0f)), + {3l}, {0.01f}); +} + +TEST_F(LuminaGlobalIndexTest, TestWriteAndReadWithTagNullAndEmptyValues) { + auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(test_root_dir); + std::string test_root = test_root_dir->Str(); + + std::map tag_options = options_; + tag_options["lumina.extension.build.tag.tag_schema"] = + R"([{"key_name":"color","type":"enum","value_type":"string"},)" + R"({"key_name":"labels","type":"enum","value_type":"string"},)" + R"({"key_name":"price","type":"range","value_type":"double"},)" + R"({"key_name":"scores","type":"enum","value_type":"double"},)" + R"({"key_name":"category","type":"enum","value_type":"int64"},)" + R"({"key_name":"category_ids","type":"enum","value_type":"int64"}])"; + + std::shared_ptr tag_data_type = arrow::struct_( + {arrow::field("f0", arrow::list(arrow::float32())), arrow::field("color", arrow::utf8()), + arrow::field("labels", arrow::list(arrow::utf8())), + arrow::field("price", arrow::float32()), + arrow::field("scores", arrow::list(arrow::float32())), + arrow::field("category", arrow::int32()), + arrow::field("category_ids", arrow::list(arrow::int32()))}); + std::shared_ptr tag_array = + arrow::ipc::internal::json::ArrayFromJSON(tag_data_type, + R"([ + [[0.0, 0.0, 0.0, 0.0], "red", ["hot"], 5.0, [0.25], 7, [1, 2]], + [null, "red", ["vip"], 8.0, [0.8], 7, [9]], + [[0.0, 1.0, 0.0, 1.0], null, null, null, null, null, null], + [[1.0, 0.0, 1.0, 0.0], " ", [], 20.0, [], 0, []], + [[1.0, 1.0, 1.0, 1.0], "blue", ["vip", null], 30.0, [null, 0.75], 3, [null, 9]] + ])") + .ValueOrDie(); + + ASSERT_OK_AND_ASSIGN( + GlobalIndexIOMeta meta, + WriteGlobalIndex(test_root, tag_data_type, tag_options, tag_array, Range(0, 4))); + ASSERT_OK_AND_ASSIGN(std::shared_ptr reader, + CreateGlobalIndexReader(test_root, data_type_, tag_options, meta)); + auto search_and_check_with_filter = + [&](VectorSearch::PreFilter pre_filter, const std::shared_ptr& predicate, + const std::vector& expected_ids, const std::vector& expected_scores) { + std::shared_ptr vector_search = std::make_shared( + /*field_name=*/"f0", /*limit=*/5, query_, pre_filter, predicate, + /*distance_type=*/std::nullopt, /*options=*/tag_options); + ASSERT_OK_AND_ASSIGN(std::shared_ptr scored_result, + reader->VisitVectorSearch(vector_search)); + CheckResult(scored_result, expected_ids, expected_scores); + }; + auto search_and_check = [&](const std::shared_ptr& predicate, + const std::vector& expected_ids, + const std::vector& expected_scores) { + search_and_check_with_filter(/*pre_filter=*/nullptr, predicate, expected_ids, + expected_scores); + }; + auto search_and_check_error = [&](const std::shared_ptr& predicate, + const std::string& expected_message) { + ASSERT_NOK_WITH_MSG( + reader->VisitVectorSearch(std::make_shared( + /*field_name=*/"f0", /*limit=*/5, query_, /*filter=*/nullptr, predicate, + /*distance_type=*/std::nullopt, /*options=*/tag_options)), + expected_message); + }; + + search_and_check(/*predicate=*/nullptr, {4l, 2l, 3l, 0l}, {0.01f, 2.01f, 2.21f, 4.21f}); + search_and_check( + PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"color", FieldType::STRING, + Literal(FieldType::STRING, "red", 3)), + {0l}, {4.21f}); + search_and_check(PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"color", + FieldType::STRING, Literal(FieldType::STRING, " ", 1)), + {3l}, {2.21f}); + search_and_check( + PredicateBuilder::In(/*field_index=*/2, /*field_name=*/"labels", FieldType::STRING, + {Literal(FieldType::STRING, "vip", 3)}), + {4l}, {0.01f}); + search_and_check(PredicateBuilder::LessOrEqual(/*field_index=*/3, /*field_name=*/"price", + FieldType::FLOAT, Literal(10.0f)), + {0l}, {4.21f}); + search_and_check(PredicateBuilder::LessThan(/*field_index=*/3, /*field_name=*/"price", + FieldType::FLOAT, Literal(10.0f)), + {0l}, {4.21f}); + search_and_check(PredicateBuilder::GreaterThan(/*field_index=*/3, /*field_name=*/"price", + FieldType::FLOAT, Literal(10.0f)), + {4l, 3l}, {0.01f, 2.21f}); + search_and_check(PredicateBuilder::In(/*field_index=*/4, /*field_name=*/"scores", + FieldType::FLOAT, {Literal(0.25f), Literal(0.75f)}), + {4l, 0l}, {0.01f, 4.21f}); + search_and_check(PredicateBuilder::Equal(/*field_index=*/5, /*field_name=*/"category", + FieldType::INT, Literal(7)), + {0l}, {4.21f}); + search_and_check(PredicateBuilder::In(/*field_index=*/6, /*field_name=*/"category_ids", + FieldType::INT, {Literal(9)}), + {4l}, {0.01f}); + search_and_check( + PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"color", FieldType::STRING, + Literal(FieldType::STRING, "green", 5)), + {}, {}); + search_and_check_error( + PredicateBuilder::NotIn(/*field_index=*/1, /*field_name=*/"color", FieldType::STRING, + {Literal(FieldType::STRING, "red", 3)}), + "lumina tag predicate does not support leaf function NotIn"); + search_and_check_error( + PredicateBuilder::NotEqual(/*field_index=*/1, /*field_name=*/"color", FieldType::STRING, + Literal(FieldType::STRING, "red", 3)), + "lumina tag predicate does not support leaf function NotEqual"); + search_and_check_error( + PredicateBuilder::Equal(/*field_index=*/7, /*field_name=*/"unknown", FieldType::STRING, + Literal(FieldType::STRING, "red", 3)), + "unknown tag key 'unknown' in label filter"); + search_and_check_error(PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"color", + FieldType::INT, Literal(1)), + "tag value type mismatch for key 'color'"); + search_and_check_error( + PredicateBuilder::Equal(/*field_index=*/3, /*field_name=*/"price", FieldType::STRING, + Literal(FieldType::STRING, "x", 1)), + "tag value type mismatch for key 'price'"); + { + std::shared_ptr predicate = PredicateBuilder::Equal( + /*field_index=*/1, /*field_name=*/"color", FieldType::STRING, + Literal(FieldType::STRING, "red", 3)); + search_and_check_with_filter([](int64_t id) { return id == 0 || id == 4; }, predicate, {0l}, + {4.21f}); + search_and_check_with_filter([](int64_t id) { return id == 4; }, predicate, {}, {}); + } + { + std::shared_ptr red_predicate = PredicateBuilder::Equal( + /*field_index=*/1, /*field_name=*/"color", FieldType::STRING, + Literal(FieldType::STRING, "red", 3)); + std::shared_ptr blue_predicate = PredicateBuilder::Equal( + /*field_index=*/1, /*field_name=*/"color", FieldType::STRING, + Literal(FieldType::STRING, "blue", 4)); + std::shared_ptr high_price_predicate = PredicateBuilder::GreaterOrEqual( + /*field_index=*/3, /*field_name=*/"price", FieldType::FLOAT, Literal(30.0f)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr blue_high_price_predicate, + PredicateBuilder::And({blue_predicate, high_price_predicate})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr compound_predicate, + PredicateBuilder::Or({red_predicate, blue_high_price_predicate})); + search_and_check(compound_predicate, {0l, 4l}, {4.21f, 0.01f}); + search_and_check_with_filter([](int64_t id) { return id == 4; }, compound_predicate, {4l}, + {0.01f}); + } +} + +TEST_F(LuminaGlobalIndexTest, TestTagSchemaValidation) { + auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); + ASSERT_TRUE(test_root_dir); + std::string index_root = test_root_dir->Str(); + + std::shared_ptr tag_data_type = arrow::struct_( + {arrow::field("f0", arrow::list(arrow::float32())), arrow::field("color", arrow::utf8())}); + + { + std::map tag_options = options_; + tag_options["lumina.extension.build.tag.tag_schema"] = + R"([{"key_name":"color","type":"range","value_type":"string"}])"; + ASSERT_NOK_WITH_MSG( + WriteGlobalIndex(index_root, tag_data_type, tag_options, array_, Range(0, 3)), + "lumina tag_schema tag[0] range type does not support string value_type"); + } + { + std::map tag_options = options_; + tag_options["lumina.extension.build.tag.tag_schema"] = + R"([{"key_name":"color","type":"enum","value_type":"int64"}])"; + ASSERT_NOK_WITH_MSG( + WriteGlobalIndex(index_root, tag_data_type, tag_options, array_, Range(0, 3)), + "lumina tag field color type string is not compatible with tag_schema value_type"); + } + { + std::shared_ptr int64_tag_data_type = + arrow::struct_({arrow::field("f0", arrow::list(arrow::float32())), + arrow::field("category", arrow::int64())}); + std::map tag_options = options_; + tag_options["lumina.extension.build.tag.tag_schema"] = + R"([{"key_name":"category","type":"enum","value_type":"int64"}])"; + ASSERT_NOK_WITH_MSG( + WriteGlobalIndex(index_root, int64_tag_data_type, tag_options, array_, Range(0, 3)), + "lumina tag field category type int64 is not compatible with tag_schema value_type"); + } + { + std::shared_ptr double_tag_data_type = + arrow::struct_({arrow::field("f0", arrow::list(arrow::float32())), + arrow::field("price", arrow::float64())}); + std::map tag_options = options_; + tag_options["lumina.extension.build.tag.tag_schema"] = + R"([{"key_name":"price","type":"range","value_type":"double"}])"; + ASSERT_NOK_WITH_MSG( + WriteGlobalIndex(index_root, double_tag_data_type, tag_options, array_, Range(0, 3)), + "lumina tag field price type double is not compatible with tag_schema value_type"); + } +} + +TEST_F(LuminaGlobalIndexTest, TestGetExtraFieldNames) { + { + LuminaGlobalIndex global_index(options_); + ASSERT_OK_AND_ASSIGN(std::optional> field_names, + global_index.GetExtraFieldNames()); + ASSERT_FALSE(field_names); + } + { + std::map tag_options = options_; + tag_options["lumina.extension.build.tag.tag_schema"] = + R"([{"key_name":"color","type":"enum","value_type":"string"},)" + R"({"key_name":"price","type":"range","value_type":"double"},)" + R"({"key_name":"category_ids","type":"enum","value_type":"int64"}])"; + LuminaGlobalIndex global_index(tag_options); + ASSERT_OK_AND_ASSIGN(std::optional> field_names, + global_index.GetExtraFieldNames()); + ASSERT_TRUE(field_names); + ASSERT_EQ(field_names.value(), + std::vector({"color", "price", "category_ids"})); + } +} + TEST_F(LuminaGlobalIndexTest, TestInvalidInputs) { auto test_root_dir = paimon::test::UniqueTestDirectory::Create(); ASSERT_TRUE(test_root_dir); @@ -372,10 +777,10 @@ TEST_F(LuminaGlobalIndexTest, TestInvalidInputs) { "f1", /*limit=*/2, query_, /*filter=*/nullptr, PredicateBuilder::Equal(/*field_index=*/1, /*field_name=*/"f0", - FieldType::BIGINT, Literal(5l)), + FieldType::INT, Literal(5)), /*distance_type=*/std::nullopt, /*options=*/std::map())), - "lumina index not support predicate in VisitVectorSearch"); + "lumina index was not built with tag"); } { ASSERT_OK_AND_ASSIGN(auto reader, diff --git a/test/inte/global_index_test.cpp b/test/inte/global_index_test.cpp index d59c11089..978d2255f 100644 --- a/test/inte/global_index_test.cpp +++ b/test/inte/global_index_test.cpp @@ -1133,6 +1133,158 @@ TEST_P(GlobalIndexTest, TestWriteCommitScanReadIndexWithScore) { ASSERT_FALSE(typed_result->bitmap_.Contains(7)); } } + +TEST_P(GlobalIndexTest, TestWriteAndQueryLuminaIndexWithTagNullAndEmptyValues) { + arrow::FieldVector fields = {arrow::field("name", arrow::utf8()), + arrow::field("embedding", arrow::list(arrow::float32())), + arrow::field("color", arrow::utf8()), + arrow::field("labels", arrow::list(arrow::utf8())), + arrow::field("price", arrow::float32()), + arrow::field("scores", arrow::list(arrow::float32())), + arrow::field("category", arrow::int32()), + arrow::field("category_ids", arrow::list(arrow::int32()))}; + auto schema = arrow::schema(fields); + std::map lumina_options = { + {"lumina.index.dimension", "4"}, + {"lumina.index.type", "bruteforce"}, + {"lumina.distance.metric", "l2"}, + {"lumina.encoding.type", "rawf32"}, + {"lumina.search.parallel_number", "10"}, + {"lumina.extension.build.tag.tag_schema", + R"([{"key_name":"color","type":"enum","value_type":"string"},)" + R"({"key_name":"labels","type":"enum","value_type":"string"},)" + R"({"key_name":"price","type":"range","value_type":"double"},)" + R"({"key_name":"scores","type":"range","value_type":"double"},)" + R"({"key_name":"category","type":"enum","value_type":"int64"},)" + R"({"key_name":"category_ids","type":"enum","value_type":"int64"}])"}}; + std::map options = {{Options::MANIFEST_FORMAT, "orc"}, + {Options::FILE_FORMAT, file_format_}, + {Options::FILE_SYSTEM, "local"}, + {Options::ROW_TRACKING_ENABLED, "true"}, + {Options::DATA_EVOLUTION_ENABLED, "true"}}; + CreateTable(/*partition_keys=*/{}, schema, options); + + std::string table_path = PathUtil::JoinPath(dir_->Str(), "foo.db/bar"); + std::vector write_cols = schema->field_names(); + auto src_array = arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([ +["row0", [0.0, 0.0, 0.0, 0.0], "red", ["hot"], 5.0, [0.25], 7, [1, 2]], +["row1", null, "red", ["vip"], 8.0, [0.8], 7, [9]], +["row2", [0.0, 1.0, 0.0, 1.0], null, null, null, null, null, null], +["row3", [1.0, 0.0, 1.0, 0.0], " ", [], 20.0, [], 0, []], +["row4", [1.0, 1.0, 1.0, 1.0], "blue", ["vip", null], 30.0, [null, 0.75], 3, [null, 9]] + ])") + .ValueOrDie(); + + ASSERT_OK_AND_ASSIGN(auto commit_msgs, WriteArray(table_path, write_cols, src_array)); + ASSERT_OK(Commit(table_path, commit_msgs)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr split, + ScanData(table_path, /*partition_filters=*/{})); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr index_commit_msg, + GlobalIndexWriteTask::WriteIndex( + table_path, "embedding", "lumina", + std::make_shared(split, std::vector({Range(0, 4)})), + /*options=*/lumina_options, pool_, fs_)); + + std::shared_ptr index_commit_msg_impl = + std::dynamic_pointer_cast(index_commit_msg); + ASSERT_TRUE(index_commit_msg_impl); + const auto& new_index_files = index_commit_msg_impl->GetNewFilesIncrement().NewIndexFiles(); + ASSERT_EQ(new_index_files.size(), 1u); + ASSERT_EQ(new_index_files[0]->IndexType(), "lumina"); + ASSERT_EQ(new_index_files[0]->RowCount(), 5); + const std::optional& global_index_meta = + new_index_files[0]->GetGlobalIndexMeta(); + ASSERT_TRUE(global_index_meta); + std::string expected_index_meta_json = + R"({"distance.metric":"l2","encoding.type":"rawf32","extension.build.tag.tag_schema":"[{\"key_name\":\"color\",\"type\":\"enum\",\"value_type\":\"string\"},{\"key_name\":\"labels\",\"type\":\"enum\",\"value_type\":\"string\"},{\"key_name\":\"price\",\"type\":\"range\",\"value_type\":\"double\"},{\"key_name\":\"scores\",\"type\":\"range\",\"value_type\":\"double\"},{\"key_name\":\"category\",\"type\":\"enum\",\"value_type\":\"int64\"},{\"key_name\":\"category_ids\",\"type\":\"enum\",\"value_type\":\"int64\"}]","index.dimension":"4","index.type":"bruteforce","search.parallel_number":"10"})"; + GlobalIndexMeta expected_global_index_meta( + /*row_range_start=*/0, /*row_range_end=*/4, /*index_field_id=*/1, + /*extra_field_ids=*/std::optional>({2, 3, 4, 5, 6, 7}), + std::make_shared(expected_index_meta_json, pool_.get())); + ASSERT_EQ(global_index_meta.value(), expected_global_index_meta); + + ASSERT_OK(Commit(table_path, {index_commit_msg})); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr global_index_scan, + GlobalIndexScan::Create(table_path, /*snapshot_id=*/std::nullopt, + /*partitions=*/std::nullopt, lumina_options, fs_, + /*executor=*/nullptr, pool_)); + ASSERT_OK_AND_ASSIGN(auto lumina_readers, + global_index_scan->CreateReaders("embedding", std::nullopt)); + ASSERT_EQ(lumina_readers.size(), 1u); + + std::vector query = {1.0f, 1.0f, 1.0f, 1.1f}; + auto search_and_check_with_filter = [&](VectorSearch::PreFilter pre_filter, + const std::shared_ptr& predicate, + const std::string& expected) { + std::shared_ptr vector_search = std::make_shared( + "embedding", /*limit=*/5, query, pre_filter, predicate, + /*distance_type=*/std::nullopt, /*options=*/lumina_options); + ASSERT_OK_AND_ASSIGN(std::shared_ptr scored_result, + lumina_readers[0]->VisitVectorSearch(vector_search)); + ASSERT_EQ(scored_result->ToString(), expected); + }; + auto search_and_check = [&](const std::shared_ptr& predicate, + const std::string& expected) { + search_and_check_with_filter(/*pre_filter=*/nullptr, predicate, expected); + }; + + search_and_check(/*predicate=*/nullptr, "row ids: {0,2,3,4}, scores: {4.21,2.01,2.21,0.01}"); + search_and_check( + PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"color", FieldType::STRING, + Literal(FieldType::STRING, "red", 3)), + "row ids: {0}, scores: {4.21}"); + search_and_check(PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"color", + FieldType::STRING, Literal(FieldType::STRING, " ", 1)), + "row ids: {3}, scores: {2.21}"); + search_and_check( + PredicateBuilder::In(/*field_index=*/3, /*field_name=*/"labels", FieldType::STRING, + {Literal(FieldType::STRING, "vip", 3)}), + "row ids: {4}, scores: {0.01}"); + search_and_check(PredicateBuilder::LessOrEqual(/*field_index=*/4, /*field_name=*/"price", + FieldType::FLOAT, Literal(10.0f)), + "row ids: {0}, scores: {4.21}"); + search_and_check(PredicateBuilder::GreaterOrEqual(/*field_index=*/5, /*field_name=*/"scores", + FieldType::FLOAT, Literal(0.5f)), + "row ids: {4}, scores: {0.01}"); + search_and_check(PredicateBuilder::Equal(/*field_index=*/6, /*field_name=*/"category", + FieldType::INT, Literal(7)), + "row ids: {0}, scores: {4.21}"); + search_and_check(PredicateBuilder::In(/*field_index=*/7, /*field_name=*/"category_ids", + FieldType::INT, {Literal(9)}), + "row ids: {4}, scores: {0.01}"); + search_and_check( + PredicateBuilder::Equal(/*field_index=*/2, /*field_name=*/"color", FieldType::STRING, + Literal(FieldType::STRING, "green", 5)), + "row ids: {}, scores: {}"); + { + std::shared_ptr predicate = PredicateBuilder::Equal( + /*field_index=*/2, /*field_name=*/"color", FieldType::STRING, + Literal(FieldType::STRING, "red", 3)); + search_and_check_with_filter([](int64_t id) { return id == 0 || id == 4; }, predicate, + "row ids: {0}, scores: {4.21}"); + search_and_check_with_filter([](int64_t id) { return id == 4; }, predicate, + "row ids: {}, scores: {}"); + } + { + std::shared_ptr red_predicate = PredicateBuilder::Equal( + /*field_index=*/2, /*field_name=*/"color", FieldType::STRING, + Literal(FieldType::STRING, "red", 3)); + std::shared_ptr blue_predicate = PredicateBuilder::Equal( + /*field_index=*/2, /*field_name=*/"color", FieldType::STRING, + Literal(FieldType::STRING, "blue", 4)); + std::shared_ptr high_price_predicate = PredicateBuilder::GreaterOrEqual( + /*field_index=*/4, /*field_name=*/"price", FieldType::FLOAT, Literal(30.0f)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr blue_high_price_predicate, + PredicateBuilder::And({blue_predicate, high_price_predicate})); + ASSERT_OK_AND_ASSIGN(std::shared_ptr compound_predicate, + PredicateBuilder::Or({red_predicate, blue_high_price_predicate})); + search_and_check(compound_predicate, "row ids: {0,4}, scores: {4.21,0.01}"); + search_and_check_with_filter([](int64_t id) { return id == 4; }, compound_predicate, + "row ids: {4}, scores: {0.01}"); + } +} #endif TEST_P(GlobalIndexTest, TestDataEvolutionBatchScan) { diff --git a/third_party/lumina/VERSION b/third_party/lumina/VERSION index 8207be260..6a02f67b3 100644 --- a/third_party/lumina/VERSION +++ b/third_party/lumina/VERSION @@ -1,2 +1,2 @@ -tag: v0.3.0-rc1 -304227808425be9b2db0ded1e2e43be0ffba42b2 \ No newline at end of file +tag: v0.3.0 +c683b5f490827cbd25051ed485e5b29e0748c31c diff --git a/third_party/lumina/include/lumina/api/Dataset.h b/third_party/lumina/include/lumina/api/Dataset.h index d45f295b1..0da8e9b7d 100755 --- a/third_party/lumina/include/lumina/api/Dataset.h +++ b/third_party/lumina/include/lumina/api/Dataset.h @@ -14,7 +14,6 @@ * limitations under the License. */ - #pragma once #include diff --git a/third_party/lumina/include/lumina/core/Constants.h b/third_party/lumina/include/lumina/core/Constants.h index 562985df7..33c39c1b4 100755 --- a/third_party/lumina/include/lumina/core/Constants.h +++ b/third_party/lumina/include/lumina/core/Constants.h @@ -99,6 +99,9 @@ constexpr std::string_view kExtensionCkptCount = "extension.build.ckpt.count"; constexpr std::string_view kExtensionGetVector = "extension.search.get_vector"; constexpr std::string_view kExtensionTagSchema = "extension.build.tag.tag_schema"; constexpr std::string_view kExtensionTagMaxRangeLabelRatio = "extension.build.tag.max_range_label_ratio"; +constexpr std::string_view kExtensionTagMinRangeLabelRatio = "extension.build.tag.min_range_label_ratio"; +constexpr std::string_view kExtensionTagRangeSegmentTree = + "extension.build.tag.range_segment_tree"; // Only valid for flatNSW, single range tag // Available extension.build.tag.tag_schema keys constexpr std::string_view kExtensionTagKName = "key_name"; constexpr std::string_view kExtensionTagType = "type"; diff --git a/third_party/lumina/include/lumina/core/MemoryResource.h b/third_party/lumina/include/lumina/core/MemoryResource.h index 362c0d5c2..d2d1f3bf9 100755 --- a/third_party/lumina/include/lumina/core/MemoryResource.h +++ b/third_party/lumina/include/lumina/core/MemoryResource.h @@ -28,13 +28,18 @@ inline std::shared_ptr RefMemoryResource(std::pmr::me } struct MemoryResourceConfig { - // huge but thread unsafe + // Long-lived storage allocations. Components that write to storage must serialize access + // unless the supplied resource is thread-safe. std::shared_ptr storage; - // thread safe + // Temporary scratch allocations. Parallel build/search paths may allocate from this + // resource concurrently, so use a thread-safe resource when enabling those paths. std::shared_ptr instant; MemoryResourceConfig() = default; + // Uses the same resource for storage and instant allocations. The resource must be valid + // for every access pattern used by the selected component; parallel paths may access it + // concurrently through instant. explicit MemoryResourceConfig(std::pmr::memory_resource* resource) { storage = RefMemoryResource(resource); diff --git a/third_party/lumina/include/lumina/core/Status.h b/third_party/lumina/include/lumina/core/Status.h index 5fc241632..44c772a76 100755 --- a/third_party/lumina/include/lumina/core/Status.h +++ b/third_party/lumina/include/lumina/core/Status.h @@ -44,6 +44,7 @@ class [[nodiscard]] Status std::string _msg; }; +// TODO(feishi.wzj) add log #define LUMINA_RETURN_IF_ERROR(expr) \ do { \ auto _s = (expr); \ diff --git a/third_party/lumina/include/lumina/distance/EncodedDistance.h b/third_party/lumina/include/lumina/distance/EncodedDistance.h index b92c48d8f..409123387 100755 --- a/third_party/lumina/include/lumina/distance/EncodedDistance.h +++ b/third_party/lumina/include/lumina/distance/EncodedDistance.h @@ -18,13 +18,14 @@ #include #include +#include +#include +#include + #include #include #include #include -#include -#include -#include namespace lumina::dist { @@ -167,48 +168,6 @@ struct GatherEvalEncodedTag { }; inline constexpr GatherEvalEncodedTag GatherEvalEncoded {}; -struct GatherEvalEncodedWithLowerBoundsTag { - /** - * @brief Evaluates distances AND returns a lower-bound estimate for each row. - * - * This is an optimization for multi-stage search (e.g., DiskANN). - * - * Lower bounds (D_lb) guarantee that D_true >= D_lb. In search algorithms, if D_lb is - * already greater than the current search radius, we can skip further processing of this - * candidate. - * - * Note: Not all encodings support lower bounds. If unsupported, use the standard GatherEvalEncoded. - * - * @param results Output span for the (potentially approximate) distances. - * @param lowerBounds Output span for the lower-bound estimates. - */ - template - requires TagInvocable, std::span, std::span> - constexpr void operator()(const M& m, const E& e, const S& s, const std::byte* recordsBase, - std::span rowIds, std::span results, - std::span lowerBounds) const - noexcept(noexcept(TagInvoke(std::declval(), m, e, s, recordsBase, rowIds, - results, lowerBounds))) - { - TagInvoke(*this, m, e, s, recordsBase, rowIds, results, lowerBounds); - } - - template - requires(encode_space::EncodedRowSource> && - TagInvocable, std::span, std::span>) - constexpr void operator()(const M& m, const E& e, const S& s, const DataSource& data, - std::span rowIds, std::span results, - std::span lowerBounds) const - noexcept(noexcept(TagInvoke(std::declval(), m, e, s, data, rowIds, results, - lowerBounds))) - { - TagInvoke(*this, m, e, s, data, rowIds, results, lowerBounds); - } -}; -inline constexpr GatherEvalEncodedWithLowerBoundsTag GatherEvalEncodedWithLowerBounds {}; - struct SymmetricEvalEncodedTag { /** * @brief Computes distance between two encoded vectors (symmetric distance). diff --git a/third_party/lumina/include/lumina/extensions/experimental/CkptManager.h b/third_party/lumina/include/lumina/extensions/experimental/CkptManager.h index b48832c40..72347e793 100755 --- a/third_party/lumina/include/lumina/extensions/experimental/CkptManager.h +++ b/third_party/lumina/include/lumina/extensions/experimental/CkptManager.h @@ -16,6 +16,7 @@ #pragma once +// Generated by NewClass.py — edit as needed. #include #include @@ -30,6 +31,8 @@ class CkptManager : public core::NoCopyable CkptManager() = default; virtual ~CkptManager() = default; + // TODO: add ckpt header or not. + // TODO: add public API. virtual std::unique_ptr CreateCkptFileWriter() = 0; // If this function return false, there is no file reader diff --git a/third_party/lumina/include/lumina/extensions/experimental/DistributeBuildCombinedExtension.h b/third_party/lumina/include/lumina/extensions/experimental/DistributeBuildCombinedExtension.h index f692822f6..f1ad468e7 100644 --- a/third_party/lumina/include/lumina/extensions/experimental/DistributeBuildCombinedExtension.h +++ b/third_party/lumina/include/lumina/extensions/experimental/DistributeBuildCombinedExtension.h @@ -110,6 +110,8 @@ class DistributeBuildExtension : public detail::Di // ===================================================================== // Mixin subclasses +// TODO: add DistributePartitionWithTagExtension — needs tag-aware +// InsertBatch/Dump in the Partition Impl before wiring TagCapabilityMixin. // ===================================================================== class DistributeBuildWithCkptExtension final : public DistributeBuildExtension, public CkptCapabilityMixin @@ -132,5 +134,6 @@ using DistributeBuildExtensionWithCkpt = std::conditional_t, DistributeBuildWithCkptExtension, DistributeBuildExtension>; +// TODO: add DistributeBuildExtensionWithTag alias once DistributePartitionWithTagExtension is implemented }} // namespace lumina::extensions::experimental diff --git a/third_party/lumina/lib/liblumina.so b/third_party/lumina/lib/liblumina.so index db4374e82..58956fa88 100755 --- a/third_party/lumina/lib/liblumina.so +++ b/third_party/lumina/lib/liblumina.so @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:b1eb2cb4b6bb6f5c2ba76000a0ebcaf7df79d3cd36a6e258c431fea2b692a7d7 -size 111890648 +oid sha256:e9855f924e2439a147063c591253e3eaf3e35ab822e6eb8168be189a16eede4d +size 118212032 diff --git a/third_party/lumina/reference/DiskANNParameters.md b/third_party/lumina/reference/DiskANNParameters.md index b8041cb0e..a636d4c48 100644 --- a/third_party/lumina/reference/DiskANNParameters.md +++ b/third_party/lumina/reference/DiskANNParameters.md @@ -152,6 +152,7 @@ v0.3.0 Release Tag (2026-06-16). ## Related +- [Release Docs](../README.md) - [Options API](../api/Options.md) - [Options Reference](OptionsReference.md) - [Builder API](../api/Builder.md) diff --git a/third_party/lumina/reference/FilteringParameters.md b/third_party/lumina/reference/FilteringParameters.md new file mode 100644 index 000000000..19597fff6 --- /dev/null +++ b/third_party/lumina/reference/FilteringParameters.md @@ -0,0 +1,325 @@ +# Filtering Parameters & Tuning + +## Overview + +This page covers the configuration and tuning of two filtering approaches in Lumina: + +1. **Custom predicate filtering** (SearchWithFilterExtension): executes a per-query custom filter function on each candidate `vector_id_t` during search. +2. **Tag filtering** (BuildWithTagExtension + SearchWithTagExtension): tag-aware index construction that incorporates tag information into the index structure at build time; queries use a TagFilter expression tree for structured filtering. + +The two approaches target different scenarios: custom filtering suits dynamic business logic that cannot be predefined; tag filtering suits predefined structured attribute filtering, where incorporating tag information at build time makes the index aware of filtered subset distributions, yielding better recall and efficiency for filtered search. + +Content is based on the current release implementation and includes parameter semantics and tuning recommendations. This is a "current implementation" usage guide, not an additional long-term compatibility contract. + +## Audience + +- Developers integrating filtering logic into vector search. +- Architects choosing between custom filtering and tag filtering. +- Developers using filtering via the C++ or Python API. + +## Backend Support Summary + +| Backend | SearchWithFilter | BuildWithTag | SearchWithTag | +|---------|-----------------|--------------|---------------| +| `bruteforce` | Stable | Experimental | Experimental | +| `diskann` | Stable | Experimental | Experimental | +| `flatnsw` | Experimental | Experimental | Experimental | +| `ivf` | Not supported | Not supported | Not supported | + +--- + +## Custom Predicate Filtering (SearchWithFilterExtension) + +### Overview + +`SearchWithFilterExtension` allows the caller to provide a `bool(vector_id_t)` predicate function per query. The backend invokes this function on each candidate to decide whether to keep it. + +### Search Parameters + +These keys belong to `api::SearchOptions`. + +| Key | Type | Description | +| --- | --- | --- | +| `search.topk` | `FieldType::kInt` | Required. Number of results to return. | +| `search.parallel_number` | `FieldType::kInt` | Query parallelism. Valid range `1..1000`. | +| `search.thread_safe_filter` | `FieldType::kBool` | MUST be `true` when `parallel_number > 1`, declaring that the filter function is thread-safe. | + +### Backend-Specific Parameters + +Backend-specific search parameters remain effective when using SearchWithFilter: + +- **DiskANN**: `diskann.search.list_size` (required), `diskann.search.io_limit`, `diskann.search.sector_aligned_read`, etc. See [DiskANN Parameters & Tuning](DiskANNParameters.md). +- **Bruteforce**: no additional search parameters. +- **FlatNSW**: `flatnsw.search.ef_search`. + +### Filter Function Requirements + +- Type: `std::function`. +- Returns `true` to keep the candidate, `false` to discard. +- MUST be thread-safe and re-entrant (may be called concurrently during parallel search). +- SHOULD be fast and side-effect free; a single query may trigger thousands of calls. Avoid IO operations or locking inside the filter. +- MUST tolerate repeated calls for the same id. +- If the filter discards too many candidates, results MAY contain fewer than `topk` hits. + +### Tuning Recommendations + +#### Bruteforce Backend + +- Bruteforce scans all vectors and invokes the filter on each one, so per-call overhead directly affects total latency. +- When the filter is lightweight (e.g., hash table or bitmap lookup), prefer increasing `search.parallel_number` for higher throughput. +- Bruteforce does not lose recall due to filtering (full scan), but if fewer vectors pass the filter than `topk`, results will be fewer than `topk`. + +#### DiskANN Backend + +- For filtered search, increase both `diskann.search.list_size` and `diskann.search.io_limit`. Candidates discarded by the filter still consume IO budget and candidate pool slots, so both need to be larger to maintain recall. `io_limit` defaults to the same value as `list_size`, but in custom predicate filtering scenarios, set `io_limit` larger than `list_size` to leave headroom for IO consumed by discarded candidates. +- If filter pass rate is very low (e.g., < 5%), expect recall to drop. Compensate by further increasing `list_size` and `io_limit`. + +#### FlatNSW Backend + +- For filtered search, increase `flatnsw.search.ef_search`. The principle is similar to DiskANN: candidates discarded by the filter do not count toward top-k, so a larger ef provides more candidates. + +--- + +## Tag Filtering + +Tag filtering has two phases: the build phase uses `BuildWithTagExtension` to inject tag data, and the query phase uses `SearchWithTagExtension` for structured filtering. + +Precondition: only indexes built with `BuildWithTagExtension` can use `SearchWithTagExtension` for search. Attaching `SearchWithTagExtension` to an index built without tags causes `Attach` to return `FailedPrecondition` on all supported backends (bruteforce, diskann, flatnsw). + +### Build Phase + +#### Builder Parameters + +These keys belong to `api::BuilderOptions`. + +| Key | Type | Required | Default | Description | +| --- | --- | --- | --- | --- | +| `extension.build.tag.tag_schema` | JSON string | Yes | — | Declares tag dimensions (key_name, type, value_type). See format below. | +| `extension.build.tag.max_range_label_ratio` | `FieldType::kDouble` | No | 0.50 | Upper bound of log-uniform sampling for range label search window during construction. Controls the widest window, covering at most 50% of the full label value span. Range `[min_ratio, 1.00]`. Only effective for DiskANN and FlatNSW backends. | +| `extension.build.tag.min_range_label_ratio` | `FieldType::kDouble` | No | 0.02 | Lower bound of log-uniform sampling for range label search window during construction. Controls the narrowest window, covering at least 2% of the full label value span. Range `[0.02, 1.00]`. Only effective for DiskANN and FlatNSW backends. | +| `extension.build.tag.range_segment_tree` | `FieldType::kBool` | No | false | Enable segment-tree multi-layer construction. Only supported by `flatnsw` backend, requires exactly 1 range dimension and 0 enum dimensions. | + +#### tag_schema Format + +`extension.build.tag.tag_schema` is a JSON string declaring tag dimensions: + +```json +[ + {"key_name": "color", "type": "enum", "value_type": "string"}, + {"key_name": "category", "type": "enum", "value_type": "int64"}, + {"key_name": "price", "type": "range", "value_type": "double"}, + {"key_name": "year", "type": "range", "value_type": "int64"} +] +``` + +| Field | Values | Description | +|-------|--------|-------------| +| `key_name` | Any string (max 65535 bytes) | Tag dimension name. Must be unique. | +| `type` | `"enum"` / `"range"` | Tag type. `enum` for equality matching (e.g., color, category); `range` for interval queries (e.g., price, timestamp). | +| `value_type` | `"string"` / `"int64"` / `"double"` | Value type. Note: range type does not support string. Range dimension values are converted to an internal order-preserving uint32 representation: `double` values are first truncated to `float` precision; `int64` values are clamped to int32 range (`[-2^31, 2^31-1]`), with out-of-range values truncated to the boundary. | + +#### Build Input (TagDimensionData) + +When inserting data via `InsertBatchWithTag` or `InsertFromWithTag`, provide tag data for each vector. Each `TagDimensionData` element covers one dimension for all vectors: + +- `tagkName`: must match a `key_name` declared in tag_schema. +- `values`: type-safe variant (`vector>`, `vector>`, or `vector>`), outer indexed by vector position, inner holds multi-value tags for that vector on that dimension. + +Example: assuming tag_schema declares `color` (enum/string) and `price` (range/double), inserting 3 vectors: + +```cpp skip +std::vector tagData(2); + +// Dimension 1: color (enum, supports multi-value) +tagData[0].tagkName = "color"; +tagData[0].values = std::vector>{ + {"red", "blue"}, // vector 0: belongs to both red and blue + {"green"}, // vector 1: belongs to green + {"red"} // vector 2: belongs to red +}; + +// Dimension 2: price (range, single value) +tagData[1].tagkName = "price"; +tagData[1].values = std::vector>{ + {9.99}, // vector 0: price 9.99 + {19.99}, // vector 1: price 19.99 + {29.99} // vector 2: price 29.99 +}; +``` + +**Empty value handling:** + +Empty values fall into two scenarios: some vectors missing tags, and all vectors in a batch having no tags. + +- **Enum**: empty inner vector `{}` means the vector has no tag on that dimension and will not be matched by any enum filter condition. +- **Range**: empty inner vector or NaN means the vector has no valid range value on that dimension and will not be matched by range filter conditions, but participates normally in search without range filtering. +- **All-empty**: if all vectors in a batch have empty values for a dimension (or all dimensions), the build still succeeds, but any filter condition targeting that dimension will match no vectors, returning empty results. + +Scenario 1: some vectors missing tags + +```cpp skip +// 6 vectors, ID 1 and 4 have no color tag +std::vector tagData(1); +tagData[0].tagkName = "color"; +tagData[0].values = std::vector>{ + {"red"}, // ID 0: matched by Eq("color", "red") + {}, // ID 1: no tag, not matched by any color filter + {"blue"}, // ID 2 + {"red"}, // ID 3 + {}, // ID 4: no tag + {"blue"} // ID 5 +}; +// Query Eq("color", "red") returns only ID 0, 3; +// ID 1, 4 are not matched by any color filter. +``` + +Scenario 2: all vectors have no tags + +```cpp skip +// 200 vectors, all with empty color and size +constexpr uint32_t kVectorCount = 200; +std::vector tagData(2); + +tagData[0].tagkName = "color"; +tagData[0].values = std::vector>(kVectorCount); // all {} + +tagData[1].tagkName = "size"; +tagData[1].values = std::vector>(kVectorCount); // all {} + +// Build succeeds, but Eq("color", "red") or any tag filter returns empty results. +// Normal search without tag filtering still works. +``` + +### Query Phase + +#### Search Parameters + +Standard search parameters remain effective when using `SearchWithTagExtension`: + +| Key | Type | Description | +| --- | --- | --- | +| `search.topk` | `FieldType::kInt` | Required. | +| `search.parallel_number` | `FieldType::kInt` | Query parallelism (tag match table is read-only, inherently thread-safe). | + +Backend-specific parameters (`diskann.search.list_size`, etc.) also apply. + +#### Query Input (TagFilter) + +Queries describe filter conditions via a `TagFilter` expression tree, supporting the following operators: + +| Operator | Semantics | Enum dimensions | Range dimensions | +|----------|-----------|-----------------|------------------| +| `Eq` | `==` | Supported | Supported | +| `Ne` | `!=` | Not supported | Not supported | +| `Gt` | `>` | Not supported | Supported | +| `Gte` | `>=` | Not supported | Supported | +| `Lt` | `<` | Not supported | Supported | +| `Lte` | `<=` | Not supported | Supported | +| `In` | Set membership | Supported | Not supported | + +Logical combinations: `And` (all children must match), `Or` (at least one child must match). + +Example: + +```cpp skip +using namespace lumina::extensions::experimental; + +// Single condition: color == "blue" +auto f1 = TagFilter::Eq("color", std::string("blue")); + +// Multi-value match: color in {"red", "blue"} +auto f2 = TagFilter::In("color", std::vector{"red", "blue"}); + +// Range query: 10.0 <= price <= 50.0 +auto f3 = TagFilter::And( + TagFilter::Gte("price", 10.0), + TagFilter::Lte("price", 50.0) +); + +// Combined: color == "blue" AND price >= 10.0 +auto f4 = TagFilter::And( + TagFilter::Eq("color", std::string("blue")), + TagFilter::Gte("price", 10.0) +); + +// Or combination: color == "red" OR color == "blue" +auto f5 = TagFilter::Or( + TagFilter::Eq("color", std::string("red")), + TagFilter::Eq("color", std::string("blue")) +); +// Equivalent to In, but In is more efficient: +// auto f5 = TagFilter::In("color", std::vector{"red", "blue"}); +``` + +Limitations: +- Empty TagFilter is not allowed. If the query has no filtering requirement, use `LuminaSearcher::Search` directly. +- `And` conditions on the same `enum` dimension return `InvalidArgument`; use `In` or `Or` instead. +- OR expansion is capped at 256 compiled queries; exceeding this returns `InvalidArgument`. + +#### SearchWithTagAndFilter + +`SearchWithTagAndFilter` combines tag filtering with custom predicate filtering via AND logic. The user filter requirements are the same as `SearchWithFilterExtension::Filter`. + +### Tuning Recommendations + +#### Build Tuning + +- **Tag Schema Design**: + - `enum` suits equality matching scenarios (e.g., color, category); `range` suits interval query scenarios (e.g., price, timestamp). + - Minimize the number of tag dimensions; each additional dimension increases memory and build time. + +- **Range Label Ratio Parameters** (DiskANN and FlatNSW backends only): + - `max_range_label_ratio` and `min_range_label_ratio` control the range of label search window sizes during construction. + - These parameters have no effect on the Bruteforce backend. + - Tuning strategy depends on query selectivity distribution: + - If query selectivity is concentrated, narrow `[minRatio, maxRatio]` toward the selectivity. For example, if most queries filter ~30%, try `[0.25, 0.35]`; if selectivity is mostly 60%-90%, try `[0.60, 0.90]`. + - If query selectivity is dispersed (wide/narrow range queries mix), widen `[minRatio, maxRatio]`; the default `[0.02, 0.50]` is appropriate. + - If query selectivity distribution is unknown, use the defaults. + +- **Segment Tree Build Mode** (FlatNSW backend only): + - When there is exactly 1 range dimension and 0 enum dimensions, set `extension.build.tag.range_segment_tree=true` to enable segment-tree multi-layer construction, suited for range filtered query scenarios. + - Strict preconditions: exactly 1 range dimension, 0 enum dimensions, and `flatnsw` backend; otherwise returns `InvalidArgument`. + - When segment tree mode is enabled, `max_range_label_ratio` and `min_range_label_ratio` have no effect. + - Enabling segment tree improves single-range-attribute filtered query performance but increases build time, index disk space, and memory usage during both build and query. + +- **Graph Backend Build Parameters** (DiskANN and FlatNSW): + - Tag filtering requires tag-aware construction where the index maintains good connectivity on filtered subsets. Consider increasing neighbor count and construction candidate set size: + - DiskANN: `diskann.build.neighbor_count` and `diskann.build.ef_construction`. + - FlatNSW: `flatnsw.build.neighbor_count` and `flatnsw.build.ef_construction`. + +#### Query Tuning + +- **TagFilter Expressions**: + - Use `In` instead of multiple `Eq` combined with `Or` for better efficiency. + - Tag filtering compiles the TagFilter into a set of queries, executing each and merging with deduplication. Complex Or expressions lead to multiple sub-queries, increasing total latency. + +- **DiskANN Backend**: + - When tag filter selectivity is low, increase `diskann.search.list_size` and `diskann.search.io_limit`, similar to custom predicate filtering. + +- **FlatNSW Backend**: + - When tag filter selectivity is low, increase `flatnsw.search.ef_search`. + +--- + +## Notes + +- `SearchWithFilterExtension` is Stable; `BuildWithTagExtension` and `SearchWithTagExtension` are currently Experimental. +- IVF backend does not support any filtering extensions. +- The custom filter Python wrapper is experimental: `from lumina.experimental import SearchWithFilterExtension`. +- Tag extension Python wrappers are experimental: `from lumina.experimental import BuildWithTagExtension, SearchWithTagExtension`. +- When using `SearchWithTagAndFilter`, tag filtering and custom filtering are combined via AND logic in a single backend search; candidates must pass both to appear in results. +- The Checkpoint extension does not persist tag data; after recovering from a checkpoint, all data must be re-inserted. + +## Status + +v0.3.0 Release Tag (2026-07-03). + +## Related + +- [Release Docs](../README.md) +- [SearchWithFilter Extension](../extensions/SearchWithFilterExtension.md) +- [BuildWithTag Extension](../extensions/BuildWithTagExtension.md) +- [SearchWithTag Extension](../extensions/SearchWithTagExtension.md) +- [DiskANN Parameters & Tuning](DiskANNParameters.md) +- [Quantization Parameters & Tuning](QuantizationParameters.md) +- [Options API](../api/Options.md) diff --git a/third_party/lumina/reference/OptionsReference.md b/third_party/lumina/reference/OptionsReference.md index f38f60f8d..3a4089225 100644 --- a/third_party/lumina/reference/OptionsReference.md +++ b/third_party/lumina/reference/OptionsReference.md @@ -120,6 +120,8 @@ Note: `io.*` options apply to built-in file reader/writer implementations only. | Key | Type | Required | Deprecated | Validator | Description | | --- | ---- | -------- | ---------- | --------- | ----------- | | `extension.build.tag.max_range_label_ratio` | `FieldType::kDouble` | `false` | `false` | `ValidatePositiveDouble` | Max range label ratio for label-aware construction. | +| `extension.build.tag.min_range_label_ratio` | `FieldType::kDouble` | `false` | `false` | `ValidatePositiveDouble` | Min range label ratio for label-aware construction. | +| `extension.build.tag.range_segment_tree` | `FieldType::kBool` | `false` | `false` | `nullptr` | Build with segment tree or not, only valid for flatNSW and single range tag. | | `extension.build.tag.tag_schema` | `FieldType::kJson` | `false` | `false` | `ValidateTagSchema` | Tag schema for tag. | ## ivf / builder @@ -148,4 +150,4 @@ Note: `io.*` options apply to built-in file reader/writer implementations only. | `encoding.rabitq.centroid_count` | `FieldType::kInt` | `false` | `false` | `ValidatePositiveInt` | RabitQ kmeans centroid count for pretrain (default 64). Larger => better fit & slower training. | | `encoding.rabitq.max_epoch` | `FieldType::kInt` | `false` | `false` | `ValidatePositiveInt` | RabitQ kmeans max epochs for pretrain (default 10). Higher => better fit & slower training. | | `encoding.rabitq.quantized_bit_count` | `FieldType::kInt` | `false` | `false` | `ValidateIntInSet<1, 4, 5, 8, 9>` | RabitQ code bit width (default 4). Supported: 1, 4, 5, 8, 9. Larger => better accuracy & bigger records. | -| `encoding.rabitq.thread_count` | `FieldType::kInt` | `false` | `false` | `ValidatePositiveInt` | RabitQ kmeans thread count for pretrain (default 1). Controls training parallelism only. | \ No newline at end of file +| `encoding.rabitq.thread_count` | `FieldType::kInt` | `false` | `false` | `ValidatePositiveInt` | RabitQ thread count for training and encode (default 1). | diff --git a/third_party/lumina/reference/Overview.md b/third_party/lumina/reference/Overview.md index f23c475a3..3067d2246 100644 --- a/third_party/lumina/reference/Overview.md +++ b/third_party/lumina/reference/Overview.md @@ -136,4 +136,5 @@ Research behind Lumina has been published at top-tier database and systems venue - [Python quick start](../PythonQuickStart.md) — run the full build → dump → open → search → stream flow in Python. - [API docs](../api/README.md) — detailed Builder, Searcher, Streamer, and Options reference. - [DiskANN tuning guide](DiskANNParameters.md) — graph build and search parameter tuning for DiskANN. +- [Filtering tuning guide](FilteringParameters.md) — custom predicate and tag filtering parameter tuning. - [Options reference](OptionsReference.md) — complete list of configuration keys. diff --git a/third_party/lumina/reference/QuantizationParameters.md b/third_party/lumina/reference/QuantizationParameters.md index ae522126e..e65fef857 100644 --- a/third_party/lumina/reference/QuantizationParameters.md +++ b/third_party/lumina/reference/QuantizationParameters.md @@ -87,7 +87,7 @@ Current behavior: `rabitq` can usually reduce memory further, but they also come with stricter constraints and more tuning overhead. If you want a clear memory reduction without adding much complexity, `sq8` is usually the first option to try. - `Computation speed`: if distance evaluation on encoded vectors becomes the query-side bottleneck, consider `sq8`, - `pq`, or `rabitq`. + `pq`, or `rabitq`. For release-level benchmark context, see [Performance Report](../benchmark/PerformanceReport.md). - `Build time`: if you care more about build speed and configuration simplicity, prefer `rawf32` or `sq8`. `encoding.pq.thread_count` and `encoding.rabitq.thread_count` mainly improve build throughput. Increasing `encoding.pq.max_epoch`, `encoding.rabitq.max_epoch`, and `encoding.rabitq.centroid_count` also increases build @@ -96,8 +96,8 @@ Current behavior: - `Accuracy`: start with `rawf32` as the baseline. When compression is needed, `sq8` is usually the safer first step. For `pq`, start with `encoding.pq.m`; if accuracy is still not enough, then consider enabling `encoding.pq.use_opq`. `encoding.pq.make_zero_mean` is only worth trying for L2 with OPQ disabled. For `rabitq`, start from the default - `quantized_bit_count = 4`; if you need higher fidelity, raise `quantized_bit_count` first, then consider increasing - `centroid_count`. + `encoding.rabitq.quantized_bit_count = 4`; if you need higher fidelity, raise + `encoding.rabitq.quantized_bit_count` first, then consider increasing `encoding.rabitq.centroid_count`. ## Current Caveats @@ -115,6 +115,7 @@ v0.3.0 Release Tag (2026-06-16). ## Related +- [Release Docs](../README.md) - [Options API](../api/Options.md) - [Options Reference](OptionsReference.md) - [Builder API](../api/Builder.md) diff --git a/third_party/lumina/releases/v0.3.0.md b/third_party/lumina/releases/v0.3.0.md index 40b3be5ff..7d913f044 100644 --- a/third_party/lumina/releases/v0.3.0.md +++ b/third_party/lumina/releases/v0.3.0.md @@ -2,10 +2,25 @@ ## Overview -v0.3.0 brings three major additions: the FlatNSW in-memory graph backend (experimental), a tag filtering system that works across -backends, and `LuminaStreamer` for real-time vector ingestion. We also promoted `SearchHit` and `SearchResult` from `LuminaSearcher` nested types to standalone types -in `lumina::api` — existing code using `LuminaSearcher::SearchHit` still compiles via type aliases, but -new code should use the standalone types. See Migration Notes. +Until now, Lumina could tell you which vectors were closest — but not whether they were the ones you actually +wanted, and not without rebuilding an index every time your data changed. v0.3.0 is about closing both gaps: +making search aware of the attributes attached to your vectors, and letting an index stay useful as new vectors arrive. + +Three experimental additions lead the release. **FlatNSW** is a new in-memory graph backend for datasets in the +million- to ten-million-vector range — for when your working set fits in RAM and you'd rather not stand up and +babysit a disk-backed index. **Tag filtering** is the standout: rather than filtering results after the fact, Lumina builds the graph itself +*tag-aware*, so on `flatnsw` and `diskann` an attribute-constrained search is designed to keep ANN-level speed and +recall even under selective filters — exactly where naive post-filtering loses recall and brute-force pre-filtering +loses speed. Enum, range, multi-dimensional range, and mixed constraints are supported (`bruteforce` gets exact tag +filtering too). **`LuminaStreamer`** covers the insert-and-search-right-now case — for now on the `bruteforce` +backend — where data doesn't sit still long enough for an offline build. + +RabitQ isn't new — it shipped in v0.2.0 — but this release improves its performance: cutoff pruning on the search +side, and multi-threaded batch encoding. Same feature, less waiting. + +One bit of housekeeping: `SearchHit` and `SearchResult` have moved out from under `LuminaSearcher` and are now +standalone types in `lumina::api`. Existing code keeps compiling through type aliases, so nothing breaks today, but +new code should prefer the standalone types — see Migration Notes. ## Audience @@ -15,11 +30,11 @@ new code should use the standalone types. See Migration Notes. ## Status -Stable (2026-06-16). +Stable. ## Compatibility -- **Public API**: stable within `include/lumina/api/**` following semantic versioning. +- **Public API**: stable within `include/lumina/api/**` following semantic versioning rules. - **Source compatibility**: this release is source-compatible with v0.2.x for most code. `SearchHit` / `SearchResult` were promoted to standalone types, but backward-compatible type aliases are provided (see Migration Notes). - **ABI compatibility**: not promised. Recompile all dependent code after upgrading. @@ -28,24 +43,38 @@ Stable (2026-06-16). - Backend-specific layouts are excluded from long-term compatibility promises unless explicitly declared stable. - IVF snapshot layout is experimental: IVFIndex. - **Extensions**: - - `SearchWithFilterExtension`: stable, supported by `diskann` and `bruteforce` searchers (not `ivf`). + - `SearchWithFilterExtension`: stable, supported by `diskann`, `bruteforce`, and `flatnsw` searchers (not `ivf`). - `BuildWithTagExtension` / `SearchWithTagExtension`: new in this release, experimental. - `DistributeBuildExtension`: new in this release, experimental. - - Checkpoint extension: experimental, backend-specific. + - Checkpoint extension: experimental, backend-specific; no long-term compatibility promise. - `GetVectorExtension`: experimental, `bruteforce` with `rawf32` only. - **Python**: experimental interface; the API may change across versions and is not covered by stability promises. ## Changes - Features: - - FlatNSW backend (experimental): new in-memory graph index based on Navigable Small World. Supports build, search, - and tag-based filtering. Suitable for million- to ten-million-scale datasets that fit in memory. API and format may change. - - Tag filtering system: tag-aware build and search across `flatnsw`, `diskann`, and `bruteforce` backends. Supports - enum tags, range tags, multi-dimensional range, and mixed mode. Tag schema uses JSON-typed options. + - FlatNSW backend (experimental): a graph index that lives entirely in memory — for when your dataset (a million to + ten million vectors) fits in RAM and you'd rather not stand up and babysit a disk-backed index. Supports + build/search, search-time predicate filtering (`SearchWithFilterExtension`), and tag filtering. API and on-disk + format may still change. + - Tag filtering (tag-aware build + search): the advantage isn't filtering results, it's that + the graph is *constructed* with your tags in mind (label-aware neighbor selection and pruning, plus a + range-optimized graph for range tags). So attribute-constrained queries are designed to stay fast and high-recall + even when the filter is selective — where naive post-filtering loses recall and brute-force pre-filtering loses speed. Handles + "nearest, but only the ones in stock" or "closest, priced 10–50" in a single pass. Supports enum tags, range tags, + multi-dimensional range, and mixed mode across the `flatnsw`, `diskann`, and `bruteforce` backends; tag schema uses + JSON-typed options. See [BuildWithTagExtension](../extensions/BuildWithTagExtension.md) and [SearchWithTagExtension](../extensions/SearchWithTagExtension.md). - - LuminaStreamer: real-time, in-memory vector indexing with single-writer `Insert` and concurrent `Search`. - Currently backed by bruteforce. See [Streamer API](../api/Streamer.md). + - FlatNSW single-range-tag acceleration (experimental): optional build mode enabled with + `extension.build.tag.range_segment_tree` (flatnsw only, one range tag) that speeds up range-filtered search. + - Tag build tuning: new `extension.build.tag.min_range_label_ratio` and `extension.build.tag.max_range_label_ratio` + options to control the label-aware graph construction search window. + - FlatNSW search stats now expose `filteredCount` and a distance-pruned count for observability. + - LuminaStreamer: insert vectors and search them right away, no rebuild in between — for when data keeps arriving. + A single writer inserts while readers search concurrently. So far, only the `bruteforce` backend supports + Streamer (raw float32, fixed capacity), so results are exact for whatever is in it right now. + See [Streamer API](../api/Streamer.md). - DistributeBuildExtension: extension framework for distributed index building. See [DistributeBuildExtension](../extensions/DistributeBuildExtension.md). - DiskANN: added `graph_node_per_sector` build option to control disk layout density. @@ -53,15 +82,14 @@ Stable (2026-06-16). - Performance: - SQ8: added `GatherBatch4` SIMD dispatch for batch distance computation. - Util: `SimpleByteMap` now uses a generation counter to avoid per-query `memset`. + - RabitQ: cutoff pruning in graph/filtered search — uses the 1-bit lower bound to skip full multi-bit distance + evaluation for candidates that cannot beat the current search radius, reducing distance computations. + - RabitQ: multi-threaded batch encoding via `encoding.rabitq.thread_count`. - API changes: - `SearchHit` and `SearchResult` are now standalone types in `lumina::api` (previously nested in `LuminaSearcher`). A new header `SearchResult.h` is provided. `LuminaSearcher::SearchHit` and `LuminaSearcher::SearchResult` remain as type aliases. - - `ISearchExtension` and `IBuildExtension` now also inherit `NoMoveable`. - - Added `IStreamExtension` interface and `StreamerOptions` type. - - `LuminaBuilder` and `LuminaSearcher` gained move-assignment operators. - - `NormalizeBuilderOptions`, `NormalizeSearcherOptions`, `NormalizeSearchOptions`, and `ValidateSearchOptions` - are now `noexcept`. + - Added `IStreamExtension` interface and `StreamerOptions` type (for `LuminaStreamer`). ## Migration Notes @@ -69,9 +97,6 @@ Stable (2026-06-16). `LuminaSearcher::SearchResult` directly, it still compiles (they are now type aliases). But if you forward-declared them or used them in template specializations, switch to `lumina::api::SearchHit` / `lumina::api::SearchResult` and include ``. -- **Extension classes are now non-moveable**: `ISearchExtension` and `IBuildExtension` subclasses can no longer be - moved. This is unlikely to affect most code, but if you stored extensions in move-requiring containers, switch to - pointers. - **Recompile required**: ABI is not preserved; rebuild all dependent binaries. ## Known Issues @@ -92,4 +117,5 @@ Stable (2026-06-16). - [CheckpointExtension](../extensions/CheckpointExtension.md) - [GetVectorExtension](../extensions/GetVectorExtension.md) - [DiskANN parameters](../reference/DiskANNParameters.md) +- [Filtering parameters](../reference/FilteringParameters.md) - [Quantization parameters](../reference/QuantizationParameters.md)