123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200 |
- // 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 integration
- import (
- "context"
- "time"
- "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
- "github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
- "github.com/milvus-io/milvus/pkg/util/testutils"
- )
- func (s *MiniClusterSuite) WaitForFlush(ctx context.Context, segIDs []int64, flushTs uint64, dbName, collectionName string) {
- flushed := func() bool {
- resp, err := s.Cluster.Proxy.GetFlushState(ctx, &milvuspb.GetFlushStateRequest{
- SegmentIDs: segIDs,
- FlushTs: flushTs,
- DbName: dbName,
- CollectionName: collectionName,
- })
- if err != nil {
- return false
- }
- return resp.GetFlushed()
- }
- for !flushed() {
- select {
- case <-ctx.Done():
- s.FailNow("failed to wait for flush until ctx done")
- return
- default:
- time.Sleep(500 * time.Millisecond)
- }
- }
- }
- func NewInt64FieldData(fieldName string, numRows int) *schemapb.FieldData {
- return &schemapb.FieldData{
- Type: schemapb.DataType_Int64,
- FieldName: fieldName,
- Field: &schemapb.FieldData_Scalars{
- Scalars: &schemapb.ScalarField{
- Data: &schemapb.ScalarField_LongData{
- LongData: &schemapb.LongArray{
- Data: GenerateInt64Array(numRows, 0),
- },
- },
- },
- },
- }
- }
- func NewInt64FieldDataWithStart(fieldName string, numRows int, start int64) *schemapb.FieldData {
- return &schemapb.FieldData{
- Type: schemapb.DataType_Int64,
- FieldName: fieldName,
- Field: &schemapb.FieldData_Scalars{
- Scalars: &schemapb.ScalarField{
- Data: &schemapb.ScalarField_LongData{
- LongData: &schemapb.LongArray{
- Data: GenerateInt64Array(numRows, start),
- },
- },
- },
- },
- }
- }
- func NewInt64FieldDataNullableWithStart(fieldName string, numRows, start int) *schemapb.FieldData {
- validData, num := GenerateBoolArray(numRows)
- return &schemapb.FieldData{
- Type: schemapb.DataType_Int64,
- FieldName: fieldName,
- Field: &schemapb.FieldData_Scalars{
- Scalars: &schemapb.ScalarField{
- Data: &schemapb.ScalarField_LongData{
- LongData: &schemapb.LongArray{
- Data: GenerateInt64Array(num, int64(start)),
- },
- },
- },
- },
- ValidData: validData,
- }
- }
- func NewInt64SameFieldData(fieldName string, numRows int, value int64) *schemapb.FieldData {
- return &schemapb.FieldData{
- Type: schemapb.DataType_Int64,
- FieldName: fieldName,
- Field: &schemapb.FieldData_Scalars{
- Scalars: &schemapb.ScalarField{
- Data: &schemapb.ScalarField_LongData{
- LongData: &schemapb.LongArray{
- Data: GenerateSameInt64Array(numRows, value),
- },
- },
- },
- },
- }
- }
- func NewVarCharSameFieldData(fieldName string, numRows int, value string) *schemapb.FieldData {
- return &schemapb.FieldData{
- Type: schemapb.DataType_String,
- FieldName: fieldName,
- Field: &schemapb.FieldData_Scalars{
- Scalars: &schemapb.ScalarField{
- Data: &schemapb.ScalarField_StringData{
- StringData: &schemapb.StringArray{
- Data: GenerateSameStringArray(numRows, value),
- },
- },
- },
- },
- }
- }
- func NewStringFieldData(fieldName string, numRows int) *schemapb.FieldData {
- return testutils.NewStringFieldData(fieldName, numRows)
- }
- func NewFloatVectorFieldData(fieldName string, numRows, dim int) *schemapb.FieldData {
- return testutils.NewFloatVectorFieldData(fieldName, numRows, dim)
- }
- func NewFloat16VectorFieldData(fieldName string, numRows, dim int) *schemapb.FieldData {
- return testutils.NewFloat16VectorFieldData(fieldName, numRows, dim)
- }
- func NewBFloat16VectorFieldData(fieldName string, numRows, dim int) *schemapb.FieldData {
- return testutils.NewBFloat16VectorFieldData(fieldName, numRows, dim)
- }
- func NewBinaryVectorFieldData(fieldName string, numRows, dim int) *schemapb.FieldData {
- return testutils.NewBinaryVectorFieldData(fieldName, numRows, dim)
- }
- func NewSparseFloatVectorFieldData(fieldName string, numRows int) *schemapb.FieldData {
- return testutils.NewSparseFloatVectorFieldData(fieldName, numRows)
- }
- func GenerateInt64Array(numRows int, start int64) []int64 {
- ret := make([]int64, numRows)
- for i := 0; i < numRows; i++ {
- ret[i] = int64(i) + start
- }
- return ret
- }
- func GenerateSameInt64Array(numRows int, value int64) []int64 {
- ret := make([]int64, numRows)
- for i := 0; i < numRows; i++ {
- ret[i] = value
- }
- return ret
- }
- func GenerateBoolArray(numRows int) ([]bool, int) {
- var num int
- ret := make([]bool, numRows)
- for i := 0; i < numRows; i++ {
- ret[i] = i%2 == 0
- if ret[i] {
- num++
- }
- }
- return ret, num
- }
- func GenerateSameStringArray(numRows int, value string) []string {
- ret := make([]string, numRows)
- for i := 0; i < numRows; i++ {
- ret[i] = value
- }
- return ret
- }
- func GenerateSparseFloatArray(numRows int) *schemapb.SparseFloatArray {
- return testutils.GenerateSparseFloatVectors(numRows)
- }
- func GenerateHashKeys(numRows int) []uint32 {
- return testutils.GenerateHashKeys(numRows)
- }
|