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:
Spade A
2026-03-31 17:19:32 +08:00
committed by GitHub
parent e1aefb24b8
commit 08cdd63f68
13 changed files with 637 additions and 125 deletions
@@ -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(
+48 -10
View File
@@ -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 "
+8 -1
View File
@@ -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);
+1
View File
@@ -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";
}
}
+63
View File
@@ -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
View File
@@ -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,
+8 -8
View File
@@ -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")
+9 -2
View File
@@ -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)
+22 -4
View File
@@ -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")