项目文件夹

文件
wehub-resource-sync 498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

3426 行
130 KiB
Go

此文件含有模棱两可的 Unicode 字符
此文件含有可能会与其他字符混淆的 Unicode 字符。 如果您是想特意这样的,可以安全地忽略该警告。 使用 Escape 按钮显示他们。
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package queryutil
import (
"context"
"encoding/binary"
"fmt"
"math"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
)
// makeNullableSchema builds a minimal CollectionSchema marking the given
// fieldID as nullable with the specified DataType. For dense vectors, pass
// dim > 0; for sparse or scalars, pass dim = 0.
func makeNullableSchema(fieldID int64, dataType schemapb.DataType, dim int64) *schemapb.CollectionSchema {
fs := &schemapb.FieldSchema{
FieldID: fieldID,
DataType: dataType,
Nullable: true,
}
if dim > 0 {
fs.TypeParams = []*commonpb.KeyValuePair{
{Key: "dim", Value: fmt.Sprintf("%d", dim)},
}
}
return &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{fs}}
}
func requireRangeSliceRetrieveResults(t *testing.T, result *internalpb.RetrieveResults, start, end int) *internalpb.RetrieveResults {
t.Helper()
sliced, err := rangeSliceRetrieveResults(result, start, end)
require.NoError(t, err)
return sliced
}
func requireRangeSliceFieldData(t *testing.T, fd *schemapb.FieldData, start, end int) *schemapb.FieldData {
t.Helper()
sliced, err := rangeSliceFieldData(fd, start, end)
require.NoError(t, err)
return sliced
}
func requireRangeSliceVectorField(t *testing.T, vf *schemapb.VectorField, start, end int, validData []bool) *schemapb.VectorField {
t.Helper()
sliced, err := rangeSliceVectorField(vf, start, end, validData)
require.NoError(t, err)
return sliced
}
func requireSliceFieldData(t *testing.T, fd *schemapb.FieldData, indices []int) *schemapb.FieldData {
t.Helper()
sliced, err := sliceFieldData(fd, indices)
require.NoError(t, err)
return sliced
}
func requireSliceVectorField(t *testing.T, vf *schemapb.VectorField, indices []int, validData []bool) *schemapb.VectorField {
t.Helper()
sliced, err := sliceVectorField(vf, indices, validData)
require.NoError(t, err)
return sliced
}
func TestGetRowCountNullableCompactVectorWithoutIDsUsesLogicalRows(t *testing.T) {
result := &internalpb.RetrieveResults{
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_FloatVector,
FieldName: "nullable_vec",
FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_FloatVector{
FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4}},
},
}},
ValidData: []bool{true, false, true},
},
},
}
assert.Equal(t, 3, getRowCount(result))
}
// =========================================================================
// buildMergedFieldData: ValidData preservation
// =========================================================================
func TestBuildMergedFieldData_ValidData(t *testing.T) {
// Two results with nullable int64 fields
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Int64, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{10, 20}}},
}},
ValidData: []bool{true, false}, // row 0 valid, row 1 null
},
},
}
r2 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{3}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Int64, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{30}}},
}},
ValidData: []bool{true},
},
},
}
results := []*internalpb.RetrieveResults{r1, r2}
selectedRows := []rowRef{
{resultIdx: 0, rowIdx: 0}, // r1 row 0 (valid)
{resultIdx: 0, rowIdx: 1}, // r1 row 1 (null)
{resultIdx: 1, rowIdx: 0}, // r2 row 0 (valid)
}
merged, err := buildMergedRetrieveResults(results, selectedRows, makeNullableSchema(100, schemapb.DataType_Int64, 0))
require.NoError(t, err)
require.Len(t, merged.FieldsData, 1)
fd := merged.FieldsData[0]
assert.Equal(t, []bool{true, false, true}, fd.ValidData)
assert.Equal(t, []int64{10, 20, 30}, fd.GetScalars().GetLongData().GetData())
}
func TestBuildMergedFieldData_ValidData_MixedNullable(t *testing.T) {
// r1 has ValidData (nullable), r2 has explicit ValidData=[true] (all valid)
// This tests that explicit all-valid ValidData is preserved correctly.
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Int64, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{10}}},
}},
ValidData: []bool{false}, // null
},
},
}
r2 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{2}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Int64, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{20}}},
}},
ValidData: []bool{true}, // explicit valid (segcore always includes ValidData for nullable fields)
},
},
}
results := []*internalpb.RetrieveResults{r1, r2}
selectedRows := []rowRef{
{resultIdx: 0, rowIdx: 0}, // r1 (null)
{resultIdx: 1, rowIdx: 0}, // r2 (valid=true)
}
merged, err := buildMergedRetrieveResults(results, selectedRows, makeNullableSchema(100, schemapb.DataType_Int64, 0))
require.NoError(t, err)
fd := merged.FieldsData[0]
assert.Equal(t, []bool{false, true}, fd.ValidData)
}
// =========================================================================
// buildMergedScalarField: ArrayData ElementType preservation
// =========================================================================
func TestBuildMergedScalarField_ArrayElementType(t *testing.T) {
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Array, FieldId: 200,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_ArrayData{ArrayData: &schemapb.ArrayArray{
Data: []*schemapb.ScalarField{
{Data: &schemapb.ScalarField_FloatData{FloatData: &schemapb.FloatArray{Data: []float32{1.1, 2.2}}}},
{Data: &schemapb.ScalarField_FloatData{FloatData: &schemapb.FloatArray{Data: []float32{3.3}}}},
},
ElementType: schemapb.DataType_Float,
}},
}},
},
},
}
results := []*internalpb.RetrieveResults{r1}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}, {resultIdx: 0, rowIdx: 1}}
merged, err := buildMergedRetrieveResults(results, selectedRows, nil)
require.NoError(t, err)
arrayField := merged.FieldsData[0].GetScalars().GetArrayData()
assert.Equal(t, schemapb.DataType_Float, arrayField.GetElementType())
assert.Len(t, arrayField.GetData(), 2)
}
// =========================================================================
// buildMergedScalarField: Geometry types
// =========================================================================
func TestBuildMergedScalarField_GeometryWktData(t *testing.T) {
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Geometry, FieldId: 300,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_GeometryWktData{GeometryWktData: &schemapb.GeometryWktArray{
Data: []string{"POINT(0 0)", "POINT(1 1)"},
}},
}},
},
},
}
results := []*internalpb.RetrieveResults{r1}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 1}}
merged, err := buildMergedRetrieveResults(results, selectedRows, nil)
require.NoError(t, err)
wktData := merged.FieldsData[0].GetScalars().GetGeometryWktData().GetData()
assert.Equal(t, []string{"POINT(1 1)"}, wktData)
}
func TestBuildMergedScalarField_GeometryData(t *testing.T) {
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Geometry, FieldId: 300,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_GeometryData{GeometryData: &schemapb.GeometryArray{
Data: [][]byte{{0x01, 0x02}, {0x03, 0x04}},
}},
}},
},
},
}
results := []*internalpb.RetrieveResults{r1}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}}
merged, err := buildMergedRetrieveResults(results, selectedRows, nil)
require.NoError(t, err)
geoData := merged.FieldsData[0].GetScalars().GetGeometryData().GetData()
assert.Equal(t, [][]byte{{0x01, 0x02}}, geoData)
}
// =========================================================================
// buildMergedScalarField: Timestamptz and Mol types
// =========================================================================
func TestBuildMergedScalarField_TimestamptzData(t *testing.T) {
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Timestamptz, FieldId: 400,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_TimestamptzData{TimestamptzData: &schemapb.TimestamptzArray{
Data: []int64{1000, 2000},
}},
}},
},
},
}
results := []*internalpb.RetrieveResults{r1}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}, {resultIdx: 0, rowIdx: 1}}
merged, err := buildMergedRetrieveResults(results, selectedRows, nil)
require.NoError(t, err)
tsData := merged.FieldsData[0].GetScalars().GetTimestamptzData().GetData()
assert.Equal(t, []int64{1000, 2000}, tsData)
}
func TestBuildMergedScalarField_MolData(t *testing.T) {
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Mol, FieldId: 500,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_MolData{MolData: &schemapb.MolArray{
Data: [][]byte{{0xAA, 0xBB}},
}},
}},
},
},
}
results := []*internalpb.RetrieveResults{r1}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}}
merged, err := buildMergedRetrieveResults(results, selectedRows, nil)
require.NoError(t, err)
molData := merged.FieldsData[0].GetScalars().GetMolData().GetData()
assert.Equal(t, [][]byte{{0xAA, 0xBB}}, molData)
}
// =========================================================================
// buildMergedVectorField: Int8Vector
// =========================================================================
func TestBuildMergedVectorField_Int8Vector(t *testing.T) {
dim := 4
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Int8Vector, FieldId: 600,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_Int8Vector{Int8Vector: []byte{1, 2, 3, 4, 5, 6, 7, 8}},
}},
},
},
}
results := []*internalpb.RetrieveResults{r1}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 1}} // select row 1 only
merged, err := buildMergedRetrieveResults(results, selectedRows, nil)
require.NoError(t, err)
vecData := merged.FieldsData[0].GetVectors().GetInt8Vector()
assert.Equal(t, []byte{5, 6, 7, 8}, vecData)
}
// =========================================================================
// rangeSliceScalarField: new types
// =========================================================================
func TestRangeSliceScalarField_ArrayElementType(t *testing.T) {
fd := &schemapb.FieldData{
Type: schemapb.DataType_Array, FieldId: 200,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_ArrayData{ArrayData: &schemapb.ArrayArray{
Data: []*schemapb.ScalarField{
{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{1}}}},
{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{2}}}},
{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{3}}}},
},
ElementType: schemapb.DataType_Int32,
}},
}},
}
sliced := requireRangeSliceFieldData(t, fd, 1, 3)
arrayField := sliced.GetScalars().GetArrayData()
assert.Equal(t, schemapb.DataType_Int32, arrayField.GetElementType())
assert.Len(t, arrayField.GetData(), 2)
}
func TestRangeSliceScalarField_GeometryWkt(t *testing.T) {
fd := &schemapb.FieldData{
Type: schemapb.DataType_Geometry, FieldId: 300,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_GeometryWktData{GeometryWktData: &schemapb.GeometryWktArray{
Data: []string{"POINT(0 0)", "POINT(1 1)", "POINT(2 2)"},
}},
}},
}
sliced := requireRangeSliceFieldData(t, fd, 1, 3)
assert.Equal(t, []string{"POINT(1 1)", "POINT(2 2)"}, sliced.GetScalars().GetGeometryWktData().GetData())
}
func TestRangeSliceScalarField_Timestamptz(t *testing.T) {
fd := &schemapb.FieldData{
Type: schemapb.DataType_Timestamptz, FieldId: 400,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_TimestamptzData{TimestamptzData: &schemapb.TimestamptzArray{
Data: []int64{100, 200, 300},
}},
}},
}
sliced := requireRangeSliceFieldData(t, fd, 0, 2)
assert.Equal(t, []int64{100, 200}, sliced.GetScalars().GetTimestamptzData().GetData())
}
// =========================================================================
// sliceScalarField (index-based): new types
// =========================================================================
func TestSliceScalarField_ArrayElementType(t *testing.T) {
fd := &schemapb.FieldData{
Type: schemapb.DataType_Array, FieldId: 200,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_ArrayData{ArrayData: &schemapb.ArrayArray{
Data: []*schemapb.ScalarField{
{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{10}}}},
{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{20}}}},
{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{30}}}},
},
ElementType: schemapb.DataType_Int64,
}},
}},
}
sliced := requireSliceFieldData(t, fd, []int{2, 0})
arrayField := sliced.GetScalars().GetArrayData()
assert.Equal(t, schemapb.DataType_Int64, arrayField.GetElementType())
assert.Len(t, arrayField.GetData(), 2)
assert.Equal(t, []int64{30}, arrayField.GetData()[0].GetLongData().GetData())
assert.Equal(t, []int64{10}, arrayField.GetData()[1].GetLongData().GetData())
}
func TestSliceScalarField_GeometryWkt(t *testing.T) {
fd := &schemapb.FieldData{
Type: schemapb.DataType_Geometry, FieldId: 300,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_GeometryWktData{GeometryWktData: &schemapb.GeometryWktArray{
Data: []string{"POINT(0 0)", "POINT(1 1)", "POINT(2 2)"},
}},
}},
}
sliced := requireSliceFieldData(t, fd, []int{2, 0})
assert.Equal(t, []string{"POINT(2 2)", "POINT(0 0)"}, sliced.GetScalars().GetGeometryWktData().GetData())
}
// =========================================================================
// calcFieldElementSize: new types
// =========================================================================
func TestCalcFieldElementSize_Geometry(t *testing.T) {
fd := &schemapb.FieldData{
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_GeometryWktData{GeometryWktData: &schemapb.GeometryWktArray{
Data: []string{"POINT(0 0)", "LINESTRING(0 0, 1 1)"},
}},
}},
}
assert.Equal(t, int64(len("POINT(0 0)")), calcFieldElementSize(fd, 0))
assert.Equal(t, int64(len("LINESTRING(0 0, 1 1)")), calcFieldElementSize(fd, 1))
}
func TestCalcFieldElementSize_Timestamptz(t *testing.T) {
fd := &schemapb.FieldData{
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_TimestamptzData{TimestamptzData: &schemapb.TimestamptzArray{
Data: []int64{1000},
}},
}},
}
assert.Equal(t, int64(8), calcFieldElementSize(fd, 0))
}
func TestCalcFieldElementSize_Int8Vector(t *testing.T) {
fd := &schemapb.FieldData{
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 128,
Data: &schemapb.VectorField_Int8Vector{Int8Vector: make([]byte, 256)},
}},
}
assert.Equal(t, int64(128), calcFieldElementSize(fd, 0))
}
// =========================================================================
// ElementIndices propagation through merge
// =========================================================================
func TestBuildMergedRetrieveResults_ElementIndices(t *testing.T) {
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Int64, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{10, 20}}},
}},
},
},
ElementLevel: true,
ElementIndices: []*internalpb.ElementIndices{{Indices: []int32{0, 1}}, {Indices: []int32{2}}},
}
results := []*internalpb.RetrieveResults{r1}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 1}, {resultIdx: 0, rowIdx: 0}}
merged, err := buildMergedRetrieveResults(results, selectedRows, nil)
require.NoError(t, err)
assert.True(t, merged.GetElementLevel())
require.Len(t, merged.GetElementIndices(), 2)
assert.Equal(t, []int32{2}, merged.GetElementIndices()[0].GetIndices())
assert.Equal(t, []int32{0, 1}, merged.GetElementIndices()[1].GetIndices())
}
// =========================================================================
// rangeSliceRetrieveResults: ElementIndices propagation
// =========================================================================
func TestRangeSliceRetrieveResults_ElementIndices(t *testing.T) {
result := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2, 3}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Int64, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{10, 20, 30}}},
}},
},
},
ElementLevel: true,
ElementIndices: []*internalpb.ElementIndices{
{Indices: []int32{0, 1}},
{Indices: []int32{2}},
{Indices: []int32{0}},
},
}
sliced := requireRangeSliceRetrieveResults(t, result, 1, 3)
assert.True(t, sliced.GetElementLevel())
require.Len(t, sliced.GetElementIndices(), 2)
assert.Equal(t, []int32{2}, sliced.GetElementIndices()[0].GetIndices())
assert.Equal(t, []int32{0}, sliced.GetElementIndices()[1].GetIndices())
}
// =========================================================================
// ConcatAndCheckPKOperator
// =========================================================================
func TestConcatAndCheckPKOperator_NoDuplicate(t *testing.T) {
op := NewConcatAndCheckPKOperator(nil)
ctx := context.Background()
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Int64, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{10, 20}}},
}},
},
},
}
r2 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{3, 4}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Int64, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{30, 40}}},
}},
},
},
}
outputs, err := op.Run(ctx, nil, []*internalpb.RetrieveResults{r1, r2})
require.NoError(t, err)
result := outputs[0].(*internalpb.RetrieveResults)
ids := result.GetIds().GetIntId().GetData()
// Concatenated in order: 1,2,3,4
assert.Equal(t, []int64{1, 2, 3, 4}, ids)
}
func TestConcatAndCheckPKOperator_DuplicateFails(t *testing.T) {
op := NewConcatAndCheckPKOperator(nil)
ctx := context.Background()
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Int64, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{10, 20}}},
}},
},
},
}
r2 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{2, 3}}}}, // PK=2 duplicate
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Int64, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{21, 30}}},
}},
},
},
}
_, err := op.Run(ctx, nil, []*internalpb.RetrieveResults{r1, r2})
require.Error(t, err)
assert.Contains(t, err.Error(), "duplicate PK")
}
// =========================================================================
// buildMergedRetrieveResults: common scalar types (LongData, DoubleData, StringData)
// =========================================================================
func TestBuildMergedFieldData_CommonScalarTypes(t *testing.T) {
// Merge two results with Int64 + Double + String fields
makeDoubleField := func(id int64, name string, vals []float64) *schemapb.FieldData {
return &schemapb.FieldData{
FieldId: id, FieldName: name, Type: schemapb.DataType_Double,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_DoubleData{DoubleData: &schemapb.DoubleArray{Data: vals}},
}},
}
}
makeStringField := func(id int64, name string, vals []string) *schemapb.FieldData {
return &schemapb.FieldData{
FieldId: id, FieldName: name, Type: schemapb.DataType_VarChar,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: vals}},
}},
}
}
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{
{FieldId: 100, Type: schemapb.DataType_Int64, Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{10, 20}}},
}}},
makeDoubleField(101, "price", []float64{1.1, 2.2}),
makeStringField(102, "name", []string{"alice", "bob"}),
},
}
r2 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{3}}}},
FieldsData: []*schemapb.FieldData{
{FieldId: 100, Type: schemapb.DataType_Int64, Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{30}}},
}}},
makeDoubleField(101, "price", []float64{3.3}),
makeStringField(102, "name", []string{"charlie"}),
},
}
selectedRows := []rowRef{{0, 0}, {1, 0}, {0, 1}} // r1[0], r2[0], r1[1]
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1, r2}, selectedRows, nil)
require.NoError(t, err)
assert.Equal(t, []int64{1, 3, 2}, merged.GetIds().GetIntId().GetData())
assert.Equal(t, []int64{10, 30, 20}, merged.GetFieldsData()[0].GetScalars().GetLongData().GetData())
assert.Equal(t, []float64{1.1, 3.3, 2.2}, merged.GetFieldsData()[1].GetScalars().GetDoubleData().GetData())
assert.Equal(t, []string{"alice", "charlie", "bob"}, merged.GetFieldsData()[2].GetScalars().GetStringData().GetData())
}
// =========================================================================
// buildMergedRetrieveResults: FloatVector merge
// =========================================================================
func TestBuildMergedFieldData_FloatVector(t *testing.T) {
dim := 4
makeVecField := func(id int64, data []float32) *schemapb.FieldData {
return &schemapb.FieldData{
FieldId: id, Type: schemapb.DataType_FloatVector,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: data}},
}},
}
}
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{
makeVecField(100, []float32{1, 2, 3, 4, 5, 6, 7, 8}), // 2 rows × dim=4
},
}
r2 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{3}}}},
FieldsData: []*schemapb.FieldData{
makeVecField(100, []float32{9, 10, 11, 12}), // 1 row × dim=4
},
}
selectedRows := []rowRef{{1, 0}, {0, 1}} // r2[0], r1[1]
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1, r2}, selectedRows, nil)
require.NoError(t, err)
vecData := merged.GetFieldsData()[0].GetVectors().GetFloatVector().GetData()
assert.Equal(t, []float32{9, 10, 11, 12, 5, 6, 7, 8}, vecData)
}
// =========================================================================
// calcRowSize
// =========================================================================
func TestCalcRowSize_BasicTypes(t *testing.T) {
r := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}},
FieldsData: []*schemapb.FieldData{
// Int64: 8 bytes
{FieldId: 100, Type: schemapb.DataType_Int64, Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{42}}},
}}},
// Double: 8 bytes
{FieldId: 101, Type: schemapb.DataType_Double, Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_DoubleData{DoubleData: &schemapb.DoubleArray{Data: []float64{3.14}}},
}}},
// String: len("hello") = 5 bytes
{FieldId: 102, Type: schemapb.DataType_VarChar, Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"hello"}}},
}}},
},
}
size := calcRowSize(r, 0)
assert.Equal(t, int64(8+8+5), size) // Int64(8) + Double(8) + String("hello"=5)
}
func TestCalcRowSize_EmptyString(t *testing.T) {
r := &internalpb.RetrieveResults{
FieldsData: []*schemapb.FieldData{
{FieldId: 100, Type: schemapb.DataType_VarChar, Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{""}}},
}}},
},
}
size := calcRowSize(r, 0)
assert.Equal(t, int64(0), size)
}
func TestCalcRowSize_FloatVector(t *testing.T) {
dim := 8
r := &internalpb.RetrieveResults{
FieldsData: []*schemapb.FieldData{
{
FieldId: 100, Type: schemapb.DataType_FloatVector,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: make([]float32, dim*2)}}, // 2 rows
}},
},
},
}
// FloatVector: dim * 4 bytes per row = 8 * 4 = 32
size := calcRowSize(r, 0)
assert.Equal(t, int64(dim*4), size)
}
func TestCalcRowSize_NullableDenseVectorUsesLogicalValidity(t *testing.T) {
tests := []struct {
name string
field *schemapb.FieldData
wantSize int64
}{
{
name: "float_vector",
field: &schemapb.FieldData{
Type: schemapb.DataType_FloatVector,
ValidData: []bool{false, true},
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2}}},
}},
},
wantSize: 8,
},
{
name: "binary_vector",
field: &schemapb.FieldData{
Type: schemapb.DataType_BinaryVector,
ValidData: []bool{false, true},
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 16,
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{0x01, 0x02}},
}},
},
wantSize: 2,
},
{
name: "float16_vector",
field: &schemapb.FieldData{
Type: schemapb.DataType_Float16Vector,
ValidData: []bool{false, true},
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{0x01, 0x02, 0x03, 0x04}},
}},
},
wantSize: 4,
},
{
name: "bfloat16_vector",
field: &schemapb.FieldData{
Type: schemapb.DataType_BFloat16Vector,
ValidData: []bool{false, true},
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{0x01, 0x02, 0x03, 0x04}},
}},
},
wantSize: 4,
},
{
name: "int8_vector",
field: &schemapb.FieldData{
Type: schemapb.DataType_Int8Vector,
ValidData: []bool{false, true},
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 4,
Data: &schemapb.VectorField_Int8Vector{Int8Vector: []byte{0x01, 0x02, 0x03, 0x04}},
}},
},
wantSize: 4,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
r := &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{tt.field}}
assert.Equal(t, int64(0), calcRowSize(r, 0))
assert.Equal(t, tt.wantSize, calcRowSize(r, 1))
})
}
}
func TestCalcRowSize_NullableSparseVectorUsesCompactMapping(t *testing.T) {
rowA := makeTestSparseVec(1, 0.5)
rowC := makeTestSparseVec(3, 1.5)
tests := []struct {
name string
validData []bool
contents [][]byte
wantSizes []int64
}{
{
name: "null_before_valid",
validData: []bool{false, true},
contents: [][]byte{rowA},
wantSizes: []int64{0, int64(len(rowA))},
},
{
name: "valid_null_valid",
validData: []bool{true, false, true},
contents: [][]byte{rowA, rowC},
wantSizes: []int64{int64(len(rowA)), 0, int64(len(rowC))},
},
{
name: "all_null",
validData: []bool{false, false},
contents: nil,
wantSizes: []int64{0, 0},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
r := &internalpb.RetrieveResults{
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_SparseFloatVector,
ValidData: tt.validData,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 4,
Data: &schemapb.VectorField_SparseFloatVector{SparseFloatVector: &schemapb.SparseFloatArray{
Dim: 4,
Contents: tt.contents,
}},
}},
},
},
}
for rowIdx, want := range tt.wantSizes {
assert.Equal(t, want, calcRowSize(r, int64(rowIdx)))
}
})
}
}
func TestCalcRowSize_MultipleFieldTypes(t *testing.T) {
r := &internalpb.RetrieveResults{
FieldsData: []*schemapb.FieldData{
// Bool: 1 byte
{
FieldId: 100, Type: schemapb.DataType_Bool,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: []bool{true}}},
}},
},
// Float: 4 bytes
{
FieldId: 101, Type: schemapb.DataType_Float,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_FloatData{FloatData: &schemapb.FloatArray{Data: []float32{1.0}}},
}},
},
// Int32: 4 bytes
{
FieldId: 102, Type: schemapb.DataType_Int32,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{42}}},
}},
},
},
}
size := calcRowSize(r, 0)
assert.Equal(t, int64(1+4+4), size) // Bool(1) + Float(4) + Int32(4)
}
// =========================================================================
// validateElementLevelConsistency
// =========================================================================
func TestValidateElementLevelConsistency_Consistent(t *testing.T) {
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}},
ElementLevel: true,
ElementIndices: []*internalpb.ElementIndices{
{Indices: []int32{0, 1}},
},
}
r2 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{2}}}},
ElementLevel: true,
ElementIndices: []*internalpb.ElementIndices{
{Indices: []int32{3}},
},
}
err := validateElementLevelConsistency([]*internalpb.RetrieveResults{r1, r2}, nil)
assert.NoError(t, err)
}
// makeStructArrayFieldData creates a FieldData wrapping a StructArrayField with scalar sub-fields.
// Used by TestGetRowCount and slice/reconstruct tests below.
func makeStructArrayFieldData(fieldID int64, name string, subFields []*schemapb.FieldData) *schemapb.FieldData {
return &schemapb.FieldData{
Type: schemapb.DataType_Array,
FieldId: fieldID,
FieldName: name,
Field: &schemapb.FieldData_StructArrays{
StructArrays: &schemapb.StructArrayField{
Fields: subFields,
},
},
}
}
// =========================================================================
// rangeSliceStructArrayField
// =========================================================================
func TestRangeSliceStructArrayField(t *testing.T) {
sa := &schemapb.StructArrayField{
Fields: []*schemapb.FieldData{
{
Type: schemapb.DataType_Int64, FieldId: 101, FieldName: "age",
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{10, 20, 30, 40}}},
}},
},
{
Type: schemapb.DataType_VarChar, FieldId: 102, FieldName: "name",
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"a", "b", "c", "d"}}},
}},
},
},
}
sliced, err := rangeSliceStructArrayField(sa, 1, 3)
require.NoError(t, err)
require.Len(t, sliced.GetFields(), 2)
assert.Equal(t, []int64{20, 30}, sliced.GetFields()[0].GetScalars().GetLongData().GetData())
assert.Equal(t, []string{"b", "c"}, sliced.GetFields()[1].GetScalars().GetStringData().GetData())
}
func TestRangeSliceStructArrayField_Nil(t *testing.T) {
sliced, err := rangeSliceStructArrayField(nil, 0, 1)
require.NoError(t, err)
assert.Nil(t, sliced)
}
// =========================================================================
// sliceStructArrayField
// =========================================================================
func TestSliceStructArrayField(t *testing.T) {
sa := &schemapb.StructArrayField{
Fields: []*schemapb.FieldData{
{
Type: schemapb.DataType_Int64, FieldId: 101, FieldName: "val",
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{100, 200, 300, 400}}},
}},
},
{
Type: schemapb.DataType_VarChar, FieldId: 102, FieldName: "tag",
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"x", "y", "z", "w"}}},
}},
},
},
}
// Select indices 3, 0, 2 (out of order)
sliced, err := sliceStructArrayField(sa, []int{3, 0, 2})
require.NoError(t, err)
require.Len(t, sliced.GetFields(), 2)
assert.Equal(t, []int64{400, 100, 300}, sliced.GetFields()[0].GetScalars().GetLongData().GetData())
assert.Equal(t, []string{"w", "x", "z"}, sliced.GetFields()[1].GetScalars().GetStringData().GetData())
}
func TestSliceStructArrayField_Nil(t *testing.T) {
sliced, err := sliceStructArrayField(nil, []int{0})
require.NoError(t, err)
assert.Nil(t, sliced)
}
// =========================================================================
// buildMergedVectorField: VectorArray (used in StructArray vector sub-fields)
// =========================================================================
func TestBuildMergedVectorField_VectorArray(t *testing.T) {
dim := int64(2)
v1 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1.0, 2.0}}}}
v2 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{3.0, 4.0}}}}
v3 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{5.0, 6.0}}}}
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_FloatVector, FieldId: 700,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: dim,
Data: []*schemapb.VectorField{v1, v2},
ElementType: schemapb.DataType_FloatVector,
}},
}},
},
},
}
r2 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{3}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_FloatVector, FieldId: 700,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: dim,
Data: []*schemapb.VectorField{v3},
ElementType: schemapb.DataType_FloatVector,
}},
}},
},
},
}
results := []*internalpb.RetrieveResults{r1, r2}
selectedRows := []rowRef{{resultIdx: 1, rowIdx: 0}, {resultIdx: 0, rowIdx: 0}}
merged, err := buildMergedRetrieveResults(results, selectedRows, nil)
require.NoError(t, err)
va := merged.FieldsData[0].GetVectors().GetVectorArray()
require.Len(t, va.GetData(), 2)
assert.Equal(t, []float32{5.0, 6.0}, va.GetData()[0].GetFloatVector().GetData())
assert.Equal(t, []float32{1.0, 2.0}, va.GetData()[1].GetFloatVector().GetData())
}
func TestBuildMergedVectorField_NullableStructVectorArray_AllNull(t *testing.T) {
const fieldID = int64(700)
dim := int64(2)
result := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_ArrayOfVector,
FieldId: fieldID,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: dim,
}},
ValidData: []bool{false},
},
},
}
schema := &schemapb.CollectionSchema{
StructArrayFields: []*schemapb.StructArrayFieldSchema{
{
Name: "struct_field",
Nullable: true,
Fields: []*schemapb.FieldSchema{
{
FieldID: fieldID,
Name: "struct_field[vec_array]",
DataType: schemapb.DataType_ArrayOfVector,
ElementType: schemapb.DataType_FloatVector,
Nullable: true,
TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "2"},
},
},
},
},
},
}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{result}, []rowRef{{resultIdx: 0, rowIdx: 0}}, schema)
require.NoError(t, err)
assert.Equal(t, []bool{false}, merged.FieldsData[0].GetValidData())
va := merged.FieldsData[0].GetVectors().GetVectorArray()
require.NotNil(t, va)
assert.Equal(t, dim, va.GetDim())
assert.Equal(t, schemapb.DataType_FloatVector, va.GetElementType())
require.Len(t, va.GetData(), 1)
assert.Empty(t, va.GetData()[0].GetFloatVector().GetData())
}
func TestBuildMergedVectorField_NullableStructVectorArray_AllNullWithEmptyVectorArray(t *testing.T) {
const fieldID = int64(700)
dim := int64(2)
result := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_ArrayOfVector,
FieldId: fieldID,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: dim,
ElementType: schemapb.DataType_FloatVector,
}},
}},
ValidData: []bool{false, false},
},
},
}
schema := makeNullableSchema(fieldID, schemapb.DataType_ArrayOfVector, dim)
schema.GetFields()[0].ElementType = schemapb.DataType_FloatVector
merged, err := buildMergedRetrieveResults(
[]*internalpb.RetrieveResults{result},
[]rowRef{{resultIdx: 0, rowIdx: 1}},
schema,
)
require.NoError(t, err)
assert.Equal(t, []bool{false}, merged.FieldsData[0].GetValidData())
va := merged.FieldsData[0].GetVectors().GetVectorArray()
require.NotNil(t, va)
assert.Equal(t, dim, va.GetDim())
assert.Equal(t, schemapb.DataType_FloatVector, va.GetElementType())
require.Len(t, va.GetData(), 1)
assert.Empty(t, va.GetData()[0].GetFloatVector().GetData())
}
func TestBuildMergedVectorField_NullableArrayOfVector_RejectsEmptyValidData(t *testing.T) {
const fieldID = int64(700)
dim := int64(2)
result := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_ArrayOfVector,
FieldId: fieldID,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: dim,
Data: []*schemapb.VectorField{{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1.0, 2.0}}}}},
ElementType: schemapb.DataType_FloatVector,
}},
}},
},
},
}
schema := makeNullableSchema(fieldID, schemapb.DataType_ArrayOfVector, dim)
schema.GetFields()[0].ElementType = schemapb.DataType_FloatVector
_, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{result}, []rowRef{{resultIdx: 0, rowIdx: 0}}, schema)
require.Error(t, err)
assert.Contains(t, err.Error(), "empty ValidData")
}
func TestBuildMergedVectorField_NullableArrayOfVector_RejectsValidDataRowOutOfRange(t *testing.T) {
const fieldID = int64(700)
dim := int64(2)
result := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_ArrayOfVector,
FieldId: fieldID,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: dim,
Data: []*schemapb.VectorField{{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1.0, 2.0}}}}, {Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{3.0, 4.0}}}}},
ElementType: schemapb.DataType_FloatVector,
}},
}},
ValidData: []bool{true},
},
},
}
schema := makeNullableSchema(fieldID, schemapb.DataType_ArrayOfVector, dim)
schema.GetFields()[0].ElementType = schemapb.DataType_FloatVector
_, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{result}, []rowRef{{resultIdx: 0, rowIdx: 1}}, schema)
require.Error(t, err)
assert.Contains(t, err.Error(), "rowIdx=1")
assert.Contains(t, err.Error(), "len(ValidData)=1")
}
func TestBuildMergedVectorField_NullableArrayOfVectorMixedAllNullSourceKeepsRowDense(t *testing.T) {
const fieldID = int64(700)
dim := int64(2)
v1 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1.0, 2.0}}}}
allNullResult := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_ArrayOfVector,
FieldId: fieldID,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: dim,
}},
ValidData: []bool{false},
},
},
}
validResult := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{2}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_ArrayOfVector,
FieldId: fieldID,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: dim,
Data: []*schemapb.VectorField{v1},
ElementType: schemapb.DataType_FloatVector,
}},
}},
ValidData: []bool{true},
},
},
}
schema := makeNullableSchema(fieldID, schemapb.DataType_ArrayOfVector, dim)
schema.GetFields()[0].ElementType = schemapb.DataType_FloatVector
merged, err := buildMergedRetrieveResults(
[]*internalpb.RetrieveResults{allNullResult, validResult},
[]rowRef{{resultIdx: 0, rowIdx: 0}, {resultIdx: 1, rowIdx: 0}},
schema,
)
require.NoError(t, err)
assert.Equal(t, []bool{false, true}, merged.FieldsData[0].GetValidData())
va := merged.FieldsData[0].GetVectors().GetVectorArray()
require.Len(t, va.GetData(), 2)
assert.Empty(t, va.GetData()[0].GetFloatVector().GetData())
assert.Equal(t, []float32{1.0, 2.0}, va.GetData()[1].GetFloatVector().GetData())
}
func TestBuildMergedVectorField_NullableArrayOfVectorUsesLogicalRowIndex(t *testing.T) {
const fieldID = int64(700)
dim := int64(2)
v1 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1.0, 2.0}}}}
empty := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{}}}
v3 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{3.0, 4.0}}}}
result := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2, 3, 4}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_ArrayOfVector,
FieldId: fieldID,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: dim,
Data: []*schemapb.VectorField{v1, empty, v3, empty},
ElementType: schemapb.DataType_FloatVector,
}},
}},
ValidData: []bool{true, false, true, false},
},
},
}
schema := makeNullableSchema(fieldID, schemapb.DataType_ArrayOfVector, dim)
schema.GetFields()[0].ElementType = schemapb.DataType_FloatVector
merged, err := buildMergedRetrieveResults(
[]*internalpb.RetrieveResults{result},
[]rowRef{{resultIdx: 0, rowIdx: 2}, {resultIdx: 0, rowIdx: 1}, {resultIdx: 0, rowIdx: 0}},
schema,
)
require.NoError(t, err)
assert.Equal(t, []bool{true, false, true}, merged.FieldsData[0].GetValidData())
va := merged.FieldsData[0].GetVectors().GetVectorArray()
require.Len(t, va.GetData(), 3)
assert.Equal(t, []float32{3.0, 4.0}, va.GetData()[0].GetFloatVector().GetData())
assert.Empty(t, va.GetData()[1].GetFloatVector().GetData())
assert.Equal(t, []float32{1.0, 2.0}, va.GetData()[2].GetFloatVector().GetData())
}
// =========================================================================
// sliceVectorField & rangeSliceVectorField: VectorArray branch
// =========================================================================
func TestSliceVectorField_VectorArray(t *testing.T) {
dim := int64(2)
v1 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1.0, 2.0}}}}
v2 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{3.0, 4.0}}}}
v3 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{5.0, 6.0}}}}
vf := &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: dim,
Data: []*schemapb.VectorField{v1, v2, v3},
ElementType: schemapb.DataType_FloatVector,
}},
}
sliced := requireSliceVectorField(t, vf, []int{2, 0}, nil)
va := sliced.GetVectorArray()
require.Len(t, va.GetData(), 2)
assert.Equal(t, []float32{5.0, 6.0}, va.GetData()[0].GetFloatVector().GetData())
assert.Equal(t, []float32{1.0, 2.0}, va.GetData()[1].GetFloatVector().GetData())
assert.Equal(t, schemapb.DataType_FloatVector, va.GetElementType())
}
func TestRangeSliceVectorField_VectorArray(t *testing.T) {
dim := int64(2)
v1 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1.0, 2.0}}}}
v2 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{3.0, 4.0}}}}
v3 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{5.0, 6.0}}}}
vf := &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: dim,
Data: []*schemapb.VectorField{v1, v2, v3},
ElementType: schemapb.DataType_FloatVector,
}},
}
sliced := requireRangeSliceVectorField(t, vf, 1, 3, nil)
va := sliced.GetVectorArray()
require.Len(t, va.GetData(), 2)
assert.Equal(t, []float32{3.0, 4.0}, va.GetData()[0].GetFloatVector().GetData())
assert.Equal(t, []float32{5.0, 6.0}, va.GetData()[1].GetFloatVector().GetData())
}
// =========================================================================
// comparePK: string PK
// =========================================================================
func TestComparePK_StringPK(t *testing.T) {
assert.Equal(t, -1, comparePK("apple", "banana"))
assert.Equal(t, 0, comparePK("same", "same"))
assert.Equal(t, 1, comparePK("zebra", "alpha"))
}
func TestComparePK_Int64PK(t *testing.T) {
assert.Equal(t, -1, comparePK(int64(1), int64(2)))
assert.Equal(t, 0, comparePK(int64(5), int64(5)))
assert.Equal(t, 1, comparePK(int64(10), int64(3)))
}
// =========================================================================
// buildMergedScalarField: additional type branches
// =========================================================================
func TestBuildMergedScalarField_BoolData(t *testing.T) {
r1 := makeInternalResultIntPK([]int64{1, 2}, &schemapb.FieldData{
Type: schemapb.DataType_Bool, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: []bool{true, false}}},
}},
})
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 1}, {resultIdx: 0, rowIdx: 0}}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, nil)
require.NoError(t, err)
assert.Equal(t, []bool{false, true}, merged.FieldsData[0].GetScalars().GetBoolData().GetData())
}
func TestBuildMergedScalarField_IntData(t *testing.T) {
r1 := makeInternalResultIntPK([]int64{1, 2}, &schemapb.FieldData{
Type: schemapb.DataType_Int32, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{10, 20}}},
}},
})
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 1}}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, nil)
require.NoError(t, err)
assert.Equal(t, []int32{20}, merged.FieldsData[0].GetScalars().GetIntData().GetData())
}
func TestBuildMergedScalarField_FloatData(t *testing.T) {
r1 := makeInternalResultIntPK([]int64{1}, &schemapb.FieldData{
Type: schemapb.DataType_Float, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_FloatData{FloatData: &schemapb.FloatArray{Data: []float32{1.5, 2.5}}},
}},
})
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, nil)
require.NoError(t, err)
assert.Equal(t, []float32{1.5}, merged.FieldsData[0].GetScalars().GetFloatData().GetData())
}
func TestBuildMergedScalarField_DoubleData(t *testing.T) {
r1 := makeInternalResultIntPK([]int64{1, 2}, &schemapb.FieldData{
Type: schemapb.DataType_Double, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_DoubleData{DoubleData: &schemapb.DoubleArray{Data: []float64{3.14, 2.71}}},
}},
})
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 1}}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, nil)
require.NoError(t, err)
assert.Equal(t, []float64{2.71}, merged.FieldsData[0].GetScalars().GetDoubleData().GetData())
}
func TestBuildMergedScalarField_BytesData(t *testing.T) {
r1 := makeInternalResultIntPK([]int64{1, 2}, &schemapb.FieldData{
Type: schemapb.DataType_BinaryVector, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_BytesData{BytesData: &schemapb.BytesArray{Data: [][]byte{{0x01}, {0x02}}}},
}},
})
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}, {resultIdx: 0, rowIdx: 1}}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, nil)
require.NoError(t, err)
assert.Equal(t, [][]byte{{0x01}, {0x02}}, merged.FieldsData[0].GetScalars().GetBytesData().GetData())
}
func TestBuildMergedScalarField_JsonData(t *testing.T) {
r1 := makeInternalResultIntPK([]int64{1, 2}, &schemapb.FieldData{
Type: schemapb.DataType_JSON, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_JsonData{JsonData: &schemapb.JSONArray{Data: [][]byte{[]byte(`{"a":1}`), []byte(`{"b":2}`)}}},
}},
})
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 1}}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, nil)
require.NoError(t, err)
assert.Equal(t, [][]byte{[]byte(`{"b":2}`)}, merged.FieldsData[0].GetScalars().GetJsonData().GetData())
}
// =========================================================================
// buildMergedIDs: string PK
// =========================================================================
func TestBuildMergedIDs_StringPK(t *testing.T) {
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_StrId{StrId: &schemapb.StringArray{Data: []string{"a", "b"}}}},
FieldsData: []*schemapb.FieldData{
makeStringField(100, "pk", []string{"a", "b"}),
},
}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 1}, {resultIdx: 0, rowIdx: 0}}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, nil)
require.NoError(t, err)
assert.Equal(t, []string{"b", "a"}, merged.GetIds().GetStrId().GetData())
}
// =========================================================================
// buildMergedVectorField: additional vector type branches
// =========================================================================
func TestBuildMergedVectorField_BinaryVector(t *testing.T) {
dim := 16 // 2 bytes per row
r1 := makeInternalResultIntPK([]int64{1, 2}, &schemapb.FieldData{
Type: schemapb.DataType_BinaryVector, FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{0xAA, 0xBB, 0xCC, 0xDD}},
}},
})
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 1}}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, nil)
require.NoError(t, err)
assert.Equal(t, []byte{0xCC, 0xDD}, merged.FieldsData[0].GetVectors().GetBinaryVector())
}
func TestBuildMergedVectorField_Float16Vector(t *testing.T) {
dim := 2 // 4 bytes per row (dim*2)
r1 := makeInternalResultIntPK([]int64{1, 2}, &schemapb.FieldData{
Type: schemapb.DataType_Float16Vector, FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{1, 2, 3, 4, 5, 6, 7, 8}},
}},
})
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 1}}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, nil)
require.NoError(t, err)
assert.Equal(t, []byte{5, 6, 7, 8}, merged.FieldsData[0].GetVectors().GetFloat16Vector())
}
func TestBuildMergedVectorField_BFloat16Vector(t *testing.T) {
dim := 2
r1 := makeInternalResultIntPK([]int64{1, 2}, &schemapb.FieldData{
Type: schemapb.DataType_BFloat16Vector, FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{10, 20, 30, 40, 50, 60, 70, 80}},
}},
})
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, nil)
require.NoError(t, err)
assert.Equal(t, []byte{10, 20, 30, 40}, merged.FieldsData[0].GetVectors().GetBfloat16Vector())
}
func TestBuildMergedVectorField_SparseFloatVector(t *testing.T) {
r1 := makeInternalResultIntPK([]int64{1, 2}, &schemapb.FieldData{
Type: schemapb.DataType_SparseFloatVector, FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 100,
Data: &schemapb.VectorField_SparseFloatVector{SparseFloatVector: &schemapb.SparseFloatArray{
Contents: [][]byte{{0x01, 0x02}, {0x03, 0x04}},
Dim: 100,
}},
}},
})
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 1}}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, nil)
require.NoError(t, err)
assert.Equal(t, [][]byte{{0x03, 0x04}}, merged.FieldsData[0].GetVectors().GetSparseFloatVector().GetContents())
}
// =========================================================================
// buildCompactIndices / getVecDataIdx: nullable vector
// =========================================================================
func TestBuildCompactIndices_NullableVector(t *testing.T) {
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2, 3}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_FloatVector, FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4}}},
}},
ValidData: []bool{true, false, true}, // row0=valid(idx0), row1=null, row2=valid(idx1)
},
},
}
indices, err := buildCompactIndices([]*internalpb.RetrieveResults{r1}, 0, true)
require.NoError(t, err)
require.NotNil(t, indices)
assert.Equal(t, []int{0, -1, 1}, indices[0])
}
func TestBuildCompactIndices_NonNullable(t *testing.T) {
r1 := &internalpb.RetrieveResults{
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_FloatVector, FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4}}},
}},
// No ValidData
},
},
}
indices, err := buildCompactIndices([]*internalpb.RetrieveResults{r1}, 0, false)
require.NoError(t, err)
assert.Nil(t, indices) // non-nullable → nil
}
func TestBuildCompactIndices_FailFast_EmptyValidData(t *testing.T) {
// nullable field with numRows > 0 but empty ValidData → must error
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_SparseFloatVector, FieldId: 100,
// ValidData intentionally empty — segcore contract violation
},
},
}
_, err := buildCompactIndices([]*internalpb.RetrieveResults{r1}, 0, true)
require.Error(t, err)
assert.Contains(t, err.Error(), "empty ValidData")
}
func TestBuildCompactIndices_FailFast_LengthMismatch(t *testing.T) {
// nullable field with len(ValidData) != numRows → must error
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2, 3}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_FloatVector, FieldId: 100,
ValidData: []bool{true, false}, // 2 entries but 3 rows → mismatch
},
},
}
_, err := buildCompactIndices([]*internalpb.RetrieveResults{r1}, 0, true)
require.Error(t, err)
assert.Contains(t, err.Error(), "len(ValidData)")
}
func TestGetVecDataIdx(t *testing.T) {
// nil compactIndices → returns rowIdx directly
assert.Equal(t, 3, getVecDataIdx(nil, rowRef{resultIdx: 0, rowIdx: 3}))
// With compact indices
compactIdx := [][]int{
{0, -1, 1}, // result 0: row0=0, row1=null, row2=1
nil, // result 1: non-nullable
}
assert.Equal(t, 0, getVecDataIdx(compactIdx, rowRef{resultIdx: 0, rowIdx: 0}))
assert.Equal(t, -1, getVecDataIdx(compactIdx, rowRef{resultIdx: 0, rowIdx: 1}))
assert.Equal(t, 1, getVecDataIdx(compactIdx, rowRef{resultIdx: 0, rowIdx: 2}))
assert.Equal(t, 5, getVecDataIdx(compactIdx, rowRef{resultIdx: 1, rowIdx: 5})) // nil slice → rowIdx
// Out of range → panics (validated upstream by buildCompactIndices)
assert.Panics(t, func() { getVecDataIdx(compactIdx, rowRef{resultIdx: 0, rowIdx: 10}) })
}
// =========================================================================
// buildMergedVectorField: truncated data error paths
// =========================================================================
func TestBuildMergedVectorField_FloatVector_TruncatedData(t *testing.T) {
// FloatVector with dim=2 but only 2 floats (enough for 1 row, not 2)
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{{
FieldId: 100, Type: schemapb.DataType_FloatVector,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2}}}, // only 1 row
}},
}},
}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}, {resultIdx: 0, rowIdx: 1}}
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
{FieldID: 100, DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "2"}}},
}}
_, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, schema)
require.Error(t, err)
assert.Contains(t, err.Error(), "truncated data")
}
func TestBuildMergedVectorField_BinaryVector_TruncatedData(t *testing.T) {
// BinaryVector with dim=16 (2 bytes/row) but only 2 bytes (1 row, not 2)
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{{
FieldId: 100, Type: schemapb.DataType_BinaryVector,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 16,
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{0xFF, 0x00}}, // 1 row
}},
}},
}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}, {resultIdx: 0, rowIdx: 1}}
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
{FieldID: 100, DataType: schemapb.DataType_BinaryVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "16"}}},
}}
_, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, schema)
require.Error(t, err)
assert.Contains(t, err.Error(), "truncated data")
}
func TestBuildMergedVectorField_SparseVector_TruncatedContents(t *testing.T) {
// Sparse with ValidData=[true, true] but Contents has only 1 entry
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{{
FieldId: 100, Type: schemapb.DataType_SparseFloatVector,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Data: &schemapb.VectorField_SparseFloatVector{SparseFloatVector: &schemapb.SparseFloatArray{
Contents: [][]byte{makeTestSparseVec(1, 0.5)}, // only 1 content
}},
}},
ValidData: []bool{true, true}, // claims 2 valid rows
}},
}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}, {resultIdx: 0, rowIdx: 1}}
_, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows,
makeNullableSchema(100, schemapb.DataType_SparseFloatVector, 0))
require.Error(t, err)
assert.Contains(t, err.Error(), "truncated data")
}
func TestBuildMergedVectorField_NullableCompact_BinaryVector(t *testing.T) {
// BinaryVector dim=8 (1 byte/row), nullable: row0=valid, row1=null, row2=valid
// Compact data: 2 bytes (only valid rows)
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2, 3}}}},
FieldsData: []*schemapb.FieldData{{
FieldId: 100, Type: schemapb.DataType_BinaryVector,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 8,
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{0xAA, 0xBB}}, // 2 compact rows
}},
ValidData: []bool{true, false, true},
}},
}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 2}, {resultIdx: 0, rowIdx: 0}, {resultIdx: 0, rowIdx: 1}}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows,
makeNullableSchema(100, schemapb.DataType_BinaryVector, 8))
require.NoError(t, err)
// row2 → compact idx 1 → 0xBB, row0 → compact idx 0 → 0xAA, row1 → null (skipped)
assert.Equal(t, []byte{0xBB, 0xAA}, merged.FieldsData[0].GetVectors().GetBinaryVector())
assert.Equal(t, []bool{true, true, false}, merged.FieldsData[0].ValidData)
}
func TestBuildMergedVectorField_NullableCompact_Float16Vector(t *testing.T) {
// Float16 dim=1 (2 bytes/row), nullable: row0=null, row1=valid
// Compact data: 2 bytes (1 valid row)
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{{
FieldId: 100, Type: schemapb.DataType_Float16Vector,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 1,
Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{0x01, 0x02}}, // 1 compact row
}},
ValidData: []bool{false, true},
}},
}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 1}, {resultIdx: 0, rowIdx: 0}}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows,
makeNullableSchema(100, schemapb.DataType_Float16Vector, 1))
require.NoError(t, err)
// row1 → compact idx 0 → [0x01, 0x02], row0 → null (skipped)
assert.Equal(t, []byte{0x01, 0x02}, merged.FieldsData[0].GetVectors().GetFloat16Vector())
assert.Equal(t, []bool{true, false}, merged.FieldsData[0].ValidData)
}
func TestBuildMergedVectorField_Float16Vector_TruncatedData(t *testing.T) {
// Float16 dim=2 (4 bytes/row) but only 4 bytes provided (1 row, not 2)
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{{
FieldId: 100, Type: schemapb.DataType_Float16Vector,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{0x01, 0x02, 0x03, 0x04}}, // 1 row
}},
}},
}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}, {resultIdx: 0, rowIdx: 1}}
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
{FieldID: 100, DataType: schemapb.DataType_Float16Vector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "2"}}},
}}
_, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, schema)
require.Error(t, err)
assert.Contains(t, err.Error(), "truncated data")
}
func TestBuildMergedVectorField_NullableCompact_BFloat16Vector(t *testing.T) {
// BFloat16 dim=1 (2 bytes/row), nullable: row0=valid, row1=null, row2=valid
// Compact data: 4 bytes (2 valid rows)
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2, 3}}}},
FieldsData: []*schemapb.FieldData{{
FieldId: 100, Type: schemapb.DataType_BFloat16Vector,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 1,
Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{0xAA, 0xBB, 0xCC, 0xDD}},
}},
ValidData: []bool{true, false, true},
}},
}
// Select row1(null), row2(valid), row0(valid)
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 1}, {resultIdx: 0, rowIdx: 2}, {resultIdx: 0, rowIdx: 0}}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows,
makeNullableSchema(100, schemapb.DataType_BFloat16Vector, 1))
require.NoError(t, err)
// row1 → null (skipped), row2 → compact idx 1 → [0xCC,0xDD], row0 → compact idx 0 → [0xAA,0xBB]
assert.Equal(t, []byte{0xCC, 0xDD, 0xAA, 0xBB}, merged.FieldsData[0].GetVectors().GetBfloat16Vector())
assert.Equal(t, []bool{false, true, true}, merged.FieldsData[0].ValidData)
}
func TestBuildMergedVectorField_BFloat16Vector_TruncatedData(t *testing.T) {
// BFloat16 dim=2 (4 bytes/row) but only 4 bytes (1 row, not 2)
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{{
FieldId: 100, Type: schemapb.DataType_BFloat16Vector,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{0x01, 0x02, 0x03, 0x04}},
}},
}},
}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}, {resultIdx: 0, rowIdx: 1}}
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
{FieldID: 100, DataType: schemapb.DataType_BFloat16Vector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "2"}}},
}}
_, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, schema)
require.Error(t, err)
assert.Contains(t, err.Error(), "truncated data")
}
func TestBuildMergedVectorField_NullableCompact_Int8Vector(t *testing.T) {
// Int8Vector dim=2 (2 bytes/row), nullable: row0=null, row1=valid
// Compact data: 2 bytes (1 valid row)
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{{
FieldId: 100, Type: schemapb.DataType_Int8Vector,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_Int8Vector{Int8Vector: []byte{0x11, 0x22}},
}},
ValidData: []bool{false, true},
}},
}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}, {resultIdx: 0, rowIdx: 1}}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows,
makeNullableSchema(100, schemapb.DataType_Int8Vector, 2))
require.NoError(t, err)
// row0 → null (skipped), row1 → compact idx 0 → [0x11,0x22]
assert.Equal(t, []byte{0x11, 0x22}, merged.FieldsData[0].GetVectors().GetInt8Vector())
assert.Equal(t, []bool{false, true}, merged.FieldsData[0].ValidData)
}
func TestBuildMergedVectorField_Int8Vector_TruncatedData(t *testing.T) {
// Int8Vector dim=3 (3 bytes/row) but only 3 bytes (1 row, not 2)
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{{
FieldId: 100, Type: schemapb.DataType_Int8Vector,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 3,
Data: &schemapb.VectorField_Int8Vector{Int8Vector: []byte{0x01, 0x02, 0x03}},
}},
}},
}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}, {resultIdx: 0, rowIdx: 1}}
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
{FieldID: 100, DataType: schemapb.DataType_Int8Vector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "3"}}},
}}
_, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, schema)
require.Error(t, err)
assert.Contains(t, err.Error(), "truncated data")
}
func TestBuildMergedVectorField_VectorArray_TruncatedData(t *testing.T) {
dim := int64(2)
v1 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{
FloatVector: &schemapb.FloatArray{Data: []float32{1.0, 2.0}},
}}
// Result has 2 IDs but VectorArray has only 1 entry
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
FieldsData: []*schemapb.FieldData{{
FieldId: 100, Type: schemapb.DataType_FloatVector,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: dim,
Data: []*schemapb.VectorField{v1}, // only 1 entry
ElementType: schemapb.DataType_FloatVector,
}},
}},
}},
}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}, {resultIdx: 0, rowIdx: 1}}
// nil schema → fallback infers VectorArray from data
_, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, nil)
require.Error(t, err)
assert.Contains(t, err.Error(), "truncated data")
}
// =========================================================================
// rangeSlice / slice: string IDs
// =========================================================================
func TestRangeSliceIDs_StringID(t *testing.T) {
ids := &schemapb.IDs{IdField: &schemapb.IDs_StrId{StrId: &schemapb.StringArray{Data: []string{"a", "b", "c", "d"}}}}
sliced := rangeSliceIDs(ids, 1, 3)
assert.Equal(t, []string{"b", "c"}, sliced.GetStrId().GetData())
}
func TestSliceIDs_StringID(t *testing.T) {
ids := &schemapb.IDs{IdField: &schemapb.IDs_StrId{StrId: &schemapb.StringArray{Data: []string{"x", "y", "z"}}}}
sliced := sliceIDs(ids, []int{2, 0})
assert.Equal(t, []string{"z", "x"}, sliced.GetStrId().GetData())
}
// =========================================================================
// rangeSliceVectorField: additional vector types
// =========================================================================
func TestRangeSliceVectorField_FloatVector(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5, 6}}},
}
sliced := requireRangeSliceVectorField(t, vf, 1, 3, nil)
assert.Equal(t, []float32{3, 4, 5, 6}, sliced.GetFloatVector().GetData())
}
func TestRangeSliceVectorField_BinaryVector(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 16, // 2 bytes per row
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF}},
}
sliced := requireRangeSliceVectorField(t, vf, 0, 2, nil)
assert.Equal(t, []byte{0xAA, 0xBB, 0xCC, 0xDD}, sliced.GetBinaryVector())
}
func TestRangeSliceVectorField_Float16Vector(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 2, // 4 bytes per row
Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{1, 2, 3, 4, 5, 6, 7, 8}},
}
sliced := requireRangeSliceVectorField(t, vf, 1, 2, nil)
assert.Equal(t, []byte{5, 6, 7, 8}, sliced.GetFloat16Vector())
}
func TestRangeSliceVectorField_SparseFloatVector(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 100,
Data: &schemapb.VectorField_SparseFloatVector{SparseFloatVector: &schemapb.SparseFloatArray{
Contents: [][]byte{{0x01}, {0x02}, {0x03}},
Dim: 100,
}},
}
sliced := requireRangeSliceVectorField(t, vf, 1, 3, nil)
assert.Equal(t, [][]byte{{0x02}, {0x03}}, sliced.GetSparseFloatVector().GetContents())
}
// =========================================================================
// sliceVectorField: additional vector types
// =========================================================================
func TestSliceVectorField_FloatVector(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5, 6}}},
}
sliced := requireSliceVectorField(t, vf, []int{2, 0}, nil)
assert.Equal(t, []float32{5, 6, 1, 2}, sliced.GetFloatVector().GetData())
}
func TestSliceVectorField_BinaryVector(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 16, // 2 bytes per row
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF}},
}
sliced := requireSliceVectorField(t, vf, []int{2, 0}, nil)
assert.Equal(t, []byte{0xEE, 0xFF, 0xAA, 0xBB}, sliced.GetBinaryVector())
}
func TestSliceVectorField_Float16Vector(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 2, // 4 bytes per row
Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{1, 2, 3, 4, 5, 6, 7, 8}},
}
sliced := requireSliceVectorField(t, vf, []int{1}, nil)
assert.Equal(t, []byte{5, 6, 7, 8}, sliced.GetFloat16Vector())
}
func TestSliceVectorField_BFloat16Vector(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 2, // 4 bytes per row
Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{10, 20, 30, 40, 50, 60, 70, 80}},
}
sliced := requireSliceVectorField(t, vf, []int{1, 0}, nil)
assert.Equal(t, []byte{50, 60, 70, 80, 10, 20, 30, 40}, sliced.GetBfloat16Vector())
}
func TestSliceVectorField_Int8Vector(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 3,
Data: &schemapb.VectorField_Int8Vector{Int8Vector: []byte{1, 2, 3, 4, 5, 6}},
}
sliced := requireSliceVectorField(t, vf, []int{1}, nil)
assert.Equal(t, []byte{4, 5, 6}, sliced.GetInt8Vector())
}
func TestSliceVectorField_SparseFloatVector(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 100,
Data: &schemapb.VectorField_SparseFloatVector{SparseFloatVector: &schemapb.SparseFloatArray{
Contents: [][]byte{{0x01}, {0x02}, {0x03}},
Dim: 100,
}},
}
sliced := requireSliceVectorField(t, vf, []int{2, 0}, nil)
assert.Equal(t, [][]byte{{0x03}, {0x01}}, sliced.GetSparseFloatVector().GetContents())
}
// =========================================================================
// rangeSliceFieldData: vector branch
// =========================================================================
func TestRangeSliceFieldData_Vector(t *testing.T) {
fd := &schemapb.FieldData{
Type: schemapb.DataType_FloatVector, FieldId: 100, FieldName: "vec",
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5, 6}}},
}},
}
sliced := requireRangeSliceFieldData(t, fd, 1, 3)
assert.Equal(t, []float32{3, 4, 5, 6}, sliced.GetVectors().GetFloatVector().GetData())
}
func TestRangeSliceFieldData_Nil(t *testing.T) {
sliced := requireRangeSliceFieldData(t, nil, 0, 1)
assert.Nil(t, sliced)
}
// =========================================================================
// calcFieldElementSize: various types
// =========================================================================
func TestCalcFieldElementSize(t *testing.T) {
tests := []struct {
name string
fd *schemapb.FieldData
rowIdx int
expected int64
}{
{
name: "bool",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: []bool{true}}},
}}},
expected: 1,
},
{
name: "int32",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{42}}},
}}},
expected: 4,
},
{
name: "int64",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{42}}},
}}},
expected: 8,
},
{
name: "float32",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_FloatData{FloatData: &schemapb.FloatArray{Data: []float32{1.0}}},
}}},
expected: 4,
},
{
name: "float64",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_DoubleData{DoubleData: &schemapb.DoubleArray{Data: []float64{1.0}}},
}}},
expected: 8,
},
{
name: "string",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"hello"}}},
}}},
expected: 5,
},
{
name: "json",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_JsonData{JsonData: &schemapb.JSONArray{Data: [][]byte{[]byte(`{"a":1}`)}}},
}}},
expected: 7,
},
{
name: "bytes",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_BytesData{BytesData: &schemapb.BytesArray{Data: [][]byte{{0x01, 0x02, 0x03}}}},
}}},
expected: 3,
},
{
name: "timestamptz",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_TimestamptzData{TimestamptzData: &schemapb.TimestamptzArray{Data: []int64{1000}}},
}}},
expected: 8,
},
{
name: "float_vector",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 4,
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4}}},
}}},
expected: 16, // 4 * 4
},
{
name: "binary_vector",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 16,
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{0, 0}},
}}},
expected: 2, // 16/8
},
{
name: "float16_vector",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 4,
Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{0, 0, 0, 0, 0, 0, 0, 0}},
}}},
expected: 8, // 4*2
},
{
name: "int8_vector",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 3,
Data: &schemapb.VectorField_Int8Vector{Int8Vector: []byte{1, 2, 3}},
}}},
expected: 3,
},
{
name: "sparse_float_vector",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 100,
Data: &schemapb.VectorField_SparseFloatVector{SparseFloatVector: &schemapb.SparseFloatArray{
Contents: [][]byte{{0x01, 0x02, 0x03, 0x04}},
Dim: 100,
}},
}}},
expected: 4,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
assert.Equal(t, tt.expected, calcFieldElementSize(tt.fd, tt.rowIdx))
})
}
}
// =========================================================================
// getFieldValue: various types
// =========================================================================
func TestGetFieldValue(t *testing.T) {
tests := []struct {
name string
fd *schemapb.FieldData
rowIdx int
expected any
isNull bool
}{
{
name: "bool",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: []bool{true, false}}},
}}},
rowIdx: 1, expected: false,
},
{
name: "int32",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{10, 20}}},
}}},
rowIdx: 0, expected: int32(10),
},
{
name: "float32",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_FloatData{FloatData: &schemapb.FloatArray{Data: []float32{1.5}}},
}}},
rowIdx: 0, expected: float32(1.5),
},
{
name: "float64",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_DoubleData{DoubleData: &schemapb.DoubleArray{Data: []float64{3.14}}},
}}},
rowIdx: 0, expected: float64(3.14),
},
{
name: "nullable_null",
fd: &schemapb.FieldData{
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{99}}},
}},
ValidData: []bool{false},
},
rowIdx: 0, expected: nil, isNull: true,
},
{
name: "nil_scalars",
fd: &schemapb.FieldData{},
rowIdx: 0, expected: nil, isNull: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
val, isNull := getFieldValue(tt.fd, tt.rowIdx)
assert.Equal(t, tt.isNull, isNull)
if !tt.isNull {
assert.Equal(t, tt.expected, val)
}
})
}
}
// =========================================================================
// getRowCount: various types
// =========================================================================
func TestGetRowCount(t *testing.T) {
tests := []struct {
name string
result *internalpb.RetrieveResults
expected int
}{
{
name: "from IDs",
result: makeInternalResultIntPK([]int64{1, 2, 3}),
expected: 3,
},
{
name: "from bool field",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: []bool{true, false}}},
}}},
}},
expected: 2,
},
{
name: "from int32 field",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{1, 2, 3}}},
}}},
}},
expected: 3,
},
{
name: "from float field",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_FloatData{FloatData: &schemapb.FloatArray{Data: []float32{1.0}}},
}}},
}},
expected: 1,
},
{
name: "from double field",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_DoubleData{DoubleData: &schemapb.DoubleArray{Data: []float64{1.0, 2.0}}},
}}},
}},
expected: 2,
},
{
name: "from json field",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_JsonData{JsonData: &schemapb.JSONArray{Data: [][]byte{[]byte(`{}`)}}},
}}},
}},
expected: 1,
},
{
name: "from float vector",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5, 6}}},
}}},
}},
expected: 3,
},
{
name: "from binary vector",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 16,
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{0, 0, 0, 0}},
}}},
}},
expected: 2,
},
{
name: "from sparse float vector",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 100,
Data: &schemapb.VectorField_SparseFloatVector{SparseFloatVector: &schemapb.SparseFloatArray{
Contents: [][]byte{{0x01}, {0x02}},
}},
}}},
}},
expected: 2,
},
{
name: "from struct array",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
makeStructArrayFieldData(100, "info", []*schemapb.FieldData{
makeInt64Field(101, "age", []int64{10, 20, 30}),
}),
}},
expected: 3,
},
{
name: "empty",
result: &internalpb.RetrieveResults{},
expected: 0,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
assert.Equal(t, tt.expected, getRowCount(tt.result))
})
}
}
// =========================================================================
// compareValues: various types
// =========================================================================
func TestCompareValues(t *testing.T) {
tests := []struct {
name string
a, b any
dataType schemapb.DataType
expected int
}{
{"bool_lt", false, true, schemapb.DataType_Bool, -1},
{"bool_gt", true, false, schemapb.DataType_Bool, 1},
{"bool_eq", true, true, schemapb.DataType_Bool, 0},
{"int32_lt", int32(1), int32(2), schemapb.DataType_Int32, -1},
{"int32_gt", int32(5), int32(3), schemapb.DataType_Int32, 1},
{"int32_eq", int32(7), int32(7), schemapb.DataType_Int32, 0},
{"int64_lt", int64(10), int64(20), schemapb.DataType_Int64, -1},
{"int64_eq", int64(5), int64(5), schemapb.DataType_Int64, 0},
{"float32_lt", float32(1.0), float32(2.0), schemapb.DataType_Float, -1},
{"float32_gt", float32(3.0), float32(1.0), schemapb.DataType_Float, 1},
{"float32_eq", float32(1.5), float32(1.5), schemapb.DataType_Float, 0},
{"float64_lt", float64(1.0), float64(2.0), schemapb.DataType_Double, -1},
{"float64_gt", float64(9.0), float64(1.0), schemapb.DataType_Double, 1},
{"float64_eq", float64(3.14), float64(3.14), schemapb.DataType_Double, 0},
{"string_lt", "apple", "banana", schemapb.DataType_VarChar, -1},
{"string_gt", "zebra", "alpha", schemapb.DataType_VarChar, 1},
{"string_eq", "same", "same", schemapb.DataType_VarChar, 0},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
assert.Equal(t, tt.expected, compareValues(tt.a, tt.b, tt.dataType))
})
}
}
// =========================================================================
// rangeSliceScalarField / sliceScalarField: more type branches
// =========================================================================
func TestRangeSliceScalarField_AllTypes(t *testing.T) {
tests := []struct {
name string
sf *schemapb.ScalarField
check func(t *testing.T, sf *schemapb.ScalarField)
}{
{
name: "bool",
sf: &schemapb.ScalarField{Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: []bool{true, false, true}}}},
check: func(t *testing.T, sf *schemapb.ScalarField) {
assert.Equal(t, []bool{false}, sf.GetBoolData().GetData())
},
},
{
name: "int32",
sf: &schemapb.ScalarField{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{10, 20, 30}}}},
check: func(t *testing.T, sf *schemapb.ScalarField) {
assert.Equal(t, []int32{20}, sf.GetIntData().GetData())
},
},
{
name: "float",
sf: &schemapb.ScalarField{Data: &schemapb.ScalarField_FloatData{FloatData: &schemapb.FloatArray{Data: []float32{1.0, 2.0, 3.0}}}},
check: func(t *testing.T, sf *schemapb.ScalarField) {
assert.Equal(t, []float32{2.0}, sf.GetFloatData().GetData())
},
},
{
name: "double",
sf: &schemapb.ScalarField{Data: &schemapb.ScalarField_DoubleData{DoubleData: &schemapb.DoubleArray{Data: []float64{1.0, 2.0, 3.0}}}},
check: func(t *testing.T, sf *schemapb.ScalarField) {
assert.Equal(t, []float64{2.0}, sf.GetDoubleData().GetData())
},
},
{
name: "bytes",
sf: &schemapb.ScalarField{Data: &schemapb.ScalarField_BytesData{BytesData: &schemapb.BytesArray{Data: [][]byte{{1}, {2}, {3}}}}},
check: func(t *testing.T, sf *schemapb.ScalarField) {
assert.Equal(t, [][]byte{{2}}, sf.GetBytesData().GetData())
},
},
{
name: "json",
sf: &schemapb.ScalarField{Data: &schemapb.ScalarField_JsonData{JsonData: &schemapb.JSONArray{Data: [][]byte{[]byte(`{}`), []byte(`[]`), []byte(`""`)}}}},
check: func(t *testing.T, sf *schemapb.ScalarField) {
assert.Equal(t, [][]byte{[]byte(`[]`)}, sf.GetJsonData().GetData())
},
},
{
name: "timestamptz",
sf: &schemapb.ScalarField{Data: &schemapb.ScalarField_TimestamptzData{TimestamptzData: &schemapb.TimestamptzArray{Data: []int64{100, 200, 300}}}},
check: func(t *testing.T, sf *schemapb.ScalarField) {
assert.Equal(t, []int64{200}, sf.GetTimestamptzData().GetData())
},
},
{
name: "mol",
sf: &schemapb.ScalarField{Data: &schemapb.ScalarField_MolData{MolData: &schemapb.MolArray{Data: [][]byte{{0xAA}, {0xBB}, {0xCC}}}}},
check: func(t *testing.T, sf *schemapb.ScalarField) {
assert.Equal(t, [][]byte{{0xBB}}, sf.GetMolData().GetData())
},
},
}
for _, tt := range tests {
t.Run("rangeSlice_"+tt.name, func(t *testing.T) {
sliced := rangeSliceScalarField(tt.sf, 1, 2)
tt.check(t, sliced)
})
t.Run("slice_"+tt.name, func(t *testing.T) {
sliced := sliceScalarField(tt.sf, []int{1})
tt.check(t, sliced)
})
}
}
// =========================================================================
// rangeSliceVectorField: more vector types
// =========================================================================
func TestRangeSliceVectorField_BFloat16(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{10, 20, 30, 40, 50, 60, 70, 80}},
}
sliced := requireRangeSliceVectorField(t, vf, 1, 2, nil)
assert.Equal(t, []byte{50, 60, 70, 80}, sliced.GetBfloat16Vector())
}
func TestRangeSliceVectorField_Int8(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 3,
Data: &schemapb.VectorField_Int8Vector{Int8Vector: []byte{1, 2, 3, 4, 5, 6}},
}
sliced := requireRangeSliceVectorField(t, vf, 0, 1, nil)
assert.Equal(t, []byte{1, 2, 3}, sliced.GetInt8Vector())
}
// =========================================================================
// getRowCount: additional vector + scalar types
// =========================================================================
func TestGetRowCount_AdditionalTypes(t *testing.T) {
tests := []struct {
name string
result *internalpb.RetrieveResults
expected int
}{
{
name: "string field",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"a", "b"}}},
}}},
}},
expected: 2,
},
{
name: "array field",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_ArrayData{ArrayData: &schemapb.ArrayArray{Data: []*schemapb.ScalarField{nil, nil, nil}}},
}}},
}},
expected: 3,
},
{
name: "geometry wkt field",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_GeometryWktData{GeometryWktData: &schemapb.GeometryWktArray{Data: []string{"POINT(0 0)"}}},
}}},
}},
expected: 1,
},
{
name: "timestamptz field",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_TimestamptzData{TimestamptzData: &schemapb.TimestamptzArray{Data: []int64{1, 2}}},
}}},
}},
expected: 2,
},
{
name: "mol field",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_MolData{MolData: &schemapb.MolArray{Data: [][]byte{{1}}}},
}}},
}},
expected: 1,
},
{
name: "float16 vector",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{1, 2, 3, 4, 5, 6, 7, 8}},
}}},
}},
expected: 2,
},
{
name: "bfloat16 vector",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{1, 2, 3, 4}},
}}},
}},
expected: 1,
},
{
name: "int8 vector",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 3,
Data: &schemapb.VectorField_Int8Vector{Int8Vector: []byte{1, 2, 3, 4, 5, 6}},
}}},
}},
expected: 2,
},
{
name: "vector array",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Data: []*schemapb.VectorField{nil, nil},
}},
}}},
}},
expected: 2,
},
{
name: "geometry field",
result: &internalpb.RetrieveResults{FieldsData: []*schemapb.FieldData{
{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_GeometryData{GeometryData: &schemapb.GeometryArray{Data: [][]byte{{1}, {2}}}},
}}},
}},
expected: 2,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
assert.Equal(t, tt.expected, getRowCount(tt.result))
})
}
}
// =========================================================================
// calcFieldElementSize: remaining branches (geometry, array, mol, vectorArray)
// =========================================================================
func TestCalcFieldElementSize_RemainingBranches(t *testing.T) {
tests := []struct {
name string
fd *schemapb.FieldData
rowIdx int
expected int64
}{
{
name: "geometry_data",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_GeometryData{GeometryData: &schemapb.GeometryArray{Data: [][]byte{{1, 2, 3}}}},
}}},
expected: 3,
},
{
name: "geometry_wkt_data",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_GeometryWktData{GeometryWktData: &schemapb.GeometryWktArray{Data: []string{"POINT(0 0)"}}},
}}},
expected: 10,
},
{
name: "mol_data",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_MolData{MolData: &schemapb.MolArray{Data: [][]byte{{0xAA, 0xBB}}}},
}}},
expected: 2,
},
{
name: "array_data",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_ArrayData{ArrayData: &schemapb.ArrayArray{
Data: []*schemapb.ScalarField{
{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{1, 2}}}},
},
ElementType: schemapb.DataType_Int32,
}},
}}},
expected: 10, // proto.Size of the ScalarField
},
{
name: "bfloat16_vector",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 4,
Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{0, 0, 0, 0, 0, 0, 0, 0}},
}}},
expected: 8, // 4*2
},
{
name: "vector_array",
fd: &schemapb.FieldData{Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Data: []*schemapb.VectorField{
{Dim: 2, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2}}}},
},
}},
}}},
},
{
name: "nil_field",
fd: &schemapb.FieldData{},
expected: 0,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := calcFieldElementSize(tt.fd, tt.rowIdx)
switch tt.name {
case "vector_array":
assert.True(t, result > 0, "vector_array size should be > 0")
case "array_data":
assert.True(t, result > 0, "array_data size should be > 0")
default:
assert.Equal(t, tt.expected, result)
}
})
}
}
// =========================================================================
// sliceFieldData / rangeSliceFieldData: vector + structArray branches
// =========================================================================
func TestSliceFieldData_Vector(t *testing.T) {
fd := &schemapb.FieldData{
Type: schemapb.DataType_FloatVector, FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5, 6}}},
}},
}
sliced := requireSliceFieldData(t, fd, []int{2, 0})
assert.Equal(t, []float32{5, 6, 1, 2}, sliced.GetVectors().GetFloatVector().GetData())
}
func TestSliceFieldData_NilOrEmpty(t *testing.T) {
sliced := requireSliceFieldData(t, nil, []int{0})
assert.Nil(t, sliced)
sliced = requireSliceFieldData(t, &schemapb.FieldData{}, nil)
assert.Nil(t, sliced)
}
func TestSliceFieldData_ValidData(t *testing.T) {
fd := &schemapb.FieldData{
Type: schemapb.DataType_Int64, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{10, 20, 30}}},
}},
ValidData: []bool{true, false, true},
}
sliced := requireSliceFieldData(t, fd, []int{2, 0})
assert.Equal(t, []bool{true, true}, sliced.GetValidData())
}
func TestRangeSliceFieldData_ValidData(t *testing.T) {
fd := &schemapb.FieldData{
Type: schemapb.DataType_Int64, FieldId: 100,
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{10, 20, 30}}},
}},
ValidData: []bool{true, false, true},
}
sliced := requireRangeSliceFieldData(t, fd, 0, 2)
assert.Equal(t, []bool{true, false}, sliced.GetValidData())
}
// =========================================================================
// sliceScalarField: remaining branches (geometry, geometryWkt)
// =========================================================================
func TestSliceScalarField_GeometryData(t *testing.T) {
sf := &schemapb.ScalarField{
Data: &schemapb.ScalarField_GeometryData{GeometryData: &schemapb.GeometryArray{Data: [][]byte{{1}, {2}, {3}}}},
}
sliced := sliceScalarField(sf, []int{2, 0})
assert.Equal(t, [][]byte{{3}, {1}}, sliced.GetGeometryData().GetData())
}
func TestSliceScalarField_GeometryWktData(t *testing.T) {
sf := &schemapb.ScalarField{
Data: &schemapb.ScalarField_GeometryWktData{GeometryWktData: &schemapb.GeometryWktArray{Data: []string{"P1", "P2", "P3"}}},
}
sliced := sliceScalarField(sf, []int{1})
assert.Equal(t, []string{"P2"}, sliced.GetGeometryWktData().GetData())
}
// =========================================================================
// rangeSliceScalarField: remaining branches (geometry)
// =========================================================================
func TestRangeSliceScalarField_GeometryData(t *testing.T) {
sf := &schemapb.ScalarField{
Data: &schemapb.ScalarField_GeometryData{GeometryData: &schemapb.GeometryArray{Data: [][]byte{{1}, {2}, {3}}}},
}
sliced := rangeSliceScalarField(sf, 0, 2)
assert.Equal(t, [][]byte{{1}, {2}}, sliced.GetGeometryData().GetData())
}
// =========================================================================
// sliceVectorField / rangeSliceVectorField: nullable compact mode branches
// =========================================================================
func TestSliceVectorField_NullableCompact(t *testing.T) {
// 3 logical rows, but only 2 valid rows in compact data
vf := &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4}}}, // 2 valid rows
}
validData := []bool{true, false, true} // row0=valid(data0), row1=null, row2=valid(data1)
sliced := requireSliceVectorField(t, vf, []int{2, 0}, validData)
assert.Equal(t, []float32{3, 4, 1, 2}, sliced.GetFloatVector().GetData())
}
func TestSliceVectorField_NullableCompact_BinaryVector(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 16, // 2 bytes per row
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{0xAA, 0xBB, 0xCC, 0xDD}}, // 2 valid rows
}
validData := []bool{true, false, true}
sliced := requireSliceVectorField(t, vf, []int{2, 0}, validData)
assert.Equal(t, []byte{0xCC, 0xDD, 0xAA, 0xBB}, sliced.GetBinaryVector())
}
func TestSliceVectorField_NullableCompact_Float16(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 2, // 4 bytes per row
Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{1, 2, 3, 4, 5, 6, 7, 8}},
}
validData := []bool{true, false, true}
sliced := requireSliceVectorField(t, vf, []int{0}, validData)
assert.Equal(t, []byte{1, 2, 3, 4}, sliced.GetFloat16Vector())
}
func TestSliceVectorField_NullableCompact_BFloat16(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{10, 20, 30, 40, 50, 60, 70, 80}},
}
validData := []bool{true, false, true}
sliced := requireSliceVectorField(t, vf, []int{2}, validData)
assert.Equal(t, []byte{50, 60, 70, 80}, sliced.GetBfloat16Vector())
}
func TestSliceVectorField_NullableCompact_Int8(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 3,
Data: &schemapb.VectorField_Int8Vector{Int8Vector: []byte{1, 2, 3, 4, 5, 6}},
}
validData := []bool{true, false, true}
sliced := requireSliceVectorField(t, vf, []int{2}, validData)
assert.Equal(t, []byte{4, 5, 6}, sliced.GetInt8Vector())
}
func TestSliceVectorField_NullableCompact_Sparse(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 100,
Data: &schemapb.VectorField_SparseFloatVector{SparseFloatVector: &schemapb.SparseFloatArray{
Contents: [][]byte{{0x01}, {0x02}}, // 2 valid rows
Dim: 100,
}},
}
validData := []bool{true, false, true}
sliced := requireSliceVectorField(t, vf, []int{2}, validData)
assert.Equal(t, [][]byte{{0x02}}, sliced.GetSparseFloatVector().GetContents())
}
func TestSliceVectorField_NullableRowDense_VectorArray_SelectValidRow(t *testing.T) {
dim := int64(2)
v1 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2}}}}
empty := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{}}}
v2 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{3, 4}}}}
vf := &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: dim,
Data: []*schemapb.VectorField{v1, empty, v2},
ElementType: schemapb.DataType_FloatVector,
}},
}
validData := []bool{true, false, true}
sliced := requireSliceVectorField(t, vf, []int{2}, validData)
assert.Len(t, sliced.GetVectorArray().GetData(), 1)
assert.Equal(t, []float32{3, 4}, sliced.GetVectorArray().GetData()[0].GetFloatVector().GetData())
}
func TestSliceVectorField_NullableRowDense_VectorArray_PreservesNullPlaceholders(t *testing.T) {
dim := int64(2)
v1 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2}}}}
empty := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{}}}
v3 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{3, 4}}}}
vf := &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: dim,
Data: []*schemapb.VectorField{v1, empty, v3},
ElementType: schemapb.DataType_FloatVector,
}},
}
validData := []bool{true, false, true}
sliced := requireSliceVectorField(t, vf, []int{2, 1, 0}, validData)
require.Len(t, sliced.GetVectorArray().GetData(), 3)
assert.Equal(t, []float32{3, 4}, sliced.GetVectorArray().GetData()[0].GetFloatVector().GetData())
assert.Empty(t, sliced.GetVectorArray().GetData()[1].GetFloatVector().GetData())
assert.Equal(t, []float32{1, 2}, sliced.GetVectorArray().GetData()[2].GetFloatVector().GetData())
}
func TestSliceVectorField_NullableRowDense_VectorArray_AllNullNoData(t *testing.T) {
dim := int64(2)
vf := &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: dim,
ElementType: schemapb.DataType_FloatVector,
}},
}
validData := []bool{false}
sliced := requireSliceVectorField(t, vf, []int{0}, validData)
require.Len(t, sliced.GetVectorArray().GetData(), 1)
assert.Empty(t, sliced.GetVectorArray().GetData()[0].GetFloatVector().GetData())
}
func TestSliceVectorField_NullableRowDense_VectorArray_UnsupportedElementTypeReturnsError(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: 2,
ElementType: schemapb.DataType_None,
}},
}
sliced, err := sliceVectorField(vf, []int{0}, []bool{false})
require.Error(t, err)
require.Nil(t, sliced)
assert.Contains(t, err.Error(), "unsupported ArrayOfVector element type")
}
func TestSliceVectorField_NullableCompact_SkipNull(t *testing.T) {
// Select a null row — it should be skipped
vf := &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2}}},
}
validData := []bool{true, false}
sliced := requireSliceVectorField(t, vf, []int{1}, validData) // row 1 is null
assert.Empty(t, sliced.GetFloatVector().GetData())
}
func TestRangeSliceVectorField_NullableCompact(t *testing.T) {
// 3 logical rows: valid, null, valid → 2 data entries in compact mode
vf := &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4}}},
}
validData := []bool{true, false, true}
// Range [0, 2) = logical row 0 and 1. Row 0 is valid(data0), row 1 is null → only data0
sliced := requireRangeSliceVectorField(t, vf, 0, 2, validData)
assert.Equal(t, []float32{1, 2}, sliced.GetFloatVector().GetData())
}
func TestRangeSliceVectorField_NullableCompact_BinaryVector(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 16,
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{0xAA, 0xBB, 0xCC, 0xDD}},
}
validData := []bool{true, false, true}
sliced := requireRangeSliceVectorField(t, vf, 0, 2, validData)
assert.Equal(t, []byte{0xAA, 0xBB}, sliced.GetBinaryVector())
}
func TestRangeSliceVectorField_NullableCompact_Float16(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{1, 2, 3, 4, 5, 6, 7, 8}},
}
validData := []bool{true, false, true}
sliced := requireRangeSliceVectorField(t, vf, 2, 3, validData)
assert.Equal(t, []byte{5, 6, 7, 8}, sliced.GetFloat16Vector())
}
func TestRangeSliceVectorField_NullableCompact_BFloat16(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{10, 20, 30, 40, 50, 60, 70, 80}},
}
validData := []bool{true, false, true}
sliced := requireRangeSliceVectorField(t, vf, 0, 1, validData)
assert.Equal(t, []byte{10, 20, 30, 40}, sliced.GetBfloat16Vector())
}
func TestRangeSliceVectorField_NullableCompact_Int8(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 3,
Data: &schemapb.VectorField_Int8Vector{Int8Vector: []byte{1, 2, 3, 4, 5, 6}},
}
validData := []bool{true, false, true}
sliced := requireRangeSliceVectorField(t, vf, 2, 3, validData)
assert.Equal(t, []byte{4, 5, 6}, sliced.GetInt8Vector())
}
func TestRangeSliceVectorField_NullableCompact_Sparse(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 100,
Data: &schemapb.VectorField_SparseFloatVector{SparseFloatVector: &schemapb.SparseFloatArray{
Contents: [][]byte{{0x01}, {0x02}},
Dim: 100,
}},
}
validData := []bool{true, false, true}
sliced := requireRangeSliceVectorField(t, vf, 2, 3, validData)
assert.Equal(t, [][]byte{{0x02}}, sliced.GetSparseFloatVector().GetContents())
}
func TestRangeSliceVectorField_NullableRowDense_VectorArray_SelectPrefix(t *testing.T) {
dim := int64(2)
v1 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2}}}}
empty := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{}}}
v2 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{3, 4}}}}
vf := &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: dim,
Data: []*schemapb.VectorField{v1, empty, v2},
ElementType: schemapb.DataType_FloatVector,
}},
}
validData := []bool{true, false, true}
sliced := requireRangeSliceVectorField(t, vf, 0, 1, validData)
assert.Len(t, sliced.GetVectorArray().GetData(), 1)
assert.Equal(t, []float32{1, 2}, sliced.GetVectorArray().GetData()[0].GetFloatVector().GetData())
}
func TestRangeSliceVectorField_NullableRowDense_VectorArray_PreservesNullPlaceholders(t *testing.T) {
dim := int64(2)
v1 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2}}}}
empty := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{}}}
v3 := &schemapb.VectorField{Dim: dim, Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{3, 4}}}}
vf := &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: dim,
Data: []*schemapb.VectorField{v1, empty, v3},
ElementType: schemapb.DataType_FloatVector,
}},
}
validData := []bool{true, false, true}
sliced := requireRangeSliceVectorField(t, vf, 2, 3, validData)
require.Len(t, sliced.GetVectorArray().GetData(), 1)
assert.Equal(t, []float32{3, 4}, sliced.GetVectorArray().GetData()[0].GetFloatVector().GetData())
sliced = requireRangeSliceVectorField(t, vf, 0, 3, validData)
require.Len(t, sliced.GetVectorArray().GetData(), 3)
assert.Equal(t, []float32{1, 2}, sliced.GetVectorArray().GetData()[0].GetFloatVector().GetData())
assert.Empty(t, sliced.GetVectorArray().GetData()[1].GetFloatVector().GetData())
assert.Equal(t, []float32{3, 4}, sliced.GetVectorArray().GetData()[2].GetFloatVector().GetData())
}
func TestRangeSliceVectorField_NullableRowDense_VectorArray_AllNullNoData(t *testing.T) {
dim := int64(2)
vf := &schemapb.VectorField{
Dim: dim,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: dim,
ElementType: schemapb.DataType_FloatVector,
}},
}
validData := []bool{false, false}
sliced := requireRangeSliceVectorField(t, vf, 0, 2, validData)
require.Len(t, sliced.GetVectorArray().GetData(), 2)
assert.Empty(t, sliced.GetVectorArray().GetData()[0].GetFloatVector().GetData())
assert.Empty(t, sliced.GetVectorArray().GetData()[1].GetFloatVector().GetData())
}
func TestRangeSliceVectorField_NullableRowDense_VectorArray_UnsupportedElementTypeReturnsError(t *testing.T) {
vf := &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
Dim: 2,
ElementType: schemapb.DataType_None,
}},
}
sliced, err := rangeSliceVectorField(vf, 0, 1, []bool{false})
require.Error(t, err)
require.Nil(t, sliced)
assert.Contains(t, err.Error(), "unsupported ArrayOfVector element type")
}
// =========================================================================
// buildMergedVectorField: nullable compact mode — null row skip
// =========================================================================
func TestBuildMergedVectorField_NullableCompact_FloatVector(t *testing.T) {
// FloatVector dim=2, nullable: row0=valid(compact idx 0), row1=null, row2=valid(compact idx 1)
// Compact data: [1,2, 3,4] — 2 valid rows only
r1 := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2, 3}}}},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_FloatVector, FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4}}},
}},
ValidData: []bool{true, false, true},
},
},
}
// Include the null row (row1) in selectedRows to exercise the di < 0 skip path
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 2}, {resultIdx: 0, rowIdx: 1}, {resultIdx: 0, rowIdx: 0}}
merged, err := buildMergedRetrieveResults([]*internalpb.RetrieveResults{r1}, selectedRows, makeNullableSchema(100, schemapb.DataType_FloatVector, 2))
require.NoError(t, err)
// row2 → compact idx 1 → [3,4], row1 → null (skipped), row0 → compact idx 0 → [1,2]
assert.Equal(t, []float32{3, 4, 1, 2}, merged.FieldsData[0].GetVectors().GetFloatVector().GetData())
assert.Equal(t, []bool{true, false, true}, merged.FieldsData[0].ValidData)
}
// =========================================================================
// getFieldValue: remaining branches (string, out-of-range)
// =========================================================================
func TestGetFieldValue_String(t *testing.T) {
fd := &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"hello", "world"}}},
}}}
val, isNull := getFieldValue(fd, 1)
assert.False(t, isNull)
assert.Equal(t, "world", val)
}
func TestGetFieldValue_OutOfRange(t *testing.T) {
fd := &schemapb.FieldData{Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{1}}},
}}}
_, isNull := getFieldValue(fd, 5) // out of range
assert.True(t, isNull)
}
// =========================================================================
// rangeSliceRetrieveResults: element-level metadata propagation
// =========================================================================
func TestRangeSliceRetrieveResults_EmptyRange(t *testing.T) {
result := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
}
sliced := requireRangeSliceRetrieveResults(t, result, 1, 1) // start == end
assert.Nil(t, sliced.GetIds())
}
func TestRangeSliceRetrieveResults_ElementLevel(t *testing.T) {
result := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2, 3}}}},
ElementLevel: true,
ElementIndices: []*internalpb.ElementIndices{
{Indices: []int32{0}},
{Indices: []int32{1, 2}},
{Indices: []int32{3}},
},
FieldsData: []*schemapb.FieldData{
makeInt64Field(100, "val", []int64{10, 20, 30}),
},
}
sliced := requireRangeSliceRetrieveResults(t, result, 1, 3)
assert.True(t, sliced.GetElementLevel())
assert.Len(t, sliced.GetElementIndices(), 2)
assert.Equal(t, []int32{1, 2}, sliced.GetElementIndices()[0].GetIndices())
}
func TestValidateElementLevelConsistency_InconsistentFlag(t *testing.T) {
r1 := &internalpb.RetrieveResults{
ElementLevel: true,
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}},
ElementIndices: []*internalpb.ElementIndices{{Indices: []int32{0}}},
}
r2 := &internalpb.RetrieveResults{
ElementLevel: false,
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{2}}}},
}
err := validateElementLevelConsistency([]*internalpb.RetrieveResults{r1, r2}, nil)
assert.Error(t, err)
assert.Contains(t, err.Error(), "inconsistent element-level flag")
}
func TestValidateElementLevelConsistency_LengthMismatch(t *testing.T) {
r := &internalpb.RetrieveResults{
Ids: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1, 2}}}},
ElementLevel: true,
ElementIndices: []*internalpb.ElementIndices{{Indices: []int32{0}}}, // 1 != 2 ids
}
err := validateElementLevelConsistency([]*internalpb.RetrieveResults{r}, nil)
assert.Error(t, err)
assert.Contains(t, err.Error(), "element_indices length")
}
// =========================================================================
// Nullable vector / scalar merge correctness tests
// =========================================================================
// makeTestIntIDs is a helper to build IDs for test results.
func makeTestIntIDs(ids ...int64) *schemapb.IDs {
return &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: ids}}}
}
// makeTestSparseVec builds a knowhere sparse vector binary: [dim uint32][val float32], little-endian.
func makeTestSparseVec(dim uint32, val float32) []byte {
b := make([]byte, 8)
binary.LittleEndian.PutUint32(b[0:4], dim)
binary.LittleEndian.PutUint32(b[4:8], math.Float32bits(val))
return b
}
// TestBuildMergedVectorField_NullableSparseVector_AllNull_EmptyValidData validates Bug 1 fix:
// when a nullable sparse vector field has empty ValidData + empty Contents, merge must not panic,
// and the merged output must correctly mark those rows as null (ValidData=false).
func TestBuildMergedVectorField_NullableSparse_RejectsEmptyValidData(t *testing.T) {
sparseSchema := makeNullableSchema(100, schemapb.DataType_SparseFloatVector, 0)
// rA: 1 row, has valid data
rA := &internalpb.RetrieveResults{
Ids: makeTestIntIDs(1),
FieldsData: []*schemapb.FieldData{{
FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Data: &schemapb.VectorField_SparseFloatVector{
SparseFloatVector: &schemapb.SparseFloatArray{
Contents: [][]byte{makeTestSparseVec(1, 0.5)},
},
},
}},
ValidData: []bool{true},
}},
}
// rB: 1 row, ValidData absent, Contents absent (segcore get_vector early return).
// This is now a contract violation: nullable field with numRows > 0 must have ValidData.
// The old code silently treated this as "all null"; the new code fails fast.
rB := &internalpb.RetrieveResults{
Ids: makeTestIntIDs(2),
FieldsData: []*schemapb.FieldData{{
FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Data: &schemapb.VectorField_SparseFloatVector{
SparseFloatVector: &schemapb.SparseFloatArray{Contents: nil},
},
}},
ValidData: nil,
}},
}
results := []*internalpb.RetrieveResults{rA, rB}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}, {resultIdx: 1, rowIdx: 0}}
_, err := buildMergedRetrieveResults(results, selectedRows, sparseSchema)
require.Error(t, err)
assert.Contains(t, err.Error(), "empty ValidData")
}
// TestBuildMergedVectorField_NullableSparse_RejectsMultipleRowsWithEmptyValidData validates that
// multiple rows from a result with absent ValidData trigger fail-fast error.
func TestBuildMergedVectorField_NullableSparse_RejectsMultipleRowsWithEmptyValidData(t *testing.T) {
sparseSchema := makeNullableSchema(100, schemapb.DataType_SparseFloatVector, 0)
rA := &internalpb.RetrieveResults{
Ids: makeTestIntIDs(1),
FieldsData: []*schemapb.FieldData{{
FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Data: &schemapb.VectorField_SparseFloatVector{
SparseFloatVector: &schemapb.SparseFloatArray{
Contents: [][]byte{makeTestSparseVec(1, 0.1)},
},
},
}},
ValidData: []bool{true},
}},
}
rB := &internalpb.RetrieveResults{
Ids: makeTestIntIDs(2, 3),
FieldsData: []*schemapb.FieldData{{
FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Data: &schemapb.VectorField_SparseFloatVector{
SparseFloatVector: &schemapb.SparseFloatArray{Contents: nil},
},
}},
ValidData: nil,
}},
}
results := []*internalpb.RetrieveResults{rA, rB}
selectedRows := []rowRef{
{resultIdx: 0, rowIdx: 0},
{resultIdx: 1, rowIdx: 0},
{resultIdx: 1, rowIdx: 1},
}
// rB has 2 rows but no ValidData → contract violation, must error
_, err := buildMergedRetrieveResults(results, selectedRows, sparseSchema)
require.Error(t, err)
assert.Contains(t, err.Error(), "empty ValidData")
}
// TestBuildMergedVectorField_NullableSparseVector_AllNullWithValidData validates compact path
// when ValidData is present and all false (explicit all-null).
func TestBuildMergedVectorField_NullableSparseVector_AllNullWithValidData(t *testing.T) {
sparseSchema := makeNullableSchema(100, schemapb.DataType_SparseFloatVector, 0)
r := &internalpb.RetrieveResults{
Ids: makeTestIntIDs(1, 2),
FieldsData: []*schemapb.FieldData{{
FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Data: &schemapb.VectorField_SparseFloatVector{
SparseFloatVector: &schemapb.SparseFloatArray{Contents: nil},
},
}},
ValidData: []bool{false, false},
}},
}
results := []*internalpb.RetrieveResults{r}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}, {resultIdx: 0, rowIdx: 1}}
merged, err := buildMergedRetrieveResults(results, selectedRows, sparseSchema)
require.NoError(t, err)
assert.Equal(t, []bool{false, false}, merged.FieldsData[0].ValidData)
assert.Empty(t, merged.FieldsData[0].GetVectors().GetSparseFloatVector().GetContents())
}
// TestBuildMergedVectorField_NullableSparseVector_Mixed validates compact index mapping when
// some rows are valid and some are null within the same result.
func TestBuildMergedVectorField_NullableSparseVector_Mixed(t *testing.T) {
sparseSchema := makeNullableSchema(100, schemapb.DataType_SparseFloatVector, 0)
// 3 rows: ValidData=[true,false,true], Contents=[vec0, vec2] (compact)
r := &internalpb.RetrieveResults{
Ids: makeTestIntIDs(1, 2, 3),
FieldsData: []*schemapb.FieldData{{
FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Data: &schemapb.VectorField_SparseFloatVector{
SparseFloatVector: &schemapb.SparseFloatArray{
Contents: [][]byte{makeTestSparseVec(1, 0.1), makeTestSparseVec(3, 0.3)},
},
},
}},
ValidData: []bool{true, false, true},
}},
}
results := []*internalpb.RetrieveResults{r}
// Select row 0 (valid) and row 1 (null)
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}, {resultIdx: 0, rowIdx: 1}}
merged, err := buildMergedRetrieveResults(results, selectedRows, sparseSchema)
require.NoError(t, err)
assert.Equal(t, []bool{true, false}, merged.FieldsData[0].ValidData)
// Only 1 non-null row in output
assert.Len(t, merged.FieldsData[0].GetVectors().GetSparseFloatVector().GetContents(), 1)
}
// TestBuildMergedFieldData_NullableScalar_EmptyValidData_FailFast validates that
// a nullable field with rows but absent ValidData triggers a fail-fast error,
// not a silent fallback to all-null. This catches segcore contract violations
// early instead of letting them propagate to downstream panics.
func TestBuildMergedVectorField_NullableVector_RejectsEmptyValidData(t *testing.T) {
// Note: buildCompactIndices only validates vector fields. For scalar fields,
// the ValidData check is still lenient (scalar fields are not compacted).
// This test uses a vector schema to trigger the vector-path validation.
sparseSchema := makeNullableSchema(100, schemapb.DataType_SparseFloatVector, 0)
rA := &internalpb.RetrieveResults{
Ids: makeTestIntIDs(1),
FieldsData: []*schemapb.FieldData{{
FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Data: &schemapb.VectorField_SparseFloatVector{
SparseFloatVector: &schemapb.SparseFloatArray{
Contents: [][]byte{makeTestSparseVec(1, 0.5)},
},
},
}},
ValidData: []bool{true},
}},
}
// rB: ValidData absent → contract violation for nullable vector
rB := &internalpb.RetrieveResults{
Ids: makeTestIntIDs(2),
FieldsData: []*schemapb.FieldData{{
FieldId: 100,
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
Data: &schemapb.VectorField_SparseFloatVector{
SparseFloatVector: &schemapb.SparseFloatArray{Contents: nil},
},
}},
ValidData: nil,
}},
}
results := []*internalpb.RetrieveResults{rA, rB}
selectedRows := []rowRef{{resultIdx: 0, rowIdx: 0}, {resultIdx: 1, rowIdx: 0}}
_, err := buildMergedRetrieveResults(results, selectedRows, sparseSchema)
require.Error(t, err)
assert.Contains(t, err.Error(), "empty ValidData")
}