mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 02:05:41 +00:00
fix: miscellaneous struct array fixes (#48349)
issue: https://github.com/milvus-io/milvus/issues/42148 Correctness: - JSON/CSV import: reject struct array elements with mismatched field count - Parquet import: error on type assertion failure in scalar and vector paths instead of silent zero-fill/data loss - JSON import: validate per-vector dimension in ArrayOfVector FloatVector path - Move element_level inference from pymilvus into C++ ParsePlaceholderGroup: pymilvus cannot reliably infer if it's elelment_level as which kinds of search on ArrayOfVector are supported are determined by metric type but pymilvus does not have this info. - Fix element-level search returning wrong row IDs on growing segments with multiple chunks by using cumulative element offset instead of row offset as begin_id in brute-force search. https://github.com/milvus-io/milvus/issues/48617 Performance: - Replace proto.Clone with in-place FieldName mutation in reconstructStructFieldData Nested index: - field name is missing in CreateIndexInfo which is needed to determine whether the index should be nested or not Tests: added for all above fixes --------- Signed-off-by: SpadeA <tangchenjie1210@gmail.com>
This commit is contained in:
@@ -83,6 +83,8 @@ ScalarIndexCreator::ScalarIndexCreator(
|
||||
.value_or("");
|
||||
|
||||
index_info.field_type = dtype_;
|
||||
index_info.field_name =
|
||||
file_manager_context.fieldDataMeta.field_schema.name();
|
||||
index_info.index_type = index_type();
|
||||
if (dtype == DataType::JSON) {
|
||||
index_info.json_cast_type = milvus::JsonCastType::FromString(
|
||||
|
||||
@@ -29,6 +29,7 @@
|
||||
#include "common/QueryInfo.h"
|
||||
#include "common/Types.h"
|
||||
#include "common/Utils.h"
|
||||
#include "knowhere/emb_list_utils.h"
|
||||
#include "common/protobuf_utils.h"
|
||||
#include "fmt/core.h"
|
||||
#include "google/protobuf/io/coded_stream.h"
|
||||
@@ -50,13 +51,24 @@ ParsePlaceholderGroup(const Plan* plan,
|
||||
placeholder_group_blob.size());
|
||||
}
|
||||
|
||||
static bool
|
||||
is_emb_list_placeholder(milvus::proto::common::PlaceholderType type) {
|
||||
using PHType = milvus::proto::common::PlaceholderType;
|
||||
return type == PHType::EmbListFloatVector ||
|
||||
type == PHType::EmbListFloat16Vector ||
|
||||
type == PHType::EmbListBFloat16Vector ||
|
||||
type == PHType::EmbListBinaryVector ||
|
||||
type == PHType::EmbListInt8Vector;
|
||||
}
|
||||
|
||||
bool
|
||||
check_data_type(
|
||||
const FieldMeta& field_meta,
|
||||
const milvus::proto::common::PlaceholderValue& placeholder_value) {
|
||||
const milvus::proto::common::PlaceholderValue& placeholder_value,
|
||||
bool element_level) {
|
||||
if (field_meta.get_data_type() == DataType::VECTOR_ARRAY) {
|
||||
if (field_meta.get_element_type() == DataType::VECTOR_FLOAT) {
|
||||
if (placeholder_value.element_level()) {
|
||||
if (element_level) {
|
||||
return placeholder_value.type() ==
|
||||
milvus::proto::common::PlaceholderType::FloatVector;
|
||||
} else {
|
||||
@@ -65,7 +77,7 @@ check_data_type(
|
||||
EmbListFloatVector;
|
||||
}
|
||||
} else if (field_meta.get_element_type() == DataType::VECTOR_FLOAT16) {
|
||||
if (placeholder_value.element_level()) {
|
||||
if (element_level) {
|
||||
return placeholder_value.type() ==
|
||||
milvus::proto::common::PlaceholderType::Float16Vector;
|
||||
} else {
|
||||
@@ -74,7 +86,7 @@ check_data_type(
|
||||
EmbListFloat16Vector;
|
||||
}
|
||||
} else if (field_meta.get_element_type() == DataType::VECTOR_BFLOAT16) {
|
||||
if (placeholder_value.element_level()) {
|
||||
if (element_level) {
|
||||
return placeholder_value.type() ==
|
||||
milvus::proto::common::PlaceholderType::BFloat16Vector;
|
||||
} else {
|
||||
@@ -83,7 +95,7 @@ check_data_type(
|
||||
EmbListBFloat16Vector;
|
||||
}
|
||||
} else if (field_meta.get_element_type() == DataType::VECTOR_BINARY) {
|
||||
if (placeholder_value.element_level()) {
|
||||
if (element_level) {
|
||||
return placeholder_value.type() ==
|
||||
milvus::proto::common::PlaceholderType::BinaryVector;
|
||||
} else {
|
||||
@@ -92,7 +104,7 @@ check_data_type(
|
||||
EmbListBinaryVector;
|
||||
}
|
||||
} else if (field_meta.get_element_type() == DataType::VECTOR_INT8) {
|
||||
if (placeholder_value.element_level()) {
|
||||
if (element_level) {
|
||||
return placeholder_value.type() ==
|
||||
milvus::proto::common::PlaceholderType::Int8Vector;
|
||||
} else {
|
||||
@@ -118,11 +130,37 @@ ParsePlaceholderGroup(const Plan* plan,
|
||||
for (auto& ph : ph_group.placeholders()) {
|
||||
Placeholder element;
|
||||
element.tag_ = ph.tag();
|
||||
element.element_level_ = ph.element_level();
|
||||
Assert(plan->tag2field_.count(element.tag_));
|
||||
auto field_id = plan->tag2field_.at(element.tag_);
|
||||
auto& field_meta = plan->schema_->operator[](field_id);
|
||||
AssertInfo(check_data_type(field_meta, ph),
|
||||
|
||||
// Determine element_level for VECTOR_ARRAY fields based on
|
||||
// metric_type and placeholder type:
|
||||
// non-embedding-list metric + plain vector type → element-level search
|
||||
// embedding-list metric + emb-list placeholder → embedding-list search
|
||||
// mismatch combinations → error
|
||||
if (field_meta.get_data_type() == DataType::VECTOR_ARRAY) {
|
||||
auto& metric = plan->plan_node_->search_info_.metric_type_;
|
||||
bool emb_list_metric =
|
||||
knowhere::get_el_metric_type(metric).has_value();
|
||||
bool emb_list_ph = is_emb_list_placeholder(ph.type());
|
||||
if (emb_list_metric != emb_list_ph) {
|
||||
ThrowInfo(DataTypeInvalid,
|
||||
fmt::format(
|
||||
"search type mismatch for VECTOR_ARRAY field {}: "
|
||||
"metric_type {} {} embedding list search, "
|
||||
"but search data is {}",
|
||||
field_meta.get_name().get(),
|
||||
metric,
|
||||
emb_list_metric ? "requires" : "does not support",
|
||||
emb_list_ph ? "embedding list" : "plain vector"));
|
||||
}
|
||||
element.element_level_ = !emb_list_metric;
|
||||
} else {
|
||||
element.element_level_ = false;
|
||||
}
|
||||
|
||||
AssertInfo(check_data_type(field_meta, ph, element.element_level_),
|
||||
"vector type must be the same, field {} - type {}, search "
|
||||
"ph type {}",
|
||||
field_meta.get_name().get(),
|
||||
@@ -139,9 +177,9 @@ ParsePlaceholderGroup(const Plan* plan,
|
||||
auto& target = element.blob_;
|
||||
|
||||
if (field_meta.get_data_type() != DataType::VECTOR_ARRAY ||
|
||||
ph.element_level()) {
|
||||
element.element_level_) {
|
||||
if (field_meta.get_sizeof() != line_size &&
|
||||
!ph.element_level()) {
|
||||
!element.element_level_) {
|
||||
ThrowInfo(DimNotMatch,
|
||||
fmt::format(
|
||||
"vector dimension mismatch, expected vector "
|
||||
|
||||
@@ -234,6 +234,12 @@ SearchOnGrowing(const segcore::SegmentGrowingImpl& segment,
|
||||
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.
|
||||
int64_t cumulative_element_offset = 0;
|
||||
|
||||
for (int chunk_id = current_chunk_id; chunk_id < max_chunk;
|
||||
++chunk_id) {
|
||||
auto chunk_data = vec_ptr->get_chunk_data(chunk_id);
|
||||
@@ -271,7 +277,8 @@ SearchOnGrowing(const segcore::SegmentGrowingImpl& segment,
|
||||
count += vec_ptr[i].length();
|
||||
}
|
||||
sub_data = query::dataset::RawDataset{
|
||||
row_begin, dim, count, buf.get()};
|
||||
cumulative_element_offset, dim, count, buf.get()};
|
||||
cumulative_element_offset += count;
|
||||
} else {
|
||||
offsets.reserve(size_per_chunk + 1);
|
||||
offsets.push_back(0);
|
||||
|
||||
@@ -1366,6 +1366,7 @@ LoadIndexData(milvus::tracer::TraceContext& ctx,
|
||||
|
||||
milvus::index::CreateIndexInfo index_info;
|
||||
index_info.field_type = load_index_info->field_type;
|
||||
index_info.field_name = load_index_info->schema.name();
|
||||
index_info.index_engine_version = engine_version;
|
||||
|
||||
auto config = milvus::index::ParseConfigFromIndexParams(
|
||||
|
||||
@@ -2639,3 +2639,234 @@ TEST(ElementFilterGroupBy, DeduplicateRowsInGroup) {
|
||||
ASSERT_LE(group_counts.size(), static_cast<size_t>(topK))
|
||||
<< "Should have at most topK distinct groups";
|
||||
}
|
||||
|
||||
TEST(ElementFilter, SearchWithNestedScalarIndex) {
|
||||
auto saved_batch_size = EXEC_EVAL_EXPR_BATCH_SIZE.load();
|
||||
EXEC_EVAL_EXPR_BATCH_SIZE.store(100);
|
||||
|
||||
int dim = 4;
|
||||
auto schema = std::make_shared<Schema>();
|
||||
auto vec_fid = schema->AddDebugVectorArrayField("structA[array_float_vec]",
|
||||
DataType::VECTOR_FLOAT,
|
||||
dim,
|
||||
knowhere::metric::L2);
|
||||
auto int_array_fid = schema->AddDebugArrayField(
|
||||
"structA[price_array]", DataType::INT32, false);
|
||||
auto int64_fid = schema->AddDebugField("id", DataType::INT64);
|
||||
schema->set_primary_field_id(int64_fid);
|
||||
|
||||
size_t N = 200;
|
||||
int array_len = 3;
|
||||
|
||||
auto raw_data = DataGen(schema, N, 42, 0, 1, array_len);
|
||||
|
||||
// Customize int_array data: doc i has elements [i*3+1, i*3+2, i*3+3]
|
||||
for (int i = 0; i < raw_data.raw_->fields_data_size(); i++) {
|
||||
auto* field_data = raw_data.raw_->mutable_fields_data(i);
|
||||
if (field_data->field_id() == int_array_fid.get()) {
|
||||
field_data->mutable_scalars()
|
||||
->mutable_array_data()
|
||||
->mutable_data()
|
||||
->Clear();
|
||||
for (size_t row = 0; row < N; row++) {
|
||||
auto* array_data = field_data->mutable_scalars()
|
||||
->mutable_array_data()
|
||||
->mutable_data()
|
||||
->Add();
|
||||
for (int elem = 0; elem < array_len; elem++) {
|
||||
int value = row * array_len + elem + 1;
|
||||
array_data->mutable_int_data()->mutable_data()->Add(value);
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
auto segment = CreateSealedWithFieldDataLoaded(schema, raw_data);
|
||||
|
||||
// Load vector index
|
||||
auto array_vec_values = raw_data.get_col<VectorFieldProto>(vec_fid);
|
||||
std::vector<float> vector_data(dim * N * array_len);
|
||||
for (size_t i = 0; i < N; i++) {
|
||||
const auto& float_vec = array_vec_values[i].float_vector().data();
|
||||
for (int j = 0; j < array_len * dim; j++) {
|
||||
vector_data[i * array_len * dim + j] = float_vec[j];
|
||||
}
|
||||
}
|
||||
auto indexing = GenVecIndexing(N * array_len,
|
||||
dim,
|
||||
vector_data.data(),
|
||||
knowhere::IndexEnum::INDEX_HNSW);
|
||||
LoadIndexInfo load_index_info;
|
||||
load_index_info.field_id = vec_fid.get();
|
||||
load_index_info.index_params = GenIndexParams(indexing.get());
|
||||
load_index_info.cache_index =
|
||||
CreateTestCacheIndex("test", std::move(indexing));
|
||||
load_index_info.index_params["metric_type"] = knowhere::metric::L2;
|
||||
load_index_info.field_type = DataType::VECTOR_ARRAY;
|
||||
load_index_info.element_type = DataType::VECTOR_FLOAT;
|
||||
segment->LoadIndex(load_index_info);
|
||||
|
||||
// Build nested scalar index with is_nested=true (the correct behavior
|
||||
// after the fix sets field_name so IndexFactory routes to CreateNestedIndex).
|
||||
std::vector<int32_t> all_elements;
|
||||
all_elements.reserve(N * array_len);
|
||||
for (size_t row = 0; row < N; row++) {
|
||||
for (int elem = 0; elem < array_len; elem++) {
|
||||
all_elements.push_back(row * array_len + elem + 1);
|
||||
}
|
||||
}
|
||||
|
||||
auto stl_index = std::make_unique<milvus::index::ScalarIndexSort<int32_t>>(
|
||||
storage::FileManagerContext(),
|
||||
true /* is_nested=true: correct after fix */);
|
||||
stl_index->Build(all_elements.size(), all_elements.data(), nullptr);
|
||||
|
||||
LoadIndexInfo nested_load_info;
|
||||
nested_load_info.field_id = int_array_fid.get();
|
||||
nested_load_info.field_type = DataType::ARRAY;
|
||||
nested_load_info.element_type = DataType::INT32;
|
||||
nested_load_info.index_params["index_type"] = milvus::index::ASCENDING_SORT;
|
||||
nested_load_info.cache_index =
|
||||
CreateTestCacheIndex("nested_test", std::move(stl_index));
|
||||
segment->LoadIndex(nested_load_info);
|
||||
|
||||
// Build query plan: element_filter with nested index in full mode
|
||||
int topK = 5;
|
||||
ScopedSchemaHandle handle(*schema);
|
||||
std::string search_params = R"({"ef": 50})";
|
||||
// High selectivity (~50%) forces full mode evaluation
|
||||
std::string expr =
|
||||
"id % 2 == 0 && element_filter(structA, 2000 > $[price_array] > 100)";
|
||||
auto plan_bytes = handle.ParseSearch(expr,
|
||||
"structA[array_float_vec]",
|
||||
topK,
|
||||
knowhere::metric::L2,
|
||||
search_params,
|
||||
3);
|
||||
auto plan =
|
||||
CreateSearchPlanByExpr(schema, plan_bytes.data(), plan_bytes.size());
|
||||
ASSERT_NE(plan, nullptr);
|
||||
|
||||
auto num_queries = 1;
|
||||
auto seed = 1024;
|
||||
auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed, true);
|
||||
auto ph_group =
|
||||
ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
|
||||
|
||||
// With is_nested=true, the index correctly returns element-level results
|
||||
// matching the element_filter expression's expectations.
|
||||
auto search_result = segment->Search(plan.get(), ph_group.get(), 1L << 63);
|
||||
ASSERT_NE(search_result, nullptr);
|
||||
ASSERT_TRUE(search_result->element_level_);
|
||||
|
||||
EXEC_EVAL_EXPR_BATCH_SIZE.store(saved_batch_size);
|
||||
}
|
||||
|
||||
// Regression test for element-level search with multiple chunks in growing
|
||||
// segments. The bug was that begin_id was set to row_begin (row offset)
|
||||
// instead of the cumulative element offset, causing wrong row IDs to be
|
||||
// returned for chunks after the first one.
|
||||
// See: https://github.com/milvus-io/milvus/issues/48617
|
||||
TEST(ElementFilter, GrowingMultiChunkElementSearch) {
|
||||
int dim = 4;
|
||||
int array_len = 3;
|
||||
size_t N = 200;
|
||||
|
||||
auto schema = std::make_shared<Schema>();
|
||||
auto vec_fid = schema->AddDebugVectorArrayField(
|
||||
"structA[array_vec]", DataType::VECTOR_FLOAT, dim, "L2");
|
||||
auto int_array_fid = schema->AddDebugArrayField(
|
||||
"structA[price_array]", DataType::INT32, false);
|
||||
auto int64_fid = schema->AddDebugField("id", DataType::INT64);
|
||||
schema->set_primary_field_id(int64_fid);
|
||||
|
||||
auto raw_data = DataGen(schema, N, 42, 0, 1, array_len);
|
||||
|
||||
// Customize int_array: doc i has elements [i*3+1, i*3+2, i*3+3]
|
||||
for (int i = 0; i < raw_data.raw_->fields_data_size(); i++) {
|
||||
auto* field_data = raw_data.raw_->mutable_fields_data(i);
|
||||
if (field_data->field_id() == int_array_fid.get()) {
|
||||
field_data->mutable_scalars()
|
||||
->mutable_array_data()
|
||||
->mutable_data()
|
||||
->Clear();
|
||||
for (size_t row = 0; row < N; row++) {
|
||||
auto* array_data = field_data->mutable_scalars()
|
||||
->mutable_array_data()
|
||||
->mutable_data()
|
||||
->Add();
|
||||
for (int elem = 0; elem < array_len; elem++) {
|
||||
array_data->mutable_int_data()->mutable_data()->Add(
|
||||
row * array_len + elem + 1);
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Use small chunk_rows=32 to force multiple chunks (200 rows / 32 = 7
|
||||
// chunks). The bug only manifests with multiple chunks.
|
||||
SegcoreConfig config;
|
||||
config.set_chunk_rows(32);
|
||||
|
||||
auto segment = CreateGrowingSegment(schema, empty_index_meta, 1, config);
|
||||
segment->PreInsert(N);
|
||||
segment->Insert(0,
|
||||
N,
|
||||
raw_data.row_ids_.data(),
|
||||
raw_data.timestamps_.data(),
|
||||
raw_data.raw_);
|
||||
|
||||
// Verify ArrayOffsets
|
||||
auto growing_impl = dynamic_cast<SegmentGrowingImpl*>(segment.get());
|
||||
ASSERT_NE(growing_impl, nullptr);
|
||||
auto offsets = growing_impl->GetArrayOffsets(vec_fid);
|
||||
ASSERT_NE(offsets, nullptr);
|
||||
ASSERT_EQ(offsets->GetRowCount(), N);
|
||||
ASSERT_EQ(offsets->GetTotalElementCount(), N * array_len);
|
||||
|
||||
// Element-level search with element_filter that selects elements from
|
||||
// rows in later chunks (value > 300 means row >= 100, which is beyond
|
||||
// chunk 0-2).
|
||||
int topK = 5;
|
||||
ScopedSchemaHandle handle(*schema);
|
||||
std::string expr =
|
||||
"element_filter(structA, $[price_array] > 300 && $[price_array] < 400)";
|
||||
auto plan_bytes = handle.ParseSearch(
|
||||
expr, "structA[array_vec]", topK, "L2", R"({"ef": 50})", 3);
|
||||
auto plan =
|
||||
CreateSearchPlanByExpr(schema, plan_bytes.data(), plan_bytes.size());
|
||||
ASSERT_NE(plan, nullptr);
|
||||
|
||||
auto ph_group_raw = CreatePlaceholderGroup(1, dim, 1024, true);
|
||||
auto ph_group =
|
||||
ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
|
||||
|
||||
auto search_result = segment->Search(plan.get(), ph_group.get(), 1L << 63);
|
||||
ASSERT_NE(search_result, nullptr);
|
||||
ASSERT_TRUE(search_result->element_level_);
|
||||
|
||||
for (size_t i = 0; i < search_result->seg_offsets_.size(); i++) {
|
||||
int64_t doc_id = search_result->seg_offsets_[i];
|
||||
int32_t elem_idx = search_result->element_indices_[i];
|
||||
if (doc_id == INVALID_SEG_OFFSET) {
|
||||
continue;
|
||||
}
|
||||
|
||||
ASSERT_GE(doc_id, 0);
|
||||
ASSERT_LT(doc_id, static_cast<int64_t>(N))
|
||||
<< "doc_id should be a valid row index";
|
||||
ASSERT_GE(elem_idx, 0);
|
||||
ASSERT_LT(elem_idx, array_len);
|
||||
|
||||
// Verify the element satisfies the filter: value in (300, 400)
|
||||
int element_value = doc_id * array_len + elem_idx + 1;
|
||||
ASSERT_GT(element_value, 300)
|
||||
<< "doc_id=" << doc_id << " elem_idx=" << elem_idx
|
||||
<< " element_value=" << element_value << " should be > 300";
|
||||
ASSERT_LT(element_value, 400)
|
||||
<< "doc_id=" << doc_id << " elem_idx=" << elem_idx
|
||||
<< " element_value=" << element_value << " should be < 400";
|
||||
}
|
||||
}
|
||||
|
||||
@@ -848,3 +848,66 @@ TEST(Query, ExecWithPredicateBinary) {
|
||||
std::cout << json.dump(2);
|
||||
// ASSERT_EQ(json.dump(2), ref.dump(2));
|
||||
}
|
||||
|
||||
TEST(Query, VectorArrayElementLevelInference) {
|
||||
auto dim = 32;
|
||||
|
||||
// Helper to create schema + plan for a VECTOR_ARRAY field with given metric
|
||||
auto make_plan = [&](const std::string& metric) {
|
||||
auto schema = std::make_shared<Schema>();
|
||||
auto int64_field = schema->AddDebugField("int64", DataType::INT64);
|
||||
schema->AddDebugVectorArrayField(
|
||||
"array_vec", DataType::VECTOR_FLOAT, dim, metric);
|
||||
schema->set_primary_field_id(int64_field);
|
||||
|
||||
ScopedSchemaHandle handle(*schema);
|
||||
auto plan_str =
|
||||
handle.ParseSearch("", "array_vec", 5, metric, R"({"nprobe": 10})");
|
||||
auto plan =
|
||||
CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
|
||||
return plan;
|
||||
};
|
||||
|
||||
int num_queries = 2;
|
||||
std::vector<float> query_vec = generate_float_vector(num_queries, dim);
|
||||
|
||||
// Case 1: MAX_SIM + EmbList → element_level=false (embedding list search)
|
||||
{
|
||||
auto plan = make_plan("MAX_SIM");
|
||||
std::vector<size_t> offsets = {0, 1, 2};
|
||||
auto ph_raw = CreatePlaceholderGroupFromBlob<EmbListFloatVector>(
|
||||
num_queries, dim, query_vec.data(), offsets);
|
||||
auto ph = ParsePlaceholderGroup(plan.get(), ph_raw.SerializeAsString());
|
||||
EXPECT_FALSE(ph->at(0).element_level_);
|
||||
}
|
||||
|
||||
// Case 2: COSINE + plain vector → element_level=true (element-level search)
|
||||
{
|
||||
auto plan = make_plan("COSINE");
|
||||
auto ph_raw =
|
||||
CreatePlaceholderGroupFromBlob(num_queries, dim, query_vec.data());
|
||||
auto ph = ParsePlaceholderGroup(plan.get(), ph_raw.SerializeAsString());
|
||||
EXPECT_TRUE(ph->at(0).element_level_);
|
||||
}
|
||||
|
||||
// Case 3: MAX_SIM + plain vector → error (mismatch)
|
||||
{
|
||||
auto plan = make_plan("MAX_SIM");
|
||||
auto ph_raw =
|
||||
CreatePlaceholderGroupFromBlob(num_queries, dim, query_vec.data());
|
||||
EXPECT_THROW(
|
||||
ParsePlaceholderGroup(plan.get(), ph_raw.SerializeAsString()),
|
||||
std::exception);
|
||||
}
|
||||
|
||||
// Case 4: COSINE + EmbList → error (mismatch)
|
||||
{
|
||||
auto plan = make_plan("COSINE");
|
||||
std::vector<size_t> offsets = {0, 1, 2};
|
||||
auto ph_raw = CreatePlaceholderGroupFromBlob<EmbListFloatVector>(
|
||||
num_queries, dim, query_vec.data(), offsets);
|
||||
EXPECT_THROW(
|
||||
ParsePlaceholderGroup(plan.get(), ph_raw.SerializeAsString()),
|
||||
std::exception);
|
||||
}
|
||||
}
|
||||
|
||||
+15
-22
@@ -2862,9 +2862,10 @@ func getCollectionTTL(pairs []*commonpb.KeyValuePair) uint64 {
|
||||
return uint64(ttl)
|
||||
}
|
||||
|
||||
// reconstructStructFieldDataCommon reconstructs struct fields from flattened sub-fields
|
||||
// It works with both QueryResults and SearchResults by operating on the common data structures
|
||||
func reconstructStructFieldDataCommon(
|
||||
// reconstructStructFieldData regroups flattened sub-fields (named "structName[fieldName]")
|
||||
// back into StructArrayField entries, restoring original field names for the user-facing response.
|
||||
// It modifies sub-field FieldName in place; callers must not reuse the input slice afterwards.
|
||||
func reconstructStructFieldData(
|
||||
fieldsData []*schemapb.FieldData,
|
||||
outputFields []string,
|
||||
schema *schemapb.CollectionSchema,
|
||||
@@ -2909,40 +2910,33 @@ func reconstructStructFieldDataCommon(
|
||||
}
|
||||
|
||||
for structFieldID, fields := range groupedStructFields {
|
||||
// Create deep copies of fields to avoid modifying original data
|
||||
// and restore original field names for user-facing response
|
||||
copiedFields := make([]*schemapb.FieldData, len(fields))
|
||||
for i, field := range fields {
|
||||
copiedFields[i] = proto.Clone(field).(*schemapb.FieldData)
|
||||
// Extract original field name from structName[fieldName] format
|
||||
originalName, err := extractOriginalFieldName(copiedFields[i].FieldName)
|
||||
// Restore original field names (from "structName[fieldName]" to "fieldName")
|
||||
// for the user-facing response.
|
||||
for _, field := range fields {
|
||||
originalName, err := extractOriginalFieldName(field.FieldName)
|
||||
if err != nil {
|
||||
// This should not happen in normal operation - indicates a bug
|
||||
log.Error("failed to extract original field name from struct field",
|
||||
zap.String("fieldName", copiedFields[i].FieldName),
|
||||
zap.String("fieldName", field.FieldName),
|
||||
zap.Error(err))
|
||||
// Keep the transformed name to avoid data corruption
|
||||
} else {
|
||||
copiedFields[i].FieldName = originalName
|
||||
field.FieldName = originalName
|
||||
}
|
||||
}
|
||||
|
||||
fieldData := &schemapb.FieldData{
|
||||
newFieldsData = append(newFieldsData, &schemapb.FieldData{
|
||||
FieldName: structFieldNames[structFieldID],
|
||||
FieldId: structFieldID,
|
||||
Type: schemapb.DataType_ArrayOfStruct,
|
||||
Field: &schemapb.FieldData_StructArrays{StructArrays: &schemapb.StructArrayField{Fields: copiedFields}},
|
||||
}
|
||||
newFieldsData = append(newFieldsData, fieldData)
|
||||
Field: &schemapb.FieldData_StructArrays{StructArrays: &schemapb.StructArrayField{Fields: fields}},
|
||||
})
|
||||
reconstructedOutputFields = append(reconstructedOutputFields, structFieldNames[structFieldID])
|
||||
}
|
||||
|
||||
return newFieldsData, reconstructedOutputFields
|
||||
}
|
||||
|
||||
// Wrapper for QueryResults
|
||||
func reconstructStructFieldDataForQuery(results *milvuspb.QueryResults, schema *schemapb.CollectionSchema) {
|
||||
fieldsData, outputFields := reconstructStructFieldDataCommon(
|
||||
fieldsData, outputFields := reconstructStructFieldData(
|
||||
results.FieldsData,
|
||||
results.OutputFields,
|
||||
schema,
|
||||
@@ -2951,12 +2945,11 @@ func reconstructStructFieldDataForQuery(results *milvuspb.QueryResults, schema *
|
||||
results.OutputFields = outputFields
|
||||
}
|
||||
|
||||
// New wrapper for SearchResults
|
||||
func reconstructStructFieldDataForSearch(results *milvuspb.SearchResults, schema *schemapb.CollectionSchema) {
|
||||
if results.Results == nil {
|
||||
return
|
||||
}
|
||||
fieldsData, outputFields := reconstructStructFieldDataCommon(
|
||||
fieldsData, outputFields := reconstructStructFieldData(
|
||||
results.Results.FieldsData,
|
||||
results.Results.OutputFields,
|
||||
schema,
|
||||
|
||||
@@ -4283,7 +4283,7 @@ func TestValidateFieldsInStruct(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func Test_reconstructStructFieldDataCommon(t *testing.T) {
|
||||
func Test_reconstructStructFieldData(t *testing.T) {
|
||||
t.Run("count(*) query - should return early", func(t *testing.T) {
|
||||
fieldsData := []*schemapb.FieldData{
|
||||
{
|
||||
@@ -4316,7 +4316,7 @@ func Test_reconstructStructFieldDataCommon(t *testing.T) {
|
||||
originalOutputFields := make([]string, len(outputFields))
|
||||
copy(originalOutputFields, outputFields)
|
||||
|
||||
resultFieldsData, resultOutputFields := reconstructStructFieldDataCommon(fieldsData, outputFields, schema)
|
||||
resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
|
||||
|
||||
// Should not modify anything for count(*) query
|
||||
assert.Equal(t, originalFieldsData, resultFieldsData)
|
||||
@@ -4366,7 +4366,7 @@ func Test_reconstructStructFieldDataCommon(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
resultFieldsData, resultOutputFields := reconstructStructFieldDataCommon(fieldsData, outputFields, schema)
|
||||
resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
|
||||
|
||||
// Should reconstruct the struct field with the restored field name
|
||||
assert.Len(t, resultFieldsData, 1)
|
||||
@@ -4407,7 +4407,7 @@ func Test_reconstructStructFieldDataCommon(t *testing.T) {
|
||||
originalOutputFields := make([]string, len(outputFields))
|
||||
copy(originalOutputFields, outputFields)
|
||||
|
||||
resultFieldsData, resultOutputFields := reconstructStructFieldDataCommon(fieldsData, outputFields, schema)
|
||||
resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
|
||||
|
||||
// Should not modify anything when no struct array fields
|
||||
assert.Equal(t, originalFieldsData, resultFieldsData)
|
||||
@@ -4494,7 +4494,7 @@ func Test_reconstructStructFieldDataCommon(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
resultFieldsData, resultOutputFields := reconstructStructFieldDataCommon(fieldsData, outputFields, schema)
|
||||
resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
|
||||
|
||||
// Check result
|
||||
assert.Len(t, resultFieldsData, 1, "Should only have one reconstructed struct field")
|
||||
@@ -4594,7 +4594,7 @@ func Test_reconstructStructFieldDataCommon(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
resultFieldsData, resultOutputFields := reconstructStructFieldDataCommon(fieldsData, outputFields, schema)
|
||||
resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
|
||||
|
||||
// Check result: should have 2 fields (1 regular + 1 reconstructed struct)
|
||||
assert.Len(t, resultFieldsData, 2)
|
||||
@@ -4713,7 +4713,7 @@ func Test_reconstructStructFieldDataCommon(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
resultFieldsData, resultOutputFields := reconstructStructFieldDataCommon(fieldsData, outputFields, schema)
|
||||
resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
|
||||
|
||||
// Check result: should have 2 struct fields
|
||||
assert.Len(t, resultFieldsData, 2)
|
||||
@@ -4815,7 +4815,7 @@ func Test_reconstructStructFieldDataCommon(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
resultFieldsData, resultOutputFields := reconstructStructFieldDataCommon(fieldsData, outputFields, schema)
|
||||
resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
|
||||
|
||||
// Check result
|
||||
assert.Len(t, resultFieldsData, 1, "Should have one reconstructed struct field")
|
||||
|
||||
@@ -170,17 +170,24 @@ func (r *rowParser) reconstructArrayForStructArray(structName string, subFieldsM
|
||||
return nil, merr.WrapErrImportFailed(fmt.Sprintf("invalid StructArray format in CSV, failed to parse JSON: %v", err))
|
||||
}
|
||||
|
||||
expectedFieldCount := len(subFieldsMap)
|
||||
flatStructs := make(map[string][]any)
|
||||
for _, elem := range structs {
|
||||
for i, elem := range structs {
|
||||
dict, ok := elem.(map[string]any)
|
||||
if !ok {
|
||||
return nil, merr.WrapErrImportFailed(fmt.Sprintf("invalid element in StructArray, expect map[string]any but got type %T", elem))
|
||||
}
|
||||
if len(dict) != expectedFieldCount {
|
||||
return nil, merr.WrapErrImportFailed(
|
||||
fmt.Sprintf("inconsistent field count in StructArray element: position=%d, actual=%d, expected=%d",
|
||||
i, len(dict), expectedFieldCount))
|
||||
}
|
||||
for key, value := range dict {
|
||||
fieldName := typeutil.ConcatStructFieldName(structName, key)
|
||||
_, ok := subFieldsMap[fieldName]
|
||||
if !ok {
|
||||
return nil, merr.WrapErrImportFailed(fmt.Sprintf("field %s not found", fieldName))
|
||||
return nil, merr.WrapErrImportFailed(
|
||||
fmt.Sprintf("unexpected field in StructArray element: field=%s, position=%d", fieldName, i))
|
||||
}
|
||||
|
||||
flatStructs[fieldName] = append(flatStructs[fieldName], value)
|
||||
|
||||
@@ -160,19 +160,33 @@ func (r *rowParser) wrapArrayValueTypeError(v any, eleType schemapb.DataType) er
|
||||
// we reconstruct it to be handled by handleField as:
|
||||
//
|
||||
// {"sub-field1": [1, 2], "sub-field2": [[1.0, 2.0], [3.0, 4.0]]}
|
||||
func reconstructArrayForStructArray(raw any) (map[string]any, error) {
|
||||
func reconstructArrayForStructArray(raw any, subFieldNames []string) (map[string]any, error) {
|
||||
rows, ok := raw.([]any)
|
||||
if !ok {
|
||||
return nil, merr.WrapErrImportFailed(fmt.Sprintf("invalid StructArray format in JSON, each row should be a key-value map, but got type %T", raw))
|
||||
}
|
||||
|
||||
expectedFields := make(map[string]struct{}, len(subFieldNames))
|
||||
for _, name := range subFieldNames {
|
||||
expectedFields[name] = struct{}{}
|
||||
}
|
||||
|
||||
buf := make(map[string][]any)
|
||||
for _, elem := range rows {
|
||||
for i, elem := range rows {
|
||||
row, ok := elem.(map[string]any)
|
||||
if !ok {
|
||||
return nil, merr.WrapErrImportFailed(fmt.Sprintf("invalid element in StructArray, expect map[string]any but got type %T", elem))
|
||||
}
|
||||
if len(row) != len(subFieldNames) {
|
||||
return nil, merr.WrapErrImportFailed(
|
||||
fmt.Sprintf("inconsistent field count in StructArray element: position=%d, actual=%d, expected=%d",
|
||||
i, len(row), len(subFieldNames)))
|
||||
}
|
||||
for key, value := range row {
|
||||
if _, ok := expectedFields[key]; !ok {
|
||||
return nil, merr.WrapErrImportFailed(
|
||||
fmt.Sprintf("unexpected field in StructArray element: field=%s, position=%d, expected fields=%v", key, i, subFieldNames))
|
||||
}
|
||||
buf[key] = append(buf[key], value)
|
||||
}
|
||||
}
|
||||
@@ -240,7 +254,7 @@ func (r *rowParser) Parse(raw any) (Row, error) {
|
||||
// read values from json file
|
||||
for key, value := range stringMap {
|
||||
if subFieldNames, ok := r.structArrays[key]; ok {
|
||||
values, err := reconstructArrayForStructArray(value)
|
||||
values, err := reconstructArrayForStructArray(value, subFieldNames)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -788,11 +802,15 @@ func (r *rowParser) arrayOfVectorToFieldData(vectors []any, field *schemapb.Fiel
|
||||
case schemapb.DataType_FloatVector:
|
||||
values := make([]float32, 0, len(vectors)*10)
|
||||
dim := r.id2Dim[fieldID]
|
||||
for _, vectorAny := range vectors {
|
||||
for i, vectorAny := range vectors {
|
||||
vector, ok := vectorAny.([]any)
|
||||
if !ok {
|
||||
return nil, merr.WrapErrImportFailed(fmt.Sprintf("expected slice as vector, but got %T", vectorAny))
|
||||
}
|
||||
if len(vector) != dim {
|
||||
return nil, merr.WrapErrImportFailed(
|
||||
fmt.Sprintf("vector at index %d has dimension %d, expected %d", i, len(vector), dim))
|
||||
}
|
||||
for _, v := range vector {
|
||||
value, ok := v.(json.Number)
|
||||
if !ok {
|
||||
|
||||
@@ -752,6 +752,84 @@ func (suite *RowParserSuite) TestParseError() {
|
||||
}
|
||||
}
|
||||
|
||||
func TestReconstructArrayForStructArray_InconsistentFields(t *testing.T) {
|
||||
subFields := []string{"sub_int", "sub_str"}
|
||||
|
||||
// Element missing a field should produce an error
|
||||
raw := []any{
|
||||
map[string]any{"sub_int": 1, "sub_str": "hello"},
|
||||
map[string]any{"sub_int": 2}, // missing "sub_str"
|
||||
}
|
||||
_, err := reconstructArrayForStructArray(raw, subFields)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "inconsistent field count in StructArray")
|
||||
|
||||
// Element with an extra field should produce an error
|
||||
raw = []any{
|
||||
map[string]any{"sub_int": 1, "sub_str": "hello", "sub_extra": true},
|
||||
}
|
||||
_, err = reconstructArrayForStructArray(raw, subFields)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "inconsistent field count in StructArray")
|
||||
|
||||
// Consistent fields should succeed
|
||||
raw = []any{
|
||||
map[string]any{"sub_int": 1, "sub_str": "hello1"},
|
||||
map[string]any{"sub_int": 2, "sub_str": "hello2"},
|
||||
}
|
||||
result, err := reconstructArrayForStructArray(raw, subFields)
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, result, 2)
|
||||
|
||||
// Empty array should succeed
|
||||
raw = []any{}
|
||||
result, err = reconstructArrayForStructArray(raw, subFields)
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, result, 0)
|
||||
|
||||
// Single element should succeed
|
||||
raw = []any{
|
||||
map[string]any{"sub_int": 1, "sub_str": "hello"},
|
||||
}
|
||||
result, err = reconstructArrayForStructArray(raw, subFields)
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, result, 2)
|
||||
}
|
||||
|
||||
func TestArrayOfVectorToFieldData_DimensionMismatch(t *testing.T) {
|
||||
dim := 3
|
||||
fieldID := int64(100)
|
||||
field := &schemapb.FieldSchema{
|
||||
FieldID: fieldID,
|
||||
Name: "test_vec",
|
||||
DataType: schemapb.DataType_ArrayOfVector,
|
||||
ElementType: schemapb.DataType_FloatVector,
|
||||
TypeParams: []*commonpb.KeyValuePair{
|
||||
{Key: common.DimKey, Value: "3"},
|
||||
},
|
||||
}
|
||||
parser := &rowParser{
|
||||
id2Dim: map[int64]int{fieldID: dim},
|
||||
id2Field: map[int64]*schemapb.FieldSchema{fieldID: field},
|
||||
}
|
||||
|
||||
// Mismatched dimension: expect 3, got 2
|
||||
vectors := []any{
|
||||
[]any{json.Number("1.0"), json.Number("2.0")}, // dim=2, expected 3
|
||||
}
|
||||
_, err := parser.arrayOfVectorToFieldData(vectors, field)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "dimension")
|
||||
|
||||
// Correct dimension should succeed
|
||||
vectors = []any{
|
||||
[]any{json.Number("1.0"), json.Number("2.0"), json.Number("3.0")},
|
||||
}
|
||||
result, err := parser.arrayOfVectorToFieldData(vectors, field)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, int64(dim), result.Dim)
|
||||
}
|
||||
|
||||
func TestJsonRowParser(t *testing.T) {
|
||||
suite.Run(t, new(RowParserSuite))
|
||||
}
|
||||
|
||||
@@ -852,3 +852,35 @@ func TestArrayNullElement(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStructFieldReader_toScalarField_TypeMismatch(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
elementType schemapb.DataType
|
||||
data []interface{}
|
||||
}{
|
||||
{"Bool_wrong_type", schemapb.DataType_Bool, []interface{}{"not_a_bool"}},
|
||||
{"Int8_wrong_type", schemapb.DataType_Int8, []interface{}{"not_an_int8"}},
|
||||
{"Int16_wrong_type", schemapb.DataType_Int16, []interface{}{3.14}},
|
||||
{"Int32_wrong_type", schemapb.DataType_Int32, []interface{}{true}},
|
||||
{"Int64_wrong_type", schemapb.DataType_Int64, []interface{}{"not_an_int64"}},
|
||||
{"Float_wrong_type", schemapb.DataType_Float, []interface{}{int32(1)}},
|
||||
{"Double_wrong_type", schemapb.DataType_Double, []interface{}{int64(1)}},
|
||||
{"VarChar_wrong_type", schemapb.DataType_VarChar, []interface{}{123}},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
reader := &StructFieldReader{
|
||||
field: &schemapb.FieldSchema{
|
||||
Name: "test_field",
|
||||
DataType: schemapb.DataType_Array,
|
||||
ElementType: tt.elementType,
|
||||
},
|
||||
}
|
||||
_, err := reader.toScalarField(tt.data)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "expected")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -109,9 +109,11 @@ func (r *StructFieldReader) toScalarField(data []interface{}) (*schemapb.ScalarF
|
||||
case schemapb.DataType_Bool:
|
||||
boolData := make([]bool, len(data))
|
||||
for i, v := range data {
|
||||
if val, ok := v.(bool); ok {
|
||||
boolData[i] = val
|
||||
val, ok := v.(bool)
|
||||
if !ok {
|
||||
return nil, merr.WrapErrImportFailed(fmt.Sprintf("expected bool for field '%s', got %T at index %d", r.field.GetName(), v, i))
|
||||
}
|
||||
boolData[i] = val
|
||||
}
|
||||
return &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_BoolData{
|
||||
@@ -121,9 +123,11 @@ func (r *StructFieldReader) toScalarField(data []interface{}) (*schemapb.ScalarF
|
||||
case schemapb.DataType_Int8:
|
||||
intData := make([]int32, len(data))
|
||||
for i, v := range data {
|
||||
if val, ok := v.(int8); ok {
|
||||
intData[i] = int32(val)
|
||||
val, ok := v.(int8)
|
||||
if !ok {
|
||||
return nil, merr.WrapErrImportFailed(fmt.Sprintf("expected int8 for field '%s', got %T at index %d", r.field.GetName(), v, i))
|
||||
}
|
||||
intData[i] = int32(val)
|
||||
}
|
||||
return &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_IntData{
|
||||
@@ -133,9 +137,11 @@ func (r *StructFieldReader) toScalarField(data []interface{}) (*schemapb.ScalarF
|
||||
case schemapb.DataType_Int16:
|
||||
intData := make([]int32, len(data))
|
||||
for i, v := range data {
|
||||
if val, ok := v.(int16); ok {
|
||||
intData[i] = int32(val)
|
||||
val, ok := v.(int16)
|
||||
if !ok {
|
||||
return nil, merr.WrapErrImportFailed(fmt.Sprintf("expected int16 for field '%s', got %T at index %d", r.field.GetName(), v, i))
|
||||
}
|
||||
intData[i] = int32(val)
|
||||
}
|
||||
return &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_IntData{
|
||||
@@ -145,9 +151,11 @@ func (r *StructFieldReader) toScalarField(data []interface{}) (*schemapb.ScalarF
|
||||
case schemapb.DataType_Int32:
|
||||
intData := make([]int32, len(data))
|
||||
for i, v := range data {
|
||||
if val, ok := v.(int32); ok {
|
||||
intData[i] = val
|
||||
val, ok := v.(int32)
|
||||
if !ok {
|
||||
return nil, merr.WrapErrImportFailed(fmt.Sprintf("expected int32 for field '%s', got %T at index %d", r.field.GetName(), v, i))
|
||||
}
|
||||
intData[i] = val
|
||||
}
|
||||
return &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_IntData{
|
||||
@@ -157,9 +165,11 @@ func (r *StructFieldReader) toScalarField(data []interface{}) (*schemapb.ScalarF
|
||||
case schemapb.DataType_Int64:
|
||||
intData := make([]int64, len(data))
|
||||
for i, v := range data {
|
||||
if val, ok := v.(int64); ok {
|
||||
intData[i] = val
|
||||
val, ok := v.(int64)
|
||||
if !ok {
|
||||
return nil, merr.WrapErrImportFailed(fmt.Sprintf("expected int64 for field '%s', got %T at index %d", r.field.GetName(), v, i))
|
||||
}
|
||||
intData[i] = val
|
||||
}
|
||||
return &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_LongData{
|
||||
@@ -169,9 +179,11 @@ func (r *StructFieldReader) toScalarField(data []interface{}) (*schemapb.ScalarF
|
||||
case schemapb.DataType_Float:
|
||||
floatData := make([]float32, len(data))
|
||||
for i, v := range data {
|
||||
if val, ok := v.(float32); ok {
|
||||
floatData[i] = val
|
||||
val, ok := v.(float32)
|
||||
if !ok {
|
||||
return nil, merr.WrapErrImportFailed(fmt.Sprintf("expected float32 for field '%s', got %T at index %d", r.field.GetName(), v, i))
|
||||
}
|
||||
floatData[i] = val
|
||||
}
|
||||
return &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_FloatData{
|
||||
@@ -181,9 +193,11 @@ func (r *StructFieldReader) toScalarField(data []interface{}) (*schemapb.ScalarF
|
||||
case schemapb.DataType_Double:
|
||||
floatData := make([]float64, len(data))
|
||||
for i, v := range data {
|
||||
if val, ok := v.(float64); ok {
|
||||
floatData[i] = val
|
||||
val, ok := v.(float64)
|
||||
if !ok {
|
||||
return nil, merr.WrapErrImportFailed(fmt.Sprintf("expected float64 for field '%s', got %T at index %d", r.field.GetName(), v, i))
|
||||
}
|
||||
floatData[i] = val
|
||||
}
|
||||
return &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_DoubleData{
|
||||
@@ -193,9 +207,11 @@ func (r *StructFieldReader) toScalarField(data []interface{}) (*schemapb.ScalarF
|
||||
case schemapb.DataType_String, schemapb.DataType_VarChar:
|
||||
strData := make([]string, len(data))
|
||||
for i, v := range data {
|
||||
if val, ok := v.(string); ok {
|
||||
strData[i] = val
|
||||
val, ok := v.(string)
|
||||
if !ok {
|
||||
return nil, merr.WrapErrImportFailed(fmt.Sprintf("expected string for field '%s', got %T at index %d", r.field.GetName(), v, i))
|
||||
}
|
||||
strData[i] = val
|
||||
}
|
||||
return &schemapb.ScalarField{
|
||||
Data: &schemapb.ScalarField_StringData{
|
||||
@@ -203,7 +219,7 @@ func (r *StructFieldReader) toScalarField(data []interface{}) (*schemapb.ScalarF
|
||||
},
|
||||
}, nil
|
||||
default:
|
||||
return nil, merr.WrapErrImportFailed(fmt.Sprintf("unsupported element type for struct field: %v", r.field.GetDataType()))
|
||||
return nil, merr.WrapErrImportFailed(fmt.Sprintf("unsupported element type for struct field: %v", r.field.GetElementType()))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -309,100 +325,126 @@ func (r *StructFieldReader) readArrayOfVectorField(chunked *arrow.Chunked) (any,
|
||||
// Extract vectors based on element type
|
||||
switch r.field.GetElementType() {
|
||||
case schemapb.DataType_FloatVector:
|
||||
floatArr, ok := fieldArray.ListValues().(*array.Float32)
|
||||
if !ok {
|
||||
return nil, nil, merr.WrapErrImportFailed(
|
||||
fmt.Sprintf("expected Float32 array for FloatVector field '%s', got %T", r.field.GetName(), fieldArray.ListValues()))
|
||||
}
|
||||
var allVectors []float32
|
||||
for structIdx := startIdx; structIdx < endIdx; structIdx++ {
|
||||
vecStart, vecEnd := fieldArray.ValueOffsets(int(structIdx))
|
||||
if floatArr, ok := fieldArray.ListValues().(*array.Float32); ok {
|
||||
for j := vecStart; j < vecEnd; j++ {
|
||||
allVectors = append(allVectors, floatArr.Value(int(j)))
|
||||
}
|
||||
if int(vecEnd-vecStart) != r.dim {
|
||||
return nil, nil, merr.WrapErrImportFailed(
|
||||
fmt.Sprintf("vector dimension mismatch for field '%s': position=%d, actual=%d, expected=%d",
|
||||
r.field.GetName(), structIdx, vecEnd-vecStart, r.dim))
|
||||
}
|
||||
for j := vecStart; j < vecEnd; j++ {
|
||||
allVectors = append(allVectors, floatArr.Value(int(j)))
|
||||
}
|
||||
}
|
||||
// struct list could be empty, len(allVectors) can be zero
|
||||
// build an empty VectorField if len(allVectors) is zero
|
||||
if len(allVectors) >= 0 {
|
||||
vectorField := &schemapb.VectorField{
|
||||
Dim: int64(r.dim),
|
||||
Data: &schemapb.VectorField_FloatVector{
|
||||
FloatVector: &schemapb.FloatArray{Data: allVectors},
|
||||
},
|
||||
}
|
||||
result = append(result, vectorField)
|
||||
}
|
||||
result = append(result, &schemapb.VectorField{
|
||||
Dim: int64(r.dim),
|
||||
Data: &schemapb.VectorField_FloatVector{
|
||||
FloatVector: &schemapb.FloatArray{Data: allVectors},
|
||||
},
|
||||
})
|
||||
|
||||
case schemapb.DataType_Float16Vector:
|
||||
uint8Arr, ok := fieldArray.ListValues().(*array.Uint8)
|
||||
if !ok {
|
||||
return nil, nil, merr.WrapErrImportFailed(
|
||||
fmt.Sprintf("expected Uint8 array for Float16Vector field '%s', got %T", r.field.GetName(), fieldArray.ListValues()))
|
||||
}
|
||||
expectedBytes := int64(r.dim * 2)
|
||||
var allVectors []byte
|
||||
for structIdx := startIdx; structIdx < endIdx; structIdx++ {
|
||||
vecStart, vecEnd := fieldArray.ValueOffsets(int(structIdx))
|
||||
if uint8Arr, ok := fieldArray.ListValues().(*array.Uint8); ok {
|
||||
allVectors = append(allVectors, uint8Arr.Uint8Values()[vecStart:vecEnd]...)
|
||||
if vecEnd-vecStart != expectedBytes {
|
||||
return nil, nil, merr.WrapErrImportFailed(
|
||||
fmt.Sprintf("vector dimension mismatch for field '%s': position=%d, actual_bytes=%d, expected_bytes=%d",
|
||||
r.field.GetName(), structIdx, vecEnd-vecStart, expectedBytes))
|
||||
}
|
||||
allVectors = append(allVectors, uint8Arr.Uint8Values()[vecStart:vecEnd]...)
|
||||
}
|
||||
if len(allVectors) >= 0 {
|
||||
vectorField := &schemapb.VectorField{
|
||||
Dim: int64(r.dim),
|
||||
Data: &schemapb.VectorField_Float16Vector{
|
||||
Float16Vector: allVectors,
|
||||
},
|
||||
}
|
||||
result = append(result, vectorField)
|
||||
}
|
||||
result = append(result, &schemapb.VectorField{
|
||||
Dim: int64(r.dim),
|
||||
Data: &schemapb.VectorField_Float16Vector{
|
||||
Float16Vector: allVectors,
|
||||
},
|
||||
})
|
||||
|
||||
case schemapb.DataType_BFloat16Vector:
|
||||
uint8Arr, ok := fieldArray.ListValues().(*array.Uint8)
|
||||
if !ok {
|
||||
return nil, nil, merr.WrapErrImportFailed(
|
||||
fmt.Sprintf("expected Uint8 array for BFloat16Vector field '%s', got %T", r.field.GetName(), fieldArray.ListValues()))
|
||||
}
|
||||
expectedBytes := int64(r.dim * 2)
|
||||
var allVectors []byte
|
||||
for structIdx := startIdx; structIdx < endIdx; structIdx++ {
|
||||
vecStart, vecEnd := fieldArray.ValueOffsets(int(structIdx))
|
||||
if uint8Arr, ok := fieldArray.ListValues().(*array.Uint8); ok {
|
||||
allVectors = append(allVectors, uint8Arr.Uint8Values()[vecStart:vecEnd]...)
|
||||
if vecEnd-vecStart != expectedBytes {
|
||||
return nil, nil, merr.WrapErrImportFailed(
|
||||
fmt.Sprintf("vector dimension mismatch for field '%s': position=%d, actual_bytes=%d, expected_bytes=%d",
|
||||
r.field.GetName(), structIdx, vecEnd-vecStart, expectedBytes))
|
||||
}
|
||||
allVectors = append(allVectors, uint8Arr.Uint8Values()[vecStart:vecEnd]...)
|
||||
}
|
||||
if len(allVectors) >= 0 {
|
||||
vectorField := &schemapb.VectorField{
|
||||
Dim: int64(r.dim),
|
||||
Data: &schemapb.VectorField_Bfloat16Vector{
|
||||
Bfloat16Vector: allVectors,
|
||||
},
|
||||
}
|
||||
result = append(result, vectorField)
|
||||
}
|
||||
result = append(result, &schemapb.VectorField{
|
||||
Dim: int64(r.dim),
|
||||
Data: &schemapb.VectorField_Bfloat16Vector{
|
||||
Bfloat16Vector: allVectors,
|
||||
},
|
||||
})
|
||||
|
||||
case schemapb.DataType_Int8Vector:
|
||||
int8Arr, ok := fieldArray.ListValues().(*array.Int8)
|
||||
if !ok {
|
||||
return nil, nil, merr.WrapErrImportFailed(
|
||||
fmt.Sprintf("expected Int8 array for Int8Vector field '%s', got %T", r.field.GetName(), fieldArray.ListValues()))
|
||||
}
|
||||
var allVectors []byte
|
||||
for structIdx := startIdx; structIdx < endIdx; structIdx++ {
|
||||
vecStart, vecEnd := fieldArray.ValueOffsets(int(structIdx))
|
||||
if int8Arr, ok := fieldArray.ListValues().(*array.Int8); ok {
|
||||
for j := vecStart; j < vecEnd; j++ {
|
||||
allVectors = append(allVectors, byte(int8Arr.Value(int(j))))
|
||||
}
|
||||
if int(vecEnd-vecStart) != r.dim {
|
||||
return nil, nil, merr.WrapErrImportFailed(
|
||||
fmt.Sprintf("vector dimension mismatch for field '%s': position=%d, actual=%d, expected=%d",
|
||||
r.field.GetName(), structIdx, vecEnd-vecStart, r.dim))
|
||||
}
|
||||
for j := vecStart; j < vecEnd; j++ {
|
||||
allVectors = append(allVectors, byte(int8Arr.Value(int(j))))
|
||||
}
|
||||
}
|
||||
if len(allVectors) >= 0 {
|
||||
vectorField := &schemapb.VectorField{
|
||||
Dim: int64(r.dim),
|
||||
Data: &schemapb.VectorField_Int8Vector{
|
||||
Int8Vector: allVectors,
|
||||
},
|
||||
}
|
||||
result = append(result, vectorField)
|
||||
}
|
||||
result = append(result, &schemapb.VectorField{
|
||||
Dim: int64(r.dim),
|
||||
Data: &schemapb.VectorField_Int8Vector{
|
||||
Int8Vector: allVectors,
|
||||
},
|
||||
})
|
||||
|
||||
case schemapb.DataType_BinaryVector:
|
||||
uint8Arr, ok := fieldArray.ListValues().(*array.Uint8)
|
||||
if !ok {
|
||||
return nil, nil, merr.WrapErrImportFailed(
|
||||
fmt.Sprintf("expected Uint8 array for BinaryVector field '%s', got %T", r.field.GetName(), fieldArray.ListValues()))
|
||||
}
|
||||
expectedBytes := int64(r.dim / 8)
|
||||
var allVectors []byte
|
||||
for structIdx := startIdx; structIdx < endIdx; structIdx++ {
|
||||
vecStart, vecEnd := fieldArray.ValueOffsets(int(structIdx))
|
||||
if uint8Arr, ok := fieldArray.ListValues().(*array.Uint8); ok {
|
||||
allVectors = append(allVectors, uint8Arr.Uint8Values()[vecStart:vecEnd]...)
|
||||
if vecEnd-vecStart != expectedBytes {
|
||||
return nil, nil, merr.WrapErrImportFailed(
|
||||
fmt.Sprintf("vector dimension mismatch for field '%s': position=%d, actual_bytes=%d, expected_bytes=%d",
|
||||
r.field.GetName(), structIdx, vecEnd-vecStart, expectedBytes))
|
||||
}
|
||||
allVectors = append(allVectors, uint8Arr.Uint8Values()[vecStart:vecEnd]...)
|
||||
}
|
||||
if len(allVectors) >= 0 {
|
||||
vectorField := &schemapb.VectorField{
|
||||
Dim: int64(r.dim),
|
||||
Data: &schemapb.VectorField_BinaryVector{
|
||||
BinaryVector: allVectors,
|
||||
},
|
||||
}
|
||||
result = append(result, vectorField)
|
||||
}
|
||||
result = append(result, &schemapb.VectorField{
|
||||
Dim: int64(r.dim),
|
||||
Data: &schemapb.VectorField_BinaryVector{
|
||||
BinaryVector: allVectors,
|
||||
},
|
||||
})
|
||||
|
||||
case schemapb.DataType_SparseFloatVector:
|
||||
return nil, nil, merr.WrapErrImportFailed("ArrayOfVector with SparseFloatVector element type is not implemented yet")
|
||||
|
||||
Reference in New Issue
Block a user