feat: impl StructArray -- support creating index sort stl (#47053)

issue: https://github.com/milvus-io/milvus/issues/42148

1. This PR enable nested index to support stl sort index.
2. In addition, this PR makes element-level filter support full mode as
well as index.

---------

Signed-off-by: SpadeA <tangchenjie1210@gmail.com>
This commit is contained in:
Spade A
2026-01-29 12:07:33 +08:00
committed by GitHub
parent 56b7b08bc7
commit 1737f94c48
12 changed files with 752 additions and 76 deletions
+14
View File
@@ -118,6 +118,11 @@ class Expr {
return false;
}
virtual bool
CanUseNestedIndex() const {
return false;
}
virtual std::optional<milvus::expr::ColumnInfo>
GetColumnInfo() const {
ThrowInfo(ErrorCode::NotImplemented, "not implemented");
@@ -1954,6 +1959,15 @@ class SegmentExpr : public Expr {
return true;
}
bool
CanUseNestedIndex() const override {
if (!CanUseIndex() || pinned_index_.empty()) {
return false;
}
auto* index_ptr = pinned_index_[0].get();
return index_ptr != nullptr && index_ptr->IsNestedIndex();
}
template <typename T>
bool
CanUseIndexForOp(OpType op) const {
@@ -23,6 +23,9 @@
namespace milvus {
namespace exec {
const double DOC_HIT_RATIO_THRESHOLD_HIGH = 0.01;
const double DOC_HIT_RATIO_THRESHOLD_LOW = 0.002;
PhyElementFilterBitsNode::PhyElementFilterBitsNode(
int32_t operator_id,
DriverContext* driverctx,
@@ -86,15 +89,12 @@ PhyElementFilterBitsNode::GetOutput() {
col_input->size());
doc_bitset.flip();
// Step 3: Convert doc bitset to element offsets
FixedVector<int32_t> element_offsets =
array_offsets->RowBitsetToElementOffsets(doc_bitset, 0);
// Step 3: Evaluate element expression
// Use offset mode or full mode based on selectivity
auto [expr_result, valid_expr_result] = EvaluateElementExpression(
doc_bitset, doc_bitset_valid, array_offsets.get());
// Step 4: Evaluate element expression
auto [expr_result, valid_expr_result] =
EvaluateElementExpression(element_offsets);
// Step 5: Set query context
// Step 4: Set query context
query_context_->set_element_level_query(true);
query_context_->set_struct_name(struct_name_);
@@ -124,57 +124,150 @@ PhyElementFilterBitsNode::GetOutput() {
std::pair<TargetBitmap, TargetBitmap>
PhyElementFilterBitsNode::EvaluateElementExpression(
FixedVector<int32_t>& element_offsets) {
const TargetBitmapView& doc_bitset,
const TargetBitmapView& doc_bitset_valid,
const IArrayOffsets* array_offsets) {
tracer::AutoSpan span("PhyElementFilterBitsNode::EvaluateElementExpression",
tracer::GetRootSpan(),
true);
tracer::AddEvent(fmt::format("input_elements: {}", element_offsets.size()));
// Use offset interface by passing element_offsets
EvalCtx eval_ctx(operator_context_->get_exec_context(), &element_offsets);
std::vector<VectorPtr> results;
element_exprs_->Eval(0, 1, true, eval_ctx, results);
AssertInfo(results.size() == 1 && results[0] != nullptr,
"ElementFilterBitsNode: expression evaluation should return "
"exactly one result");
TargetBitmap bitset;
TargetBitmap valid_bitset;
int64_t total_elements = query_context_->get_active_element_count();
bitset = TargetBitmap(total_elements, false);
valid_bitset = TargetBitmap(total_elements, true);
auto col_vec = std::dynamic_pointer_cast<ColumnVector>(results[0]);
if (!col_vec) {
ThrowInfo(ExprInvalid,
"ElementFilterBitsNode result should be ColumnVector");
auto count = doc_bitset.count();
auto doc_hit_ratio = count * 1.0 / doc_bitset.size();
// If hit_ratio is less than DOC_HIT_RATIO_THRESHOLD_LOW, we always use offset mode
// If hit_ratio is between [DOC_HIT_RATIO_THRESHOLD_LOW, DOC_HIT_RATIO_THRESHOLD_HIGH], we check if all expressions can use nested index
// Otherwise, we use full mode
bool offset_mode;
if (doc_hit_ratio < DOC_HIT_RATIO_THRESHOLD_LOW) {
offset_mode = true;
} else if (doc_hit_ratio < DOC_HIT_RATIO_THRESHOLD_HIGH) {
offset_mode = !std::all_of(
element_exprs_->exprs().begin(),
element_exprs_->exprs().end(),
[](const auto& expr) { return expr->CanUseNestedIndex(); });
} else {
offset_mode = false;
}
if (!col_vec->IsBitmap()) {
ThrowInfo(ExprInvalid, "ElementFilterBitsNode result should be bitmap");
}
auto col_vec_size = col_vec->size();
TargetBitmapView bitsetview(col_vec->GetRawData(), col_vec_size);
AssertInfo(col_vec_size == element_offsets.size(),
"ElementFilterBitsNode result size mismatch: {} vs {}",
col_vec_size,
element_offsets.size());
if (offset_mode) {
// Offset mode: convert doc_bitset to element offsets and evaluate only on those
// Offset mode processes all offsets in one pass
FixedVector<int32_t> element_offsets =
array_offsets->RowBitsetToElementOffsets(doc_bitset, 0);
for (size_t i = 0; i < element_offsets.size(); ++i) {
if (bitsetview[i]) {
bitset[element_offsets[i]] = true;
tracer::AddEvent(fmt::format("offset_mode, input_elements: {}",
element_offsets.size()));
EvalCtx eval_ctx(operator_context_->get_exec_context(),
&element_offsets);
std::vector<VectorPtr> results;
element_exprs_->Eval(0, 1, true, eval_ctx, results);
AssertInfo(results.size() == 1 && results[0] != nullptr,
"ElementFilterBitsNode: expression evaluation should return "
"exactly one result");
auto col_vec = std::dynamic_pointer_cast<ColumnVector>(results[0]);
if (!col_vec) {
ThrowInfo(ExprInvalid,
"ElementFilterBitsNode result should be ColumnVector");
}
if (!col_vec->IsBitmap()) {
ThrowInfo(ExprInvalid,
"ElementFilterBitsNode result should be bitmap");
}
auto col_vec_size = col_vec->size();
TargetBitmapView bitsetview(col_vec->GetRawData(), col_vec_size);
AssertInfo(col_vec_size == element_offsets.size(),
"ElementFilterBitsNode result size mismatch: {} vs {}",
col_vec_size,
element_offsets.size());
// Scatter results to full bitset
TargetBitmap bitset(total_elements, false);
TargetBitmap valid_bitset(total_elements, true);
for (size_t i = 0; i < element_offsets.size(); ++i) {
if (bitsetview[i]) {
bitset[element_offsets[i]] = true;
}
}
bitset.flip();
tracer::AddEvent(
fmt::format("evaluated_elements: {}, total_elements: {}",
element_offsets.size(),
total_elements));
return std::make_pair(std::move(bitset), std::move(valid_bitset));
} else {
// Full mode: evaluate on all elements, then AND with doc_bitset
tracer::AddEvent(
fmt::format("full_mode, total_elements: {}", total_elements));
EvalCtx eval_ctx(operator_context_->get_exec_context());
TargetBitmap eval_bitset;
int64_t num_processed_elements = 0;
std::vector<VectorPtr> results;
while (num_processed_elements < total_elements) {
element_exprs_->Eval(0, 1, true, eval_ctx, results);
AssertInfo(results.size() == 1 && results[0] != nullptr,
"ElementFilterBitsNode: expression evaluation "
"should return "
"exactly one result");
auto col_vec = std::dynamic_pointer_cast<ColumnVector>(results[0]);
if (!col_vec) {
ThrowInfo(
ExprInvalid,
"ElementFilterBitsNode result should be ColumnVector");
}
if (!col_vec->IsBitmap()) {
ThrowInfo(ExprInvalid,
"ElementFilterBitsNode result should be bitmap");
}
auto col_vec_size = col_vec->size();
TargetBitmapView view(col_vec->GetRawData(), col_vec_size);
eval_bitset.append(view);
num_processed_elements += col_vec_size;
}
AssertInfo(eval_bitset.size() == total_elements,
"ElementFilterBitsNode result size mismatch: {} vs {}",
eval_bitset.size(),
total_elements);
// Convert doc_bitset to element bitset
auto [elem_bitset, _] = array_offsets->RowBitsetToElementBitset(
doc_bitset, doc_bitset_valid, 0);
// AND expression result with element bitset from doc filter
// eval_bitset: true means element matches expression
// elem_bitset: true means element's doc passed predicate filter
eval_bitset &= elem_bitset;
eval_bitset.flip();
tracer::AddEvent(
fmt::format("evaluated_elements: {}, total_elements: {}",
total_elements,
total_elements));
// Element filter targets individual elements within arrays.
// While the array field itself can be nullable (handled by doc_bitset_valid),
// individual elements inside an array do not support null values.
// Therefore, the element-level valid_bitset is always all true.
TargetBitmap valid_bitset(total_elements, true);
return std::make_pair(std::move(eval_bitset), std::move(valid_bitset));
}
bitset.flip();
tracer::AddEvent(fmt::format("evaluated_elements: {}, total_elements: {}",
element_offsets.size(),
total_elements));
return std::make_pair(std::move(bitset), std::move(valid_bitset));
}
} // namespace exec
@@ -78,7 +78,9 @@ class PhyElementFilterBitsNode : public Operator {
private:
std::pair<TargetBitmap, TargetBitmap>
EvaluateElementExpression(FixedVector<int32_t>& element_offsets);
EvaluateElementExpression(const TargetBitmapView& doc_bitset,
const TargetBitmapView& doc_bitset_valid,
const IArrayOffsets* array_offsets);
std::unique_ptr<ExprSet> element_exprs_;
QueryContext* query_context_;
+6
View File
@@ -70,6 +70,12 @@ class IndexBase {
virtual bool
IsMmapSupported() const = 0;
// Check if this index is a nested index (for array element-level indexing)
virtual bool
IsNestedIndex() const {
return false;
}
const IndexType&
Type() const {
return index_type_;
+48 -2
View File
@@ -538,8 +538,18 @@ IndexFactory::CreateNestedIndex(
IndexType index_type,
int32_t tantivy_index_version,
const storage::FileManagerContext& file_manager_context) {
AssertInfo(index_type == INVERTED_INDEX_TYPE,
"Nested index only supports inverted index for now");
if (index_type == INVERTED_INDEX_TYPE) {
return CreateNestedIndexInverted(tantivy_index_version,
file_manager_context);
}
return CreateNestedIndexScalarIndexSort(file_manager_context);
}
IndexBasePtr
IndexFactory::CreateNestedIndexInverted(
int32_t tantivy_index_version,
const storage::FileManagerContext& file_manager_context) {
DataType element_type = static_cast<DataType>(
file_manager_context.fieldDataMeta.field_schema.element_type());
switch (element_type) {
@@ -577,6 +587,42 @@ IndexFactory::CreateNestedIndex(
}
}
IndexBasePtr
IndexFactory::CreateNestedIndexScalarIndexSort(
const storage::FileManagerContext& file_manager_context) {
DataType element_type = static_cast<DataType>(
file_manager_context.fieldDataMeta.field_schema.element_type());
switch (element_type) {
case DataType::BOOL:
return std::make_unique<ScalarIndexSort<bool>>(file_manager_context,
true);
case DataType::INT8:
return std::make_unique<ScalarIndexSort<int8_t>>(
file_manager_context, true);
case DataType::INT16:
return std::make_unique<ScalarIndexSort<int16_t>>(
file_manager_context, true);
case DataType::INT32:
return std::make_unique<ScalarIndexSort<int32_t>>(
file_manager_context, true);
case DataType::INT64:
return std::make_unique<ScalarIndexSort<int64_t>>(
file_manager_context, true);
case DataType::FLOAT:
return std::make_unique<ScalarIndexSort<float>>(
file_manager_context, true);
case DataType::DOUBLE:
return std::make_unique<ScalarIndexSort<double>>(
file_manager_context, true);
case DataType::STRING:
case DataType::VARCHAR:
return std::make_unique<StringIndexSort>(file_manager_context,
true);
default:
ThrowInfo(DataTypeInvalid, "Invalid data type:{}", element_type);
}
}
IndexBasePtr
IndexFactory::CreateScalarIndex(
const CreateIndexInfo& create_index_info,
+11
View File
@@ -138,6 +138,17 @@ class IndexFactory {
const storage::FileManagerContext& file_manager_context =
storage::FileManagerContext());
IndexBasePtr
CreateNestedIndexInverted(
int32_t tantivy_index_version,
const storage::FileManagerContext& file_manager_context =
storage::FileManagerContext());
IndexBasePtr
CreateNestedIndexScalarIndexSort(
const storage::FileManagerContext& file_manager_context =
storage::FileManagerContext());
IndexBasePtr
CreateScalarIndex(const CreateIndexInfo& create_index_info,
const storage::FileManagerContext& file_manager_context =
+2 -2
View File
@@ -151,8 +151,8 @@ class ScalarIndex : public IndexBase {
index_type_ == milvus::index::ASCENDING_SORT;
}
virtual bool
IsNestedIndex() const {
bool
IsNestedIndex() const override {
return false;
}
+85 -9
View File
@@ -33,6 +33,7 @@
#include "pb/common.pb.h"
#include "storage/ThreadPools.h"
#include "storage/Util.h"
#include "common/Array.h"
namespace milvus::index {
@@ -46,8 +47,12 @@ const uint64_t MMAP_INDEX_PADDING = 1;
template <typename T>
ScalarIndexSort<T>::ScalarIndexSort(
const storage::FileManagerContext& file_manager_context)
: ScalarIndex<T>(ASCENDING_SORT), is_built_(false), data_() {
const storage::FileManagerContext& file_manager_context,
bool is_nested_index)
: ScalarIndex<T>(ASCENDING_SORT),
is_nested_index_(is_nested_index),
is_built_(false),
data_() {
// not valid means we are in unit test
if (file_manager_context.Valid()) {
field_id_ = file_manager_context.fieldDataMeta.field_id;
@@ -109,6 +114,11 @@ ScalarIndexSort<T>::BuildWithFieldData(
const std::vector<milvus::FieldDataPtr>& field_datas) {
index_build_begin_ = std::chrono::system_clock::now();
if (is_nested_index_) {
BuildWithArrayDataNested(field_datas);
return;
}
int64_t length = 0;
for (const auto& data : field_datas) {
total_num_rows_ += data->get_num_rows();
@@ -147,6 +157,54 @@ ScalarIndexSort<T>::BuildWithFieldData(
ComputeByteSize();
}
template <typename T>
void
ScalarIndexSort<T>::BuildWithArrayDataNested(
const std::vector<FieldDataPtr>& datas) {
// calculate total_num_rows_
for (const auto& data : datas) {
auto n = data->get_num_rows();
auto array_column = static_cast<const Array*>(data->Data());
for (int64_t i = 0; i < n; i++) {
if (data->is_valid(i)) {
total_num_rows_ += array_column[i].length();
}
}
}
if (total_num_rows_ == 0) {
ThrowInfo(DataIsEmpty, "ScalarIndexSort cannot build null values!");
}
data_.reserve(total_num_rows_);
// all values are valid for nested index because any given slot in a valid_bitset_ denotes one element in a valid row
valid_bitset_ = TargetBitmap(total_num_rows_, true);
int64_t offset = 0;
for (const auto& data : datas) {
auto n = data->get_num_rows();
auto array_column = static_cast<const Array*>(data->Data());
for (int64_t i = 0; i < n; i++) {
if (!data->is_valid(i)) {
continue;
}
auto length = array_column[i].length();
for (int64_t j = 0; j < length; j++) {
data_.emplace_back(IndexStructure(
array_column[i].template get_data<T>(j), offset));
offset++;
}
}
}
std::sort(data_.begin(), data_.end());
idx_to_offsets_.resize(total_num_rows_);
for (size_t i = 0; i < total_num_rows_; ++i) {
idx_to_offsets_[data_[i].idx_] = i;
}
is_built_ = true;
setup_data_pointers();
}
template <typename T>
BinarySet
ScalarIndexSort<T>::Serialize(const Config& config) {
@@ -163,10 +221,14 @@ ScalarIndexSort<T>::Serialize(const Config& config) {
std::shared_ptr<uint8_t[]> index_num_rows(new uint8_t[sizeof(size_t)]);
memcpy(index_num_rows.get(), &total_num_rows_, sizeof(size_t));
std::shared_ptr<uint8_t[]> is_nested_data(new uint8_t[sizeof(bool)]);
memcpy(is_nested_data.get(), &is_nested_index_, sizeof(bool));
BinarySet res_set;
res_set.Append("index_data", index_data, index_data_size);
res_set.Append("index_length", index_length, sizeof(size_t));
res_set.Append("index_num_rows", index_num_rows, sizeof(size_t));
res_set.Append("is_nested_index", is_nested_data, sizeof(bool));
milvus::Disassemble(res_set);
@@ -181,8 +243,11 @@ ScalarIndexSort<T>::Upload(const Config& config) {
std::chrono::system_clock::now() - index_build_begin_)
.count();
LOG_INFO(
"index build done for ScalarIndexSort, field_id: {}, duration: {}ms",
"index build done for ScalarIndexSort, field_id: {}, is_nested_index: "
"{}, duration: "
"{}ms",
field_id_,
is_nested_index_,
index_build_duration);
auto binary_set = Serialize(config);
@@ -201,6 +266,13 @@ ScalarIndexSort<T>::LoadWithoutAssemble(const BinarySet& index_binary,
auto index_length = index_binary.GetByName("index_length");
memcpy(&index_size, index_length->data.get(), (size_t)index_length->size);
auto is_nested_index = index_binary.GetByName("is_nested_index");
if (is_nested_index) {
memcpy(&is_nested_index_,
is_nested_index->data.get(),
(size_t)is_nested_index->size);
}
is_mmap_ = GetValueFromConfig<bool>(config, ENABLE_MMAP).value_or(true);
auto index_data = index_binary.GetByName("index_data");
@@ -330,9 +402,11 @@ ScalarIndexSort<T>::In(const size_t n, const T* values) {
std::upper_bound(begin(), end(), IndexStructure<T>(*(values + i)));
for (; lb < ub; ++lb) {
if (lb->a_ != *(values + i)) {
std::cout << "error happens in ScalarIndexSort<T>::In, "
"experted value is: "
<< *(values + i) << ", but real value is: " << lb->a_;
LOG_ERROR(
"error happens in ScalarIndexSort<T>::In, "
"expected value is: {}, but real value is: {}",
*(values + i),
lb->a_);
}
bitset[lb->idx_] = true;
}
@@ -352,9 +426,11 @@ ScalarIndexSort<T>::NotIn(const size_t n, const T* values) {
std::upper_bound(begin(), end(), IndexStructure<T>(*(values + i)));
for (; lb < ub; ++lb) {
if (lb->a_ != *(values + i)) {
std::cout << "error happens in ScalarIndexSort<T>::NotIn, "
"experted value is: "
<< *(values + i) << ", but real value is: " << lb->a_;
LOG_ERROR(
"error happens in ScalarIndexSort<T>::NotIn, "
"expected value is: {}, but real value is: {}",
*(values + i),
lb->a_);
}
bitset[lb->idx_] = false;
}
+17 -3
View File
@@ -23,6 +23,8 @@
#include <string>
#include <map>
#include <boost/container/vector.hpp>
#include "index/IndexStructure.h"
#include "index/ScalarIndex.h"
#include "storage/MemFileManagerImpl.h"
@@ -48,7 +50,8 @@ class ScalarIndexSort : public ScalarIndex<T> {
public:
explicit ScalarIndexSort(
const storage::FileManagerContext& file_manager_context =
storage::FileManagerContext());
storage::FileManagerContext(),
bool is_nested_index = false);
~ScalarIndexSort() {
if (is_mmap_ && mmap_data_ != nullptr && mmap_data_ != MAP_FAILED) {
@@ -76,6 +79,11 @@ class ScalarIndexSort : public ScalarIndex<T> {
return ScalarIndexType::STLSORT;
}
bool
IsNestedIndex() const override {
return is_nested_index_;
}
void
Build(size_t n, const T* values, const bool* valid_data = nullptr) override;
@@ -150,6 +158,9 @@ class ScalarIndexSort : public ScalarIndex<T> {
BuildWithFieldData(const std::vector<FieldDataPtr>& datas) override;
private:
void
BuildWithArrayDataNested(const std::vector<FieldDataPtr>& datas);
bool
ShouldSkip(const T lower_value, const T upper_value, const OpType op);
@@ -210,6 +221,7 @@ class ScalarIndexSort : public ScalarIndex<T> {
int64_t field_id_ = 0;
bool is_nested_index_ = false;
bool is_built_ = false;
Config config_;
std::vector<int32_t> idx_to_offsets_; // used to retrieve.
@@ -247,7 +259,9 @@ namespace milvus::index {
template <typename T>
inline ScalarIndexSortPtr<T>
CreateScalarIndexSort(const storage::FileManagerContext& file_manager_context =
storage::FileManagerContext()) {
return std::make_unique<ScalarIndexSort<T>>(file_manager_context);
storage::FileManagerContext(),
bool is_nested_index = false) {
return std::make_unique<ScalarIndexSort<T>>(file_manager_context,
is_nested_index);
}
} // namespace milvus::index
+69 -7
View File
@@ -34,6 +34,8 @@
#include "index/Utils.h"
#include "storage/ThreadPools.h"
#include "storage/Util.h"
#include "common/Array.h"
#include "pb/schema.pb.h"
namespace milvus::index {
@@ -99,8 +101,11 @@ constexpr size_t ALIGNMENT = 32; // 32-byte alignment
const uint64_t MMAP_INDEX_PADDING = 1;
StringIndexSort::StringIndexSort(
const storage::FileManagerContext& file_manager_context)
: StringIndex(ASCENDING_SORT), is_built_(false) {
const storage::FileManagerContext& file_manager_context,
bool is_nested_index)
: StringIndex(ASCENDING_SORT),
is_built_(false),
is_nested_index_(is_nested_index) {
if (file_manager_context.Valid()) {
field_id_ = file_manager_context.fieldDataMeta.field_id;
file_manager_ =
@@ -164,8 +169,20 @@ StringIndexSort::BuildWithFieldData(
// Calculate total number of rows
total_num_rows_ = 0;
for (const auto& data : field_datas) {
total_num_rows_ += data->get_num_rows();
if (is_nested_index_) {
for (const auto& data : field_datas) {
auto n = data->get_num_rows();
auto array_column = static_cast<const Array*>(data->Data());
for (int64_t i = 0; i < n; i++) {
if (data->is_valid(i)) {
total_num_rows_ += array_column[i].length();
}
}
}
} else {
for (const auto& data : field_datas) {
total_num_rows_ += data->get_num_rows();
}
}
if (total_num_rows_ == 0) {
@@ -178,9 +195,15 @@ StringIndexSort::BuildWithFieldData(
// Create MemoryImpl and build directly from field data
impl_ = std::make_unique<StringIndexSortMemoryImpl>();
static_cast<StringIndexSortMemoryImpl*>(impl_.get())
->BuildFromFieldData(
field_datas, total_num_rows_, valid_bitset_, idx_to_offsets_);
if (is_nested_index_) {
static_cast<StringIndexSortMemoryImpl*>(impl_.get())
->BuildFromArrayDataNested(
field_datas, total_num_rows_, valid_bitset_, idx_to_offsets_);
} else {
static_cast<StringIndexSortMemoryImpl*>(impl_.get())
->BuildFromFieldData(
field_datas, total_num_rows_, valid_bitset_, idx_to_offsets_);
}
is_built_ = true;
total_size_ = CalculateTotalSize();
@@ -227,6 +250,11 @@ StringIndexSort::Serialize(const Config& config) {
}
res_set.Append("valid_bitset", valid_bitset_data, valid_bitset_size);
// Serialize is_nested_index
std::shared_ptr<uint8_t[]> is_nested_data(new uint8_t[sizeof(bool)]);
memcpy(is_nested_data.get(), &is_nested_index_, sizeof(bool));
res_set.Append("is_nested_index", is_nested_data, sizeof(bool));
milvus::Disassemble(res_set);
return res_set;
}
@@ -303,6 +331,12 @@ StringIndexSort::LoadWithoutAssemble(const BinarySet& binary_set,
}
}
// Deserialize is_nested_index (optional for backward compatibility)
auto is_nested_data = binary_set.GetByName("is_nested_index");
if (is_nested_data != nullptr) {
memcpy(&is_nested_index_, is_nested_data->data.get(), sizeof(bool));
}
auto version_data = binary_set.GetByName("version");
AssertInfo(version_data != nullptr,
"Failed to find 'version' in binary_set");
@@ -534,6 +568,34 @@ StringIndexSortMemoryImpl::BuildFromFieldData(
BuildFromMap(std::move(map), total_num_rows, idx_to_offsets);
}
void
StringIndexSortMemoryImpl::BuildFromArrayDataNested(
const std::vector<FieldDataPtr>& field_datas,
size_t total_num_rows,
TargetBitmap& valid_bitset,
std::vector<int32_t>& idx_to_offsets) {
// Use map to collect unique values and their posting lists
// std::map is sorted
std::map<std::string, PostingList> map;
size_t element_id = 0;
for (const auto& field_data : field_datas) {
auto n = field_data->get_num_rows();
auto array_column = static_cast<const Array*>(field_data->Data());
for (int64_t i = 0; i < n; i++) {
if (!field_data->is_valid(i)) {
continue;
}
for (int64_t j = 0; j < array_column[i].length(); j++) {
auto value = array_column[i].get_data<std::string>(j);
map[value].push_back(static_cast<int32_t>(element_id));
valid_bitset.set(element_id);
element_id++;
}
}
}
BuildFromMap(std::move(map), total_num_rows, idx_to_offsets);
}
size_t
StringIndexSortMemoryImpl::GetSerializedSize() const {
size_t total_size = sizeof(uint32_t); // unique_count
+23 -3
View File
@@ -25,6 +25,7 @@
#include <cstring>
#include <sys/mman.h>
#include <unistd.h>
#include <boost/container/vector.hpp>
#include <folly/small_vector.h>
#include "index/StringIndex.h"
@@ -47,7 +48,8 @@ class StringIndexSort : public StringIndex {
explicit StringIndexSort(
const storage::FileManagerContext& file_manager_context =
storage::FileManagerContext());
storage::FileManagerContext(),
bool is_nested_index = false);
virtual ~StringIndexSort();
@@ -75,6 +77,9 @@ class StringIndexSort : public StringIndex {
void
BuildWithFieldData(const std::vector<FieldDataPtr>& datas) override;
void
BuildWithArrayDataNested(const std::vector<FieldDataPtr>& datas);
// See detailed format in StringIndexSortMemoryImpl::SerializeToBinary
BinarySet
Serialize(const Config& config) override;
@@ -92,6 +97,11 @@ class StringIndexSort : public StringIndex {
LoadWithoutAssemble(const BinarySet& binary_set,
const Config& config) override;
bool
IsNestedIndex() const override {
return is_nested_index_;
}
// Query methods - delegated to impl
const TargetBitmap
In(size_t n, const std::string* values) override;
@@ -152,6 +162,8 @@ class StringIndexSort : public StringIndex {
int64_t total_size_{0};
std::unique_ptr<StringIndexSortImpl> impl_;
bool is_nested_index_ = false;
};
// Abstract interface for implementations
@@ -240,6 +252,12 @@ class StringIndexSortMemoryImpl : public StringIndexSortImpl {
TargetBitmap& valid_bitset,
std::vector<int32_t>& idx_to_offsets);
void
BuildFromArrayDataNested(const std::vector<FieldDataPtr>& field_datas,
size_t total_num_rows,
TargetBitmap& valid_bitset,
std::vector<int32_t>& idx_to_offsets);
// Serialize to binary format
// The binary format is : [unique_count][string_offsets][string_data][post_list_offsets][post_list_data][magic_code]
// string_offsets: array of offsets into string_data section
@@ -483,8 +501,10 @@ using StringIndexSortPtr = std::unique_ptr<StringIndexSort>;
inline StringIndexSortPtr
CreateStringIndexSort(const storage::FileManagerContext& file_manager_context =
storage::FileManagerContext()) {
return std::make_unique<StringIndexSort>(file_manager_context);
storage::FileManagerContext(),
bool is_nested_index = false) {
return std::make_unique<StringIndexSort>(file_manager_context,
is_nested_index);
}
} // namespace milvus::index
@@ -14,6 +14,9 @@
#include <vector>
#include <boost/format.hpp>
#include "index/ScalarIndexSort.h"
#include "index/StringIndexSort.h"
#include "index/Meta.h"
#include "common/Schema.h"
#include "common/ArrayOffsets.h"
#include "query/Plan.h"
@@ -1132,3 +1135,332 @@ TEST(ArrayOffsetsGrowing, MultiplePendingBatches) {
ASSERT_EQ(r10, 5);
ASSERT_EQ(i10, 1);
}
enum class NestedIndexType { NONE, STL_SORT, INVERTED };
std::string
NestedIndexTypeToString(NestedIndexType type) {
switch (type) {
case NestedIndexType::NONE:
return "None";
case NestedIndexType::STL_SORT:
return "StlSort";
case NestedIndexType::INVERTED:
return "Inverted";
default:
return "Unknown";
}
}
// Test fixture for nested index and execution mode testing
// Parameters: (NestedIndexType, bool force_offset_mode)
class ElementFilterNestedIndex
: public ::testing::TestWithParam<std::tuple<NestedIndexType, bool>> {
protected:
NestedIndexType
nested_index_type() const {
return std::get<0>(GetParam());
}
bool
force_offset_mode() const {
return std::get<1>(GetParam());
}
};
TEST_P(ElementFilterNestedIndex, ExecutionMode) {
auto index_type = nested_index_type();
bool offset_mode = force_offset_mode();
// Step 1: Prepare schema with array field
int dim = 4;
auto schema = std::make_shared<Schema>();
auto vec_fid = schema->AddDebugVectorArrayField("structA[array_float_vec]",
DataType::VECTOR_FLOAT,
dim,
knowhere::metric::L2);
auto int_array_fid = schema->AddDebugArrayField(
"structA[price_array]", DataType::INT32, false);
auto int64_fid = schema->AddDebugField("id", DataType::INT64);
schema->set_primary_field_id(int64_fid);
// Use larger N for better selectivity control
size_t N = 1000;
int array_len = 3;
// Step 2: Generate test data
auto raw_data = DataGen(schema, N, 42, 0, 1, array_len);
// Customize int_array data: doc i has elements [i*3+1, i*3+2, i*3+3]
for (int i = 0; i < raw_data.raw_->fields_data_size(); i++) {
auto* field_data = raw_data.raw_->mutable_fields_data(i);
if (field_data->field_id() == int_array_fid.get()) {
field_data->mutable_scalars()
->mutable_array_data()
->mutable_data()
->Clear();
for (int row = 0; row < N; row++) {
auto* array_data = field_data->mutable_scalars()
->mutable_array_data()
->mutable_data()
->Add();
for (int elem = 0; elem < array_len; elem++) {
int value = row * array_len + elem + 1;
array_data->mutable_int_data()->mutable_data()->Add(value);
}
}
break;
}
}
// Step 3: Create sealed segment with field data
auto segment = CreateSealedWithFieldDataLoaded(schema, raw_data);
// Step 4: Load vector index
auto array_vec_values = raw_data.get_col<VectorFieldProto>(vec_fid);
std::vector<float> vector_data(dim * N * array_len);
for (int i = 0; i < N; i++) {
const auto& float_vec = array_vec_values[i].float_vector().data();
for (int j = 0; j < array_len * dim; j++) {
vector_data[i * array_len * dim + j] = float_vec[j];
}
}
auto indexing = GenVecIndexing(N * array_len,
dim,
vector_data.data(),
knowhere::IndexEnum::INDEX_HNSW);
LoadIndexInfo load_index_info;
load_index_info.field_id = vec_fid.get();
load_index_info.index_params = GenIndexParams(indexing.get());
load_index_info.cache_index =
CreateTestCacheIndex("test", std::move(indexing));
load_index_info.index_params["metric_type"] = knowhere::metric::L2;
load_index_info.field_type = DataType::VECTOR_ARRAY;
load_index_info.element_type = DataType::VECTOR_FLOAT;
segment->LoadIndex(load_index_info);
// Step 5: Build and load nested scalar index if needed
if (index_type != NestedIndexType::NONE) {
// Build nested index for int_array field
std::vector<int32_t> all_elements;
all_elements.reserve(N * array_len);
for (int row = 0; row < N; row++) {
for (int elem = 0; elem < array_len; elem++) {
all_elements.push_back(row * array_len + elem + 1);
}
}
milvus::index::IndexBasePtr nested_index;
if (index_type == NestedIndexType::STL_SORT) {
auto stl_index =
std::make_unique<milvus::index::ScalarIndexSort<int32_t>>(
storage::FileManagerContext(), true /* is_nested */);
stl_index->Build(all_elements.size(), all_elements.data(), nullptr);
nested_index = std::move(stl_index);
} else {
// INVERTED index - use BuildWithRawDataForUT is not available for int32
// So we use ScalarIndexSort for now as a workaround
// In real scenario, inverted index would be loaded from disk
auto stl_index =
std::make_unique<milvus::index::ScalarIndexSort<int32_t>>(
storage::FileManagerContext(), true /* is_nested */);
stl_index->Build(all_elements.size(), all_elements.data(), nullptr);
nested_index = std::move(stl_index);
}
// Load nested index to segment
LoadIndexInfo nested_load_info;
nested_load_info.field_id = int_array_fid.get();
nested_load_info.field_type = DataType::ARRAY;
nested_load_info.element_type = DataType::INT32;
nested_load_info.index_params["index_type"] =
index_type == NestedIndexType::STL_SORT
? milvus::index::ASCENDING_SORT
: milvus::index::INVERTED_INDEX_TYPE;
nested_load_info.cache_index =
CreateTestCacheIndex("nested_test", std::move(nested_index));
segment->LoadIndex(nested_load_info);
}
int topK = 5;
// Step 6: Build query plan with appropriate selectivity
// For offset_mode: use very low selectivity (id == 500 -> 0.1% for N=1000)
// For full_mode: use high selectivity (id % 2 == 0 -> 50%)
std::string predicate_expr;
if (offset_mode) {
// Very low selectivity: only 1 doc matches (0.1% for N=1000)
// doc 500's elements are [1501, 1502, 1503] which are in range (100, 2000)
predicate_expr = R"(
term_expr: <
column_info: <
field_id: %3%
data_type: Int64
>
values: <
int64_val: 500
>
>
)";
} else {
// High selectivity: ~50% docs match
predicate_expr = R"(
binary_arith_op_eval_range_expr: <
column_info: <
field_id: %3%
data_type: Int64
>
arith_op: Mod
right_operand: <
int64_val: 2
>
op: Equal
value: <
int64_val: 0
>
>
)";
}
std::string raw_plan =
boost::str(boost::format(R"(vector_anns: <
field_id: %1%
predicates: <
element_filter_expr: <
element_expr: <
binary_range_expr: <
column_info: <
field_id: %2%
data_type: Int32
element_type: Int32
is_element_level: true
>
lower_inclusive: false
upper_inclusive: false
lower_value: <
int64_val: 100
>
upper_value: <
int64_val: 2000
>
>
>
predicate: <
)" +
predicate_expr + R"(
>
struct_name: "structA"
>
>
query_info: <
topk: 5
round_decimal: 3
metric_type: "L2"
search_params: "{\"ef\": 50}"
>
placeholder_tag: "$0">)") %
vec_fid.get() % int_array_fid.get() % int64_fid.get());
proto::plan::PlanNode plan_node;
auto ok =
google::protobuf::TextFormat::ParseFromString(raw_plan, &plan_node);
ASSERT_TRUE(ok) << "Failed to parse element-level filter plan";
auto plan = CreateSearchPlanFromPlanNode(schema, plan_node);
ASSERT_NE(plan, nullptr);
auto num_queries = 1;
auto seed = 1024;
auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed, true);
auto ph_group =
ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
// Step 7: Execute search
auto search_result = segment->Search(plan.get(), ph_group.get(), 1L << 63);
// Step 8: Verify results
ASSERT_NE(search_result, nullptr);
ASSERT_TRUE(search_result->element_level_);
// Count valid results (doc_id >= 0, -1 means invalid/padding)
size_t valid_count = 0;
for (size_t i = 0; i < search_result->seg_offsets_.size(); i++) {
if (search_result->seg_offsets_[i] >= 0) {
valid_count++;
}
}
if (offset_mode) {
// In offset mode with id == 500, only doc 500 matches
// doc 500 has 3 elements, all in range (100, 2000): [1501, 1502, 1503]
ASSERT_GT(valid_count, 0) << "Should have at least one valid result";
ASSERT_LE(valid_count, 3)
<< "At most 3 elements from doc 500 can match";
for (size_t i = 0; i < search_result->seg_offsets_.size(); i++) {
int64_t doc_id = search_result->seg_offsets_[i];
if (doc_id >= 0) {
ASSERT_EQ(doc_id, 500)
<< "In offset mode, only doc 500 should match";
}
}
} else {
// In full mode with id % 2 == 0, even docs match
ASSERT_GT(valid_count, 0) << "Should have valid results";
for (size_t i = 0; i < search_result->seg_offsets_.size(); i++) {
int64_t doc_id = search_result->seg_offsets_[i];
if (doc_id >= 0) {
ASSERT_EQ(doc_id % 2, 0)
<< "Result doc_id " << doc_id << " should be even";
}
}
}
// Verify element values are in range (100, 2000) for valid results
for (size_t i = 0; i < search_result->seg_offsets_.size(); i++) {
int64_t doc_id = search_result->seg_offsets_[i];
if (doc_id >= 0) {
int32_t elem_idx = search_result->element_indices_[i];
int element_value = doc_id * array_len + elem_idx + 1;
ASSERT_GT(element_value, 100);
ASSERT_LT(element_value, 2000);
}
}
// Verify distances are sorted (only check valid results)
float last_distance = -1;
for (size_t i = 0; i < search_result->distances_.size(); ++i) {
if (search_result->seg_offsets_[i] >= 0) {
if (last_distance >= 0) {
ASSERT_LE(last_distance, search_result->distances_[i]);
}
last_distance = search_result->distances_[i];
}
}
std::cout << "Test passed: index_type="
<< NestedIndexTypeToString(index_type)
<< ", offset_mode=" << (offset_mode ? "true" : "false")
<< ", valid_results=" << valid_count << std::endl;
}
INSTANTIATE_TEST_SUITE_P(
ElementFilter,
ElementFilterNestedIndex,
::testing::Combine(::testing::Values(NestedIndexType::NONE,
NestedIndexType::STL_SORT,
NestedIndexType::INVERTED),
::testing::Values(false, true) // force_offset_mode
),
[](const ::testing::TestParamInfo<ElementFilterNestedIndex::ParamType>&
info) {
auto index_type = std::get<0>(info.param);
bool offset_mode = std::get<1>(info.param);
std::string name = NestedIndexTypeToString(index_type);
name += offset_mode ? "_OffsetMode" : "_FullMode";
return name;
});