diff --git a/internal/storage/primary_keys.go b/internal/storage/primary_keys.go new file mode 100644 index 0000000000..4f6be2e3a4 --- /dev/null +++ b/internal/storage/primary_keys.go @@ -0,0 +1,158 @@ +// 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 storage + +import ( + "github.com/samber/lo" + + "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" + "github.com/milvus-io/milvus/pkg/util/merr" +) + +// PrimaryKeys is the interface holding a slice of PrimaryKey +type PrimaryKeys interface { + Append(pks ...PrimaryKey) error + MustAppend(pks ...PrimaryKey) + Get(idx int) PrimaryKey + Type() schemapb.DataType + Size() int64 + Len() int + MustMerge(pks PrimaryKeys) +} + +type Int64PrimaryKeys struct { + values []int64 +} + +func NewInt64PrimaryKeys(cap int) *Int64PrimaryKeys { + return &Int64PrimaryKeys{values: make([]int64, 0, cap)} +} + +func (pks *Int64PrimaryKeys) AppendRaw(values ...int64) { + pks.values = append(pks.values, values...) +} + +func (pks *Int64PrimaryKeys) Append(values ...PrimaryKey) error { + iValues := make([]int64, 0, len(values)) + for _, pk := range values { + iPk, ok := pk.(*Int64PrimaryKey) + if !ok { + return merr.WrapErrParameterInvalid("Int64PrimaryKey", "non-int64 pk") + } + iValues = append(iValues, iPk.Value) + } + + pks.AppendRaw(iValues...) + return nil +} + +func (pks *Int64PrimaryKeys) MustAppend(values ...PrimaryKey) { + err := pks.Append(values...) + if err != nil { + panic(err) + } +} + +func (pks *Int64PrimaryKeys) Get(idx int) PrimaryKey { + return NewInt64PrimaryKey(pks.values[idx]) +} + +func (pks *Int64PrimaryKeys) Type() schemapb.DataType { + return schemapb.DataType_Int64 +} + +func (pks *Int64PrimaryKeys) Len() int { + return len(pks.values) +} + +func (pks *Int64PrimaryKeys) Size() int64 { + return int64(pks.Len()) * 8 +} + +func (pks *Int64PrimaryKeys) MustMerge(another PrimaryKeys) { + aPks, ok := another.(*Int64PrimaryKeys) + if !ok { + panic("cannot merge different kind of pks") + } + + pks.values = append(pks.values, aPks.values...) +} + +type VarcharPrimaryKeys struct { + values []string + size int64 +} + +func NewVarcharPrimaryKeys(cap int) *VarcharPrimaryKeys { + return &VarcharPrimaryKeys{ + values: make([]string, 0, cap), + } +} + +func (pks *VarcharPrimaryKeys) AppendRaw(values ...string) { + pks.values = append(pks.values, values...) + lo.ForEach(values, func(str string, _ int) { + pks.size += int64(len(str)) + 16 + }) +} + +func (pks *VarcharPrimaryKeys) Append(values ...PrimaryKey) error { + sValues := make([]string, 0, len(values)) + for _, pk := range values { + iPk, ok := pk.(*VarCharPrimaryKey) + if !ok { + return merr.WrapErrParameterInvalid("Int64PrimaryKey", "non-int64 pk") + } + sValues = append(sValues, iPk.Value) + } + + pks.AppendRaw(sValues...) + return nil +} + +func (pks *VarcharPrimaryKeys) MustAppend(values ...PrimaryKey) { + err := pks.Append(values...) + if err != nil { + panic(err) + } +} + +func (pks *VarcharPrimaryKeys) Get(idx int) PrimaryKey { + return NewVarCharPrimaryKey(pks.values[idx]) +} + +func (pks *VarcharPrimaryKeys) Type() schemapb.DataType { + return schemapb.DataType_VarChar +} + +func (pks *VarcharPrimaryKeys) Len() int { + return len(pks.values) +} + +func (pks *VarcharPrimaryKeys) Size() int64 { + return pks.size +} + +func (pks *VarcharPrimaryKeys) MustMerge(another PrimaryKeys) { + aPks, ok := another.(*VarcharPrimaryKeys) + if !ok { + panic("cannot merge different kind of pks") + } + + pks.values = append(pks.values, aPks.values...) + pks.size += aPks.size +} diff --git a/internal/storage/primary_keys_test.go b/internal/storage/primary_keys_test.go new file mode 100644 index 0000000000..486b7f8879 --- /dev/null +++ b/internal/storage/primary_keys_test.go @@ -0,0 +1,137 @@ +// 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 storage + +import ( + "testing" + + "github.com/stretchr/testify/suite" + + "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" +) + +type PrimaryKeysSuite struct { + suite.Suite +} + +func (s *PrimaryKeysSuite) TestAppend() { + s.Run("IntAppend", func() { + intPks := NewInt64PrimaryKeys(10) + s.Equal(schemapb.DataType_Int64, intPks.Type()) + s.EqualValues(0, intPks.Len()) + s.EqualValues(0, intPks.Size()) + + err := intPks.Append(NewInt64PrimaryKey(1)) + s.NoError(err) + s.EqualValues(1, intPks.Len()) + s.EqualValues(8, intPks.Size()) + + val := intPks.Get(0) + pk, ok := val.(*Int64PrimaryKey) + s.Require().True(ok) + s.EqualValues(1, pk.Value) + + err = intPks.Append(NewVarCharPrimaryKey("1")) + s.Error(err) + }) + + s.Run("VarcharAppend", func() { + strPks := NewVarcharPrimaryKeys(10) + s.Equal(schemapb.DataType_VarChar, strPks.Type()) + s.EqualValues(0, strPks.Len()) + s.EqualValues(0, strPks.Size()) + + err := strPks.Append(NewVarCharPrimaryKey("1")) + s.NoError(err) + s.EqualValues(1, strPks.Len()) + s.EqualValues(17, strPks.Size()) + val := strPks.Get(0) + pk, ok := val.(*VarCharPrimaryKey) + s.Require().True(ok) + s.EqualValues("1", pk.Value) + + err = strPks.Append(NewInt64PrimaryKey(1)) + s.Error(err) + }) + + s.Run("IntMustAppend", func() { + intPks := NewInt64PrimaryKeys(10) + + s.NotPanics(func() { + intPks.MustAppend(NewInt64PrimaryKey(1)) + }) + s.Panics(func() { + intPks.MustAppend(NewVarCharPrimaryKey("1")) + }) + }) + + s.Run("VarcharMustAppend", func() { + strPks := NewVarcharPrimaryKeys(10) + + s.NotPanics(func() { + strPks.MustAppend(NewVarCharPrimaryKey("1")) + }) + s.Panics(func() { + strPks.MustAppend(NewInt64PrimaryKey(1)) + }) + }) +} + +func (s *PrimaryKeysSuite) TestMustMerge() { + s.Run("IntPksMustMerge", func() { + intPks := NewInt64PrimaryKeys(10) + intPks.AppendRaw(1, 2, 3) + + anotherPks := NewInt64PrimaryKeys(10) + anotherPks.AppendRaw(4, 5, 6) + + strPks := NewVarcharPrimaryKeys(10) + strPks.AppendRaw("1", "2", "3") + + s.NotPanics(func() { + intPks.MustMerge(anotherPks) + + s.Equal(6, intPks.Len()) + }) + + s.Panics(func() { + intPks.MustMerge(strPks) + }) + }) + + s.Run("StrPksMustMerge", func() { + strPks := NewVarcharPrimaryKeys(10) + strPks.AppendRaw("1", "2", "3") + intPks := NewInt64PrimaryKeys(10) + intPks.AppendRaw(1, 2, 3) + anotherPks := NewVarcharPrimaryKeys(10) + anotherPks.AppendRaw("4", "5", "6") + + s.NotPanics(func() { + strPks.MustMerge(anotherPks) + s.Equal(6, strPks.Len()) + }) + + s.Panics(func() { + strPks.MustMerge(intPks) + }) + }) +} + +func TestPrimaryKeys(t *testing.T) { + suite.Run(t, new(PrimaryKeysSuite)) +}