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
402 lines
13 KiB
Go
402 lines
13 KiB
Go
// 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 index
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/bytedance/mockey"
|
|
"github.com/cockroachdb/errors"
|
|
"github.com/samber/lo"
|
|
"github.com/stretchr/testify/mock"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/stretchr/testify/suite"
|
|
|
|
"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/internal/datanode/compactor"
|
|
"github.com/milvus-io/milvus/internal/mocks"
|
|
"github.com/milvus-io/milvus/internal/mocks/flushcommon/mock_util"
|
|
"github.com/milvus-io/milvus/internal/storage"
|
|
"github.com/milvus-io/milvus/internal/util/indexcgowrapper"
|
|
"github.com/milvus-io/milvus/pkg/v3/common"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/indexcgopb"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/workerpb"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
|
|
)
|
|
|
|
func TestTaskStatsSuite(t *testing.T) {
|
|
suite.Run(t, new(TaskStatsSuite))
|
|
}
|
|
|
|
type TaskStatsSuite struct {
|
|
suite.Suite
|
|
|
|
collectionID int64
|
|
partitionID int64
|
|
clusterID string
|
|
schema *schemapb.CollectionSchema
|
|
|
|
mockBinlogIO *mock_util.MockBinlogIO
|
|
mockChunkManager *mocks.ChunkManager
|
|
segWriter *compactor.SegmentWriter
|
|
}
|
|
|
|
func (s *TaskStatsSuite) SetupSuite() {
|
|
s.collectionID = 100
|
|
s.partitionID = 101
|
|
s.clusterID = "102"
|
|
}
|
|
|
|
func (s *TaskStatsSuite) SetupSubTest() {
|
|
paramtable.Init()
|
|
s.mockBinlogIO = mock_util.NewMockBinlogIO(s.T())
|
|
s.mockChunkManager = mocks.NewChunkManager(s.T())
|
|
}
|
|
|
|
func (s *TaskStatsSuite) GenSegmentWriterWithBM25(magic int64) {
|
|
segWriter, err := compactor.NewSegmentWriter(s.schema, 100, statsBatchSize, magic, s.partitionID, s.collectionID, []int64{102})
|
|
s.Require().NoError(err)
|
|
|
|
v := storage.Value{
|
|
PK: storage.NewInt64PrimaryKey(magic),
|
|
Timestamp: int64(tsoutil.ComposeTSByTime(getMilvusBirthday())),
|
|
Value: genRowWithBM25(magic),
|
|
}
|
|
err = segWriter.Write(&v)
|
|
s.Require().NoError(err)
|
|
segWriter.FlushAndIsFull()
|
|
|
|
s.segWriter = segWriter
|
|
}
|
|
|
|
func (s *TaskStatsSuite) TestSortSegmentWithBM25() {
|
|
s.Run("normal case", func() {
|
|
s.schema = genCollectionSchemaWithBM25()
|
|
s.GenSegmentWriterWithBM25(0)
|
|
_, kvs, fBinlogs, err := serializeWrite(context.TODO(), "root_path", 0, s.segWriter)
|
|
s.NoError(err)
|
|
s.mockBinlogIO.EXPECT().Download(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, paths []string) ([][]byte, error) {
|
|
result := make([][]byte, len(paths))
|
|
for i, path := range paths {
|
|
result[i] = kvs[path]
|
|
}
|
|
return result, nil
|
|
})
|
|
s.mockBinlogIO.EXPECT().Upload(mock.Anything, mock.Anything).Return(nil)
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
|
|
testTaskKey := Key{ClusterID: s.clusterID, TaskID: 100}
|
|
manager := NewTaskManager(ctx)
|
|
manager.LoadOrStoreStatsTask(s.clusterID, testTaskKey.TaskID, &StatsTaskInfo{SegID: 1})
|
|
task := NewStatsTask(ctx, cancel, &workerpb.CreateStatsRequest{
|
|
CollectionID: s.collectionID,
|
|
PartitionID: s.partitionID,
|
|
ClusterID: s.clusterID,
|
|
TaskID: testTaskKey.TaskID,
|
|
TargetSegmentID: 1,
|
|
InsertLogs: lo.Values(fBinlogs),
|
|
Schema: s.schema,
|
|
NumRows: 1,
|
|
StartLogID: 0,
|
|
EndLogID: 7,
|
|
BinlogMaxSize: 64 * 1024 * 1024,
|
|
StorageConfig: &indexpb.StorageConfig{
|
|
RootPath: "root_path",
|
|
},
|
|
}, manager, s.mockChunkManager)
|
|
task.binlogIO = s.mockBinlogIO
|
|
|
|
err = task.PreExecute(ctx)
|
|
s.Require().NoError(err)
|
|
binlog, err := task.sort(ctx)
|
|
s.Require().NoError(err)
|
|
s.Equal(5, len(binlog))
|
|
|
|
// check bm25 log
|
|
s.Equal(1, len(manager.statsTasks))
|
|
for key, task := range manager.statsTasks {
|
|
s.Equal(testTaskKey.ClusterID, key.ClusterID)
|
|
s.Equal(testTaskKey.TaskID, key.TaskID)
|
|
s.Equal(1, len(task.Bm25Logs))
|
|
}
|
|
})
|
|
|
|
s.Run("upload bm25 binlog failed", func() {
|
|
s.schema = genCollectionSchemaWithBM25()
|
|
s.GenSegmentWriterWithBM25(0)
|
|
_, kvs, fBinlogs, err := serializeWrite(context.TODO(), "root_path", 0, s.segWriter)
|
|
s.NoError(err)
|
|
s.mockBinlogIO.EXPECT().Download(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, paths []string) ([][]byte, error) {
|
|
result := make([][]byte, len(paths))
|
|
for i, path := range paths {
|
|
result[i] = kvs[path]
|
|
}
|
|
return result, nil
|
|
})
|
|
s.mockBinlogIO.EXPECT().Upload(mock.Anything, mock.Anything).Return(errors.New("mock error")).Once()
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
|
|
testTaskKey := Key{ClusterID: s.clusterID, TaskID: 100}
|
|
manager := NewTaskManager(ctx)
|
|
manager.LoadOrStoreStatsTask(s.clusterID, testTaskKey.TaskID, &StatsTaskInfo{SegID: 1})
|
|
task := NewStatsTask(ctx, cancel, &workerpb.CreateStatsRequest{
|
|
CollectionID: s.collectionID,
|
|
PartitionID: s.partitionID,
|
|
ClusterID: s.clusterID,
|
|
TaskID: testTaskKey.TaskID,
|
|
TargetSegmentID: 1,
|
|
InsertLogs: lo.Values(fBinlogs),
|
|
Schema: s.schema,
|
|
NumRows: 1,
|
|
StartLogID: 0,
|
|
EndLogID: 7,
|
|
BinlogMaxSize: 64 * 1024 * 1024,
|
|
StorageConfig: &indexpb.StorageConfig{
|
|
RootPath: "root_path",
|
|
},
|
|
}, manager, s.mockChunkManager)
|
|
task.binlogIO = s.mockBinlogIO
|
|
|
|
err = task.PreExecute(ctx)
|
|
s.Require().NoError(err)
|
|
_, err = task.sort(ctx)
|
|
s.Error(err)
|
|
})
|
|
}
|
|
|
|
func (s *TaskStatsSuite) TestBuildIndexParams() {
|
|
s.Run("test storage v2 index params", func() {
|
|
req := &workerpb.CreateStatsRequest{
|
|
TaskID: 1,
|
|
CollectionID: 2,
|
|
PartitionID: 3,
|
|
TargetSegmentID: 4,
|
|
TaskVersion: 5,
|
|
CurrentScalarIndexVersion: int32(1),
|
|
StorageVersion: storage.StorageV2,
|
|
InsertLogs: []*datapb.FieldBinlog{},
|
|
StorageConfig: &indexpb.StorageConfig{RootPath: "/test/path"},
|
|
}
|
|
|
|
options := &BuildIndexOptions{
|
|
TantivyMemory: 0,
|
|
JSONStatsMaxShreddingColumns: 256,
|
|
JSONStatsShreddingRatio: 0.3,
|
|
JSONStatsWriteBatchSize: 81920,
|
|
}
|
|
params := buildIndexParams(req, []string{"file1", "file2"}, nil, &indexcgopb.StorageConfig{}, options, "")
|
|
|
|
s.Equal(storage.StorageV2, params.StorageVersion)
|
|
s.NotNil(params.SegmentInsertFiles)
|
|
})
|
|
|
|
s.Run("test external source spec params", func() {
|
|
req := &workerpb.CreateStatsRequest{
|
|
TaskID: 1,
|
|
CollectionID: 2,
|
|
PartitionID: 3,
|
|
TargetSegmentID: 4,
|
|
TaskVersion: 5,
|
|
CurrentScalarIndexVersion: int32(1),
|
|
StorageVersion: storage.StorageV3,
|
|
InsertLogs: []*datapb.FieldBinlog{},
|
|
StorageConfig: &indexpb.StorageConfig{RootPath: "/test/path"},
|
|
Schema: &schemapb.CollectionSchema{
|
|
ExternalSource: "minio://localhost:9000/a-bucket/external",
|
|
ExternalSpec: `{"format":"parquet"}`,
|
|
},
|
|
}
|
|
|
|
params := buildIndexParams(req, nil, nil, &indexcgopb.StorageConfig{}, nil, "")
|
|
|
|
s.Equal(req.GetSchema().GetExternalSource(), params.GetExternalSource())
|
|
s.Equal(req.GetSchema().GetExternalSpec(), params.GetExternalSpec())
|
|
})
|
|
}
|
|
|
|
func genCollectionSchemaWithBM25() *schemapb.CollectionSchema {
|
|
return &schemapb.CollectionSchema{
|
|
Name: "schema",
|
|
Description: "schema",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: common.RowIDField,
|
|
Name: "row_id",
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
{
|
|
FieldID: common.TimeStampField,
|
|
Name: "Timestamp",
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
{
|
|
FieldID: 100,
|
|
Name: "pk",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPrimaryKey: true,
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "text",
|
|
DataType: schemapb.DataType_VarChar,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: common.MaxLengthKey,
|
|
Value: "8",
|
|
},
|
|
},
|
|
},
|
|
{
|
|
FieldID: 102,
|
|
Name: "sparse",
|
|
DataType: schemapb.DataType_SparseFloatVector,
|
|
},
|
|
},
|
|
Functions: []*schemapb.FunctionSchema{{
|
|
Name: "BM25",
|
|
Id: 100,
|
|
Type: schemapb.FunctionType_BM25,
|
|
InputFieldNames: []string{"text"},
|
|
InputFieldIds: []int64{101},
|
|
OutputFieldNames: []string{"sparse"},
|
|
OutputFieldIds: []int64{102},
|
|
}},
|
|
}
|
|
}
|
|
|
|
func genRowWithBM25(magic int64) map[int64]interface{} {
|
|
ts := tsoutil.ComposeTSByTime(getMilvusBirthday())
|
|
return map[int64]interface{}{
|
|
common.RowIDField: magic,
|
|
common.TimeStampField: int64(ts),
|
|
100: magic,
|
|
101: "varchar",
|
|
102: typeutil.CreateAndSortSparseFloatRow(map[uint32]float32{1: 1}),
|
|
}
|
|
}
|
|
|
|
func getMilvusBirthday() time.Time {
|
|
return time.Date(2019, time.Month(5), 30, 0, 0, 0, 0, time.UTC)
|
|
}
|
|
|
|
// nullable JSON may have no insert column binlog; getInsertFiles should allow empty paths (aligned with text index).
|
|
func TestCreateJSONKeyStats_NullableJSONMissingFieldBinlog(t *testing.T) {
|
|
paramtable.Init()
|
|
ctx := context.Background()
|
|
mgr := NewTaskManager(ctx)
|
|
mgr.LoadOrStoreStatsTask("c1", 1, &StatsTaskInfo{SegID: 10})
|
|
|
|
req := &workerpb.CreateStatsRequest{
|
|
ClusterID: "c1",
|
|
TaskID: 1,
|
|
CollectionID: 100,
|
|
PartitionID: 101,
|
|
TargetSegmentID: 102,
|
|
SegmentID: 103,
|
|
InsertChannel: "ch",
|
|
TaskVersion: 1,
|
|
JsonKeyStatsDataFormat: common.JSONStatsDataFormatVersion,
|
|
StorageConfig: &indexpb.StorageConfig{RootPath: "/root"},
|
|
SubJobType: indexpb.StatsSubJob_JsonKeyIndexJob,
|
|
StorageVersion: 1,
|
|
NumRows: 10,
|
|
Schema: &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64},
|
|
{FieldID: 201, Name: "j", DataType: schemapb.DataType_JSON, Nullable: true},
|
|
},
|
|
},
|
|
}
|
|
ctx2, cancel := context.WithCancel(ctx)
|
|
defer cancel()
|
|
st := NewStatsTask(ctx2, cancel, req, mgr, nil)
|
|
|
|
insertBinlogs := []*datapb.FieldBinlog{
|
|
{FieldID: 100, Binlogs: []*datapb.Binlog{{LogID: 1}}},
|
|
}
|
|
|
|
var gotInsertFiles []string
|
|
var gotNumRows int64
|
|
m := mockey.Mock(indexcgowrapper.CreateJSONKeyStats).To(func(_ context.Context, info *indexcgopb.BuildIndexInfo) (*indexcgowrapper.JSONKeyStatsResult, error) {
|
|
gotInsertFiles = info.InsertFiles
|
|
gotNumRows = info.GetNumRows()
|
|
return &indexcgowrapper.JSONKeyStatsResult{Files: map[string]int64{}}, nil
|
|
}).Build()
|
|
defer m.UnPatch()
|
|
|
|
err := st.createJSONKeyStats(ctx, req.GetStorageConfig(),
|
|
req.GetCollectionID(), req.GetPartitionID(), req.GetTargetSegmentID(),
|
|
req.GetTaskVersion(), req.GetTaskID(),
|
|
common.JSONStatsDataFormatVersion,
|
|
insertBinlogs, 256, 0.3, 81920)
|
|
require.NoError(t, err)
|
|
require.Empty(t, gotInsertFiles)
|
|
require.Equal(t, int64(10), gotNumRows)
|
|
}
|
|
|
|
func TestCreateJSONKeyStats_NonNullableJSONMissingFieldBinlog(t *testing.T) {
|
|
paramtable.Init()
|
|
ctx := context.Background()
|
|
mgr := NewTaskManager(ctx)
|
|
mgr.LoadOrStoreStatsTask("c1", 2, &StatsTaskInfo{SegID: 10})
|
|
|
|
req := &workerpb.CreateStatsRequest{
|
|
ClusterID: "c1",
|
|
TaskID: 2,
|
|
CollectionID: 100,
|
|
PartitionID: 101,
|
|
TargetSegmentID: 102,
|
|
SegmentID: 103,
|
|
InsertChannel: "ch",
|
|
TaskVersion: 1,
|
|
JsonKeyStatsDataFormat: common.JSONStatsDataFormatVersion,
|
|
StorageConfig: &indexpb.StorageConfig{RootPath: "/root"},
|
|
SubJobType: indexpb.StatsSubJob_JsonKeyIndexJob,
|
|
StorageVersion: 1,
|
|
Schema: &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64},
|
|
{FieldID: 201, Name: "j", DataType: schemapb.DataType_JSON, Nullable: false},
|
|
},
|
|
},
|
|
}
|
|
ctx2, cancel := context.WithCancel(ctx)
|
|
defer cancel()
|
|
st := NewStatsTask(ctx2, cancel, req, mgr, nil)
|
|
|
|
insertBinlogs := []*datapb.FieldBinlog{
|
|
{FieldID: 100, Binlogs: []*datapb.Binlog{{LogID: 1}}},
|
|
}
|
|
|
|
err := st.createJSONKeyStats(ctx, req.GetStorageConfig(),
|
|
req.GetCollectionID(), req.GetPartitionID(), req.GetTargetSegmentID(),
|
|
req.GetTaskVersion(), req.GetTaskID(),
|
|
common.JSONStatsDataFormatVersion,
|
|
insertBinlogs, 256, 0.3, 81920)
|
|
require.Error(t, err)
|
|
require.Contains(t, err.Error(), "field binlog not found for field 201")
|
|
}
|