Files
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
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

296 lines
10 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
package packed
import (
"path"
"testing"
"github.com/apache/arrow/go/v17/arrow"
"github.com/apache/arrow/go/v17/arrow/array"
"github.com/apache/arrow/go/v17/arrow/memory"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus/internal/storagecommon"
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
// manifestTestStorageConfig returns a storage config wired to a per-test
// temp dir for use with CommitManifestUpdates / FFIPackedWriter.
func manifestTestStorageConfig(t *testing.T) *indexpb.StorageConfig {
t.Helper()
paramtable.Init()
pt := paramtable.Get()
pt.Save(pt.CommonCfg.StorageType.Key, "local")
dir := t.TempDir()
pt.Save(pt.LocalStorageCfg.Path.Key, dir)
t.Cleanup(func() {
pt.Reset(pt.CommonCfg.StorageType.Key)
pt.Reset(pt.LocalStorageCfg.Path.Key)
})
return &indexpb.StorageConfig{StorageType: "local", RootPath: dir}
}
// TestCommitManifestUpdates_EmptyShortCircuit verifies that an empty
// ManifestUpdates returns the unchanged manifest path without opening a
// loon transaction (and therefore without bumping the version).
func TestCommitManifestUpdates_EmptyShortCircuit(t *testing.T) {
cfg := manifestTestStorageConfig(t)
basePath := "files/commit_empty_test/seg1"
cases := []struct {
name string
u *ManifestUpdates
}{
{"nil-updates", nil},
{"zero-updates", &ManifestUpdates{}},
{"empty-slices", &ManifestUpdates{DeltaLogs: []DeltaLogEntry{}, Stats: []StatEntry{}}},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
got, err := CommitManifestUpdates(basePath, ManifestEarliest, cfg, tc.u)
require.NoError(t, err)
require.Equal(t, MarshalManifestPath(basePath, ManifestEarliest), got,
"empty updates must return the unchanged manifest path")
})
}
}
// TestCommitManifestUpdates_StatsOnly exercises the stats-without-inserts
// path: a stat blob is written to storage and registered in the manifest
// via a single CommitManifestUpdates call.
func TestCommitManifestUpdates_StatsOnly(t *testing.T) {
cfg := manifestTestStorageConfig(t)
basePath := "files/commit_stats_only/seg1"
statPath := path.Join(cfg.RootPath, basePath, "_stats/bloom_filter.100/1")
require.NoError(t, WriteFile(cfg, statPath, []byte("bloom-blob")))
got, err := CommitManifestUpdates(basePath, ManifestEarliest, cfg, &ManifestUpdates{
Stats: []StatEntry{{
Key: "bloom_filter.100",
Files: []string{statPath},
Metadata: map[string]string{"memory_size": "10"},
}},
})
require.NoError(t, err)
_, version, err := UnmarshalManifestPath(got)
require.NoError(t, err)
require.Equal(t, int64(1), version, "stats-only commit must bump version exactly once")
stats, err := GetManifestStats(got, cfg)
require.NoError(t, err)
require.Contains(t, stats, "bloom_filter.100")
require.Equal(t, "10", stats["bloom_filter.100"].Metadata["memory_size"])
}
// TestCommitManifestUpdates_DeltaOnly exercises the delta-only path.
func TestCommitManifestUpdates_DeltaOnly(t *testing.T) {
cfg := manifestTestStorageConfig(t)
basePath := "files/commit_delta_only/seg1"
deltaPath := path.Join(cfg.RootPath, basePath, "_delta/delta-1")
require.NoError(t, WriteFile(cfg, deltaPath, []byte("delta-payload")))
got, err := CommitManifestUpdates(basePath, ManifestEarliest, cfg, &ManifestUpdates{
DeltaLogs: []DeltaLogEntry{{Path: deltaPath, NumEntries: 7}},
})
require.NoError(t, err)
_, version, err := UnmarshalManifestPath(got)
require.NoError(t, err)
require.Equal(t, int64(1), version)
paths, err := GetDeltaLogPathsFromManifest(got, cfg)
require.NoError(t, err)
require.Len(t, paths, 1)
}
// TestCommitManifestUpdates_AllSections is the integration variant: one
// commit covering inserts (column-groups from a real FFIPackedWriter),
// stats, and delta entries — asserts a single version bump for the whole
// bundle.
func TestCommitManifestUpdates_AllSections(t *testing.T) {
cfg := manifestTestStorageConfig(t)
basePath := "files/commit_all_sections/seg1"
schema := arrow.NewSchema([]arrow.Field{
{
Name: "pk",
Type: arrow.PrimitiveTypes.Int64,
Nullable: false,
Metadata: arrow.NewMetadata([]string{ArrowFieldIdMetadataKey}, []string{"100"}),
},
}, nil)
b := array.NewRecordBuilder(memory.DefaultAllocator, schema)
defer b.Release()
pkb := b.Field(0).(*array.Int64Builder)
for i := 0; i < 4; i++ {
pkb.Append(int64(i))
}
rec := b.NewRecord()
defer rec.Release()
columnGroups := []storagecommon.ColumnGroup{
{Columns: []int{0}, GroupID: storagecommon.DefaultShortColumnGroupID},
}
w, err := NewFFIPackedWriter(basePath, schema, columnGroups, cfg, nil)
require.NoError(t, err)
require.NoError(t, w.WriteRecordBatch(rec))
out, err := w.Close()
require.NoError(t, err)
defer out.Destroy()
statPath := path.Join(cfg.RootPath, basePath, "_stats/bloom_filter.100/1")
require.NoError(t, WriteFile(cfg, statPath, []byte("bloom-blob")))
deltaPath := path.Join(cfg.RootPath, basePath, "_delta/delta-1")
require.NoError(t, WriteFile(cfg, deltaPath, []byte("delta-payload")))
got, err := CommitManifestUpdates(basePath, ManifestEarliest, cfg, &ManifestUpdates{
NewFiles: out,
DeltaLogs: []DeltaLogEntry{{Path: deltaPath, NumEntries: 4}},
Stats: []StatEntry{{
Key: "bloom_filter.100",
Files: []string{statPath},
Metadata: map[string]string{"memory_size": "10"},
}},
})
require.NoError(t, err)
_, version, err := UnmarshalManifestPath(got)
require.NoError(t, err)
require.Equal(t, int64(1), version,
"inserts + delta + stats together must bump version exactly once")
}
// TestCommitManifestUpdates_AddNewColumnGroups exercises the
// add_column_group (function-backfill) branch in applyManifestUpdates.
func TestCommitManifestUpdates_AddNewColumnGroups(t *testing.T) {
cfg := manifestTestStorageConfig(t)
basePath := "files/commit_new_cgs/seg1"
schema := arrow.NewSchema([]arrow.Field{
{
Name: "pk",
Type: arrow.PrimitiveTypes.Int64,
Nullable: false,
Metadata: arrow.NewMetadata([]string{ArrowFieldIdMetadataKey}, []string{"100"}),
},
}, nil)
b := array.NewRecordBuilder(memory.DefaultAllocator, schema)
defer b.Release()
b.Field(0).(*array.Int64Builder).Append(int64(1))
rec := b.NewRecord()
defer rec.Release()
columnGroups := []storagecommon.ColumnGroup{
{Columns: []int{0}, GroupID: storagecommon.DefaultShortColumnGroupID},
}
w, err := NewFFIPackedWriter(basePath, schema, columnGroups, cfg, nil)
require.NoError(t, err)
w.AsNewColumnGroups()
require.NoError(t, w.WriteRecordBatch(rec))
out, err := w.Close()
require.NoError(t, err)
defer out.Destroy()
cgs, ok := out.(*ColumnGroups)
require.True(t, ok, "Close should return *ColumnGroups for FFIPackedWriter")
require.True(t, cgs.addNewColumnGroups,
"AsNewColumnGroups should propagate into ColumnGroups handle")
got, err := CommitManifestUpdates(basePath, ManifestEarliest, cfg,
&ManifestUpdates{NewFiles: out})
require.NoError(t, err)
_, version, err := UnmarshalManifestPath(got)
require.NoError(t, err)
require.Equal(t, int64(1), version)
}
// TestColumnGroups_DestroyIdempotent verifies that calling Destroy on a
// ColumnGroups handle multiple times (and on a nil receiver) is safe.
func TestColumnGroups_DestroyIdempotent(t *testing.T) {
cfg := manifestTestStorageConfig(t)
basePath := "files/cg_destroy_test/seg1"
schema := arrow.NewSchema([]arrow.Field{
{
Name: "pk",
Type: arrow.PrimitiveTypes.Int64,
Nullable: false,
Metadata: arrow.NewMetadata([]string{ArrowFieldIdMetadataKey}, []string{"100"}),
},
}, nil)
b := array.NewRecordBuilder(memory.DefaultAllocator, schema)
defer b.Release()
b.Field(0).(*array.Int64Builder).Append(int64(0))
rec := b.NewRecord()
defer rec.Release()
w, err := NewFFIPackedWriter(basePath, schema,
[]storagecommon.ColumnGroup{{Columns: []int{0}, GroupID: storagecommon.DefaultShortColumnGroupID}},
cfg, nil)
require.NoError(t, err)
require.NoError(t, w.WriteRecordBatch(rec))
out, err := w.Close()
require.NoError(t, err)
cgs, ok := out.(*ColumnGroups)
require.True(t, ok, "Close should return *ColumnGroups")
require.NotNil(t, cgs.cColumnGroups, "Close should yield non-nil cColumnGroups")
cgs.Destroy()
require.Nil(t, cgs.cColumnGroups, "first Destroy must clear cColumnGroups")
cgs.Destroy() // second call must be a no-op
assert.Nil(t, cgs.cColumnGroups)
var nilHandle *ColumnGroups
assert.NotPanics(t, func() { nilHandle.Destroy() },
"Destroy on nil receiver must be safe")
}
// TestFFIPackedWriter_DoubleCloseRejected verifies that the second Close
// after a successful first Close returns an error rather than re-running
// the close FFI on an exhausted handle.
func TestFFIPackedWriter_DoubleCloseRejected(t *testing.T) {
cfg := manifestTestStorageConfig(t)
basePath := "files/ffi_double_close/seg1"
schema := arrow.NewSchema([]arrow.Field{
{
Name: "pk",
Type: arrow.PrimitiveTypes.Int64,
Nullable: false,
Metadata: arrow.NewMetadata([]string{ArrowFieldIdMetadataKey}, []string{"100"}),
},
}, nil)
b := array.NewRecordBuilder(memory.DefaultAllocator, schema)
defer b.Release()
b.Field(0).(*array.Int64Builder).Append(int64(0))
rec := b.NewRecord()
defer rec.Release()
w, err := NewFFIPackedWriter(basePath, schema,
[]storagecommon.ColumnGroup{{Columns: []int{0}, GroupID: storagecommon.DefaultShortColumnGroupID}},
cfg, nil)
require.NoError(t, err)
require.NoError(t, w.WriteRecordBatch(rec))
out, err := w.Close()
require.NoError(t, err)
defer out.Destroy()
_, err = w.Close()
require.Error(t, err, "second Close must return an error")
require.Contains(t, err.Error(), "already closed")
}