From a55019d50477d3498152b99a2ff9a2b002672c9d Mon Sep 17 00:00:00 2001 From: Spade A <71589810+SpadeA-Tang@users.noreply.github.com> Date: Wed, 22 Apr 2026 10:17:44 +0800 Subject: [PATCH] feat: impl StructArray --- support range search and iterator search for element-level (#49182) issue: https://github.com/milvus-io/milvus/issues/42148 design doc: https://github.com/milvus-io/milvus-design-docs/blob/main/design_docs/20260306-struct.md --------- Signed-off-by: SpadeA --- internal/core/src/common/QueryResult.h | 35 +- internal/core/src/exec/Driver.cpp | 20 +- ...ode.cpp => IterativeElementFilterNode.cpp} | 51 +- ...terNode.h => IterativeElementFilterNode.h} | 13 +- .../src/exec/operator/IterativeFilterNode.cpp | 2 +- .../src/exec/operator/IterativeFilterNode.h | 2 +- internal/core/src/exec/operator/Utils.h | 9 +- .../src/exec/operator/VectorSearchNode.cpp | 2 +- internal/core/src/plan/PlanNode.h | 26 +- .../core/src/query/CachedSearchIterator.cpp | 45 +- .../core/src/query/CachedSearchIterator.h | 5 + internal/core/src/query/PlanProto.cpp | 8 +- internal/core/src/query/SearchBruteForce.cpp | 4 - internal/core/src/query/SearchOnGrowing.cpp | 72 +- internal/core/src/query/SearchOnSealed.cpp | 80 +- internal/core/src/query/SubSearchResult.h | 19 +- internal/core/src/query/Utils.h | 61 + .../core/unittest/test_element_filter.cpp | 486 ++++ internal/core/unittest/test_utils/DataGen.h | 2 +- internal/proxy/search_util.go | 29 +- internal/proxy/task_search.go | 53 +- internal/proxy/task_search_test.go | 276 +- .../hellomilvus/range_search_test.go | 135 + tests/integration/util_query.go | 45 + .../test_milvus_client_struct_array.py | 855 +++--- ...lvus_client_struct_array_element_search.py | 2313 ++++++++++------- 26 files changed, 3020 insertions(+), 1628 deletions(-) rename internal/core/src/exec/operator/{ElementFilterNode.cpp => IterativeElementFilterNode.cpp} (85%) rename internal/core/src/exec/operator/{ElementFilterNode.h => IterativeElementFilterNode.h} (87%) diff --git a/internal/core/src/common/QueryResult.h b/internal/core/src/common/QueryResult.h index 1c4e764d26..bca67231b7 100644 --- a/internal/core/src/common/QueryResult.h +++ b/internal/core/src/common/QueryResult.h @@ -30,7 +30,6 @@ #include "common/FieldMeta.h" #include "common/ArrayOffsets.h" #include "common/OffsetMapping.h" -#include "query/Utils.h" #include "pb/schema.pb.h" #include "knowhere/index/index_node.h" @@ -157,11 +156,12 @@ class VectorIterator { // returning results in distance-sorted order. class ChunkMergeIterator : public VectorIterator { public: + // Pass nullptr for VECTOR_ARRAY element-level search: the iterator + // returns element IDs, which will be processed by IArrayOffsets. ChunkMergeIterator(int chunk_count, - const milvus::OffsetMapping& offset_mapping, - const std::vector& total_rows_until_chunk = {}, + const milvus::OffsetMapping* offset_mapping, bool larger_is_closer = false) - : offset_mapping_(&offset_mapping), + : offset_mapping_(offset_mapping), heap_(OffsetDisPairComparator(larger_is_closer)) { iterators_.reserve(chunk_count); } @@ -214,22 +214,6 @@ class ChunkMergeIterator : public VectorIterator { } } - private: - int64_t - convert_to_segment_offset(int64_t chunk_offset, int chunk_idx) { - if (total_rows_until_chunk_.size() == 0) { - AssertInfo( - iterators_.size() == 1, - "Wrong state for vectorIterators, which having incorrect " - "kw_iterator count:{} " - "without setting value for chunk_rows, " - "cannot convert chunk_offset to segment_offset correctly", - iterators_.size()); - return chunk_offset; - } - return total_rows_until_chunk_[chunk_idx] + chunk_offset; - } - private: std::vector iterators_; std::priority_queue, @@ -238,7 +222,6 @@ class ChunkMergeIterator : public VectorIterator { heap_; bool sealed = false; const milvus::OffsetMapping* offset_mapping_ = nullptr; - std::vector total_rows_until_chunk_; //currently, ChunkMergeIterator is guaranteed to be used serially without concurrent problem, in the future //we may need to add mutex to protect the variable sealed }; @@ -262,9 +245,8 @@ struct SearchResult { AssembleChunkVectorIterators( int64_t nq, int chunk_count, - const std::vector& total_rows_until_chunk, const std::vector& kw_iterators, - const milvus::OffsetMapping& offset_mapping, + const milvus::OffsetMapping* offset_mapping, bool larger_is_closer = false) { AssertInfo(kw_iterators.size() == nq * chunk_count, "kw_iterators count:{} is not equal to nq*chunk_count:{}, " @@ -276,11 +258,8 @@ struct SearchResult { for (int i = 0, vec_iter_idx = 0; i < kw_iterators.size(); i++) { vec_iter_idx = vec_iter_idx % nq; if (vector_iterators.size() < nq) { - auto chunk_merge_iter = - std::make_shared(chunk_count, - offset_mapping, - total_rows_until_chunk, - larger_is_closer); + auto chunk_merge_iter = std::make_shared( + chunk_count, offset_mapping, larger_is_closer); vector_iterators.emplace_back(chunk_merge_iter); } const auto& kw_iterator = kw_iterators[i]; diff --git a/internal/core/src/exec/Driver.cpp b/internal/core/src/exec/Driver.cpp index 1ff8e02656..159139ffe5 100644 --- a/internal/core/src/exec/Driver.cpp +++ b/internal/core/src/exec/Driver.cpp @@ -34,7 +34,7 @@ #include "exec/operator/AggregationNode.h" #include "exec/operator/CallbackSink.h" #include "exec/operator/ElementFilterBitsNode.h" -#include "exec/operator/ElementFilterNode.h" +#include "exec/operator/IterativeElementFilterNode.h" #include "exec/operator/FilterBitsNode.h" #include "exec/operator/IterativeFilterNode.h" #include "exec/operator/MvccNode.h" @@ -80,10 +80,9 @@ DriverFactory::CreateDriver( tracer::AddEvent("create_operator: FilterBitsNode"); operators.push_back(std::make_unique( id, ctx.get(), filterbitsnode)); - } else if (auto filternode = - std::dynamic_pointer_cast( - plannode)) { - tracer::AddEvent("create_operator: FilterNode"); + } else if (auto filternode = std::dynamic_pointer_cast< + const plan::IterativeFilterNode>(plannode)) { + tracer::AddEvent("create_operator: IterativeFilterNode"); operators.push_back(std::make_unique( id, ctx.get(), filternode)); } else if (auto mvccnode = @@ -134,12 +133,11 @@ DriverFactory::CreateDriver( tracer::AddEvent("create_operator: RescoresNode"); operators.push_back( std::make_unique(id, ctx.get(), rescoresnode)); - } else if (auto node = - std::dynamic_pointer_cast( - plannode)) { - tracer::AddEvent("create_operator: ElementFilterNode"); - operators.push_back( - std::make_unique(id, ctx.get(), node)); + } else if (auto node = std::dynamic_pointer_cast< + const plan::IterativeElementFilterNode>(plannode)) { + tracer::AddEvent("create_operator: IterativeElementFilterNode"); + operators.push_back(std::make_unique( + id, ctx.get(), node)); } else if (auto node = std::dynamic_pointer_cast< const plan::ElementFilterBitsNode>(plannode)) { tracer::AddEvent("create_operator: ElementFilterBitsNode"); diff --git a/internal/core/src/exec/operator/ElementFilterNode.cpp b/internal/core/src/exec/operator/IterativeElementFilterNode.cpp similarity index 85% rename from internal/core/src/exec/operator/ElementFilterNode.cpp rename to internal/core/src/exec/operator/IterativeElementFilterNode.cpp index 091c7c3327..bfee38fa80 100644 --- a/internal/core/src/exec/operator/ElementFilterNode.cpp +++ b/internal/core/src/exec/operator/IterativeElementFilterNode.cpp @@ -14,7 +14,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "ElementFilterNode.h" +#include "IterativeElementFilterNode.h" #include "exec/operator/Utils.h" #include @@ -43,15 +43,16 @@ namespace milvus { namespace exec { -PhyElementFilterNode::PhyElementFilterNode( +PhyIterativeElementFilterNode::PhyIterativeElementFilterNode( int32_t operator_id, DriverContext* driverctx, - const std::shared_ptr& element_filter_node) + const std::shared_ptr& + element_filter_node) : Operator(driverctx, element_filter_node->output_type(), operator_id, element_filter_node->id(), - "PhyElementFilterNode"), + "PhyIterativeElementFilterNode"), struct_name_(element_filter_node->struct_name()), has_doc_predicate_(element_filter_node->has_doc_predicate()) { ExecContext* exec_context = operator_context_->get_exec_context(); @@ -62,18 +63,19 @@ PhyElementFilterNode::PhyElementFilterNode( } void -PhyElementFilterNode::AddInput(RowVectorPtr& input) { +PhyIterativeElementFilterNode::AddInput(RowVectorPtr& input) { input_ = std::move(input); } RowVectorPtr -PhyElementFilterNode::GetOutput() { +PhyIterativeElementFilterNode::GetOutput() { if (is_finished_ || !no_more_input_) { return nullptr; } - tracer::AutoSpan span( - "PhyElementFilterNode::GetOutput", tracer::GetRootSpan(), true); + tracer::AutoSpan span("PhyIterativeElementFilterNode::GetOutput", + tracer::GetRootSpan(), + true); DeferLambda([&]() { is_finished_ = true; }); @@ -89,13 +91,14 @@ PhyElementFilterNode::GetOutput() { if (!search_result.element_level_) { ThrowInfo(ExprInvalid, - "PhyElementFilterNode expects element-level search result"); + "PhyIterativeElementFilterNode expects element-level search " + "result"); } if (!search_result.vector_iterators_.has_value()) { - ThrowInfo( - ExprInvalid, - "PhyElementFilterNode expects vector_iterators in search result"); + ThrowInfo(ExprInvalid, + "PhyIterativeElementFilterNode expects vector_iterators in " + "search result"); } auto segment = query_context_->get_segment(); @@ -132,7 +135,7 @@ PhyElementFilterNode::GetOutput() { : 0; // Step 4: If no doc-level predicate, collect results directly - // (otherwise, downstream FilterNode will do this) + // (otherwise, downstream IterativeFilterNode will do this) if (!has_doc_predicate_) { CollectResults(search_result, array_offsets.get()); } @@ -146,21 +149,21 @@ PhyElementFilterNode::GetOutput() { std::chrono::duration(end_time - start_time) .count(); - tracer::AddEvent( - fmt::format("PhyElementFilterNode: wrapped {} iterators, struct_name: " - "{}, has_doc_predicate: {}, cost_us: {}", - num_iterators, - struct_name_, - has_doc_predicate_, - cost)); + tracer::AddEvent(fmt::format( + "PhyIterativeElementFilterNode: wrapped {} iterators, struct_name: " + "{}, has_doc_predicate: {}, cost_us: {}", + num_iterators, + struct_name_, + has_doc_predicate_, + cost)); // Pass through input to downstream return input_; } void -PhyElementFilterNode::CollectResults(SearchResult& search_result, - const IArrayOffsets* array_offsets) { +PhyIterativeElementFilterNode::CollectResults( + SearchResult& search_result, const IArrayOffsets* array_offsets) { // When there's no doc-level predicate, we need to consume the iterators // and collect the top-K results ourselves. // @@ -231,7 +234,9 @@ PhyElementFilterNode::CollectResults(SearchResult& search_result, search_result.vector_iterators_.reset(); tracer::AddEvent(fmt::format( - "PhyElementFilterNode::CollectResults: nq={}, topk={}", nq, topk)); + "PhyIterativeElementFilterNode::CollectResults: nq={}, topk={}", + nq, + topk)); } } // namespace exec diff --git a/internal/core/src/exec/operator/ElementFilterNode.h b/internal/core/src/exec/operator/IterativeElementFilterNode.h similarity index 87% rename from internal/core/src/exec/operator/ElementFilterNode.h rename to internal/core/src/exec/operator/IterativeElementFilterNode.h index fa9a1dd277..3ac2ed8468 100644 --- a/internal/core/src/exec/operator/ElementFilterNode.h +++ b/internal/core/src/exec/operator/IterativeElementFilterNode.h @@ -32,12 +32,13 @@ namespace milvus { namespace exec { -class PhyElementFilterNode : public Operator { +class PhyIterativeElementFilterNode : public Operator { public: - PhyElementFilterNode(int32_t operator_id, - DriverContext* ctx, - const std::shared_ptr& - element_filter_node); + PhyIterativeElementFilterNode( + int32_t operator_id, + DriverContext* ctx, + const std::shared_ptr& + element_filter_node); bool IsFilter() const override { @@ -75,7 +76,7 @@ class PhyElementFilterNode : public Operator { std::string ToString() const override { - return "PhyElementFilterNode"; + return "PhyIterativeElementFilterNode"; } private: diff --git a/internal/core/src/exec/operator/IterativeFilterNode.cpp b/internal/core/src/exec/operator/IterativeFilterNode.cpp index 7acc035358..c635f15b3c 100644 --- a/internal/core/src/exec/operator/IterativeFilterNode.cpp +++ b/internal/core/src/exec/operator/IterativeFilterNode.cpp @@ -54,7 +54,7 @@ namespace exec { PhyIterativeFilterNode::PhyIterativeFilterNode( int32_t operator_id, DriverContext* driverctx, - const std::shared_ptr& filter) + const std::shared_ptr& filter) : Operator(driverctx, filter->output_type(), operator_id, diff --git a/internal/core/src/exec/operator/IterativeFilterNode.h b/internal/core/src/exec/operator/IterativeFilterNode.h index 9e5cfd0257..affd4a37cf 100644 --- a/internal/core/src/exec/operator/IterativeFilterNode.h +++ b/internal/core/src/exec/operator/IterativeFilterNode.h @@ -39,7 +39,7 @@ class PhyIterativeFilterNode : public Operator { PhyIterativeFilterNode( int32_t operator_id, DriverContext* ctx, - const std::shared_ptr& filter); + const std::shared_ptr& filter); bool IsFilter() const override { diff --git a/internal/core/src/exec/operator/Utils.h b/internal/core/src/exec/operator/Utils.h index 96aabd95b6..eb5578b473 100644 --- a/internal/core/src/exec/operator/Utils.h +++ b/internal/core/src/exec/operator/Utils.h @@ -71,12 +71,17 @@ PrepareVectorIteratorsFromIndex(const SearchInfo& search_info, if (iterators_val.has_value()) { bool larger_is_closer = PositivelyRelated(search_info.metric_type_); + // Element-level search skips row-level mapping (element IDs + // are not row-aligned); see ChunkMergeIterator ctor. + const milvus::OffsetMapping* iter_offset_mapping = + search_info.array_offsets_ != nullptr + ? nullptr + : &index.GetOffsetMapping(); search_result.AssembleChunkVectorIterators( nq, 1, - {0}, iterators_val.value(), - index.GetOffsetMapping(), + iter_offset_mapping, larger_is_closer); } else { std::string operator_type = ""; diff --git a/internal/core/src/exec/operator/VectorSearchNode.cpp b/internal/core/src/exec/operator/VectorSearchNode.cpp index 7c44e39230..8000765244 100644 --- a/internal/core/src/exec/operator/VectorSearchNode.cpp +++ b/internal/core/src/exec/operator/VectorSearchNode.cpp @@ -121,7 +121,7 @@ PhyVectorSearchNode::GetOutput() { } else { // There are two types of execution: pre-filter and iterative filter // For **pre-filter**: FilterBitsNode -> MvccNode -> ElementFilterBitsNode -> VectorSearchNode -> ... - // For **iterative filter**: MvccNode -> VectorSearchNode -> ElementFilterNode -> FilterNode -> ... + // For **iterative filter**: MvccNode -> VectorSearchNode -> IterativeElementFilterNode -> IterativeFilterNode -> ... // // When element_level_ is true, we need to transform doc-level bitset // to element-level bitset. In pre-filter path, ElementFilterBitsNode diff --git a/internal/core/src/plan/PlanNode.h b/internal/core/src/plan/PlanNode.h index d1280719e0..ba025afc9f 100644 --- a/internal/core/src/plan/PlanNode.h +++ b/internal/core/src/plan/PlanNode.h @@ -89,11 +89,11 @@ class PlanNode { }; using PlanNodePtr = std::shared_ptr; -class FilterNode : public PlanNode { +class IterativeFilterNode : public PlanNode { public: - FilterNode(const PlanNodeId& id, - expr::TypedExprPtr filter, - std::vector sources) + IterativeFilterNode(const PlanNodeId& id, + expr::TypedExprPtr filter, + std::vector sources) : PlanNode(id), sources_{std::move(sources)}, filter_(std::move(filter)) { @@ -120,7 +120,7 @@ class FilterNode : public PlanNode { std::string_view name() const override { - return "Filter"; + return "IterativeFilter"; } std::string @@ -186,13 +186,13 @@ class FilterBitsNode : public PlanNode { const expr::TypedExprPtr filter_; }; -class ElementFilterNode : public PlanNode { +class IterativeElementFilterNode : public PlanNode { public: - ElementFilterNode(const PlanNodeId& id, - expr::TypedExprPtr element_filter, - std::string struct_name, - std::vector sources, - bool has_doc_predicate = true) + IterativeElementFilterNode(const PlanNodeId& id, + expr::TypedExprPtr element_filter, + std::string struct_name, + std::vector sources, + bool has_doc_predicate = true) : PlanNode(id), sources_{std::move(sources)}, element_filter_(std::move(element_filter)), @@ -232,13 +232,13 @@ class ElementFilterNode : public PlanNode { std::string_view name() const override { - return "ElementFilter"; + return "IterativeElementFilter"; } std::string ToString() const override { return fmt::format( - "ElementFilterNode:[struct_name:{}, element_filter:{}, " + "IterativeElementFilterNode:[struct_name:{}, element_filter:{}, " "has_doc_predicate:{}]", struct_name_, element_filter_->ToString(), diff --git a/internal/core/src/query/CachedSearchIterator.cpp b/internal/core/src/query/CachedSearchIterator.cpp index fb6f99088f..22d444e97c 100644 --- a/internal/core/src/query/CachedSearchIterator.cpp +++ b/internal/core/src/query/CachedSearchIterator.cpp @@ -20,6 +20,7 @@ #include "common/QueryInfo.h" #include "common/QueryResult.h" #include "common/Utils.h" +#include "common/VectorArray.h" #include "index/Utils.h" #include "index/VectorIndex.h" #include "knowhere/expected.h" @@ -113,6 +114,17 @@ CachedSearchIterator::CachedSearchIterator( nq_ = query_ds.num_queries; Init(search_info); + // VECTOR_ARRAY element-level search: growing stores each row as a + // separate VectorArray with its own backing allocation, so we must + // flatten per-chunk into a contiguous buffer that knowhere can read. + // array_offsets_ != nullptr is the element-level signal (multi-search- + // multi emb-list iterator is rejected upstream, so we don't branch on + // it here). + const bool is_element_level = search_info.array_offsets_ != nullptr; + if (is_element_level) { + chunk_buffers_.reserve(num_chunks_); + } + iterators_.reserve(nq_ * num_chunks_); InitializeChunkedIterators( query_ds, @@ -120,12 +132,32 @@ CachedSearchIterator::CachedSearchIterator( index_info, bitset, data_type, - [&vec_data, vec_size_per_chunk, row_count](int64_t chunk_id) { + [this, &vec_data, vec_size_per_chunk, row_count, is_element_level]( + int64_t chunk_id) { const void* chunk_data = vec_data->get_chunk_data(chunk_id); // no need to store a PinWrapper for growing, because vec_data is guaranteed to not be evicted. int64_t chunk_size = std::min( vec_size_per_chunk, row_count - chunk_id * vec_size_per_chunk); - return std::make_pair(chunk_data, chunk_size); + if (!is_element_level) { + return std::make_pair(chunk_data, chunk_size); + } + + auto va_ptr = reinterpret_cast(chunk_data); + int64_t total_bytes = 0; + int64_t total_elements = 0; + for (int64_t i = 0; i < chunk_size; ++i) { + total_bytes += va_ptr[i].byte_size(); + total_elements += va_ptr[i].length(); + } + auto buf = std::make_unique(total_bytes); + auto* ptr = buf.get(); + for (int64_t i = 0; i < chunk_size; ++i) { + memcpy(ptr, va_ptr[i].data(), va_ptr[i].byte_size()); + ptr += va_ptr[i].byte_size(); + } + const void* flat_data = buf.get(); + chunk_buffers_.emplace_back(std::move(buf)); + return std::make_pair(flat_data, total_elements); }); } @@ -155,7 +187,7 @@ CachedSearchIterator::CachedSearchIterator( index_info, bitset, data_type, - [this, column](int64_t chunk_id) { + [this, column, &search_info](int64_t chunk_id) { auto pw = column->DataOfChunk(nullptr, chunk_id) .transform([](const auto& x) { return static_cast(x); @@ -166,6 +198,13 @@ CachedSearchIterator::CachedSearchIterator( if (offset_mapping.IsEnabled() && !valid_count_per_chunk.empty()) { chunk_size = valid_count_per_chunk[chunk_id]; } + // For element-level search on vector array field, chunk_size + // must be the element count in this chunk, not the row count. + if (search_info.array_offsets_ != nullptr) { + auto elem_offsets_pw = + column->VectorArrayOffsets(nullptr, chunk_id); + chunk_size = elem_offsets_pw.get()[chunk_size]; + } // pw guarantees chunk_data is kept alive. auto chunk_data = pw.get(); pin_wrappers_.emplace_back(std::move(pw)); diff --git a/internal/core/src/query/CachedSearchIterator.h b/internal/core/src/query/CachedSearchIterator.h index e9a856fafe..cbfe821dd1 100644 --- a/internal/core/src/query/CachedSearchIterator.h +++ b/internal/core/src/query/CachedSearchIterator.h @@ -94,6 +94,11 @@ class CachedSearchIterator { // used only for sealed segment with chunked data std::vector> pin_wrappers_; + // used only for growing segment with VECTOR_ARRAY element-level search: + // knowhere needs a contiguous flat buffer, but growing stores each row as + // a separate VectorArray with its own backing allocation, so we flatten + // per-chunk here and keep the buffers alive alongside the iterators. + std::vector> chunk_buffers_; int64_t batch_size_ = 0; std::vector iterators_; int8_t sign_ = 1; diff --git a/internal/core/src/query/PlanProto.cpp b/internal/core/src/query/PlanProto.cpp index 74c021c7b1..c60fd6de46 100644 --- a/internal/core/src/query/PlanProto.cpp +++ b/internal/core/src/query/PlanProto.cpp @@ -537,7 +537,7 @@ ProtoParser::PlanNodeFromProto(const planpb::PlanNode& plan_node_proto) { // Add element-level filter if needed if (is_element_level) { bool has_doc_pred = (doc_expr != nullptr); - plannode = std::make_shared( + plannode = std::make_shared( milvus::plan::GetNextPlanNodeId(), element_expr, struct_name, @@ -548,7 +548,7 @@ ProtoParser::PlanNodeFromProto(const planpb::PlanNode& plan_node_proto) { // Add doc-level filter if present if (doc_expr) { - plannode = std::make_shared( + plannode = std::make_shared( milvus::plan::GetNextPlanNodeId(), doc_expr, sources); sources = std::vector{plannode}; } @@ -598,8 +598,8 @@ ProtoParser::PlanNodeFromProto(const planpb::PlanNode& plan_node_proto) { } } else { // No user filter. Force non-iterative: the iterative path requires - // a FilterNode for row-by-row post-filtering, which doesn't apply - // here (TTL uses bitmap pre-filtering via FilterBitsNode). + // an IterativeFilterNode for row-by-row post-filtering, which doesn't + // apply here (TTL uses bitmap pre-filtering via FilterBitsNode). plan_node->search_info_.iterative_filter_execution = false; // When entity-level TTL is enabled, add an AlwaysTrueExpr so that diff --git a/internal/core/src/query/SearchBruteForce.cpp b/internal/core/src/query/SearchBruteForce.cpp index f66864d6c3..73ceafa5b2 100644 --- a/internal/core/src/query/SearchBruteForce.cpp +++ b/internal/core/src/query/SearchBruteForce.cpp @@ -180,10 +180,6 @@ BruteForceSearch(const dataset::SearchDataset& query_ds, } if (search_cfg.contains(RADIUS)) { - AssertInfo(data_type != DataType::VECTOR_ARRAY, - "Vector array(embedding list) is not supported for range " - "search"); - if (search_cfg.contains(RANGE_FILTER)) { CheckRangeSearchParam(search_cfg[RADIUS], search_cfg[RANGE_FILTER], diff --git a/internal/core/src/query/SearchOnGrowing.cpp b/internal/core/src/query/SearchOnGrowing.cpp index 9000852907..c4fc30bc33 100644 --- a/internal/core/src/query/SearchOnGrowing.cpp +++ b/internal/core/src/query/SearchOnGrowing.cpp @@ -205,10 +205,22 @@ SearchOnGrowing(const segcore::SegmentGrowingImpl& segment, return; } + // Element-level search (embedding-search-embedding): knowhere sees + // a scalar vector type and the per-chunk size must be measured in + // elements. Compute this before the iterator_v2 branch so both + // paths share the substitution. Emb-list (multi-search-multi) + // iterator is rejected by the proxy; the assert below is + // defense-in-depth. + const bool is_element_level_search = + data_type == DataType::VECTOR_ARRAY && + info.array_offsets_ != nullptr; + const auto iter_data_type = + is_element_level_search ? element_type : data_type; + if (info.iterator_v2_info_.has_value()) { - AssertInfo(data_type != DataType::VECTOR_ARRAY, - "vector array(embedding list) is not supported for " - "vector iterator"); + AssertInfo(iter_data_type != DataType::VECTOR_ARRAY, + "embedding list (multi-search-multi) iterator is not " + "supported on vector array fields"); CachedSearchIterator cached_iter(search_dataset, vec_ptr, @@ -216,28 +228,19 @@ SearchOnGrowing(const segcore::SegmentGrowingImpl& segment, info, index_info, search_bitset, - data_type); + iter_data_type); cached_iter.NextBatch(info, search_result); - if (offset_mapping.IsEnabled()) { - TransformOffset(search_result.seg_offsets_, offset_mapping); - } + FinalizeVectorSearchOffsets( + search_result, offset_mapping, info.array_offsets_.get()); return; } auto vec_size_per_chunk = vec_ptr->get_size_per_chunk(); auto max_chunk = upper_div(active_count, vec_size_per_chunk); - // embedding search embedding on embedding list - bool embedding_search = false; - if (data_type == DataType::VECTOR_ARRAY && - info.array_offsets_ != nullptr) { - embedding_search = true; - } - - // Track cumulative element offset for embedding search. - // For embedding_search, begin_id must be the cumulative element - // count (not row offset), because ArrayOffsets maps global - // element IDs to row IDs. + // Track cumulative element offset for element-level search. + // begin_id must be the cumulative element count (not row offset), + // because ArrayOffsets maps global element IDs to row IDs. int64_t cumulative_element_offset = 0; std::vector offsets; @@ -268,7 +271,7 @@ SearchOnGrowing(const segcore::SegmentGrowingImpl& segment, buf = std::make_unique(size); - if (embedding_search) { + if (is_element_level_search) { auto count = 0; auto ptr = buf.get(); for (int i = 0; i < size_per_chunk; ++i) { @@ -301,13 +304,8 @@ SearchOnGrowing(const segcore::SegmentGrowingImpl& segment, } } - auto vector_type = data_type; - if (embedding_search) { - vector_type = element_type; - } - if (milvus::exec::UseVectorIterator(info)) { - AssertInfo(vector_type != DataType::VECTOR_ARRAY, + AssertInfo(iter_data_type != DataType::VECTOR_ARRAY, "vector array(embedding list) is not supported for " "vector iterator"); @@ -321,7 +319,7 @@ SearchOnGrowing(const segcore::SegmentGrowingImpl& segment, info, index_info, search_bitset, - vector_type); + iter_data_type); final_qr.merge(sub_qr); } else { auto sub_qr = BruteForceSearch(search_dataset, @@ -329,41 +327,41 @@ SearchOnGrowing(const segcore::SegmentGrowingImpl& segment, info, index_info, search_bitset, - vector_type, + iter_data_type, element_type, op_context); final_qr.merge(sub_qr); } } if (milvus::exec::UseVectorIterator(info)) { - std::vector chunk_rows(max_chunk, 0); - for (int i = 1; i < max_chunk; ++i) { - chunk_rows[i] = i * vec_size_per_chunk; - } bool larger_is_closer = PositivelyRelated(info.metric_type_); + // Element-level search skips row-level mapping (element IDs are + // not row-aligned); see ChunkMergeIterator ctor. + const milvus::OffsetMapping* iter_offset_mapping = + is_element_level_search ? nullptr : &offset_mapping; search_result.AssembleChunkVectorIterators( num_queries, max_chunk, - chunk_rows, final_qr.chunk_iterators(), - offset_mapping, + iter_offset_mapping, larger_is_closer); } else { + // See FinalizeVectorSearchOffsets for the rationale: + // element-level and row-level remapping are mutually exclusive. if (info.array_offsets_ != nullptr) { auto [seg_offsets, elem_indicies] = final_qr.convert_to_element_offsets( info.array_offsets_.get()); search_result.seg_offsets_ = std::move(seg_offsets); search_result.element_indices_ = std::move(elem_indicies); - search_result.element_level_ = true; } else { + if (offset_mapping.IsEnabled()) { + TransformOffset(final_qr.mutable_offsets(), offset_mapping); + } search_result.seg_offsets_ = std::move(final_qr.mutable_offsets()); } search_result.distances_ = std::move(final_qr.mutable_distances()); - if (offset_mapping.IsEnabled()) { - TransformOffset(search_result.seg_offsets_, offset_mapping); - } } search_result.unity_topK_ = topk; search_result.total_nq_ = num_queries; diff --git a/internal/core/src/query/SearchOnSealed.cpp b/internal/core/src/query/SearchOnSealed.cpp index 60f1a157d8..5bec7f53c6 100644 --- a/internal/core/src/query/SearchOnSealed.cpp +++ b/internal/core/src/query/SearchOnSealed.cpp @@ -115,7 +115,8 @@ SearchOnSealedIndex(const Schema& schema, CachedSearchIterator cached_iter( *vec_index, dataset, search_info, search_bitset); cached_iter.NextBatch(search_info, search_result); - TransformOffset(search_result.seg_offsets_, offset_mapping); + FinalizeVectorSearchOffsets( + search_result, offset_mapping, search_info.array_offsets_.get()); return; } @@ -138,30 +139,11 @@ SearchOnSealedIndex(const Schema& schema, std::round(distances[i] * multiplier) / multiplier; } } - - // Handle element-level conversion if needed - if (search_info.array_offsets_ != nullptr) { - std::vector element_ids = - std::move(search_result.seg_offsets_); - search_result.seg_offsets_.resize(element_ids.size()); - search_result.element_indices_.resize(element_ids.size()); - - for (size_t i = 0; i < element_ids.size(); i++) { - if (element_ids[i] == INVALID_SEG_OFFSET) { - search_result.seg_offsets_[i] = INVALID_SEG_OFFSET; - search_result.element_indices_[i] = -1; - } else { - auto [doc_id, elem_index] = - search_info.array_offsets_->ElementIDToRowID( - element_ids[i]); - search_result.seg_offsets_[i] = doc_id; - search_result.element_indices_[i] = elem_index; - } - } - search_result.element_level_ = true; - } } - TransformOffset(search_result.seg_offsets_, offset_mapping); + FinalizeVectorSearchOffsets( + search_result, + offset_mapping, + use_iterator ? nullptr : search_info.array_offsets_.get()); search_result.total_nq_ = num_queries; search_result.unity_topK_ = topK; } @@ -218,10 +200,24 @@ SearchOnSealedColumn(const Schema& schema, } } + // For element-level search (embedding-search-embedding), the underlying + // knowhere search is keyed by the scalar element type rather than + // VECTOR_ARRAY, and per-chunk sizes must be counted in elements. + bool is_element_level_search = + field.get_data_type() == DataType::VECTOR_ARRAY && + query_offsets == nullptr; + if (is_element_level_search) { + data_type = element_type; + } + if (search_info.iterator_v2_info_.has_value()) { + // Element-level search has already replaced data_type with + // element_type above. This assert still guards emb-list + // (multi-search-multi) iterator, which proxy should have rejected + // — keep it as a defense-in-depth check. AssertInfo(data_type != DataType::VECTOR_ARRAY, - "vector array(embedding list) is not supported for " - "vector iterator"); + "embedding list (multi-search-multi) iterator is not " + "supported on vector array fields"); CachedSearchIterator cached_iter(column, query_dataset, @@ -230,9 +226,8 @@ SearchOnSealedColumn(const Schema& schema, search_bitview, data_type); cached_iter.NextBatch(search_info, result); - if (offset_mapping.IsEnabled()) { - TransformOffset(result.seg_offsets_, offset_mapping); - } + FinalizeVectorSearchOffsets( + result, offset_mapping, search_info.array_offsets_.get()); return; } @@ -243,17 +238,6 @@ SearchOnSealedColumn(const Schema& schema, search_info.metric_type_, search_info.round_decimal_); - // For element-level search (embedding-search-embedding), we need to use - // element count instead of row count - bool is_element_level_search = - field.get_data_type() == DataType::VECTOR_ARRAY && - query_offsets == nullptr; - - if (is_element_level_search) { - // embedding-search-embedding on embedding list pattern - data_type = element_type; - } - auto offset = 0; auto vector_chunks = column->GetAllChunks(op_context); const auto& valid_count_per_chunk = column->GetValidCountPerChunk(); @@ -312,27 +296,31 @@ SearchOnSealedColumn(const Schema& schema, } if (milvus::exec::UseVectorIterator(search_info)) { bool larger_is_closer = PositivelyRelated(search_info.metric_type_); + // Element-level search skips row-level mapping (element IDs are + // not row-aligned); see ChunkMergeIterator ctor. + const milvus::OffsetMapping* iter_offset_mapping = + search_info.array_offsets_ != nullptr ? nullptr : &offset_mapping; result.AssembleChunkVectorIterators(num_queries, num_chunk, - column->GetNumRowsUntilChunk(), final_qr.chunk_iterators(), - offset_mapping, + iter_offset_mapping, larger_is_closer); } else { + // See FinalizeVectorSearchOffsets for the rationale: element-level + // and row-level remapping are mutually exclusive. if (search_info.array_offsets_ != nullptr) { auto [seg_offsets, elem_indicies] = final_qr.convert_to_element_offsets( search_info.array_offsets_.get()); result.seg_offsets_ = std::move(seg_offsets); result.element_indices_ = std::move(elem_indicies); - result.element_level_ = true; } else { + if (offset_mapping.IsEnabled()) { + TransformOffset(final_qr.mutable_offsets(), offset_mapping); + } result.seg_offsets_ = std::move(final_qr.mutable_offsets()); } result.distances_ = std::move(final_qr.mutable_distances()); - if (offset_mapping.IsEnabled()) { - TransformOffset(result.seg_offsets_, offset_mapping); - } } result.unity_topK_ = query_dataset.topk; result.total_nq_ = query_dataset.num_queries; diff --git a/internal/core/src/query/SubSearchResult.h b/internal/core/src/query/SubSearchResult.h index dca20c5862..6896165352 100644 --- a/internal/core/src/query/SubSearchResult.h +++ b/internal/core/src/query/SubSearchResult.h @@ -19,6 +19,7 @@ #include "common/Utils.h" #include "knowhere/index/index_node.h" #include "common/ArrayOffsets.h" +#include "query/Utils.h" namespace milvus::query { class SubSearchResult { @@ -119,23 +120,7 @@ class SubSearchResult { std::pair, std::vector> convert_to_element_offsets(const IArrayOffsets* array_offsets) { - std::vector doc_offsets; - std::vector element_indices; - doc_offsets.reserve(offsets_.size()); - element_indices.reserve(offsets_.size()); - for (size_t i = 0; i < offsets_.size(); i++) { - if (offsets_[i] == INVALID_SEG_OFFSET) { - doc_offsets.push_back(INVALID_SEG_OFFSET); - element_indices.push_back(-1); - } else { - auto [doc_id, elem_index] = - array_offsets->ElementIDToRowID(offsets_[i]); - doc_offsets.push_back(doc_id); - element_indices.push_back(elem_index); - } - } - return std::make_pair(std::move(doc_offsets), - std::move(element_indices)); + return milvus::query::ApplyElementIDMapping(offsets_, *array_offsets); } private: diff --git a/internal/core/src/query/Utils.h b/internal/core/src/query/Utils.h index 80aed027a3..4d8c726c1d 100644 --- a/internal/core/src/query/Utils.h +++ b/internal/core/src/query/Utils.h @@ -14,12 +14,44 @@ #include #include +#include +#include + +#include "common/ArrayOffsets.h" #include "common/BitsetView.h" +#include "common/Consts.h" #include "common/OffsetMapping.h" +#include "common/QueryResult.h" #include "common/Types.h" #include "common/Utils.h" namespace milvus::query { +// Map logical element IDs returned by knowhere to (doc_id, elem_idx) pairs +// via ArrayOffsets. Caller must ensure the input `element_ids` are already +// in logical space (i.e. apply TransformOffset first when offset_mapping is +// enabled), since ArrayOffsets::ElementIDToRowID is keyed on logical +// element IDs. +inline std::pair, std::vector> +ApplyElementIDMapping(const std::vector& element_ids, + const milvus::IArrayOffsets& array_offsets) { + std::vector doc_offsets; + std::vector element_indices; + doc_offsets.reserve(element_ids.size()); + element_indices.reserve(element_ids.size()); + for (size_t i = 0; i < element_ids.size(); i++) { + if (element_ids[i] == INVALID_SEG_OFFSET) { + doc_offsets.push_back(INVALID_SEG_OFFSET); + element_indices.push_back(-1); + } else { + auto [doc_id, elem_index] = + array_offsets.ElementIDToRowID(element_ids[i]); + doc_offsets.push_back(doc_id); + element_indices.push_back(elem_index); + } + } + return std::make_pair(std::move(doc_offsets), std::move(element_indices)); +} + inline TargetBitmap TransformBitset(const BitsetView& bitset, const milvus::OffsetMapping& mapping) { @@ -46,6 +78,35 @@ TransformOffset(std::vector& seg_offsets, } } +// Map knowhere's raw offsets back to logical space. The two inputs are +// mutually exclusive: +// +// - array_offsets != nullptr (VECTOR_ARRAY element-level search): +// knowhere returns physical element IDs. ArrayOffsets is built by +// walking every row in the segment and advancing the row counter on +// every row (including empty/null rows, which occupy a zero-length +// element range), so ElementIDToRowID produces (logical_row_id, +// elem_idx) directly. No OffsetMapping pass is needed. +// +// - array_offsets == nullptr (plain vector field): when OffsetMapping is +// enabled, the index/chunk was built over valid rows only, so +// knowhere's physical row IDs must be remapped to logical via +// OffsetMapping. When OffsetMapping is disabled, TransformOffset is a +// no-op. +inline void +FinalizeVectorSearchOffsets(SearchResult& result, + const milvus::OffsetMapping& offset_mapping, + const milvus::IArrayOffsets* array_offsets) { + if (array_offsets != nullptr) { + auto [doc_offsets, elem_indices] = + ApplyElementIDMapping(result.seg_offsets_, *array_offsets); + result.seg_offsets_ = std::move(doc_offsets); + result.element_indices_ = std::move(elem_indices); + } else { + TransformOffset(result.seg_offsets_, offset_mapping); + } +} + template inline bool Match(const T& x, const U& y, OpType op) { diff --git a/internal/core/unittest/test_element_filter.cpp b/internal/core/unittest/test_element_filter.cpp index 7164942cc5..deb20cbe04 100644 --- a/internal/core/unittest/test_element_filter.cpp +++ b/internal/core/unittest/test_element_filter.cpp @@ -16,6 +16,8 @@ #include #include #include +#include +#include #include #include #include @@ -2870,3 +2872,487 @@ TEST(ElementFilter, GrowingMultiChunkElementSearch) { << " element_value=" << element_value << " should be < 400"; } } + +namespace { + +constexpr int kElemDim = 4; +constexpr int kElemArrayLen = 3; +constexpr size_t kElemN = 200; +constexpr uint64_t kFixtureVecSeed = 0x5EEDBEEFULL; +constexpr int kElemTopK = 10; +// Ground-truth target: the query vector equals row kElemTargetDoc's element +// kElemTargetElem, so an exact (BF) search must return it first with +// distance ~0. A mid-range row + non-zero elem index makes "always returns +// the first element" bugs visible. +constexpr int64_t kElemTargetDoc = 5; +constexpr int32_t kElemTargetElem = 1; + +struct ElementSearchFixture { + SchemaPtr schema; + FieldId vec_fid; + FieldId int64_fid; + GeneratedData raw_data; + std::vector flat_data; // kN * kArrayLen * kDim + std::vector query_data; // kDim +}; + +inline ElementSearchFixture +MakeElementSearchFixture() { + ElementSearchFixture f; + f.schema = std::make_shared(); + f.vec_fid = f.schema->AddDebugVectorArrayField( + "structA[array_vec]", DataType::VECTOR_FLOAT, kElemDim, "L2"); + f.int64_fid = f.schema->AddDebugField("id", DataType::INT64); + f.schema->set_primary_field_id(f.int64_fid); + + f.raw_data = DataGen(f.schema, kElemN, 42, 0, 1, kElemArrayLen); + + // DataGen's VECTOR_ARRAY branch uses the same seed for every row, so all + // rows share identical element vectors. Overwrite the whole vector-array + // field with per-(row, elem) unique values so that brute-force search has + // no ties and iterator_v2 batch boundaries see strictly increasing + // distances. + { + std::mt19937_64 rng(kFixtureVecSeed); + std::normal_distribution distr(0.0f, 1.0f); + for (int i = 0; i < f.raw_data.raw_->fields_data_size(); ++i) { + auto* fd = f.raw_data.raw_->mutable_fields_data(i); + if (fd->field_id() != f.vec_fid.get()) { + continue; + } + auto* vec_array = fd->mutable_vectors()->mutable_vector_array(); + for (int row = 0; row < vec_array->data_size(); ++row) { + auto* row_data = vec_array->mutable_data(row) + ->mutable_float_vector() + ->mutable_data(); + for (int k = 0; k < kElemArrayLen * kElemDim; ++k) { + row_data->Set(k, distr(rng)); + } + } + // Pin the target element to a unique sentinel so exact search + // deterministically returns (kElemTargetDoc, kElemTargetElem). + auto* target_row = vec_array->mutable_data(kElemTargetDoc) + ->mutable_float_vector() + ->mutable_data(); + for (int k = 0; k < kElemDim; ++k) { + target_row->Set(kElemTargetElem * kElemDim + k, 7.0f); + } + break; + } + } + + auto array_vec_values = f.raw_data.get_col(f.vec_fid); + f.flat_data.resize(array_vec_values.size() * kElemArrayLen * kElemDim); + for (size_t i = 0; i < array_vec_values.size(); ++i) { + const auto& float_vec = array_vec_values[i].float_vector().data(); + for (int j = 0; j < kElemArrayLen * kElemDim; ++j) { + f.flat_data[i * kElemArrayLen * kElemDim + j] = float_vec[j]; + } + } + const size_t target_off = + (kElemTargetDoc * kElemArrayLen + kElemTargetElem) * kElemDim; + f.query_data.assign(f.flat_data.begin() + target_off, + f.flat_data.begin() + target_off + kElemDim); + return f; +} + +inline proto::common::PlaceholderGroup +MakeElementLevelPlaceholder(const std::vector& query_data) { + auto raw = CreatePlaceholderGroupFromBlob( + /*num_queries=*/1, kElemDim, query_data.data()); + raw.mutable_placeholders(0)->set_element_level(true); + return raw; +} + +inline void +LoadElementHnswIndex(SegmentSealed* segment, + FieldId vec_fid, + const std::vector& flat_data) { + auto indexing = GenVecIndexing(kElemN * kElemArrayLen, + kElemDim, + flat_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); +} + +inline void +ExpectElementLevelShape(const milvus::SearchResult& sr) { + ASSERT_TRUE(sr.element_level_); + ASSERT_FALSE(sr.seg_offsets_.empty()); + ASSERT_EQ(sr.element_indices_.size(), sr.seg_offsets_.size()); + for (size_t i = 0; i < sr.seg_offsets_.size(); ++i) { + if (sr.seg_offsets_[i] < 0) { + continue; + } + ASSERT_GE(sr.element_indices_[i], 0); + ASSERT_LT(sr.element_indices_[i], kElemArrayLen); + } +} + +inline void +ExpectSortedAscending(const milvus::SearchResult& sr) { + int valid = 0; + for (size_t i = 0; i < sr.distances_.size(); ++i) { + if (sr.seg_offsets_[i] < 0) { + continue; + } + ++valid; + if (i > 0 && sr.seg_offsets_[i - 1] >= 0) { + ASSERT_LE(sr.distances_[i - 1], sr.distances_[i]); + } + } + ASSERT_GT(valid, 0); +} + +inline void +ExpectWithinRadius(const milvus::SearchResult& sr, float radius) { + for (size_t i = 0; i < sr.seg_offsets_.size(); ++i) { + if (sr.seg_offsets_[i] < 0) { + continue; + } + ASSERT_LT(sr.distances_[i], radius); + ASSERT_GE(sr.distances_[i], 0.0f); + } +} + +// Exact path (BF): the target element must sit at index 0 with distance ~0. +inline void +ExpectTopOneIsTarget(const milvus::SearchResult& sr) { + ASSERT_FALSE(sr.seg_offsets_.empty()); + ASSERT_EQ(sr.seg_offsets_[0], kElemTargetDoc); + ASSERT_EQ(sr.element_indices_[0], kElemTargetElem); + ASSERT_NEAR(sr.distances_[0], 0.0f, 1e-5f); +} + +// Approximate path (HNSW): the target must appear in the top-K, but may not +// be strictly first. +inline void +ExpectTargetInTopK(const milvus::SearchResult& sr) { + bool found = false; + for (size_t i = 0; i < sr.seg_offsets_.size(); ++i) { + if (sr.seg_offsets_[i] == kElemTargetDoc && + sr.element_indices_[i] == kElemTargetElem) { + found = true; + EXPECT_NEAR(sr.distances_[i], 0.0f, 1e-3f); + break; + } + } + ASSERT_TRUE(found) << "Target (doc=" << kElemTargetDoc + << ", elem=" << kElemTargetElem + << ") not found in approximate search result"; +} + +} // namespace + +TEST(ElementVectorSearch, GrowingBruteForce_RangeSearch) { + auto f = MakeElementSearchFixture(); + auto segment = CreateGrowingSegment(f.schema, empty_index_meta); + segment->PreInsert(kElemN); + segment->Insert(0, + kElemN, + f.raw_data.row_ids_.data(), + f.raw_data.timestamps_.data(), + f.raw_data.raw_); + + ScopedSchemaHandle handle(*f.schema); + const float radius = 1e6f; + const std::string search_params = + R"({"radius": )" + std::to_string(radius) + R"(, "range_filter": 0.0})"; + auto plan_bytes = handle.ParseSearch( + "", "structA[array_vec]", kElemTopK, "L2", search_params, 3); + auto plan = + CreateSearchPlanByExpr(f.schema, plan_bytes.data(), plan_bytes.size()); + + auto ph_group_raw = MakeElementLevelPlaceholder(f.query_data); + auto ph_group = + ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString()); + auto sr = segment->Search(plan.get(), ph_group.get(), 1L << 63); + ASSERT_NE(sr, nullptr); + ExpectElementLevelShape(*sr); + ExpectWithinRadius(*sr, radius); + ExpectTopOneIsTarget(*sr); +} + +TEST(ElementVectorSearch, GrowingBruteForce_IteratorV2) { + auto f = MakeElementSearchFixture(); + auto segment = CreateGrowingSegment(f.schema, empty_index_meta); + segment->PreInsert(kElemN); + segment->Insert(0, + kElemN, + f.raw_data.row_ids_.data(), + f.raw_data.timestamps_.data(), + f.raw_data.raw_); + + ScopedSchemaHandle handle(*f.schema); + auto plan_bytes = + handle.ParseSearchIterator("", + "structA[array_vec]", + kElemTopK, + "L2", + R"({"ef": 50})", + static_cast(kElemTopK), + "", + std::nullopt, + 3); + auto plan = + CreateSearchPlanByExpr(f.schema, plan_bytes.data(), plan_bytes.size()); + + auto ph_group_raw = MakeElementLevelPlaceholder(f.query_data); + auto ph_group = + ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString()); + auto sr = segment->Search(plan.get(), ph_group.get(), 1L << 63); + ASSERT_NE(sr, nullptr); + ExpectElementLevelShape(*sr); + ExpectSortedAscending(*sr); + ExpectTopOneIsTarget(*sr); +} + +TEST(ElementVectorSearch, SealedBruteForce_RangeSearch) { + auto f = MakeElementSearchFixture(); + auto segment = CreateSealedWithFieldDataLoaded(f.schema, f.raw_data); + + ScopedSchemaHandle handle(*f.schema); + const float radius = 1e6f; + const std::string search_params = + R"({"radius": )" + std::to_string(radius) + R"(, "range_filter": 0.0})"; + auto plan_bytes = handle.ParseSearch( + "", "structA[array_vec]", kElemTopK, "L2", search_params, 3); + auto plan = + CreateSearchPlanByExpr(f.schema, plan_bytes.data(), plan_bytes.size()); + + auto ph_group_raw = MakeElementLevelPlaceholder(f.query_data); + auto ph_group = + ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString()); + auto sr = segment->Search(plan.get(), ph_group.get(), 1L << 63); + ASSERT_NE(sr, nullptr); + ExpectElementLevelShape(*sr); + ExpectWithinRadius(*sr, radius); + ExpectTopOneIsTarget(*sr); +} + +TEST(ElementVectorSearch, SealedBruteForce_IteratorV2) { + auto f = MakeElementSearchFixture(); + auto segment = CreateSealedWithFieldDataLoaded(f.schema, f.raw_data); + + ScopedSchemaHandle handle(*f.schema); + auto plan_bytes = + handle.ParseSearchIterator("", + "structA[array_vec]", + kElemTopK, + "L2", + R"({"ef": 50})", + static_cast(kElemTopK), + "", + std::nullopt, + 3); + auto plan = + CreateSearchPlanByExpr(f.schema, plan_bytes.data(), plan_bytes.size()); + + auto ph_group_raw = MakeElementLevelPlaceholder(f.query_data); + auto ph_group = + ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString()); + auto sr = segment->Search(plan.get(), ph_group.get(), 1L << 63); + ASSERT_NE(sr, nullptr); + ExpectElementLevelShape(*sr); + ExpectSortedAscending(*sr); + ExpectTopOneIsTarget(*sr); +} + +TEST(ElementVectorSearch, SealedIndex_RangeSearch) { + auto f = MakeElementSearchFixture(); + auto segment = CreateSealedWithFieldDataLoaded(f.schema, f.raw_data); + LoadElementHnswIndex(segment.get(), f.vec_fid, f.flat_data); + + ScopedSchemaHandle handle(*f.schema); + const float radius = 1e6f; + const std::string search_params = R"({"ef": 50, "radius": )" + + std::to_string(radius) + + R"(, "range_filter": 0.0})"; + auto plan_bytes = handle.ParseSearch( + "", "structA[array_vec]", kElemTopK, "L2", search_params, 3); + auto plan = + CreateSearchPlanByExpr(f.schema, plan_bytes.data(), plan_bytes.size()); + + auto ph_group_raw = MakeElementLevelPlaceholder(f.query_data); + auto ph_group = + ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString()); + auto sr = segment->Search(plan.get(), ph_group.get(), 1L << 63); + ASSERT_NE(sr, nullptr); + ExpectElementLevelShape(*sr); + ExpectWithinRadius(*sr, radius); + ExpectTargetInTopK(*sr); +} + +TEST(ElementVectorSearch, SealedIndex_IteratorV2) { + auto f = MakeElementSearchFixture(); + auto segment = CreateSealedWithFieldDataLoaded(f.schema, f.raw_data); + LoadElementHnswIndex(segment.get(), f.vec_fid, f.flat_data); + + ScopedSchemaHandle handle(*f.schema); + auto plan_bytes = + handle.ParseSearchIterator("", + "structA[array_vec]", + kElemTopK, + "L2", + R"({"ef": 50})", + static_cast(kElemTopK), + "", + std::nullopt, + 3); + auto plan = + CreateSearchPlanByExpr(f.schema, plan_bytes.data(), plan_bytes.size()); + + auto ph_group_raw = MakeElementLevelPlaceholder(f.query_data); + auto ph_group = + ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString()); + auto sr = segment->Search(plan.get(), ph_group.get(), 1L << 63); + ASSERT_NE(sr, nullptr); + ExpectElementLevelShape(*sr); + ExpectSortedAscending(*sr); + ExpectTargetInTopK(*sr); +} + +// ---- Multi-chunk growing ------------------------------------------------- +// Default chunk_rows is 1024, so the 6 tests above fit in one chunk. Here +// we shrink chunk_rows so N=200 spans ~4 chunks, exercising the per-chunk +// flatten loop in SearchOnGrowing.cpp and the chunk merge in +// ChunkMergeIterator. + +namespace { + +constexpr int64_t kElemChunkRows = 64; + +inline SegmentGrowingPtr +CreateMultiChunkGrowingSegment(const ElementSearchFixture& f) { + auto config = SegcoreConfig::default_config(); + config.set_chunk_rows(kElemChunkRows); + auto segment = CreateGrowingSegment(f.schema, empty_index_meta, 1, config); + segment->PreInsert(kElemN); + segment->Insert(0, + kElemN, + f.raw_data.row_ids_.data(), + f.raw_data.timestamps_.data(), + f.raw_data.raw_); + static_assert(kElemN > kElemChunkRows, + "N must exceed chunk_rows to exercise multi-chunk paths"); + return segment; +} + +} // namespace + +TEST(ElementVectorSearch, GrowingMultiChunk_RangeSearch) { + auto f = MakeElementSearchFixture(); + auto segment = CreateMultiChunkGrowingSegment(f); + + ScopedSchemaHandle handle(*f.schema); + const float radius = 1e6f; + const std::string search_params = + R"({"radius": )" + std::to_string(radius) + R"(, "range_filter": 0.0})"; + auto plan_bytes = handle.ParseSearch( + "", "structA[array_vec]", kElemTopK, "L2", search_params, 3); + auto plan = + CreateSearchPlanByExpr(f.schema, plan_bytes.data(), plan_bytes.size()); + + auto ph_group_raw = MakeElementLevelPlaceholder(f.query_data); + auto ph_group = + ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString()); + auto sr = segment->Search(plan.get(), ph_group.get(), 1L << 63); + ASSERT_NE(sr, nullptr); + ExpectElementLevelShape(*sr); + ExpectWithinRadius(*sr, radius); + ExpectTopOneIsTarget(*sr); +} + +TEST(ElementVectorSearch, GrowingMultiChunk_IteratorV2) { + auto f = MakeElementSearchFixture(); + auto segment = CreateMultiChunkGrowingSegment(f); + + ScopedSchemaHandle handle(*f.schema); + auto plan_bytes = + handle.ParseSearchIterator("", + "structA[array_vec]", + kElemTopK, + "L2", + R"({"ef": 50})", + static_cast(kElemTopK), + "", + std::nullopt, + 3); + auto plan = + CreateSearchPlanByExpr(f.schema, plan_bytes.data(), plan_bytes.size()); + + auto ph_group_raw = MakeElementLevelPlaceholder(f.query_data); + auto ph_group = + ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString()); + auto sr = segment->Search(plan.get(), ph_group.get(), 1L << 63); + ASSERT_NE(sr, nullptr); + ExpectElementLevelShape(*sr); + ExpectSortedAscending(*sr); + ExpectTopOneIsTarget(*sr); +} + +// ---- Iterator_v2: consecutive NextBatch rounds --------------------------- +// Each search call runs exactly one NextBatch. Chain rounds by feeding the +// previous batch's last distance as last_bound: post-round distances must +// stay strictly greater than that bound (L2, lower-is-better). + +TEST(ElementVectorSearch, SealedBruteForce_IteratorV2_MultiBatch) { + auto f = MakeElementSearchFixture(); + auto segment = CreateSealedWithFieldDataLoaded(f.schema, f.raw_data); + + ScopedSchemaHandle handle(*f.schema); + std::optional last_bound; + std::set> seen; + const int kRounds = 3; + for (int round = 0; round < kRounds; ++round) { + auto plan_bytes = + handle.ParseSearchIterator("", + "structA[array_vec]", + kElemTopK, + "L2", + R"({"ef": 50})", + static_cast(kElemTopK), + "", + last_bound, + 3); + auto plan = CreateSearchPlanByExpr( + f.schema, plan_bytes.data(), plan_bytes.size()); + auto ph_group_raw = MakeElementLevelPlaceholder(f.query_data); + auto ph_group = + ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString()); + auto sr = segment->Search(plan.get(), ph_group.get(), 1L << 63); + ASSERT_NE(sr, nullptr); + ExpectElementLevelShape(*sr); + ExpectSortedAscending(*sr); + + float round_last = last_bound.value_or(-1.0f); + for (size_t i = 0; i < sr->seg_offsets_.size(); ++i) { + if (sr->seg_offsets_[i] < 0) { + continue; + } + if (last_bound.has_value()) { + ASSERT_GT(sr->distances_[i], *last_bound) + << "round " << round + << " returned a distance not strictly past last_bound"; + } + auto key = + std::make_pair(sr->seg_offsets_[i], sr->element_indices_[i]); + ASSERT_TRUE(seen.insert(key).second) + << "iterator returned duplicate (doc=" << key.first + << ", elem=" << key.second << ") across rounds"; + round_last = sr->distances_[i]; + } + last_bound = round_last; + } + ASSERT_GE(seen.size(), static_cast(kElemTopK)) + << "multi-batch iterator should accumulate strictly more results " + << "than a single batch"; +} diff --git a/internal/core/unittest/test_utils/DataGen.h b/internal/core/unittest/test_utils/DataGen.h index e99b1e52c6..e08d79cfaa 100644 --- a/internal/core/unittest/test_utils/DataGen.h +++ b/internal/core/unittest/test_utils/DataGen.h @@ -61,7 +61,7 @@ namespace milvus::segcore { struct GeneratedData { std::vector row_ids_; std::vector timestamps_; - InsertRecordProto* raw_; + InsertRecordProto* raw_ = nullptr; std::vector field_ids; SchemaPtr schema_; diff --git a/internal/proxy/search_util.go b/internal/proxy/search_util.go index 643946ce98..a1cf8ef492 100644 --- a/internal/proxy/search_util.go +++ b/internal/proxy/search_util.go @@ -568,38 +568,13 @@ func parseSearchInfo(searchParamsPair []*commonpb.KeyValuePair, schema *schemapb return nil, fmt.Errorf("parse iterator v2 info failed: %w", err) } - // 7. check search for embedding list - annsFieldName, _ := funcutil.GetAttrByKeyFromRepeatedKV(AnnsFieldKey, searchParamsPair) - if annsFieldName != "" { - annField := typeutil.GetFieldByName(schema, annsFieldName) - if annField != nil && annField.GetDataType() == schemapb.DataType_ArrayOfVector { - if isRangeSearch { - return nil, merr.WrapErrParameterInvalid("", "", - "range search is not supported for vector array (embedding list) fields, fieldName:", annsFieldName) - } - - // For hybrid search (rankInfo != nil), group by on vector array is not supported. - // For simple search, group by validation is handled in initSearchRequest() - // where placeholder type is available to distinguish element-level vs embedding list. - if groupByFieldId > 0 && rankParams != nil { - return nil, merr.WrapErrParameterInvalid("", "", - "group by search is not supported for vector array (embedding list) fields, fieldName:", annsFieldName) - } - - if isIterator { - return nil, merr.WrapErrParameterInvalid("", "", - "search iterator is not supported for vector array (embedding list) fields, fieldName:", annsFieldName) - } - } - } - - // 8. parse order_by_fields + // 7. parse order_by_fields orderByFields, err := parseOrderByFields(searchParamsPair, schema) if err != nil { return nil, err } - // 9. validate iterator + order_by combination is not allowed + // 8. validate iterator + order_by combination is not allowed if isIterator && len(orderByFields) > 0 { return nil, merr.WrapErrParameterInvalid("", "", "order_by is not supported when using search iterator") diff --git a/internal/proxy/task_search.go b/internal/proxy/task_search.go index fe507be9dd..883ca0f95e 100644 --- a/internal/proxy/task_search.go +++ b/internal/proxy/task_search.go @@ -10,6 +10,7 @@ import ( "github.com/cockroachdb/errors" "github.com/samber/lo" + "github.com/tidwall/gjson" "go.opentelemetry.io/otel" "go.uber.org/zap" "google.golang.org/protobuf/proto" @@ -462,11 +463,32 @@ func (t *searchTask) initAdvancedSearchRequest(ctx context.Context) error { queryFieldIDs := []int64{} for index, subReq := range t.request.GetSubReqs() { // For hybrid search, order_by_fields comes from main search params, not sub-search params - plan, queryInfo, offset, _, _, searchType, err := t.tryGeneratePlan(subReq.GetSearchParams(), subReq.GetDsl(), subReq.GetExprTemplateValues()) + plan, queryInfo, offset, subIsIterator, _, searchType, err := t.tryGeneratePlan(subReq.GetSearchParams(), subReq.GetDsl(), subReq.GetExprTemplateValues()) if err != nil { return err } + // Hybrid search does not yet support vector array (embedding list) fields: + // placeholder type is per sub-request and element-level vs embedding-list-level + // differentiation is not wired up on this path. Reject range search / iterator / + // group by on such fields here; plain top-K stays allowed for now (unchanged + // behavior). + annsField := typeutil.GetField(t.schema.CollectionSchema, queryInfo.GetQueryFieldId()) + if annsField != nil && annsField.GetDataType() == schemapb.DataType_ArrayOfVector { + if gjson.Get(queryInfo.GetSearchParams(), radiusKey).Exists() { + return merr.WrapErrParameterInvalid("", "", + "range search is not supported for vector array (embedding list) fields in hybrid search, fieldName:"+annsField.GetName()) + } + if t.rankParams.GetGroupByFieldId() > 0 { + return merr.WrapErrParameterInvalid("", "", + "group by search is not supported for vector array (embedding list) fields in hybrid search, fieldName:"+annsField.GetName()) + } + if subIsIterator { + return merr.WrapErrParameterInvalid("", "", + "search iterator is not supported for vector array (embedding list) fields in hybrid search, fieldName:"+annsField.GetName()) + } + } + ignoreGrowing := t.IgnoreGrowing if !ignoreGrowing { // fetch ignore_growing from sub search param if not set in search request @@ -775,13 +797,28 @@ func (t *searchTask) initSearchRequest(ctx context.Context) error { return err } - // For ArrayOfVector fields with group by: - // 1. Embedding list search does not support group by - // 2. Element-level search only supports group by PK (doc-level dedup) - if queryInfo.GetGroupByFieldId() > 0 { - annsField := typeutil.GetField(t.schema.CollectionSchema, t.FieldId) - if annsField != nil && annsField.GetDataType() == schemapb.DataType_ArrayOfVector { - if isEmbeddingListPlaceholderType(placeholderType) { + // For ArrayOfVector fields, the placeholder type decides the search semantics: + // - Element-level (plain vector placeholder): behaves like a normal single-vector + // search; supports range search, iterator, and group by primary key. + // - Embedding-list-level (multi-search-multi): does not support range search, + // iterator, or group by (other than the PK case above). + annsField := typeutil.GetField(t.schema.CollectionSchema, t.FieldId) + if annsField != nil && annsField.GetDataType() == schemapb.DataType_ArrayOfVector { + isEmbList := isEmbeddingListPlaceholderType(placeholderType) + + if isEmbList { + if gjson.Get(queryInfo.GetSearchParams(), radiusKey).Exists() { + return merr.WrapErrParameterInvalid("", "", + "range search is not supported for multi-search-multi on embedding list fields") + } + if t.isIterator { + return merr.WrapErrParameterInvalid("", "", + "search iterator is not supported for multi-search-multi on embedding list fields") + } + } + + if queryInfo.GetGroupByFieldId() > 0 { + if isEmbList { return merr.WrapErrParameterInvalid("", "", "group by is not supported for multi-search-multi on embedding list fields") } diff --git a/internal/proxy/task_search_test.go b/internal/proxy/task_search_test.go index 1516cfeb0c..6c32658ec0 100644 --- a/internal/proxy/task_search_test.go +++ b/internal/proxy/task_search_test.go @@ -3729,7 +3729,12 @@ func TestSearchTask_parseSearchInfo(t *testing.T) { } } - t.Run("vector array with range search", func(t *testing.T) { + // ArrayOfVector-specific validation (range / iterator / group-by) has been + // moved out of parseSearchInfo. The checks below verify that parseSearchInfo + // now passes regardless of those features; the real rejections are asserted + // in TestSearchTask_ArrayOfVectorSimpleSearch and + // TestSearchTask_ArrayOfVectorHybridSearch. + t.Run("vector array with range search passes parseSearchInfo", func(t *testing.T) { schema := createSchemaWithVectorArray("embeddings_list") params := createSearchParams("embeddings_list") @@ -3737,11 +3742,8 @@ func TestSearchTask_parseSearchInfo(t *testing.T) { resetSearchParamsValue(params, ParamsKey, `{"nprobe": 10, "radius": 0.2}`) searchInfo, err := parseSearchInfo(params, schema, nil, false) - assert.Error(t, err) - assert.Nil(t, searchInfo) - assert.ErrorIs(t, err, merr.ErrParameterInvalid) - fmt.Println(err.Error()) - assert.Contains(t, err.Error(), "range search is not supported for vector array (embedding list) fields") + assert.NoError(t, err) + assert.NotNil(t, searchInfo) }) t.Run("vector array with group by passes parseSearchInfo", func(t *testing.T) { @@ -3761,7 +3763,7 @@ func TestSearchTask_parseSearchInfo(t *testing.T) { assert.NotNil(t, searchInfo) }) - t.Run("vector array with iterator", func(t *testing.T) { + t.Run("vector array with iterator passes parseSearchInfo", func(t *testing.T) { schema := createSchemaWithVectorArray("embeddings_list") params := createSearchParams("embeddings_list") @@ -3772,14 +3774,11 @@ func TestSearchTask_parseSearchInfo(t *testing.T) { }) searchInfo, err := parseSearchInfo(params, schema, nil, false) - assert.Error(t, err) - assert.Nil(t, searchInfo) - assert.ErrorIs(t, err, merr.ErrParameterInvalid) - assert.Contains(t, err.Error(), "search iterator is not supported for vector array (embedding list) fields") - assert.Contains(t, err.Error(), "embeddings_list") + assert.NoError(t, err) + assert.NotNil(t, searchInfo) }) - t.Run("vector array with iterator v2", func(t *testing.T) { + t.Run("vector array with iterator v2 passes parseSearchInfo", func(t *testing.T) { schema := createSchemaWithVectorArray("embeddings_list") params := createSearchParams("embeddings_list") @@ -3800,11 +3799,8 @@ func TestSearchTask_parseSearchInfo(t *testing.T) { ) searchInfo, err := parseSearchInfo(params, schema, nil, false) - assert.Error(t, err) - assert.Nil(t, searchInfo) - assert.ErrorIs(t, err, merr.ErrParameterInvalid) - assert.Contains(t, err.Error(), "search iterator is not supported for vector array (embedding list) fields") - assert.Contains(t, err.Error(), "embeddings_list") + assert.NoError(t, err) + assert.NotNil(t, searchInfo) }) t.Run("normal search on vector array should succeed", func(t *testing.T) { @@ -3863,18 +3859,6 @@ func TestSearchTask_parseSearchInfo(t *testing.T) { assert.NotNil(t, searchInfo.planInfo) }) - t.Run("vector array with range search", func(t *testing.T) { - schema := createSchemaWithVectorArray("embeddings_list") - params := createSearchParams("embeddings_list") - resetSearchParamsValue(params, ParamsKey, `{"nprobe": 10, "radius": 0.2}`) - - searchInfo, err := parseSearchInfo(params, schema, nil, false) - assert.Error(t, err) - assert.Nil(t, searchInfo) - // Should fail on range search first - assert.Contains(t, err.Error(), "range search is not supported for vector array (embedding list) fields") - }) - t.Run("no anns field specified", func(t *testing.T) { schema := createSchemaWithVectorArray("embeddings_list") params := getValidSearchParams() @@ -3896,7 +3880,9 @@ func TestSearchTask_parseSearchInfo(t *testing.T) { // Should not trigger vector array validation for non-existent field }) - t.Run("hybrid search with outer group by on vector array", func(t *testing.T) { + t.Run("hybrid search with outer group by on vector array passes parseSearchInfo", func(t *testing.T) { + // Hybrid + ArrayOfVector + group-by is rejected in initAdvancedSearchRequest, + // not in parseSearchInfo. This subtest verifies the parse step no longer fails. schema := createSchemaWithVectorArray("embeddings_list") // Create rank params with group by @@ -3916,12 +3902,9 @@ func TestSearchTask_parseSearchInfo(t *testing.T) { assert.NoError(t, err) searchParams := createSearchParams("embeddings_list") - // Parse search info with rank params searchInfo, err := parseSearchInfo(searchParams, schema, parsedRankParams, false) - assert.Error(t, err) - assert.Nil(t, searchInfo) - assert.ErrorIs(t, err, merr.ErrParameterInvalid) - assert.Contains(t, err.Error(), "group by search is not supported for vector array (embedding list) fields") + assert.NoError(t, err) + assert.NotNil(t, searchInfo) }) }) @@ -5852,6 +5835,227 @@ func TestSearchTask_ArrayOfVectorGroupBy(t *testing.T) { }) } +// TestSearchTask_ArrayOfVectorSimpleSearch verifies the placeholder-type-aware +// validation in initSearchRequest for range search and search iterator on +// ArrayOfVector fields: +// - element-level (plain vector placeholder) is allowed for both features. +// - embedding-list-level (EmbList placeholder) is rejected for both. +func TestSearchTask_ArrayOfVectorSimpleSearch(t *testing.T) { + paramtable.Init() + ctx := context.Background() + + schema := &schemapb.CollectionSchema{ + Name: "test_collection", + Fields: []*schemapb.FieldSchema{ + {FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true}, + {FieldID: 101, Name: "regular_vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "4"}}}, + }, + StructArrayFields: []*schemapb.StructArrayFieldSchema{ + { + FieldID: 103, + Name: "struct_array", + Fields: []*schemapb.FieldSchema{ + {FieldID: 104, Name: "emb_vec", DataType: schemapb.DataType_ArrayOfVector, ElementType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "4"}}}, + }, + }, + }, + } + schemaInfo := newSchemaInfo(schema) + + makePlaceholderGroup := func(phType commonpb.PlaceholderType) []byte { + phg := &commonpb.PlaceholderGroup{ + Placeholders: []*commonpb.PlaceholderValue{{ + Tag: "$0", + Type: phType, + Values: [][]byte{{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}}, // 4 floats + }}, + } + bs, _ := proto.Marshal(phg) + return bs + } + + // paramsJSON is inlined into the ParamsKey field; setting withIterator adds + // the iterator flag as a separate KV pair. + makeTask := func(annsField string, phType commonpb.PlaceholderType, paramsJSON string, withIterator bool) *searchTask { + params := []*commonpb.KeyValuePair{ + {Key: AnnsFieldKey, Value: annsField}, + {Key: TopKKey, Value: "10"}, + {Key: common.MetricTypeKey, Value: metric.L2}, + {Key: ParamsKey, Value: paramsJSON}, + } + if withIterator { + params = append(params, &commonpb.KeyValuePair{Key: IteratorField, Value: "True"}) + } + + phgBytes := makePlaceholderGroup(phType) + return &searchTask{ + ctx: ctx, + collectionName: "test_collection", + SearchRequest: &internalpb.SearchRequest{ + CollectionID: 1, + PartitionIDs: []int64{1}, + OutputFieldsId: []int64{100}, + PlaceholderGroup: nil, + DslType: commonpb.DslType_BoolExprV1, + }, + request: &milvuspb.SearchRequest{ + CollectionName: "test_collection", + OutputFields: []string{"pk"}, + SearchParams: params, + SearchInput: &milvuspb.SearchRequest_PlaceholderGroup{ + PlaceholderGroup: phgBytes, + }, + Nq: 1, + ConsistencyLevel: commonpb.ConsistencyLevel_Session, + }, + schema: schemaInfo, + translatedOutputFields: []string{"pk"}, + tr: timerecord.NewTimeRecorder("test"), + queryInfos: []*planpb.QueryInfo{{}}, + } + } + + const rangeParams = `{"nprobe": 10, "radius": 0.2}` + const plainParams = `{"nprobe": 10}` + + t.Run("element-level range search should succeed", func(t *testing.T) { + task := makeTask("emb_vec", commonpb.PlaceholderType_FloatVector, rangeParams, false) + err := task.initSearchRequest(ctx) + assert.NoError(t, err) + }) + + t.Run("element-level iterator should succeed", func(t *testing.T) { + task := makeTask("emb_vec", commonpb.PlaceholderType_FloatVector, plainParams, true) + err := task.initSearchRequest(ctx) + assert.NoError(t, err) + }) + + t.Run("emblist range search should fail", func(t *testing.T) { + task := makeTask("emb_vec", commonpb.PlaceholderType_EmbListFloatVector, rangeParams, false) + err := task.initSearchRequest(ctx) + assert.Error(t, err) + assert.ErrorIs(t, err, merr.ErrParameterInvalid) + assert.Contains(t, err.Error(), "range search is not supported for multi-search-multi") + }) + + t.Run("emblist iterator should fail", func(t *testing.T) { + task := makeTask("emb_vec", commonpb.PlaceholderType_EmbListFloatVector, plainParams, true) + err := task.initSearchRequest(ctx) + assert.Error(t, err) + assert.ErrorIs(t, err, merr.ErrParameterInvalid) + assert.Contains(t, err.Error(), "search iterator is not supported for multi-search-multi") + }) + + t.Run("regular vector range search should succeed", func(t *testing.T) { + // Regression: new checks must not impact plain FloatVector fields. + task := makeTask("regular_vec", commonpb.PlaceholderType_FloatVector, rangeParams, false) + err := task.initSearchRequest(ctx) + assert.NoError(t, err) + }) +} + +// TestSearchTask_ArrayOfVectorHybridSearch verifies that hybrid search rejects +// range search, search iterator, and group-by on ArrayOfVector fields directly +// at the initAdvancedSearchRequest layer (hybrid does not yet support the +// element-level/embedding-list-level split, so all three are rejected). +func TestSearchTask_ArrayOfVectorHybridSearch(t *testing.T) { + paramtable.Init() + ctx := context.Background() + + schema := &schemapb.CollectionSchema{ + Name: "test_collection", + Fields: []*schemapb.FieldSchema{ + {FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true}, + {FieldID: 101, Name: "regular_vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "4"}}}, + {FieldID: 102, Name: "scalar_field", DataType: schemapb.DataType_VarChar}, + }, + StructArrayFields: []*schemapb.StructArrayFieldSchema{ + { + FieldID: 103, + Name: "struct_array", + Fields: []*schemapb.FieldSchema{ + {FieldID: 104, Name: "emb_vec", DataType: schemapb.DataType_ArrayOfVector, ElementType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "4"}}}, + }, + }, + }, + } + schemaInfo := newSchemaInfo(schema) + + // buildHybridTask constructs a hybrid-search task with a single sub-request. + // rangeRadius != "" attaches a radius param to the sub-request; withIterator + // appends the iterator flag; groupByField != "" adds GroupByFieldKey to the + // outer rank params. + buildHybridTask := func(annsField string, rangeRadius string, withIterator bool, groupByField string) *searchTask { + paramsJSON := `{"nprobe": 10}` + if rangeRadius != "" { + paramsJSON = `{"nprobe": 10, "radius": ` + rangeRadius + `}` + } + subParams := []*commonpb.KeyValuePair{ + {Key: common.MetricTypeKey, Value: metric.L2}, + {Key: ParamsKey, Value: paramsJSON}, + {Key: AnnsFieldKey, Value: annsField}, + {Key: TopKKey, Value: "10"}, + } + if withIterator { + subParams = append(subParams, &commonpb.KeyValuePair{Key: IteratorField, Value: "True"}) + } + + outerParams := []*commonpb.KeyValuePair{ + {Key: LimitKey, Value: "10"}, + } + if groupByField != "" { + outerParams = append(outerParams, &commonpb.KeyValuePair{Key: GroupByFieldKey, Value: groupByField}) + } + + return &searchTask{ + ctx: ctx, + SearchRequest: &internalpb.SearchRequest{ + Base: &commonpb.MsgBase{ + MsgType: commonpb.MsgType_Search, + Timestamp: uint64(time.Now().UnixNano()), + }, + }, + request: &milvuspb.SearchRequest{ + CollectionName: "test_collection", + SearchParams: outerParams, + SubReqs: []*milvuspb.SubSearchRequest{ + { + Dsl: "", + PlaceholderGroup: nil, + SearchParams: subParams, + }, + }, + }, + schema: schemaInfo, + tr: timerecord.NewTimeRecorder("test"), + } + } + + t.Run("hybrid with ArrayOfVector range search should fail", func(t *testing.T) { + qt := buildHybridTask("emb_vec", "0.2", false, "") + err := qt.initAdvancedSearchRequest(ctx) + assert.Error(t, err) + assert.ErrorIs(t, err, merr.ErrParameterInvalid) + assert.Contains(t, err.Error(), "range search is not supported for vector array (embedding list) fields in hybrid search") + }) + + t.Run("hybrid with ArrayOfVector iterator should fail", func(t *testing.T) { + qt := buildHybridTask("emb_vec", "", true, "") + err := qt.initAdvancedSearchRequest(ctx) + assert.Error(t, err) + assert.ErrorIs(t, err, merr.ErrParameterInvalid) + assert.Contains(t, err.Error(), "search iterator is not supported for vector array (embedding list) fields in hybrid search") + }) + + t.Run("hybrid with ArrayOfVector group by should fail", func(t *testing.T) { + qt := buildHybridTask("emb_vec", "", false, "scalar_field") + err := qt.initAdvancedSearchRequest(ctx) + assert.Error(t, err) + assert.ErrorIs(t, err, merr.ErrParameterInvalid) + assert.Contains(t, err.Error(), "group by search is not supported for vector array (embedding list) fields in hybrid search") + }) +} + func TestSearchTask_SearchRequeryPolicy(t *testing.T) { paramtable.Init() ctx := context.Background() diff --git a/tests/integration/hellomilvus/range_search_test.go b/tests/integration/hellomilvus/range_search_test.go index 55a11fe4d3..d67f5058cb 100644 --- a/tests/integration/hellomilvus/range_search_test.go +++ b/tests/integration/hellomilvus/range_search_test.go @@ -30,6 +30,7 @@ import ( "github.com/milvus-io/milvus/pkg/v2/util/funcutil" "github.com/milvus-io/milvus/pkg/v2/util/merr" "github.com/milvus-io/milvus/pkg/v2/util/metric" + "github.com/milvus-io/milvus/pkg/v2/util/typeutil" "github.com/milvus-io/milvus/tests/integration" ) @@ -329,3 +330,137 @@ func (s *HelloMilvusSuite) TestRangeSearchL2() { log.Info("=========================") log.Info("=========================") } + +// TestRangeSearchElementLevelL2 runs range search on a VECTOR_ARRAY sub-field +// in element-level mode: HNSW + L2 index on the sub-field, single-vector +// placeholder with element_level=true, radius/range_filter per L2 semantics. +func (s *HelloMilvusSuite) TestRangeSearchElementLevelL2() { + c := s.Cluster + ctx, cancel := context.WithCancel(c.GetContext()) + defer cancel() + + prefix := "TestRangeSearchElementLevelL2" + dbName := "" + collectionName := prefix + funcutil.GenRandomStr() + dim := 32 + rowNum := 1000 + + schema := integration.ConstructSchemaOfVecDataTypeWithStruct(collectionName, dim, true) + marshaledSchema, err := proto.Marshal(schema) + s.NoError(err) + + createCollectionStatus, err := c.MilvusClient.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{ + DbName: dbName, + CollectionName: collectionName, + Schema: marshaledSchema, + ShardsNum: common.DefaultShardsNum, + }) + s.NoError(err) + s.NoError(merr.Error(createCollectionStatus)) + + fVecColumn := integration.NewFloatVectorFieldData(integration.FloatVecField, rowNum, dim) + structColumn := integration.NewStructArrayFieldData(schema.StructArrayFields[0], integration.StructArrayField, rowNum, dim) + hashKeys := integration.GenerateHashKeys(rowNum) + insertResult, err := c.MilvusClient.Insert(ctx, &milvuspb.InsertRequest{ + DbName: dbName, + CollectionName: collectionName, + FieldsData: []*schemapb.FieldData{fVecColumn, structColumn}, + HashKeys: hashKeys, + NumRows: uint32(rowNum), + }) + s.NoError(err) + s.True(merr.Ok(insertResult.GetStatus())) + + flushResp, err := c.MilvusClient.Flush(ctx, &milvuspb.FlushRequest{ + DbName: dbName, + CollectionNames: []string{collectionName}, + }) + s.NoError(err) + segmentIDs, has := flushResp.GetCollSegIDs()[collectionName] + s.Require().True(has) + s.Require().NotEmpty(segmentIDs) + flushTs, has := flushResp.GetCollFlushTs()[collectionName] + s.Require().True(has) + s.WaitForFlush(ctx, segmentIDs.GetData(), flushTs, dbName, collectionName) + + // regular float vector index, required for loading + createIndexStatus, err := c.MilvusClient.CreateIndex(ctx, &milvuspb.CreateIndexRequest{ + CollectionName: collectionName, + FieldName: integration.FloatVecField, + IndexName: "_default", + ExtraParams: integration.ConstructIndexParam(dim, integration.IndexFaissIvfFlat, metric.L2), + }) + s.NoError(err) + s.NoError(merr.Error(createIndexStatus)) + s.WaitForIndexBuilt(ctx, collectionName, integration.FloatVecField) + + // non-EmbList metric on ArrayOfVector → each element vector is indexed + // independently (proxy task_index.go:658). + subFieldName := typeutil.ConcatStructFieldName(integration.StructArrayField, integration.StructSubFloatVecField) + createIndexStatus, err = c.MilvusClient.CreateIndex(ctx, &milvuspb.CreateIndexRequest{ + CollectionName: collectionName, + FieldName: subFieldName, + IndexName: "struct_vec_element_index", + ExtraParams: integration.ConstructIndexParam(dim, integration.IndexHNSW, metric.L2), + }) + s.NoError(err) + s.NoError(merr.Error(createIndexStatus)) + s.WaitForIndexBuilt(ctx, collectionName, subFieldName) + + loadStatus, err := c.MilvusClient.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{ + DbName: dbName, + CollectionName: collectionName, + }) + s.NoError(err) + s.NoError(merr.Error(loadStatus)) + s.WaitForLoad(ctx, collectionName) + + // L2 semantics: radius is upper bound, range_filter is lower bound. + expr := fmt.Sprintf("%s > 0", integration.Int64Field) + nq := 2 + topk := 10 + roundDecimal := -1 + radius := 20 + filter := 10 + + params := integration.GetSearchParams(integration.IndexHNSW, metric.L2) + + // only radius + params["radius"] = radius + searchReq := integration.ConstructElementLevelSearchRequest("", collectionName, expr, + subFieldName, schemapb.DataType_FloatVector, []string{integration.StructArrayField}, + metric.L2, params, nq, dim, topk, roundDecimal) + searchResult, err := c.MilvusClient.Search(ctx, searchReq) + s.NoError(err) + if err := merr.Error(searchResult.GetStatus()); err != nil { + log.Warn("element-level range search with radius only failed", zap.Error(err)) + } + s.NoError(merr.Error(searchResult.GetStatus())) + + // radius + range_filter + params["range_filter"] = filter + searchReq = integration.ConstructElementLevelSearchRequest("", collectionName, expr, + subFieldName, schemapb.DataType_FloatVector, []string{integration.StructArrayField}, + metric.L2, params, nq, dim, topk, roundDecimal) + searchResult, err = c.MilvusClient.Search(ctx, searchReq) + s.NoError(err) + if err := merr.Error(searchResult.GetStatus()); err != nil { + log.Warn("element-level range search with radius+range_filter failed", zap.Error(err)) + } + s.NoError(merr.Error(searchResult.GetStatus())) + + // illegal: for L2, radius must be > range_filter + params["radius"] = filter + params["range_filter"] = radius + searchReq = integration.ConstructElementLevelSearchRequest("", collectionName, expr, + subFieldName, schemapb.DataType_FloatVector, []string{integration.StructArrayField}, + metric.L2, params, nq, dim, topk, roundDecimal) + searchResult, _ = c.MilvusClient.Search(ctx, searchReq) + s.Error(merr.Error(searchResult.GetStatus())) + + log.Info("=========================") + log.Info("=========================") + log.Info("TestRangeSearchElementLevelL2 succeed") + log.Info("=========================") + log.Info("=========================") +} diff --git a/tests/integration/util_query.go b/tests/integration/util_query.go index cc50df6fa0..c4416d7ca5 100644 --- a/tests/integration/util_query.go +++ b/tests/integration/util_query.go @@ -190,6 +190,51 @@ func ConstructEmbeddingListSearchRequest( return constructSearchRequest(dbName, collectionName, expr, vecField, true, vectorType, outputFields, metricType, params, nq, dim, topk, roundDecimal) } +// ConstructElementLevelSearchRequest builds a search request that runs +// element-level search against a VECTOR_ARRAY field. The placeholder carries +// a single query vector per nq row with element_level=true. +func ConstructElementLevelSearchRequest( + dbName, collectionName string, + expr string, + vecField string, + vectorType schemapb.DataType, + outputFields []string, + metricType string, + params map[string]any, + nq, dim int, topk, roundDecimal int, +) *milvuspb.SearchRequest { + b, err := json.Marshal(params) + if err != nil { + panic(err) + } + plg := constructPlaceholderGroup(nq, dim, vectorType, false) + for _, ph := range plg.Placeholders { + ph.ElementLevel = true + } + plgBs, err := proto.Marshal(plg) + if err != nil { + panic(err) + } + return &milvuspb.SearchRequest{ + DbName: dbName, + CollectionName: collectionName, + Dsl: expr, + SearchInput: &milvuspb.SearchRequest_PlaceholderGroup{ + PlaceholderGroup: plgBs, + }, + DslType: commonpb.DslType_BoolExprV1, + OutputFields: outputFields, + SearchParams: []*commonpb.KeyValuePair{ + {Key: common.MetricTypeKey, Value: metricType}, + {Key: proxy.ParamsKey, Value: string(b)}, + {Key: AnnsFieldKey, Value: vecField}, + {Key: common.TopKKey, Value: strconv.Itoa(topk)}, + {Key: RoundDecimalKey, Value: strconv.Itoa(roundDecimal)}, + }, + Nq: int64(nq), + } +} + func constructSearchRequest( dbName, collectionName string, expr string, diff --git a/tests/python_client/milvus_client/test_milvus_client_struct_array.py b/tests/python_client/milvus_client/test_milvus_client_struct_array.py index 9773313383..9d9e54b672 100644 --- a/tests/python_client/milvus_client/test_milvus_client_struct_array.py +++ b/tests/python_client/milvus_client/test_milvus_client_struct_array.py @@ -1,28 +1,32 @@ -import pytest -import numpy as np +import json +import os import random import time -import os -import json from pathlib import Path -from typing import List, Dict, Any -from deepdiff import DeepDiff +from typing import Any +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import pytest from base.client_v2_base import TestMilvusClientV2Base -from utils.util_log import test_log as log +from check.param_check import compare_lists_with_epsilon_ignore_dict_order from common import common_func as cf from common import common_type as ct from common.common_type import CaseLabel, CheckTasks -from check.param_check import compare_lists_with_epsilon_ignore_dict_order -from utils.util_pymilvus import * -from pymilvus import DataType, MilvusClient, AnnSearchRequest, RRFRanker, WeightedRanker -from pymilvus.client.embedding_list import EmbeddingList -from pymilvus.bulk_writer import bulk_import, get_import_progress, LocalBulkWriter, BulkFileType -import pyarrow as pa -import pyarrow.parquet as pq +from deepdiff import DeepDiff from minio import Minio from minio.error import S3Error - +from pymilvus import AnnSearchRequest, DataType, MilvusClient, RRFRanker, WeightedRanker +from pymilvus.bulk_writer import ( + BulkFileType, + LocalBulkWriter, + bulk_import, + get_import_progress, +) +from pymilvus.client.embedding_list import EmbeddingList +from utils.util_log import test_log as log +from utils.util_pymilvus import * # noqa: F403 prefix = "struct_array" epsilon = 0.001 @@ -64,17 +68,46 @@ EMB_LIST_INDEX_TYPES = list(EMB_LIST_INDEX_CONFIGS.keys()) # Supported vector types per emb list index type (for MaxSim metrics) EMB_LIST_VECTOR_TYPES = { - "HNSW": [DataType.FLOAT_VECTOR, DataType.FLOAT16_VECTOR, DataType.BFLOAT16_VECTOR, - DataType.INT8_VECTOR, DataType.BINARY_VECTOR], - "HNSW_SQ": [DataType.FLOAT_VECTOR, DataType.FLOAT16_VECTOR, DataType.BFLOAT16_VECTOR, - DataType.INT8_VECTOR], - "HNSW_PQ": [DataType.FLOAT_VECTOR, DataType.FLOAT16_VECTOR, DataType.BFLOAT16_VECTOR, - DataType.INT8_VECTOR], - "HNSW_PRQ": [DataType.FLOAT_VECTOR, DataType.FLOAT16_VECTOR, DataType.BFLOAT16_VECTOR, - DataType.INT8_VECTOR], - "IVF_FLAT": [DataType.FLOAT_VECTOR, DataType.FLOAT16_VECTOR, DataType.BFLOAT16_VECTOR], - "IVF_FLAT_CC": [DataType.FLOAT_VECTOR, DataType.FLOAT16_VECTOR, DataType.BFLOAT16_VECTOR], - "DISKANN": [DataType.FLOAT_VECTOR, DataType.FLOAT16_VECTOR, DataType.BFLOAT16_VECTOR], + "HNSW": [ + DataType.FLOAT_VECTOR, + DataType.FLOAT16_VECTOR, + DataType.BFLOAT16_VECTOR, + DataType.INT8_VECTOR, + DataType.BINARY_VECTOR, + ], + "HNSW_SQ": [ + DataType.FLOAT_VECTOR, + DataType.FLOAT16_VECTOR, + DataType.BFLOAT16_VECTOR, + DataType.INT8_VECTOR, + ], + "HNSW_PQ": [ + DataType.FLOAT_VECTOR, + DataType.FLOAT16_VECTOR, + DataType.BFLOAT16_VECTOR, + DataType.INT8_VECTOR, + ], + "HNSW_PRQ": [ + DataType.FLOAT_VECTOR, + DataType.FLOAT16_VECTOR, + DataType.BFLOAT16_VECTOR, + DataType.INT8_VECTOR, + ], + "IVF_FLAT": [ + DataType.FLOAT_VECTOR, + DataType.FLOAT16_VECTOR, + DataType.BFLOAT16_VECTOR, + ], + "IVF_FLAT_CC": [ + DataType.FLOAT_VECTOR, + DataType.FLOAT16_VECTOR, + DataType.BFLOAT16_VECTOR, + ], + "DISKANN": [ + DataType.FLOAT_VECTOR, + DataType.FLOAT16_VECTOR, + DataType.BFLOAT16_VECTOR, + ], } # Dim for emb list index tests (smaller for faster index building) @@ -91,7 +124,7 @@ class TestMilvusClientStructArrayBasic(TestMilvusClientV2Base): def generate_struct_array_data( self, num_rows: int, dim: int = default_dim, capacity: int = default_capacity - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """Generate test data for struct array""" data = [] for i in range(num_rows): @@ -149,12 +182,10 @@ class TestMilvusClientStructArrayBasic(TestMilvusClientV2Base): field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=dim, - mmap_enabled=mmap_enabled + mmap_enabled=mmap_enabled, ) else: - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=dim) # Create struct schema with optional mmap for sub-fields struct_schema = client.create_struct_field_schema() @@ -165,7 +196,7 @@ class TestMilvusClientStructArrayBasic(TestMilvusClientV2Base): "clip_embedding1", DataType.FLOAT_VECTOR, dim=dim, - mmap_enabled=subfield1_mmap + mmap_enabled=subfield1_mmap, ) else: struct_schema.add_field("clip_embedding1", DataType.FLOAT_VECTOR, dim=dim) @@ -175,7 +206,7 @@ class TestMilvusClientStructArrayBasic(TestMilvusClientV2Base): "clip_embedding2", DataType.FLOAT_VECTOR, dim=dim, - mmap_enabled=subfield2_mmap + mmap_enabled=subfield2_mmap, ) else: struct_schema.add_field("clip_embedding2", DataType.FLOAT_VECTOR, dim=dim) @@ -238,9 +269,7 @@ class TestMilvusClientStructArrayBasic(TestMilvusClientV2Base): # Create schema schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) # Create struct schema with only scalar fields struct_schema = client.create_struct_field_schema() @@ -290,9 +319,7 @@ class TestMilvusClientStructArrayBasic(TestMilvusClientV2Base): # Create schema schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) # Create struct schema with the specific scalar field type struct_schema = client.create_struct_field_schema() @@ -325,9 +352,7 @@ class TestMilvusClientStructArrayBasic(TestMilvusClientV2Base): # Create schema schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) # Create struct schema with all supported scalar field types struct_schema = client.create_struct_field_schema() @@ -471,9 +496,7 @@ class TestMilvusClientStructArrayBasic(TestMilvusClientV2Base): # Create schema with mixed field types schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) # Create struct schema with mixed types struct_schema = client.create_struct_field_schema() @@ -511,9 +534,7 @@ class TestMilvusClientStructArrayBasic(TestMilvusClientV2Base): # Create schema schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) # Create first struct schema struct_schema1 = client.create_struct_field_schema() @@ -561,9 +582,7 @@ class TestMilvusClientStructArrayBasic(TestMilvusClientV2Base): # Create schema schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) # Create first struct schema with common subfield names struct_schema1 = client.create_struct_field_schema() @@ -573,9 +592,7 @@ class TestMilvusClientStructArrayBasic(TestMilvusClientV2Base): # Create second struct schema with same subfield names but potentially different semantics struct_schema2 = client.create_struct_field_schema() - struct_schema2.add_field( - "embedding", DataType.FLOAT_VECTOR, dim=128 - ) # Same name + struct_schema2.add_field("embedding", DataType.FLOAT_VECTOR, dim=128) # Same name struct_schema2.add_field("label", DataType.VARCHAR, max_length=128) # Same name struct_schema2.add_field("score", DataType.FLOAT) # Same name @@ -621,9 +638,7 @@ class TestMilvusClientStructArrayBasic(TestMilvusClientV2Base): { "embedding": [random.random() for _ in range(128)], "label": f"text_label_{i}_{k}", - "score": float( - i * 100 + k + 0.2 - ), # Deterministic score, different from image + "score": float(i * 100 + k + 0.2), # Deterministic score, different from image } ) @@ -703,12 +718,8 @@ class TestMilvusClientStructArrayBasic(TestMilvusClientV2Base): expected_score = original_data["image_features"][idx]["score"] expected_label = original_data["image_features"][idx]["label"] - assert "score" in img_feat, ( - f"score missing in image_features for id {result_id}" - ) - assert "label" in img_feat, ( - f"label missing in image_features for id {result_id}" - ) + assert "score" in img_feat, f"score missing in image_features for id {result_id}" + assert "label" in img_feat, f"label missing in image_features for id {result_id}" assert abs(img_feat["score"] - expected_score) < 0.0001, ( f"image_features[{idx}].score mismatch: expected {expected_score}, got {img_feat['score']}" ) @@ -722,12 +733,8 @@ class TestMilvusClientStructArrayBasic(TestMilvusClientV2Base): expected_score = original_data["text_features"][idx]["score"] expected_label = original_data["text_features"][idx]["label"] - assert "score" in txt_feat, ( - f"score missing in text_features for id {result_id}" - ) - assert "label" in txt_feat, ( - f"label missing in text_features for id {result_id}" - ) + assert "score" in txt_feat, f"score missing in text_features for id {result_id}" + assert "label" in txt_feat, f"label missing in text_features for id {result_id}" assert abs(txt_feat["score"] - expected_score) < 0.0001, ( f"text_features[{idx}].score mismatch: expected {expected_score}, got {txt_feat['score']}" ) @@ -759,9 +766,7 @@ class TestMilvusClientStructArrayBasic(TestMilvusClientV2Base): hit_id = hit_data["id"] original_data = data_map[hit_id] - assert "image_features" in hit_data, ( - f"image_features missing for id {hit_id}" - ) + assert "image_features" in hit_data, f"image_features missing for id {hit_id}" # Verify each element in image_features array for idx, img_feat in enumerate(hit_data["image_features"]): @@ -998,14 +1003,19 @@ class TestMilvusClientStructArrayBasic(TestMilvusClientV2Base): assert check res, check = self.load_collection(client, collection_name) assert check - results, check = self.query(client, collection_name, filter="id >= 0", output_fields=["id", "clips"], limit=10) + results, check = self.query( + client, + collection_name, + filter="id >= 0", + output_fields=["id", "clips"], + limit=10, + ) assert check assert len(results) == 10 for result in results: assert "clips" in result assert len(result["clips"]) == 0 - @pytest.mark.tags(CaseLabel.L1) def test_insert_struct_array_different_lengths(self): """ @@ -1152,9 +1162,7 @@ class TestMilvusClientStructArrayIndex(TestMilvusClientV2Base): # Create schema schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=dim) # Create struct schema with vector field struct_schema = client.create_struct_field_schema() @@ -1280,9 +1288,7 @@ class TestMilvusClientStructArrayIndex(TestMilvusClientV2Base): # Create schema with multiple vector fields in struct schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) # Create struct schema with multiple vector fields struct_schema = client.create_struct_field_schema() @@ -1487,9 +1493,7 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): # Create schema schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=dim) # Create struct schema struct_schema = client.create_struct_field_schema() @@ -1585,9 +1589,7 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): # Create schema schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=dim) # Create struct schema struct_schema = client.create_struct_field_schema() @@ -1676,7 +1678,7 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): client = self._client() # Create collection with data and index - data = self.create_collection_with_index(client, collection_name) + self.create_collection_with_index(client, collection_name) # Create search vector search_vector = [random.random() for _ in range(default_dim)] @@ -1860,19 +1862,13 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): client = self._client() # Create collection - collection_name = cf.gen_unique_str( - f"{prefix}_metric_search_{metric_type.lower()}" - ) + collection_name = cf.gen_unique_str(f"{prefix}_metric_search_{metric_type.lower()}") schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) struct_schema = client.create_struct_field_schema() - struct_schema.add_field( - "clip_embedding1", DataType.FLOAT_VECTOR, dim=default_dim - ) + struct_schema.add_field("clip_embedding1", DataType.FLOAT_VECTOR, dim=default_dim) schema.add_field( "clips", datatype=DataType.ARRAY, @@ -1887,9 +1883,7 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): # Insert data data = [] for i in range(20): - struct_array = [ - {"clip_embedding1": [random.random() for _ in range(default_dim)]} - ] + struct_array = [{"clip_embedding1": [random.random() for _ in range(default_dim)]}] row = { "id": i, "normal_vector": [random.random() for _ in range(default_dim)], @@ -1983,15 +1977,20 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): self.create_collection_with_index(client, collection_name) # Search using EmbeddingList - error = {ct.err_code: 999, - ct.err_msg: "array of vector is not supported for search by IDs"} - self.search(client, - collection_name, - ids=[0, 1], - anns_field="clips[clip_embedding1]", - search_params={"metric_type": "MAX_SIM_COSINE"}, - limit=10, - check_task=CheckTasks.err_res, check_items=error) + error = { + ct.err_code: 999, + ct.err_msg: "array of vector is not supported for search by IDs", + } + self.search( + client, + collection_name, + ids=[0, 1], + anns_field="clips[clip_embedding1]", + search_params={"metric_type": "MAX_SIM_COSINE"}, + limit=10, + check_task=CheckTasks.err_res, + check_items=error, + ) @pytest.mark.tags(CaseLabel.L1) @pytest.mark.parametrize("retrieval_ann_ratio", [1.0, 3.0, 5.0, 10.0]) @@ -2021,7 +2020,7 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): anns_field="clips[clip_embedding1]", search_params={ "metric_type": "MAX_SIM_COSINE", - "params": {"retrieval_ann_ratio": retrieval_ann_ratio} + "params": {"retrieval_ann_ratio": retrieval_ann_ratio}, }, limit=10, ) @@ -2069,15 +2068,15 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): min_patches = 100 # Min patches per document max_patches = 300 # Max patches per document - log.info(f"Creating collection with {nb} docs, {num_query_vectors} query vectors, " - f"{min_patches}-{max_patches} patches per doc") + log.info( + f"Creating collection with {nb} docs, {num_query_vectors} query vectors, " + f"{min_patches}-{max_patches} patches per doc" + ) # Create schema schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=dim) # Create struct schema struct_schema = client.create_struct_field_schema() @@ -2126,7 +2125,7 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): doc_embeddings[i] = np.array(embeddings_list) self.insert(client, collection_name, data) - log.info(f"Inserted batch {batch_start//batch_size + 1}/{(nb + batch_size - 1)//batch_size}") + log.info(f"Inserted batch {batch_start // batch_size + 1}/{(nb + batch_size - 1) // batch_size}") # Create indexes index_params = client.prepare_index_params() @@ -2184,13 +2183,13 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): log.info(f"Ground truth top-{limit} with scores:") for i, (doc_id, score) in enumerate(ground_truth_scores[:limit]): num_patches = len(doc_embeddings[doc_id]) - log.info(f" GT rank {i+1}: id={doc_id}, score={score:.6f}, num_patches={num_patches}") + log.info(f" GT rank {i + 1}: id={doc_id}, score={score:.6f}, num_patches={num_patches}") # Search with different retrieval_ann_ratio values retrieval_ann_ratios = [0.1, 1.0, 3.0, 5.0] recall_results = {} # Track recall for each ratio for retrieval_ann_ratio in retrieval_ann_ratios: - log.info(f"\n{'='*50}") + log.info(f"\n{'=' * 50}") log.info(f"Testing retrieval_ann_ratio={retrieval_ann_ratio}") results, _ = self.search( @@ -2200,7 +2199,7 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): anns_field="clips[clip_embedding1]", search_params={ "metric_type": "MAX_SIM_COSINE", - "params": {"retrieval_ann_ratio": retrieval_ann_ratio} + "params": {"retrieval_ann_ratio": retrieval_ann_ratio}, }, limit=limit, ) @@ -2214,19 +2213,28 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): recall = recall_hits / len(ground_truth_ids) recall_results[retrieval_ann_ratio] = recall - log.info(f"retrieval_ann_ratio={retrieval_ann_ratio}, recall={recall:.4f}, " - f"recall_hits={recall_hits}/{limit}") + log.info( + f"retrieval_ann_ratio={retrieval_ann_ratio}, recall={recall:.4f}, recall_hits={recall_hits}/{limit}" + ) log.info(f"Milvus result IDs: {sorted(milvus_result_ids)}") # Log detailed comparison for Milvus results - log.info(f"Milvus results with ground truth comparison:") + log.info("Milvus results with ground truth comparison:") gt_ranks_for_milvus = [] for i, hit in enumerate(results[0][:limit]): - gt_score = next((s for doc_id, s in ground_truth_scores if doc_id == hit["id"]), None) - gt_rank = next((idx+1 for idx, (doc_id, _) in enumerate(ground_truth_scores) if doc_id == hit["id"]), None) + gt_score = next( + (s for doc_id, s in ground_truth_scores if doc_id == hit["id"]), + None, + ) + gt_rank = next( + (idx + 1 for idx, (doc_id, _) in enumerate(ground_truth_scores) if doc_id == hit["id"]), + None, + ) gt_ranks_for_milvus.append(gt_rank) - log.info(f" Milvus rank {i+1}: id={hit['id']}, distance={hit['distance']:.6f}, " - f"gt_score={gt_score:.6f}, gt_rank={gt_rank}") + log.info( + f" Milvus rank {i + 1}: id={hit['id']}, distance={hit['distance']:.6f}, " + f"gt_score={gt_score:.6f}, gt_rank={gt_rank}" + ) # Calculate recall at different K values for k in [1, 5, 10]: @@ -2247,17 +2255,18 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): for i in range(len(retrieval_ann_ratios) - 1): ratio_curr = retrieval_ann_ratios[i] ratio_next = retrieval_ann_ratios[i + 1] - assert recall_results[ratio_next] >= recall_results[ratio_curr], \ - f"Recall should increase with higher retrieval_ann_ratio: " \ - f"ratio {ratio_curr} has recall {recall_results[ratio_curr]}, " \ + assert recall_results[ratio_next] >= recall_results[ratio_curr], ( + f"Recall should increase with higher retrieval_ann_ratio: " + f"ratio {ratio_curr} has recall {recall_results[ratio_curr]}, " f"but ratio {ratio_next} has recall {recall_results[ratio_next]}" + ) # Verify that recall >= 0.8 when retrieval_ann_ratio >= 3 for ratio, recall in recall_results.items(): if ratio >= 3: - assert recall >= 0.8, \ - f"Recall should be >= 0.8 when retrieval_ann_ratio >= 3, " \ - f"but ratio {ratio} has recall {recall}" + assert recall >= 0.8, ( + f"Recall should be >= 0.8 when retrieval_ann_ratio >= 3, but ratio {ratio} has recall {recall}" + ) @pytest.mark.tags(CaseLabel.L1) @pytest.mark.parametrize("index_type", EMB_LIST_INDEX_TYPES) @@ -2267,9 +2276,7 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): method: insert (flushed + growing) → index → load → search → upsert → delete → search → verify expected: all CRUD operations and search work correctly for each index type """ - collection_name = cf.gen_unique_str( - f"{prefix}_search_{index_type.lower()}" - ) + collection_name = cf.gen_unique_str(f"{prefix}_search_{index_type.lower()}") client = self._client() config = EMB_LIST_INDEX_CONFIGS[index_type] @@ -2277,9 +2284,7 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): # Create collection with full schema schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=EMB_LIST_DIM - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=EMB_LIST_DIM) struct_schema = client.create_struct_field_schema() struct_schema.add_field("clip_embedding1", DataType.FLOAT_VECTOR, dim=EMB_LIST_DIM) @@ -2430,7 +2435,9 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): assert check results, check = self.query( - client, collection_name, filter="id < 10", + client, + collection_name, + filter="id < 10", output_fields=["id", "category", "clips"], ) assert check @@ -2448,9 +2455,7 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): res, check = self.flush(client, collection_name) assert check - results, check = self.query( - client, collection_name, filter="id >= 0", output_fields=["id"] - ) + results, check = self.query(client, collection_name, filter="id >= 0", output_fields=["id"]) assert check remaining_ids = {r["id"] for r in results} for del_id in delete_ids: @@ -2478,8 +2483,12 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): @pytest.mark.tags(CaseLabel.L2) @pytest.mark.parametrize( "vector_type", - [DataType.FLOAT16_VECTOR, DataType.BFLOAT16_VECTOR, DataType.INT8_VECTOR, - DataType.BINARY_VECTOR], + [ + DataType.FLOAT16_VECTOR, + DataType.BFLOAT16_VECTOR, + DataType.INT8_VECTOR, + DataType.BINARY_VECTOR, + ], ) def test_search_emb_list_with_different_vector_types(self, vector_type): """ @@ -2500,9 +2509,7 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): # Create schema with specified vector type schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=EMB_LIST_DIM - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=EMB_LIST_DIM) struct_schema = client.create_struct_field_schema() struct_schema.add_field("clip_embedding1", vector_type, dim=EMB_LIST_DIM) @@ -2614,9 +2621,7 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): [np.frombuffer(v, dtype=np.uint8) if isinstance(v, bytes) else v for v in search_vecs] ) else: - embedding_list = EmbeddingList( - [np.array(v) if not isinstance(v, np.ndarray) else v for v in search_vecs] - ) + embedding_list = EmbeddingList([np.array(v) if not isinstance(v, np.ndarray) else v for v in search_vecs]) results, check = self.search( client, @@ -2659,7 +2664,9 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): assert check results, check = self.query( - client, collection_name, filter="id < 10", + client, + collection_name, + filter="id < 10", output_fields=["id", "category", "clips"], ) assert check @@ -2672,18 +2679,14 @@ class TestMilvusClientStructArraySearch(TestMilvusClientV2Base): delete_flushed_ids = [10, 11, 12, 13, 14] delete_growing_ids = list(range(nb_flushed, nb_flushed + 3)) all_delete_ids = delete_flushed_ids + delete_growing_ids - res, check = self.delete( - client, collection_name, filter=f"id in {all_delete_ids}" - ) + res, check = self.delete(client, collection_name, filter=f"id in {all_delete_ids}") assert check # Verify deletion res, check = self.flush(client, collection_name) assert check - results, check = self.query( - client, collection_name, filter="id >= 0", output_fields=["id"] - ) + results, check = self.query(client, collection_name, filter="id >= 0", output_fields=["id"]) assert check remaining_ids = {r["id"] for r in results} for del_id in all_delete_ids: @@ -2721,20 +2724,14 @@ class TestMilvusClientStructArrayHybridSearch(TestMilvusClientV2Base): # Create schema with both normal vector and struct array schema = client.create_schema(auto_id=False, enable_dynamic_field=False) - schema.add_field( - field_name="pk", datatype=DataType.VARCHAR, is_primary=True, max_length=100 - ) + schema.add_field(field_name="pk", datatype=DataType.VARCHAR, is_primary=True, max_length=100) schema.add_field(field_name="random", datatype=DataType.DOUBLE) - schema.add_field( - field_name="embeddings", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="embeddings", datatype=DataType.FLOAT_VECTOR, dim=default_dim) # Create struct schema with vector field struct_schema = client.create_struct_field_schema() struct_schema.add_field("clip_str", DataType.VARCHAR, max_length=65535) - struct_schema.add_field( - "clip_embedding", DataType.FLOAT_VECTOR, dim=default_dim - ) + struct_schema.add_field("clip_embedding", DataType.FLOAT_VECTOR, dim=default_dim) schema.add_field( "clips", @@ -2759,9 +2756,7 @@ class TestMilvusClientStructArrayHybridSearch(TestMilvusClientV2Base): params=INDEX_PARAMS, ) - res, check = self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + res, check = self.create_collection(client, collection_name, schema=schema, index_params=index_params) assert check # Insert data @@ -2848,18 +2843,12 @@ class TestMilvusClientStructArrayHybridSearch(TestMilvusClientV2Base): # Create schema schema = client.create_schema(auto_id=False, enable_dynamic_field=False) - schema.add_field( - field_name="pk", datatype=DataType.VARCHAR, is_primary=True, max_length=100 - ) + schema.add_field(field_name="pk", datatype=DataType.VARCHAR, is_primary=True, max_length=100) schema.add_field(field_name="random", datatype=DataType.DOUBLE) - schema.add_field( - field_name="embeddings", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="embeddings", datatype=DataType.FLOAT_VECTOR, dim=default_dim) struct_schema = client.create_struct_field_schema() - struct_schema.add_field( - "clip_embedding", DataType.FLOAT_VECTOR, dim=default_dim - ) + struct_schema.add_field("clip_embedding", DataType.FLOAT_VECTOR, dim=default_dim) schema.add_field( "clips", datatype=DataType.ARRAY, @@ -2883,9 +2872,7 @@ class TestMilvusClientStructArrayHybridSearch(TestMilvusClientV2Base): params=INDEX_PARAMS, ) - res, check = self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + res, check = self.create_collection(client, collection_name, schema=schema, index_params=index_params) assert check # Insert data @@ -2952,9 +2939,7 @@ class TestMilvusClientStructArrayHybridSearch(TestMilvusClientV2Base): assert len(hybrid_res) > 0 assert len(hybrid_res[0]) > 0 - log.info( - f"Hybrid search with WeightedRanker returned {len(hybrid_res[0])} results" - ) + log.info(f"Hybrid search with WeightedRanker returned {len(hybrid_res[0])} results") @pytest.mark.tags(CaseLabel.L1) def test_hybrid_search_multiple_struct_array_vectors(self): @@ -2968,18 +2953,12 @@ class TestMilvusClientStructArrayHybridSearch(TestMilvusClientV2Base): # Create schema with both struct array vector and normal vector schema = client.create_schema(auto_id=False, enable_dynamic_field=False) - schema.add_field( - field_name="pk", datatype=DataType.VARCHAR, is_primary=True, max_length=100 - ) + schema.add_field(field_name="pk", datatype=DataType.VARCHAR, is_primary=True, max_length=100) schema.add_field(field_name="random", datatype=DataType.DOUBLE) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) struct_schema = client.create_struct_field_schema() - struct_schema.add_field( - "clip_embedding", DataType.FLOAT_VECTOR, dim=default_dim - ) + struct_schema.add_field("clip_embedding", DataType.FLOAT_VECTOR, dim=default_dim) schema.add_field( "clips", datatype=DataType.ARRAY, @@ -3003,9 +2982,7 @@ class TestMilvusClientStructArrayHybridSearch(TestMilvusClientV2Base): params=INDEX_PARAMS, ) - res, check = self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + res, check = self.create_collection(client, collection_name, schema=schema, index_params=index_params) assert check # Insert data @@ -3071,14 +3048,10 @@ class TestMilvusClientStructArrayHybridSearch(TestMilvusClientV2Base): assert len(hybrid_res) > 0 assert len(hybrid_res[0]) > 0 - log.info( - f"Hybrid search with multiple vector sources returned {len(hybrid_res[0])} results" - ) + log.info(f"Hybrid search with multiple vector sources returned {len(hybrid_res[0])} results") @pytest.mark.tags(CaseLabel.L2) - @pytest.mark.parametrize( - "metric_type", ["MAX_SIM_L2", "MAX_SIM_IP", "MAX_SIM_COSINE"] - ) + @pytest.mark.parametrize("metric_type", ["MAX_SIM_L2", "MAX_SIM_IP", "MAX_SIM_COSINE"]) def test_hybrid_search_struct_array_different_metrics(self, metric_type): """ target: test hybrid search with different metric types for struct array @@ -3090,17 +3063,11 @@ class TestMilvusClientStructArrayHybridSearch(TestMilvusClientV2Base): # Create schema schema = client.create_schema(auto_id=False, enable_dynamic_field=False) - schema.add_field( - field_name="pk", datatype=DataType.VARCHAR, is_primary=True, max_length=100 - ) - schema.add_field( - field_name="embeddings", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="pk", datatype=DataType.VARCHAR, is_primary=True, max_length=100) + schema.add_field(field_name="embeddings", datatype=DataType.FLOAT_VECTOR, dim=default_dim) struct_schema = client.create_struct_field_schema() - struct_schema.add_field( - "clip_embedding", DataType.FLOAT_VECTOR, dim=default_dim - ) + struct_schema.add_field("clip_embedding", DataType.FLOAT_VECTOR, dim=default_dim) schema.add_field( "clips", datatype=DataType.ARRAY, @@ -3124,9 +3091,7 @@ class TestMilvusClientStructArrayHybridSearch(TestMilvusClientV2Base): params=INDEX_PARAMS, ) - res, check = self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + res, check = self.create_collection(client, collection_name, schema=schema, index_params=index_params) assert check # Insert data @@ -3179,9 +3144,7 @@ class TestMilvusClientStructArrayHybridSearch(TestMilvusClientV2Base): ) # Perform hybrid search - hybrid_res = client.hybrid_search( - collection_name, req_list, RRFRanker(), default_limit - ) + hybrid_res = client.hybrid_search(collection_name, req_list, RRFRanker(), default_limit) assert len(hybrid_res) > 0 assert len(hybrid_res[0]) > 0 @@ -3190,22 +3153,16 @@ class TestMilvusClientStructArrayHybridSearch(TestMilvusClientV2Base): class TestMilvusClientStructArrayQuery(TestMilvusClientV2Base): """Test case of struct array query functionality""" - def create_collection_with_data( - self, client: MilvusClient, collection_name: str, nb: int = 50 - ): + def create_collection_with_data(self, client: MilvusClient, collection_name: str, nb: int = 50): """Create collection with struct array and insert data""" # Create schema schema = client.create_schema(auto_id=False, enable_dynamic_field=True) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) # Create struct schema struct_schema = client.create_struct_field_schema() - struct_schema.add_field( - "clip_embedding1", DataType.FLOAT_VECTOR, dim=default_dim - ) + struct_schema.add_field("clip_embedding1", DataType.FLOAT_VECTOR, dim=default_dim) struct_schema.add_field("scalar_field", DataType.INT64) struct_schema.add_field("category", DataType.VARCHAR, max_length=128) struct_schema.add_field("score", DataType.FLOAT) @@ -3293,7 +3250,7 @@ class TestMilvusClientStructArrayQuery(TestMilvusClientV2Base): client = self._client() # Create collection with data - data = self.create_collection_with_data(client, collection_name) + self.create_collection_with_data(client, collection_name) # Query all data results, check = self.query(client, collection_name, filter="id >= 0", limit=10) @@ -3366,9 +3323,7 @@ class TestMilvusClientStructArrayQuery(TestMilvusClientV2Base): for result in results: struct_array = result["clips"] # At least one element in the struct array should satisfy the filter - has_matching_element = any( - elem["scalar_field"] < 50 for elem in struct_array - ) + has_matching_element = any(elem["scalar_field"] < 50 for elem in struct_array) assert has_matching_element @pytest.mark.tags(CaseLabel.L1) @@ -3414,12 +3369,10 @@ class TestMilvusClientStructArrayQuery(TestMilvusClientV2Base): client = self._client() # Create collection with data - data = self.create_collection_with_data(client, collection_name, nb=20) + self.create_collection_with_data(client, collection_name, nb=20) # Count all records - count_result = client.query( - collection_name=collection_name, filter="", output_fields=["count(*)"] - ) + count_result = client.query(collection_name=collection_name, filter="", output_fields=["count(*)"]) assert len(count_result) == 1 assert count_result[0]["count(*)"] == 20 @@ -3475,13 +3428,14 @@ class TestMilvusClientStructArrayCRUD(TestMilvusClientV2Base): schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim, nullable=True + field_name="normal_vector", + datatype=DataType.FLOAT_VECTOR, + dim=default_dim, + nullable=True, ) struct_schema = client.create_struct_field_schema() - struct_schema.add_field( - "clip_embedding1", DataType.FLOAT_VECTOR, dim=default_dim - ) + struct_schema.add_field("clip_embedding1", DataType.FLOAT_VECTOR, dim=default_dim) struct_schema.add_field("scalar_field", DataType.INT64) struct_schema.add_field("label", DataType.VARCHAR, max_length=128) @@ -3957,16 +3911,13 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base): ) else: struct_schema.add_field("unsupported_field", unsupported_type) - assert False, ( - f"Expected ParamError when adding {unsupported_type} field to struct schema" - ) + assert False, f"Expected ParamError when adding {unsupported_type} field to struct schema" except Exception as e: # Verify the error message indicates the type is not supported error_msg = str(e) # Different error messages for different types assert ( - "Struct field schema does not support Array, ArrayOfVector or Struct" - in error_msg + "Struct field schema does not support Array, ArrayOfVector or Struct" in error_msg or "is not supported" in error_msg or "not supported for fields in struct" in error_msg ) @@ -3985,9 +3936,7 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base): # Create schema schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) struct_schema = client.create_struct_field_schema() struct_schema.add_field("sparse_vector_field", DataType.SPARSE_FLOAT_VECTOR) @@ -4010,7 +3959,6 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base): check_items=error, ) - @pytest.mark.tags(CaseLabel.L2) def test_struct_with_json_field(self): """ @@ -4025,9 +3973,7 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base): # Create schema schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) struct_schema = client.create_struct_field_schema() struct_schema.add_field("json_field", DataType.JSON) @@ -4064,9 +4010,7 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base): # Create schema schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) struct_schema = client.create_struct_field_schema() struct_schema.add_field("geometry_field", DataType.GEOMETRY) @@ -4103,9 +4047,7 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base): # Create collection with small capacity schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) struct_schema = client.create_struct_field_schema() struct_schema.add_field("scalar_field", DataType.INT64) @@ -4157,18 +4099,12 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base): schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) struct_schema = client.create_struct_field_schema() - struct_schema.add_field( - "clip_embedding1", DataType.FLOAT_VECTOR, dim=0 - ) # Invalid + struct_schema.add_field("clip_embedding1", DataType.FLOAT_VECTOR, dim=0) # Invalid - struct_schema.add_field( - "clip_embedding2", DataType.FLOAT_VECTOR, dim=-1 - ) # Invalid + struct_schema.add_field("clip_embedding2", DataType.FLOAT_VECTOR, dim=-1) # Invalid schema.add_field( "struct_array", datatype=DataType.ARRAY, @@ -4200,9 +4136,7 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base): # Create collection schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) struct_schema = client.create_struct_field_schema() struct_schema.add_field("field1", DataType.INT64) @@ -4277,9 +4211,7 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base): # Create schema schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) # Try to create struct with unsupported vector type (sparse vector) struct_schema = client.create_struct_field_schema() @@ -4294,7 +4226,10 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base): ) # Should fail - sparse vectors are not supported in struct array - error = {ct.err_code: 65535, ct.err_msg: "only fixed dimension vector types are supported"} + error = { + ct.err_code: 65535, + ct.err_msg: "only fixed dimension vector types are supported", + } self.create_collection( client, collection_name, @@ -4311,15 +4246,16 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base): collection_name = cf.gen_collection_name_by_testcase_name() schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) struct_schema = client.create_struct_field_schema() - struct_schema.add_field("clip_embedding1", DataType.FLOAT_VECTOR, dim=default_dim, - nullable=("clip_embedding1" == nullable_field)) - struct_schema.add_field("scalar_field", DataType.INT64, - nullable=("scalar_field" == nullable_field)) + struct_schema.add_field( + "clip_embedding1", + DataType.FLOAT_VECTOR, + dim=default_dim, + nullable=("clip_embedding1" == nullable_field), + ) + struct_schema.add_field("scalar_field", DataType.INT64, nullable=("scalar_field" == nullable_field)) struct_schema.add_field("label", DataType.VARCHAR, max_length=128) schema.add_field( @@ -4329,10 +4265,17 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base): struct_schema=struct_schema, max_capacity=100, ) - error = {ct.err_code: 999, - ct.err_msg: f"nullable is not supported for fields in struct array now, fieldName = {nullable_field}"} - self.create_collection(client, collection_name, schema=schema, - check_task=CheckTasks.err_res, check_items=error) + error = { + ct.err_code: 999, + ct.err_msg: f"nullable is not supported for fields in struct array now, fieldName = {nullable_field}", + } + self.create_collection( + client, + collection_name, + schema=schema, + check_task=CheckTasks.err_res, + check_items=error, + ) @pytest.mark.tags(CaseLabel.L0) def test_struct_array_range_search_not_supported(self): @@ -4348,13 +4291,9 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base): # Create collection with data and index using COSINE metric schema = client.create_schema(auto_id=False, enable_dynamic_field=False) schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim) struct_schema = client.create_struct_field_schema() - struct_schema.add_field( - "clip_embedding1", DataType.FLOAT_VECTOR, dim=default_dim - ) + struct_schema.add_field("clip_embedding1", DataType.FLOAT_VECTOR, dim=default_dim) schema.add_field( "clips", datatype=DataType.ARRAY, @@ -4375,18 +4314,14 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base): metric_type="MAX_SIM_COSINE", params=INDEX_PARAMS, ) - res, check = self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + res, check = self.create_collection(client, collection_name, schema=schema, index_params=index_params) assert check nb = 3000 data = [] for i in range(nb): struct_array = [] for j in range(random.randint(1, 10)): - struct_array.append( - {"clip_embedding1": [random.random() for _ in range(default_dim)]} - ) + struct_array.append({"clip_embedding1": [random.random() for _ in range(default_dim)]}) tmp_data = { "id": i, "normal_vector": [random.random() for _ in range(default_dim)], @@ -4414,11 +4349,7 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base): assert len(initial_results[0]) > 0 # Get distances and sort them - distances = [ - hit.get("distance") - for hit in initial_results[0] - if hit.get("distance") is not None - ] + distances = [hit.get("distance") for hit in initial_results[0] if hit.get("distance") is not None] distances.sort() # Select range from sorted distances (e.g., from 20th percentile to 80th percentile) @@ -4432,8 +4363,8 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base): # Step 2: Perform range search with selected radius and range_filter # For COSINE: radius < distance <= range_filter error = { - ct.err_code: 65535, - ct.err_msg: "range search is not supported for vector array", + ct.err_code: 1100, + ct.err_msg: "range search is not supported for multi-search-multi on embedding list fields", } self.search( client, @@ -4468,10 +4399,7 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): self.minio_endpoint = f"{minio_host}:9000" def gen_file_with_local_bulk_writer( - self, - schema, - data: List[Dict[str, Any]], - file_type: str = "PARQUET" + self, schema, data: list[dict[str, Any]], file_type: str = "PARQUET" ) -> tuple[str, dict]: """ Generate import file using LocalBulkWriter from insert-format data @@ -4492,7 +4420,7 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): schema=schema, local_path=self.LOCAL_FILES_PATH, segment_size=512 * 1024 * 1024, # 512MB - file_type=bulk_file_type + file_type=bulk_file_type, ) log.info(f"Creating {file_type} file using LocalBulkWriter with {len(data)} rows") @@ -4516,10 +4444,10 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): return batch_files, { "id": id_arr, "float_vector": float_vector_arr, - "struct_array": struct_arr + "struct_array": struct_arr, } - def upload_to_minio(self, local_file_path: str) -> List[List[str]]: + def upload_to_minio(self, local_file_path: str) -> list[list[str]]: """ Upload parquet file to MinIO @@ -4537,7 +4465,7 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): endpoint=self.minio_endpoint, access_key=self.MINIO_ACCESS_KEY, secret_key=self.MINIO_SECRET_KEY, - secure=False + secure=False, ) # Check if bucket exists @@ -4555,7 +4483,7 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): except S3Error as e: raise Exception(f"Failed to connect MinIO server {self.minio_endpoint}, error: {e}") - def call_bulkinsert(self, collection_name: str, batch_files: List[List[str]]): + def call_bulkinsert(self, collection_name: str, batch_files: list[list[str]]): """ Call bulk import API and wait for completion @@ -4572,7 +4500,7 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): files=batch_files, ) - job_id = resp.json()['data']['jobId'] + job_id = resp.json()["data"]["jobId"] log.info(f"Bulk import job created, job_id: {job_id}") # Wait for import to complete @@ -4582,15 +4510,15 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): time.sleep(5) resp = get_import_progress(url=url, job_id=job_id) - state = resp.json()['data']['state'] - progress = resp.json()['data']['progress'] + state = resp.json()["data"]["state"] + progress = resp.json()["data"]["progress"] log.info(f"Import job {job_id} - state: {state}, progress: {progress}%") if state == "Importing": continue elif state == "Failed": - reason = resp.json()['data']['reason'] + reason = resp.json()["data"]["reason"] raise Exception(f"Bulk import job {job_id} failed: {reason}") elif state == "Completed" and progress == 100: log.info(f"Bulk import job {job_id} completed successfully") @@ -4619,7 +4547,7 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): collection_name=collection_name, filter=f"id >= {min(original_data['id'])}", output_fields=["*"], - limit=num_rows + 100 # Add buffer to ensure all data is retrieved + limit=num_rows + 100, # Add buffer to ensure all data is retrieved ) log.info(f"Query returned {len(results)} rows") @@ -4632,11 +4560,13 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): original_rows = [] for i in range(num_rows): row_id = int(original_data["id"][i]) - original_rows.append({ - "id": row_id, - "float_vector": [float(x) for x in original_data["float_vector"][i]], - "struct_array": original_data["struct_array"][i] - }) + original_rows.append( + { + "id": row_id, + "float_vector": [float(x) for x in original_data["float_vector"][i]], + "struct_array": original_data["struct_array"][i], + } + ) # Convert query results to comparable format query_rows_formatted = [] @@ -4644,20 +4574,23 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): formatted_row = { "id": row["id"], "float_vector": [float(x) for x in row["float_vector"]], - "struct_array": row["struct_array"] + "struct_array": row["struct_array"], } query_rows_formatted.append(formatted_row) # Use compare_lists_with_epsilon_ignore_dict_order for comparison # This function handles floating-point tolerance and dict order - is_equal = compare_lists_with_epsilon_ignore_dict_order( - original_rows, - query_rows_formatted, - epsilon=epsilon - ) + is_equal = compare_lists_with_epsilon_ignore_dict_order(original_rows, query_rows_formatted, epsilon=epsilon) if not is_equal: - deepdiff = DeepDiff(original_rows, query_rows_formatted, ignore_order=True, significant_digits=3), + deepdiff = ( + DeepDiff( + original_rows, + query_rows_formatted, + ignore_order=True, + significant_digits=3, + ), + ) log.info(f"DeepDiff: {deepdiff}") assert False, "Data verification failed: original data and query results do not match" @@ -4725,7 +4658,7 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): field_name="float_vector", index_type="HNSW", metric_type="COSINE", - params={"M": 16, "efConstruction": 200} + params={"M": 16, "efConstruction": 200}, ) # Index for vector field inside struct array @@ -4734,15 +4667,11 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): field_name="struct_array[struct_float_vec]", index_type="HNSW", metric_type="MAX_SIM_COSINE", - params={"M": 16, "efConstruction": 200} + params={"M": 16, "efConstruction": 200}, ) # Create collection - client.create_collection( - collection_name=c_name, - schema=schema, - index_params=index_params - ) + client.create_collection(collection_name=c_name, schema=schema, index_params=index_params) log.info(f"Collection '{c_name}' created") # Step 2: Generate parquet file with all supported types @@ -4760,7 +4689,7 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): # Step 5: Verify number of entities stats = client.get_collection_stats(collection_name=c_name) - num_entities = stats['row_count'] + num_entities = stats["row_count"] log.info(f"Collection entities: {num_entities}") assert num_entities == entities, f"Expected {entities} entities, got {num_entities}" @@ -4832,7 +4761,7 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): field_name="float_vector", index_type="HNSW", metric_type="COSINE", - params={"M": 16, "efConstruction": 200} + params={"M": 16, "efConstruction": 200}, ) # Index for vector field inside struct array @@ -4841,15 +4770,11 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): field_name="struct_array[struct_float_vec]", index_type="HNSW", metric_type="MAX_SIM_COSINE", - params={"M": 16, "efConstruction": 200} + params={"M": 16, "efConstruction": 200}, ) # Create collection - client.create_collection( - collection_name=c_name, - schema=schema, - index_params=index_params - ) + client.create_collection(collection_name=c_name, schema=schema, index_params=index_params) log.info(f"Collection '{c_name}' created") # Step 2: Generate JSON file with all supported types @@ -4867,7 +4792,7 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): # Step 5: Verify number of entities stats = client.get_collection_stats(collection_name=c_name) - num_entities = stats['row_count'] + num_entities = stats["row_count"] log.info(f"Collection entities: {num_entities}") assert num_entities == entities, f"Expected {entities} entities, got {num_entities}" @@ -4942,7 +4867,7 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): field_name="float_vector", index_type="HNSW", metric_type="COSINE", - params={"M": 16, "efConstruction": 200} + params={"M": 16, "efConstruction": 200}, ) # Index for vector field inside struct array @@ -4950,15 +4875,11 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): field_name="struct_array[struct_float_vec]", index_type="HNSW", metric_type="MAX_SIM_COSINE", - params={"M": 16, "efConstruction": 200} + params={"M": 16, "efConstruction": 200}, ) # Create collection - client.create_collection( - collection_name=c_name, - schema=schema, - index_params=index_params - ) + client.create_collection(collection_name=c_name, schema=schema, index_params=index_params) log.info(f"Collection '{c_name}' created") # Step 2: Generate insert-format data with all supported types @@ -4984,14 +4905,14 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): "struct_float": random.random() * 100, "struct_double": random.random() * 1000, "struct_bool": random.choice([True, False]), - "struct_float_vec": [random.random() for _ in range(dim)] + "struct_float_vec": [random.random() for _ in range(dim)], } struct_list.append(struct_obj) row = { "id": i, "float_vector": [random.random() for _ in range(dim)], - "struct_array": struct_list + "struct_array": struct_list, } insert_data.append(row) log.info(f"Generated {empty_array_count} rows with empty struct_array") @@ -4999,9 +4920,7 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): # Step 3: Use LocalBulkWriter to convert data to import files log.info(f"Using LocalBulkWriter to generate {file_type} files") batch_files, original_data = self.gen_file_with_local_bulk_writer( - schema=schema, - data=insert_data, - file_type=file_type + schema=schema, data=insert_data, file_type=file_type ) # Step 4: Upload generated files to MinIO @@ -5023,7 +4942,7 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): # Step 6: Verify number of entities stats = client.get_collection_stats(collection_name=c_name) - num_entities = stats['row_count'] + num_entities = stats["row_count"] log.info(f"Collection entities: {num_entities}") assert num_entities == entities, f"Expected {entities} entities, got {num_entities}" @@ -5073,32 +4992,34 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): "struct_float": np.float32(random.random() * 100), "struct_double": np.float64(random.random() * 1000), "struct_bool": bool(random.choice([True, False])), - "struct_float_vec": [random.random() for _ in range(dim)] + "struct_float_vec": [random.random() for _ in range(dim)], } struct_list.append(struct_obj) struct_arr.append(struct_list) # Define PyArrow schema for struct field with all types - struct_type = pa.struct([ - pa.field("struct_varchar", pa.string()), - pa.field("struct_int8", pa.int8()), - pa.field("struct_int16", pa.int16()), - pa.field("struct_int32", pa.int32()), - pa.field("struct_int64", pa.int64()), - pa.field("struct_float", pa.float32()), - pa.field("struct_double", pa.float64()), - pa.field("struct_bool", pa.bool_()), - pa.field("struct_float_vec", pa.list_(pa.float32())) - ]) + struct_type = pa.struct( + [ + pa.field("struct_varchar", pa.string()), + pa.field("struct_int8", pa.int8()), + pa.field("struct_int16", pa.int16()), + pa.field("struct_int32", pa.int32()), + pa.field("struct_int64", pa.int64()), + pa.field("struct_float", pa.float32()), + pa.field("struct_double", pa.float64()), + pa.field("struct_bool", pa.bool_()), + pa.field("struct_float_vec", pa.list_(pa.float32())), + ] + ) # Build PyArrow arrays with explicit types pa_arrays = { "id": pa.array(id_arr, type=pa.int64()), "float_vector": pa.array( [np.array(v, dtype=np.float32) for v in float_vector_arr], - type=pa.list_(pa.float32()) + type=pa.list_(pa.float32()), ), - "struct_array": pa.array(struct_arr, type=pa.list_(struct_type)) + "struct_array": pa.array(struct_arr, type=pa.list_(struct_type)), } # Create PyArrow table and write to Parquet @@ -5110,7 +5031,7 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): return { "id": id_arr, "float_vector": float_vector_arr, - "struct_array": struct_arr + "struct_array": struct_arr, } def gen_json_file(self, num_rows: int, dim: int, file_path: str) -> dict: @@ -5151,17 +5072,13 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): "struct_float": random.random() * 100, "struct_double": random.random() * 1000, "struct_bool": random.choice([True, False]), - "struct_float_vec": [random.random() for _ in range(dim)] + "struct_float_vec": [random.random() for _ in range(dim)], } struct_list.append(struct_obj) struct_arr.append(struct_list) # Build row object - row = { - "id": i, - "float_vector": float_vector, - "struct_array": struct_list - } + row = {"id": i, "float_vector": float_vector, "struct_array": struct_list} rows.append(row) # Write to JSON file @@ -5173,17 +5090,16 @@ class TestMilvusClientStructArrayImport(TestMilvusClientV2Base): return { "id": id_arr, "float_vector": float_vector_arr, - "struct_array": struct_arr + "struct_array": struct_arr, } - class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): """Test case of struct array mmap functionality""" def generate_struct_array_data( self, num_rows: int, dim: int = default_dim, capacity: int = default_capacity - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """Generate test data for struct array""" data = [] for i in range(num_rows): @@ -5241,12 +5157,10 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=dim, - mmap_enabled=mmap_enabled + mmap_enabled=mmap_enabled, ) else: - schema.add_field( - field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=dim - ) + schema.add_field(field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=dim) # Create struct schema with optional mmap for sub-fields struct_schema = client.create_struct_field_schema() @@ -5257,7 +5171,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): "clip_embedding1", DataType.FLOAT_VECTOR, dim=dim, - mmap_enabled=subfield1_mmap + mmap_enabled=subfield1_mmap, ) else: struct_schema.add_field("clip_embedding1", DataType.FLOAT_VECTOR, dim=dim) @@ -5267,7 +5181,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): "clip_embedding2", DataType.FLOAT_VECTOR, dim=dim, - mmap_enabled=subfield2_mmap + mmap_enabled=subfield2_mmap, ) else: struct_schema.add_field("clip_embedding2", DataType.FLOAT_VECTOR, dim=dim) @@ -5333,9 +5247,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): ) # Create collection - self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Insert data nb = 3000 @@ -5346,11 +5258,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): self.release_collection(client, collection_name) # Enable mmap on collection - self.alter_collection_properties( - client, - collection_name, - properties={"mmap.enabled": True} - ) + self.alter_collection_properties(client, collection_name, properties={"mmap.enabled": True}) log.info(f"Enabled mmap on collection {collection_name}") # Verify mmap is enabled via describe_collection @@ -5373,7 +5281,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful with mmap enabled on collection") + log.info("Search successful with mmap enabled on collection") @pytest.mark.tags(CaseLabel.L1) def test_mmap_enable_on_struct_array_field(self): @@ -5414,9 +5322,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): ) # Create collection - self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Insert data nb = 3000 @@ -5431,9 +5337,9 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): client, collection_name, field_name="clips", - field_params={"mmap_enabled": True} + field_params={"mmap_enabled": True}, ) - log.info(f"Enabled mmap on struct array field 'clips'") + log.info("Enabled mmap on struct array field 'clips'") # Load collection self.load_collection(client, collection_name) @@ -5451,7 +5357,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful with mmap enabled on struct array field") + log.info("Search successful with mmap enabled on struct array field") @pytest.mark.tags(CaseLabel.L1) def test_mmap_enable_on_struct_subfield(self): @@ -5492,9 +5398,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): ) # Create collection - self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Insert data nb = 3000 @@ -5509,9 +5413,9 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): client, collection_name, field_name="clips[clip_embedding1]", - field_params={"mmap_enabled": True} + field_params={"mmap_enabled": True}, ) - log.info(f"Enabled mmap on struct sub-field 'clips[clip_embedding1]'") + log.info("Enabled mmap on struct sub-field 'clips[clip_embedding1]'") # Load collection self.load_collection(client, collection_name) @@ -5529,7 +5433,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful on field with mmap enabled") + log.info("Search successful on field with mmap enabled") # Verify search also works on clip_embedding2 (without mmap) self.search( @@ -5540,7 +5444,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful on field without mmap") + log.info("Search successful on field without mmap") @pytest.mark.tags(CaseLabel.L1) def test_mmap_enable_on_emb_list_index(self): @@ -5582,9 +5486,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): ) # Create collection - self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Insert data nb = 3000 @@ -5599,12 +5501,17 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): client, collection_name, index_name="clips[clip_embedding1]", - properties={"mmap.enabled": True} + properties={"mmap.enabled": True}, ) - log.info(f"Enabled mmap on embedding list index 'clips[clip_embedding1]'") + log.info("Enabled mmap on embedding list index 'clips[clip_embedding1]'") # Verify mmap is enabled via describe_index - index_info = self.describe_index(client, collection_name, index_name="clips[clip_embedding1]", check_task="check_nothing") + index_info = self.describe_index( + client, + collection_name, + index_name="clips[clip_embedding1]", + check_task="check_nothing", + ) log.info(f"Index info after enabling mmap: {index_info}") # Load collection @@ -5623,7 +5530,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful with mmap enabled on embedding list index") + log.info("Search successful with mmap enabled on embedding list index") @pytest.mark.tags(CaseLabel.L2) def test_mmap_enable_multiple_levels(self): @@ -5665,9 +5572,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): ) # Create collection - self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Insert data nb = 3000 @@ -5678,30 +5583,26 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): self.release_collection(client, collection_name) # Enable mmap at collection level - self.alter_collection_properties( - client, - collection_name, - properties={"mmap.enabled": True} - ) - log.info(f"Enabled mmap at collection level") + self.alter_collection_properties(client, collection_name, properties={"mmap.enabled": True}) + log.info("Enabled mmap at collection level") # Enable mmap at field level for struct array field self.alter_collection_field( client, collection_name, field_name="clips", - field_params={"mmap_enabled": True} + field_params={"mmap_enabled": True}, ) - log.info(f"Enabled mmap at field level for 'clips'") + log.info("Enabled mmap at field level for 'clips'") # Enable mmap at index level for embedding list index self.alter_index_properties( client, collection_name, index_name="clips[clip_embedding1]", - properties={"mmap.enabled": True} + properties={"mmap.enabled": True}, ) - log.info(f"Enabled mmap at index level for 'clips[clip_embedding1]'") + log.info("Enabled mmap at index level for 'clips[clip_embedding1]'") # Load collection self.load_collection(client, collection_name) @@ -5740,7 +5641,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "COSINE"}, limit=10, ) - log.info(f"All searches successful with mmap enabled at multiple levels") + log.info("All searches successful with mmap enabled at multiple levels") @pytest.mark.tags(CaseLabel.L1) def test_mmap_collection_enable_then_disable(self): @@ -5780,9 +5681,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): ) # Create collection - self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Insert data nb = 3000 @@ -5791,12 +5690,8 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): # Release and enable mmap self.release_collection(client, collection_name) - self.alter_collection_properties( - client, - collection_name, - properties={"mmap.enabled": True} - ) - log.info(f"Enabled mmap on collection") + self.alter_collection_properties(client, collection_name, properties={"mmap.enabled": True}) + log.info("Enabled mmap on collection") # Load and verify search works self.load_collection(client, collection_name) @@ -5812,16 +5707,12 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful with mmap enabled") + log.info("Search successful with mmap enabled") # Now disable mmap self.release_collection(client, collection_name) - self.alter_collection_properties( - client, - collection_name, - properties={"mmap.enabled": False} - ) - log.info(f"Disabled mmap on collection") + self.alter_collection_properties(client, collection_name, properties={"mmap.enabled": False}) + log.info("Disabled mmap on collection") # Load and verify search still works self.load_collection(client, collection_name) @@ -5833,7 +5724,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful after disabling mmap") + log.info("Search successful after disabling mmap") @pytest.mark.tags(CaseLabel.L1) def test_mmap_index_disable_then_enable(self): @@ -5873,9 +5764,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): ) # Create collection - self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Insert data nb = 3000 @@ -5895,7 +5784,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful without mmap") + log.info("Search successful without mmap") # Now enable mmap on index self.release_collection(client, collection_name) @@ -5903,9 +5792,9 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): client, collection_name, index_name="clips[clip_embedding1]", - properties={"mmap.enabled": True} + properties={"mmap.enabled": True}, ) - log.info(f"Enabled mmap on index 'clips[clip_embedding1]'") + log.info("Enabled mmap on index 'clips[clip_embedding1]'") # Load and verify search works self.load_collection(client, collection_name) @@ -5917,7 +5806,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful after enabling mmap on index") + log.info("Search successful after enabling mmap on index") @pytest.mark.tags(CaseLabel.L2) def test_mmap_index_enable_then_disable(self): @@ -5957,9 +5846,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): ) # Create collection - self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Insert data nb = 3000 @@ -5979,7 +5866,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful with mmap enabled on index (from schema)") + log.info("Search successful with mmap enabled on index (from schema)") # Disable mmap on index self.release_collection(client, collection_name) @@ -5987,9 +5874,9 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): client, collection_name, index_name="clips[clip_embedding1]", - properties={"mmap.enabled": False} + properties={"mmap.enabled": False}, ) - log.info(f"Disabled mmap on index") + log.info("Disabled mmap on index") # Load and verify search still works self.load_collection(client, collection_name) @@ -6001,7 +5888,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful after disabling mmap on index") + log.info("Search successful after disabling mmap on index") @pytest.mark.tags(CaseLabel.L1) def test_mmap_struct_subfield_enable_then_disable(self): @@ -6043,9 +5930,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): ) # Create collection - self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Verify mmap is enabled on subfield via describe_collection collection_info = self.describe_collection(client, collection_name, check_task="check_nothing") @@ -6071,7 +5956,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful with mmap enabled on struct sub-field (from schema)") + log.info("Search successful with mmap enabled on struct sub-field (from schema)") # Disable mmap on struct sub-field self.release_collection(client, collection_name) @@ -6079,9 +5964,9 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): client, collection_name, field_name="clips[clip_embedding1]", - field_params={"mmap_enabled": False} + field_params={"mmap_enabled": False}, ) - log.info(f"Disabled mmap on struct sub-field 'clips[clip_embedding1]'") + log.info("Disabled mmap on struct sub-field 'clips[clip_embedding1]'") # Verify mmap is disabled via describe_collection collection_info = self.describe_collection(client, collection_name, check_task="check_nothing") @@ -6097,7 +5982,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful after disabling mmap on struct sub-field") + log.info("Search successful after disabling mmap on struct sub-field") @pytest.mark.tags(CaseLabel.L1) def test_mmap_struct_subfield_disable_then_enable(self): @@ -6137,9 +6022,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): ) # Create collection - self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Insert data nb = 3000 @@ -6159,7 +6042,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful without mmap on struct sub-field") + log.info("Search successful without mmap on struct sub-field") # Now enable mmap on struct sub-field self.release_collection(client, collection_name) @@ -6167,9 +6050,9 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): client, collection_name, field_name="clips[clip_embedding1]", - field_params={"mmap_enabled": True} + field_params={"mmap_enabled": True}, ) - log.info(f"Enabled mmap on struct sub-field 'clips[clip_embedding1]'") + log.info("Enabled mmap on struct sub-field 'clips[clip_embedding1]'") # Load and verify search works self.load_collection(client, collection_name) @@ -6181,7 +6064,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful after enabling mmap on struct sub-field") + log.info("Search successful after enabling mmap on struct sub-field") @pytest.mark.tags(CaseLabel.L2) def test_mmap_struct_array_field_enable_then_disable(self): @@ -6221,9 +6104,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): ) # Create collection - self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Insert data nb = 3000 @@ -6243,7 +6124,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful with mmap enabled on struct array field (from schema)") + log.info("Search successful with mmap enabled on struct array field (from schema)") # Disable mmap on struct array field self.release_collection(client, collection_name) @@ -6251,9 +6132,9 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): client, collection_name, field_name="clips", - field_params={"mmap_enabled": False} + field_params={"mmap_enabled": False}, ) - log.info(f"Disabled mmap on struct array field 'clips'") + log.info("Disabled mmap on struct array field 'clips'") # Load and verify search still works self.load_collection(client, collection_name) @@ -6265,7 +6146,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful after disabling mmap on struct array field") + log.info("Search successful after disabling mmap on struct array field") @pytest.mark.tags(CaseLabel.L2) def test_mmap_struct_array_field_disable_then_enable(self): @@ -6305,9 +6186,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): ) # Create collection - self.create_collection( - client, collection_name, schema=schema, index_params=index_params - ) + self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Insert data nb = 3000 @@ -6327,7 +6206,7 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful without mmap on struct array field") + log.info("Search successful without mmap on struct array field") # Now enable mmap on struct array field self.release_collection(client, collection_name) @@ -6335,9 +6214,9 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): client, collection_name, field_name="clips", - field_params={"mmap_enabled": True} + field_params={"mmap_enabled": True}, ) - log.info(f"Enabled mmap on struct array field 'clips'") + log.info("Enabled mmap on struct array field 'clips'") # Load and verify search works self.load_collection(client, collection_name) @@ -6349,4 +6228,4 @@ class TestMilvusClientStructArrayMmap(TestMilvusClientV2Base): search_params={"metric_type": "MAX_SIM_COSINE"}, limit=10, ) - log.info(f"Search successful after enabling mmap on struct array field") + log.info("Search successful after enabling mmap on struct array field") diff --git a/tests/python_client/milvus_client/test_milvus_client_struct_array_element_search.py b/tests/python_client/milvus_client/test_milvus_client_struct_array_element_search.py index a996cda01c..bea1234224 100644 --- a/tests/python_client/milvus_client/test_milvus_client_struct_array_element_search.py +++ b/tests/python_client/milvus_client/test_milvus_client_struct_array_element_search.py @@ -1,16 +1,16 @@ -import pytest -import numpy as np import random +import numpy as np +import pytest from base.client_v2_base import TestMilvusClientV2Base -from utils.util_log import test_log as log from common import common_func as cf from common import common_type as ct from common.common_type import CaseLabel, CheckTasks -from utils.util_pymilvus import * from common.constants import * # noqa -from pymilvus import DataType, AnnSearchRequest, RRFRanker, WeightedRanker +from pymilvus import AnnSearchRequest, DataType, RRFRanker, WeightedRanker from pymilvus.client.embedding_list import EmbeddingList +from utils.util_log import test_log as log +from utils.util_pymilvus import * # noqa: F403 prefix = "struct_elem_search" epsilon = 0.001 @@ -38,6 +38,7 @@ def _seed_vector(seed, dim=default_dim): # ========== Ground Truth Computation ========== + def _cosine_sim(v1, v2): """Compute cosine similarity between two vectors.""" a = np.array(v1, dtype=np.float32) @@ -75,8 +76,14 @@ def _is_descending(metric_type): return metric_type in ("COSINE", "IP") -def gt_element_filter_search(data, query_vector, elem_filter_fn, metric_type="COSINE", - limit=10, doc_filter_fn=None): +def gt_element_filter_search( + data, + query_vector, + elem_filter_fn, + metric_type="COSINE", + limit=10, + doc_filter_fn=None, +): """ Ground truth for element_filter search. @@ -98,8 +105,11 @@ def gt_element_filter_search(data, query_vector, elem_filter_fn, metric_type="CO for j, elem in enumerate(row["structA"]): if elem_filter_fn(elem): score = _compute_similarity(query_vector, elem["embedding"], metric_type) - if best_score is None or (_is_descending(metric_type) and score > best_score) or \ - (not _is_descending(metric_type) and score < best_score): + if ( + best_score is None + or (_is_descending(metric_type) and score > best_score) + or (not _is_descending(metric_type) and score < best_score) + ): best_score = score best_offset = j if best_offset >= 0: @@ -109,8 +119,7 @@ def gt_element_filter_search(data, query_vector, elem_filter_fn, metric_type="CO return row_scores[:limit] -def gt_match_query(data, match_type, elem_filter_fn, threshold=None, - doc_filter_fn=None): +def gt_match_query(data, match_type, elem_filter_fn, threshold=None, doc_filter_fn=None): """ Ground truth for MATCH family query. @@ -126,23 +135,31 @@ def gt_match_query(data, match_type, elem_filter_fn, threshold=None, total = len(row["structA"]) matched = False if match_type == "MATCH_ALL": - matched = (count == total) + matched = count == total elif match_type == "MATCH_ANY": - matched = (count >= 1) + matched = count >= 1 elif match_type == "MATCH_LEAST": - matched = (count >= threshold) + matched = count >= threshold elif match_type == "MATCH_MOST": - matched = (count <= threshold) + matched = count <= threshold elif match_type == "MATCH_EXACT": - matched = (count == threshold) + matched = count == threshold if matched: matching.append(row["id"]) return set(matching) -def gt_match_search(data, query_vector, match_type, elem_filter_fn, - threshold=None, metric_type="COSINE", limit=10, - anns_field="normal_vector", doc_filter_fn=None): +def gt_match_search( + data, + query_vector, + match_type, + elem_filter_fn, + threshold=None, + metric_type="COSINE", + limit=10, + anns_field="normal_vector", + doc_filter_fn=None, +): """ Ground truth for MATCH + vector search. @@ -164,8 +181,7 @@ def gt_match_search(data, query_vector, match_type, elem_filter_fn, return row_scores[:limit] -def assert_result_ids_match(milvus_results, gt_results, recall_threshold=0.9, - is_search=True, check_order=False): +def assert_result_ids_match(milvus_results, gt_results, recall_threshold=0.9, is_search=True, check_order=False): """ Compare Milvus results against ground truth. @@ -192,17 +208,17 @@ def assert_result_ids_match(milvus_results, gt_results, recall_threshold=0.9, if len(gt_ids) > 0: overlap = milvus_id_set & gt_ids recall = len(overlap) / min(len(gt_ids), len(milvus_ids)) if milvus_ids else 0 - assert recall >= recall_threshold, \ - f"Recall {recall:.2f} < {recall_threshold}. " \ - f"GT IDs (top-{len(gt_ids)}): {sorted(gt_ids)[:20]}, " \ + assert recall >= recall_threshold, ( + f"Recall {recall:.2f} < {recall_threshold}. " + f"GT IDs (top-{len(gt_ids)}): {sorted(gt_ids)[:20]}, " f"Milvus IDs: {sorted(milvus_id_set)[:20]}" + ) # Check order (top-1 should match in most cases) if check_order and gt_results and milvus_ids: gt_top1 = gt_results[0][0] if isinstance(gt_results, list) else None if gt_top1 is not None: - assert milvus_ids[0] == gt_top1, \ - f"Top-1 mismatch: Milvus={milvus_ids[0]}, GT={gt_top1}" + assert milvus_ids[0] == gt_top1, f"Top-1 mismatch: Milvus={milvus_ids[0]}, GT={gt_top1}" def _assert_distance_order(results, metric_type): @@ -212,12 +228,14 @@ def _assert_distance_order(results, metric_type): distances = [hit["distance"] for hit in results[0]] if _is_descending(metric_type): for k in range(len(distances) - 1): - assert distances[k] >= distances[k + 1] - epsilon, \ - f"Results not sorted descending: distances[{k}]={distances[k]} < distances[{k+1}]={distances[k+1]}" + assert distances[k] >= distances[k + 1] - epsilon, ( + f"Results not sorted descending: distances[{k}]={distances[k]} < distances[{k + 1}]={distances[k + 1]}" + ) else: for k in range(len(distances) - 1): - assert distances[k] <= distances[k + 1] + epsilon, \ - f"Results not sorted ascending: distances[{k}]={distances[k]} > distances[{k+1}]={distances[k+1]}" + assert distances[k] <= distances[k + 1] + epsilon, ( + f"Results not sorted ascending: distances[{k}]={distances[k]} > distances[{k + 1}]={distances[k + 1]}" + ) def _generate_float16_vector(dim, seed=None): @@ -229,6 +247,7 @@ def _generate_float16_vector(dim, seed=None): def _generate_bfloat16_vector(dim, seed=None): """Generate BFloat16 vector as np.ndarray(dtype=bfloat16) via ml_dtypes.""" import ml_dtypes # noqa: F811 + rng = np.random.RandomState(seed) return rng.rand(dim).astype(ml_dtypes.bfloat16) @@ -288,28 +307,38 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): emb = [1.0] + [0.0] * (dim - 1) else: emb = _seed_vector(i * 1000 + j, dim) - struct_array.append({ - "embedding": emb, - "int_val": i * 100 + j, - "str_val": f"row_{i}_elem_{j}", - "float_val": float(i + j * 0.1), - "color": COLORS[j % 3], - "size": SIZES[(i + j) % 4], - }) - data.append({ - "id": i, - "doc_int": i, - "doc_varchar": f"cat_{i % 10}", - "normal_vector": _seed_vector(i + 999999, dim), - "structA": struct_array, - }) + struct_array.append( + { + "embedding": emb, + "int_val": i * 100 + j, + "str_val": f"row_{i}_elem_{j}", + "float_val": float(i + j * 0.1), + "color": COLORS[j % 3], + "size": SIZES[(i + j) % 4], + } + ) + data.append( + { + "id": i, + "doc_int": i, + "doc_varchar": f"cat_{i % 10}", + "normal_vector": _seed_vector(i + 999999, dim), + "structA": struct_array, + } + ) return data - def _create_collection_and_insert(self, client, collection_name, nb=1000, - dim=default_dim, flush=True, - metric_type="COSINE", - create_nested_index=False, - index_type="HNSW"): + def _create_collection_and_insert( + self, + client, + collection_name, + nb=1000, + dim=default_dim, + flush=True, + metric_type="COSINE", + create_nested_index=False, + index_type="HNSW", + ): """Helper: create collection, insert data, build index, load.""" schema = self._create_schema(client, dim=dim) index_params = client.prepare_index_params() @@ -362,11 +391,12 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): # Use a known element vector as query query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=10, output_fields=["id", "structA"], ) @@ -374,11 +404,9 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): assert len(results) > 0 # Scalar condition: each hit must have at least one element with int_val >= 0 for hit in results[0]: - assert any(e["int_val"] >= 0 for e in hit["structA"]), \ - f"Row {hit['id']} has no element with int_val >= 0" + assert any(e["int_val"] >= 0 for e in hit["structA"]), f"Row {hit['id']} has no element with int_val >= 0" # Top-1: querying row 0's own vector should return row 0 - assert results[0][0]["id"] == 0, \ - f"Top-1 should be row 0 (queried its own vector), got {results[0][0]['id']}" + assert results[0][0]["id"] == 0, f"Top-1 should be row 0 (queried its own vector), got {results[0][0]['id']}" # Distance ordering _assert_distance_order(results, "COSINE") @@ -396,11 +424,12 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "L2"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=10, output_fields=["id", "structA"], ) @@ -408,11 +437,9 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): assert len(results) > 0 # Scalar condition verification for hit in results[0]: - assert any(e["int_val"] >= 0 for e in hit["structA"]), \ - f"Row {hit['id']} has no element with int_val >= 0" + assert any(e["int_val"] >= 0 for e in hit["structA"]), f"Row {hit['id']} has no element with int_val >= 0" # Top-1: querying row 0's own vector should return row 0 - assert results[0][0]["id"] == 0, \ - f"Top-1 should be row 0, got {results[0][0]['id']}" + assert results[0][0]["id"] == 0, f"Top-1 should be row 0, got {results[0][0]['id']}" # Distance ordering (L2: ascending) _assert_distance_order(results, "L2") @@ -429,7 +456,8 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): query_vector = data[200]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -442,8 +470,7 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): for hit in results[0]: assert hit["doc_int"] > 100 # Top-1: querying row 200's vector with its specific str_val filter - assert results[0][0]["id"] == 200, \ - f"Top-1 should be row 200, got {results[0][0]['id']}" + assert results[0][0]["id"] == 200, f"Top-1 should be row 200, got {results[0][0]['id']}" @pytest.mark.tags(CaseLabel.L0) def test_element_filter_search_compound_same_element_semantic(self): @@ -461,12 +488,16 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): schema = self._create_schema(client) index_params = client.prepare_index_params() index_params.add_index( - field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS, + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, ) index_params.add_index( - field_name="structA[embedding]", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS, + field_name="structA[embedding]", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, ) res, check = self.create_collection(client, collection_name, schema=schema, index_params=index_params) assert check @@ -475,29 +506,59 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): target_vec = _seed_vector(77777) data = [ { - "id": 0, "doc_int": 0, "doc_varchar": "cat_0", + "id": 0, + "doc_int": 0, + "doc_varchar": "cat_0", "normal_vector": _seed_vector(99990), "structA": [ - {"embedding": _seed_vector(0), "int_val": 1, "str_val": "a", - "float_val": 0.1, "color": "Red", "size": "S"}, - {"embedding": target_vec, "int_val": 2, "str_val": "b", - "float_val": 0.2, "color": "Blue", "size": "L"}, + { + "embedding": _seed_vector(0), + "int_val": 1, + "str_val": "a", + "float_val": 0.1, + "color": "Red", + "size": "S", + }, + { + "embedding": target_vec, + "int_val": 2, + "str_val": "b", + "float_val": 0.2, + "color": "Blue", + "size": "L", + }, ], }, { - "id": 1, "doc_int": 1, "doc_varchar": "cat_1", + "id": 1, + "doc_int": 1, + "doc_varchar": "cat_1", "normal_vector": _seed_vector(99991), "structA": [ - {"embedding": target_vec, "int_val": 10, "str_val": "x", - "float_val": 1.0, "color": "Red", "size": "L"}, + { + "embedding": target_vec, + "int_val": 10, + "str_val": "x", + "float_val": 1.0, + "color": "Red", + "size": "L", + }, ], }, { - "id": 2, "doc_int": 2, "doc_varchar": "cat_2", + "id": 2, + "doc_int": 2, + "doc_varchar": "cat_2", "normal_vector": _seed_vector(99992), "structA": [ - {"embedding": _seed_vector(20), "int_val": 20, "str_val": "p", - "float_val": 2.0, "color": "Blue", "size": "S"}, + { + "embedding": _seed_vector(20), + "int_val": 20, + "str_val": "p", + "float_val": 2.0, + "color": "Blue", + "size": "S", + }, ], }, ] @@ -507,7 +568,8 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): # Search using target_vec so Row 0 (elem[1]) and Row 1 (elem[0]) are closest results, check = self.search( - client, collection_name, + client, + collection_name, data=[target_vec], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -518,16 +580,16 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): assert check matched_ids = {hit["id"] for hit in results[0]} # Row 0: Red on elem[0](size=S), L on elem[1](color=Blue) → NO match - assert 0 not in matched_ids, \ - "Row 0 should NOT match (Red and L are on different elements)" + assert 0 not in matched_ids, "Row 0 should NOT match (Red and L are on different elements)" # Row 1: elem[0]={Red,L} → YES - assert 1 in matched_ids, \ - "Row 1 should match (Red+L on same element)" + assert 1 in matched_ids, "Row 1 should match (Red+L on same element)" # Ground truth verification gt = gt_element_filter_search( - data, target_vec, + data, + target_vec, elem_filter_fn=lambda e: e["color"] == "Red" and e["size"] == "L", - metric_type="COSINE", limit=10, + metric_type="COSINE", + limit=10, ) assert_result_ids_match(results, gt, recall_threshold=1.0) @@ -549,7 +611,8 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): query_vector = data[target_row]["structA"][target_elem]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -563,8 +626,7 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): top_hit = results[0] assert top_hit["id"] == target_row # Verify offset field exists - assert "offset" in top_hit or hasattr(top_hit, "offset"), \ - "element_indices (offset) not exposed in pymilvus" + assert "offset" in top_hit or hasattr(top_hit, "offset"), "element_indices (offset) not exposed in pymilvus" # ---- L1 tests ---- @@ -581,19 +643,19 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "IP"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=10, output_fields=["id", "structA"], ) assert check assert len(results) > 0 for hit in results[0]: - assert any(e["int_val"] >= 0 for e in hit["structA"]), \ - f"Row {hit['id']} has no element with int_val >= 0" + assert any(e["int_val"] >= 0 for e in hit["structA"]), f"Row {hit['id']} has no element with int_val >= 0" _assert_distance_order(results, "IP") @pytest.mark.tags(CaseLabel.L1) @@ -609,11 +671,12 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): query_vector = data[10]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] > 50)', + filter="element_filter(structA, $[int_val] > 50)", limit=10, output_fields=["id", "structA"], ) @@ -621,8 +684,7 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): assert len(results) > 0 # Scalar condition: each hit must have element with int_val > 50 for hit in results[0]: - assert any(e["int_val"] > 50 for e in hit["structA"]), \ - f"Row {hit['id']} has no element with int_val > 50" + assert any(e["int_val"] > 50 for e in hit["structA"]), f"Row {hit['id']} has no element with int_val > 50" _assert_distance_order(results, "COSINE") @pytest.mark.tags(CaseLabel.L1) @@ -638,7 +700,8 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -650,8 +713,9 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): assert len(results) > 0 # CRITICAL: BOTH conditions must hold on the SAME element for hit in results[0]: - assert any(e["color"] == "Red" and e["int_val"] > 10 for e in hit["structA"]), \ + assert any(e["color"] == "Red" and e["int_val"] > 10 for e in hit["structA"]), ( f"Row {hit['id']} has no element with color==Red AND int_val>10 on same element" + ) _assert_distance_order(results, "COSINE") @pytest.mark.tags(CaseLabel.L1) @@ -667,7 +731,8 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -678,8 +743,9 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): assert check assert len(results) > 0 for hit in results[0]: - assert any(e["color"] == "Red" or e["int_val"] > 99900 for e in hit["structA"]), \ + assert any(e["color"] == "Red" or e["int_val"] > 99900 for e in hit["structA"]), ( f"Row {hit['id']} has no element matching OR condition" + ) _assert_distance_order(results, "COSINE") @pytest.mark.tags(CaseLabel.L1) @@ -695,11 +761,12 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): query_vector = data[1]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] > 100 && $[int_val] < 200)', + filter="element_filter(structA, $[int_val] > 100 && $[int_val] < 200)", limit=10, output_fields=["id", "structA"], ) @@ -707,8 +774,9 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): assert len(results[0]) > 0, "Range filter should match row 1 elements (int_val 101-109)" # Scalar condition verification: each hit must have element in range for hit in results[0]: - assert any(100 < e["int_val"] < 200 for e in hit["structA"]), \ + assert any(100 < e["int_val"] < 200 for e in hit["structA"]), ( f"Row {hit['id']} has no element in range (100, 200)" + ) _assert_distance_order(results, "COSINE") @pytest.mark.tags(CaseLabel.L1) @@ -724,7 +792,8 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): query_vector = data[1]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -751,19 +820,19 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): query_vector = data[5]["structA"][0]["embedding"] # All elements of row 5 have int_val = 500..509, all > 100 results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] > 100)', + filter="element_filter(structA, $[int_val] > 100)", limit=10, output_fields=["id", "structA"], ) assert check assert len(results) > 0 for hit in results[0]: - assert any(e["int_val"] > 100 for e in hit["structA"]), \ - f"Row {hit['id']} has no element with int_val > 100" + assert any(e["int_val"] > 100 for e in hit["structA"]), f"Row {hit['id']} has no element with int_val > 100" _assert_distance_order(results, "COSINE") @pytest.mark.tags(CaseLabel.L1) @@ -780,11 +849,12 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] # int_val max is 999*100+9 = 99909, so > 999999 matches nothing results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] > 999999)', + filter="element_filter(structA, $[int_val] > 999999)", limit=10, ) assert check @@ -799,17 +869,16 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): """ client = self._client() collection_name = cf.gen_unique_str(f"{prefix}_ef_growing") - data = self._create_collection_and_insert( - client, collection_name, nb=1000, metric_type="COSINE", flush=False - ) + data = self._create_collection_and_insert(client, collection_name, nb=1000, metric_type="COSINE", flush=False) query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=10, output_fields=["id", "structA"], ) @@ -818,8 +887,7 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): for hit in results[0]: assert any(e["int_val"] >= 0 for e in hit["structA"]) # Top-1: searching row 0's own vector should return row 0 - assert results[0][0]["id"] == 0, \ - f"Top-1 should be row 0 (self-vector), got {results[0][0]['id']}" + assert results[0][0]["id"] == 0, f"Top-1 should be row 0 (self-vector), got {results[0][0]['id']}" _assert_distance_order(results, "COSINE") @pytest.mark.xfail(reason="flaky: element-level search on sealed segment returns wrong element-to-row mapping") @@ -832,17 +900,16 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): """ client = self._client() collection_name = cf.gen_unique_str(f"{prefix}_ef_sealed") - data = self._create_collection_and_insert( - client, collection_name, nb=1000, metric_type="COSINE", flush=True - ) + data = self._create_collection_and_insert(client, collection_name, nb=1000, metric_type="COSINE", flush=True) query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=10, output_fields=["id", "structA"], ) @@ -851,8 +918,7 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): for hit in results[0]: assert any(e["int_val"] >= 0 for e in hit["structA"]) # Top-1: searching row 0's own vector should return row 0 - assert results[0][0]["id"] == 0, \ - f"Top-1 should be row 0 (self-vector), got {results[0][0]['id']}" + assert results[0][0]["id"] == 0, f"Top-1 should be row 0 (self-vector), got {results[0][0]['id']}" _assert_distance_order(results, "COSINE") @pytest.mark.tags(CaseLabel.L1) @@ -868,11 +934,12 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=10, output_fields=["id", "doc_int", "structA[int_val]", "structA[color]"], ) @@ -882,7 +949,7 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): for hit in results[0]: assert hit["id"] is not None # Access doc_int via [] (pymilvus proxies to entity) - assert hit["doc_int"] is not None, f"doc_int missing from hit" + assert hit["doc_int"] is not None, "doc_int missing from hit" # Verify structA sub-fields accessible assert hit["structA"] is not None _assert_distance_order(results, "COSINE") @@ -925,11 +992,12 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=10, output_fields=["id", "structA"], ) @@ -938,8 +1006,7 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): for hit in results[0]: assert any(e["int_val"] >= 0 for e in hit["structA"]) # Top-1: searching row 0's own vector should return row 0 - assert results[0][0]["id"] == 0, \ - f"Top-1 should be row 0 (self-vector), got {results[0][0]['id']}" + assert results[0][0]["id"] == 0, f"Top-1 should be row 0 (self-vector), got {results[0][0]['id']}" _assert_distance_order(results, "COSINE") @pytest.mark.tags(CaseLabel.L1) @@ -957,7 +1024,8 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE", "params": {"ef": 200}}, @@ -968,8 +1036,7 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): assert check assert len(results) > 0 for hit in results[0]: - assert any(e["color"] == "Red" for e in hit["structA"]), \ - f"Row {hit['id']} has no Red element" + assert any(e["color"] == "Red" for e in hit["structA"]), f"Row {hit['id']} has no Red element" _assert_distance_order(results, "COSINE") # ---- L2 tests ---- @@ -987,12 +1054,16 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): schema = self._create_schema(client) index_params = client.prepare_index_params() index_params.add_index( - field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS, + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, ) index_params.add_index( - field_name="structA[embedding]", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS, + field_name="structA[embedding]", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, ) res, check = self.create_collection(client, collection_name, schema=schema, index_params=index_params) assert check @@ -1009,30 +1080,36 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): num_elems = rng.randint(3, 10) struct_array = [] for j in range(num_elems): - struct_array.append({ - "embedding": _seed_vector(i * 1000 + j), - "int_val": i * 100 + j, - "str_val": f"row_{i}_elem_{j}", - "float_val": float(i + j * 0.1), - "color": COLORS[j % 3], - "size": SIZES[(i + j) % 4], - }) - growing_data.append({ - "id": i, "doc_int": i, "doc_varchar": f"cat_{i % 10}", - "normal_vector": _seed_vector(i + 999999), - "structA": struct_array, - }) + struct_array.append( + { + "embedding": _seed_vector(i * 1000 + j), + "int_val": i * 100 + j, + "str_val": f"row_{i}_elem_{j}", + "float_val": float(i + j * 0.1), + "color": COLORS[j % 3], + "size": SIZES[(i + j) % 4], + } + ) + growing_data.append( + { + "id": i, + "doc_int": i, + "doc_varchar": f"cat_{i % 10}", + "normal_vector": _seed_vector(i + 999999), + "structA": struct_array, + } + ) self.insert(client, collection_name, growing_data) self.load_collection(client, collection_name) - all_data = sealed_data + growing_data query_vector = sealed_data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=20, output_fields=["id", "structA"], ) @@ -1056,21 +1133,23 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] # First page results_p1, _ = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=5, offset=0, ) # Second page results_p2, _ = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=5, offset=5, ) @@ -1085,8 +1164,7 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): # Last score of page 1 should be >= first score of page 2 (COSINE: higher is better) p1_last = results_p1[0][-1]["distance"] p2_first = results_p2[0][0]["distance"] - assert p1_last >= p2_first - 1e-4, \ - f"Page 1 last distance ({p1_last}) should >= page 2 first ({p2_first})" + assert p1_last >= p2_first - 1e-4, f"Page 1 last distance ({p1_last}) should >= page 2 first ({p2_first})" @pytest.mark.tags(CaseLabel.L2) def test_element_filter_search_nq_multiple(self): @@ -1099,15 +1177,14 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): collection_name = cf.gen_unique_str(f"{prefix}_ef_nq") data = self._create_collection_and_insert(client, collection_name, nb=1000, metric_type="COSINE") - query_vectors = [ - data[i]["structA"][0]["embedding"] for i in range(3) - ] + query_vectors = [data[i]["structA"][0]["embedding"] for i in range(3)] results, check = self.search( - client, collection_name, + client, + collection_name, data=query_vectors, anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=5, output_fields=["id", "structA"], ) @@ -1121,8 +1198,9 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): # Verify distance ordering within each nq distances = [hit["distance"] for hit in results[nq_idx]] for k in range(len(distances) - 1): - assert distances[k] >= distances[k + 1] - 1e-4, \ + assert distances[k] >= distances[k + 1] - 1e-4, ( f"nq[{nq_idx}] distances not sorted: {distances[k]} < {distances[k + 1]}" + ) @pytest.mark.tags(CaseLabel.L2) def test_element_filter_search_negation(self): @@ -1137,11 +1215,12 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, !($[int_val] < 0))', + filter="element_filter(structA, !($[int_val] < 0))", limit=10, output_fields=["id", "structA"], ) @@ -1162,17 +1241,16 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): """ client = self._client() collection_name = cf.gen_unique_str(f"{prefix}_ef_{metric_type.lower()}") - data = self._create_collection_and_insert( - client, collection_name, nb=500, metric_type=metric_type - ) + data = self._create_collection_and_insert(client, collection_name, nb=500, metric_type=metric_type) query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": metric_type}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=10, output_fields=["id", "structA"], ) @@ -1181,8 +1259,7 @@ class TestMilvusClientStructArrayElementFilterSearch(TestMilvusClientV2Base): for hit in results[0]: assert any(e["int_val"] >= 0 for e in hit["structA"]) # Top-1: searching row 0's own vector should return row 0 - assert results[0][0]["id"] == 0, \ - f"Top-1 should be row 0 (self-vector), got {results[0][0]['id']}" + assert results[0][0]["id"] == 0, f"Top-1 should be row 0 (self-vector), got {results[0][0]['id']}" _assert_distance_order(results, metric_type) @@ -1225,19 +1302,24 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): emb = [1.0] + [0.0] * (dim - 1) else: emb = _seed_vector(i * 1000 + j, dim) - struct_array.append({ - "embedding": emb, - "int_val": i * 100 + j, - "str_val": f"row_{i}_elem_{j}", - "float_val": float(i + j * 0.1), - "color": COLORS[j % 3], - "size": SIZES[(i + j) % 4], - }) - data.append({ - "id": i, "doc_int": i, - "normal_vector": _seed_vector(i + 999999, dim), - "structA": struct_array, - }) + struct_array.append( + { + "embedding": emb, + "int_val": i * 100 + j, + "str_val": f"row_{i}_elem_{j}", + "float_val": float(i + j * 0.1), + "color": COLORS[j % 3], + "size": SIZES[(i + j) % 4], + } + ) + data.append( + { + "id": i, + "doc_int": i, + "normal_vector": _seed_vector(i + 999999, dim), + "structA": struct_array, + } + ) return data def _generate_semantic_data(self, client, dim=default_dim): @@ -1248,52 +1330,108 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): """ data = [ { - "id": 0, "doc_int": 0, + "id": 0, + "doc_int": 0, "normal_vector": _seed_vector(99990, dim), "structA": [ - {"embedding": _seed_vector(0, dim), "int_val": 1, "str_val": "a", - "float_val": 0.1, "color": "Red", "size": "S"}, - {"embedding": _seed_vector(1, dim), "int_val": 2, "str_val": "b", - "float_val": 0.2, "color": "Blue", "size": "L"}, - {"embedding": _seed_vector(2, dim), "int_val": 3, "str_val": "c", - "float_val": 0.3, "color": "Green", "size": "M"}, + { + "embedding": _seed_vector(0, dim), + "int_val": 1, + "str_val": "a", + "float_val": 0.1, + "color": "Red", + "size": "S", + }, + { + "embedding": _seed_vector(1, dim), + "int_val": 2, + "str_val": "b", + "float_val": 0.2, + "color": "Blue", + "size": "L", + }, + { + "embedding": _seed_vector(2, dim), + "int_val": 3, + "str_val": "c", + "float_val": 0.3, + "color": "Green", + "size": "M", + }, ], }, { - "id": 1, "doc_int": 1, + "id": 1, + "doc_int": 1, "normal_vector": _seed_vector(99991, dim), "structA": [ - {"embedding": _seed_vector(10, dim), "int_val": 10, "str_val": "x", - "float_val": 1.0, "color": "Red", "size": "L"}, - {"embedding": _seed_vector(11, dim), "int_val": 11, "str_val": "y", - "float_val": 1.1, "color": "Red", "size": "L"}, + { + "embedding": _seed_vector(10, dim), + "int_val": 10, + "str_val": "x", + "float_val": 1.0, + "color": "Red", + "size": "L", + }, + { + "embedding": _seed_vector(11, dim), + "int_val": 11, + "str_val": "y", + "float_val": 1.1, + "color": "Red", + "size": "L", + }, ], }, { - "id": 2, "doc_int": 2, + "id": 2, + "doc_int": 2, "normal_vector": _seed_vector(99992, dim), "structA": [ - {"embedding": _seed_vector(20, dim), "int_val": 20, "str_val": "p", - "float_val": 2.0, "color": "Blue", "size": "S"}, - {"embedding": _seed_vector(21, dim), "int_val": 21, "str_val": "q", - "float_val": 2.1, "color": "Green", "size": "XL"}, + { + "embedding": _seed_vector(20, dim), + "int_val": 20, + "str_val": "p", + "float_val": 2.0, + "color": "Blue", + "size": "S", + }, + { + "embedding": _seed_vector(21, dim), + "int_val": 21, + "str_val": "q", + "float_val": 2.1, + "color": "Green", + "size": "XL", + }, ], }, ] return data - def _setup_collection(self, client, collection_name, data, dim=default_dim, - flush=True, create_nested_index=False): + def _setup_collection( + self, + client, + collection_name, + data, + dim=default_dim, + flush=True, + create_nested_index=False, + ): """Helper to create collection, insert data, build index, load.""" schema = self._create_schema(client, dim=dim) index_params = client.prepare_index_params() index_params.add_index( - field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS, + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, ) index_params.add_index( - field_name="structA[embedding]", index_type="HNSW", - metric_type="MAX_SIM_COSINE", params=INDEX_PARAMS, + field_name="structA[embedding]", + index_type="HNSW", + metric_type="MAX_SIM_COSINE", + params=INDEX_PARAMS, ) if create_nested_index: index_params.add_index(field_name="structA[int_val]", index_type="INVERTED") @@ -1323,7 +1461,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ALL(structA, $[color] == "Red")', output_fields=["id", "structA"], limit=100, @@ -1353,7 +1492,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ALL(structA, $[color] == "Red" && $[size] == "L")', output_fields=["id", "structA"], limit=100, @@ -1378,7 +1518,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ANY(structA, $[color] == "Blue")', output_fields=["id", "structA"], limit=100, @@ -1410,7 +1551,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ANY(structA, $[color] == "Red" && $[size] == "L")', output_fields=["id", "structA"], limit=100, @@ -1437,7 +1579,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ANY(structA, $[color] == "Red" && $[size] == "L")', output_fields=["id", "structA"], limit=100, @@ -1460,8 +1603,9 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results, check = self.query( - client, collection_name, - filter='MATCH_LEAST(structA, $[int_val] > 5, threshold=3)', + client, + collection_name, + filter="MATCH_LEAST(structA, $[int_val] > 5, threshold=3)", output_fields=["id", "structA"], limit=100, ) @@ -1489,13 +1633,15 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results_least, _ = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_LEAST(structA, $[color] == "Blue", threshold=1)', output_fields=["id"], limit=500, ) results_any, _ = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ANY(structA, $[color] == "Blue")', output_fields=["id"], limit=500, @@ -1517,8 +1663,9 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results, check = self.query( - client, collection_name, - filter='MATCH_MOST(structA, $[int_val] > 5, threshold=2)', + client, + collection_name, + filter="MATCH_MOST(structA, $[int_val] > 5, threshold=2)", output_fields=["id", "structA"], limit=100, ) @@ -1546,7 +1693,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_EXACT(structA, $[color] == "Red", threshold=2)', output_fields=["id", "structA"], limit=100, @@ -1575,7 +1723,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results, check = self.query( - client, collection_name, + client, + collection_name, filter='doc_int > 10 && MATCH_ANY(structA, $[color] == "Red")', output_fields=["id", "doc_int", "structA"], limit=100, @@ -1586,7 +1735,9 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): assert any(e["color"] == "Red" for e in hit["structA"]) # GT completeness check with doc-level filter gt_ids = gt_match_query( - data, "MATCH_ANY", lambda e: e["color"] == "Red", + data, + "MATCH_ANY", + lambda e: e["color"] == "Red", doc_filter_fn=lambda row: row["doc_int"] > 10, ) milvus_ids = {r["id"] for r in results} @@ -1608,7 +1759,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): query_vector = data[0]["normal_vector"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="normal_vector", search_params={"metric_type": "COSINE"}, @@ -1622,8 +1774,13 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): assert any(e["color"] == "Red" for e in hit["structA"]) # GT: MATCH_ANY filter + normal_vector search gt = gt_match_search( - data, query_vector, "MATCH_ANY", lambda e: e["color"] == "Red", - metric_type="COSINE", limit=10, anns_field="normal_vector", + data, + query_vector, + "MATCH_ANY", + lambda e: e["color"] == "Red", + metric_type="COSINE", + limit=10, + anns_field="normal_vector", ) assert_result_ids_match(results, gt, recall_threshold=0.8) _assert_distance_order(results, "COSINE") @@ -1641,7 +1798,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data, flush=False) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ANY(structA, $[color] == "Red")', output_fields=["id"], limit=100, @@ -1662,7 +1820,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data, flush=True) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ANY(structA, $[color] == "Blue")', output_fields=["id"], limit=100, @@ -1683,8 +1842,9 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data, create_nested_index=False) results, check = self.query( - client, collection_name, - filter='MATCH_ANY(structA, $[int_val] > 500)', + client, + collection_name, + filter="MATCH_ANY(structA, $[int_val] > 500)", output_fields=["id", "structA"], limit=100, ) @@ -1710,7 +1870,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data, create_nested_index=True) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ANY(structA, $[color] == "Green")', output_fields=["id", "structA"], limit=100, @@ -1732,12 +1893,16 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): schema = self._create_schema(client) index_params = client.prepare_index_params() index_params.add_index( - field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS, + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, ) index_params.add_index( - field_name="structA[embedding]", index_type="HNSW", - metric_type="MAX_SIM_COSINE", params=INDEX_PARAMS, + field_name="structA[embedding]", + index_type="HNSW", + metric_type="MAX_SIM_COSINE", + params=INDEX_PARAMS, ) index_params.add_index(field_name="structA[int_val]", index_type="STL_SORT") @@ -1750,8 +1915,9 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self.load_collection(client, collection_name) results, check = self.query( - client, collection_name, - filter='MATCH_ANY(structA, $[int_val] > 1000)', + client, + collection_name, + filter="MATCH_ANY(structA, $[int_val] > 1000)", output_fields=["id", "structA"], limit=100, ) @@ -1774,7 +1940,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_MOST(structA, $[color] == "Red", threshold=0)', output_fields=["id", "structA"], limit=100, @@ -1802,7 +1969,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_EXACT(structA, $[color] == "Red", threshold=0)', output_fields=["id", "structA"], limit=100, @@ -1825,7 +1993,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ALL(structA, $[color] == "Red" || $[size] == "L")', output_fields=["id", "structA"], limit=100, @@ -1835,7 +2004,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): for elem in hit["structA"]: assert elem["color"] == "Red" or elem["size"] == "L" gt_ids = gt_match_query( - data, "MATCH_ALL", + data, + "MATCH_ALL", lambda e: e["color"] == "Red" or e["size"] == "L", ) milvus_ids = {r["id"] for r in results} @@ -1856,14 +2026,18 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results_upper, _ = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ANY(structA, $[color] == "Red")', - output_fields=["id"], limit=500, + output_fields=["id"], + limit=500, ) results_lower, _ = self.query( - client, collection_name, + client, + collection_name, filter='match_any(structA, $[color] == "Red")', - output_fields=["id"], limit=500, + output_fields=["id"], + limit=500, ) ids_upper = sorted([r["id"] for r in results_upper]) ids_lower = sorted([r["id"] for r in results_lower]) @@ -1882,9 +2056,11 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results, check = self.query( - client, collection_name, - filter='MATCH_LEAST(structA, $[int_val] >= 0, threshold=100)', - output_fields=["id"], limit=100, + client, + collection_name, + filter="MATCH_LEAST(structA, $[int_val] >= 0, threshold=100)", + output_fields=["id"], + limit=100, ) assert check assert len(results) == 0, "No row should have 100+ matching elements" @@ -1902,27 +2078,39 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): schema = self._create_schema(client) index_params = client.prepare_index_params() index_params.add_index( - field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS, + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, ) index_params.add_index( - field_name="structA[embedding]", index_type="HNSW", - metric_type="MAX_SIM_COSINE", params=INDEX_PARAMS, + field_name="structA[embedding]", + index_type="HNSW", + metric_type="MAX_SIM_COSINE", + params=INDEX_PARAMS, ) self.create_collection(client, collection_name, schema=schema, index_params=index_params) data = [ { - "id": 0, "doc_int": 0, + "id": 0, + "doc_int": 0, "normal_vector": _seed_vector(0), "structA": [], # empty array }, { - "id": 1, "doc_int": 1, + "id": 1, + "doc_int": 1, "normal_vector": _seed_vector(1), "structA": [ - {"embedding": _seed_vector(10), "int_val": 1, "str_val": "a", - "float_val": 0.1, "color": "Red", "size": "S"}, + { + "embedding": _seed_vector(10), + "int_val": 1, + "str_val": "a", + "float_val": 0.1, + "color": "Red", + "size": "S", + }, ], }, ] @@ -1931,7 +2119,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self.load_collection(client, collection_name) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ALL(structA, $[color] == "Red")', output_fields=["id", "structA"], limit=100, @@ -1959,7 +2148,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ALL(structA, $[int_val] > 0) && MATCH_ANY(structA, $[color] == "Blue")', output_fields=["id", "structA"], limit=100, @@ -1990,7 +2180,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): self._setup_collection(client, collection_name, data) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ANY(structA, $[color] == "Red") || doc_int > 900', output_fields=["id", "doc_int", "structA"], limit=500, @@ -2026,10 +2217,8 @@ class TestMilvusClientStructArrayElementMatchSearch(TestMilvusClientV2Base): filter_expr = 'MATCH_ANY(structA, $[color] == "Red" && $[int_val] > 500)' - results_bf, _ = self.query(client, name_bf, filter=filter_expr, - output_fields=["id"], limit=500) - results_idx, _ = self.query(client, name_idx, filter=filter_expr, - output_fields=["id"], limit=500) + results_bf, _ = self.query(client, name_bf, filter=filter_expr, output_fields=["id"], limit=500) + results_idx, _ = self.query(client, name_idx, filter=filter_expr, output_fields=["id"], limit=500) ids_bf = sorted([r["id"] for r in results_bf]) ids_idx = sorted([r["id"] for r in results_idx]) @@ -2057,7 +2246,8 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): struct_schema.add_field("bool_val", DataType.BOOL) schema.add_field( - "structA", datatype=DataType.ARRAY, + "structA", + datatype=DataType.ARRAY, element_type=DataType.STRUCT, struct_schema=struct_schema, max_capacity=default_capacity, @@ -2087,25 +2277,39 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): if add_bool: elem["bool_val"] = (i + j) % 2 == 0 struct_array.append(elem) - data.append({ - "id": i, "doc_int": i, - "normal_vector": _seed_vector(i + 999999, dim), - "structA": struct_array, - }) + data.append( + { + "id": i, + "doc_int": i, + "normal_vector": _seed_vector(i + 999999, dim), + "structA": struct_array, + } + ) return data - def _setup_base_collection(self, client, collection_name, dim=default_dim, - nb=500, add_bool=False, extra_index_fn=None): + def _setup_base_collection( + self, + client, + collection_name, + dim=default_dim, + nb=500, + add_bool=False, + extra_index_fn=None, + ): """Create collection with base indexes, optionally add extra indexes.""" schema = self._create_schema(client, dim=dim, add_bool=add_bool) index_params = client.prepare_index_params() index_params.add_index( - field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS, + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, ) index_params.add_index( - field_name="structA[embedding]", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS, + field_name="structA[embedding]", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, ) if extra_index_fn: extra_index_fn(index_params) @@ -2187,8 +2391,7 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): def add_nested(ip): ip.add_index(field_name="structA[bool_val]", index_type="INVERTED") - self._setup_base_collection(client, collection_name, add_bool=True, - extra_index_fn=add_nested) + self._setup_base_collection(client, collection_name, add_bool=True, extra_index_fn=add_nested) @pytest.mark.tags(CaseLabel.L1) def test_create_inverted_index_on_struct_float(self): @@ -2252,19 +2455,19 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] > 100)', + filter="element_filter(structA, $[int_val] > 100)", limit=10, output_fields=["id", "structA"], ) assert check assert len(results) > 0 for hit in results[0]: - assert any(e["int_val"] > 100 for e in hit["structA"]), \ - f"Row {hit['id']} has no element with int_val > 100" + assert any(e["int_val"] > 100 for e in hit["structA"]), f"Row {hit['id']} has no element with int_val > 100" _assert_distance_order(results, "COSINE") @pytest.mark.tags(CaseLabel.L1) @@ -2283,7 +2486,8 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): self._setup_base_collection(client, collection_name, extra_index_fn=add_nested) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ANY(structA, $[color] == "Blue")', output_fields=["id", "structA"], limit=100, @@ -2309,11 +2513,12 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] > 100)', + filter="element_filter(structA, $[int_val] > 100)", limit=10, ) assert check @@ -2334,8 +2539,9 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): self._setup_base_collection(client, collection_name, extra_index_fn=add_nested) results, check = self.query( - client, collection_name, - filter='MATCH_ANY(structA, $[int_val] > 1000)', + client, + collection_name, + filter="MATCH_ANY(structA, $[int_val] > 1000)", output_fields=["id", "structA"], limit=100, ) @@ -2361,7 +2567,8 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): # Query using both indexed fields results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ANY(structA, $[color] == "Red" && $[int_val] > 100)', output_fields=["id"], limit=100, @@ -2384,8 +2591,9 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): self._setup_base_collection(client, collection_name, extra_index_fn=add_nested) results, check = self.query( - client, collection_name, - filter='MATCH_ANY(structA, $[int_val] > 100)', + client, + collection_name, + filter="MATCH_ANY(structA, $[int_val] > 100)", output_fields=["id"], limit=100, ) @@ -2405,8 +2613,7 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): def add_nested(ip): ip.add_index(field_name="structA[int_val]", index_type="INVERTED") - self._setup_base_collection(client, collection_name, nb=500, - extra_index_fn=add_nested) + self._setup_base_collection(client, collection_name, nb=500, extra_index_fn=add_nested) # Insert more data without flush (growing) extra_data = [] @@ -2415,25 +2622,31 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): num_elems = rng.randint(3, 10) struct_array = [] for j in range(num_elems): - struct_array.append({ - "embedding": _seed_vector(i * 1000 + j), - "int_val": i * 100 + j, - "str_val": f"row_{i}_elem_{j}", - "float_val": float(i + j * 0.1), - "color": COLORS[j % 3], - "size": SIZES[(i + j) % 4], - }) - extra_data.append({ - "id": i, "doc_int": i, - "normal_vector": _seed_vector(i + 999999), - "structA": struct_array, - }) + struct_array.append( + { + "embedding": _seed_vector(i * 1000 + j), + "int_val": i * 100 + j, + "str_val": f"row_{i}_elem_{j}", + "float_val": float(i + j * 0.1), + "color": COLORS[j % 3], + "size": SIZES[(i + j) % 4], + } + ) + extra_data.append( + { + "id": i, + "doc_int": i, + "normal_vector": _seed_vector(i + 999999), + "structA": struct_array, + } + ) self.insert(client, collection_name, extra_data) # No flush - growing segment results, check = self.query( - client, collection_name, - filter='MATCH_ANY(structA, $[int_val] > 60000)', + client, + collection_name, + filter="MATCH_ANY(structA, $[int_val] > 60000)", output_fields=["id"], limit=100, ) @@ -2456,8 +2669,11 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): collection_name = cf.gen_unique_str(f"{prefix}_ni_rebuild") def add_nested(ip): - ip.add_index(field_name="structA[int_val]", index_type="INVERTED", - index_name="nested_int_idx") + ip.add_index( + field_name="structA[int_val]", + index_type="INVERTED", + index_name="nested_int_idx", + ) self._setup_base_collection(client, collection_name, extra_index_fn=add_nested) @@ -2467,14 +2683,18 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): # Rebuild index_params = client.prepare_index_params() - index_params.add_index(field_name="structA[int_val]", index_type="INVERTED", - index_name="nested_int_idx_v2") + index_params.add_index( + field_name="structA[int_val]", + index_type="INVERTED", + index_name="nested_int_idx_v2", + ) self.create_index(client, collection_name, index_params) self.load_collection(client, collection_name) results, check = self.query( - client, collection_name, - filter='MATCH_ANY(structA, $[int_val] > 100)', + client, + collection_name, + filter="MATCH_ANY(structA, $[int_val] > 100)", output_fields=["id"], limit=100, ) @@ -2513,17 +2733,26 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): struct_schema.add_field("test_field", scalar_type, **scalar_params) schema.add_field( - "structA", datatype=DataType.ARRAY, + "structA", + datatype=DataType.ARRAY, element_type=DataType.STRUCT, struct_schema=struct_schema, max_capacity=default_capacity, ) index_params = client.prepare_index_params() - index_params.add_index(field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) - index_params.add_index(field_name="structA[embedding]", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) + index_params.add_index( + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + index_params.add_index( + field_name="structA[embedding]", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) index_params.add_index(field_name="structA[test_field]", index_type="INVERTED") res, check = self.create_collection(client, collection_name, schema=schema, index_params=index_params) @@ -2536,15 +2765,19 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): num_elems = rng.randint(2, 5) struct_array = [] for j in range(num_elems): - struct_array.append({ - "embedding": _seed_vector(i * 100 + j), - "test_field": gen_val(i, j), - }) - data.append({ - "id": i, - "normal_vector": _seed_vector(i + 999999), - "structA": struct_array, - }) + struct_array.append( + { + "embedding": _seed_vector(i * 100 + j), + "test_field": gen_val(i, j), + } + ) + data.append( + { + "id": i, + "normal_vector": _seed_vector(i + 999999), + "structA": struct_array, + } + ) self.insert(client, collection_name, data) self.flush(client, collection_name) self.load_collection(client, collection_name) @@ -2552,8 +2785,11 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): # Verify index is functional with a query results, check = self.query( - client, collection_name, filter="id >= 0", - output_fields=["id"], limit=10, + client, + collection_name, + filter="id >= 0", + output_fields=["id"], + limit=10, ) assert check assert len(results) > 0, "Should be able to query data after nested index creation" @@ -2570,18 +2806,33 @@ class TestMilvusClientStructArrayElementNestedIndex(TestMilvusClientV2Base): schema = self._create_schema(client) index_params = client.prepare_index_params() - index_params.add_index(field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) - index_params.add_index(field_name="structA[embedding]", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) + index_params.add_index( + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + index_params.add_index( + field_name="structA[embedding]", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) # Try INVERTED on vector sub-field - index_params.add_index(field_name="structA[embedding]", index_type="INVERTED", - index_name="invalid_inv_on_vec") + index_params.add_index( + field_name="structA[embedding]", + index_type="INVERTED", + index_name="invalid_inv_on_vec", + ) error = {ct.err_code: 65535, ct.err_msg: ""} self.create_collection( - client, collection_name, schema=schema, index_params=index_params, - check_task=CheckTasks.err_res, check_items=error, + client, + collection_name, + schema=schema, + index_params=index_params, + check_task=CheckTasks.err_res, + check_items=error, ) @@ -2601,7 +2852,8 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): struct_schema.add_field("str_val", DataType.VARCHAR, max_length=256) schema.add_field( - "structA", datatype=DataType.ARRAY, + "structA", + datatype=DataType.ARRAY, element_type=DataType.STRUCT, struct_schema=struct_schema, max_capacity=default_capacity, @@ -2652,32 +2904,48 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): num_elems = rng.randint(2, 5) struct_array = [] for j in range(num_elems): - struct_array.append({ - "embedding": self._gen_vec_by_type(vec_type, dim, seed=i * 100 + j), - "int_val": i * 10 + j, - "str_val": f"row_{i}_elem_{j}", - }) - data.append({ - "id": i, "doc_int": i, - "normal_vector": _seed_vector(i + 999999, dim), - "structA": struct_array, - }) + struct_array.append( + { + "embedding": self._gen_vec_by_type(vec_type, dim, seed=i * 100 + j), + "int_val": i * 10 + j, + "str_val": f"row_{i}_elem_{j}", + } + ) + data.append( + { + "id": i, + "doc_int": i, + "normal_vector": _seed_vector(i + 999999, dim), + "structA": struct_array, + } + ) return data - def _setup_collection(self, client, collection_name, vec_type, nb=200, dim=default_dim, - use_max_sim=True): + def _setup_collection( + self, + client, + collection_name, + vec_type, + nb=200, + dim=default_dim, + use_max_sim=True, + ): """Create collection with non-float vector type, insert data, load.""" schema = self._create_schema_with_vec_type(client, vec_type, dim) metric = self._get_max_sim_metric(vec_type) if use_max_sim else self._get_metric_for_type(vec_type) index_params = client.prepare_index_params() index_params.add_index( - field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS, + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, ) index_params.add_index( - field_name="structA[embedding]", index_type="HNSW", - metric_type=metric, params=INDEX_PARAMS, + field_name="structA[embedding]", + index_type="HNSW", + metric_type=metric, + params=INDEX_PARAMS, ) res, check = self.create_collection(client, collection_name, schema=schema, index_params=index_params) @@ -2704,10 +2972,18 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): schema = self._create_schema_with_vec_type(client, DataType.FLOAT16_VECTOR) index_params = client.prepare_index_params() - index_params.add_index(field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) - index_params.add_index(field_name="structA[embedding]", index_type="HNSW", - metric_type="MAX_SIM_L2", params=INDEX_PARAMS) + index_params.add_index( + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + index_params.add_index( + field_name="structA[embedding]", + index_type="HNSW", + metric_type="MAX_SIM_L2", + params=INDEX_PARAMS, + ) res, check = self.create_collection(client, collection_name, schema=schema, index_params=index_params) assert check @@ -2724,10 +3000,18 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): schema = self._create_schema_with_vec_type(client, DataType.BFLOAT16_VECTOR) index_params = client.prepare_index_params() - index_params.add_index(field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) - index_params.add_index(field_name="structA[embedding]", index_type="HNSW", - metric_type="MAX_SIM_IP", params=INDEX_PARAMS) + index_params.add_index( + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + index_params.add_index( + field_name="structA[embedding]", + index_type="HNSW", + metric_type="MAX_SIM_IP", + params=INDEX_PARAMS, + ) res, check = self.create_collection(client, collection_name, schema=schema, index_params=index_params) assert check @@ -2768,8 +3052,11 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): self._setup_collection(client, collection_name, DataType.FLOAT16_VECTOR, nb=200) results, check = self.query( - client, collection_name, filter="id < 10", - output_fields=["id", "structA"], limit=10, + client, + collection_name, + filter="id < 10", + output_fields=["id", "structA"], + limit=10, ) assert check assert len(results) == 10 @@ -2786,8 +3073,11 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): self._setup_collection(client, collection_name, DataType.BFLOAT16_VECTOR, nb=200) results, check = self.query( - client, collection_name, filter="id < 10", - output_fields=["id", "structA"], limit=10, + client, + collection_name, + filter="id < 10", + output_fields=["id", "structA"], + limit=10, ) assert check assert len(results) == 10 @@ -2804,8 +3094,11 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): self._setup_collection(client, collection_name, DataType.INT8_VECTOR, nb=200) results, check = self.query( - client, collection_name, filter="id < 10", - output_fields=["id", "structA"], limit=10, + client, + collection_name, + filter="id < 10", + output_fields=["id", "structA"], + limit=10, ) assert check assert len(results) == 10 @@ -2822,8 +3115,11 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): self._setup_collection(client, collection_name, DataType.BINARY_VECTOR, nb=200) results, check = self.query( - client, collection_name, filter="id < 10", - output_fields=["id", "structA"], limit=10, + client, + collection_name, + filter="id < 10", + output_fields=["id", "structA"], + limit=10, ) assert check assert len(results) == 10 @@ -2843,7 +3139,8 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): search_tensor.add(self._gen_vec_by_type(DataType.FLOAT16_VECTOR, default_dim, seed=42)) results, check = self.search( - client, collection_name, + client, + collection_name, data=[search_tensor], anns_field="structA[embedding]", search_params={"metric_type": "MAX_SIM_L2"}, @@ -2867,7 +3164,8 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): search_tensor.add(self._gen_vec_by_type(DataType.BFLOAT16_VECTOR, default_dim, seed=42)) results, check = self.search( - client, collection_name, + client, + collection_name, data=[search_tensor], anns_field="structA[embedding]", search_params={"metric_type": "MAX_SIM_IP"}, @@ -2891,7 +3189,8 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): search_tensor.add(self._gen_vec_by_type(DataType.INT8_VECTOR, default_dim, seed=42)) results, check = self.search( - client, collection_name, + client, + collection_name, data=[search_tensor], anns_field="structA[embedding]", search_params={"metric_type": "MAX_SIM_COSINE"}, @@ -2915,7 +3214,8 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): search_tensor.add(self._gen_vec_by_type(DataType.BINARY_VECTOR, default_dim, seed=42)) results, check = self.search( - client, collection_name, + client, + collection_name, data=[search_tensor], anns_field="structA[embedding]", search_params={"metric_type": "MAX_SIM_HAMMING"}, @@ -2933,16 +3233,16 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): """ client = self._client() collection_name = cf.gen_unique_str(f"{prefix}_nf_f16_ef") - data = self._setup_collection(client, collection_name, DataType.FLOAT16_VECTOR, nb=200, - use_max_sim=False) + data = self._setup_collection(client, collection_name, DataType.FLOAT16_VECTOR, nb=200, use_max_sim=False) query_vec = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vec], anns_field="structA[embedding]", search_params={"metric_type": "L2"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=10, ) assert check @@ -2957,16 +3257,16 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): """ client = self._client() collection_name = cf.gen_unique_str(f"{prefix}_nf_bf16_ef") - data = self._setup_collection(client, collection_name, DataType.BFLOAT16_VECTOR, nb=200, - use_max_sim=False) + data = self._setup_collection(client, collection_name, DataType.BFLOAT16_VECTOR, nb=200, use_max_sim=False) query_vec = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vec], anns_field="structA[embedding]", search_params={"metric_type": "IP"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=10, ) assert check @@ -3003,7 +3303,8 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): struct_schema.add_field("sparse_vec", DataType.SPARSE_FLOAT_VECTOR) schema.add_field( - "structA", datatype=DataType.ARRAY, + "structA", + datatype=DataType.ARRAY, element_type=DataType.STRUCT, struct_schema=struct_schema, max_capacity=default_capacity, @@ -3014,8 +3315,11 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): ct.err_msg: "only fixed dimension vector types are supported", } self.create_collection( - client, collection_name, schema=schema, - check_task=CheckTasks.err_res, check_items=error, + client, + collection_name, + schema=schema, + check_task=CheckTasks.err_res, + check_items=error, ) # ---- L2 tests ---- @@ -3029,16 +3333,16 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): """ client = self._client() collection_name = cf.gen_unique_str(f"{prefix}_nf_int8_ef") - data = self._setup_collection(client, collection_name, DataType.INT8_VECTOR, nb=200, - use_max_sim=False) + data = self._setup_collection(client, collection_name, DataType.INT8_VECTOR, nb=200, use_max_sim=False) query_vec = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vec], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=10, ) assert check @@ -3053,16 +3357,16 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): """ client = self._client() collection_name = cf.gen_unique_str(f"{prefix}_nf_bin_ef") - data = self._setup_collection(client, collection_name, DataType.BINARY_VECTOR, nb=200, - use_max_sim=False) + data = self._setup_collection(client, collection_name, DataType.BINARY_VECTOR, nb=200, use_max_sim=False) query_vec = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vec], anns_field="structA[embedding]", search_params={"metric_type": "HAMMING"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=10, ) assert check @@ -3088,19 +3392,32 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): struct_schema.add_field("int_val", DataType.INT64) schema.add_field( - "structA", datatype=DataType.ARRAY, + "structA", + datatype=DataType.ARRAY, element_type=DataType.STRUCT, struct_schema=struct_schema, max_capacity=default_capacity, ) index_params = client.prepare_index_params() - index_params.add_index(field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) - index_params.add_index(field_name="structA[float_emb]", index_type="HNSW", - metric_type="MAX_SIM_COSINE", params=INDEX_PARAMS) - index_params.add_index(field_name="structA[float16_emb]", index_type="HNSW", - metric_type="MAX_SIM_L2", params=INDEX_PARAMS) + index_params.add_index( + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + index_params.add_index( + field_name="structA[float_emb]", + index_type="HNSW", + metric_type="MAX_SIM_COSINE", + params=INDEX_PARAMS, + ) + index_params.add_index( + field_name="structA[float16_emb]", + index_type="HNSW", + metric_type="MAX_SIM_L2", + params=INDEX_PARAMS, + ) res, check = self.create_collection(client, collection_name, schema=schema, index_params=index_params) assert check @@ -3111,16 +3428,20 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): num_elems = rng.randint(2, 5) struct_array = [] for j in range(num_elems): - struct_array.append({ - "float_emb": _seed_vector(i * 100 + j), - "float16_emb": _generate_float16_vector(default_dim, seed=i * 100 + j), - "int_val": i * 10 + j, - }) - data.append({ - "id": i, - "normal_vector": _seed_vector(i + 999999), - "structA": struct_array, - }) + struct_array.append( + { + "float_emb": _seed_vector(i * 100 + j), + "float16_emb": _generate_float16_vector(default_dim, seed=i * 100 + j), + "int_val": i * 10 + j, + } + ) + data.append( + { + "id": i, + "normal_vector": _seed_vector(i + 999999), + "structA": struct_array, + } + ) res, check = self.insert(client, collection_name, data) assert check self.flush(client, collection_name) @@ -3130,7 +3451,8 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): search_tensor = EmbeddingList() search_tensor.add(_seed_vector(42)) results, check = self.search( - client, collection_name, + client, + collection_name, data=[search_tensor], anns_field="structA[float_emb]", search_params={"metric_type": "MAX_SIM_COSINE"}, @@ -3139,6 +3461,7 @@ class TestMilvusClientStructArrayElementNonFloatVectors(TestMilvusClientV2Base): assert check assert len(results) > 0 + class TestMilvusClientStructArrayElementHybridSearch(TestMilvusClientV2Base): """Test element_filter search + hybrid search (4 cases)""" @@ -3155,7 +3478,8 @@ class TestMilvusClientStructArrayElementHybridSearch(TestMilvusClientV2Base): struct_schema.add_field("color", DataType.VARCHAR, max_length=128) schema.add_field( - "structA", datatype=DataType.ARRAY, + "structA", + datatype=DataType.ARRAY, element_type=DataType.STRUCT, struct_schema=struct_schema, max_capacity=default_capacity, @@ -3174,27 +3498,46 @@ class TestMilvusClientStructArrayElementHybridSearch(TestMilvusClientV2Base): emb = [1.0] + [0.0] * (dim - 1) else: emb = _seed_vector(i * 1000 + j, dim) - struct_array.append({ - "embedding": emb, - "int_val": i * 100 + j, - "color": COLORS[j % 3], - }) - data.append({ - "id": i, "doc_int": i, - "normal_vector": _seed_vector(i + 999999, dim), - "structA": struct_array, - }) + struct_array.append( + { + "embedding": emb, + "int_val": i * 100 + j, + "color": COLORS[j % 3], + } + ) + data.append( + { + "id": i, + "doc_int": i, + "normal_vector": _seed_vector(i + 999999, dim), + "structA": struct_array, + } + ) return data - def _setup_collection(self, client, collection_name, nb=500, dim=default_dim, - struct_metric_type="COSINE"): + def _setup_collection( + self, + client, + collection_name, + nb=500, + dim=default_dim, + struct_metric_type="COSINE", + ): """Helper to setup collection for hybrid search.""" schema = self._create_schema(client, dim) index_params = client.prepare_index_params() - index_params.add_index(field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) - index_params.add_index(field_name="structA[embedding]", index_type="HNSW", - metric_type=struct_metric_type, params=INDEX_PARAMS) + index_params.add_index( + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + index_params.add_index( + field_name="structA[embedding]", + index_type="HNSW", + metric_type=struct_metric_type, + params=INDEX_PARAMS, + ) self.create_collection(client, collection_name, schema=schema, index_params=index_params) data = self._generate_data(nb=nb, dim=dim) @@ -3217,25 +3560,32 @@ class TestMilvusClientStructArrayElementHybridSearch(TestMilvusClientV2Base): data = self._setup_collection(client, collection_name, nb=500) # Request 1: element_filter search on struct vector - req1 = AnnSearchRequest(**{ - "data": [data[0]["structA"][0]["embedding"]], - "anns_field": "structA[embedding]", - "param": {"metric_type": "COSINE"}, - "limit": 10, - "expr": 'element_filter(structA, $[int_val] >= 0)', - }) + req1 = AnnSearchRequest( + **{ + "data": [data[0]["structA"][0]["embedding"]], + "anns_field": "structA[embedding]", + "param": {"metric_type": "COSINE"}, + "limit": 10, + "expr": "element_filter(structA, $[int_val] >= 0)", + } + ) # Request 2: normal vector search - req2 = AnnSearchRequest(**{ - "data": [data[0]["normal_vector"]], - "anns_field": "normal_vector", - "param": {"metric_type": "COSINE"}, - "limit": 10, - }) + req2 = AnnSearchRequest( + **{ + "data": [data[0]["normal_vector"]], + "anns_field": "normal_vector", + "param": {"metric_type": "COSINE"}, + "limit": 10, + } + ) results, check = self.hybrid_search( - client, collection_name, [req1, req2], - ranker=RRFRanker(), limit=10, + client, + collection_name, + [req1, req2], + ranker=RRFRanker(), + limit=10, output_fields=["id"], ) assert check @@ -3252,24 +3602,31 @@ class TestMilvusClientStructArrayElementHybridSearch(TestMilvusClientV2Base): collection_name = cf.gen_unique_str(f"{prefix}_hyb_weighted") data = self._setup_collection(client, collection_name, nb=500) - req1 = AnnSearchRequest(**{ - "data": [data[0]["structA"][0]["embedding"]], - "anns_field": "structA[embedding]", - "param": {"metric_type": "COSINE"}, - "limit": 10, - "expr": 'element_filter(structA, $[int_val] >= 0)', - }) + req1 = AnnSearchRequest( + **{ + "data": [data[0]["structA"][0]["embedding"]], + "anns_field": "structA[embedding]", + "param": {"metric_type": "COSINE"}, + "limit": 10, + "expr": "element_filter(structA, $[int_val] >= 0)", + } + ) - req2 = AnnSearchRequest(**{ - "data": [data[0]["normal_vector"]], - "anns_field": "normal_vector", - "param": {"metric_type": "COSINE"}, - "limit": 10, - }) + req2 = AnnSearchRequest( + **{ + "data": [data[0]["normal_vector"]], + "anns_field": "normal_vector", + "param": {"metric_type": "COSINE"}, + "limit": 10, + } + ) results, check = self.hybrid_search( - client, collection_name, [req1, req2], - ranker=WeightedRanker(0.7, 0.3), limit=10, + client, + collection_name, + [req1, req2], + ranker=WeightedRanker(0.7, 0.3), + limit=10, output_fields=["id"], ) assert check @@ -3277,7 +3634,9 @@ class TestMilvusClientStructArrayElementHybridSearch(TestMilvusClientV2Base): # ---- L2 tests ---- - @pytest.mark.xfail(reason="element_filter(COSINE) + EmbeddingList(MAX_SIM) require different index metrics on same field") + @pytest.mark.xfail( + reason="element_filter(COSINE) + EmbeddingList(MAX_SIM) require different index metrics on same field" + ) @pytest.mark.tags(CaseLabel.L2) def test_hybrid_element_filter_and_max_sim(self): """ @@ -3287,31 +3646,37 @@ class TestMilvusClientStructArrayElementHybridSearch(TestMilvusClientV2Base): """ client = self._client() collection_name = cf.gen_unique_str(f"{prefix}_hyb_maxsim") - data = self._setup_collection(client, collection_name, nb=500, - struct_metric_type="MAX_SIM_COSINE") + data = self._setup_collection(client, collection_name, nb=500, struct_metric_type="MAX_SIM_COSINE") # element_filter search (single vector) - req1 = AnnSearchRequest(**{ - "data": [data[0]["structA"][0]["embedding"]], - "anns_field": "structA[embedding]", - "param": {"metric_type": "COSINE"}, - "limit": 10, - "expr": 'element_filter(structA, $[color] == "Red")', - }) + req1 = AnnSearchRequest( + **{ + "data": [data[0]["structA"][0]["embedding"]], + "anns_field": "structA[embedding]", + "param": {"metric_type": "COSINE"}, + "limit": 10, + "expr": 'element_filter(structA, $[color] == "Red")', + } + ) # MAX_SIM search (EmbeddingList) search_tensor = EmbeddingList() search_tensor.add(_seed_vector(42)) - req2 = AnnSearchRequest(**{ - "data": [search_tensor], - "anns_field": "structA[embedding]", - "param": {"metric_type": "MAX_SIM_COSINE"}, - "limit": 10, - }) + req2 = AnnSearchRequest( + **{ + "data": [search_tensor], + "anns_field": "structA[embedding]", + "param": {"metric_type": "MAX_SIM_COSINE"}, + "limit": 10, + } + ) results, check = self.hybrid_search( - client, collection_name, [req1, req2], - ranker=RRFRanker(), limit=10, + client, + collection_name, + [req1, req2], + ranker=RRFRanker(), + limit=10, output_fields=["id"], ) assert check @@ -3326,29 +3691,35 @@ class TestMilvusClientStructArrayElementHybridSearch(TestMilvusClientV2Base): """ client = self._client() collection_name = cf.gen_unique_str(f"{prefix}_hyb_match") - data = self._setup_collection(client, collection_name, nb=500, - struct_metric_type="MAX_SIM_COSINE") + data = self._setup_collection(client, collection_name, nb=500, struct_metric_type="MAX_SIM_COSINE") - req1 = AnnSearchRequest(**{ - "data": [data[0]["normal_vector"]], - "anns_field": "normal_vector", - "param": {"metric_type": "COSINE"}, - "limit": 10, - "expr": 'MATCH_ANY(structA, $[color] == "Red")', - }) + req1 = AnnSearchRequest( + **{ + "data": [data[0]["normal_vector"]], + "anns_field": "normal_vector", + "param": {"metric_type": "COSINE"}, + "limit": 10, + "expr": 'MATCH_ANY(structA, $[color] == "Red")', + } + ) search_tensor = EmbeddingList() search_tensor.add(_seed_vector(42)) - req2 = AnnSearchRequest(**{ - "data": [search_tensor], - "anns_field": "structA[embedding]", - "param": {"metric_type": "MAX_SIM_COSINE"}, - "limit": 10, - }) + req2 = AnnSearchRequest( + **{ + "data": [search_tensor], + "anns_field": "structA[embedding]", + "param": {"metric_type": "MAX_SIM_COSINE"}, + "limit": 10, + } + ) results, check = self.hybrid_search( - client, collection_name, [req1, req2], - ranker=RRFRanker(), limit=10, + client, + collection_name, + [req1, req2], + ranker=RRFRanker(), + limit=10, output_fields=["id"], ) assert check @@ -3373,17 +3744,26 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): struct_schema.add_field("size", DataType.VARCHAR, max_length=128) schema.add_field( - "structA", datatype=DataType.ARRAY, + "structA", + datatype=DataType.ARRAY, element_type=DataType.STRUCT, struct_schema=struct_schema, max_capacity=default_capacity, ) index_params = client.prepare_index_params() - index_params.add_index(field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) - index_params.add_index(field_name="structA[embedding]", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) + index_params.add_index( + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + index_params.add_index( + field_name="structA[embedding]", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) self.create_collection(client, collection_name, schema=schema, index_params=index_params) data = [] @@ -3392,18 +3772,23 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): num_elems = rng.randint(3, 8) struct_array = [] for j in range(num_elems): - struct_array.append({ - "embedding": _seed_vector(i * 100 + j), - "int_val": i * 100 + j, - "str_val": f"row_{i}_elem_{j}", - "color": COLORS[j % 3], - "size": SIZES[(i + j) % 4], - }) - data.append({ - "id": i, "doc_int": i, - "normal_vector": _seed_vector(i + 999999), - "structA": struct_array, - }) + struct_array.append( + { + "embedding": _seed_vector(i * 100 + j), + "int_val": i * 100 + j, + "str_val": f"row_{i}_elem_{j}", + "color": COLORS[j % 3], + "size": SIZES[(i + j) % 4], + } + ) + data.append( + { + "id": i, + "doc_int": i, + "normal_vector": _seed_vector(i + 999999), + "structA": struct_array, + } + ) self.insert(client, collection_name, data) self.flush(client, collection_name) self.load_collection(client, collection_name) @@ -3423,13 +3808,15 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] error = {ct.err_code: 65535, ct.err_msg: ""} self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, element_filter(structA, $[int_val] > 0))', + filter="element_filter(structA, element_filter(structA, $[int_val] > 0))", limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3446,13 +3833,15 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] error = {ct.err_code: 65535, ct.err_msg: ""} self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] > 0) && doc_int > 10', + filter="element_filter(structA, $[int_val] > 0) && doc_int > 10", limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3469,13 +3858,15 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] error = {ct.err_code: 65535, ct.err_msg: ""} self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] > 0) || doc_int > 10', + filter="element_filter(structA, $[int_val] > 0) || doc_int > 10", limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3492,13 +3883,15 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] error = {ct.err_code: 65535, ct.err_msg: ""} self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(nonexistent_field, $[int_val] > 0)', + filter="element_filter(nonexistent_field, $[int_val] > 0)", limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3515,13 +3908,15 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] error = {ct.err_code: 65535, ct.err_msg: ""} self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[nonexistent_field] > 0)', + filter="element_filter(structA, $[nonexistent_field] > 0)", limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3538,13 +3933,15 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] error = {ct.err_code: 65535, ct.err_msg: ""} self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter()', + filter="element_filter()", limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3560,11 +3957,13 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): error = {ct.err_code: 65535, ct.err_msg: ""} self.query( - client, collection_name, - filter='$[int_val] > 1', + client, + collection_name, + filter="$[int_val] > 1", output_fields=["id"], limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3581,13 +3980,15 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] error = {ct.err_code: 65535, ct.err_msg: "metric type not match"} self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "MAX_SIM_COSINE"}, - filter='element_filter(structA, $[int_val] > 0)', + filter="element_filter(structA, $[int_val] > 0)", limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3605,13 +4006,15 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): search_tensor.add(_seed_vector(42)) error = {ct.err_code: 65535, ct.err_msg: ""} self.search( - client, collection_name, + client, + collection_name, data=[search_tensor], anns_field="structA[embedding]", search_params={"metric_type": "MAX_SIM_COSINE"}, - filter='element_filter(structA, $[int_val] > 0)', + filter="element_filter(structA, $[int_val] > 0)", limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3627,11 +4030,13 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): error = {ct.err_code: 65535, ct.err_msg: ""} self.query( - client, collection_name, - filter='MATCH_ANY(structA, MATCH_ALL(structA, $[int_val] > 0))', + client, + collection_name, + filter="MATCH_ANY(structA, MATCH_ALL(structA, $[int_val] > 0))", output_fields=["id"], limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3647,11 +4052,13 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): error = {ct.err_code: 65535, ct.err_msg: ""} self.query( - client, collection_name, - filter='MATCH_LEAST(structA, $[int_val] > 0, threshold=-1)', + client, + collection_name, + filter="MATCH_LEAST(structA, $[int_val] > 0, threshold=-1)", output_fields=["id"], limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3667,11 +4074,13 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): error = {ct.err_code: 65535, ct.err_msg: ""} self.query( - client, collection_name, - filter='MATCH_LEAST(structA, $[int_val] > 0, threshold=0)', + client, + collection_name, + filter="MATCH_LEAST(structA, $[int_val] > 0, threshold=0)", output_fields=["id"], limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3687,11 +4096,13 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): error = {ct.err_code: 65535, ct.err_msg: ""} self.query( - client, collection_name, - filter='MATCH_MOST(structA, $[int_val] > 0, threshold=-1)', + client, + collection_name, + filter="MATCH_MOST(structA, $[int_val] > 0, threshold=-1)", output_fields=["id"], limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3707,11 +4118,13 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): error = {ct.err_code: 65535, ct.err_msg: ""} self.query( - client, collection_name, - filter='MATCH_EXACT(structA, $[int_val] > 0, threshold=-1)', + client, + collection_name, + filter="MATCH_EXACT(structA, $[int_val] > 0, threshold=-1)", output_fields=["id"], limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3727,11 +4140,13 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): error = {ct.err_code: 65535, ct.err_msg: ""} self.query( - client, collection_name, - filter='MATCH_ANY(nonexistent_struct, $[int_val] > 0)', + client, + collection_name, + filter="MATCH_ANY(nonexistent_struct, $[int_val] > 0)", output_fields=["id"], limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3747,11 +4162,13 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): error = {ct.err_code: 65535, ct.err_msg: ""} self.query( - client, collection_name, - filter='MATCH_ANY(structA, $[nonexistent_field] > 0)', + client, + collection_name, + filter="MATCH_ANY(structA, $[nonexistent_field] > 0)", output_fields=["id"], limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3771,21 +4188,37 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): struct_schema = client.create_struct_field_schema() struct_schema.add_field("embedding", DataType.FLOAT_VECTOR, dim=default_dim) struct_schema.add_field("int_val", DataType.INT64) - schema.add_field("structA", datatype=DataType.ARRAY, - element_type=DataType.STRUCT, - struct_schema=struct_schema, max_capacity=10) + schema.add_field( + "structA", + datatype=DataType.ARRAY, + element_type=DataType.STRUCT, + struct_schema=struct_schema, + max_capacity=10, + ) index_params = client.prepare_index_params() - index_params.add_index(field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) - index_params.add_index(field_name="structA[embedding]", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) + index_params.add_index( + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + index_params.add_index( + field_name="structA[embedding]", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) index_params.add_index(field_name="structA[nonexistent]", index_type="INVERTED") error = {ct.err_code: 65535, ct.err_msg: ""} self.create_collection( - client, collection_name, schema=schema, index_params=index_params, - check_task=CheckTasks.err_res, check_items=error, + client, + collection_name, + schema=schema, + index_params=index_params, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3805,22 +4238,41 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): struct_schema = client.create_struct_field_schema() struct_schema.add_field("embedding", DataType.FLOAT_VECTOR, dim=default_dim) struct_schema.add_field("int_val", DataType.INT64) - schema.add_field("structA", datatype=DataType.ARRAY, - element_type=DataType.STRUCT, - struct_schema=struct_schema, max_capacity=10) + schema.add_field( + "structA", + datatype=DataType.ARRAY, + element_type=DataType.STRUCT, + struct_schema=struct_schema, + max_capacity=10, + ) index_params = client.prepare_index_params() - index_params.add_index(field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) - index_params.add_index(field_name="structA[embedding]", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) - index_params.add_index(field_name="structA[embedding]", index_type="STL_SORT", - index_name="invalid_stl_on_vec") + index_params.add_index( + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + index_params.add_index( + field_name="structA[embedding]", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + index_params.add_index( + field_name="structA[embedding]", + index_type="STL_SORT", + index_name="invalid_stl_on_vec", + ) error = {ct.err_code: 65535, ct.err_msg: ""} self.create_collection( - client, collection_name, schema=schema, index_params=index_params, - check_task=CheckTasks.err_res, check_items=error, + client, + collection_name, + schema=schema, + index_params=index_params, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3837,13 +4289,15 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] error = {ct.err_code: 65535, ct.err_msg: ""} self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA)', + filter="element_filter(structA)", limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @pytest.mark.tags(CaseLabel.L2) @@ -3859,11 +4313,13 @@ class TestMilvusClientStructArrayElementSearchInvalid(TestMilvusClientV2Base): error = {ct.err_code: 65535, ct.err_msg: ""} self.query( - client, collection_name, - filter='MATCH_ALL(structA)', + client, + collection_name, + filter="MATCH_ALL(structA)", output_fields=["id"], limit=10, - check_task=CheckTasks.err_res, check_items=error, + check_task=CheckTasks.err_res, + check_items=error, ) @@ -3884,7 +4340,8 @@ class TestMilvusClientStructArrayElementSearchCRUD(TestMilvusClientV2Base): struct_schema.add_field("color", DataType.VARCHAR, max_length=128) schema.add_field( - "structA", datatype=DataType.ARRAY, + "structA", + datatype=DataType.ARRAY, element_type=DataType.STRUCT, struct_schema=struct_schema, max_capacity=default_capacity, @@ -3895,10 +4352,18 @@ class TestMilvusClientStructArrayElementSearchCRUD(TestMilvusClientV2Base): """Helper: create collection, insert data, index, load.""" schema = self._create_schema(client, dim) index_params = client.prepare_index_params() - index_params.add_index(field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) - index_params.add_index(field_name="structA[embedding]", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) + index_params.add_index( + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + index_params.add_index( + field_name="structA[embedding]", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) self.create_collection(client, collection_name, schema=schema, index_params=index_params) data = [] @@ -3907,17 +4372,22 @@ class TestMilvusClientStructArrayElementSearchCRUD(TestMilvusClientV2Base): num_elems = rng.randint(3, 8) struct_array = [] for j in range(num_elems): - struct_array.append({ - "embedding": _seed_vector(i * 1000 + j, dim), - "int_val": i * 100 + j, - "str_val": f"row_{i}_elem_{j}", - "color": COLORS[j % 3], - }) - data.append({ - "id": i, "doc_int": i, - "normal_vector": _seed_vector(i + 999999, dim), - "structA": struct_array, - }) + struct_array.append( + { + "embedding": _seed_vector(i * 1000 + j, dim), + "int_val": i * 100 + j, + "str_val": f"row_{i}_elem_{j}", + "color": COLORS[j % 3], + } + ) + data.append( + { + "id": i, + "doc_int": i, + "normal_vector": _seed_vector(i + 999999, dim), + "structA": struct_array, + } + ) self.insert(client, collection_name, data) self.flush(client, collection_name) self.load_collection(client, collection_name) @@ -3937,23 +4407,27 @@ class TestMilvusClientStructArrayElementSearchCRUD(TestMilvusClientV2Base): # Upsert rows 0-9 with color="Purple" upsert_data = [] for i in range(10): - upsert_data.append({ - "id": i, "doc_int": i, - "normal_vector": _seed_vector(i + 999999), - "structA": [ - { - "embedding": _seed_vector(i * 1000), - "int_val": i * 100 + 9999, - "str_val": f"upserted_{i}", - "color": "Purple", - } - ], - }) + upsert_data.append( + { + "id": i, + "doc_int": i, + "normal_vector": _seed_vector(i + 999999), + "structA": [ + { + "embedding": _seed_vector(i * 1000), + "int_val": i * 100 + 9999, + "str_val": f"upserted_{i}", + "color": "Purple", + } + ], + } + ) self.upsert(client, collection_name, upsert_data) self.flush(client, collection_name) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ANY(structA, $[color] == "Purple")', output_fields=["id", "structA"], limit=100, @@ -3979,11 +4453,12 @@ class TestMilvusClientStructArrayElementSearchCRUD(TestMilvusClientV2Base): query_vector = data[15]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=100, output_fields=["id"], ) @@ -4006,8 +4481,9 @@ class TestMilvusClientStructArrayElementSearchCRUD(TestMilvusClientV2Base): self.delete(client, collection_name, ids=list(range(10))) results, check = self.query( - client, collection_name, - filter='MATCH_ANY(structA, $[int_val] >= 0)', + client, + collection_name, + filter="MATCH_ANY(structA, $[int_val] >= 0)", output_fields=["id"], limit=500, ) @@ -4028,24 +4504,28 @@ class TestMilvusClientStructArrayElementSearchCRUD(TestMilvusClientV2Base): self._setup_collection(client, collection_name, nb=500) # Upsert row 50 with known struct - upsert_data = [{ - "id": 50, "doc_int": 50, - "normal_vector": _seed_vector(50 + 999999), - "structA": [ - { - "embedding": _seed_vector(99999), - "int_val": 99999, - "str_val": "upserted_target", - "color": "Gold", - } - ], - }] + upsert_data = [ + { + "id": 50, + "doc_int": 50, + "normal_vector": _seed_vector(50 + 999999), + "structA": [ + { + "embedding": _seed_vector(99999), + "int_val": 99999, + "str_val": "upserted_target", + "color": "Gold", + } + ], + } + ] self.upsert(client, collection_name, upsert_data) self.flush(client, collection_name) query_vector = _seed_vector(99999) results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -4075,25 +4555,29 @@ class TestMilvusClientStructArrayElementSearchCRUD(TestMilvusClientV2Base): # Re-insert new data new_data = [] for i in range(100): - new_data.append({ - "id": i + 10000, "doc_int": i + 10000, - "normal_vector": _seed_vector(i + 888888), - "structA": [ - { - "embedding": _seed_vector(i * 500), - "int_val": i + 50000, - "str_val": f"new_{i}", - "color": "Silver", - } - ], - }) + new_data.append( + { + "id": i + 10000, + "doc_int": i + 10000, + "normal_vector": _seed_vector(i + 888888), + "structA": [ + { + "embedding": _seed_vector(i * 500), + "int_val": i + 50000, + "str_val": f"new_{i}", + "color": "Silver", + } + ], + } + ) self.insert(client, collection_name, new_data) self.flush(client, collection_name) self.load_collection(client, collection_name) query_vector = new_data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -4123,17 +4607,26 @@ class TestMilvusClientStructArrayElementSearchIterator(TestMilvusClientV2Base): struct_schema.add_field("color", DataType.VARCHAR, max_length=128) schema.add_field( - "structA", datatype=DataType.ARRAY, + "structA", + datatype=DataType.ARRAY, element_type=DataType.STRUCT, struct_schema=struct_schema, max_capacity=default_capacity, ) index_params = client.prepare_index_params() - index_params.add_index(field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) - index_params.add_index(field_name="structA[embedding]", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) + index_params.add_index( + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + index_params.add_index( + field_name="structA[embedding]", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) self.create_collection(client, collection_name, schema=schema, index_params=index_params) data = [] @@ -4142,16 +4635,21 @@ class TestMilvusClientStructArrayElementSearchIterator(TestMilvusClientV2Base): num_elems = rng.randint(3, 8) struct_array = [] for j in range(num_elems): - struct_array.append({ - "embedding": _seed_vector(i * 1000 + j), - "int_val": i * 100 + j, - "color": COLORS[j % 3], - }) - data.append({ - "id": i, "doc_int": i, - "normal_vector": _seed_vector(i + 999999), - "structA": struct_array, - }) + struct_array.append( + { + "embedding": _seed_vector(i * 1000 + j), + "int_val": i * 100 + j, + "color": COLORS[j % 3], + } + ) + data.append( + { + "id": i, + "doc_int": i, + "normal_vector": _seed_vector(i + 999999), + "structA": struct_array, + } + ) self.insert(client, collection_name, data) self.flush(client, collection_name) self.load_collection(client, collection_name) @@ -4171,12 +4669,13 @@ class TestMilvusClientStructArrayElementSearchIterator(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] iterator, _ = self.search_iterator( - client, collection_name, + client, + collection_name, data=[query_vector], batch_size=50, anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=200, output_fields=["id"], ) @@ -4206,12 +4705,13 @@ class TestMilvusClientStructArrayElementSearchIterator(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] iterator, _ = self.search_iterator( - client, collection_name, + client, + collection_name, data=[query_vector], batch_size=batch_size, anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=200, output_fields=["id"], ) @@ -4247,52 +4747,80 @@ class TestMilvusClientStructArrayElementSearchIterator(TestMilvusClientV2Base): struct_schema.add_field("color", DataType.VARCHAR, max_length=128) schema.add_field( - "structA", datatype=DataType.ARRAY, + "structA", + datatype=DataType.ARRAY, element_type=DataType.STRUCT, struct_schema=struct_schema, max_capacity=default_capacity, ) index_params = client.prepare_index_params() - index_params.add_index(field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) - index_params.add_index(field_name="structA[embedding]", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) + index_params.add_index( + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + index_params.add_index( + field_name="structA[embedding]", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) self.create_collection(client, collection_name, schema=schema, index_params=index_params) # Sealed segment sealed_data = [] for i in range(300): random.Random(i) - struct_array = [{"embedding": _seed_vector(i * 1000), "int_val": i * 100, "color": "Red"}] - sealed_data.append({ - "id": i, "doc_int": i, - "normal_vector": _seed_vector(i + 999999), - "structA": struct_array, - }) + struct_array = [ + { + "embedding": _seed_vector(i * 1000), + "int_val": i * 100, + "color": "Red", + } + ] + sealed_data.append( + { + "id": i, + "doc_int": i, + "normal_vector": _seed_vector(i + 999999), + "structA": struct_array, + } + ) self.insert(client, collection_name, sealed_data) self.flush(client, collection_name) # Growing segment growing_data = [] for i in range(300, 500): - struct_array = [{"embedding": _seed_vector(i * 1000), "int_val": i * 100, "color": "Blue"}] - growing_data.append({ - "id": i, "doc_int": i, - "normal_vector": _seed_vector(i + 999999), - "structA": struct_array, - }) + struct_array = [ + { + "embedding": _seed_vector(i * 1000), + "int_val": i * 100, + "color": "Blue", + } + ] + growing_data.append( + { + "id": i, + "doc_int": i, + "normal_vector": _seed_vector(i + 999999), + "structA": struct_array, + } + ) self.insert(client, collection_name, growing_data) self.load_collection(client, collection_name) query_vector = sealed_data[0]["structA"][0]["embedding"] iterator, _ = self.search_iterator( - client, collection_name, + client, + collection_name, data=[query_vector], batch_size=50, anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=400, output_fields=["id"], ) @@ -4390,17 +4918,26 @@ class TestMilvusClientStructArrayElementSearchMmap(TestMilvusClientV2Base): struct_schema.add_field("color", DataType.VARCHAR, max_length=128) schema.add_field( - "structA", datatype=DataType.ARRAY, + "structA", + datatype=DataType.ARRAY, element_type=DataType.STRUCT, struct_schema=struct_schema, max_capacity=default_capacity, ) index_params = client.prepare_index_params() - index_params.add_index(field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) - index_params.add_index(field_name="structA[embedding]", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) + index_params.add_index( + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + index_params.add_index( + field_name="structA[embedding]", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) self.create_collection(client, collection_name, schema=schema, index_params=index_params) data = [] @@ -4409,16 +4946,21 @@ class TestMilvusClientStructArrayElementSearchMmap(TestMilvusClientV2Base): num_elems = rng.randint(3, 8) struct_array = [] for j in range(num_elems): - struct_array.append({ - "embedding": _seed_vector(i * 1000 + j), - "int_val": i * 100 + j, - "color": COLORS[j % 3], - }) - data.append({ - "id": i, "doc_int": i, - "normal_vector": _seed_vector(i + 999999), - "structA": struct_array, - }) + struct_array.append( + { + "embedding": _seed_vector(i * 1000 + j), + "int_val": i * 100 + j, + "color": COLORS[j % 3], + } + ) + data.append( + { + "id": i, + "doc_int": i, + "normal_vector": _seed_vector(i + 999999), + "structA": struct_array, + } + ) self.insert(client, collection_name, data) self.flush(client, collection_name) return data @@ -4437,18 +4979,20 @@ class TestMilvusClientStructArrayElementSearchMmap(TestMilvusClientV2Base): # Release before enabling mmap (cannot alter mmap on loaded collection) self.release_collection(client, collection_name) self.alter_collection_properties( - client, collection_name, + client, + collection_name, properties={"mmap.enabled": True}, ) self.load_collection(client, collection_name) query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] >= 0)', + filter="element_filter(structA, $[int_val] >= 0)", limit=10, output_fields=["id"], ) @@ -4469,13 +5013,15 @@ class TestMilvusClientStructArrayElementSearchMmap(TestMilvusClientV2Base): # Release before enabling mmap (cannot alter mmap on loaded collection) self.release_collection(client, collection_name) self.alter_collection_properties( - client, collection_name, + client, + collection_name, properties={"mmap.enabled": True}, ) self.load_collection(client, collection_name) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ANY(structA, $[color] == "Red")', output_fields=["id", "structA"], limit=100, @@ -4506,17 +5052,26 @@ class TestMilvusClientStructArrayElementSearchMmap(TestMilvusClientV2Base): struct_schema.add_field("color", DataType.VARCHAR, max_length=128) schema.add_field( - "structA", datatype=DataType.ARRAY, + "structA", + datatype=DataType.ARRAY, element_type=DataType.STRUCT, struct_schema=struct_schema, max_capacity=default_capacity, ) index_params = client.prepare_index_params() - index_params.add_index(field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) - index_params.add_index(field_name="structA[embedding]", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS) + index_params.add_index( + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + index_params.add_index( + field_name="structA[embedding]", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) index_params.add_index(field_name="structA[int_val]", index_type="INVERTED") index_params.add_index(field_name="structA[color]", index_type="INVERTED") @@ -4528,29 +5083,36 @@ class TestMilvusClientStructArrayElementSearchMmap(TestMilvusClientV2Base): num_elems = rng.randint(3, 8) struct_array = [] for j in range(num_elems): - struct_array.append({ - "embedding": _seed_vector(i * 1000 + j), - "int_val": i * 100 + j, - "color": COLORS[j % 3], - }) - data.append({ - "id": i, "doc_int": i, - "normal_vector": _seed_vector(i + 999999), - "structA": struct_array, - }) + struct_array.append( + { + "embedding": _seed_vector(i * 1000 + j), + "int_val": i * 100 + j, + "color": COLORS[j % 3], + } + ) + data.append( + { + "id": i, + "doc_int": i, + "normal_vector": _seed_vector(i + 999999), + "structA": struct_array, + } + ) self.insert(client, collection_name, data) self.flush(client, collection_name) # Release before enabling mmap (cannot alter mmap on loaded collection) self.release_collection(client, collection_name) self.alter_collection_properties( - client, collection_name, + client, + collection_name, properties={"mmap.enabled": True}, ) self.load_collection(client, collection_name) results, check = self.query( - client, collection_name, + client, + collection_name, filter='MATCH_ANY(structA, $[color] == "Blue" && $[int_val] > 100)', output_fields=["id", "structA"], limit=100, @@ -4558,8 +5120,10 @@ class TestMilvusClientStructArrayElementSearchMmap(TestMilvusClientV2Base): assert check assert len(results) > 0 + # ==================== Test Case 3: Group By + Element-level Search ==================== + @pytest.mark.xdist_group("TestMilvusClientStructArrayElementGroupBySearch") class TestMilvusClientStructArrayElementGroupBySearch(TestMilvusClientV2Base): """Test group_by_field with element-level search (PR #47252). @@ -4586,7 +5150,8 @@ class TestMilvusClientStructArrayElementGroupBySearch(TestMilvusClientV2Base): struct_schema.add_field("color", DataType.VARCHAR, max_length=128) schema.add_field( - "structA", datatype=DataType.ARRAY, + "structA", + datatype=DataType.ARRAY, element_type=DataType.STRUCT, struct_schema=struct_schema, max_capacity=default_capacity, @@ -4601,21 +5166,25 @@ class TestMilvusClientStructArrayElementGroupBySearch(TestMilvusClientV2Base): num_elems = rng.randint(3, 8) struct_array = [] for j in range(num_elems): - struct_array.append({ - "embedding": _seed_vector(i * 1000 + j, dim), - "int_val": i * 100 + j, - "str_val": f"row_{i}_elem_{j}", - "float_val": float(i + j * 0.1), - "color": COLORS[j % 3], - }) - data.append({ - "id": i, - "doc_int": i, - "doc_category": CATEGORIES[i % 4], - "doc_group": i % 5, - "normal_vector": _seed_vector(i + 999999, dim), - "structA": struct_array, - }) + struct_array.append( + { + "embedding": _seed_vector(i * 1000 + j, dim), + "int_val": i * 100 + j, + "str_val": f"row_{i}_elem_{j}", + "float_val": float(i + j * 0.1), + "color": COLORS[j % 3], + } + ) + data.append( + { + "id": i, + "doc_int": i, + "doc_category": CATEGORIES[i % 4], + "doc_group": i % 5, + "normal_vector": _seed_vector(i + 999999, dim), + "structA": struct_array, + } + ) return data @pytest.fixture(scope="class", autouse=True) @@ -4627,15 +5196,24 @@ class TestMilvusClientStructArrayElementGroupBySearch(TestMilvusClientV2Base): schema = self._create_schema(client) index_params = client.prepare_index_params() index_params.add_index( - field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS, + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, ) index_params.add_index( - field_name="structA[embedding]", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS, + field_name="structA[embedding]", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + self.create_collection( + client, + collection_name, + schema=schema, + index_params=index_params, + force_teardown=False, ) - self.create_collection(client, collection_name, schema=schema, - index_params=index_params, force_teardown=False) # Sealed: 3000 rows in batches data = [] @@ -4675,11 +5253,12 @@ class TestMilvusClientStructArrayElementGroupBySearch(TestMilvusClientV2Base): query_vector = data[0]["normal_vector"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="normal_vector", search_params={"metric_type": "COSINE"}, - filter='MATCH_ANY(structA, $[int_val] > 100)', + filter="MATCH_ANY(structA, $[int_val] > 100)", group_by_field="doc_category", limit=10, output_fields=["id", "doc_category", "structA"], @@ -4692,12 +5271,10 @@ class TestMilvusClientStructArrayElementGroupBySearch(TestMilvusClientV2Base): for hit in results[0]: cat = hit.get("doc_category") if cat is not None: - assert cat not in seen_categories, \ - f"Duplicate category '{cat}' in grouped results" + assert cat not in seen_categories, f"Duplicate category '{cat}' in grouped results" seen_categories.append(cat) # Verify element filter - assert any(e["int_val"] > 100 for e in hit["structA"]), \ - f"Row {hit['id']}: no element with int_val > 100" + assert any(e["int_val"] > 100 for e in hit["structA"]), f"Row {hit['id']}: no element with int_val > 100" # ---- 3.2 Group by + MATCH_ALL filter ---- @@ -4714,11 +5291,12 @@ class TestMilvusClientStructArrayElementGroupBySearch(TestMilvusClientV2Base): query_vector = data[0]["normal_vector"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="normal_vector", search_params={"metric_type": "COSINE"}, - filter='MATCH_ALL(structA, $[int_val] > 0)', + filter="MATCH_ALL(structA, $[int_val] > 0)", group_by_field="doc_category", limit=10, output_fields=["id", "doc_category", "structA"], @@ -4727,8 +5305,7 @@ class TestMilvusClientStructArrayElementGroupBySearch(TestMilvusClientV2Base): assert len(results) > 0 for hit in results[0]: - assert all(e["int_val"] > 0 for e in hit["structA"]), \ - f"Row {hit['id']}: not all elements have int_val > 0" + assert all(e["int_val"] > 0 for e in hit["structA"]), f"Row {hit['id']}: not all elements have int_val > 0" # ---- 3.3 Group by + group_size ---- @@ -4745,7 +5322,8 @@ class TestMilvusClientStructArrayElementGroupBySearch(TestMilvusClientV2Base): query_vector = data[0]["normal_vector"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="normal_vector", search_params={"metric_type": "COSINE"}, @@ -4765,8 +5343,7 @@ class TestMilvusClientStructArrayElementGroupBySearch(TestMilvusClientV2Base): if gval is not None: groups.setdefault(gval, []).append(hit) for gval, hits in groups.items(): - assert len(hits) <= 3, \ - f"Group '{gval}' has {len(hits)} results, expected <= 3" + assert len(hits) <= 3, f"Group '{gval}' has {len(hits)} results, expected <= 3" # ---- 3.4 Group by on INT32 field + struct filter ---- @@ -4783,11 +5360,12 @@ class TestMilvusClientStructArrayElementGroupBySearch(TestMilvusClientV2Base): query_vector = data[0]["normal_vector"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="normal_vector", search_params={"metric_type": "COSINE"}, - filter='MATCH_ANY(structA, $[int_val] > 0)', + filter="MATCH_ANY(structA, $[int_val] > 0)", group_by_field="doc_group", limit=10, output_fields=["id", "doc_group"], @@ -4800,8 +5378,7 @@ class TestMilvusClientStructArrayElementGroupBySearch(TestMilvusClientV2Base): for hit in results[0]: gval = hit.get("doc_group") if gval is not None: - assert gval not in seen_groups, \ - f"Duplicate doc_group '{gval}' in grouped results" + assert gval not in seen_groups, f"Duplicate doc_group '{gval}' in grouped results" seen_groups.append(gval) @pytest.mark.tags(CaseLabel.L1) @@ -4817,7 +5394,8 @@ class TestMilvusClientStructArrayElementGroupBySearch(TestMilvusClientV2Base): query_vector = data[0]["normal_vector"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="normal_vector", search_params={"metric_type": "COSINE"}, @@ -4833,8 +5411,7 @@ class TestMilvusClientStructArrayElementGroupBySearch(TestMilvusClientV2Base): for hit in results[0]: cat = hit.get("doc_category") if cat is not None: - assert cat not in seen_categories, \ - f"Duplicate doc_category '{cat}' in grouped results" + assert cat not in seen_categories, f"Duplicate doc_category '{cat}' in grouped results" seen_categories.append(cat) colors = [e["color"] for e in hit["structA"]] assert "Red" in colors @@ -4854,7 +5431,8 @@ class TestMilvusClientStructArrayElementGroupBySearch(TestMilvusClientV2Base): query_vector = data[200]["normal_vector"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="normal_vector", search_params={"metric_type": "COSINE"}, @@ -4867,9 +5445,9 @@ class TestMilvusClientStructArrayElementGroupBySearch(TestMilvusClientV2Base): assert len(results) > 0 - # ==================== Test Case 9: Element-Level Search Without Filter ==================== + def gt_element_search_no_filter(data, query_vector, metric_type="COSINE", limit=10): """Ground truth for element-level search without filter. Flattens struct array so each element competes independently by its embedding. @@ -4880,8 +5458,11 @@ def gt_element_search_no_filter(data, query_vector, metric_type="COSINE", limit= best_offset = -1 for j, elem in enumerate(row["structA"]): score = _compute_similarity(query_vector, elem["embedding"], metric_type) - if best_score is None or (_is_descending(metric_type) and score > best_score) or \ - (not _is_descending(metric_type) and score < best_score): + if ( + best_score is None + or (_is_descending(metric_type) and score > best_score) + or (not _is_descending(metric_type) and score < best_score) + ): best_score = score best_offset = j if best_offset >= 0: @@ -4918,7 +5499,8 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): struct_schema.add_field("color", DataType.VARCHAR, max_length=128) schema.add_field( - "structA", datatype=DataType.ARRAY, + "structA", + datatype=DataType.ARRAY, element_type=DataType.STRUCT, struct_schema=struct_schema, max_capacity=default_capacity, @@ -4926,34 +5508,48 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): index_params = client.prepare_index_params() index_params.add_index( - field_name="normal_vector", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS, + field_name="normal_vector", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, ) index_params.add_index( - field_name="structA[embedding]", index_type="HNSW", - metric_type="COSINE", params=INDEX_PARAMS, + field_name="structA[embedding]", + index_type="HNSW", + metric_type="COSINE", + params=INDEX_PARAMS, + ) + self.create_collection( + client, + collection_name, + schema=schema, + index_params=index_params, + force_teardown=False, ) - self.create_collection(client, collection_name, schema=schema, - index_params=index_params, force_teardown=False) def _make_rows(start, count): rows = [] for i in range(start, start + count): rng = random.Random(i) num_elems = rng.randint(2, 6) - struct_array = [{ - "embedding": _seed_vector(i * 1000 + j, default_dim), - "int_val": i * 100 + j, - "str_val": f"row_{i}_elem_{j}", - "color": COLORS[j % 3], - } for j in range(num_elems)] - rows.append({ - "id": i, - "doc_int": i, - "doc_category": CATEGORIES[i % 4], - "normal_vector": _seed_vector(i + 999999, default_dim), - "structA": struct_array, - }) + struct_array = [ + { + "embedding": _seed_vector(i * 1000 + j, default_dim), + "int_val": i * 100 + j, + "str_val": f"row_{i}_elem_{j}", + "color": COLORS[j % 3], + } + for j in range(num_elems) + ] + rows.append( + { + "id": i, + "doc_int": i, + "doc_category": CATEGORIES[i % 4], + "normal_vector": _seed_vector(i + 999999, default_dim), + "structA": struct_array, + } + ) return rows # Sealed: default_nb rows in batches @@ -4996,7 +5592,8 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -5008,13 +5605,11 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): assert len(results[0]) == 10 # First result should be the query vector itself (self-match) - assert results[0][0]["id"] == 0, \ - f"Expected row 0 as top hit (self-match), got {results[0][0]['id']}" + assert results[0][0]["id"] == 0, f"Expected row 0 as top hit (self-match), got {results[0][0]['id']}" # Distance should be ~1.0 for cosine self-match top_distance = results[0][0]["distance"] - assert abs(top_distance - 1.0) < epsilon, \ - f"Self-match distance should be ~1.0, got {top_distance}" + assert abs(top_distance - 1.0) < epsilon, f"Self-match distance should be ~1.0, got {top_distance}" # ---- 9.2 Ground truth ranking verification ---- @@ -5032,7 +5627,8 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): query_vector = data[42]["structA"][1]["embedding"] limit = 20 results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -5047,14 +5643,12 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): milvus_ids = [hit["id"] for hit in results[0]] # Top-1 must match - assert milvus_ids[0] == gt_ids[0], \ - f"Top-1 mismatch: Milvus={milvus_ids[0]}, GT={gt_ids[0]}" + assert milvus_ids[0] == gt_ids[0], f"Top-1 mismatch: Milvus={milvus_ids[0]}, GT={gt_ids[0]}" # At least 90% overlap in top-K (HNSW recall tolerance) overlap = len(set(milvus_ids) & set(gt_ids)) recall = overlap / limit - assert recall >= 0.9, \ - f"Recall too low: {recall:.2f} ({overlap}/{limit})" + assert recall >= 0.9, f"Recall too low: {recall:.2f} ({overlap}/{limit})" # ---- 9.3 Distance monotonicity ---- @@ -5072,7 +5666,8 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] limit = 50 results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -5084,8 +5679,9 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): distances = [hit["distance"] for hit in results[0]] for i in range(len(distances) - 1): - assert distances[i] >= distances[i + 1] - epsilon, \ - f"Distance not monotonic at position {i}: {distances[i]} < {distances[i+1]}" + assert distances[i] >= distances[i + 1] - epsilon, ( + f"Distance not monotonic at position {i}: {distances[i]} < {distances[i + 1]}" + ) # ---- 9.4 Offset / pagination without filter ---- @@ -5103,7 +5699,8 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): query_vector = data[5]["structA"][0]["embedding"] results_full, check1 = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -5113,7 +5710,8 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): assert check1 results_offset, check2 = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -5125,8 +5723,7 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): full_tail_ids = [hit["id"] for hit in results_full[0][10:20]] offset_ids = [hit["id"] for hit in results_offset[0]] - assert full_tail_ids == offset_ids, \ - f"Offset mismatch:\n full[10:20]={full_tail_ids}\n offset={offset_ids}" + assert full_tail_ids == offset_ids, f"Offset mismatch:\n full[10:20]={full_tail_ids}\n offset={offset_ids}" # ---- 9.5 Multiple query vectors (batch search) ---- @@ -5147,7 +5744,8 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): data[200]["structA"][0]["embedding"], ] results, check = self.search( - client, collection_name, + client, + collection_name, data=query_vectors, anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -5160,8 +5758,9 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): # Each query should find its own source row as top-1 expected_top1 = [0, 100, 200] for i, expected_id in enumerate(expected_top1): - assert results[i][0]["id"] == expected_id, \ + assert results[i][0]["id"] == expected_id, ( f"Query {i}: expected top-1 id={expected_id}, got {results[i][0]['id']}" + ) # ---- 9.6 Large limit ---- @@ -5179,7 +5778,8 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): query_vector = data[50]["structA"][0]["embedding"] limit = 200 results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -5191,14 +5791,12 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): # No duplicate IDs (each row appears at most once) ids = [hit["id"] for hit in results[0]] - assert len(ids) == len(set(ids)), \ - f"Duplicate IDs in results: {len(ids)} total, {len(set(ids))} unique" + assert len(ids) == len(set(ids)), f"Duplicate IDs in results: {len(ids)} total, {len(set(ids))} unique" # Distances monotonically decreasing distances = [hit["distance"] for hit in results[0]] for i in range(len(distances) - 1): - assert distances[i] >= distances[i + 1] - epsilon, \ - f"Distance not monotonic at position {i}" + assert distances[i] >= distances[i + 1] - epsilon, f"Distance not monotonic at position {i}" # ---- 9.7 Element-level vs normal vector search independence ---- @@ -5211,12 +5809,12 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): """ client = self.shared_client collection_name = self.shared_collection - data = self.shared_data query_vector = _seed_vector(12345, default_dim) results_elem, check1 = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -5226,7 +5824,8 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): assert check1 results_normal, check2 = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="normal_vector", search_params={"metric_type": "COSINE"}, @@ -5240,39 +5839,7 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): # Results should differ because element embeddings and normal vectors # are generated from different seeds - assert elem_ids != normal_ids, \ - "Element-level and normal vector search returned identical results" - - # ---- 9.8 Range search not supported on element-level embedding ---- - - @pytest.mark.tags(CaseLabel.L1) - def test_element_search_no_filter_range_search_not_supported(self): - """ - target: range search on structA[embedding] should be rejected - method: search with radius and range_filter params on element-level field - expected: server returns error (range search not supported for vector array) - """ - client = self.shared_client - collection_name = self.shared_collection - data = self.shared_data - - query_vector = data[0]["structA"][0]["embedding"] - radius = 0.5 - range_filter = 1.0 - results, check = self.search( - client, collection_name, - data=[query_vector], - anns_field="structA[embedding]", - search_params={ - "metric_type": "COSINE", - "params": {"radius": radius, "range_filter": range_filter}, - }, - limit=50, - output_fields=["id"], - check_task=CheckTasks.err_res, - check_items={"err_code": 1100, - "err_msg": "range search is not supported for vector array"}, - ) + assert elem_ids != normal_ids, "Element-level and normal vector search returned identical results" # ---- 9.9 Growing segment search without filter ---- @@ -5292,7 +5859,8 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): query_vector = data[growing_row_id]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -5303,8 +5871,9 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): assert len(results[0]) > 0 # Self-match should be top-1 - assert results[0][0]["id"] == growing_row_id, \ + assert results[0][0]["id"] == growing_row_id, ( f"Expected growing row {growing_row_id} as top-1, got {results[0][0]['id']}" + ) assert abs(results[0][0]["distance"] - 1.0) < epsilon # Verify results include both sealed and growing rows @@ -5331,7 +5900,8 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] limit = 20 results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -5344,8 +5914,7 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): # No duplicate PK — group_by PK ensures each row at most once ids = [hit["id"] for hit in results[0]] - assert len(ids) == len(set(ids)), \ - f"Duplicate PKs: {len(ids)} total, {len(set(ids))} unique" + assert len(ids) == len(set(ids)), f"Duplicate PKs: {len(ids)} total, {len(set(ids))} unique" # Self-match should be top-1 assert ids[0] == 0, f"Expected PK 0 as top-1 (self-match), got {ids[0]}" @@ -5353,8 +5922,9 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): # Distance monotonically decreasing distances = [hit["distance"] for hit in results[0]] for i in range(len(distances) - 1): - assert distances[i] >= distances[i + 1] - epsilon, \ - f"Distance not monotonic at position {i}: {distances[i]} < {distances[i+1]}" + assert distances[i] >= distances[i + 1] - epsilon, ( + f"Distance not monotonic at position {i}: {distances[i]} < {distances[i + 1]}" + ) # ---- 9.11 Group by PK with group_size on element-level search ---- @@ -5373,7 +5943,8 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): limit = 10 group_size = 3 results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -5391,14 +5962,12 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): groups.setdefault(hit["id"], []).append(hit) for pk, hits in groups.items(): - assert len(hits) <= group_size, \ - f"PK {pk}: {len(hits)} hits exceeds group_size={group_size}" + assert len(hits) <= group_size, f"PK {pk}: {len(hits)} hits exceeds group_size={group_size}" # Distance monotonically decreasing distances = [hit["distance"] for hit in results[0]] for i in range(len(distances) - 1): - assert distances[i] >= distances[i + 1] - epsilon, \ - f"Distance not monotonic at position {i}" + assert distances[i] >= distances[i + 1] - epsilon, f"Distance not monotonic at position {i}" # ---- 9.12 Group by PK with element_filter ---- @@ -5416,11 +5985,12 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] limit = 20 results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, - filter='element_filter(structA, $[int_val] > 500)', + filter="element_filter(structA, $[int_val] > 500)", group_by_field="id", limit=limit, output_fields=["id", "structA"], @@ -5434,8 +6004,7 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): # Every returned row must have at least one element with int_val > 500 for hit in results[0]: - assert any(e["int_val"] > 500 for e in hit["structA"]), \ - f"Row {hit['id']}: no element with int_val > 500" + assert any(e["int_val"] > 500 for e in hit["structA"]), f"Row {hit['id']}: no element with int_val > 500" # ---- 9.13 Group by PK ground truth verification ---- @@ -5453,7 +6022,8 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): query_vector = data[50]["structA"][1]["embedding"] limit = 20 results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -5470,14 +6040,12 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): milvus_ids = [hit["id"] for hit in results[0]] # Top-1 must match - assert milvus_ids[0] == gt_ids[0], \ - f"Top-1 mismatch: Milvus={milvus_ids[0]}, GT={gt_ids[0]}" + assert milvus_ids[0] == gt_ids[0], f"Top-1 mismatch: Milvus={milvus_ids[0]}, GT={gt_ids[0]}" # At least 90% recall overlap = len(set(milvus_ids) & set(gt_ids)) recall = overlap / limit - assert recall >= 0.9, \ - f"Group-by PK recall too low: {recall:.2f} ({overlap}/{limit})" + assert recall >= 0.9, f"Group-by PK recall too low: {recall:.2f} ({overlap}/{limit})" # ---- 9.14 Group by non-PK field not supported on element-level search ---- @@ -5494,7 +6062,8 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): query_vector = data[0]["structA"][0]["embedding"] results, check = self.search( - client, collection_name, + client, + collection_name, data=[query_vector], anns_field="structA[embedding]", search_params={"metric_type": "COSINE"}, @@ -5502,6 +6071,8 @@ class TestMilvusClientStructArrayElementSearchNoFilter(TestMilvusClientV2Base): limit=10, output_fields=["id"], check_task=CheckTasks.err_res, - check_items={"err_code": 1100, - "err_msg": "only group by primary key is supported for element-level search"}, + check_items={ + "err_code": 1100, + "err_msg": "only group by primary key is supported for element-level search", + }, )