From 37464a281dfd4aa9b4aaa2d4e109747ba8c13e42 Mon Sep 17 00:00:00 2001 From: congqixia Date: Tue, 21 Jun 2022 11:38:12 +0800 Subject: [PATCH] Fix datasorter swap vector logic (#17661) Signed-off-by: Congqi Xia --- internal/storage/data_sorter.go | 10 ++-- internal/storage/data_sorter_test.go | 70 ++++++++++++++-------------- 2 files changed, 41 insertions(+), 39 deletions(-) diff --git a/internal/storage/data_sorter.go b/internal/storage/data_sorter.go index c98eb4d8ab..db9be80723 100644 --- a/internal/storage/data_sorter.go +++ b/internal/storage/data_sorter.go @@ -83,14 +83,16 @@ func (ds *DataSorter) Swap(i, j int) { case schemapb.DataType_BinaryVector: data := singleData.(*BinaryVectorFieldData).Data dim := singleData.(*BinaryVectorFieldData).Dim - for i := 0; i < dim/8; i++ { - data[i], data[i+dim/8] = data[i+dim/8], data[i] + // dim for binary vector must be multiplier of 8, simple verion for swapping: + steps := dim / 8 + for idx := 0; idx < steps; idx++ { + data[i*steps+idx], data[j*steps+idx] = data[j*steps+idx], data[i*steps+idx] } case schemapb.DataType_FloatVector: data := singleData.(*FloatVectorFieldData).Data dim := singleData.(*FloatVectorFieldData).Dim - for i := 0; i < dim; i++ { - data[i], data[i+dim] = data[i+dim], data[i] + for idx := 0; idx < dim; idx++ { + data[i*dim+idx], data[j*dim+idx] = data[j*dim+idx], data[i*dim+idx] } default: errMsg := "undefined data type " + string(field.DataType) diff --git a/internal/storage/data_sorter_test.go b/internal/storage/data_sorter_test.go index af6713c45b..7705f19ede 100644 --- a/internal/storage/data_sorter_test.go +++ b/internal/storage/data_sorter_test.go @@ -128,53 +128,53 @@ func TestDataSorter(t *testing.T) { insertDataFirst := &InsertData{ Data: map[int64]FieldData{ 0: &Int64FieldData{ - NumRows: []int64{2}, - Data: []int64{6, 4}, + NumRows: []int64{3}, + Data: []int64{3, 4, 2}, }, 1: &Int64FieldData{ - NumRows: []int64{2}, - Data: []int64{3, 4}, + NumRows: []int64{3}, + Data: []int64{3, 4, 5}, }, 100: &BoolFieldData{ - NumRows: []int64{2}, - Data: []bool{true, false}, + NumRows: []int64{3}, + Data: []bool{true, false, true}, }, 101: &Int8FieldData{ - NumRows: []int64{2}, - Data: []int8{3, 4}, + NumRows: []int64{3}, + Data: []int8{3, 4, 5}, }, 102: &Int16FieldData{ - NumRows: []int64{2}, - Data: []int16{3, 4}, + NumRows: []int64{3}, + Data: []int16{3, 4, 5}, }, 103: &Int32FieldData{ - NumRows: []int64{2}, - Data: []int32{3, 4}, + NumRows: []int64{3}, + Data: []int32{3, 4, 5}, }, 104: &Int64FieldData{ - NumRows: []int64{2}, - Data: []int64{3, 4}, + NumRows: []int64{3}, + Data: []int64{3, 4, 5}, }, 105: &FloatFieldData{ - NumRows: []int64{2}, - Data: []float32{3, 4}, + NumRows: []int64{3}, + Data: []float32{3, 4, 5}, }, 106: &DoubleFieldData{ - NumRows: []int64{2}, - Data: []float64{3, 4}, + NumRows: []int64{3}, + Data: []float64{3, 4, 5}, }, 107: &StringFieldData{ NumRows: []int64{2}, - Data: []string{"3", "4"}, + Data: []string{"3", "4", "5"}, }, 108: &BinaryVectorFieldData{ - NumRows: []int64{2}, - Data: []byte{0, 255}, + NumRows: []int64{3}, + Data: []byte{0, 255, 128}, Dim: 8, }, 109: &FloatVectorFieldData{ - NumRows: []int64{2}, - Data: []float32{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}, + NumRows: []int64{3}, + Data: []float32{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23}, Dim: 8, }, }, @@ -226,18 +226,18 @@ func TestDataSorter(t *testing.T) { // } // } - assert.Equal(t, []int64{4, 6}, dataSorter.InsertData.Data[0].(*Int64FieldData).Data) - assert.Equal(t, []int64{4, 3}, dataSorter.InsertData.Data[1].(*Int64FieldData).Data) - assert.Equal(t, []bool{false, true}, dataSorter.InsertData.Data[100].(*BoolFieldData).Data) - assert.Equal(t, []int8{4, 3}, dataSorter.InsertData.Data[101].(*Int8FieldData).Data) - assert.Equal(t, []int16{4, 3}, dataSorter.InsertData.Data[102].(*Int16FieldData).Data) - assert.Equal(t, []int32{4, 3}, dataSorter.InsertData.Data[103].(*Int32FieldData).Data) - assert.Equal(t, []int64{4, 3}, dataSorter.InsertData.Data[104].(*Int64FieldData).Data) - assert.Equal(t, []float32{4, 3}, dataSorter.InsertData.Data[105].(*FloatFieldData).Data) - assert.Equal(t, []float64{4, 3}, dataSorter.InsertData.Data[106].(*DoubleFieldData).Data) - assert.Equal(t, []string{"4", "3"}, dataSorter.InsertData.Data[107].(*StringFieldData).Data) - assert.Equal(t, []byte{255, 0}, dataSorter.InsertData.Data[108].(*BinaryVectorFieldData).Data) - assert.Equal(t, []float32{8, 9, 10, 11, 12, 13, 14, 15, 0, 1, 2, 3, 4, 5, 6, 7}, dataSorter.InsertData.Data[109].(*FloatVectorFieldData).Data) + assert.Equal(t, []int64{2, 3, 4}, dataSorter.InsertData.Data[0].(*Int64FieldData).Data) + assert.Equal(t, []int64{5, 3, 4}, dataSorter.InsertData.Data[1].(*Int64FieldData).Data) + assert.Equal(t, []bool{true, true, false}, dataSorter.InsertData.Data[100].(*BoolFieldData).Data) + assert.Equal(t, []int8{5, 3, 4}, dataSorter.InsertData.Data[101].(*Int8FieldData).Data) + assert.Equal(t, []int16{5, 3, 4}, dataSorter.InsertData.Data[102].(*Int16FieldData).Data) + assert.Equal(t, []int32{5, 3, 4}, dataSorter.InsertData.Data[103].(*Int32FieldData).Data) + assert.Equal(t, []int64{5, 3, 4}, dataSorter.InsertData.Data[104].(*Int64FieldData).Data) + assert.Equal(t, []float32{5, 3, 4}, dataSorter.InsertData.Data[105].(*FloatFieldData).Data) + assert.Equal(t, []float64{5, 3, 4}, dataSorter.InsertData.Data[106].(*DoubleFieldData).Data) + assert.Equal(t, []string{"5", "3", "4"}, dataSorter.InsertData.Data[107].(*StringFieldData).Data) + assert.Equal(t, []byte{128, 0, 255}, dataSorter.InsertData.Data[108].(*BinaryVectorFieldData).Data) + assert.Equal(t, []float32{16, 17, 18, 19, 20, 21, 22, 23, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}, dataSorter.InsertData.Data[109].(*FloatVectorFieldData).Data) } func TestDataSorter_Len(t *testing.T) {