From b39de6391b727e6745d7433ba348236508f1c0a2 Mon Sep 17 00:00:00 2001 From: Spade A <71589810+SpadeA-Tang@users.noreply.github.com> Date: Thu, 15 Jan 2026 10:13:27 +0800 Subject: [PATCH] feat: impl StructArray -- add more tests for non-float embedding lists (#46686) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit issue: https://github.com/milvus-io/milvus/issues/42148 - 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. --------- Signed-off-by: SpadeA --- internal/core/src/common/ChunkWriterTest.cpp | 197 ++++--- internal/core/src/query/Plan.cpp | 41 +- internal/core/src/segcore/Utils.cpp | 24 + .../core/unittest/test_element_filter.cpp | 496 ++++++++++++++---- internal/core/unittest/test_utils/DataGen.h | 153 ++++++ internal/proxy/util.go | 9 +- internal/proxy/util_test.go | 2 +- internal/storage/payload_writer_test.go | 51 ++ internal/storage/serde.go | 37 +- internal/storage/serde_test.go | 374 +++++++++++-- .../util/importutilv2/parquet/reader_test.go | 196 +++---- .../parquet/struct_field_reader.go | 74 ++- internal/util/importutilv2/parquet/util.go | 2 +- internal/util/testutil/test_util.go | 110 +++- pkg/util/testutils/gen_data.go | 144 +++++ .../getvector/array_struct_test.go | 36 ++ tests/integration/import/vector_array_test.go | 27 +- tests/integration/util_query.go | 51 +- .../test_milvus_client_struct_array.py | 49 +- 19 files changed, 1665 insertions(+), 408 deletions(-) diff --git a/internal/core/src/common/ChunkWriterTest.cpp b/internal/core/src/common/ChunkWriterTest.cpp index 8b15837464..3542f0b0f7 100644 --- a/internal/core/src/common/ChunkWriterTest.cpp +++ b/internal/core/src/common/ChunkWriterTest.cpp @@ -18,6 +18,7 @@ #include #include #include +#include #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 -BuildVectorArrayListArray(const std::vector& 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& 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& vectors_per_row, int dim) { auto& fsb_builder = dynamic_cast( *list_builder.value_builder()); - float counter = 0.0f; + std::default_random_engine gen(42); + std::uniform_int_distribution 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 vector_data(dim); - for (int d = 0; d < dim; ++d) { - vector_data[d] = counter++; + std::vector vector_data(byte_width); + for (int d = 0; d < byte_width; ++d) { + vector_data[d] = static_cast(dist(gen)); } - EXPECT_TRUE(fsb_builder - .Append(reinterpret_cast( - vector_data.data())) - .ok()); + EXPECT_TRUE(fsb_builder.Append(vector_data.data()).ok()); } } @@ -74,27 +95,52 @@ BuildVectorArrayListArray(const std::vector& vectors_per_row, int dim) { return std::static_pointer_cast(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 { + 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 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(dim, - row_count, - data, - calculated_size, - milvus::DataType::VECTOR_FLOAT, - nullptr); + auto chunk = std::make_unique( + 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 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(dim, - row_count, - data, - calculated_size, - milvus::DataType::VECTOR_FLOAT, - nullptr); + auto chunk = std::make_unique( + 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 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 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(dim, - row_count, - data, - calculated_size, - milvus::DataType::VECTOR_FLOAT, - nullptr); + auto chunk = std::make_unique( + 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 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 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 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> 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( 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(dim, - row_count, - data, - calculated_size, - milvus::DataType::VECTOR_FLOAT, - nullptr); + auto chunk = std::make_unique( + 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& info) { + return info.param.test_name; + }); diff --git a/internal/core/src/query/Plan.cpp b/internal/core/src/query/Plan.cpp index 81f133c269..6e3c15222f 100644 --- a/internal/core/src/query/Plan.cpp +++ b/internal/core/src/query/Plan.cpp @@ -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; } diff --git a/internal/core/src/segcore/Utils.cpp b/internal/core/src/segcore/Utils.cpp index f9a93d64de..ff7781b6a1 100644 --- a/internal/core/src/segcore/Utils.cpp +++ b/internal/core/src/segcore/Utils.cpp @@ -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 {}", diff --git a/internal/core/unittest/test_element_filter.cpp b/internal/core/unittest/test_element_filter.cpp index 6585a82e86..4e67bac348 100644 --- a/internal/core/unittest/test_element_filter.cpp +++ b/internal/core/unittest/test_element_filter.cpp @@ -25,8 +25,12 @@ using namespace milvus; using namespace milvus::query; using namespace milvus::segcore; +// Test parameter: +using ElementFilterSealedParam = + std::tuple; + class ElementFilterSealed - : public ::testing::TestWithParam> { + : public ::testing::TestWithParam { 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( + num_queries, dim, seed, true); + } else if (element_type() == DataType::VECTOR_FLOAT16) { + return CreatePlaceholderGroup( + num_queries, dim, seed, true); + } else if (element_type() == DataType::VECTOR_BFLOAT16) { + return CreatePlaceholderGroup( + num_queries, dim, seed, true); + } else if (element_type() == DataType::VECTOR_INT8) { + return CreatePlaceholderGroup( + num_queries, dim, seed, true); + } else { + // VECTOR_FLOAT + return CreatePlaceholderGroup( + 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(); - 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(vec_fid); - // DataGen generates VECTOR_ARRAY with data in float_vector (flattened), - // not in vector_array (nested structure) - std::vector 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 indexing; + std::string actual_metric; + + if (elem_type == DataType::VECTOR_FLOAT) { + std::vector 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 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(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 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(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 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(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 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(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(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(); - 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(vec_fid); - // DataGen generates VECTOR_ARRAY with data in float_vector (flattened), - // not in vector_array (nested structure) - std::vector 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 indexing; + std::string actual_metric; + + if (elem_type == DataType::VECTOR_FLOAT) { + std::vector 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 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(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 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(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 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(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 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(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(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(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& 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& 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 { +// Test parameter for Growing: +using ElementFilterGrowingParam = std::tuple; + +class ElementFilterGrowing + : public ::testing::TestWithParam { 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( + num_queries, dim, seed, true); + } else if (element_type() == DataType::VECTOR_FLOAT16) { + return CreatePlaceholderGroup( + num_queries, dim, seed, true); + } else if (element_type() == DataType::VECTOR_BFLOAT16) { + return CreatePlaceholderGroup( + num_queries, dim, seed, true); + } else if (element_type() == DataType::VECTOR_INT8) { + return CreatePlaceholderGroup( + num_queries, dim, seed, true); + } else { + // VECTOR_FLOAT + return CreatePlaceholderGroup( + 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(); - 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(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(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& 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& 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 diff --git a/internal/core/unittest/test_utils/DataGen.h b/internal/core/unittest/test_utils/DataGen.h index 018ebaa8ed..35bf599a90 100644 --- a/internal/core/unittest/test_utils/DataGen.h +++ b/internal/core/unittest/test_utils/DataGen.h @@ -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 +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>( + 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 +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>( + 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 +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>( + 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 +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>( + 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 inline index::IndexBasePtr GenScalarIndexing(int64_t N, const T* data) { diff --git a/internal/proxy/util.go b/internal/proxy/util.go index 41cacc1313..b585341012 100644 --- a/internal/proxy/util.go +++ b/internal/proxy/util.go @@ -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 diff --git a/internal/proxy/util_test.go b/internal/proxy/util_test.go index 7b8cb110e5..4db6a77f9a 100644 --- a/internal/proxy/util_test.go +++ b/internal/proxy/util_test.go @@ -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) { diff --git a/internal/storage/payload_writer_test.go b/internal/storage/payload_writer_test.go index f26cb3fa78..ebe3def93c 100644 --- a/internal/storage/payload_writer_test.go +++ b/internal/storage/payload_writer_test.go @@ -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) { diff --git a/internal/storage/serde.go b/internal/storage/serde.go index f11d9f811a..61d3d7a134 100644 --- a/internal/storage/serde.go +++ b/internal/storage/serde.go @@ -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) diff --git a/internal/storage/serde_test.go b/internal/storage/serde_test.go index a6973e11a8..bdd4a68cad 100644 --- a/internal/storage/serde_test.go +++ b/internal/storage/serde_test.go @@ -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) { diff --git a/internal/util/importutilv2/parquet/reader_test.go b/internal/util/importutilv2/parquet/reader_test.go index 670ebe6721..0f087b8344 100644 --- a/internal/util/importutilv2/parquet/reader_test.go +++ b/internal/util/importutilv2/parquet/reader_test.go @@ -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) { diff --git a/internal/util/importutilv2/parquet/struct_field_reader.go b/internal/util/importutilv2/parquet/struct_field_reader.go index ca607bebf9..6210e7c3f4 100644 --- a/internal/util/importutilv2/parquet/struct_field_reader.go +++ b/internal/util/importutilv2/parquet/struct_field_reader.go @@ -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") diff --git a/internal/util/importutilv2/parquet/util.go b/internal/util/importutilv2/parquet/util.go index 5dae2cdd6e..5378afc70d 100644 --- a/internal/util/importutilv2/parquet/util.go +++ b/internal/util/importutilv2/parquet/util.go @@ -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: diff --git a/internal/util/testutil/test_util.go b/internal/util/testutil/test_util.go index 314f0248ae..6dbefa2e7c 100644 --- a/internal/util/testutil/test_util.go +++ b/internal/util/testutil/test_util.go @@ -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] } } } diff --git a/pkg/util/testutils/gen_data.go b/pkg/util/testutils/gen_data.go index 222bfd63cf..2e66e2a8f6 100644 --- a/pkg/util/testutils/gen_data.go +++ b/pkg/util/testutils/gen_data.go @@ -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)) } diff --git a/tests/integration/getvector/array_struct_test.go b/tests/integration/getvector/array_struct_test.go index 751a37282d..e23a2b1aee 100644 --- a/tests/integration/getvector/array_struct_test.go +++ b/tests/integration/getvector/array_struct_test.go @@ -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)) diff --git a/tests/integration/import/vector_array_test.go b/tests/integration/import/vector_array_test.go index c987af8271..9e016bbdce 100644 --- a/tests/integration/import/vector_array_test.go +++ b/tests/integration/import/vector_array_test.go @@ -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() + } } } diff --git a/tests/integration/util_query.go b/tests/integration/util_query.go index 82527c6c66..cc50df6fa0 100644 --- a/tests/integration/util_query.go +++ b/tests/integration/util_query.go @@ -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: diff --git a/tests/python_client/milvus_client/test_milvus_client_struct_array.py b/tests/python_client/milvus_client/test_milvus_client_struct_array.py index c8d19cdfe2..eb6b4fca72 100644 --- a/tests/python_client/milvus_client/test_milvus_client_struct_array.py +++ b/tests/python_client/milvus_client/test_milvus_client_struct_array.py @@ -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,