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
1806 lines
60 KiB
Go
1806 lines
60 KiB
Go
// Copyright 2023 Zilliz
|
|
//
|
|
// Licensed 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.
|
|
|
|
//go:build dynamic
|
|
|
|
package packed
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"sync/atomic"
|
|
"testing"
|
|
|
|
"github.com/bytedance/mockey"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"google.golang.org/protobuf/encoding/protojson"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
|
"github.com/milvus-io/milvus/internal/snapshotio"
|
|
"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/indexpb"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
|
|
)
|
|
|
|
func TestExploreFiles_EmptyColumns(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: tmpDir,
|
|
RootPath: tmpDir,
|
|
}
|
|
|
|
// Empty columns should still work
|
|
files, _, err := ExploreFilesReturnManifestPath([]string{}, "parquet", "/tmp", "/nonexistent", config, ExternalSpecContext{})
|
|
|
|
// Expect error due to nonexistent directory
|
|
assert.Error(t, err)
|
|
assert.Nil(t, files)
|
|
}
|
|
|
|
func TestExploreFiles_InvalidDirectory(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: tmpDir,
|
|
RootPath: tmpDir,
|
|
}
|
|
|
|
files, _, err := ExploreFilesReturnManifestPath(
|
|
[]string{"col1", "col2"},
|
|
"parquet",
|
|
"/nonexistent/base",
|
|
"/nonexistent/explore",
|
|
config,
|
|
ExternalSpecContext{},
|
|
)
|
|
|
|
assert.Error(t, err)
|
|
assert.Nil(t, files)
|
|
}
|
|
|
|
func TestGetFileInfo_NonexistentFile(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: tmpDir,
|
|
RootPath: tmpDir,
|
|
}
|
|
|
|
info, err := GetFileInfo("parquet", "/nonexistent/file.parquet", config, ExternalSpecContext{})
|
|
|
|
assert.Error(t, err)
|
|
assert.Nil(t, info)
|
|
}
|
|
|
|
func TestGetFileInfo_InvalidFormat(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: tmpDir,
|
|
RootPath: tmpDir,
|
|
}
|
|
|
|
// Create a temp file
|
|
tmpFile := filepath.Join(tmpDir, "test.txt")
|
|
err := os.WriteFile(tmpFile, []byte("not a parquet file"), 0o600)
|
|
require.NoError(t, err)
|
|
|
|
info, err := GetFileInfo("parquet", tmpFile, config, ExternalSpecContext{})
|
|
|
|
// Should fail because it's not a valid parquet file
|
|
assert.Error(t, err)
|
|
assert.Nil(t, info)
|
|
}
|
|
|
|
func TestCreateManifestForSegment_EmptyFragments(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: tmpDir,
|
|
RootPath: tmpDir,
|
|
}
|
|
|
|
manifestPath, err := CreateManifestForSegment(
|
|
"/test",
|
|
[]string{"col1"},
|
|
"parquet",
|
|
[]Fragment{},
|
|
config,
|
|
)
|
|
|
|
assert.Error(t, err)
|
|
assert.Empty(t, manifestPath)
|
|
assert.Contains(t, err.Error(), "fragments cannot be empty")
|
|
}
|
|
|
|
func TestCreateManifestForSegment_InvalidBasePath(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: tmpDir,
|
|
RootPath: tmpDir,
|
|
}
|
|
|
|
fragments := []Fragment{
|
|
{
|
|
FragmentID: 0,
|
|
FilePath: "/data/file.parquet",
|
|
StartRow: 0,
|
|
EndRow: 1000,
|
|
RowCount: 1000,
|
|
},
|
|
}
|
|
|
|
// NOTE: C FFI does not validate base path at creation time (paths are virtual in object storage),
|
|
// so this may succeed. We only verify no panic occurs.
|
|
_, _ = CreateManifestForSegment(
|
|
"/test/path",
|
|
[]string{"col1"},
|
|
"parquet",
|
|
fragments,
|
|
config,
|
|
)
|
|
}
|
|
|
|
func TestManifestColumnGroupsToFragments_DeduplicatesByPathRange(t *testing.T) {
|
|
groups := []manifestColumnGroup{
|
|
{
|
|
Columns: []string{"id", "vector"},
|
|
Fragments: []Fragment{
|
|
{FilePath: "a.parquet", StartRow: 0, EndRow: 10, RowCount: 10},
|
|
},
|
|
},
|
|
{
|
|
Columns: []string{"score"},
|
|
Fragments: []Fragment{
|
|
{FilePath: "a.parquet", StartRow: 0, EndRow: 10, RowCount: 10},
|
|
},
|
|
},
|
|
}
|
|
|
|
fragments := manifestColumnGroupsToFragments(groups)
|
|
assert.Len(t, fragments, 1)
|
|
assert.Equal(t, "a.parquet", fragments[0].FilePath)
|
|
assert.Equal(t, int64(0), fragments[0].StartRow)
|
|
assert.Equal(t, int64(10), fragments[0].EndRow)
|
|
}
|
|
|
|
func TestColumnsToAppend_SkipsExistingSameFragments(t *testing.T) {
|
|
existing := []manifestColumnGroup{
|
|
{
|
|
Columns: []string{"id"},
|
|
Fragments: []Fragment{
|
|
{FilePath: "a.parquet", StartRow: 0, EndRow: 10, RowCount: 10},
|
|
},
|
|
},
|
|
}
|
|
requested := []string{"id", "score"}
|
|
fragments := []Fragment{{FilePath: "a.parquet", StartRow: 0, EndRow: 10, RowCount: 10}}
|
|
|
|
columns, err := columnsToAppend(existing, requested, fragments)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, []string{"score"}, columns)
|
|
}
|
|
|
|
func TestColumnsToAppendRejectsExistingDifferentFragments(t *testing.T) {
|
|
existing := []manifestColumnGroup{
|
|
{
|
|
Columns: []string{"id"},
|
|
Fragments: []Fragment{
|
|
{FilePath: "a.parquet", StartRow: 0, EndRow: 5, RowCount: 5},
|
|
},
|
|
},
|
|
}
|
|
requested := []string{"id"}
|
|
fragments := []Fragment{{FilePath: "a.parquet", StartRow: 0, EndRow: 10, RowCount: 10}}
|
|
|
|
_, err := columnsToAppend(existing, requested, fragments)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "already exists with different fragments")
|
|
}
|
|
|
|
func TestColumnsToAppendRejectsDuplicateExistingColumnWithDifferentFragments(t *testing.T) {
|
|
existing := []manifestColumnGroup{
|
|
{
|
|
Columns: []string{"id"},
|
|
Fragments: []Fragment{
|
|
{FilePath: "a.parquet", StartRow: 0, EndRow: 10, RowCount: 10},
|
|
},
|
|
},
|
|
{
|
|
Columns: []string{"id"},
|
|
Fragments: []Fragment{
|
|
{FilePath: "a.parquet", StartRow: 10, EndRow: 20, RowCount: 10},
|
|
},
|
|
},
|
|
}
|
|
requested := []string{"id"}
|
|
fragments := []Fragment{{FilePath: "a.parquet", StartRow: 0, EndRow: 10, RowCount: 10}}
|
|
|
|
_, err := columnsToAppend(existing, requested, fragments)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "column id already exists with different fragments")
|
|
}
|
|
|
|
func TestColumnsToAppend_DeduplicatesDuplicateRequestedMissingColumns(t *testing.T) {
|
|
requested := []string{"score", "id", "score", "id"}
|
|
fragments := []Fragment{{FilePath: "a.parquet", StartRow: 0, EndRow: 10, RowCount: 10}}
|
|
|
|
columns, err := columnsToAppend(nil, requested, fragments)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, []string{"score", "id"}, columns)
|
|
}
|
|
|
|
func TestAppendSegmentManifestColumnsUsesCommitManifestUpdates(t *testing.T) {
|
|
ctx := context.Background()
|
|
config := &indexpb.StorageConfig{StorageType: "local"}
|
|
oldManifestPath := MarshalManifestPath("segment", 3)
|
|
fragments := []Fragment{{FilePath: "a.parquet", StartRow: 0, EndRow: 10, RowCount: 10}}
|
|
|
|
mockRead := mockey.Mock(readColumnGroupsFromManifest).Return(nil, nil).Build()
|
|
defer mockRead.UnPatch()
|
|
|
|
var gotBasePath string
|
|
var gotVersion int64
|
|
var gotAddNewColumnGroups bool
|
|
mockCommit := mockey.Mock(CommitManifestUpdates).
|
|
To(func(basePath string, baseVersion int64, storageConfig *indexpb.StorageConfig, updates *ManifestUpdates) (string, error) {
|
|
gotBasePath = basePath
|
|
gotVersion = baseVersion
|
|
cgs, ok := updates.NewFiles.(*ColumnGroups)
|
|
require.True(t, ok)
|
|
require.NotNil(t, cgs.cColumnGroups)
|
|
gotAddNewColumnGroups = cgs.addNewColumnGroups
|
|
return MarshalManifestPath(basePath, baseVersion+1), nil
|
|
}).Build()
|
|
defer mockCommit.UnPatch()
|
|
|
|
got, err := AppendSegmentManifestColumns(ctx, oldManifestPath, "parquet", []string{"score"}, fragments, config)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, MarshalManifestPath("segment", 4), got)
|
|
assert.Equal(t, "segment", gotBasePath)
|
|
assert.Equal(t, int64(3), gotVersion)
|
|
assert.True(t, gotAddNewColumnGroups)
|
|
}
|
|
|
|
func TestAppendSegmentManifestColumnsValidationPaths(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
_, err := AppendSegmentManifestColumns(ctx, "manifest", "parquet", []string{"score"}, []Fragment{
|
|
{FilePath: "a.parquet", StartRow: 0, EndRow: 10, RowCount: 10},
|
|
}, &indexpb.StorageConfig{StorageType: "local"})
|
|
assert.ErrorIs(t, err, context.Canceled)
|
|
|
|
oldManifestPath := MarshalManifestPath("segment", 3)
|
|
got, err := AppendSegmentManifestColumns(context.Background(), oldManifestPath, "parquet", nil, []Fragment{
|
|
{FilePath: "a.parquet", StartRow: 0, EndRow: 10, RowCount: 10},
|
|
}, &indexpb.StorageConfig{StorageType: "local"})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, oldManifestPath, got)
|
|
|
|
_, err = AppendSegmentManifestColumns(context.Background(), oldManifestPath, "parquet", []string{"score"}, nil,
|
|
&indexpb.StorageConfig{StorageType: "local"})
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "fragments cannot be empty")
|
|
}
|
|
|
|
func TestAppendSegmentManifestColumnsReadAndNoopPaths(t *testing.T) {
|
|
config := &indexpb.StorageConfig{StorageType: "local"}
|
|
oldManifestPath := MarshalManifestPath("segment", 3)
|
|
fragments := []Fragment{{FilePath: "a.parquet", StartRow: 0, EndRow: 10, RowCount: 10}}
|
|
|
|
mockReadErr := mockey.Mock(readColumnGroupsFromManifest).Return(nil, fmt.Errorf("read failed")).Build()
|
|
_, err := AppendSegmentManifestColumns(context.Background(), oldManifestPath, "parquet", []string{"score"}, fragments, config)
|
|
mockReadErr.UnPatch()
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "failed to read manifest column groups")
|
|
|
|
existing := []manifestColumnGroup{{
|
|
Columns: []string{"score"},
|
|
Fragments: fragments,
|
|
}}
|
|
mockReadNoop := mockey.Mock(readColumnGroupsFromManifest).Return(existing, nil).Build()
|
|
got, err := AppendSegmentManifestColumns(context.Background(), oldManifestPath, "parquet", []string{"score"}, fragments, config)
|
|
mockReadNoop.UnPatch()
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, oldManifestPath, got)
|
|
|
|
mockReadParse := mockey.Mock(readColumnGroupsFromManifest).Return(nil, nil).Build()
|
|
_, err = AppendSegmentManifestColumns(context.Background(), "not-json", "parquet", []string{"score"}, fragments, config)
|
|
mockReadParse.UnPatch()
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "failed to parse manifest path")
|
|
}
|
|
|
|
func TestFileInfo_Struct(t *testing.T) {
|
|
info := FileInfo{
|
|
FilePath: "/data/test.parquet",
|
|
NumRows: 1000000,
|
|
}
|
|
|
|
assert.Equal(t, "/data/test.parquet", info.FilePath)
|
|
assert.Equal(t, int64(1000000), info.NumRows)
|
|
}
|
|
|
|
func TestFragment_Struct(t *testing.T) {
|
|
fragment := Fragment{
|
|
FragmentID: 1,
|
|
FilePath: "/data/large.parquet",
|
|
StartRow: 1000000,
|
|
EndRow: 2000000,
|
|
RowCount: 1000000,
|
|
}
|
|
|
|
assert.Equal(t, int64(1), fragment.FragmentID)
|
|
assert.Equal(t, "/data/large.parquet", fragment.FilePath)
|
|
assert.Equal(t, int64(1000000), fragment.StartRow)
|
|
assert.Equal(t, int64(2000000), fragment.EndRow)
|
|
assert.Equal(t, int64(1000000), fragment.RowCount)
|
|
}
|
|
|
|
func TestFragment_RowCountConsistency(t *testing.T) {
|
|
// Test that RowCount should equal EndRow - StartRow
|
|
testCases := []struct {
|
|
startRow int64
|
|
endRow int64
|
|
}{
|
|
{0, 1000},
|
|
{1000, 2000},
|
|
{0, 1000000},
|
|
{5000000, 10000000},
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
fragment := Fragment{
|
|
FragmentID: 0,
|
|
FilePath: "/data/test.parquet",
|
|
StartRow: tc.startRow,
|
|
EndRow: tc.endRow,
|
|
RowCount: tc.endRow - tc.startRow,
|
|
}
|
|
|
|
assert.Equal(t, fragment.EndRow-fragment.StartRow, fragment.RowCount)
|
|
}
|
|
}
|
|
|
|
func TestNewFragmentIDGenerator(t *testing.T) {
|
|
gen := NewFragmentIDGenerator(0)
|
|
assert.Equal(t, int64(0), gen())
|
|
assert.Equal(t, int64(1), gen())
|
|
assert.Equal(t, int64(2), gen())
|
|
|
|
gen2 := NewFragmentIDGenerator(100)
|
|
assert.Equal(t, int64(100), gen2())
|
|
assert.Equal(t, int64(101), gen2())
|
|
}
|
|
|
|
func TestSplitFileToFragments_SmallFile(t *testing.T) {
|
|
gen := NewFragmentIDGenerator(0)
|
|
fragments := SplitFileToFragments("/data/small.parquet", 500, 1000, gen)
|
|
|
|
assert.Len(t, fragments, 1)
|
|
assert.Equal(t, int64(0), fragments[0].FragmentID)
|
|
assert.Equal(t, "/data/small.parquet", fragments[0].FilePath)
|
|
assert.Equal(t, int64(0), fragments[0].StartRow)
|
|
assert.Equal(t, int64(500), fragments[0].EndRow)
|
|
assert.Equal(t, int64(500), fragments[0].RowCount)
|
|
}
|
|
|
|
func TestSplitFileToFragments_ExactLimit(t *testing.T) {
|
|
gen := NewFragmentIDGenerator(0)
|
|
fragments := SplitFileToFragments("/data/exact.parquet", 1000, 1000, gen)
|
|
|
|
assert.Len(t, fragments, 1)
|
|
assert.Equal(t, int64(0), fragments[0].StartRow)
|
|
assert.Equal(t, int64(1000), fragments[0].EndRow)
|
|
}
|
|
|
|
func TestSplitFileToFragments_LargeFile(t *testing.T) {
|
|
gen := NewFragmentIDGenerator(0)
|
|
fragments := SplitFileToFragments("/data/large.parquet", 2500, 1000, gen)
|
|
|
|
assert.Len(t, fragments, 3)
|
|
// Fragment 0: [0, 1000)
|
|
assert.Equal(t, int64(0), fragments[0].FragmentID)
|
|
assert.Equal(t, int64(0), fragments[0].StartRow)
|
|
assert.Equal(t, int64(1000), fragments[0].EndRow)
|
|
assert.Equal(t, int64(1000), fragments[0].RowCount)
|
|
// Fragment 1: [1000, 2000)
|
|
assert.Equal(t, int64(1), fragments[1].FragmentID)
|
|
assert.Equal(t, int64(1000), fragments[1].StartRow)
|
|
assert.Equal(t, int64(2000), fragments[1].EndRow)
|
|
// Fragment 2: [2000, 2500) — partial last fragment
|
|
assert.Equal(t, int64(2), fragments[2].FragmentID)
|
|
assert.Equal(t, int64(2000), fragments[2].StartRow)
|
|
assert.Equal(t, int64(2500), fragments[2].EndRow)
|
|
assert.Equal(t, int64(500), fragments[2].RowCount)
|
|
}
|
|
|
|
func TestSplitFileToFragments_FragmentIDContinuity(t *testing.T) {
|
|
gen := NewFragmentIDGenerator(10)
|
|
fragments := SplitFileToFragments("/data/f.parquet", 3000, 1000, gen)
|
|
|
|
assert.Len(t, fragments, 3)
|
|
assert.Equal(t, int64(10), fragments[0].FragmentID)
|
|
assert.Equal(t, int64(11), fragments[1].FragmentID)
|
|
assert.Equal(t, int64(12), fragments[2].FragmentID)
|
|
}
|
|
|
|
func TestGetColumnNamesFromSchema_Nil(t *testing.T) {
|
|
assert.Nil(t, GetColumnNamesFromSchema(nil))
|
|
}
|
|
|
|
func TestGetColumnNamesFromSchema_NoExternalField(t *testing.T) {
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{Name: "id"},
|
|
{Name: "vector"},
|
|
},
|
|
}
|
|
columns := GetColumnNamesFromSchema(schema)
|
|
assert.Equal(t, []string{"id", "vector"}, columns)
|
|
}
|
|
|
|
func TestGetColumnNamesFromSchema_WithExternalField(t *testing.T) {
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{Name: "id", ExternalField: "external_id"},
|
|
{Name: "vector"},
|
|
{Name: "text", ExternalField: "raw_text"},
|
|
},
|
|
}
|
|
columns := GetColumnNamesFromSchema(schema)
|
|
assert.Equal(t, []string{"external_id", "vector", "raw_text"}, columns)
|
|
}
|
|
|
|
func TestGetColumnNamesFromSchema_MilvusTableUsesSourceFieldIDs(t *testing.T) {
|
|
schema := &schemapb.CollectionSchema{
|
|
ExternalSource: "s3://bucket/snapshots/100/metadata/200.json",
|
|
ExternalSpec: `{"format":"milvus-table"}`,
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 0, Name: "__virtual_pk__"},
|
|
{FieldID: 100, Name: "target_pk", ExternalField: "pk"},
|
|
{FieldID: 101, Name: "target_vector", ExternalField: "vector"},
|
|
},
|
|
}
|
|
columns := GetColumnNamesFromSchema(schema)
|
|
assert.Equal(t, []string{"100", "101"}, columns)
|
|
}
|
|
|
|
func TestGetColumnNamesFromSchema_MilvusTableSourceSchemaUsesFieldIDs(t *testing.T) {
|
|
schema := &schemapb.CollectionSchema{
|
|
ExternalSource: "s3://bucket/snapshots/100/metadata/200.json",
|
|
ExternalSpec: `{"format":"milvus-table"}`,
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 100, Name: "pk"},
|
|
{FieldID: 101, Name: "vector"},
|
|
},
|
|
}
|
|
columns := GetColumnNamesFromSchema(schema)
|
|
assert.Equal(t, []string{"100", "101"}, columns)
|
|
}
|
|
|
|
func TestGetColumnNamesFromSchema_MilvusTableRealPKIncludesTimestamp(t *testing.T) {
|
|
schema := &schemapb.CollectionSchema{
|
|
ExternalSource: "s3://bucket/snapshots/100/metadata/200.json",
|
|
ExternalSpec: `{"format":"milvus-table"}`,
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 100, Name: "pk", IsPrimaryKey: true},
|
|
{FieldID: 101, Name: "vector"},
|
|
{FieldID: common.TimeStampField, Name: common.TimeStampFieldName, DataType: schemapb.DataType_Int64},
|
|
{FieldID: common.RowIDField, Name: common.RowIDFieldName, DataType: schemapb.DataType_Int64},
|
|
},
|
|
}
|
|
|
|
columns := GetColumnNamesFromSchema(schema)
|
|
assert.Equal(t, []string{"100", "101", "1"}, columns)
|
|
}
|
|
|
|
func TestGetColumnNamesFromSchema_EmptyFields(t *testing.T) {
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{},
|
|
}
|
|
columns := GetColumnNamesFromSchema(schema)
|
|
assert.Nil(t, columns)
|
|
}
|
|
|
|
func TestGetColumnNamesFromSchema_ExternalSkipsFunctionOutputAndSystemFields(t *testing.T) {
|
|
schema := &schemapb.CollectionSchema{
|
|
ExternalSource: "s3://bucket/path",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{Name: "id", ExternalField: "external_id"},
|
|
{Name: "virtual_pk"},
|
|
{Name: "sparse", IsFunctionOutput: true},
|
|
},
|
|
}
|
|
columns := GetColumnNamesFromSchema(schema)
|
|
assert.Equal(t, []string{"external_id"}, columns)
|
|
}
|
|
|
|
func TestGetColumnNamesFromSchema_MilvusTableSkipsFunctionOutput(t *testing.T) {
|
|
schema := &schemapb.CollectionSchema{
|
|
ExternalSource: "s3://bucket/snapshots/100/metadata/200.json",
|
|
ExternalSpec: `{"format":"milvus-table"}`,
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 100, Name: "text", ExternalField: "text"},
|
|
{FieldID: 101, Name: "vec", ExternalField: "vec"},
|
|
{FieldID: 102, Name: "sparse", DataType: schemapb.DataType_SparseFloatVector, IsFunctionOutput: true},
|
|
},
|
|
Functions: []*schemapb.FunctionSchema{
|
|
{
|
|
Name: "bm25",
|
|
InputFieldNames: []string{"text"},
|
|
OutputFieldNames: []string{"sparse"},
|
|
},
|
|
},
|
|
}
|
|
|
|
columns := GetColumnNamesFromSchema(schema)
|
|
assert.Equal(t, []string{"100", "101"}, columns)
|
|
}
|
|
|
|
func TestHasExternalPrimaryKey(t *testing.T) {
|
|
assert.False(t, HasExternalPrimaryKey(nil))
|
|
|
|
assert.True(t, HasExternalPrimaryKey(&schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{Name: "id", IsPrimaryKey: true, ExternalField: "source_id"},
|
|
},
|
|
}))
|
|
|
|
assert.False(t, HasExternalPrimaryKey(&schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{Name: "__virtual_pk__", IsPrimaryKey: true},
|
|
{Name: "id", ExternalField: "source_id"},
|
|
},
|
|
}))
|
|
|
|
assert.False(t, HasExternalPrimaryKey(&schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{Name: "id", ExternalField: "source_id"},
|
|
},
|
|
}))
|
|
}
|
|
|
|
func TestBuildCurrentSegmentFragments_NoManifest(t *testing.T) {
|
|
segments := []*datapb.SegmentInfo{
|
|
{ID: 1, NumOfRows: 1000},
|
|
{ID: 2, NumOfRows: 2000},
|
|
}
|
|
result, err := BuildCurrentSegmentFragments(segments, nil, nil)
|
|
assert.NoError(t, err)
|
|
assert.Len(t, result, 2)
|
|
|
|
// Check virtual fragment for segment 1
|
|
assert.Len(t, result[1], 1)
|
|
assert.Equal(t, int64(1), result[1][0].FragmentID)
|
|
assert.Equal(t, "", result[1][0].FilePath)
|
|
assert.Equal(t, int64(0), result[1][0].StartRow)
|
|
assert.Equal(t, int64(1000), result[1][0].EndRow)
|
|
assert.Equal(t, int64(1000), result[1][0].RowCount)
|
|
|
|
// Check virtual fragment for segment 2
|
|
assert.Len(t, result[2], 1)
|
|
assert.Equal(t, int64(2000), result[2][0].RowCount)
|
|
}
|
|
|
|
func TestBuildCurrentSegmentFragments_EmptySegments(t *testing.T) {
|
|
result, err := BuildCurrentSegmentFragments(nil, nil, nil)
|
|
assert.NoError(t, err)
|
|
assert.Empty(t, result)
|
|
}
|
|
|
|
func TestBuildCurrentSegmentFragments_ManifestPathWithNilConfig(t *testing.T) {
|
|
// When storageConfig is nil, manifest path should be skipped (use virtual fragment)
|
|
segments := []*datapb.SegmentInfo{
|
|
{ID: 1, NumOfRows: 500, ManifestPath: "some/path"},
|
|
}
|
|
result, err := BuildCurrentSegmentFragments(segments, nil, nil)
|
|
assert.NoError(t, err)
|
|
assert.Len(t, result[1], 1)
|
|
assert.Equal(t, int64(500), result[1][0].RowCount)
|
|
}
|
|
|
|
func TestBuildCurrentSegmentFragments_PassesColumns(t *testing.T) {
|
|
storageConfig := &indexpb.StorageConfig{RootPath: "files", StorageType: "local"}
|
|
segments := []*datapb.SegmentInfo{
|
|
{ID: 1, NumOfRows: 500, ManifestPath: MarshalManifestPath("files/insert_log/1/2/3", 1)},
|
|
}
|
|
expected := []Fragment{{FragmentID: 7, FilePath: "/data/file.parquet", RowCount: 500}}
|
|
|
|
called := false
|
|
mockRead := mockey.Mock(ReadFragmentsFromManifest).
|
|
To(func(manifestPath string, storageConfig *indexpb.StorageConfig, columns []string) ([]Fragment, error) {
|
|
assert.Equal(t, []string{"text_col"}, append([]string(nil), columns...))
|
|
called = true
|
|
return expected, nil
|
|
}).Build()
|
|
defer mockRead.UnPatch()
|
|
|
|
result, err := BuildCurrentSegmentFragments(segments, storageConfig, []string{"text_col"})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, expected, result[1])
|
|
assert.True(t, called)
|
|
}
|
|
|
|
func TestBuildCurrentSegmentFragments_ManifestErrorsAndEmptyResult(t *testing.T) {
|
|
storageConfig := &indexpb.StorageConfig{RootPath: "files", StorageType: "local"}
|
|
segments := []*datapb.SegmentInfo{
|
|
{ID: 1, NumOfRows: 500, ManifestPath: MarshalManifestPath("files/insert_log/1/2/3", 1)},
|
|
}
|
|
|
|
mockRead := mockey.Mock(ReadFragmentsFromManifest).Return(nil, fmt.Errorf("read failed")).Build()
|
|
_, err := BuildCurrentSegmentFragments(segments, storageConfig, []string{"text_col"})
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "failed to read manifest for segment 1")
|
|
mockRead.UnPatch()
|
|
|
|
mockRead = mockey.Mock(ReadFragmentsFromManifest).Return(nil, nil).Build()
|
|
defer mockRead.UnPatch()
|
|
result, err := BuildCurrentSegmentFragments(segments, storageConfig, []string{"missing_col"})
|
|
require.NoError(t, err)
|
|
require.Len(t, result[1], 1)
|
|
assert.Equal(t, int64(1), result[1][0].FragmentID)
|
|
assert.Equal(t, int64(500), result[1][0].RowCount)
|
|
}
|
|
|
|
func TestReadFragmentsFromManifest_FiltersColumns(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{StorageType: "local", BucketName: tmpDir, RootPath: tmpDir}
|
|
fragments := []Fragment{
|
|
{FragmentID: 0, FilePath: filepath.Join(tmpDir, "data-1.parquet"), StartRow: 0, EndRow: 10, RowCount: 10},
|
|
{FragmentID: 1, FilePath: filepath.Join(tmpDir, "data-2.parquet"), StartRow: 10, EndRow: 20, RowCount: 10},
|
|
}
|
|
|
|
manifestPath, err := CreateManifestForSegment(
|
|
filepath.Join(tmpDir, "segment"),
|
|
[]string{"text_col", "101"},
|
|
"parquet",
|
|
fragments,
|
|
config,
|
|
)
|
|
require.NoError(t, err)
|
|
|
|
got, err := ReadFragmentsFromManifest(manifestPath, config, []string{"text_col"})
|
|
require.NoError(t, err)
|
|
require.Len(t, got, 2)
|
|
assert.Equal(t, fragments[0].FilePath, got[0].FilePath)
|
|
assert.Equal(t, int64(0), got[0].StartRow)
|
|
assert.Equal(t, int64(10), got[0].EndRow)
|
|
assert.Equal(t, int64(10), got[0].RowCount)
|
|
assert.Equal(t, fragments[1].FilePath, got[1].FilePath)
|
|
|
|
got, err = ReadFragmentsFromManifest(manifestPath, config, []string{"missing_col"})
|
|
require.NoError(t, err)
|
|
assert.Empty(t, got)
|
|
|
|
has, err := ManifestHasColumns(manifestPath, config, []string{"text_col", "101"})
|
|
require.NoError(t, err)
|
|
assert.True(t, has)
|
|
|
|
has, err = ManifestHasColumns(manifestPath, config, []string{"text_col", "missing_col"})
|
|
require.NoError(t, err)
|
|
assert.False(t, has)
|
|
|
|
has, err = ManifestHasColumns(manifestPath, config, nil)
|
|
require.NoError(t, err)
|
|
assert.True(t, has)
|
|
|
|
_, err = ReadFragmentsFromManifest("bad manifest", config, nil)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "failed to parse manifest path")
|
|
}
|
|
|
|
func TestMarshalUnmarshalManifestPath(t *testing.T) {
|
|
original := "/base/path/to/manifest"
|
|
version := int64(42)
|
|
|
|
marshaled := MarshalManifestPath(original, version)
|
|
assert.NotEmpty(t, marshaled)
|
|
|
|
basePath, ver, err := UnmarshalManifestPath(marshaled)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, original, basePath)
|
|
assert.Equal(t, version, ver)
|
|
}
|
|
|
|
func TestUnmarshalManifestPath_InvalidJSON(t *testing.T) {
|
|
_, _, err := UnmarshalManifestPath("not valid json")
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestUnmarshalManifestPath_EmptyString(t *testing.T) {
|
|
_, _, err := UnmarshalManifestPath("")
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestMarshalManifestPath_EmptyBasePath(t *testing.T) {
|
|
marshaled := MarshalManifestPath("", 0)
|
|
basePath, ver, err := UnmarshalManifestPath(marshaled)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, "", basePath)
|
|
assert.Equal(t, int64(0), ver)
|
|
}
|
|
|
|
func TestCreateSegmentManifestWithBasePath_CanceledContext(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
|
|
_, err := CreateSegmentManifestWithBasePath(ctx, "/base", "parquet", []string{"col1"}, nil, nil)
|
|
assert.ErrorIs(t, err, context.Canceled)
|
|
}
|
|
|
|
// The cgo bridge (loon_properties_inject_external_spec) and Tier-1/2 endpoint
|
|
// derivation logic are covered by the C++ gtest suite in test_external_take.cpp.
|
|
// Go-side coverage is exercised indirectly through callers that invoke
|
|
// injectExternalSpecProperties (GetFileInfo, ExploreFilesReturnManifestPath,
|
|
// SampleExternalFieldSizes).
|
|
|
|
func TestExtfsPrefixForCollection(t *testing.T) {
|
|
assert.Equal(t, "extfs.42.", ExtfsPrefixForCollection(42))
|
|
assert.Equal(t, "extfs.0.", ExtfsPrefixForCollection(0))
|
|
}
|
|
|
|
// NOTE: the Go-side BuildExtfsOverrides / NormalizeExternalSource /
|
|
// effectiveSpecAddress helpers were collapsed into the C++ InjectExternalSpecProperties
|
|
// pipeline. Tier-1/2 endpoint derivation and AWS-form rewriting are covered by
|
|
// the C++ gtest suite in test_external_take.cpp. Go callers exercise the cgo
|
|
// bridge indirectly through the public FFI wrappers.
|
|
|
|
func TestFilterFileInfosByFormat(t *testing.T) {
|
|
files := []FileInfo{
|
|
{FilePath: "data/file1.parquet", NumRows: 100},
|
|
{FilePath: "data/file2.parquet", NumRows: 200},
|
|
{FilePath: "data/_SUCCESS", NumRows: 0},
|
|
{FilePath: "data/metadata.json", NumRows: 0},
|
|
{FilePath: "data/file3.csv", NumRows: 50},
|
|
{FilePath: "data/FILE4.PARQUET", NumRows: 300},
|
|
}
|
|
|
|
// Parquet: keep .parquet only (case-insensitive)
|
|
filtered, skipped := filterFileInfosByFormat(files, "parquet")
|
|
assert.Equal(t, 3, len(filtered))
|
|
assert.Equal(t, 3, skipped)
|
|
assert.Equal(t, "data/file1.parquet", filtered[0].FilePath)
|
|
assert.Equal(t, "data/FILE4.PARQUET", filtered[2].FilePath)
|
|
|
|
// Vortex: keep .vortex only
|
|
vortexFiles := []FileInfo{
|
|
{FilePath: "a.vortex"},
|
|
{FilePath: "b.parquet"},
|
|
{FilePath: "c.vortex"},
|
|
}
|
|
filtered, skipped = filterFileInfosByFormat(vortexFiles, "vortex")
|
|
assert.Equal(t, 2, len(filtered))
|
|
assert.Equal(t, 1, skipped)
|
|
|
|
// Lance-table: keep all (directory-based, no extension filtering)
|
|
filtered, skipped = filterFileInfosByFormat(files, "lance-table")
|
|
assert.Equal(t, len(files), len(filtered))
|
|
assert.Equal(t, 0, skipped)
|
|
|
|
// Unknown format: keep all
|
|
filtered, skipped = filterFileInfosByFormat(files, "unknown")
|
|
assert.Equal(t, len(files), len(filtered))
|
|
assert.Equal(t, 0, skipped)
|
|
|
|
// Empty list
|
|
filtered, skipped = filterFileInfosByFormat(nil, "parquet")
|
|
assert.Equal(t, 0, len(filtered))
|
|
assert.Equal(t, 0, skipped)
|
|
}
|
|
|
|
// Regression for index-drift bug: DataCoord and DataNode both call
|
|
// NormalizeFileInfos on the same raw manifest and MUST produce the
|
|
// identical indexed view, regardless of input ordering. Without sort,
|
|
// arrow's GetFileInfo can return files in different orders on each
|
|
// process and the [fileIndexBegin, fileIndexEnd) slice picks different
|
|
// files on the two sides, leading to silent data loss or "Invalid
|
|
// parquet magic" failures on Spark `_SUCCESS`/`.crc` strays.
|
|
func TestNormalizeFileInfos_StableSortAndFilter(t *testing.T) {
|
|
rawA := []FileInfo{
|
|
{FilePath: "data/_SUCCESS"},
|
|
{FilePath: "data/part-2.parquet"},
|
|
{FilePath: "data/.crc"},
|
|
{FilePath: "data/part-1.parquet"},
|
|
{FilePath: "data/README.md"},
|
|
}
|
|
// Reordered "DataNode side" view (arrow GetFileInfo gives no order).
|
|
rawB := []FileInfo{
|
|
{FilePath: "data/part-1.parquet"},
|
|
{FilePath: "data/README.md"},
|
|
{FilePath: "data/_SUCCESS"},
|
|
{FilePath: "data/part-2.parquet"},
|
|
{FilePath: "data/.crc"},
|
|
}
|
|
|
|
gotA, skippedA := NormalizeFileInfos(rawA, "parquet")
|
|
gotB, skippedB := NormalizeFileInfos(rawB, "parquet")
|
|
|
|
assert.Equal(t, 2, len(gotA))
|
|
assert.Equal(t, 3, skippedA)
|
|
assert.Equal(t, gotA, gotB, "two views of same manifest must produce identical indexed slice")
|
|
assert.Equal(t, "data/part-1.parquet", gotA[0].FilePath)
|
|
assert.Equal(t, "data/part-2.parquet", gotA[1].FilePath)
|
|
assert.Equal(t, skippedA, skippedB)
|
|
}
|
|
|
|
func TestMakePropertiesFromStorageConfig_ExtraKVsOverride(t *testing.T) {
|
|
// Test that extraKVs can add per-collection extfs properties
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "minio",
|
|
BucketName: "original-bucket",
|
|
Address: "localhost:9000",
|
|
RootPath: "/data",
|
|
}
|
|
prefix := ExtfsPrefixForCollection(42)
|
|
extra := map[string]string{
|
|
prefix + "bucket_name": "external-bucket",
|
|
prefix + "address": "http://s3.amazonaws.com",
|
|
}
|
|
props, err := MakePropertiesFromStorageConfig(config, extra)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, props)
|
|
defer FreeProperties(props)
|
|
}
|
|
|
|
func TestMakePropertiesFromStorageConfig_UsesConfiguredStorageFormat(t *testing.T) {
|
|
paramtable.Init()
|
|
params := paramtable.Get()
|
|
require.NoError(t, params.Save(params.DataNodeCfg.StorageFormat.Key, "parquet"))
|
|
t.Cleanup(func() {
|
|
_ = params.Reset(params.DataNodeCfg.StorageFormat.Key)
|
|
})
|
|
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: t.TempDir(),
|
|
RootPath: t.TempDir(),
|
|
}
|
|
|
|
props, err := MakePropertiesFromStorageConfig(config, nil)
|
|
require.NoError(t, err)
|
|
defer FreeProperties(props)
|
|
assert.Equal(t, "parquet", loonPropertyString(props, PropertyWriterFormat))
|
|
}
|
|
|
|
func TestMakePropertiesFromStorageConfig_ExtraKVsOverrideStorageFormat(t *testing.T) {
|
|
paramtable.Init()
|
|
params := paramtable.Get()
|
|
require.NoError(t, params.Save(params.DataNodeCfg.StorageFormat.Key, "vortex"))
|
|
t.Cleanup(func() {
|
|
_ = params.Reset(params.DataNodeCfg.StorageFormat.Key)
|
|
})
|
|
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: t.TempDir(),
|
|
RootPath: t.TempDir(),
|
|
}
|
|
|
|
props, err := MakePropertiesFromStorageConfig(config, map[string]string{
|
|
PropertyWriterFormat: "parquet",
|
|
})
|
|
require.NoError(t, err)
|
|
defer FreeProperties(props)
|
|
assert.Equal(t, "parquet", loonPropertyString(props, PropertyWriterFormat))
|
|
}
|
|
|
|
func TestNormalizeExternalPathForStorage_UsesInjectedExtfsEndpoint(t *testing.T) {
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: "/tmp",
|
|
RootPath: "/tmp",
|
|
}
|
|
extfs := ExternalSpecContext{
|
|
CollectionID: 42,
|
|
Source: "s3://liyiyang-test/all_types_v2/",
|
|
Spec: `{"format":"parquet","extfs":{"cloud_provider":"aws","region":"us-west-2","access_key_id":"ak","access_key_value":"sk"}}`,
|
|
}
|
|
|
|
props, err := MakePropertiesFromStorageConfig(config, nil)
|
|
require.NoError(t, err)
|
|
defer FreeProperties(props)
|
|
require.NoError(t, injectExternalSpecProperties(props, extfs.CollectionID, extfs.Source, extfs.Spec))
|
|
|
|
got, err := normalizeExternalPathForStorage(extfs.Source, props, extfs)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "s3://s3.us-west-2.amazonaws.com/liyiyang-test/all_types_v2/", got)
|
|
|
|
got, err = normalizeExternalPathForStorage("s3://liyiyang-test/a/b/file.parquet?versionId=1", props, extfs)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "s3://s3.us-west-2.amazonaws.com/liyiyang-test/a/b/file.parquet?versionId=1", got)
|
|
}
|
|
|
|
func TestNormalizeExternalPathForStorage_EndpointFormUnchanged(t *testing.T) {
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: "/tmp",
|
|
RootPath: "/tmp",
|
|
}
|
|
source := "s3://s3.us-west-2.amazonaws.com/liyiyang-test/all_types_v2/"
|
|
extfs := ExternalSpecContext{
|
|
CollectionID: 42,
|
|
Source: source,
|
|
Spec: `{"format":"parquet","extfs":{"cloud_provider":"aws","region":"us-west-2","access_key_id":"ak","access_key_value":"sk"}}`,
|
|
}
|
|
|
|
props, err := MakePropertiesFromStorageConfig(config, nil)
|
|
require.NoError(t, err)
|
|
defer FreeProperties(props)
|
|
require.NoError(t, injectExternalSpecProperties(props, extfs.CollectionID, extfs.Source, extfs.Spec))
|
|
|
|
got, err := normalizeExternalPathForStorage(source, props, extfs)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, source, got)
|
|
}
|
|
|
|
func TestNormalizeExternalPathForStorage_NoRewriteCases(t *testing.T) {
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: "/tmp",
|
|
RootPath: "/tmp",
|
|
}
|
|
prefix := ExtfsPrefixForCollection(42)
|
|
props, err := MakePropertiesFromStorageConfig(config, map[string]string{
|
|
prefix + "address": "https://s3.us-west-2.amazonaws.com",
|
|
prefix + "bucket_name": "liyiyang-test",
|
|
})
|
|
require.NoError(t, err)
|
|
defer FreeProperties(props)
|
|
|
|
path := "s3://liyiyang-test/all_types_v2/"
|
|
|
|
got, err := normalizeExternalPathForStorage(path, props, ExternalSpecContext{CollectionID: 42})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, path, got)
|
|
|
|
got, err = normalizeExternalPathForStorage(
|
|
path,
|
|
nil,
|
|
ExternalSpecContext{CollectionID: 42, Source: "s3://liyiyang-test/all_types_v2/"},
|
|
)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, path, got)
|
|
|
|
got, err = normalizeExternalPathForStorage(
|
|
"relative/path",
|
|
props,
|
|
ExternalSpecContext{CollectionID: 42, Source: "s3://liyiyang-test/all_types_v2/"},
|
|
)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "relative/path", got)
|
|
|
|
got, err = normalizeExternalPathForStorage(
|
|
"s3://other-bucket/all_types_v2/",
|
|
props,
|
|
ExternalSpecContext{CollectionID: 42, Source: "s3://liyiyang-test/all_types_v2/"},
|
|
)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "s3://other-bucket/all_types_v2/", got)
|
|
}
|
|
|
|
func TestNormalizeExternalPathForStorage_BareAddress(t *testing.T) {
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: "/tmp",
|
|
RootPath: "/tmp",
|
|
}
|
|
prefix := ExtfsPrefixForCollection(42)
|
|
props, err := MakePropertiesFromStorageConfig(config, map[string]string{
|
|
prefix + "address": "s3.us-west-2.amazonaws.com",
|
|
prefix + "bucket_name": "liyiyang-test",
|
|
})
|
|
require.NoError(t, err)
|
|
defer FreeProperties(props)
|
|
|
|
got, err := normalizeExternalPathForStorage(
|
|
"s3://liyiyang-test/all_types_v2/",
|
|
props,
|
|
ExternalSpecContext{CollectionID: 42, Source: "s3://liyiyang-test/all_types_v2/"},
|
|
)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "s3://s3.us-west-2.amazonaws.com/liyiyang-test/all_types_v2/", got)
|
|
}
|
|
|
|
func TestResolveExternalSourceRelativePath_UsesSourceBucketRoot(t *testing.T) {
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: "/tmp",
|
|
RootPath: "/tmp",
|
|
}
|
|
extfs := ExternalSpecContext{
|
|
CollectionID: 42,
|
|
Source: "s3://source-bucket/snapshots/100/metadata/200.json",
|
|
Spec: `{"format":"milvus-table","extfs":{"cloud_provider":"aws","region":"us-west-2","access_key_id":"ak","access_key_value":"sk"}}`,
|
|
}
|
|
|
|
props, err := MakePropertiesFromStorageConfig(config, nil)
|
|
require.NoError(t, err)
|
|
defer FreeProperties(props)
|
|
require.NoError(t, injectExternalSpecProperties(props, extfs.CollectionID, extfs.Source, extfs.Spec))
|
|
|
|
got, err := resolveExternalSourceRelativePath("files/insert_log/1/2/3", props, extfs)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "s3://s3.us-west-2.amazonaws.com/source-bucket/files/insert_log/1/2/3", got)
|
|
}
|
|
|
|
func TestNormalizeExternalPathForFilesystem_RemoteURI(t *testing.T) {
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: "/tmp",
|
|
RootPath: "/tmp",
|
|
}
|
|
prefix := ExtfsPrefixForCollection(42)
|
|
props, err := MakePropertiesFromStorageConfig(config, map[string]string{
|
|
prefix + "address": "http://localhost:9000",
|
|
prefix + "bucket_name": "a-bucket",
|
|
})
|
|
require.NoError(t, err)
|
|
defer FreeProperties(props)
|
|
|
|
lookupPath, filePath, err := normalizeExternalPathForFilesystem(
|
|
"s3://localhost:9000/a-bucket/files/snapshots/metadata.json",
|
|
props,
|
|
ExternalSpecContext{CollectionID: 42, Source: "s3://localhost:9000/a-bucket/files/"},
|
|
)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "s3://localhost:9000/a-bucket/files/snapshots/metadata.json", lookupPath)
|
|
assert.Equal(t, "files/snapshots/metadata.json", filePath)
|
|
|
|
lookupPath, filePath, err = normalizeExternalPathForFilesystem(
|
|
"s3://a-bucket/files/snapshots/metadata.json",
|
|
props,
|
|
ExternalSpecContext{CollectionID: 42, Source: "s3://a-bucket/files/"},
|
|
)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "s3://localhost:9000/a-bucket/files/snapshots/metadata.json", lookupPath)
|
|
assert.Equal(t, "files/snapshots/metadata.json", filePath)
|
|
|
|
lookupPath, filePath, err = normalizeExternalPathForFilesystem(
|
|
"s3://other-bucket/files/snapshots/metadata.json",
|
|
props,
|
|
ExternalSpecContext{CollectionID: 42, Source: "s3://a-bucket/files/"},
|
|
)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "s3://other-bucket/files/snapshots/metadata.json", lookupPath)
|
|
assert.Equal(t, "s3://other-bucket/files/snapshots/metadata.json", filePath)
|
|
|
|
lookupPath, filePath, err = normalizeExternalPathForFilesystem(
|
|
"files/source/_delta/9001",
|
|
props,
|
|
ExternalSpecContext{CollectionID: 42, Source: "s3://a-bucket/snapshots/100/metadata/200.json"},
|
|
)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "s3://localhost:9000/a-bucket/files/source/_delta/9001", lookupPath)
|
|
assert.Equal(t, "files/source/_delta/9001", filePath)
|
|
}
|
|
|
|
// ==================== SampleExternalFieldSizes Tests ====================
|
|
|
|
func TestSampleExternalFieldSizes_NilStorageConfig(t *testing.T) {
|
|
result, err := SampleExternalFieldSizes(
|
|
`{"base_path":"/tmp","ver":1}`, 100, 42,
|
|
"s3://bucket/data/", `{"format":"parquet"}`,
|
|
nil,
|
|
nil,
|
|
)
|
|
assert.Nil(t, result)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "storageConfig is required")
|
|
}
|
|
|
|
func TestSampleExternalFieldSizes_EmptyManifestPath(t *testing.T) {
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: "/tmp",
|
|
RootPath: "/tmp",
|
|
}
|
|
result, err := SampleExternalFieldSizes(
|
|
"", 100, 42,
|
|
"", "",
|
|
nil,
|
|
config,
|
|
)
|
|
assert.Nil(t, result)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "manifest_path is empty")
|
|
}
|
|
|
|
func TestSampleExternalFieldSizes_InvalidManifestPath(t *testing.T) {
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: "/tmp",
|
|
RootPath: "/tmp",
|
|
}
|
|
// Valid JSON but manifest file doesn't exist
|
|
result, err := SampleExternalFieldSizes(
|
|
`{"base_path":"/nonexistent/path","ver":1}`, 100, 42,
|
|
"", "",
|
|
nil,
|
|
config,
|
|
)
|
|
assert.Nil(t, result)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestSampleExternalFieldSizes_WithSpec(t *testing.T) {
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: "/tmp",
|
|
RootPath: "/tmp",
|
|
}
|
|
// Will fail at manifest read (no real file), but validates that
|
|
// properties construction with extfs injection via cgo doesn't panic.
|
|
result, err := SampleExternalFieldSizes(
|
|
`{"base_path":"/nonexistent","ver":1}`, 100, 42,
|
|
"s3://s3.us-west-2.amazonaws.com/ext-bucket/data/",
|
|
`{"format":"parquet","extfs":{"region":"us-west-2","use_ssl":"true"}}`,
|
|
nil,
|
|
config,
|
|
)
|
|
assert.Nil(t, result)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
// ==================== ExploreFilesReturnManifestPath Tests ====================
|
|
|
|
func TestExploreFilesReturnManifestPath_InvalidDir(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{StorageType: "local", BucketName: tmpDir, RootPath: tmpDir}
|
|
|
|
_, _, err := ExploreFilesReturnManifestPath(
|
|
[]string{"col1"}, "parquet", "/nonexistent/base", "/nonexistent/explore", config, ExternalSpecContext{},
|
|
)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestExploreFilesReturnManifestPath_EmptyColumns(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{StorageType: "local", BucketName: tmpDir, RootPath: tmpDir}
|
|
|
|
_, _, err := ExploreFilesReturnManifestPath(
|
|
[]string{}, "parquet", "/nonexistent", "/nonexistent", config, ExternalSpecContext{},
|
|
)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestExploreFilesReturnManifestPath_PropertiesError(t *testing.T) {
|
|
_, _, err := ExploreFilesReturnManifestPath(
|
|
[]string{"col1"}, "parquet", "/base", "/explore", nil, ExternalSpecContext{},
|
|
)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "properties")
|
|
}
|
|
|
|
func TestResolveMilvusTableSnapshotMetadataPathRejectsNonJSONPath(t *testing.T) {
|
|
spec := `{"format":"milvus-table"}`
|
|
|
|
t.Run("accepts snapshot metadata JSON path", func(t *testing.T) {
|
|
metadataPath, err := resolveMilvusTableSnapshotMetadataPath("s3://bucket/snapshots/100/metadata/200.json", spec)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "s3://bucket/snapshots/100/metadata/200.json", metadataPath)
|
|
})
|
|
|
|
t.Run("rejects base path", func(t *testing.T) {
|
|
_, err := resolveMilvusTableSnapshotMetadataPath("s3://bucket", spec)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "snapshot metadata JSON path")
|
|
})
|
|
}
|
|
|
|
// ==================== ReadFileInfosFromManifestPath Tests ====================
|
|
|
|
func TestReadFileInfosFromManifestPath_InvalidPath(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{StorageType: "local", BucketName: tmpDir, RootPath: tmpDir}
|
|
|
|
_, err := ReadFileInfosFromManifestPath("/nonexistent/manifest.json", config)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestReadFileInfosFromManifestPath_NilConfig(t *testing.T) {
|
|
_, err := ReadFileInfosFromManifestPath("/manifest.json", nil)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "properties")
|
|
}
|
|
|
|
func TestBuildMilvusTableFileInfosFromSnapshotMetadata(t *testing.T) {
|
|
metadata := &datapb.SnapshotMetadata{
|
|
FormatVersion: 2,
|
|
ManifestList: []string{
|
|
"segment-20.avro",
|
|
"segment-10.avro",
|
|
},
|
|
Storagev2ManifestList: []*datapb.StorageV2SegmentManifest{
|
|
{SegmentId: 20, Manifest: `{"base_path":"source/20","ver":1}`},
|
|
{SegmentId: 10, Manifest: `{"base_path":"source/10","ver":2}`},
|
|
},
|
|
}
|
|
metadataBytes, err := protojson.MarshalOptions{UseProtoNames: true}.Marshal(metadata)
|
|
require.NoError(t, err)
|
|
|
|
segments := map[string]*datapb.SegmentDescription{
|
|
"segment-10.avro": {SegmentId: 10, NumOfRows: 100},
|
|
"segment-20.avro": {SegmentId: 20, NumOfRows: 200},
|
|
}
|
|
fileInfos, err := buildMilvusTableFileInfosFromSnapshotMetadata(metadataBytes, func(manifestPath string, formatVersion int32) (*datapb.SegmentDescription, error) {
|
|
require.Equal(t, int32(2), formatVersion)
|
|
return segments[manifestPath], nil
|
|
}, nil)
|
|
|
|
require.NoError(t, err)
|
|
require.Equal(t, []FileInfo{
|
|
{FilePath: `{"base_path":"source/10","ver":2}`, NumRows: 100, SourceSegmentID: 10},
|
|
{FilePath: `{"base_path":"source/20","ver":1}`, NumRows: 200, SourceSegmentID: 20},
|
|
}, fileInfos)
|
|
}
|
|
|
|
func TestBuildMilvusTableFileInfosFromSnapshotMetadata_ResolvesManifestBasePath(t *testing.T) {
|
|
metadata := &datapb.SnapshotMetadata{
|
|
FormatVersion: 2,
|
|
ManifestList: []string{"segment-manifest.avro"},
|
|
Storagev2ManifestList: []*datapb.StorageV2SegmentManifest{
|
|
{SegmentId: 10, Manifest: `{"base_path":"files/insert_log/10","ver":2}`},
|
|
},
|
|
}
|
|
metadataBytes, err := protojson.MarshalOptions{UseProtoNames: true}.Marshal(metadata)
|
|
require.NoError(t, err)
|
|
|
|
resolvedManifest := MarshalManifestPath("s3://source-bucket/files/insert_log/10", 2)
|
|
fileInfos, err := buildMilvusTableFileInfosFromSnapshotMetadata(
|
|
metadataBytes,
|
|
func(manifestPath string, formatVersion int32) (*datapb.SegmentDescription, error) {
|
|
require.Equal(t, "segment-manifest.avro", manifestPath)
|
|
require.Equal(t, int32(2), formatVersion)
|
|
return &datapb.SegmentDescription{
|
|
SegmentId: 10,
|
|
NumOfRows: 100,
|
|
}, nil
|
|
},
|
|
func(manifestPath string) (string, error) {
|
|
return resolveMilvusTableSourceManifestPath(manifestPath, func(sourcePath string) (string, error) {
|
|
return "s3://source-bucket/" + sourcePath, nil
|
|
})
|
|
},
|
|
)
|
|
|
|
require.NoError(t, err)
|
|
require.Equal(t, []FileInfo{
|
|
{FilePath: resolvedManifest, NumRows: 100, SourceSegmentID: 10},
|
|
}, fileInfos)
|
|
}
|
|
|
|
func TestResolveMilvusTableSegmentDeltalogPaths(t *testing.T) {
|
|
segment := &datapb.SegmentDescription{
|
|
Deltalogs: []*datapb.FieldBinlog{{
|
|
Binlogs: []*datapb.Binlog{{
|
|
LogPath: "files/delta_log/10/1",
|
|
}},
|
|
}},
|
|
}
|
|
|
|
err := resolveMilvusTableSegmentDeltalogPaths(segment, func(sourcePath string) (string, error) {
|
|
return "s3://source-bucket/" + sourcePath, nil
|
|
})
|
|
|
|
require.NoError(t, err)
|
|
require.Equal(t, "s3://source-bucket/files/delta_log/10/1", segment.GetDeltalogs()[0].GetBinlogs()[0].GetLogPath())
|
|
}
|
|
|
|
func TestBuildMilvusTableFileInfosFromSnapshotMetadata_NoStorageV2(t *testing.T) {
|
|
metadataBytes, err := protojson.MarshalOptions{UseProtoNames: true}.Marshal(&datapb.SnapshotMetadata{})
|
|
require.NoError(t, err)
|
|
|
|
_, err = buildMilvusTableFileInfosFromSnapshotMetadata(metadataBytes, nil, nil)
|
|
|
|
require.Error(t, err)
|
|
assert.True(t, IsMilvusTableStorageV2ManifestListMissing(err))
|
|
assert.Contains(t, err.Error(), "storagev2_manifest_list")
|
|
assert.Contains(t, err.Error(), "common.storage.useLoonFFI=true")
|
|
}
|
|
|
|
func TestBuildMilvusTableFileInfosFromSnapshotMetadata_SourceSegmentDeltalogsSkipped(t *testing.T) {
|
|
metadata := &datapb.SnapshotMetadata{
|
|
FormatVersion: 2,
|
|
ManifestList: []string{"segment-manifest.avro"},
|
|
Storagev2ManifestList: []*datapb.StorageV2SegmentManifest{
|
|
{SegmentId: 10, Manifest: `{"base_path":"source/10","ver":2}`},
|
|
},
|
|
}
|
|
metadataBytes, err := protojson.MarshalOptions{UseProtoNames: true}.Marshal(metadata)
|
|
require.NoError(t, err)
|
|
|
|
expected := []*datapb.FieldBinlog{{
|
|
Binlogs: []*datapb.Binlog{{LogPath: "files/source/10/_delta/1", EntriesNum: 8}},
|
|
}}
|
|
fileInfos, err := buildMilvusTableFileInfosFromSnapshotMetadata(metadataBytes, func(manifestPath string, formatVersion int32) (*datapb.SegmentDescription, error) {
|
|
require.Equal(t, "segment-manifest.avro", manifestPath)
|
|
require.Equal(t, int32(2), formatVersion)
|
|
return &datapb.SegmentDescription{
|
|
SegmentId: 10,
|
|
SegmentLevel: datapb.SegmentLevel_L1,
|
|
NumOfRows: 100,
|
|
Deltalogs: expected,
|
|
}, nil
|
|
}, nil)
|
|
|
|
require.NoError(t, err)
|
|
require.Len(t, fileInfos, 1)
|
|
require.Equal(t, int64(10), fileInfos[0].SourceSegmentID)
|
|
require.Empty(t, fileInfos[0].Deltalogs)
|
|
}
|
|
|
|
func TestBuildMilvusTableFileInfosFromSnapshotMetadata_L0Deltalogs(t *testing.T) {
|
|
metadata := &datapb.SnapshotMetadata{
|
|
FormatVersion: 2,
|
|
ManifestList: []string{"source-manifest.avro", "l0-manifest.avro"},
|
|
Storagev2ManifestList: []*datapb.StorageV2SegmentManifest{
|
|
{SegmentId: 10, Manifest: `{"base_path":"source/10","ver":2}`},
|
|
},
|
|
}
|
|
metadataBytes, err := protojson.MarshalOptions{UseProtoNames: true}.Marshal(metadata)
|
|
require.NoError(t, err)
|
|
|
|
l0Deltalogs := []*datapb.FieldBinlog{{
|
|
FieldID: 100,
|
|
Binlogs: []*datapb.Binlog{{
|
|
LogPath: "files/delta_log/1",
|
|
EntriesNum: 8,
|
|
}},
|
|
}}
|
|
fileInfos, err := buildMilvusTableFileInfosFromSnapshotMetadata(metadataBytes, func(manifestPath string, formatVersion int32) (*datapb.SegmentDescription, error) {
|
|
require.Equal(t, int32(2), formatVersion)
|
|
switch manifestPath {
|
|
case "source-manifest.avro":
|
|
return &datapb.SegmentDescription{
|
|
SegmentId: 10,
|
|
SegmentLevel: datapb.SegmentLevel_L1,
|
|
NumOfRows: 100,
|
|
}, nil
|
|
case "l0-manifest.avro":
|
|
return &datapb.SegmentDescription{
|
|
SegmentId: 20,
|
|
SegmentLevel: datapb.SegmentLevel_L0,
|
|
Deltalogs: l0Deltalogs,
|
|
}, nil
|
|
default:
|
|
return nil, fmt.Errorf("unexpected manifest path %s", manifestPath)
|
|
}
|
|
}, nil)
|
|
|
|
require.NoError(t, err)
|
|
require.Len(t, fileInfos, 1)
|
|
require.Equal(t, int64(10), fileInfos[0].SourceSegmentID)
|
|
require.Equal(t, l0Deltalogs, fileInfos[0].Deltalogs)
|
|
}
|
|
|
|
func TestBuildMilvusTableFileInfosFromSnapshotMetadata_DeltalogsBySegmentLevel(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name string
|
|
manifestList []string
|
|
wantDeltalogs bool
|
|
}{
|
|
{name: "source_segment", manifestList: []string{"source-manifest.avro"}, wantDeltalogs: false},
|
|
{name: "l0_overlay", manifestList: []string{"source-manifest.avro", "l0-manifest.avro"}, wantDeltalogs: true},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
metadata := &datapb.SnapshotMetadata{
|
|
FormatVersion: 2,
|
|
ManifestList: tc.manifestList,
|
|
Storagev2ManifestList: []*datapb.StorageV2SegmentManifest{
|
|
{SegmentId: 10, Manifest: `{"base_path":"source/10","ver":2}`},
|
|
},
|
|
}
|
|
metadataBytes, err := protojson.MarshalOptions{UseProtoNames: true}.Marshal(metadata)
|
|
require.NoError(t, err)
|
|
|
|
fileInfos, err := buildMilvusTableFileInfosFromSnapshotMetadata(
|
|
metadataBytes,
|
|
func(manifestPath string, formatVersion int32) (*datapb.SegmentDescription, error) {
|
|
require.Equal(t, int32(2), formatVersion)
|
|
switch manifestPath {
|
|
case "source-manifest.avro":
|
|
return &datapb.SegmentDescription{
|
|
SegmentId: 10,
|
|
SegmentLevel: datapb.SegmentLevel_L1,
|
|
NumOfRows: 100,
|
|
}, nil
|
|
case "l0-manifest.avro":
|
|
return &datapb.SegmentDescription{
|
|
SegmentId: 20,
|
|
SegmentLevel: datapb.SegmentLevel_L0,
|
|
Deltalogs: []*datapb.FieldBinlog{{
|
|
FieldID: 100,
|
|
Binlogs: []*datapb.Binlog{{
|
|
LogPath: "files/delta_log/1",
|
|
EntriesNum: 1,
|
|
}},
|
|
}},
|
|
}, nil
|
|
default:
|
|
return nil, fmt.Errorf("unexpected manifest path %s", manifestPath)
|
|
}
|
|
},
|
|
nil,
|
|
)
|
|
|
|
require.NoError(t, err)
|
|
require.Len(t, fileInfos, 1)
|
|
if tc.wantDeltalogs {
|
|
require.Len(t, fileInfos[0].Deltalogs, 1)
|
|
require.Len(t, fileInfos[0].Deltalogs[0].GetBinlogs(), 1)
|
|
assert.Equal(t, "files/delta_log/1", fileInfos[0].Deltalogs[0].GetBinlogs()[0].GetLogPath())
|
|
} else {
|
|
assert.Empty(t, fileInfos[0].Deltalogs)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestReadMilvusSnapshotSegmentManifest_Deltalogs(t *testing.T) {
|
|
data, err := snapshotio.MarshalSegmentManifest(&datapb.SegmentDescription{
|
|
SegmentId: 100,
|
|
PartitionId: 10,
|
|
SegmentLevel: datapb.SegmentLevel_L1,
|
|
ChannelName: "by-dev-rootcoord-dml_0_100v0",
|
|
NumOfRows: 128,
|
|
StorageVersion: 3,
|
|
IsSorted: true,
|
|
Deltalogs: []*datapb.FieldBinlog{{
|
|
FieldID: 100,
|
|
Binlogs: []*datapb.Binlog{{
|
|
EntriesNum: 8,
|
|
TimestampFrom: 1,
|
|
TimestampTo: 2,
|
|
LogPath: "delta/100/1",
|
|
LogSize: 64,
|
|
LogID: 1,
|
|
MemorySize: 64,
|
|
}},
|
|
}},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
segment, err := readMilvusSnapshotSegmentManifest("manifest.avro", snapshotio.SnapshotFormatVersion, func(path string) ([]byte, error) {
|
|
require.Equal(t, "manifest.avro", path)
|
|
return data, nil
|
|
})
|
|
require.NoError(t, err)
|
|
require.Equal(t, int64(100), segment.GetSegmentId())
|
|
require.Len(t, segment.GetDeltalogs(), 1)
|
|
require.Len(t, segment.GetDeltalogs()[0].GetBinlogs(), 1)
|
|
require.Equal(t, "delta/100/1", segment.GetDeltalogs()[0].GetBinlogs()[0].GetLogPath())
|
|
}
|
|
|
|
func TestReadMilvusSnapshotSegmentManifest_ParseErrorIncludesPath(t *testing.T) {
|
|
_, err := readMilvusSnapshotSegmentManifest("bad-manifest.avro", snapshotio.SnapshotFormatVersion, func(path string) ([]byte, error) {
|
|
require.Equal(t, "bad-manifest.avro", path)
|
|
return []byte("not a segment manifest"), nil
|
|
})
|
|
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "parse source segment manifest bad-manifest.avro")
|
|
}
|
|
|
|
// ==================== GetFileInfo Additional Tests ====================
|
|
|
|
func TestGetFileInfo_PropertiesError(t *testing.T) {
|
|
_, err := GetFileInfo("parquet", "/some/file.parquet", nil, ExternalSpecContext{})
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "properties")
|
|
}
|
|
|
|
func TestMakePropertiesFromStorageConfig_NilConfig(t *testing.T) {
|
|
_, err := MakePropertiesFromStorageConfig(nil, nil)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "storageConfig is required")
|
|
}
|
|
|
|
func TestMakePropertiesFromStorageConfig_LocalType(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: tmpDir,
|
|
RootPath: tmpDir,
|
|
}
|
|
props, err := MakePropertiesFromStorageConfig(config, nil)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, props)
|
|
defer FreeProperties(props)
|
|
}
|
|
|
|
func TestMakePropertiesFromStorageConfig_AllFields(t *testing.T) {
|
|
config := &indexpb.StorageConfig{
|
|
Address: "localhost:9000",
|
|
BucketName: "test-bucket",
|
|
AccessKeyID: "access-key",
|
|
SecretAccessKey: "secret-key",
|
|
RootPath: "/data",
|
|
StorageType: "minio",
|
|
CloudProvider: "aws",
|
|
IAMEndpoint: "http://iam.example.com",
|
|
Region: "us-east-1",
|
|
SslCACert: "cert-content",
|
|
GcpCredentialJSON: `{"type":"service_account"}`,
|
|
UseSSL: true,
|
|
UseIAM: true,
|
|
UseVirtualHost: true,
|
|
RequestTimeoutMs: 5000,
|
|
}
|
|
props, err := MakePropertiesFromStorageConfig(config, nil)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, props)
|
|
defer FreeProperties(props)
|
|
}
|
|
|
|
// ==================== FetchFragmentsFromExternalSourceWithRange Tests ====================
|
|
|
|
func TestFetchFragmentsFromExternalSourceWithRange_HappyPath(t *testing.T) {
|
|
ctx := context.Background()
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{StorageType: "local", BucketName: tmpDir, RootPath: tmpDir}
|
|
|
|
// Mock ReadFileInfosFromManifestPath to return 3 files
|
|
mockRead := mockey.Mock(ReadFileInfosFromManifestPath).Return([]FileInfo{
|
|
{FilePath: "f1.parquet", NumRows: 1000},
|
|
{FilePath: "f2.parquet", NumRows: 2000},
|
|
{FilePath: "f3.parquet", NumRows: 3000},
|
|
}, nil).Build()
|
|
defer mockRead.UnPatch()
|
|
|
|
fragments, err := FetchFragmentsFromExternalSourceWithRange(
|
|
ctx, "parquet", []string{"col1"}, "s3://endpoint/bucket/path", config,
|
|
0, 3, "/manifest.json", ExternalFetchOptions{CollectionID: 1, RowLimit: 5000},
|
|
)
|
|
assert.NoError(t, err)
|
|
assert.NotEmpty(t, fragments)
|
|
}
|
|
|
|
func TestFetchFragmentsFromExternalSourceWithRange_EmptyManifest(t *testing.T) {
|
|
ctx := context.Background()
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{StorageType: "local", BucketName: tmpDir, RootPath: tmpDir}
|
|
|
|
_, err := FetchFragmentsFromExternalSourceWithRange(
|
|
ctx, "parquet", nil, "", config,
|
|
0, 10, "", ExternalFetchOptions{},
|
|
)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "manifest path is required")
|
|
}
|
|
|
|
func TestFetchFragmentsFromExternalSourceWithRange_ReadManifestFailed(t *testing.T) {
|
|
ctx := context.Background()
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{StorageType: "local", BucketName: tmpDir, RootPath: tmpDir}
|
|
|
|
mockRead := mockey.Mock(ReadFileInfosFromManifestPath).
|
|
Return(nil, fmt.Errorf("manifest read error")).Build()
|
|
defer mockRead.UnPatch()
|
|
|
|
_, err := FetchFragmentsFromExternalSourceWithRange(
|
|
ctx, "parquet", nil, "", config,
|
|
0, 10, "/manifest.json", ExternalFetchOptions{},
|
|
)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "manifest")
|
|
}
|
|
|
|
func TestFetchFragmentsFromExternalSourceWithRange_IndexOutOfRange(t *testing.T) {
|
|
ctx := context.Background()
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{StorageType: "local", BucketName: tmpDir, RootPath: tmpDir}
|
|
|
|
mockRead := mockey.Mock(ReadFileInfosFromManifestPath).Return([]FileInfo{
|
|
{FilePath: "f1.parquet", NumRows: 100},
|
|
}, nil).Build()
|
|
defer mockRead.UnPatch()
|
|
|
|
_, err := FetchFragmentsFromExternalSourceWithRange(
|
|
ctx, "parquet", nil, "", config,
|
|
5, 10, "/manifest.json", ExternalFetchOptions{},
|
|
)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "fileIndexBegin")
|
|
}
|
|
|
|
func TestFetchFragmentsFromExternalSourceWithRange_NeedsFileInfo(t *testing.T) {
|
|
ctx := context.Background()
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{StorageType: "local", BucketName: tmpDir, RootPath: tmpDir}
|
|
|
|
// File with NumRows=0 needs GetFileInfo
|
|
mockRead := mockey.Mock(ReadFileInfosFromManifestPath).Return([]FileInfo{
|
|
{FilePath: "f1.parquet", NumRows: 0},
|
|
}, nil).Build()
|
|
defer mockRead.UnPatch()
|
|
|
|
mockGetInfo := mockey.Mock(GetFileInfo).Return(&FileInfo{
|
|
FilePath: "f1.parquet", NumRows: 500,
|
|
}, nil).Build()
|
|
defer mockGetInfo.UnPatch()
|
|
|
|
fragments, err := FetchFragmentsFromExternalSourceWithRange(
|
|
ctx, "parquet", nil, "", config,
|
|
0, 1, "/manifest.json", ExternalFetchOptions{CollectionID: 1, RowLimit: 1000},
|
|
)
|
|
assert.NoError(t, err)
|
|
assert.Len(t, fragments, 1)
|
|
assert.Equal(t, int64(500), fragments[0].RowCount)
|
|
}
|
|
|
|
func TestFetchFragmentsFromExternalSourceWithRange_GetFileInfoFailed(t *testing.T) {
|
|
ctx := context.Background()
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{StorageType: "local", BucketName: tmpDir, RootPath: tmpDir}
|
|
|
|
mockRead := mockey.Mock(ReadFileInfosFromManifestPath).Return([]FileInfo{
|
|
{FilePath: "f1.parquet", NumRows: 0},
|
|
}, nil).Build()
|
|
defer mockRead.UnPatch()
|
|
|
|
mockGetInfo := mockey.Mock(GetFileInfo).
|
|
Return(nil, fmt.Errorf("get file info error")).Build()
|
|
defer mockGetInfo.UnPatch()
|
|
|
|
_, err := FetchFragmentsFromExternalSourceWithRange(
|
|
ctx, "parquet", nil, "", config,
|
|
0, 1, "/manifest.json", ExternalFetchOptions{CollectionID: 1},
|
|
)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "get file info")
|
|
}
|
|
|
|
func TestFetchFragmentsFromExternalSourceWithRange_FileIndexEndClamped(t *testing.T) {
|
|
ctx := context.Background()
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{StorageType: "local", BucketName: tmpDir, RootPath: tmpDir}
|
|
|
|
mockRead := mockey.Mock(ReadFileInfosFromManifestPath).Return([]FileInfo{
|
|
{FilePath: "f1.parquet", NumRows: 100},
|
|
{FilePath: "f2.parquet", NumRows: 200},
|
|
}, nil).Build()
|
|
defer mockRead.UnPatch()
|
|
|
|
// fileIndexEnd (99) > len(fileInfos) (2), should be clamped
|
|
fragments, err := FetchFragmentsFromExternalSourceWithRange(
|
|
ctx, "parquet", nil, "", config,
|
|
0, 99, "/manifest.json", ExternalFetchOptions{CollectionID: 1, RowLimit: 1000},
|
|
)
|
|
assert.NoError(t, err)
|
|
assert.NotEmpty(t, fragments)
|
|
}
|
|
|
|
func TestMakePropertiesFromStorageConfig_WithExtraKVs(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
config := &indexpb.StorageConfig{
|
|
StorageType: "local",
|
|
BucketName: tmpDir,
|
|
RootPath: tmpDir,
|
|
}
|
|
extra := map[string]string{
|
|
"writer.policy": "split",
|
|
"custom.key": "custom.value",
|
|
}
|
|
props, err := MakePropertiesFromStorageConfig(config, extra)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, props)
|
|
defer FreeProperties(props)
|
|
}
|
|
|
|
// ==================== rowLimitOrDefault Tests ====================
|
|
|
|
func TestRowLimitOrDefault(t *testing.T) {
|
|
t.Run("zero falls back to default", func(t *testing.T) {
|
|
o := ExternalFetchOptions{RowLimit: 0}
|
|
assert.Equal(t, int64(DefaultFragmentRowLimit), o.rowLimitOrDefault())
|
|
})
|
|
t.Run("negative falls back to default", func(t *testing.T) {
|
|
o := ExternalFetchOptions{RowLimit: -1}
|
|
assert.Equal(t, int64(DefaultFragmentRowLimit), o.rowLimitOrDefault())
|
|
})
|
|
t.Run("positive used as-is", func(t *testing.T) {
|
|
o := ExternalFetchOptions{RowLimit: 1234}
|
|
assert.Equal(t, int64(1234), o.rowLimitOrDefault())
|
|
})
|
|
}
|
|
|
|
// ==================== fetchRowCountsConcurrently Tests ====================
|
|
|
|
func TestFetchRowCountsConcurrently_AllKnownSkipsFFI(t *testing.T) {
|
|
// All fileInfos already have NumRows > 0 → no GetFileInfo call.
|
|
mockCalled := int32(0)
|
|
m := mockey.Mock(GetFileInfo).To(func(format, path string, cfg *indexpb.StorageConfig, extfs ExternalSpecContext) (*FileInfo, error) {
|
|
atomic.AddInt32(&mockCalled, 1)
|
|
return &FileInfo{FilePath: path, NumRows: 99}, nil
|
|
}).Build()
|
|
defer m.UnPatch()
|
|
|
|
fileInfos := []FileInfo{
|
|
{FilePath: "a", NumRows: 10},
|
|
{FilePath: "b", NumRows: 20},
|
|
{FilePath: "c", NumRows: 30},
|
|
}
|
|
got, err := fetchRowCountsConcurrently(context.Background(), "parquet", fileInfos, &indexpb.StorageConfig{}, ExternalSpecContext{})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, []int64{10, 20, 30}, got)
|
|
assert.Equal(t, int32(0), atomic.LoadInt32(&mockCalled))
|
|
}
|
|
|
|
func TestFetchRowCountsConcurrently_MissingFilled(t *testing.T) {
|
|
m := mockey.Mock(GetFileInfo).To(func(format, path string, cfg *indexpb.StorageConfig, extfs ExternalSpecContext) (*FileInfo, error) {
|
|
return &FileInfo{FilePath: path, NumRows: 777}, nil
|
|
}).Build()
|
|
defer m.UnPatch()
|
|
|
|
fileInfos := []FileInfo{
|
|
{FilePath: "a", NumRows: 10},
|
|
{FilePath: "b", NumRows: 0},
|
|
{FilePath: "c", NumRows: -1},
|
|
}
|
|
got, err := fetchRowCountsConcurrently(context.Background(), "parquet", fileInfos, &indexpb.StorageConfig{}, ExternalSpecContext{})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, []int64{10, 777, 777}, got)
|
|
}
|
|
|
|
func TestFetchRowCountsConcurrently_WorkerErrorPropagated(t *testing.T) {
|
|
m := mockey.Mock(GetFileInfo).Return(nil, fmt.Errorf("fake ffi failure")).Build()
|
|
defer m.UnPatch()
|
|
|
|
_, err := fetchRowCountsConcurrently(context.Background(), "parquet",
|
|
[]FileInfo{{FilePath: "x", NumRows: 0}}, &indexpb.StorageConfig{}, ExternalSpecContext{})
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "failed to get file info")
|
|
}
|
|
|
|
func TestFetchRowCountsConcurrently_CtxPreCancelled(t *testing.T) {
|
|
// Mock returns nil synchronously; ctx is canceled before launching workers.
|
|
m := mockey.Mock(GetFileInfo).Return(&FileInfo{NumRows: 10}, nil).Build()
|
|
defer m.UnPatch()
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
|
|
_, err := fetchRowCountsConcurrently(ctx, "parquet",
|
|
[]FileInfo{{FilePath: "x", NumRows: 0}}, &indexpb.StorageConfig{}, ExternalSpecContext{})
|
|
require.Error(t, err)
|
|
assert.ErrorIs(t, err, context.Canceled)
|
|
}
|
|
|
|
// TestFetchRowCountsConcurrently_CtxCancelledDuringRun verifies the
|
|
// post-AwaitAll ctx.Err() guard: workers ignore ctx and return success, but
|
|
// the ctx is canceled mid-run — the function must still return ctx.Err()
|
|
// rather than a half-filled rowCounts slice.
|
|
func TestFetchRowCountsConcurrently_CtxCancelledDuringRun(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
m := mockey.Mock(GetFileInfo).To(func(_ string, _ string, _ *indexpb.StorageConfig, _ ExternalSpecContext) (*FileInfo, error) {
|
|
cancel()
|
|
return &FileInfo{NumRows: 10}, nil
|
|
}).Build()
|
|
defer m.UnPatch()
|
|
|
|
_, err := fetchRowCountsConcurrently(ctx, "parquet",
|
|
[]FileInfo{{FilePath: "x", NumRows: 0}}, &indexpb.StorageConfig{}, ExternalSpecContext{})
|
|
require.Error(t, err)
|
|
assert.ErrorIs(t, err, context.Canceled)
|
|
}
|