furyhawk/milvus
1
0
Fork 0
mirror of https://github.com/milvus-io/milvus.git synced 2026-07-21 10:15:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: impl StructArray -- add more tests for non-float embedding lists (#46686)

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

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
- Core invariant: StructArray now accepts only fixed-dimension vector
element types (typeutil.IsFixDimVectorType) for array/struct vector
fields — i.e., dense/quantized fixed-dim vectors (FloatVector,
Float16Vector, BFloat16Vector, Int8Vector, BinaryVector) are supported
while sparse vectors remain rejected at validation
(internal/proxy/util.go → ValidateFieldsInStruct).
- Logic removed / simplified: the previous FloatVector-only special-case
was removed and replaced by a single, generic fixed-dimension check;
numerous duplicated test branches were collapsed into parameterized
fixtures (internal/core/src/common/ChunkWriterTest.cpp: GetByteWidth +
TEST_P/INSTANTIATE_TEST_SUITE_P) and ElementFilter tests were unified
under CreatePlaceholderGroupForType()/ElementFilterSealedParam,
eliminating repetitive type-specific code.
- No data loss / regression rationale (concrete code paths): existing
VECTOR_FLOAT code paths and behavior are preserved — tests explicitly
instantiate VECTOR_FLOAT cases in ChunkWriterTest and ElementFilter;
deserialization now explicitly handles empty ArrayOfVector cases and
returns the correct VectorField by elementType
(internal/storage/serde.go: deserializeArrayOfVector), and StructArray
readers were extended with byte-based branches for
float16/bfloat16/int8/binary without altering FloatVector logic
(internal/util/importutilv2/parquet/struct_field_reader.go), making all
additions additive and failing unsupported sparse vectors early at proxy
validation rather than silently.
- New capability added: full StructArray support for additional
fixed-dimension vector element types (Float16Vector, BFloat16Vector,
Int8Vector, BinaryVector) across import, serde, parquet struct readers,
test generators, and end-to-end tests (files touched include
internal/util/importutilv2/parquet/struct_field_reader.go,
internal/storage/serde.go, pkg/util/testutils/gen_data.go, various
tests), enabling non-float embedding-list and byte-based vector
workflows.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: SpadeA <tangchenjie1210@gmail.com>
This commit is contained in:
Spade A
2026-01-15 10:13:27 +08:00
committed by GitHub
parent 93735186fb
commit b39de6391b
19 changed files with 1665 additions and 408 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/core/src/common/ChunkWriterTest.cpp
+122 -75
View File
@@ -18,6 +18,7 @@
#include <cstdint>
#include <memory>
#include <vector>
#include <random>
#include "arrow/api.h"
#include "arrow/array/array_nested.h"
@@ -29,6 +30,7 @@
#include "common/ChunkWriter.h"
#include "common/Types.h"
using milvus::DataType;
using milvus::MemChunkTarget;
using milvus::MMAP_ARRAY_PADDING;
using milvus::VectorArrayChunk;
@@ -36,14 +38,34 @@ using milvus::VectorArrayChunkWriter;
namespace {
// Calculate byte width for a single vector based on data type and dimension
int
GetByteWidth(DataType data_type, int dim) {
switch (data_type) {
case DataType::VECTOR_FLOAT:
return dim * sizeof(float);
case DataType::VECTOR_FLOAT16:
case DataType::VECTOR_BFLOAT16:
return dim * 2;
case DataType::VECTOR_INT8:
return dim;
case DataType::VECTOR_BINARY:
return (dim + 7) / 8;
default:
return dim * sizeof(float);
}
}
// Helper function to build a ListArray of FixedSizeBinary (vector array)
// Each row contains a variable number of vectors
// vectors_per_row: specifies how many vectors each row contains
// dim: dimension of each vector (number of floats)
// dim: dimension of each vector
// data_type: the vector data type
std::shared_ptr<arrow::ListArray>
BuildVectorArrayListArray(const std::vector<int>& vectors_per_row, int dim) {
// Each vector is stored as FixedSizeBinary with size = dim * sizeof(float)
int byte_width = dim * sizeof(float);
BuildVectorArrayListArray(const std::vector<int>& vectors_per_row,
int dim,
DataType data_type = DataType::VECTOR_FLOAT) {
int byte_width = GetByteWidth(data_type, dim);
auto value_type = arrow::fixed_size_binary(byte_width);
arrow::FixedSizeBinaryBuilder value_builder(value_type);
@@ -54,18 +76,17 @@ BuildVectorArrayListArray(const std::vector<int>& vectors_per_row, int dim) {
auto& fsb_builder = dynamic_cast<arrow::FixedSizeBinaryBuilder&>(
*list_builder.value_builder());
float counter = 0.0f;
std::default_random_engine gen(42);
std::uniform_int_distribution<int> dist(0, 255);
for (size_t row = 0; row < vectors_per_row.size(); ++row) {
EXPECT_TRUE(list_builder.Append().ok());
for (int vec = 0; vec < vectors_per_row[row]; ++vec) {
std::vector<float> vector_data(dim);
for (int d = 0; d < dim; ++d) {
vector_data[d] = counter++;
std::vector<uint8_t> vector_data(byte_width);
for (int d = 0; d < byte_width; ++d) {
vector_data[d] = static_cast<uint8_t>(dist(gen));
}
EXPECT_TRUE(fsb_builder
.Append(reinterpret_cast<const uint8_t*>(
vector_data.data()))
.ok());
EXPECT_TRUE(fsb_builder.Append(vector_data.data()).ok());
}
}
@@ -74,27 +95,52 @@ BuildVectorArrayListArray(const std::vector<int>& vectors_per_row, int dim) {
return std::static_pointer_cast<arrow::ListArray>(result);
}
// Test parameter structure for parameterized tests
struct VectorArrayWriterTestParam {
DataType data_type;
int dim;
std::string test_name;
};
} // namespace
// Test basic functionality without slicing
TEST(VectorArrayChunkWriterTest, BasicNoSlice) {
const int dim = 4;
// Parameterized test class for VectorArrayChunkWriter
class VectorArrayChunkWriterParameterizedTest
: public ::testing::TestWithParam<VectorArrayWriterTestParam> {
protected:
DataType
data_type() const {
return GetParam().data_type;
}
int
dim() const {
return GetParam().dim;
}
int
byte_width() const {
return GetByteWidth(data_type(), dim());
}
};
// Test basic functionality without slicing - parameterized version
TEST_P(VectorArrayChunkWriterParameterizedTest, BasicNoSlice) {
// 5 rows with varying number of vectors per row
std::vector<int> vectors_per_row = {2, 3, 1, 4, 2}; // Total: 12 vectors
auto list_array = BuildVectorArrayListArray(vectors_per_row, dim);
auto list_array =
BuildVectorArrayListArray(vectors_per_row, dim(), data_type());
ASSERT_EQ(list_array->length(), 5);
arrow::ArrayVector vec{list_array};
VectorArrayChunkWriter writer(dim, milvus::DataType::VECTOR_FLOAT);
VectorArrayChunkWriter writer(dim(), data_type());
auto [calculated_size, row_count] = writer.calculate_size(vec);
// Expected size:
// - 12 vectors * 4 floats * 4 bytes = 192 bytes for data
// - 12 vectors * byte_width bytes for data
// - (5 * 2 + 1) * 4 bytes = 44 bytes for offsets and lengths
// - MMAP_ARRAY_PADDING (1) byte for padding
int expected_data_size = 12 * dim * sizeof(float); // 192
int expected_data_size = 12 * byte_width();
int expected_overhead =
sizeof(uint32_t) * (5 * 2 + 1) + MMAP_ARRAY_PADDING; // 44 + 1 = 45
EXPECT_EQ(calculated_size, expected_data_size + expected_overhead);
@@ -106,24 +152,19 @@ TEST(VectorArrayChunkWriterTest, BasicNoSlice) {
// Create chunk from target data
auto* data = target->release();
auto chunk =
std::make_unique<VectorArrayChunk>(dim,
row_count,
data,
calculated_size,
milvus::DataType::VECTOR_FLOAT,
nullptr);
auto chunk = std::make_unique<VectorArrayChunk>(
dim(), row_count, data, calculated_size, data_type(), nullptr);
ASSERT_NE(chunk, nullptr);
EXPECT_EQ(chunk->RowNums(), 5);
}
// Test with sliced ListArray - THIS IS THE KEY TEST FOR THE BUG
TEST(VectorArrayChunkWriterTest, SlicedListArray) {
const int dim = 4;
TEST_P(VectorArrayChunkWriterParameterizedTest, SlicedListArray) {
// Original: 10 rows with 2 vectors each = 20 vectors total
std::vector<int> vectors_per_row(10, 2);
auto original_array = BuildVectorArrayListArray(vectors_per_row, dim);
auto original_array =
BuildVectorArrayListArray(vectors_per_row, dim(), data_type());
ASSERT_EQ(original_array->length(), 10);
// Slice: take rows 3-6 (4 rows, should have 8 vectors)
@@ -144,14 +185,14 @@ TEST(VectorArrayChunkWriterTest, SlicedListArray) {
arrow::ArrayVector vec{sliced_array};
VectorArrayChunkWriter writer(dim, milvus::DataType::VECTOR_FLOAT);
VectorArrayChunkWriter writer(dim(), data_type());
auto [calculated_size, row_count] = writer.calculate_size(vec);
// Expected size with the fix:
// - 8 vectors * 4 floats * 4 bytes = 128 bytes for data
// - 8 vectors * byte_width bytes for data
// - (4 * 2 + 1) * 4 bytes = 36 bytes for offsets and lengths
// - MMAP_ARRAY_PADDING (1) byte for padding
int expected_data_size = 8 * dim * sizeof(float); // 128
int expected_data_size = 8 * byte_width();
int expected_overhead =
sizeof(uint32_t) * (4 * 2 + 1) + MMAP_ARRAY_PADDING; // 36 + 1 = 37
EXPECT_EQ(calculated_size, expected_data_size + expected_overhead);
@@ -163,28 +204,23 @@ TEST(VectorArrayChunkWriterTest, SlicedListArray) {
// Create chunk from target data
auto* data = target->release();
auto chunk =
std::make_unique<VectorArrayChunk>(dim,
row_count,
data,
calculated_size,
milvus::DataType::VECTOR_FLOAT,
nullptr);
auto chunk = std::make_unique<VectorArrayChunk>(
dim(), row_count, data, calculated_size, data_type(), nullptr);
ASSERT_NE(chunk, nullptr);
EXPECT_EQ(chunk->RowNums(), 4);
}
// Test with multiple sliced arrays in array_vec
TEST(VectorArrayChunkWriterTest, MultipleSlicedArrays) {
const int dim = 4;
TEST_P(VectorArrayChunkWriterParameterizedTest, MultipleSlicedArrays) {
// First array: 8 rows with varying vectors
std::vector<int> vectors_per_row1 = {1, 2, 3, 2, 1, 2, 3, 2}; // 16 total
auto array1 = BuildVectorArrayListArray(vectors_per_row1, dim);
auto array1 =
BuildVectorArrayListArray(vectors_per_row1, dim(), data_type());
// Second array: 6 rows with 2 vectors each
std::vector<int> vectors_per_row2(6, 2); // 12 total
auto array2 = BuildVectorArrayListArray(vectors_per_row2, dim);
auto array2 =
BuildVectorArrayListArray(vectors_per_row2, dim(), data_type());
// Slice both: first array rows 2-5 (4 rows), second array rows 1-4 (4 rows)
auto sliced1 =
@@ -203,11 +239,10 @@ TEST(VectorArrayChunkWriterTest, MultipleSlicedArrays) {
arrow::ArrayVector vec{sliced1, sliced2};
VectorArrayChunkWriter writer(dim, milvus::DataType::VECTOR_FLOAT);
VectorArrayChunkWriter writer(dim(), data_type());
auto [calculated_size, row_count] = writer.calculate_size(vec);
int expected_data_size =
expected_vectors * dim * sizeof(float); // 16 * 16 = 256
int expected_data_size = expected_vectors * byte_width();
int expected_overhead = sizeof(uint32_t) * (expected_rows * 2 + 1) +
MMAP_ARRAY_PADDING; // 17 * 4 + 1 = 69
EXPECT_EQ(calculated_size, expected_data_size + expected_overhead);
@@ -219,23 +254,18 @@ TEST(VectorArrayChunkWriterTest, MultipleSlicedArrays) {
// Create chunk from target data
auto* data = target->release();
auto chunk =
std::make_unique<VectorArrayChunk>(dim,
row_count,
data,
calculated_size,
milvus::DataType::VECTOR_FLOAT,
nullptr);
auto chunk = std::make_unique<VectorArrayChunk>(
dim(), row_count, data, calculated_size, data_type(), nullptr);
ASSERT_NE(chunk, nullptr);
EXPECT_EQ(chunk->RowNums(), expected_rows);
}
// Test edge case: slice from the beginning
TEST(VectorArrayChunkWriterTest, SliceFromBeginning) {
const int dim = 4;
TEST_P(VectorArrayChunkWriterParameterizedTest, SliceFromBeginning) {
std::vector<int> vectors_per_row = {3, 2, 1, 4, 2}; // 12 vectors total
auto original_array = BuildVectorArrayListArray(vectors_per_row, dim);
auto original_array =
BuildVectorArrayListArray(vectors_per_row, dim(), data_type());
// Slice first 2 rows (should have 3+2=5 vectors)
auto sliced =
@@ -244,10 +274,10 @@ TEST(VectorArrayChunkWriterTest, SliceFromBeginning) {
arrow::ArrayVector vec{sliced};
VectorArrayChunkWriter writer(dim, milvus::DataType::VECTOR_FLOAT);
VectorArrayChunkWriter writer(dim(), data_type());
auto [calculated_size, row_count] = writer.calculate_size(vec);
int expected_data_size = 5 * dim * sizeof(float); // 80
int expected_data_size = 5 * byte_width();
int expected_overhead =
sizeof(uint32_t) * (2 * 2 + 1) + MMAP_ARRAY_PADDING; // 20 + 1 = 21
EXPECT_EQ(calculated_size, expected_data_size + expected_overhead);
@@ -255,11 +285,11 @@ TEST(VectorArrayChunkWriterTest, SliceFromBeginning) {
}
// Test edge case: slice to the end
TEST(VectorArrayChunkWriterTest, SliceToEnd) {
const int dim = 4;
TEST_P(VectorArrayChunkWriterParameterizedTest, SliceToEnd) {
std::vector<int> vectors_per_row = {3, 2, 1, 4, 2}; // 12 vectors total
auto original_array = BuildVectorArrayListArray(vectors_per_row, dim);
auto original_array =
BuildVectorArrayListArray(vectors_per_row, dim(), data_type());
// Slice last 2 rows (should have 4+2=6 vectors)
auto sliced =
@@ -268,10 +298,10 @@ TEST(VectorArrayChunkWriterTest, SliceToEnd) {
arrow::ArrayVector vec{sliced};
VectorArrayChunkWriter writer(dim, milvus::DataType::VECTOR_FLOAT);
VectorArrayChunkWriter writer(dim(), data_type());
auto [calculated_size, row_count] = writer.calculate_size(vec);
int expected_data_size = 6 * dim * sizeof(float); // 96
int expected_data_size = 6 * byte_width();
int expected_overhead =
sizeof(uint32_t) * (2 * 2 + 1) + MMAP_ARRAY_PADDING; // 20 + 1 = 21
EXPECT_EQ(calculated_size, expected_data_size + expected_overhead);
@@ -279,12 +309,15 @@ TEST(VectorArrayChunkWriterTest, SliceToEnd) {
}
// Test that calculate_size and write_to_target produce consistent results
TEST(VectorArrayChunkWriterTest, SizeConsistencyWithSlice) {
const int dim = 8;
TEST_P(VectorArrayChunkWriterParameterizedTest, SizeConsistencyWithSlice) {
// Use dim() from parameter, but for consistency test use a fixed dim=8
// to match the original test's vectors_per_row expectations
const int test_dim = 8;
std::vector<int> vectors_per_row = {
1, 3, 2, 4, 1, 2, 3, 1}; // 17 vectors total
auto original_array = BuildVectorArrayListArray(vectors_per_row, dim);
auto original_array =
BuildVectorArrayListArray(vectors_per_row, test_dim, data_type());
// Try various slices and verify size consistency
std::vector<std::pair<int64_t, int64_t>> slices = {
@@ -297,6 +330,8 @@ TEST(VectorArrayChunkWriterTest, SizeConsistencyWithSlice) {
{7, 1}, // last row
};
int test_byte_width = GetByteWidth(data_type(), test_dim);
for (const auto& [offset, length] : slices) {
auto sliced = std::static_pointer_cast<arrow::ListArray>(
original_array->Slice(offset, length));
@@ -304,7 +339,7 @@ TEST(VectorArrayChunkWriterTest, SizeConsistencyWithSlice) {
arrow::ArrayVector vec{sliced};
VectorArrayChunkWriter writer(dim, milvus::DataType::VECTOR_FLOAT);
VectorArrayChunkWriter writer(test_dim, data_type());
auto [calculated_size, row_count] = writer.calculate_size(vec);
EXPECT_EQ(row_count, length);
@@ -315,15 +350,27 @@ TEST(VectorArrayChunkWriterTest, SizeConsistencyWithSlice) {
// Create chunk from target data
auto* data = target->release();
auto chunk =
std::make_unique<VectorArrayChunk>(dim,
row_count,
data,
calculated_size,
milvus::DataType::VECTOR_FLOAT,
nullptr);
auto chunk = std::make_unique<VectorArrayChunk>(
test_dim, row_count, data, calculated_size, data_type(), nullptr);
ASSERT_NE(chunk, nullptr)
<< "Failed for slice(" << offset << ", " << length << ")";
EXPECT_EQ(chunk->RowNums(), length);
}
}
// Instantiate parameterized tests for all vector types
INSTANTIATE_TEST_SUITE_P(
VectorTypes,
VectorArrayChunkWriterParameterizedTest,
::testing::Values(
VectorArrayWriterTestParam{DataType::VECTOR_FLOAT, 4, "FloatVector"},
VectorArrayWriterTestParam{
DataType::VECTOR_FLOAT16, 4, "Float16Vector"},
VectorArrayWriterTestParam{
DataType::VECTOR_BFLOAT16, 4, "BFloat16Vector"},
VectorArrayWriterTestParam{DataType::VECTOR_INT8, 4, "Int8Vector"},
VectorArrayWriterTestParam{
DataType::VECTOR_BINARY, 32, "BinaryVector"}),
[](const ::testing::TestParamInfo<VectorArrayWriterTestParam>& info) {
return info.param.test_name;
});
internal/core/src/query/Plan.cpp
+32 -9
View File
@@ -45,18 +45,41 @@ check_data_type(
EmbListFloatVector;
}
} else if (field_meta.get_element_type() == DataType::VECTOR_FLOAT16) {
return placeholder_value.type() ==
milvus::proto::common::PlaceholderType::EmbListFloat16Vector;
if (placeholder_value.element_level()) {
return placeholder_value.type() ==
milvus::proto::common::PlaceholderType::Float16Vector;
} else {
return placeholder_value.type() ==
milvus::proto::common::PlaceholderType::
EmbListFloat16Vector;
}
} else if (field_meta.get_element_type() == DataType::VECTOR_BFLOAT16) {
return placeholder_value.type() ==
milvus::proto::common::PlaceholderType::
EmbListBFloat16Vector;
if (placeholder_value.element_level()) {
return placeholder_value.type() ==
milvus::proto::common::PlaceholderType::BFloat16Vector;
} else {
return placeholder_value.type() ==
milvus::proto::common::PlaceholderType::
EmbListBFloat16Vector;
}
} else if (field_meta.get_element_type() == DataType::VECTOR_BINARY) {
return placeholder_value.type() ==
milvus::proto::common::PlaceholderType::EmbListBinaryVector;
if (placeholder_value.element_level()) {
return placeholder_value.type() ==
milvus::proto::common::PlaceholderType::BinaryVector;
} else {
return placeholder_value.type() ==
milvus::proto::common::PlaceholderType::
EmbListBinaryVector;
}
} else if (field_meta.get_element_type() == DataType::VECTOR_INT8) {
return placeholder_value.type() ==
milvus::proto::common::PlaceholderType::EmbListInt8Vector;
if (placeholder_value.element_level()) {
return placeholder_value.type() ==
milvus::proto::common::PlaceholderType::Int8Vector;
} else {
return placeholder_value.type() ==
milvus::proto::common::PlaceholderType::
EmbListInt8Vector;
}
}
return false;
}
internal/core/src/segcore/Utils.cpp
+24
View File
@@ -244,6 +244,30 @@ GetRawDataSizeOfDataArray(const DataArray* data,
}
break;
}
case DataType::VECTOR_FLOAT16: {
for (auto& e : obj) {
result += e.float16_vector().size();
}
break;
}
case DataType::VECTOR_BFLOAT16: {
for (auto& e : obj) {
result += e.bfloat16_vector().size();
}
break;
}
case DataType::VECTOR_INT8: {
for (auto& e : obj) {
result += e.int8_vector().size();
}
break;
}
case DataType::VECTOR_BINARY: {
for (auto& e : obj) {
result += e.binary_vector().size();
}
break;
}
default: {
ThrowInfo(NotImplemented,
fmt::format("not implemented vector type {}",
internal/core/unittest/test_element_filter.cpp
+390 -106
View File
@@ -25,8 +25,12 @@ using namespace milvus;
using namespace milvus::query;
using namespace milvus::segcore;
// Test parameter: <use_hints, load_index, element_type, metric_type, dim>
using ElementFilterSealedParam =
std::tuple<bool, bool, DataType, std::string, int>;
class ElementFilterSealed
: public ::testing::TestWithParam<std::tuple<bool, bool>> {
: public ::testing::TestWithParam<ElementFilterSealedParam> {
protected:
bool
use_hints() const {
@@ -36,18 +40,54 @@ class ElementFilterSealed
load_index() const {
return std::get<1>(GetParam());
}
DataType
element_type() const {
return std::get<2>(GetParam());
}
std::string
metric_type() const {
return std::get<3>(GetParam());
}
int
vec_dim() const {
return std::get<4>(GetParam());
}
// Create placeholder group with element_level = true for element-level search
// Uses regular vector types (not EmbList), as query is single embedding per query
proto::common::PlaceholderGroup
CreatePlaceholderGroupForType(int num_queries, int dim, int seed) {
if (element_type() == DataType::VECTOR_BINARY) {
return CreatePlaceholderGroup<milvus::BinaryVector>(
num_queries, dim, seed, true);
} else if (element_type() == DataType::VECTOR_FLOAT16) {
return CreatePlaceholderGroup<milvus::Float16Vector>(
num_queries, dim, seed, true);
} else if (element_type() == DataType::VECTOR_BFLOAT16) {
return CreatePlaceholderGroup<milvus::BFloat16Vector>(
num_queries, dim, seed, true);
} else if (element_type() == DataType::VECTOR_INT8) {
return CreatePlaceholderGroup<milvus::Int8Vector>(
num_queries, dim, seed, true);
} else {
// VECTOR_FLOAT
return CreatePlaceholderGroup<milvus::FloatVector>(
num_queries, dim, seed, true);
}
}
};
TEST_P(ElementFilterSealed, RangeExpr) {
bool with_hints = use_hints();
bool with_load_index = load_index();
DataType elem_type = element_type();
std::string metric = metric_type();
int dim = vec_dim();
// Step 1: Prepare schema with array field
int dim = 4;
auto schema = std::make_shared<Schema>();
auto vec_fid = schema->AddDebugVectorArrayField("structA[array_float_vec]",
DataType::VECTOR_FLOAT,
dim,
knowhere::metric::L2);
auto vec_fid = schema->AddDebugVectorArrayField(
"structA[array_vec]", elem_type, dim, metric);
auto int_array_fid = schema->AddDebugArrayField(
"structA[price_array]", DataType::INT32, false);
@@ -68,7 +108,7 @@ TEST_P(ElementFilterSealed, RangeExpr) {
->mutable_data()
->Clear();
for (int row = 0; row < N; row++) {
for (size_t row = 0; row < N; row++) {
auto* array_data = field_data->mutable_scalars()
->mutable_array_data()
->mutable_data()
@@ -89,30 +129,95 @@ TEST_P(ElementFilterSealed, RangeExpr) {
// Step 4: Load vector index for element-level search
auto array_vec_values = raw_data.get_col<VectorFieldProto>(vec_fid);
// DataGen generates VECTOR_ARRAY with data in float_vector (flattened),
// not in vector_array (nested structure)
std::vector<float> vector_data(dim * N * array_len);
for (int i = 0; i < N; i++) {
const auto& float_vec = array_vec_values[i].float_vector().data();
// float_vec contains array_len * dim floats
for (int j = 0; j < array_len * dim; j++) {
vector_data[i * array_len * dim + j] = float_vec[j];
// Flatten vector data and build index based on element type
std::unique_ptr<milvus::index::VectorIndex> indexing;
std::string actual_metric;
if (elem_type == DataType::VECTOR_FLOAT) {
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];
}
}
indexing = GenVecIndexing(N * array_len,
dim,
vector_data.data(),
knowhere::IndexEnum::INDEX_HNSW);
actual_metric = knowhere::metric::L2;
} else if (elem_type == DataType::VECTOR_FLOAT16) {
std::vector<knowhere::fp16> vector_data(dim * N * array_len);
for (size_t i = 0; i < N; i++) {
const auto& fp16_data = array_vec_values[i].float16_vector();
const knowhere::fp16* src =
reinterpret_cast<const knowhere::fp16*>(fp16_data.data());
for (int j = 0; j < array_len * dim; j++) {
vector_data[i * array_len * dim + j] = src[j];
}
}
indexing = GenVecIndexingFloat16(N * array_len,
dim,
vector_data.data(),
knowhere::IndexEnum::INDEX_HNSW);
actual_metric = knowhere::metric::L2;
} else if (elem_type == DataType::VECTOR_BFLOAT16) {
std::vector<knowhere::bf16> vector_data(dim * N * array_len);
for (size_t i = 0; i < N; i++) {
const auto& bf16_data = array_vec_values[i].bfloat16_vector();
const knowhere::bf16* src =
reinterpret_cast<const knowhere::bf16*>(bf16_data.data());
for (int j = 0; j < array_len * dim; j++) {
vector_data[i * array_len * dim + j] = src[j];
}
}
indexing = GenVecIndexingBFloat16(N * array_len,
dim,
vector_data.data(),
knowhere::IndexEnum::INDEX_HNSW);
actual_metric = knowhere::metric::L2;
} else if (elem_type == DataType::VECTOR_INT8) {
std::vector<int8_t> vector_data(dim * N * array_len);
for (size_t i = 0; i < N; i++) {
const auto& int8_data = array_vec_values[i].int8_vector();
const int8_t* src =
reinterpret_cast<const int8_t*>(int8_data.data());
for (int j = 0; j < array_len * dim; j++) {
vector_data[i * array_len * dim + j] = src[j];
}
}
indexing = GenVecIndexingInt8(N * array_len,
dim,
vector_data.data(),
knowhere::IndexEnum::INDEX_HNSW);
actual_metric = knowhere::metric::L2;
} else if (elem_type == DataType::VECTOR_BINARY) {
int byte_dim = (dim + 7) / 8;
std::vector<uint8_t> vector_data(byte_dim * N * array_len);
for (size_t i = 0; i < N; i++) {
const auto& bin_data = array_vec_values[i].binary_vector();
const uint8_t* src =
reinterpret_cast<const uint8_t*>(bin_data.data());
for (int j = 0; j < array_len * byte_dim; j++) {
vector_data[i * array_len * byte_dim + j] = src[j];
}
}
indexing =
GenVecIndexingBinary(N * array_len,
dim,
vector_data.data(),
knowhere::IndexEnum::INDEX_FAISS_BIN_IDMAP);
actual_metric = knowhere::metric::HAMMING;
}
// For element-level search, index all elements (N * array_len vectors)
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.index_params["metric_type"] = actual_metric;
load_index_info.field_type = DataType::VECTOR_ARRAY;
load_index_info.element_type = DataType::VECTOR_FLOAT;
load_index_info.element_type = elem_type;
if (with_load_index) {
segment->LoadIndex(load_index_info);
}
@@ -124,7 +229,8 @@ TEST_P(ElementFilterSealed, RangeExpr) {
{
std::string hints_line =
with_hints ? R"(hints: "iterative_filter")" : "";
std::string raw_plan = boost::str(boost::format(R"(vector_anns: <
std::string raw_plan =
boost::str(boost::format(R"(vector_anns: <
field_id: %1%
predicates: <
element_filter_expr: <
@@ -168,13 +274,13 @@ TEST_P(ElementFilterSealed, RangeExpr) {
query_info: <
topk: 5
round_decimal: 3
metric_type: "L2"
%4%
metric_type: "%4%"
%5%
search_params: "{\"ef\": 50}"
>
placeholder_tag: "$0">)") %
vec_fid.get() % int_array_fid.get() %
int64_fid.get() % hints_line);
vec_fid.get() % int_array_fid.get() % int64_fid.get() %
metric % hints_line);
proto::plan::PlanNode plan_node;
auto ok =
@@ -187,7 +293,7 @@ TEST_P(ElementFilterSealed, RangeExpr) {
auto num_queries = 1;
auto seed = 1024;
auto ph_group_raw =
CreatePlaceholderGroup(num_queries, dim, seed, true);
CreatePlaceholderGroupForType(num_queries, dim, seed);
auto ph_group =
ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
@@ -206,10 +312,10 @@ TEST_P(ElementFilterSealed, RangeExpr) {
search_result->seg_offsets_.size());
// Should have topK results per query
ASSERT_LE(search_result->element_indices_.size(), topK * num_queries);
ASSERT_LE(search_result->element_indices_.size(),
static_cast<size_t>(topK * num_queries));
std::cout << "Element-level search returned:" << std::endl;
for (auto i = 0; i < search_result->seg_offsets_.size(); i++) {
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];
float distance = search_result->distances_[i];
@@ -230,7 +336,7 @@ TEST_P(ElementFilterSealed, RangeExpr) {
<< "Element value " << element_value << " should be < 400";
}
// Verify distances are sorted (ascending for L2)
// Verify distances are sorted
for (size_t i = 1; i < search_result->distances_.size(); ++i) {
ASSERT_LE(search_result->distances_[i - 1],
search_result->distances_[i])
@@ -242,13 +348,14 @@ TEST_P(ElementFilterSealed, RangeExpr) {
TEST_P(ElementFilterSealed, UnaryExpr) {
bool with_hints = use_hints();
bool with_load_index = load_index();
DataType elem_type = element_type();
std::string metric = metric_type();
int dim = vec_dim();
// Step 1: Prepare schema with array field
int dim = 4;
auto schema = std::make_shared<Schema>();
auto vec_fid = schema->AddDebugVectorArrayField("structA[array_float_vec]",
DataType::VECTOR_FLOAT,
dim,
knowhere::metric::L2);
auto vec_fid = schema->AddDebugVectorArrayField(
"structA[array_vec]", elem_type, dim, metric);
auto int_array_fid = schema->AddDebugArrayField(
"structA[price_array]", DataType::INT32, false);
@@ -269,7 +376,7 @@ TEST_P(ElementFilterSealed, UnaryExpr) {
->mutable_data()
->Clear();
for (int row = 0; row < N; row++) {
for (size_t row = 0; row < N; row++) {
auto* array_data = field_data->mutable_scalars()
->mutable_array_data()
->mutable_data()
@@ -290,30 +397,95 @@ TEST_P(ElementFilterSealed, UnaryExpr) {
// Step 4: Load vector index for element-level search
auto array_vec_values = raw_data.get_col<VectorFieldProto>(vec_fid);
// DataGen generates VECTOR_ARRAY with data in float_vector (flattened),
// not in vector_array (nested structure)
std::vector<float> vector_data(dim * N * array_len);
for (int i = 0; i < N; i++) {
const auto& float_vec = array_vec_values[i].float_vector().data();
// float_vec contains array_len * dim floats
for (int j = 0; j < array_len * dim; j++) {
vector_data[i * array_len * dim + j] = float_vec[j];
// Flatten vector data and build index based on element type
std::unique_ptr<milvus::index::VectorIndex> indexing;
std::string actual_metric;
if (elem_type == DataType::VECTOR_FLOAT) {
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];
}
}
indexing = GenVecIndexing(N * array_len,
dim,
vector_data.data(),
knowhere::IndexEnum::INDEX_HNSW);
actual_metric = knowhere::metric::L2;
} else if (elem_type == DataType::VECTOR_FLOAT16) {
std::vector<knowhere::fp16> vector_data(dim * N * array_len);
for (size_t i = 0; i < N; i++) {
const auto& fp16_data = array_vec_values[i].float16_vector();
const knowhere::fp16* src =
reinterpret_cast<const knowhere::fp16*>(fp16_data.data());
for (int j = 0; j < array_len * dim; j++) {
vector_data[i * array_len * dim + j] = src[j];
}
}
indexing = GenVecIndexingFloat16(N * array_len,
dim,
vector_data.data(),
knowhere::IndexEnum::INDEX_HNSW);
actual_metric = knowhere::metric::L2;
} else if (elem_type == DataType::VECTOR_BFLOAT16) {
std::vector<knowhere::bf16> vector_data(dim * N * array_len);
for (size_t i = 0; i < N; i++) {
const auto& bf16_data = array_vec_values[i].bfloat16_vector();
const knowhere::bf16* src =
reinterpret_cast<const knowhere::bf16*>(bf16_data.data());
for (int j = 0; j < array_len * dim; j++) {
vector_data[i * array_len * dim + j] = src[j];
}
}
indexing = GenVecIndexingBFloat16(N * array_len,
dim,
vector_data.data(),
knowhere::IndexEnum::INDEX_HNSW);
actual_metric = knowhere::metric::L2;
} else if (elem_type == DataType::VECTOR_INT8) {
std::vector<int8_t> vector_data(dim * N * array_len);
for (size_t i = 0; i < N; i++) {
const auto& int8_data = array_vec_values[i].int8_vector();
const int8_t* src =
reinterpret_cast<const int8_t*>(int8_data.data());
for (int j = 0; j < array_len * dim; j++) {
vector_data[i * array_len * dim + j] = src[j];
}
}
indexing = GenVecIndexingInt8(N * array_len,
dim,
vector_data.data(),
knowhere::IndexEnum::INDEX_HNSW);
actual_metric = knowhere::metric::L2;
} else if (elem_type == DataType::VECTOR_BINARY) {
int byte_dim = (dim + 7) / 8;
std::vector<uint8_t> vector_data(byte_dim * N * array_len);
for (size_t i = 0; i < N; i++) {
const auto& bin_data = array_vec_values[i].binary_vector();
const uint8_t* src =
reinterpret_cast<const uint8_t*>(bin_data.data());
for (int j = 0; j < array_len * byte_dim; j++) {
vector_data[i * array_len * byte_dim + j] = src[j];
}
}
indexing =
GenVecIndexingBinary(N * array_len,
dim,
vector_data.data(),
knowhere::IndexEnum::INDEX_FAISS_BIN_IDMAP);
actual_metric = knowhere::metric::HAMMING;
}
// For element-level search, index all elements (N * array_len vectors)
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.index_params["metric_type"] = actual_metric;
load_index_info.field_type = DataType::VECTOR_ARRAY;
load_index_info.element_type = DataType::VECTOR_FLOAT;
load_index_info.element_type = elem_type;
if (with_load_index) {
segment->LoadIndex(load_index_info);
}
@@ -321,11 +493,12 @@ TEST_P(ElementFilterSealed, UnaryExpr) {
int topK = 5;
// Step 5: Test with element-level filter
// Query: Search array elements, filter by element_value < 10
// Query: Search array elements, filter by element_value > 10
{
std::string hints_line =
with_hints ? R"(hints: "iterative_filter")" : "";
std::string raw_plan = boost::str(boost::format(R"(vector_anns: <
std::string raw_plan =
boost::str(boost::format(R"(vector_anns: <
field_id: %1%
predicates: <
element_filter_expr: <
@@ -365,13 +538,13 @@ TEST_P(ElementFilterSealed, UnaryExpr) {
query_info: <
topk: 5
round_decimal: 3
metric_type: "L2"
%4%
metric_type: "%4%"
%5%
search_params: "{\"ef\": 50}"
>
placeholder_tag: "$0">)") %
vec_fid.get() % int_array_fid.get() %
int64_fid.get() % hints_line);
vec_fid.get() % int_array_fid.get() % int64_fid.get() %
metric % hints_line);
proto::plan::PlanNode plan_node;
auto ok =
@@ -384,7 +557,7 @@ TEST_P(ElementFilterSealed, UnaryExpr) {
auto num_queries = 1;
auto seed = 1024;
auto ph_group_raw =
CreatePlaceholderGroup(num_queries, dim, seed, true);
CreatePlaceholderGroupForType(num_queries, dim, seed);
auto ph_group =
ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
@@ -397,24 +570,24 @@ TEST_P(ElementFilterSealed, UnaryExpr) {
// In element-level mode, results should be element indices, not doc offsets
ASSERT_TRUE(search_result->element_level_);
ASSERT_FALSE(search_result->element_indices_.empty());
// Also check seg_offsets_ which stores the doc IDs
ASSERT_FALSE(search_result->seg_offsets_.empty());
ASSERT_EQ(search_result->element_indices_.size(),
search_result->seg_offsets_.size());
// Should have topK results per query
ASSERT_LE(search_result->element_indices_.size(), topK * num_queries);
ASSERT_LE(search_result->element_indices_.size(),
static_cast<size_t>(topK * num_queries));
std::cout << "Element-level search returned:" << std::endl;
for (auto i = 0; i < search_result->seg_offsets_.size(); i++) {
std::cout << "Element-level search returned ("
<< static_cast<int>(elem_type) << "):" << std::endl;
for (size_t i = 0; i < search_result->seg_offsets_.size(); i++) {
std::cout << "doc_id: " << search_result->seg_offsets_[i]
<< ", element_index: "
<< search_result->element_indices_[i] << std::endl;
std::cout << "distance: " << search_result->distances_[i]
<< search_result->element_indices_[i]
<< ", distance: " << search_result->distances_[i]
<< std::endl;
}
// Verify distances are sorted (ascending for L2)
// Verify distances are sorted
for (size_t i = 1; i < search_result->distances_.size(); ++i) {
ASSERT_LE(search_result->distances_[i - 1],
search_result->distances_[i])
@@ -426,16 +599,53 @@ TEST_P(ElementFilterSealed, UnaryExpr) {
INSTANTIATE_TEST_SUITE_P(
ElementFilter,
ElementFilterSealed,
::testing::Combine(::testing::Bool(), // with_hints: true/false
::testing::Bool() // with_load_index: true/false
),
[](const ::testing::TestParamInfo<ElementFilterSealed::ParamType>& info) {
::testing::Values(
// FloatVector with L2
std::make_tuple(false, false, DataType::VECTOR_FLOAT, "L2", 4),
std::make_tuple(false, true, DataType::VECTOR_FLOAT, "L2", 4),
std::make_tuple(true, false, DataType::VECTOR_FLOAT, "L2", 4),
std::make_tuple(true, true, DataType::VECTOR_FLOAT, "L2", 4),
// Float16Vector with L2
std::make_tuple(false, true, DataType::VECTOR_FLOAT16, "L2", 4),
std::make_tuple(true, true, DataType::VECTOR_FLOAT16, "L2", 4),
// BFloat16Vector with L2
std::make_tuple(false, true, DataType::VECTOR_BFLOAT16, "L2", 4),
std::make_tuple(true, true, DataType::VECTOR_BFLOAT16, "L2", 4),
// Int8Vector with L2
std::make_tuple(false, true, DataType::VECTOR_INT8, "L2", 4),
std::make_tuple(true, true, DataType::VECTOR_INT8, "L2", 4),
// BinaryVector with HAMMING (no hints - BIN_FLAT doesn't support iterative filter)
std::make_tuple(false, true, DataType::VECTOR_BINARY, "HAMMING", 32)),
[](const ::testing::TestParamInfo<ElementFilterSealedParam>& info) {
bool with_hints = std::get<0>(info.param);
bool with_load_index = std::get<1>(info.param);
std::string name = "";
name += with_hints ? "WithHints" : "WithoutHints";
name += "_";
name += with_load_index ? "WithLoadIndex" : "WithoutLoadIndex";
DataType elem_type = std::get<2>(info.param);
std::string metric = std::get<3>(info.param);
std::string type_name;
switch (elem_type) {
case DataType::VECTOR_FLOAT:
type_name = "Float";
break;
case DataType::VECTOR_FLOAT16:
type_name = "Float16";
break;
case DataType::VECTOR_BFLOAT16:
type_name = "BFloat16";
break;
case DataType::VECTOR_INT8:
type_name = "Int8";
break;
case DataType::VECTOR_BINARY:
type_name = "Binary";
break;
default:
type_name = "Unknown";
}
std::string name = type_name + "_" + metric;
name += with_hints ? "_WithHints" : "_NoHints";
name += with_load_index ? "_WithIndex" : "_NoIndex";
return name;
});
@@ -619,23 +829,62 @@ TEST(ElementFilter, GrowingSegmentOutOfOrderInsert) {
}
}
// Parameterized test fixture for GrowingIterativeRangeExpr
class ElementFilterGrowing : public ::testing::TestWithParam<bool> {
// Test parameter for Growing: <use_hints, element_type, metric_type, dim>
using ElementFilterGrowingParam = std::tuple<bool, DataType, std::string, int>;
class ElementFilterGrowing
: public ::testing::TestWithParam<ElementFilterGrowingParam> {
protected:
bool
use_hints() const {
return GetParam();
return std::get<0>(GetParam());
}
DataType
element_type() const {
return std::get<1>(GetParam());
}
std::string
metric_type() const {
return std::get<2>(GetParam());
}
int
vec_dim() const {
return std::get<3>(GetParam());
}
// Create placeholder group with element_level = true for element-level search
// Uses regular vector types (not EmbList), as query is single embedding per query
proto::common::PlaceholderGroup
CreatePlaceholderGroupForType(int num_queries, int dim, int seed) {
if (element_type() == DataType::VECTOR_BINARY) {
return CreatePlaceholderGroup<milvus::BinaryVector>(
num_queries, dim, seed, true);
} else if (element_type() == DataType::VECTOR_FLOAT16) {
return CreatePlaceholderGroup<milvus::Float16Vector>(
num_queries, dim, seed, true);
} else if (element_type() == DataType::VECTOR_BFLOAT16) {
return CreatePlaceholderGroup<milvus::BFloat16Vector>(
num_queries, dim, seed, true);
} else if (element_type() == DataType::VECTOR_INT8) {
return CreatePlaceholderGroup<milvus::Int8Vector>(
num_queries, dim, seed, true);
} else {
// VECTOR_FLOAT
return CreatePlaceholderGroup<milvus::FloatVector>(
num_queries, dim, seed, true);
}
}
};
TEST_P(ElementFilterGrowing, RangeExpr) {
bool with_hints = use_hints();
int dim = 4;
DataType elem_type = element_type();
std::string metric = metric_type();
int dim = vec_dim();
auto schema = std::make_shared<Schema>();
auto vec_fid = schema->AddDebugVectorArrayField("structA[array_float_vec]",
DataType::VECTOR_FLOAT,
dim,
knowhere::metric::L2);
auto vec_fid = schema->AddDebugVectorArrayField(
"structA[array_vec]", elem_type, dim, metric);
auto int_array_fid = schema->AddDebugArrayField(
"structA[price_array]", DataType::INT32, false);
@@ -657,7 +906,7 @@ TEST_P(ElementFilterGrowing, RangeExpr) {
->mutable_data()
->Clear();
for (int row = 0; row < N; row++) {
for (size_t row = 0; row < N; row++) {
auto* array_data = field_data->mutable_scalars()
->mutable_array_data()
->mutable_data()
@@ -692,11 +941,11 @@ TEST_P(ElementFilterGrowing, RangeExpr) {
int topK = 5;
// Execute element-level search with iterative filter
// Query: Search array elements where (id % 2 == 0) AND (price_array element in range (100, 400))
{
std::string hints_line =
with_hints ? R"(hints: "iterative_filter")" : "";
std::string raw_plan = boost::str(boost::format(R"(vector_anns: <
std::string raw_plan =
boost::str(boost::format(R"(vector_anns: <
field_id: %1%
predicates: <
element_filter_expr: <
@@ -740,13 +989,13 @@ TEST_P(ElementFilterGrowing, RangeExpr) {
query_info: <
topk: 5
round_decimal: 3
metric_type: "L2"
%4%
metric_type: "%4%"
%5%
search_params: "{\"ef\": 50}"
>
placeholder_tag: "$0">)") %
vec_fid.get() % int_array_fid.get() %
int64_fid.get() % hints_line);
vec_fid.get() % int_array_fid.get() % int64_fid.get() %
metric % hints_line);
proto::plan::PlanNode plan_node;
auto ok =
@@ -759,7 +1008,7 @@ TEST_P(ElementFilterGrowing, RangeExpr) {
auto num_queries = 1;
auto seed = 1024;
auto ph_group_raw =
CreatePlaceholderGroup(num_queries, dim, seed, true);
CreatePlaceholderGroupForType(num_queries, dim, seed);
auto ph_group =
ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
@@ -769,7 +1018,6 @@ TEST_P(ElementFilterGrowing, RangeExpr) {
// Verify results
ASSERT_NE(search_result, nullptr);
// In element-level mode, results should contain element indices
ASSERT_TRUE(search_result->element_level_)
<< "Search should be in element-level mode";
ASSERT_FALSE(search_result->element_indices_.empty())
@@ -780,12 +1028,12 @@ TEST_P(ElementFilterGrowing, RangeExpr) {
search_result->seg_offsets_.size())
<< "Element indices and doc offsets should match in size";
// Should have topK results per query
ASSERT_LE(search_result->element_indices_.size(), topK * num_queries)
ASSERT_LE(search_result->element_indices_.size(),
static_cast<size_t>(topK * num_queries))
<< "Should not exceed topK results";
std::cout << "Growing segment element-level search results:"
<< std::endl;
std::cout << "Growing segment element-level search ("
<< static_cast<int>(elem_type) << "):" << std::endl;
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];
@@ -795,17 +1043,13 @@ TEST_P(ElementFilterGrowing, RangeExpr) {
<< ", element_index=" << elem_idx
<< ", distance=" << distance << std::endl;
// Verify the doc_id satisfies the predicate (id % 2 == 0)
ASSERT_EQ(doc_id % 2, 0) << "Result doc_id " << doc_id
<< " should satisfy (id % 2 == 0)";
// Verify element_idx is valid
ASSERT_GE(elem_idx, 0) << "Element index should be >= 0";
ASSERT_LT(elem_idx, array_len)
<< "Element index should be < array_len";
// Verify element value is in range (100, 400)
// Element value = doc_id * array_len + elem_idx + 1
int element_value = doc_id * array_len + elem_idx + 1;
ASSERT_GT(element_value, 100)
<< "Element value " << element_value << " should be > 100";
@@ -813,7 +1057,7 @@ TEST_P(ElementFilterGrowing, RangeExpr) {
<< "Element value " << element_value << " should be < 400";
}
// Verify distances are sorted (ascending for L2)
// Verify distances are sorted
for (size_t i = 1; i < search_result->distances_.size(); ++i) {
ASSERT_LE(search_result->distances_[i - 1],
search_result->distances_[i])
@@ -825,10 +1069,50 @@ TEST_P(ElementFilterGrowing, RangeExpr) {
INSTANTIATE_TEST_SUITE_P(
ElementFilter,
ElementFilterGrowing,
::testing::Bool(), // with_hints: true/false
[](const ::testing::TestParamInfo<ElementFilterGrowing::ParamType>& info) {
bool with_hints = info.param;
return with_hints ? "WithHints" : "WithoutHints";
::testing::Values(
// FloatVector with L2
std::make_tuple(false, DataType::VECTOR_FLOAT, "L2", 4),
std::make_tuple(true, DataType::VECTOR_FLOAT, "L2", 4),
// Float16Vector with L2
std::make_tuple(false, DataType::VECTOR_FLOAT16, "L2", 4),
std::make_tuple(true, DataType::VECTOR_FLOAT16, "L2", 4),
// BFloat16Vector with L2
std::make_tuple(false, DataType::VECTOR_BFLOAT16, "L2", 4),
std::make_tuple(true, DataType::VECTOR_BFLOAT16, "L2", 4),
// Int8Vector with L2
std::make_tuple(false, DataType::VECTOR_INT8, "L2", 4),
std::make_tuple(true, DataType::VECTOR_INT8, "L2", 4),
// BinaryVector with HAMMING (no hints - brute force doesn't support iterative filter for binary)
std::make_tuple(false, DataType::VECTOR_BINARY, "HAMMING", 32)),
[](const ::testing::TestParamInfo<ElementFilterGrowingParam>& info) {
bool with_hints = std::get<0>(info.param);
DataType elem_type = std::get<1>(info.param);
std::string metric = std::get<2>(info.param);
std::string type_name;
switch (elem_type) {
case DataType::VECTOR_FLOAT:
type_name = "Float";
break;
case DataType::VECTOR_FLOAT16:
type_name = "Float16";
break;
case DataType::VECTOR_BFLOAT16:
type_name = "BFloat16";
break;
case DataType::VECTOR_INT8:
type_name = "Int8";
break;
case DataType::VECTOR_BINARY:
type_name = "Binary";
break;
default:
type_name = "Unknown";
}
std::string name = type_name + "_" + metric;
name += with_hints ? "_WithHints" : "_NoHints";
return name;
});
// Unit tests for ArrayOffsetsGrowing
internal/core/unittest/test_utils/DataGen.h
+153
View File
@@ -34,6 +34,7 @@
#include "segcore/Utils.h"
#include "knowhere/comp/index_param.h"
#include "knowhere/operands.h"
#include "PbHelper.h"
#include "segcore/collection_c.h"
@@ -1715,6 +1716,158 @@ GenVecIndexing(int64_t N,
return indexing;
}
// GenVecIndexing for Float16Vector
inline std::unique_ptr<milvus::index::VectorIndex>
GenVecIndexingFloat16(int64_t N,
int64_t dim,
const knowhere::fp16* vec,
const char* index_type,
bool use_knowhere_build_pool = true) {
auto conf =
knowhere::Json{{knowhere::meta::METRIC_TYPE, knowhere::metric::L2},
{knowhere::meta::DIM, std::to_string(dim)},
{knowhere::indexparam::NLIST, "1024"},
{knowhere::meta::DEVICE_ID, 0}};
auto database = knowhere::GenDataSet(N, dim, vec);
milvus::storage::FieldDataMeta field_data_meta{1, 2, 3, 100};
milvus::storage::IndexMeta index_meta{3, 100, 1000, 1};
milvus::storage::StorageConfig storage_config;
storage_config.storage_type = "local";
storage_config.root_path = TestRemotePath;
auto chunk_manager = milvus::storage::CreateChunkManager(storage_config);
auto fs = milvus::storage::InitArrowFileSystem(storage_config);
milvus::storage::FileManagerContext file_manager_context(
field_data_meta, index_meta, chunk_manager, fs);
auto indexing = std::make_unique<index::VectorMemIndex<knowhere::fp16>>(
DataType::NONE,
index_type,
knowhere::metric::L2,
knowhere::Version::GetCurrentVersion().VersionNumber(),
use_knowhere_build_pool,
file_manager_context);
indexing->BuildWithDataset(database, conf);
auto create_index_result = indexing->Upload();
auto index_files = create_index_result->GetIndexFiles();
conf["index_files"] = index_files;
conf[milvus::LOAD_PRIORITY] = milvus::proto::common::LoadPriority::HIGH;
indexing->Load(milvus::tracer::TraceContext{}, conf);
return indexing;
}
// GenVecIndexing for BFloat16Vector
inline std::unique_ptr<milvus::index::VectorIndex>
GenVecIndexingBFloat16(int64_t N,
int64_t dim,
const knowhere::bf16* vec,
const char* index_type,
bool use_knowhere_build_pool = true) {
auto conf =
knowhere::Json{{knowhere::meta::METRIC_TYPE, knowhere::metric::L2},
{knowhere::meta::DIM, std::to_string(dim)},
{knowhere::indexparam::NLIST, "1024"},
{knowhere::meta::DEVICE_ID, 0}};
auto database = knowhere::GenDataSet(N, dim, vec);
milvus::storage::FieldDataMeta field_data_meta{1, 2, 3, 100};
milvus::storage::IndexMeta index_meta{3, 100, 1000, 1};
milvus::storage::StorageConfig storage_config;
storage_config.storage_type = "local";
storage_config.root_path = TestRemotePath;
auto chunk_manager = milvus::storage::CreateChunkManager(storage_config);
auto fs = milvus::storage::InitArrowFileSystem(storage_config);
milvus::storage::FileManagerContext file_manager_context(
field_data_meta, index_meta, chunk_manager, fs);
auto indexing = std::make_unique<index::VectorMemIndex<knowhere::bf16>>(
DataType::NONE,
index_type,
knowhere::metric::L2,
knowhere::Version::GetCurrentVersion().VersionNumber(),
use_knowhere_build_pool,
file_manager_context);
indexing->BuildWithDataset(database, conf);
auto create_index_result = indexing->Upload();
auto index_files = create_index_result->GetIndexFiles();
conf["index_files"] = index_files;
conf[milvus::LOAD_PRIORITY] = milvus::proto::common::LoadPriority::HIGH;
indexing->Load(milvus::tracer::TraceContext{}, conf);
return indexing;
}
// GenVecIndexing for Int8Vector
inline std::unique_ptr<milvus::index::VectorIndex>
GenVecIndexingInt8(int64_t N,
int64_t dim,
const int8_t* vec,
const char* index_type,
bool use_knowhere_build_pool = true) {
auto conf =
knowhere::Json{{knowhere::meta::METRIC_TYPE, knowhere::metric::L2},
{knowhere::meta::DIM, std::to_string(dim)},
{knowhere::indexparam::NLIST, "1024"},
{knowhere::meta::DEVICE_ID, 0}};
auto database = knowhere::GenDataSet(N, dim, vec);
milvus::storage::FieldDataMeta field_data_meta{1, 2, 3, 100};
milvus::storage::IndexMeta index_meta{3, 100, 1000, 1};
milvus::storage::StorageConfig storage_config;
storage_config.storage_type = "local";
storage_config.root_path = TestRemotePath;
auto chunk_manager = milvus::storage::CreateChunkManager(storage_config);
auto fs = milvus::storage::InitArrowFileSystem(storage_config);
milvus::storage::FileManagerContext file_manager_context(
field_data_meta, index_meta, chunk_manager, fs);
auto indexing = std::make_unique<index::VectorMemIndex<int8_t>>(
DataType::NONE,
index_type,
knowhere::metric::L2,
knowhere::Version::GetCurrentVersion().VersionNumber(),
use_knowhere_build_pool,
file_manager_context);
indexing->BuildWithDataset(database, conf);
auto create_index_result = indexing->Upload();
auto index_files = create_index_result->GetIndexFiles();
conf["index_files"] = index_files;
conf[milvus::LOAD_PRIORITY] = milvus::proto::common::LoadPriority::HIGH;
indexing->Load(milvus::tracer::TraceContext{}, conf);
return indexing;
}
// GenVecIndexing for BinaryVector
inline std::unique_ptr<milvus::index::VectorIndex>
GenVecIndexingBinary(int64_t N,
int64_t dim,
const uint8_t* vec,
const char* index_type,
bool use_knowhere_build_pool = true) {
auto conf =
knowhere::Json{{knowhere::meta::METRIC_TYPE, knowhere::metric::HAMMING},
{knowhere::meta::DIM, std::to_string(dim)},
{knowhere::indexparam::NLIST, "1024"},
{knowhere::meta::DEVICE_ID, 0}};
auto database = knowhere::GenDataSet(N, dim, vec);
milvus::storage::FieldDataMeta field_data_meta{1, 2, 3, 100};
milvus::storage::IndexMeta index_meta{3, 100, 1000, 1};
milvus::storage::StorageConfig storage_config;
storage_config.storage_type = "local";
storage_config.root_path = TestRemotePath;
auto chunk_manager = milvus::storage::CreateChunkManager(storage_config);
auto fs = milvus::storage::InitArrowFileSystem(storage_config);
milvus::storage::FileManagerContext file_manager_context(
field_data_meta, index_meta, chunk_manager, fs);
auto indexing = std::make_unique<index::VectorMemIndex<uint8_t>>(
DataType::NONE,
index_type,
knowhere::metric::HAMMING,
knowhere::Version::GetCurrentVersion().VersionNumber(),
use_knowhere_build_pool,
file_manager_context);
indexing->BuildWithDataset(database, conf);
auto create_index_result = indexing->Upload();
auto index_files = create_index_result->GetIndexFiles();
conf["index_files"] = index_files;
conf[milvus::LOAD_PRIORITY] = milvus::proto::common::LoadPriority::HIGH;
indexing->Load(milvus::tracer::TraceContext{}, conf);
return indexing;
}
template <typename T>
inline index::IndexBasePtr
GenScalarIndexing(int64_t N, const T* data) {
internal/proxy/util.go
+3 -6
View File
@@ -654,14 +654,11 @@ func ValidateFieldsInStruct(field *schemapb.FieldSchema, schema *schemapb.Collec
return err
}
} else {
// TODO(SpadeA): only support float vector now
if field.GetElementType() != schemapb.DataType_FloatVector {
return fmt.Errorf("Unsupported element type of array field %s, now only float vector is supported", field.Name)
// ArrayOfVector: support FloatVector, Float16Vector, BFloat16Vector, Int8Vector, BinaryVector
if !typeutil.IsFixDimVectorType(field.GetElementType()) {
return fmt.Errorf("Unsupported element type %s of ArrayOfVector field %s, only fixed dimension vector types are supported", field.GetElementType().String(), field.Name)
}
// if !typeutil.IsVectorType(field.GetElementType()) {
// return fmt.Errorf("Inconsistent schema: element type of array field %s is not a vector type", field.Name)
// }
err = validateDimension(field)
if err != nil {
return err
internal/proxy/util_test.go
+1 -1
View File
@@ -4055,7 +4055,7 @@ func TestValidateFieldsInStruct(t *testing.T) {
}
err := ValidateFieldsInStruct(field, schema)
assert.Error(t, err)
assert.Contains(t, err.Error(), "Unsupported element type of array field array_vector_with_scalar, now only float vector is supported")
assert.Contains(t, err.Error(), "only fixed dimension vector types are supported")
})
t.Run("array of vector missing dimension", func(t *testing.T) {
internal/storage/payload_writer_test.go
+51
View File
@@ -641,6 +641,57 @@ func TestPayloadWriter_ArrayOfVector(t *testing.T) {
require.NoError(t, err)
require.Equal(t, numRows, length)
})
t.Run("Test ArrayOfInt8Vector - Basic", func(t *testing.T) {
dim := 64
numRows := 50
vectorsPerRow := 4
// Create test data
vectorArrayData := &VectorArrayFieldData{
Data: make([]*schemapb.VectorField, numRows),
ElementType: schemapb.DataType_Int8Vector,
Dim: int64(dim),
}
for i := 0; i < numRows; i++ {
// Int8 vectors are stored as bytes (1 byte per element)
byteData := make([]byte, vectorsPerRow*dim)
for j := 0; j < len(byteData); j++ {
byteData[j] = byte((i*50 + j) % 256)
}
vectorArrayData.Data[i] = &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_Int8Vector{
Int8Vector: byteData,
},
}
}
w, err := NewPayloadWriter(
schemapb.DataType_ArrayOfVector,
WithDim(dim),
WithElementType(schemapb.DataType_Int8Vector),
)
require.NoError(t, err)
require.NotNil(t, w)
err = w.AddVectorArrayFieldDataToPayload(vectorArrayData)
require.NoError(t, err)
err = w.FinishPayloadWriter()
require.NoError(t, err)
// Verify results
buffer, err := w.GetPayloadBufferFromWriter()
require.NoError(t, err)
require.NotEmpty(t, buffer)
length, err := w.GetPayloadLengthFromWriter()
require.NoError(t, err)
require.Equal(t, numRows, length)
})
}
func TestParquetEncoding(t *testing.T) {
internal/storage/serde.go
+33 -4
View File
@@ -877,7 +877,7 @@ func getArrayOfVectorArrowType(elementType schemapb.DataType, dim int) arrow.Dat
}
// deserializeArrayOfVector deserializes ArrayOfVector data with known element type
func deserializeArrayOfVector(a arrow.Array, i int, elementType schemapb.DataType, dim int64, shouldCopy bool) (any, bool) {
func deserializeArrayOfVector(a arrow.Array, i int, elementType schemapb.DataType, dim int64, _ bool) (any, bool) {
if a.IsNull(i) {
return nil, true
}
@@ -889,15 +889,44 @@ func deserializeArrayOfVector(a arrow.Array, i int, elementType schemapb.DataTyp
start, end := arr.ValueOffsets(i)
totalElements := end - start
// Handle empty vector array - return empty VectorField based on element type
if totalElements == 0 {
return nil, false
switch elementType {
case schemapb.DataType_FloatVector:
return &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{}}},
}, true
case schemapb.DataType_BinaryVector:
return &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{}},
}, true
case schemapb.DataType_Float16Vector:
return &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{}},
}, true
case schemapb.DataType_BFloat16Vector:
return &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{}},
}, true
case schemapb.DataType_Int8Vector:
return &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_Int8Vector{Int8Vector: []byte{}},
}, true
default:
return nil, false
}
}
valuesArray := arr.ListValues()
binaryArray, ok := valuesArray.(*array.FixedSizeBinary)
if !ok {
// empty array
return nil, true
return nil, false
}
numVectors := int(totalElements)
internal/storage/serde_test.go
+332 -42
View File
@@ -17,6 +17,7 @@
package storage
import (
"fmt"
"io"
"reflect"
"testing"
@@ -353,6 +354,82 @@ func TestArrayOfVectorSerialization(t *testing.T) {
},
},
},
{
name: "Float16Vector array",
elementType: schemapb.DataType_Float16Vector,
dim: 4,
vectors: []*schemapb.VectorField{
{
Dim: 4,
Data: &schemapb.VectorField_Float16Vector{
Float16Vector: []byte{1, 2, 3, 4, 5, 6, 7, 8}, // 4 dims * 2 bytes
},
},
{
Dim: 4,
Data: &schemapb.VectorField_Float16Vector{
Float16Vector: []byte{9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24}, // 8 dims * 2 bytes (2 vectors)
},
},
},
},
{
name: "BFloat16Vector array",
elementType: schemapb.DataType_BFloat16Vector,
dim: 4,
vectors: []*schemapb.VectorField{
{
Dim: 4,
Data: &schemapb.VectorField_Bfloat16Vector{
Bfloat16Vector: []byte{1, 2, 3, 4, 5, 6, 7, 8}, // 4 dims * 2 bytes
},
},
{
Dim: 4,
Data: &schemapb.VectorField_Bfloat16Vector{
Bfloat16Vector: []byte{9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24}, // 8 dims * 2 bytes (2 vectors)
},
},
},
},
{
name: "Int8Vector array",
elementType: schemapb.DataType_Int8Vector,
dim: 4,
vectors: []*schemapb.VectorField{
{
Dim: 4,
Data: &schemapb.VectorField_Int8Vector{
Int8Vector: []byte{1, 2, 3, 4}, // 4 dims * 1 byte
},
},
{
Dim: 4,
Data: &schemapb.VectorField_Int8Vector{
Int8Vector: []byte{5, 6, 7, 8, 9, 10, 11, 12}, // 8 dims * 1 byte (2 vectors)
},
},
},
},
{
name: "BinaryVector array",
elementType: schemapb.DataType_BinaryVector,
dim: 32, // Must be multiple of 8
vectors: []*schemapb.VectorField{
{
Dim: 32,
Data: &schemapb.VectorField_BinaryVector{
BinaryVector: []byte{1, 2, 3, 4}, // 32 dims / 8 = 4 bytes per vector
},
},
{
Dim: 32,
Data: &schemapb.VectorField_BinaryVector{
BinaryVector: []byte{5, 6, 7, 8, 9, 10, 11, 12}, // 2 vectors * 4 bytes
},
},
},
},
}
for _, tt := range tests {
@@ -386,10 +463,27 @@ func TestArrayOfVectorSerialization(t *testing.T) {
assert.Equal(t, expectedVector.GetDim(), resultVector.GetDim())
if tt.elementType == schemapb.DataType_FloatVector {
switch tt.elementType {
case schemapb.DataType_FloatVector:
expectedData := expectedVector.GetFloatVector().GetData()
resultData := resultVector.GetFloatVector().GetData()
assert.Equal(t, expectedData, resultData)
case schemapb.DataType_Float16Vector:
expectedData := expectedVector.GetFloat16Vector()
resultData := resultVector.GetFloat16Vector()
assert.Equal(t, expectedData, resultData)
case schemapb.DataType_BFloat16Vector:
expectedData := expectedVector.GetBfloat16Vector()
resultData := resultVector.GetBfloat16Vector()
assert.Equal(t, expectedData, resultData)
case schemapb.DataType_Int8Vector:
expectedData := expectedVector.GetInt8Vector()
resultData := resultVector.GetInt8Vector()
assert.Equal(t, expectedData, resultData)
case schemapb.DataType_BinaryVector:
expectedData := expectedVector.GetBinaryVector()
resultData := resultVector.GetBinaryVector()
assert.Equal(t, expectedData, resultData)
}
}
}
@@ -397,29 +491,109 @@ func TestArrayOfVectorSerialization(t *testing.T) {
}
}
func TestArrayOfVectorIntegration(t *testing.T) {
// Test the full integration with BuildRecord
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "vec_array",
DataType: schemapb.DataType_ArrayOfVector,
ElementType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "4"},
},
},
},
func TestArrayOfVectorEmptyArray(t *testing.T) {
tests := []struct {
name string
elementType schemapb.DataType
dim int
}{
{"FloatVector empty", schemapb.DataType_FloatVector, 4},
{"Float16Vector empty", schemapb.DataType_Float16Vector, 4},
{"BFloat16Vector empty", schemapb.DataType_BFloat16Vector, 4},
{"Int8Vector empty", schemapb.DataType_Int8Vector, 4},
{"BinaryVector empty", schemapb.DataType_BinaryVector, 32},
}
// Create insert data
insertData := &InsertData{
Data: map[FieldID]FieldData{
100: &VectorArrayFieldData{
Data: []*schemapb.VectorField{
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
entry := serdeMap[schemapb.DataType_ArrayOfVector]
arrowType := entry.arrowType(tt.dim, tt.elementType)
assert.NotNil(t, arrowType)
// Create empty VectorField based on element type
var emptyVector *schemapb.VectorField
switch tt.elementType {
case schemapb.DataType_FloatVector:
emptyVector = &schemapb.VectorField{
Dim: int64(tt.dim),
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{}}},
}
case schemapb.DataType_Float16Vector:
emptyVector = &schemapb.VectorField{
Dim: int64(tt.dim),
Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{}},
}
case schemapb.DataType_BFloat16Vector:
emptyVector = &schemapb.VectorField{
Dim: int64(tt.dim),
Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{}},
}
case schemapb.DataType_Int8Vector:
emptyVector = &schemapb.VectorField{
Dim: int64(tt.dim),
Data: &schemapb.VectorField_Int8Vector{Int8Vector: []byte{}},
}
case schemapb.DataType_BinaryVector:
emptyVector = &schemapb.VectorField{
Dim: int64(tt.dim),
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{}},
}
}
builder := array.NewBuilder(memory.DefaultAllocator, arrowType)
defer builder.Release()
// Serialize empty vector
ok := entry.serialize(builder, emptyVector, tt.elementType)
assert.True(t, ok)
arr := builder.NewArray()
defer arr.Release()
// Deserialize and verify
result, ok := entry.deserialize(arr, 0, tt.elementType, tt.dim, false)
assert.True(t, ok, "deserialize should succeed for empty array")
assert.NotNil(t, result)
resultVector, ok := result.(*schemapb.VectorField)
assert.True(t, ok)
assert.Equal(t, int64(tt.dim), resultVector.GetDim())
// Verify data is empty
switch tt.elementType {
case schemapb.DataType_FloatVector:
assert.Empty(t, resultVector.GetFloatVector().GetData())
case schemapb.DataType_Float16Vector:
assert.Empty(t, resultVector.GetFloat16Vector())
case schemapb.DataType_BFloat16Vector:
assert.Empty(t, resultVector.GetBfloat16Vector())
case schemapb.DataType_Int8Vector:
assert.Empty(t, resultVector.GetInt8Vector())
case schemapb.DataType_BinaryVector:
assert.Empty(t, resultVector.GetBinaryVector())
}
})
}
}
func TestArrayOfVectorIntegration(t *testing.T) {
tests := []struct {
name string
elementType schemapb.DataType
dim int
elementTypeCode string // Expected element type code in metadata
createVectors func(dim int) []*schemapb.VectorField
}{
{
name: "FloatVector",
elementType: schemapb.DataType_FloatVector,
dim: 4,
elementTypeCode: "101",
createVectors: func(dim int) []*schemapb.VectorField {
return []*schemapb.VectorField{
{
Dim: 4,
Dim: int64(dim),
Data: &schemapb.VectorField_FloatVector{
FloatVector: &schemapb.FloatArray{
Data: []float32{1.0, 2.0, 3.0, 4.0},
@@ -427,44 +601,160 @@ func TestArrayOfVectorIntegration(t *testing.T) {
},
},
{
Dim: 4,
Dim: int64(dim),
Data: &schemapb.VectorField_FloatVector{
FloatVector: &schemapb.FloatArray{
Data: []float32{5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0},
},
},
},
},
}
},
},
{
name: "Float16Vector",
elementType: schemapb.DataType_Float16Vector,
dim: 4,
elementTypeCode: "102",
createVectors: func(dim int) []*schemapb.VectorField {
return []*schemapb.VectorField{
{
Dim: int64(dim),
Data: &schemapb.VectorField_Float16Vector{
Float16Vector: []byte{1, 2, 3, 4, 5, 6, 7, 8},
},
},
{
Dim: int64(dim),
Data: &schemapb.VectorField_Float16Vector{
Float16Vector: []byte{9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24},
},
},
}
},
},
{
name: "BFloat16Vector",
elementType: schemapb.DataType_BFloat16Vector,
dim: 4,
elementTypeCode: "103",
createVectors: func(dim int) []*schemapb.VectorField {
return []*schemapb.VectorField{
{
Dim: int64(dim),
Data: &schemapb.VectorField_Bfloat16Vector{
Bfloat16Vector: []byte{1, 2, 3, 4, 5, 6, 7, 8},
},
},
{
Dim: int64(dim),
Data: &schemapb.VectorField_Bfloat16Vector{
Bfloat16Vector: []byte{9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24},
},
},
}
},
},
{
name: "Int8Vector",
elementType: schemapb.DataType_Int8Vector,
dim: 4,
elementTypeCode: "105",
createVectors: func(dim int) []*schemapb.VectorField {
return []*schemapb.VectorField{
{
Dim: int64(dim),
Data: &schemapb.VectorField_Int8Vector{
Int8Vector: []byte{1, 2, 3, 4},
},
},
{
Dim: int64(dim),
Data: &schemapb.VectorField_Int8Vector{
Int8Vector: []byte{5, 6, 7, 8, 9, 10, 11, 12},
},
},
}
},
},
{
name: "BinaryVector",
elementType: schemapb.DataType_BinaryVector,
dim: 32,
elementTypeCode: "100",
createVectors: func(dim int) []*schemapb.VectorField {
return []*schemapb.VectorField{
{
Dim: int64(dim),
Data: &schemapb.VectorField_BinaryVector{
BinaryVector: []byte{1, 2, 3, 4},
},
},
{
Dim: int64(dim),
Data: &schemapb.VectorField_BinaryVector{
BinaryVector: []byte{5, 6, 7, 8, 9, 10, 11, 12},
},
},
}
},
},
}
arrowSchema, err := ConvertToArrowSchema(schema, false)
assert.NoError(t, err)
assert.NotNil(t, arrowSchema)
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "vec_array",
DataType: schemapb.DataType_ArrayOfVector,
ElementType: tt.elementType,
TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: fmt.Sprintf("%d", tt.dim)},
},
},
},
}
recordBuilder := array.NewRecordBuilder(memory.DefaultAllocator, arrowSchema)
defer recordBuilder.Release()
insertData := &InsertData{
Data: map[FieldID]FieldData{
100: &VectorArrayFieldData{
Data: tt.createVectors(tt.dim),
ElementType: tt.elementType,
Dim: int64(tt.dim),
},
},
}
err = BuildRecord(recordBuilder, insertData, schema)
assert.NoError(t, err)
arrowSchema, err := ConvertToArrowSchema(schema, false)
assert.NoError(t, err)
assert.NotNil(t, arrowSchema)
record := recordBuilder.NewRecord()
defer record.Release()
recordBuilder := array.NewRecordBuilder(memory.DefaultAllocator, arrowSchema)
defer recordBuilder.Release()
assert.Equal(t, int64(2), record.NumRows())
assert.Equal(t, int64(1), record.NumCols())
err = BuildRecord(recordBuilder, insertData, schema)
assert.NoError(t, err)
field := arrowSchema.Field(0)
assert.True(t, field.HasMetadata())
record := recordBuilder.NewRecord()
defer record.Release()
elementTypeStr, ok := field.Metadata.GetValue("elementType")
assert.True(t, ok)
assert.Equal(t, "101", elementTypeStr) // FloatVector = 101
assert.Equal(t, int64(2), record.NumRows())
assert.Equal(t, int64(1), record.NumCols())
dimStr, ok := field.Metadata.GetValue("dim")
assert.True(t, ok)
assert.Equal(t, "4", dimStr)
field := arrowSchema.Field(0)
assert.True(t, field.HasMetadata())
elementTypeStr, ok := field.Metadata.GetValue("elementType")
assert.True(t, ok)
assert.Equal(t, tt.elementTypeCode, elementTypeStr)
dimStr, ok := field.Metadata.GetValue("dim")
assert.True(t, ok)
assert.Equal(t, fmt.Sprintf("%d", tt.dim), dimStr)
})
}
}
func TestActualSizeInBytesSlicedFixedSizeBinary(t *testing.T) {
internal/util/importutilv2/parquet/reader_test.go
+105 -91
View File
@@ -617,115 +617,129 @@ func TestParquetReader(t *testing.T) {
func TestParquetReaderWithStructArray(t *testing.T) {
ctx := context.Background()
t.Run("test struct array field reading", func(t *testing.T) {
// Create schema with StructArrayField
schema := &schemapb.CollectionSchema{
Name: "test_struct_array",
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "id",
IsPrimaryKey: true,
DataType: schemapb.DataType_Int64,
},
{
FieldID: 101,
Name: "varchar_field",
DataType: schemapb.DataType_VarChar,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.MaxLengthKey, Value: "100"},
vectorTypeTests := []struct {
name string
elementType schemapb.DataType
dim string
}{
{"FloatVector", schemapb.DataType_FloatVector, "4"},
{"Float16Vector", schemapb.DataType_Float16Vector, "4"},
{"BFloat16Vector", schemapb.DataType_BFloat16Vector, "4"},
{"Int8Vector", schemapb.DataType_Int8Vector, "4"},
{"BinaryVector", schemapb.DataType_BinaryVector, "32"},
}
for _, vt := range vectorTypeTests {
t.Run("test struct array with "+vt.name, func(t *testing.T) {
// Create schema with StructArrayField
schema := &schemapb.CollectionSchema{
Name: "test_struct_array_" + vt.name,
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "id",
IsPrimaryKey: true,
DataType: schemapb.DataType_Int64,
},
{
FieldID: 101,
Name: "varchar_field",
DataType: schemapb.DataType_VarChar,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.MaxLengthKey, Value: "100"},
},
},
},
},
StructArrayFields: []*schemapb.StructArrayFieldSchema{
{
FieldID: 200,
Name: "struct_array",
Fields: []*schemapb.FieldSchema{
{
FieldID: 201,
Name: "struct_array[int_array]",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int32,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.MaxCapacityKey, Value: "20"},
StructArrayFields: []*schemapb.StructArrayFieldSchema{
{
FieldID: 200,
Name: "struct_array",
Fields: []*schemapb.FieldSchema{
{
FieldID: 201,
Name: "struct_array[int_array]",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int32,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.MaxCapacityKey, Value: "20"},
},
},
},
{
FieldID: 202,
Name: "struct_array[float_array]",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Float,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.MaxCapacityKey, Value: "20"},
{
FieldID: 202,
Name: "struct_array[float_array]",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Float,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.MaxCapacityKey, Value: "20"},
},
},
},
{
FieldID: 203,
Name: "struct_array[vector_array]",
DataType: schemapb.DataType_ArrayOfVector,
ElementType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "4"},
{Key: common.MaxCapacityKey, Value: "20"},
{
FieldID: 203,
Name: "struct_array[vector_array]",
DataType: schemapb.DataType_ArrayOfVector,
ElementType: vt.elementType,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: vt.dim},
{Key: common.MaxCapacityKey, Value: "20"},
},
},
},
},
},
},
}
}
// Create test data file
filePath := fmt.Sprintf("/tmp/test_struct_array_%d.parquet", rand.Int())
defer os.Remove(filePath)
// Create test data file
filePath := fmt.Sprintf("/tmp/test_struct_array_%s_%d.parquet", vt.name, rand.Int())
defer os.Remove(filePath)
numRows := 50
f, err := os.Create(filePath)
assert.NoError(t, err)
numRows := 50
f, err := os.Create(filePath)
assert.NoError(t, err)
// Use writeParquet to create test file
insertData, err := writeParquet(f, schema, numRows, 0)
assert.NoError(t, err)
f.Close()
// Use writeParquet to create test file
insertData, err := writeParquet(f, schema, numRows, 0)
assert.NoError(t, err)
f.Close()
// Verify the insert data contains struct fields
assert.Contains(t, insertData.Data, int64(201)) // int_array field
assert.Contains(t, insertData.Data, int64(202)) // float_array field
assert.Contains(t, insertData.Data, int64(203)) // vector_array field
// Verify the insert data contains struct fields
assert.Contains(t, insertData.Data, int64(201)) // int_array field
assert.Contains(t, insertData.Data, int64(202)) // float_array field
assert.Contains(t, insertData.Data, int64(203)) // vector_array field
// Now test reading the file using ChunkManager
factory := storage.NewChunkManagerFactory("local", objectstorage.RootPath("/tmp"))
cm, err := factory.NewPersistentStorageChunkManager(ctx)
assert.NoError(t, err)
// Now test reading the file using ChunkManager
factory := storage.NewChunkManagerFactory("local", objectstorage.RootPath("/tmp"))
cm, err := factory.NewPersistentStorageChunkManager(ctx)
assert.NoError(t, err)
reader, err := NewReader(ctx, cm, schema, filePath, 64*1024*1024)
assert.NoError(t, err)
defer reader.Close()
reader, err := NewReader(ctx, cm, schema, filePath, 64*1024*1024)
assert.NoError(t, err)
defer reader.Close()
// Read data
readData, err := reader.Read()
assert.NoError(t, err)
assert.NotNil(t, readData)
// Read data
readData, err := reader.Read()
assert.NoError(t, err)
assert.NotNil(t, readData)
// Verify the data includes struct fields
assert.Contains(t, readData.Data, int64(201)) // int_array field ID
assert.Contains(t, readData.Data, int64(202)) // float_array field ID
assert.Contains(t, readData.Data, int64(203)) // vector_array field ID
// Verify the data includes struct fields
assert.Contains(t, readData.Data, int64(201)) // int_array field ID
assert.Contains(t, readData.Data, int64(202)) // float_array field ID
assert.Contains(t, readData.Data, int64(203)) // vector_array field ID
// Check row count matches
assert.Equal(t, numRows, readData.Data[100].RowNum()) // id field
assert.Equal(t, numRows, readData.Data[101].RowNum()) // varchar_field
assert.Equal(t, numRows, readData.Data[201].RowNum()) // int_array
assert.Equal(t, numRows, readData.Data[202].RowNum()) // float_array
assert.Equal(t, numRows, readData.Data[203].RowNum()) // vector_array
// Check row count matches
assert.Equal(t, numRows, readData.Data[100].RowNum()) // id field
assert.Equal(t, numRows, readData.Data[101].RowNum()) // varchar_field
assert.Equal(t, numRows, readData.Data[201].RowNum()) // int_array
assert.Equal(t, numRows, readData.Data[202].RowNum()) // float_array
assert.Equal(t, numRows, readData.Data[203].RowNum()) // vector_array
// Verify data content matches
for fieldID, originalData := range insertData.Data {
readFieldData, ok := readData.Data[fieldID]
assert.True(t, ok, "field %d not found in read data", fieldID)
assert.Equal(t, originalData.RowNum(), readFieldData.RowNum(), "row count mismatch for field %d", fieldID)
}
})
// Verify data content matches
for fieldID, originalData := range insertData.Data {
readFieldData, ok := readData.Data[fieldID]
assert.True(t, ok, "field %d not found in read data", fieldID)
assert.Equal(t, originalData.RowNum(), readFieldData.RowNum(), "row count mismatch for field %d", fieldID)
}
})
}
}
func TestParquetReaderError(t *testing.T) {
internal/util/importutilv2/parquet/struct_field_reader.go
+68 -6
View File
@@ -329,17 +329,79 @@ func (r *StructFieldReader) readArrayOfVectorField(chunked *arrow.Chunked) (any,
result = append(result, vectorField)
}
case schemapb.DataType_BinaryVector:
return nil, nil, merr.WrapErrImportFailed("ArrayOfVector with BinaryVector element type is not implemented yet")
case schemapb.DataType_Float16Vector:
return nil, nil, merr.WrapErrImportFailed("ArrayOfVector with Float16Vector element type is not implemented yet")
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 len(allVectors) >= 0 {
vectorField := &schemapb.VectorField{
Dim: int64(r.dim),
Data: &schemapb.VectorField_Float16Vector{
Float16Vector: allVectors,
},
}
result = append(result, vectorField)
}
case schemapb.DataType_BFloat16Vector:
return nil, nil, merr.WrapErrImportFailed("ArrayOfVector with BFloat16Vector element type is not implemented yet")
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 len(allVectors) >= 0 {
vectorField := &schemapb.VectorField{
Dim: int64(r.dim),
Data: &schemapb.VectorField_Bfloat16Vector{
Bfloat16Vector: allVectors,
},
}
result = append(result, vectorField)
}
case schemapb.DataType_Int8Vector:
return nil, nil, merr.WrapErrImportFailed("ArrayOfVector with Int8Vector element type is not implemented yet")
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 len(allVectors) >= 0 {
vectorField := &schemapb.VectorField{
Dim: int64(r.dim),
Data: &schemapb.VectorField_Int8Vector{
Int8Vector: allVectors,
},
}
result = append(result, vectorField)
}
case schemapb.DataType_BinaryVector:
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 len(allVectors) >= 0 {
vectorField := &schemapb.VectorField{
Dim: int64(r.dim),
Data: &schemapb.VectorField_BinaryVector{
BinaryVector: allVectors,
},
}
result = append(result, vectorField)
}
case schemapb.DataType_SparseFloatVector:
return nil, nil, merr.WrapErrImportFailed("ArrayOfVector with SparseFloatVector element type is not implemented yet")
internal/util/importutilv2/parquet/util.go
+1 -1
View File
@@ -330,7 +330,7 @@ func convertElementTypeOfVectorArrayToArrowType(field *schemapb.FieldSchema) (ar
case schemapb.DataType_BinaryVector:
elemType = arrow.ListOf(arrow.PrimitiveTypes.Uint8)
case schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
elemType = arrow.ListOf(arrow.PrimitiveTypes.Float32)
elemType = arrow.ListOf(arrow.PrimitiveTypes.Uint8)
case schemapb.DataType_Int8Vector:
elemType = arrow.ListOf(arrow.PrimitiveTypes.Int8)
default:
internal/util/testutil/test_util.go
+94 -16
View File
@@ -271,6 +271,14 @@ func CreateInsertData(schema *schemapb.CollectionSchema, rows int, nullPercent .
switch f.GetElementType() {
case schemapb.DataType_FloatVector:
insertData.Data[f.FieldID].AppendDataRows(testutils.GenerateArrayOfFloatVectorArray(rows, int(dim)))
case schemapb.DataType_Float16Vector:
insertData.Data[f.FieldID].AppendDataRows(testutils.GenerateArrayOfFloat16VectorArray(rows, int(dim)))
case schemapb.DataType_BFloat16Vector:
insertData.Data[f.FieldID].AppendDataRows(testutils.GenerateArrayOfBFloat16VectorArray(rows, int(dim)))
case schemapb.DataType_Int8Vector:
insertData.Data[f.FieldID].AppendDataRows(testutils.GenerateArrayOfInt8VectorArray(rows, int(dim)))
case schemapb.DataType_BinaryVector:
insertData.Data[f.FieldID].AppendDataRows(testutils.GenerateArrayOfBinaryVectorArray(rows, int(dim)))
default:
panic(fmt.Sprintf("unimplemented data type: %s", f.GetElementType().String()))
}
@@ -1060,7 +1068,7 @@ func BuildArrayData(schema *schemapb.CollectionSchema, insertData *storage.Inser
case schemapb.DataType_BinaryVector:
arrType = arrow.ListOf(arrow.PrimitiveTypes.Uint8)
case schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
arrType = arrow.ListOf(arrow.PrimitiveTypes.Float32)
arrType = arrow.ListOf(arrow.PrimitiveTypes.Uint8)
case schemapb.DataType_Int8Vector:
arrType = arrow.ListOf(arrow.PrimitiveTypes.Int8)
default:
@@ -1118,9 +1126,31 @@ func BuildArrayData(schema *schemapb.CollectionSchema, insertData *storage.Inser
}
case schemapb.DataType_ArrayOfVector:
vectorField := rowData.(*schemapb.VectorField)
if vectorField.GetFloatVector() != nil {
dim, _ := typeutil.GetDim(subField)
arrayLen = len(vectorField.GetFloatVector().Data) / int(dim)
dim, _ := typeutil.GetDim(subField)
switch subField.ElementType {
case schemapb.DataType_FloatVector:
if vectorField.GetFloatVector() != nil {
arrayLen = len(vectorField.GetFloatVector().Data) / int(dim)
}
case schemapb.DataType_BinaryVector:
if vectorField.GetBinaryVector() != nil {
bytesPerVector := int(dim) / 8
arrayLen = len(vectorField.GetBinaryVector()) / bytesPerVector
}
case schemapb.DataType_Float16Vector:
if vectorField.GetFloat16Vector() != nil {
bytesPerVector := int(dim) * 2
arrayLen = len(vectorField.GetFloat16Vector()) / bytesPerVector
}
case schemapb.DataType_BFloat16Vector:
if vectorField.GetBfloat16Vector() != nil {
bytesPerVector := int(dim) * 2
arrayLen = len(vectorField.GetBfloat16Vector()) / bytesPerVector
}
case schemapb.DataType_Int8Vector:
if vectorField.GetInt8Vector() != nil {
arrayLen = len(vectorField.GetInt8Vector()) / int(dim)
}
}
}
@@ -1179,15 +1209,65 @@ func BuildArrayData(schema *schemapb.CollectionSchema, insertData *storage.Inser
vectorField := rowData.(*schemapb.VectorField)
listBuilder := fieldBuilder.(*array.ListBuilder)
listBuilder.Append(true)
dim, _ := typeutil.GetDim(subField)
if floatVectors := vectorField.GetFloatVector(); floatVectors != nil {
dim, _ := typeutil.GetDim(subField)
floatBuilder := listBuilder.ValueBuilder().(*array.Float32Builder)
start := j * int(dim)
end := start + int(dim)
if end <= len(floatVectors.Data) {
for k := start; k < end; k++ {
floatBuilder.Append(floatVectors.Data[k])
switch subField.ElementType {
case schemapb.DataType_FloatVector:
if floatVectors := vectorField.GetFloatVector(); floatVectors != nil {
floatBuilder := listBuilder.ValueBuilder().(*array.Float32Builder)
start := j * int(dim)
end := start + int(dim)
if end <= len(floatVectors.Data) {
for k := start; k < end; k++ {
floatBuilder.Append(floatVectors.Data[k])
}
}
}
case schemapb.DataType_BinaryVector:
if binaryVectors := vectorField.GetBinaryVector(); binaryVectors != nil {
uint8Builder := listBuilder.ValueBuilder().(*array.Uint8Builder)
bytesPerVector := int(dim) / 8
start := j * bytesPerVector
end := start + bytesPerVector
if end <= len(binaryVectors) {
for k := start; k < end; k++ {
uint8Builder.Append(binaryVectors[k])
}
}
}
case schemapb.DataType_Float16Vector:
if float16Vectors := vectorField.GetFloat16Vector(); float16Vectors != nil {
uint8Builder := listBuilder.ValueBuilder().(*array.Uint8Builder)
bytesPerVector := int(dim) * 2
start := j * bytesPerVector
end := start + bytesPerVector
if end <= len(float16Vectors) {
for k := start; k < end; k++ {
uint8Builder.Append(float16Vectors[k])
}
}
}
case schemapb.DataType_BFloat16Vector:
if bfloat16Vectors := vectorField.GetBfloat16Vector(); bfloat16Vectors != nil {
uint8Builder := listBuilder.ValueBuilder().(*array.Uint8Builder)
bytesPerVector := int(dim) * 2
start := j * bytesPerVector
end := start + bytesPerVector
if end <= len(bfloat16Vectors) {
for k := start; k < end; k++ {
uint8Builder.Append(bfloat16Vectors[k])
}
}
}
case schemapb.DataType_Int8Vector:
if int8Vectors := vectorField.GetInt8Vector(); int8Vectors != nil {
int8Builder := listBuilder.ValueBuilder().(*array.Int8Builder)
start := j * int(dim)
end := start + int(dim)
if end <= len(int8Vectors) {
for k := start; k < end; k++ {
int8Builder.Append(int8(int8Vectors[k]))
}
}
}
}
@@ -1372,8 +1452,7 @@ func reconstructStructArrayForJSON(structField *schemapb.StructArrayFieldSchema,
startIdx := j * bytesPerVector
endIdx := startIdx + bytesPerVector
if endIdx <= len(data) {
// Convert Float16 bytes to float32 for JSON representation
structElem[subField.GetName()] = typeutil.Float16BytesToFloat32Vector(data[startIdx:endIdx])
structElem[subField.GetName()] = data[startIdx:endIdx]
}
}
}
@@ -1385,8 +1464,7 @@ func reconstructStructArrayForJSON(structField *schemapb.StructArrayFieldSchema,
startIdx := j * bytesPerVector
endIdx := startIdx + bytesPerVector
if endIdx <= len(data) {
// Convert BFloat16 bytes to float32 for JSON representation
structElem[subField.GetName()] = typeutil.BFloat16BytesToFloat32Vector(data[startIdx:endIdx])
structElem[subField.GetName()] = data[startIdx:endIdx]
}
}
}
pkg/util/testutils/gen_data.go
+144
View File
@@ -294,6 +294,58 @@ func GenerateArrayOfFloatVectorArray(numRows int, dim int) []*schemapb.VectorFie
return ret
}
func GenerateArrayOfFloat16VectorArray(numRows int, dim int) []*schemapb.VectorField {
ret := make([]*schemapb.VectorField, 0, numRows)
for i := 0; i < numRows; i++ {
ret = append(ret, &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_Float16Vector{
Float16Vector: GenerateFloat16Vectors(ElemCountOfArray, dim),
},
})
}
return ret
}
func GenerateArrayOfBFloat16VectorArray(numRows int, dim int) []*schemapb.VectorField {
ret := make([]*schemapb.VectorField, 0, numRows)
for i := 0; i < numRows; i++ {
ret = append(ret, &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_Bfloat16Vector{
Bfloat16Vector: GenerateBFloat16Vectors(ElemCountOfArray, dim),
},
})
}
return ret
}
func GenerateArrayOfInt8VectorArray(numRows int, dim int) []*schemapb.VectorField {
ret := make([]*schemapb.VectorField, 0, numRows)
for i := 0; i < numRows; i++ {
ret = append(ret, &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_Int8Vector{
Int8Vector: typeutil.Int8ArrayToBytes(GenerateInt8Vectors(ElemCountOfArray, dim)),
},
})
}
return ret
}
func GenerateArrayOfBinaryVectorArray(numRows int, dim int) []*schemapb.VectorField {
ret := make([]*schemapb.VectorField, 0, numRows)
for i := 0; i < numRows; i++ {
ret = append(ret, &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_BinaryVector{
BinaryVector: GenerateBinaryVectors(ElemCountOfArray, dim),
},
})
}
return ret
}
func GenerateArrayOfStringArray(numRows int) []*schemapb.ScalarField {
ret := make([]*schemapb.ScalarField, 0, numRows)
for i := 0; i < numRows; i++ {
@@ -817,6 +869,82 @@ func NewVectorArrayFieldData(fieldName string, numRows, dim int) *schemapb.Field
}
}
func NewFloat16VectorArrayFieldData(fieldName string, numRows, dim int) *schemapb.FieldData {
return &schemapb.FieldData{
Type: schemapb.DataType_ArrayOfVector,
FieldName: fieldName,
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_VectorArray{
VectorArray: &schemapb.VectorArray{
Data: GenerateArrayOfFloat16VectorArray(numRows, dim),
ElementType: schemapb.DataType_Float16Vector,
Dim: int64(dim),
},
},
},
},
}
}
func NewBFloat16VectorArrayFieldData(fieldName string, numRows, dim int) *schemapb.FieldData {
return &schemapb.FieldData{
Type: schemapb.DataType_ArrayOfVector,
FieldName: fieldName,
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_VectorArray{
VectorArray: &schemapb.VectorArray{
Data: GenerateArrayOfBFloat16VectorArray(numRows, dim),
ElementType: schemapb.DataType_BFloat16Vector,
Dim: int64(dim),
},
},
},
},
}
}
func NewInt8VectorArrayFieldData(fieldName string, numRows, dim int) *schemapb.FieldData {
return &schemapb.FieldData{
Type: schemapb.DataType_ArrayOfVector,
FieldName: fieldName,
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_VectorArray{
VectorArray: &schemapb.VectorArray{
Data: GenerateArrayOfInt8VectorArray(numRows, dim),
ElementType: schemapb.DataType_Int8Vector,
Dim: int64(dim),
},
},
},
},
}
}
func NewBinaryVectorArrayFieldData(fieldName string, numRows, dim int) *schemapb.FieldData {
return &schemapb.FieldData{
Type: schemapb.DataType_ArrayOfVector,
FieldName: fieldName,
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_VectorArray{
VectorArray: &schemapb.VectorArray{
Data: GenerateArrayOfBinaryVectorArray(numRows, dim),
ElementType: schemapb.DataType_BinaryVector,
Dim: int64(dim),
},
},
},
},
}
}
func NewArrayFieldDataWithValue(fieldName string, fieldValue interface{}) *schemapb.FieldData {
return &schemapb.FieldData{
Type: schemapb.DataType_Array,
@@ -1181,6 +1309,22 @@ func GenerateArrayOfStructArray(schema *schemapb.StructArrayFieldSchema, numRows
fieldData := NewVectorArrayFieldData(field.Name, numRows, dim)
fieldData.FieldId = field.FieldID
ret = append(ret, fieldData)
case schemapb.DataType_Float16Vector:
fieldData := NewFloat16VectorArrayFieldData(field.Name, numRows, dim)
fieldData.FieldId = field.FieldID
ret = append(ret, fieldData)
case schemapb.DataType_BFloat16Vector:
fieldData := NewBFloat16VectorArrayFieldData(field.Name, numRows, dim)
fieldData.FieldId = field.FieldID
ret = append(ret, fieldData)
case schemapb.DataType_Int8Vector:
fieldData := NewInt8VectorArrayFieldData(field.Name, numRows, dim)
fieldData.FieldId = field.FieldID
ret = append(ret, fieldData)
case schemapb.DataType_BinaryVector:
fieldData := NewBinaryVectorArrayFieldData(field.Name, numRows, dim)
fieldData.FieldId = field.FieldID
ret = append(ret, fieldData)
default:
panic(fmt.Sprintf("unimplemented data type: %s", field.ElementType))
}
tests/integration/getvector/array_struct_test.go
+36
View File
@@ -246,6 +246,42 @@ func (s *TestArrayStructSuite) TestGetVector_ArrayStruct_FloatVector() {
s.run()
}
func (s *TestArrayStructSuite) TestGetVector_ArrayStruct_Float16Vector() {
s.nq = 10
s.topK = 10
s.indexType = integration.IndexHNSW
s.metricType = metric.MaxSim
s.vecType = schemapb.DataType_Float16Vector
s.run()
}
func (s *TestArrayStructSuite) TestGetVector_ArrayStruct_BFloat16Vector() {
s.nq = 10
s.topK = 10
s.indexType = integration.IndexHNSW
s.metricType = metric.MaxSim
s.vecType = schemapb.DataType_BFloat16Vector
s.run()
}
func (s *TestArrayStructSuite) TestGetVector_ArrayStruct_Int8Vector() {
s.nq = 10
s.topK = 10
s.indexType = integration.IndexHNSW
s.metricType = metric.MaxSim
s.vecType = schemapb.DataType_Int8Vector
s.run()
}
func (s *TestArrayStructSuite) TestGetVector_ArrayStruct_BinaryVector() {
s.nq = 10
s.topK = 10
s.indexType = integration.IndexHNSW
s.metricType = metric.MaxSimHamming
s.vecType = schemapb.DataType_BinaryVector
s.run()
}
func TestGetVectorArrayStruct(t *testing.T) {
// t.Skip("Skip integration test, need to refactor integration test framework.")
suite.Run(t, new(TestArrayStructSuite))
tests/integration/import/vector_array_test.go
+21 -6
View File
@@ -158,7 +158,7 @@ func (s *BulkInsertSuite) runForStructArray() {
Name: "vector_array_field",
IsPrimaryKey: false,
DataType: schemapb.DataType_ArrayOfVector,
ElementType: schemapb.DataType_FloatVector,
ElementType: s.vecType,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
@@ -302,11 +302,26 @@ func (s *BulkInsertSuite) runForStructArray() {
func (s *BulkInsertSuite) TestImportWithVectorArray() {
fileTypeArr := []importutilv2.FileType{importutilv2.CSV, importutilv2.JSON, importutilv2.Parquet}
vectorTypeConfigs := []struct {
vecType schemapb.DataType
indexType string
metricType string
}{
{schemapb.DataType_FloatVector, integration.IndexHNSW, metric.MaxSim},
{schemapb.DataType_Float16Vector, integration.IndexHNSW, metric.MaxSim},
{schemapb.DataType_BFloat16Vector, integration.IndexHNSW, metric.MaxSim},
{schemapb.DataType_Int8Vector, integration.IndexHNSW, metric.MaxSim},
{schemapb.DataType_BinaryVector, integration.IndexHNSW, metric.MaxSimHamming},
}
for _, fileType := range fileTypeArr {
s.fileType = fileType
s.vecType = schemapb.DataType_FloatVector
s.indexType = integration.IndexHNSW
s.metricType = metric.MaxSim
s.runForStructArray()
for _, vtConfig := range vectorTypeConfigs {
s.fileType = fileType
s.vecType = vtConfig.vecType
s.indexType = vtConfig.indexType
s.metricType = vtConfig.metricType
s.runForStructArray()
}
}
}
tests/integration/util_query.go
+41 -10
View File
@@ -342,9 +342,16 @@ func constructPlaceholderGroup(nq, dim int, vectorType schemapb.DataType, isEmbe
values = append(values, bs)
}
case schemapb.DataType_BinaryVector:
placeholderType = commonpb.PlaceholderType_BinaryVector
if !isEmbeddingList {
placeholderType = commonpb.PlaceholderType_BinaryVector
} else {
placeholderType = commonpb.PlaceholderType_EmbListBinaryVector
}
for i := 0; i < nq; i++ {
total := dim / 8
if isEmbeddingList {
total = total * (rand.Intn(10) + 3)
}
ret := make([]byte, total)
_, err := rand.Read(ret)
if err != nil {
@@ -353,17 +360,33 @@ func constructPlaceholderGroup(nq, dim int, vectorType schemapb.DataType, isEmbe
values = append(values, ret)
}
case schemapb.DataType_Float16Vector:
placeholderType = commonpb.PlaceholderType_Float16Vector
data := testutils.GenerateFloat16Vectors(nq, dim)
if !isEmbeddingList {
placeholderType = commonpb.PlaceholderType_Float16Vector
} else {
placeholderType = commonpb.PlaceholderType_EmbListFloat16Vector
}
vecCount := dim
if isEmbeddingList {
vecCount = vecCount * (rand.Intn(10) + 3)
}
data := testutils.GenerateFloat16Vectors(nq, vecCount)
for i := 0; i < nq; i++ {
rowBytes := dim * 2
rowBytes := vecCount * 2
values = append(values, data[rowBytes*i:rowBytes*(i+1)])
}
case schemapb.DataType_BFloat16Vector:
placeholderType = commonpb.PlaceholderType_BFloat16Vector
data := testutils.GenerateBFloat16Vectors(nq, dim)
if !isEmbeddingList {
placeholderType = commonpb.PlaceholderType_BFloat16Vector
} else {
placeholderType = commonpb.PlaceholderType_EmbListBFloat16Vector
}
vecCount := dim
if isEmbeddingList {
vecCount = vecCount * (rand.Intn(10) + 3)
}
data := testutils.GenerateBFloat16Vectors(nq, vecCount)
for i := 0; i < nq; i++ {
rowBytes := dim * 2
rowBytes := vecCount * 2
values = append(values, data[rowBytes*i:rowBytes*(i+1)])
}
case schemapb.DataType_SparseFloatVector:
@@ -373,10 +396,18 @@ func constructPlaceholderGroup(nq, dim int, vectorType schemapb.DataType, isEmbe
sparseVecs := GenerateSparseFloatArray(nq)
values = append(values, sparseVecs.Contents...)
case schemapb.DataType_Int8Vector:
placeholderType = commonpb.PlaceholderType_Int8Vector
data := testutils.GenerateInt8Vectors(nq, dim)
if !isEmbeddingList {
placeholderType = commonpb.PlaceholderType_Int8Vector
} else {
placeholderType = commonpb.PlaceholderType_EmbListInt8Vector
}
vecCount := dim
if isEmbeddingList {
vecCount = vecCount * (rand.Intn(10) + 3)
}
data := testutils.GenerateInt8Vectors(nq, vecCount)
for i := 0; i < nq; i++ {
rowBytes := dim
rowBytes := vecCount
values = append(values, typeutil.Int8ArrayToBytes(data[rowBytes*i:rowBytes*(i+1)]))
}
default:
tests/python_client/milvus_client/test_milvus_client_struct_array.py
+14 -35
View File
@@ -3383,9 +3383,9 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base):
@pytest.mark.tags(CaseLabel.L2)
def test_struct_with_unsupported_vector_field(self):
"""
target: test creating struct with BinaryVector field (should fail)
method: attempt to create struct with BinaryVector field
expected: creation should fail
target: test creating struct with SparseFloatVector field (should fail)
method: attempt to create struct with SparseFloatVector field
expected: creation should fail (sparse vectors not supported in struct)
"""
collection_name = cf.gen_unique_str(f"{prefix}_invalid")
@@ -3399,7 +3399,7 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base):
)
struct_schema = client.create_struct_field_schema()
struct_schema.add_field("binary_vector_field", DataType.BINARY_VECTOR, dim=default_dim)
struct_schema.add_field("sparse_vector_field", DataType.SPARSE_FLOAT_VECTOR)
schema.add_field(
"struct_array",
datatype=DataType.ARRAY,
@@ -3409,7 +3409,7 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base):
)
error = {
ct.err_code: 65535,
ct.err_msg: "now only float vector is supported",
ct.err_msg: "only fixed dimension vector types are supported",
}
self.create_collection(
client,
@@ -3672,22 +3672,12 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base):
)
@pytest.mark.tags(CaseLabel.L2)
@pytest.mark.parametrize(
"vector_type",
[
DataType.BINARY_VECTOR,
DataType.FLOAT16_VECTOR,
DataType.BFLOAT16_VECTOR,
DataType.SPARSE_FLOAT_VECTOR,
DataType.INT8_VECTOR,
],
)
def test_struct_array_with_unsupported_vector_types(self, vector_type):
def test_struct_array_with_unsupported_vector_types(self):
"""
target: test creating struct array with unsupported vector types (non-FLOAT_VECTOR)
method: attempt to create struct array with BINARY_VECTOR, FLOAT16_VECTOR,
BFLOAT16_VECTOR, SPARSE_FLOAT_VECTOR, INT8_VECTOR vector types
expected: creation should fail as only FLOAT_VECTOR is supported in struct array
target: test creating struct array with unsupported vector types
method: attempt to create struct array with SPARSE_FLOAT_VECTOR
expected: creation should fail as only fixed dimension vector types are supported
note: FLOAT_VECTOR, FLOAT16_VECTOR, BFLOAT16_VECTOR, BINARY_VECTOR, INT8_VECTOR are supported
"""
collection_name = cf.gen_unique_str(f"{prefix}_invalid")
@@ -3700,20 +3690,9 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base):
field_name="normal_vector", datatype=DataType.FLOAT_VECTOR, dim=default_dim
)
# Try to create struct with unsupported vector type
# Try to create struct with unsupported vector type (sparse vector)
struct_schema = client.create_struct_field_schema()
# SPARSE_FLOAT_VECTOR doesn't need dim parameter
if vector_type == DataType.SPARSE_FLOAT_VECTOR:
struct_schema.add_field("unsupported_vector", vector_type)
else:
# BINARY_VECTOR needs dim to be multiple of 8
if vector_type == DataType.BINARY_VECTOR:
struct_schema.add_field("unsupported_vector", vector_type, dim=128)
else:
struct_schema.add_field(
"unsupported_vector", vector_type, dim=default_dim
)
struct_schema.add_field("unsupported_vector", DataType.SPARSE_FLOAT_VECTOR)
schema.add_field(
"struct_array",
@@ -3723,8 +3702,8 @@ class TestMilvusClientStructArrayInvalid(TestMilvusClientV2Base):
max_capacity=100,
)
# Should fail - only FLOAT_VECTOR is supported in struct array
error = {ct.err_code: 65535, ct.err_msg: "now only float vector is supported"}
# Should fail - sparse vectors are not supported in struct array
error = {ct.err_code: 65535, ct.err_msg: "only fixed dimension vector types are supported"}
self.create_collection(
client,
collection_name,