furyhawk/milvus
1
0
Fork 0
mirror of https://github.com/milvus-io/milvus.git synced 2026-07-21 02:05:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: make BulkPackWriterV3.Write atomic with single manifest commit (#49711)

Browse Source 
issue: #49710
related: #49069

BulkPackWriterV3.Write currently bumps the loon manifest version up to
four times per call (inserts, stats, delta, bm25 each open and commit
their own transaction). That breaks atomicity for concurrent readers,
leaves orphan manifest versions behind on partial failure, and churns
loon's manifest history.

This PR refactors the V3 sync write path into a do-then-commit shape:

  Phase 1 (slow, runs once outside the retry loop): writeInserts,
  writeStats, writeDelta, writeBM25Stasts write all the data files
  (parquet/LOB via the loon writer, deltalog, stat blobs) and each
  helper returns its contribution (ColumnGroups, SegmentOutput,
  StatEntry list, DeltaLogEntry list) without touching shared state.

  Phase 2 (fast, retried on transient loon errors): Write assembles
  one packed.ManifestUpdates from the helper returns and calls a
  single new primitive packed.CommitManifestUpdates, which opens a
  short loon transaction, stages every change, and commits once.

To make this clean the loon FFI writers are decoupled from manifest
concerns entirely:

  FFIPackedWriter and FFISegmentWriter no longer take a baseVersion
  at construction. Their Close primitives return data carriers
  (packed.ColumnGroups, packed.SegmentOutput) that own the C output.
  packed.CommitManifestUpdates is the only Go-level entry point that
  mutates a manifest. Empty payload short-circuits to the unchanged
  manifest path without opening a transaction.

The same do-then-commit shape applies to the V2-style binlog import
writers (PackedManifestRecordWriter, PackedTextManifestRecordWriter):
they assemble inserts plus bloom-filter and BM25 stat blobs into one
ManifestUpdates and call CommitManifestUpdates once.

Internal renames: the low-level adapters packedRecordManifestWriter
and packedTextManifestWriter are renamed to packedRecordBatchWriter
and packedTextBatchWriter because they don't touch the manifest
themselves — they translate storage.Record into arrow.Record and
track per-column-group sizes. The interface manifestRecordWriter is
renamed to packedBatchWriter to match.

New tests in pack_writer_v3_test.go:

  TestWrite_SingleVersionBumpAcrossSections drives an insert + delta
  + flush so all four sections fire and asserts the new manifest
  version equals baseVer + 1.

  TestWrite_RetryDoesNotLeakVersionBumps mockey-patches
  CommitManifestUpdates to fail once with packed.ErrLoonTransient
  and asserts the eventual successful commit still yields baseVer + 1
  rather than baseVer + N_retries.

All 11 V3 tests pass. Existing packed and storage tests pass.

Out of scope (filed as follow-ups in #49710's discussion):

  backfill_compactor.go's runBm25Function still issues a separate
  AddStatsToManifest after the column-groups commit, producing two
  version bumps per backfill. The architecture supports folding the
  two but the change isn't strictly needed for the V3 sync atomicity
  fix.

  FFIPackedWriter lacks a Destroy method; on loon_writer_close
  failure the handle and cProperties leak. Pre-existing on master.

---------

Signed-off-by: Ted Xu <ted.xu@zilliz.com>
This commit is contained in:
Ted Xu
2026-05-27 16:42:27 +08:00
committed by GitHub
parent 9a455d61ec
commit c3207d5963
16 changed files with 1357 additions and 534 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/compaction/common_test.go
+7 -2
View File
@@ -386,7 +386,7 @@ func (s *CommonSuite) createBaseManifest(basePath string, storageConfig *indexpb
{Columns: []int{0, 1}, GroupID: storagecommon.DefaultShortColumnGroupID},
}
pw, err := packed.NewFFIPackedWriter(basePath, 0, arrowSchema, columnGroups, storageConfig, nil)
pw, err := packed.NewFFIPackedWriter(basePath, arrowSchema, columnGroups, storageConfig, nil)
require.NoError(t, err)
// Write minimal data to create a valid manifest
@@ -401,7 +401,12 @@ func (s *CommonSuite) createBaseManifest(basePath string, storageConfig *indexpb
err = pw.WriteRecordBatch(rec)
require.NoError(t, err)
manifestPath, err := pw.Close()
out, err := pw.Close()
require.NoError(t, err)
defer out.Destroy()
manifestPath, err := packed.CommitManifestUpdates(basePath, packed.ManifestEarliest, storageConfig,
&packed.ManifestUpdates{NewFiles: out})
require.NoError(t, err)
return manifestPath
internal/datanode/compactor/backfill_compactor.go
+27 -8
View File
@@ -63,11 +63,16 @@ type backfillWriter interface {
Manifest() string
}
// ffiWriterWrapper wraps packed.FFIPackedWriter to implement backfillWriter,
// capturing the manifest string returned by Close.
// ffiWriterWrapper wraps packed.FFIPackedWriter to implement backfillWriter.
// FFIPackedWriter itself knows nothing about manifests; this wrapper closes
// the writer to collect column groups, then commits them to the existing
// manifest via packed.CommitManifestUpdates.
type ffiWriterWrapper struct {
writer *packed.FFIPackedWriter
manifest string
writer *packed.FFIPackedWriter
basePath string
baseVersion int64
storageConfig *indexpb.StorageConfig
manifest string
}
func (w *ffiWriterWrapper) WriteRecordBatch(r arrow.Record) error {
@@ -75,11 +80,20 @@ func (w *ffiWriterWrapper) WriteRecordBatch(r arrow.Record) error {
}
func (w *ffiWriterWrapper) Close() error {
manifest, err := w.writer.Close()
out, err := w.writer.Close()
if err != nil {
return err
}
w.manifest = manifest
if out == nil {
return nil
}
defer out.Destroy()
newPath, err := packed.CommitManifestUpdates(w.basePath, w.baseVersion, w.storageConfig,
&packed.ManifestUpdates{NewFiles: out})
if err != nil {
return err
}
w.manifest = newPath
return nil
}
@@ -389,7 +403,7 @@ func (t *backfillCompactionTask) setupWriter(outputField *schemapb.FieldSchema,
if err != nil {
return nil, merr.WrapErrServiceInternal("failed to parse existing manifest for V3 backfill", err.Error())
}
ffiWriter, err := packed.NewFFIPackedWriter(basePath, existingVersion, arrowSchema, newColumnGroups, t.compactionParams.StorageConfig, pluginContext)
ffiWriter, err := packed.NewFFIPackedWriter(basePath, arrowSchema, newColumnGroups, t.compactionParams.StorageConfig, pluginContext)
if err != nil {
return nil, err
}
@@ -397,7 +411,12 @@ func (t *backfillCompactionTask) setupWriter(outputField *schemapb.FieldSchema,
// Use AddColumnGroup semantics so Loon does not require the count to match
// the existing groups in the manifest.
ffiWriter.AsNewColumnGroups()
result.writer = &ffiWriterWrapper{writer: ffiWriter}
result.writer = &ffiWriterWrapper{
writer: ffiWriter,
basePath: basePath,
baseVersion: existingVersion,
storageConfig: t.compactionParams.StorageConfig,
}
result.basePath = basePath
} else {
// V2: use PackedWriter with explicit file paths.
internal/flushcommon/syncmgr/pack_writer_v3.go
+181 -223
View File
@@ -33,7 +33,6 @@ import (
"github.com/milvus-io/milvus/internal/storagev2/packed"
"github.com/milvus-io/milvus/pkg/v3/log"
"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/util/metautil"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
@@ -41,14 +40,17 @@ import (
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
// manifestRecordWriter is the common interface for both packedRecordManifestWriter
// and packedTextManifestWriter used for V3 storage writes.
type manifestRecordWriter interface {
storage.RecordWriter
// packedBatchWriter is the common subset both storage.packedRecordBatchWriter
// and storage.packedTextBatchWriter expose to the V3 sync path. Close
// returns a packed.WriterOutput regardless of which concrete writer is
// in use; the V3 sync path treats both uniformly.
type packedBatchWriter interface {
Write(r storage.Record) error
Close() (packed.WriterOutput, error)
GetWrittenUncompressed() uint64
GetColumnGroupWrittenCompressed(columnGroup typeutil.UniqueID) uint64
GetColumnGroupWrittenUncompressed(columnGroup typeutil.UniqueID) uint64
GetWrittenPaths(columnGroup typeutil.UniqueID) string
GetWrittenManifest() string
GetWrittenRowNum() int64
}
@@ -57,23 +59,22 @@ type BulkPackWriterV3 struct {
manifestPath string
// initialManifestPath captures the manifest path observed when Write was
// first invoked. resetForRetry restores manifestPath to this value so each
// retry attempt restarts from the same base manifest version. Touching this
// invariant breaks the retry correctness argument — see
// ccmd/pack_writer_v3_retry_plan.md.
// initialManifestPath captures the manifest path observed when Write is
// invoked. Read-based merges (existing bloom filter / BM25 file lists)
// reference this stable value instead of bw.manifestPath so they cannot
// drift after Phase 2's commit updates manifestPath.
initialManifestPath string
// pendingMetaCacheActions accumulates RollStats / MergeBm25Stats actions
// that writeStats / writeBM25Stasts would otherwise apply directly to the
// metaCache. They are deferred and drained by SyncTask.Run only after a
// successful Write, so that retried attempts do not double-apply.
// metaCache. They are drained by SyncTask.Run only after a successful
// Write, so a failed commit retry never observes partial metaCache state.
pendingMetaCacheActions []metacache.SegmentAction
// singlePKStats holds the per-batch PK statistic computed once per Write
// call. It is reused across retry attempts and passed explicitly into the
// serializer's *With variants so merged-stats computation can include this
// batch without requiring the metaCache to be updated yet.
// singlePKStats holds the per-batch PK statistic computed during
// writeStats and passed explicitly into the serializer's *With variants
// so merged-stats computation can include this batch without requiring
// the metaCache to be updated yet.
singlePKStats *storage.PrimaryKeyStats
}
@@ -89,17 +90,20 @@ func NewBulkPackWriterV3(metaCache metacache.MetaCache, schema *schemapb.Collect
}
}
// Write performs the four manifest-mutating steps for a SyncPack inside a
// retry loop. The retry handles loon transaction conflicts (currently
// surfaced as packed.ErrLoonTransient — see internal/storagev2/packed
// /ffi_common.go for the reason this is coarse-grained today).
// Write executes a SyncPack as do-then-commit:
//
// Each attempt restarts from bw.initialManifestPath via resetForRetry so the
// manifest read base does not drift across retries. metaCache mutations
// produced by writeStats / writeBM25Stasts are accumulated in
// bw.pendingMetaCacheActions and applied via a deferred drainer that runs
// only when Write returns nil — this is what makes retries idempotent on
// metaCache without exposing the pending list to callers.
// 1. Phase 1 (slow, runs once): write all data files — parquet/LOB via the
// loon writer, deltalog via the deltalog writer, stat blobs via the
// filesystem FFI — and assemble a single packed.ManifestUpdates payload
// describing every change the segment needs to register.
// 2. Phase 2 (fast, retried on transient loon errors): call
// packed.CommitManifestUpdates once. The loon transaction handle is
// opened, all changes are staged, and the transaction is committed in
// one shot, producing exactly one manifest version bump.
//
// metaCache mutations produced by Phase 1's stats helpers are accumulated
// in bw.pendingMetaCacheActions and only drained on success, so a failed
// retry cycle leaves the metaCache untouched.
func (bw *BulkPackWriterV3) Write(ctx context.Context, pack *SyncPack) (
inserts map[int64]*datapb.FieldBinlog,
deltas *datapb.FieldBinlog,
@@ -112,9 +116,9 @@ func (bw *BulkPackWriterV3) Write(ctx context.Context, pack *SyncPack) (
log := log.Ctx(ctx)
bw.initialManifestPath = bw.manifestPath
// Drain deferred metaCache actions on successful Write only. If the retry
// loop terminates with an error we leave the metaCache untouched, so a
// failed sync does not pollute bloom-filter / BM25 history.
// Drain deferred metaCache actions on successful Write only. If we
// return an error we leave the metaCache untouched, so a failed sync
// does not pollute bloom-filter / BM25 history.
defer func() {
if err != nil {
return
@@ -128,41 +132,66 @@ func (bw *BulkPackWriterV3) Write(ctx context.Context, pack *SyncPack) (
)
}()
var deltaSummary *datapb.FieldBinlog
basePath, baseVersion, parseErr := packed.UnmarshalManifestPath(bw.initialManifestPath)
if parseErr != nil {
err = parseErr
return
}
// Phase 1: write files. Each helper returns its contribution to the
// final ManifestUpdates instead of mutating shared state.
var (
insertFiles packed.WriterOutput
statEntries []packed.StatEntry
deltaEntries []packed.DeltaLogEntry
bm25Entries []packed.StatEntry
)
defer func() {
if insertFiles != nil {
insertFiles.Destroy()
}
}()
if inserts, insertFiles, err = bw.writeInserts(ctx, pack, basePath); err != nil {
log.Warn("failed to write insert data", zap.Error(err))
return
}
if stats, statEntries, err = bw.writeStats(ctx, pack, basePath); err != nil {
log.Warn("failed to process stats blob", zap.Error(err))
return
}
if deltas, deltaEntries, err = bw.writeDelta(ctx, pack, basePath); err != nil {
log.Warn("failed to process delta blob", zap.Error(err))
return
}
if bm25Stats, bm25Entries, err = bw.writeBM25Stasts(ctx, pack, basePath); err != nil {
log.Warn("failed to process bm25 stats blob", zap.Error(err))
return
}
updates := &packed.ManifestUpdates{
NewFiles: insertFiles,
DeltaLogs: deltaEntries,
Stats: append(statEntries, bm25Entries...),
}
// Phase 2: commit the assembled updates. CommitManifestUpdates short-
// circuits to the unchanged manifest path when updates carry nothing.
// The outer retry.Do handles transient FFI errors classified as
// packed.ErrLoonTransient; loon's own optimistic retry covers
// manifest-version conflicts within a single attempt.
err = retry.Do(ctx, func() error {
bw.resetForRetry()
var innerErr error
if inserts, manifest, innerErr = bw.writeInserts(ctx, pack); innerErr != nil {
log.Warn("failed to write insert data", zap.Error(innerErr))
return classifyLoonErr(innerErr)
}
// Update manifestPath after writeInserts
bw.manifestPath = manifest
// writeStats for V3 adds bloom filter stats to manifest
if stats, innerErr = bw.writeStats(ctx, pack); innerErr != nil {
log.Warn("failed to process stats blob", zap.Error(innerErr))
return classifyLoonErr(innerErr)
}
// writeDelta for V3 updates manifest and returns delta summary
if manifest, deltaSummary, innerErr = bw.writeDelta(ctx, pack); innerErr != nil {
log.Warn("failed to process delta blob", zap.Error(innerErr))
return classifyLoonErr(innerErr)
}
// writeBM25Stasts for V3 adds BM25 stats to manifest
if bm25Stats, innerErr = bw.writeBM25Stasts(ctx, pack); innerErr != nil {
log.Warn("failed to process bm25 stats blob", zap.Error(innerErr))
return classifyLoonErr(innerErr)
newPath, commitErr := packed.CommitManifestUpdates(basePath, baseVersion, bw.storageConfig, updates)
if commitErr != nil {
return classifyLoonErr(commitErr)
}
bw.manifestPath = newPath
return nil
}, bw.writeRetryOpts...)
if err != nil {
return
}
// For V3, stats are in manifest; delta summary is returned for compaction trigger
deltas = deltaSummary
manifest = bw.manifestPath
size = bw.sizeWritten
return
@@ -186,70 +215,69 @@ func classifyLoonErr(err error) error {
return retry.Unrecoverable(err)
}
// resetForRetry restores per-attempt state. It MUST restore manifestPath to
// initialManifestPath — see the field doc on initialManifestPath for why.
func (bw *BulkPackWriterV3) resetForRetry() {
bw.manifestPath = bw.initialManifestPath
bw.sizeWritten = 0
bw.pendingMetaCacheActions = bw.pendingMetaCacheActions[:0]
bw.singlePKStats = nil
// Files written under _stats/ and _delta/ by previous failed attempts are
// intentionally NOT cleaned up here: they are unreferenced by any
// committed manifest and will be reclaimed by loon GC. See
// ccmd/pack_writer_v3_retry_plan.md decision (2).
}
func (bw *BulkPackWriterV3) writeInserts(ctx context.Context, pack *SyncPack) (map[int64]*datapb.FieldBinlog, string, error) {
// writeInserts writes the insert data files for the SyncPack and returns
// the produced WriterOutput plus the per-column-group binlog metadata.
// The returned WriterOutput is owned by the caller and must be Destroy'd
// after the surrounding commit (success or failure).
func (bw *BulkPackWriterV3) writeInserts(ctx context.Context, pack *SyncPack, basePath string) (logs map[int64]*datapb.FieldBinlog, files packed.WriterOutput, err error) {
if len(pack.insertData) == 0 {
return make(map[int64]*datapb.FieldBinlog), bw.manifestPath, nil
return make(map[int64]*datapb.FieldBinlog), nil, nil
}
rec, err := bw.serializeBinlog(ctx, pack)
if err != nil {
return nil, "", err
return nil, nil, err
}
defer rec.Release()
tsFrom, tsTo := bw.getTsRange(rec)
pluginContextPtr := bw.getPluginContext(pack.collectionID)
// NOTE: this used to wrap the call below in its own retry.Do; the outer
// BulkPackWriterV3.Write retry loop now covers it, so we drop the inner
// retry to avoid Attempts^2 amplification.
logs, manifestPath, err := bw.writeInsertsIntoStorage(ctx, pluginContextPtr, rec, tsFrom, tsTo)
// LOB base path is at partition level: {basePath}/.. = {root}/insert_log/{coll}/{part}
partitionBasePath := path.Dir(basePath)
textColumnConfigs := buildTextColumnConfigs(bw.schema, partitionBasePath)
var w packedBatchWriter
if len(textColumnConfigs) > 0 {
log.Ctx(ctx).Info("using TEXT-aware writer for import",
zap.Int("textFieldCount", len(textColumnConfigs)),
zap.String("basePath", basePath))
w, err = storage.NewPackedTextBatchWriter("", basePath, bw.schema,
bw.bufferSize, bw.multiPartUploadSize, bw.columnGroups, bw.storageConfig, textColumnConfigs)
} else {
w, err = storage.NewPackedRecordBatchWriter(basePath, bw.schema,
bw.bufferSize, bw.multiPartUploadSize, bw.columnGroups, bw.storageConfig, pluginContextPtr)
}
if err != nil {
log.Ctx(ctx).Warn("failed to write inserts into storage",
return nil, nil, err
}
// Ensure the FFI writer's C resources are reclaimed even if Write fails
// before we reach the explicit Close below. Once Close has been called
// (success or failure) the writer's own defer has already released its
// handle and properties, so closeAttempted gates against a redundant
// second Close from this defer.
closeAttempted := false
defer func() {
if !closeAttempted {
if out, closeErr := w.Close(); closeErr == nil && out != nil {
out.Destroy()
}
}
}()
if err = w.Write(rec); err != nil {
log.Ctx(ctx).Warn("failed to write inserts",
zap.Int64("collectionID", pack.collectionID),
zap.Int64("segmentID", pack.segmentID),
zap.Error(err))
return nil, "", err
return nil, nil, err
}
return logs, manifestPath, nil
}
func (bw *BulkPackWriterV3) writeInsertsIntoStorage(ctx context.Context,
pluginContextPtr *indexcgopb.StoragePluginContext,
rec storage.Record,
tsFrom typeutil.Timestamp,
tsTo typeutil.Timestamp,
) (map[int64]*datapb.FieldBinlog, string, error) {
log := log.Ctx(ctx)
logs := make(map[int64]*datapb.FieldBinlog)
columnGroups := bw.columnGroups
var err error
doWrite := func(w manifestRecordWriter) error {
if err = w.Write(rec); err != nil {
if closeErr := w.Close(); closeErr != nil {
log.Error("failed to close writer after write failed", zap.Error(closeErr))
}
return err
}
// close first the get stats & output
return w.Close()
closeAttempted = true
if files, err = w.Close(); err != nil {
return nil, nil, err
}
var manifestPath string
getFieldNullCounts := func(columnGroup storagecommon.ColumnGroup) map[int64]int64 {
result := make(map[int64]int64, len(columnGroup.Fields))
for _, fieldID := range columnGroup.Fields {
@@ -260,31 +288,8 @@ func (bw *BulkPackWriterV3) writeInsertsIntoStorage(ctx context.Context,
return result
}
basePath, version, err := packed.UnmarshalManifestPath(bw.manifestPath)
if err != nil {
return nil, "", err
}
// LOB base path is at partition level: {basePath}/.. = {root}/insert_log/{coll}/{part}
partitionBasePath := path.Dir(basePath)
textColumnConfigs := buildTextColumnConfigs(bw.schema, partitionBasePath)
var w manifestRecordWriter
if len(textColumnConfigs) > 0 {
log.Info("using TEXT-aware writer for import",
zap.Int("textFieldCount", len(textColumnConfigs)),
zap.String("basePath", basePath))
w, err = storage.NewPackedTextManifestWriter("", basePath, version, bw.schema,
bw.bufferSize, bw.multiPartUploadSize, columnGroups, bw.storageConfig, textColumnConfigs)
} else {
w, err = storage.NewPackedRecordManifestWriter(basePath, version, bw.schema, bw.bufferSize, bw.multiPartUploadSize, columnGroups, bw.storageConfig, pluginContextPtr)
}
if err != nil {
return nil, "", err
}
if err = doWrite(w); err != nil {
return nil, "", err
}
for _, columnGroup := range columnGroups {
logs = make(map[int64]*datapb.FieldBinlog)
for _, columnGroup := range bw.columnGroups {
columnGroupID := columnGroup.GroupID
logs[columnGroupID] = &datapb.FieldBinlog{
FieldID: columnGroupID,
@@ -302,75 +307,53 @@ func (bw *BulkPackWriterV3) writeInsertsIntoStorage(ctx context.Context,
},
}
}
manifestPath = w.GetWrittenManifest()
return logs, manifestPath, nil
return logs, files, nil
}
// writeDelta writes deltalog to storage and updates the manifest.
// Returns the updated manifest path and a pathless delta summary FieldBinlog
// containing only EntriesNum and MemorySize for compaction trigger decisions.
func (bw *BulkPackWriterV3) writeDelta(ctx context.Context, pack *SyncPack) (string, *datapb.FieldBinlog, error) {
// writeDelta writes the deltalog file and returns the DeltaLogEntry list
// the caller will fold into ManifestUpdates plus a pathless delta summary
// FieldBinlog (EntriesNum + MemorySize) used by compaction-trigger
// decisions.
func (bw *BulkPackWriterV3) writeDelta(ctx context.Context, pack *SyncPack, basePath string) (*datapb.FieldBinlog, []packed.DeltaLogEntry, error) {
if pack.deltaData == nil || pack.deltaData.RowCount == 0 {
return bw.manifestPath, nil, nil
return nil, nil, nil
}
pkField, err := typeutil.GetPrimaryFieldSchema(bw.schema)
if err != nil {
return "", nil, fmt.Errorf("primary key field not found: %w", err)
return nil, nil, fmt.Errorf("primary key field not found: %w", err)
}
// Allocate log ID for deltalog
logID, err := bw.allocator.AllocOne()
if err != nil {
return "", nil, err
}
// Build deltalog path under basePath/_delta/
basePath, _, err := packed.UnmarshalManifestPath(bw.manifestPath)
if err != nil {
return "", nil, fmt.Errorf("failed to parse manifest path: %w", err)
return nil, nil, err
}
deltaPath := metautil.BuildDeltaLogPathV3(basePath, logID)
// Create deltalog writer with V2 storage
writer, err := storage.NewDeltalogWriter(
ctx, pack.collectionID, pack.partitionID, pack.segmentID, logID, pkField.DataType, deltaPath,
storage.WithVersion(storage.StorageV2),
storage.WithStorageConfig(bw.storageConfig),
)
if err != nil {
return "", nil, fmt.Errorf("failed to create deltalog writer: %w", err)
return nil, nil, fmt.Errorf("failed to create deltalog writer: %w", err)
}
// Build Arrow record from delete data using existing utility
record, _, _, err := storage.BuildDeleteRecord(pack.deltaData.Pks, pack.deltaData.Tss)
if err != nil {
return "", nil, fmt.Errorf("failed to build delete record: %w", err)
return nil, nil, fmt.Errorf("failed to build delete record: %w", err)
}
defer record.Release()
// Write and close
if err := writer.Write(record); err != nil {
return "", nil, fmt.Errorf("failed to write delta record: %w", err)
return nil, nil, fmt.Errorf("failed to write delta record: %w", err)
}
if err := writer.Close(); err != nil {
return "", nil, fmt.Errorf("failed to close delta writer: %w", err)
return nil, nil, fmt.Errorf("failed to close delta writer: %w", err)
}
// Update manifest with the new deltalog
newManifest, err := packed.AddDeltaLogsToManifest(
bw.manifestPath,
bw.storageConfig,
[]packed.DeltaLogEntry{{Path: deltaPath, NumEntries: pack.deltaData.RowCount}},
)
if err != nil {
return "", nil, fmt.Errorf("failed to add deltalog to manifest: %w", err)
}
bw.manifestPath = newManifest
bw.sizeWritten += pack.deltaData.Size()
// Return delta summary for compaction trigger decisions (no path, only stats)
summary := &datapb.FieldBinlog{
Binlogs: []*datapb.Binlog{{
LogID: logID,
@@ -378,25 +361,24 @@ func (bw *BulkPackWriterV3) writeDelta(ctx context.Context, pack *SyncPack) (str
MemorySize: int64(writer.GetWrittenUncompressed()),
}},
}
return newManifest, summary, nil
return summary, []packed.DeltaLogEntry{{Path: deltaPath, NumEntries: pack.deltaData.RowCount}}, nil
}
// writeStats overrides the base class to write bloom filter stats into the
// manifest instead of separate binlog files. The stat files are written
// under _stats/ relative to the manifest base path, then registered in
// the manifest via a transaction.
func (bw *BulkPackWriterV3) writeStats(ctx context.Context, pack *SyncPack) (map[int64]*datapb.FieldBinlog, error) {
// writeStats writes bloom filter stat blobs under basePath/_stats and
// returns the resulting StatEntry list. The caller folds the entries into
// a ManifestUpdates that commits atomically with inserts / delta / bm25.
func (bw *BulkPackWriterV3) writeStats(ctx context.Context, pack *SyncPack, basePath string) (map[int64]*datapb.FieldBinlog, []packed.StatEntry, error) {
if len(pack.insertData) == 0 {
return make(map[int64]*datapb.FieldBinlog), nil
return make(map[int64]*datapb.FieldBinlog), nil, nil
}
serializer, err := NewStorageSerializer(bw.metaCache, bw.schema)
if err != nil {
return nil, err
return nil, nil, err
}
singlePKStats, batchStatsBlob, err := serializer.serializeStatslog(pack)
if err != nil {
return nil, err
return nil, nil, err
}
bw.singlePKStats = singlePKStats
@@ -409,10 +391,6 @@ func (bw *BulkPackWriterV3) writeStats(ctx context.Context, pack *SyncPack) (map
metacache.RollStats(singlePKStats))
pkFieldID := serializer.pkField.GetFieldID()
basePath, _, err := packed.UnmarshalManifestPath(bw.manifestPath)
if err != nil {
return nil, err
}
var files []string
var memorySize int64
@@ -421,7 +399,7 @@ func (bw *BulkPackWriterV3) writeStats(ctx context.Context, pack *SyncPack) (map
// loon_transaction_update_stat uses replace semantics, so we must
// merge previously written files into the new entry.
statKey := fmt.Sprintf("bloom_filter.%d", pkFieldID)
existingStats, err := packed.GetManifestStats(bw.manifestPath, bw.storageConfig)
existingStats, err := packed.GetManifestStats(bw.initialManifestPath, bw.storageConfig)
if err == nil {
if existing, ok := existingStats[statKey]; ok && len(existing.Paths) > 0 {
files = append(files, existing.Paths...)
@@ -435,12 +413,12 @@ func (bw *BulkPackWriterV3) writeStats(ctx context.Context, pack *SyncPack) (map
// Write batch stats blob via filesystem FFI.
id, err := bw.allocator.AllocOne()
if err != nil {
return nil, err
return nil, nil, err
}
relPath := fmt.Sprintf("_stats/bloom_filter.%d/%d", pkFieldID, id)
fullPath := path.Join(basePath, relPath)
if err := packed.WriteFile(bw.storageConfig, fullPath, batchStatsBlob.Value); err != nil {
return nil, err
return nil, nil, err
}
bw.sizeWritten += int64(len(batchStatsBlob.Value))
memorySize += int64(len(batchStatsBlob.Value))
@@ -452,54 +430,43 @@ func (bw *BulkPackWriterV3) writeStats(ctx context.Context, pack *SyncPack) (map
// since the corresponding RollStats has been deferred above.
mergedStatsBlob, err := serializer.serializeMergedPkStatsWith(pack, singlePKStats)
if err != nil {
return nil, err
return nil, nil, err
}
mergedRelPath := fmt.Sprintf("_stats/bloom_filter.%d/%d", pkFieldID, int64(storage.CompoundStatsType))
mergedFullPath := path.Join(basePath, mergedRelPath)
if err := packed.WriteFile(bw.storageConfig, mergedFullPath, mergedStatsBlob.Value); err != nil {
return nil, err
return nil, nil, err
}
bw.sizeWritten += int64(len(mergedStatsBlob.Value))
memorySize += int64(len(mergedStatsBlob.Value))
files = append(files, mergedFullPath)
}
// Register stats in manifest
newManifest, err := packed.AddStatsToManifest(bw.manifestPath, bw.storageConfig, []packed.StatEntry{
{
Key: statKey,
Files: files,
Metadata: map[string]string{"memory_size": fmt.Sprintf("%d", memorySize)},
},
})
if err != nil {
return nil, fmt.Errorf("failed to add stats to manifest: %w", err)
}
bw.manifestPath = newManifest
entries := []packed.StatEntry{{
Key: statKey,
Files: files,
Metadata: map[string]string{"memory_size": fmt.Sprintf("%d", memorySize)},
}}
// Return empty map - stats are in manifest
return make(map[int64]*datapb.FieldBinlog), nil
// Return empty map - stats are embedded in the manifest, not separate binlogs
return make(map[int64]*datapb.FieldBinlog), entries, nil
}
// writeBM25Stasts overrides the base class to write BM25 stats into the
// manifest instead of separate binlog files.
func (bw *BulkPackWriterV3) writeBM25Stasts(ctx context.Context, pack *SyncPack) (map[int64]*datapb.FieldBinlog, error) {
// writeBM25Stasts writes BM25 stat blobs under basePath/_stats and returns
// the resulting StatEntry list. The caller folds the entries into a
// ManifestUpdates that commits atomically with inserts / delta / stats.
func (bw *BulkPackWriterV3) writeBM25Stasts(ctx context.Context, pack *SyncPack, basePath string) (map[int64]*datapb.FieldBinlog, []packed.StatEntry, error) {
if len(pack.bm25Stats) == 0 {
return make(map[int64]*datapb.FieldBinlog), nil
return make(map[int64]*datapb.FieldBinlog), nil, nil
}
serializer, err := NewStorageSerializer(bw.metaCache, bw.schema)
if err != nil {
return nil, err
return nil, nil, err
}
bm25Blobs, err := serializer.serializeBM25Stats(pack)
if err != nil {
return nil, err
}
basePath, _, err := packed.UnmarshalManifestPath(bw.manifestPath)
if err != nil {
return nil, err
return nil, nil, err
}
// Track per-field files and memory sizes, then build stat entries at the end
@@ -510,7 +477,7 @@ func (bw *BulkPackWriterV3) writeBM25Stasts(ctx context.Context, pack *SyncPack)
fieldMap := make(map[int64]*fieldStats)
// Preserve existing BM25 stat files from previous batches.
existingStats, err := packed.GetManifestStats(bw.manifestPath, bw.storageConfig)
existingStats, err := packed.GetManifestStats(bw.initialManifestPath, bw.storageConfig)
if err == nil {
for key, existing := range existingStats {
prefix, fieldID, ok := packed.ParseStatKey(key)
@@ -528,13 +495,13 @@ func (bw *BulkPackWriterV3) writeBM25Stasts(ctx context.Context, pack *SyncPack)
for fieldID, blob := range bm25Blobs {
id, err := bw.allocator.AllocOne()
if err != nil {
return nil, err
return nil, nil, err
}
relPath := fmt.Sprintf("_stats/bm25.%d/%d", fieldID, id)
fullPath := path.Join(basePath, relPath)
if err := packed.WriteFile(bw.storageConfig, fullPath, blob.Value); err != nil {
return nil, err
return nil, nil, err
}
bw.sizeWritten += int64(len(blob.Value))
@@ -560,13 +527,13 @@ func (bw *BulkPackWriterV3) writeBM25Stasts(ctx context.Context, pack *SyncPack)
// since the corresponding MergeBm25Stats has been deferred above.
mergedBM25Blob, err := serializer.serializeMergedBM25StatsWith(pack, pack.bm25Stats)
if err != nil {
return nil, err
return nil, nil, err
}
for fieldID, blob := range mergedBM25Blob {
mergedRelPath := fmt.Sprintf("_stats/bm25.%d/%d", fieldID, int64(storage.CompoundStatsType))
mergedFullPath := path.Join(basePath, mergedRelPath)
if err := packed.WriteFile(bw.storageConfig, mergedFullPath, blob.Value); err != nil {
return nil, err
return nil, nil, err
}
bw.sizeWritten += int64(len(blob.Value))
@@ -580,26 +547,17 @@ func (bw *BulkPackWriterV3) writeBM25Stasts(ctx context.Context, pack *SyncPack)
}
}
// Build stat entries with memory_size metadata
var statEntries []packed.StatEntry
var entries []packed.StatEntry
for fieldID, fs := range fieldMap {
statEntries = append(statEntries, packed.StatEntry{
entries = append(entries, packed.StatEntry{
Key: fmt.Sprintf("bm25.%d", fieldID),
Files: fs.files,
Metadata: map[string]string{"memory_size": fmt.Sprintf("%d", fs.memorySize)},
})
}
if len(statEntries) > 0 {
newManifest, err := packed.AddStatsToManifest(bw.manifestPath, bw.storageConfig, statEntries)
if err != nil {
return nil, fmt.Errorf("failed to add BM25 stats to manifest: %w", err)
}
bw.manifestPath = newManifest
}
// Return empty map - stats are in manifest
return make(map[int64]*datapb.FieldBinlog), nil
// Return empty map - stats are embedded in the manifest, not separate binlogs
return make(map[int64]*datapb.FieldBinlog), entries, nil
}
// buildTextColumnConfigs builds TextColumnConfig for all TEXT fields in the schema.
internal/flushcommon/syncmgr/pack_writer_v3_helpers_test.go
+87
View File
@@ -0,0 +1,87 @@
// 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 syncmgr
import (
"fmt"
"testing"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/storagev2/packed"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/retry"
)
func TestClassifyLoonErr(t *testing.T) {
t.Run("nil", func(t *testing.T) {
assert.NoError(t, classifyLoonErr(nil))
})
t.Run("transient stays retryable", func(t *testing.T) {
wrapped := fmt.Errorf("loon failed: %w", packed.ErrLoonTransient)
out := classifyLoonErr(wrapped)
require.Error(t, out)
assert.True(t, errors.Is(out, packed.ErrLoonTransient),
"transient error must remain retryable")
assert.True(t, retry.IsRecoverable(out),
"transient error must remain recoverable so retry.Do retries")
})
t.Run("non-transient becomes unrecoverable", func(t *testing.T) {
out := classifyLoonErr(errors.New("permanent disk fault"))
require.Error(t, out)
assert.False(t, retry.IsRecoverable(out),
"non-transient error must be wrapped Unrecoverable so retry stops")
})
}
func TestBuildTextColumnConfigs(t *testing.T) {
// Required for paramtable.Get().DataNodeCfg.* to return defaults.
paramtable.Init()
t.Run("no text fields returns nil", func(t *testing.T) {
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "varchar", DataType: schemapb.DataType_VarChar},
{FieldID: 102, Name: "vector", DataType: schemapb.DataType_FloatVector},
}}
got := buildTextColumnConfigs(schema, "/seg/base")
assert.Empty(t, got, "schema without TEXT fields should produce no configs")
})
t.Run("one text field produces one config with derived lob path", func(t *testing.T) {
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "doc", DataType: schemapb.DataType_Text},
}}
got := buildTextColumnConfigs(schema, "files/insert_log/1/2")
require.Len(t, got, 1)
assert.Equal(t, int64(101), got[0].FieldID)
assert.Equal(t, "files/insert_log/1/2/lobs/101", got[0].LobBasePath)
})
t.Run("multiple text fields produce one config each", func(t *testing.T) {
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "title", DataType: schemapb.DataType_Text},
{FieldID: 102, Name: "body", DataType: schemapb.DataType_Text},
}}
got := buildTextColumnConfigs(schema, "p")
require.Len(t, got, 2)
ids := []int64{got[0].FieldID, got[1].FieldID}
assert.ElementsMatch(t, []int64{101, 102}, ids)
for _, cfg := range got {
assert.Contains(t, cfg.LobBasePath, "lobs/")
}
})
}
internal/flushcommon/syncmgr/pack_writer_v3_test.go
+110
View File
@@ -24,6 +24,7 @@ import (
"sync/atomic"
"testing"
"github.com/bytedance/mockey"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
@@ -552,3 +553,112 @@ func (s *PackWriterV3Suite) TestMultiBatchBM25StatsAccumulation() {
s.Greater(count2, count1, "batch 2 should accumulate more bm25 files than batch 1")
s.NotEmpty(stats2[bm25Key].Metadata["memory_size"], "memory_size metadata should be set")
}
// TestWrite_SingleVersionBumpAcrossSections verifies that one Write call
// bumps the manifest version exactly once even when inserts, stats, delta,
// and bm25 are all present — the atomicity guarantee this refactor adds.
func (s *PackWriterV3Suite) TestWrite_SingleVersionBumpAcrossSections() {
collectionID := int64(123)
partitionID := int64(456)
segmentID := int64(789)
channelName := fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", collectionID)
rows := 10
bfs := pkoracle.NewBloomFilterSet()
k := metautil.JoinIDPath(collectionID, partitionID, segmentID)
basePath := path.Join(common.SegmentInsertLogPath, k)
manifestPath := packed.MarshalManifestPath(basePath, packed.ManifestEarliest)
seg := metacache.NewSegmentInfo(&datapb.SegmentInfo{ManifestPath: manifestPath}, bfs, nil)
metacache.UpdateNumOfRows(1000)(seg)
mc := metacache.NewMockMetaCache(s.T())
mc.EXPECT().Collection().Return(collectionID).Maybe()
mc.EXPECT().GetSchema(mock.Anything).Return(s.schema).Maybe()
mc.EXPECT().GetSegmentByID(segmentID).Return(seg, true).Maybe()
mc.EXPECT().GetSegmentsBy(mock.Anything, mock.Anything).Return([]*metacache.SegmentInfo{seg}).Maybe()
mc.EXPECT().UpdateSegments(mock.Anything, mock.Anything).
Run(func(action metacache.SegmentAction, filters ...metacache.SegmentFilter) { action(seg) }).
Return().Maybe()
deletes := &storage.DeleteData{}
for i := 0; i < rows; i++ {
deletes.Append(storage.NewInt64PrimaryKey(int64(i+1)), uint64(100+i))
}
pack := new(SyncPack).
WithCollectionID(collectionID).
WithPartitionID(partitionID).
WithSegmentID(segmentID).
WithChannelName(channelName).
WithInsertData(genInsertData(rows, s.schema)).
WithDeleteData(deletes).
WithFlush()
_, baseVer, err := packed.UnmarshalManifestPath(manifestPath)
s.Require().NoError(err)
bw := NewBulkPackWriterV3(mc, s.schema, s.cm, s.logIDAlloc,
packed.DefaultWriteBufferSize, 0, s.storageConfig, s.currentSplit, manifestPath)
_, _, _, _, writtenManifestPath, _, err := bw.Write(context.Background(), pack)
s.Require().NoError(err)
_, newVer, err := packed.UnmarshalManifestPath(writtenManifestPath)
s.Require().NoError(err)
s.Equalf(baseVer+1, newVer,
"Write must produce exactly one version bump; baseVer=%d newVer=%d", baseVer, newVer)
}
// TestWrite_RetryDoesNotLeakVersionBumps verifies that when a transient
// commit failure forces the retry loop to re-run, the eventual successful
// commit produces only one version bump, not one per attempt.
func (s *PackWriterV3Suite) TestWrite_RetryDoesNotLeakVersionBumps() {
collectionID := int64(123)
partitionID := int64(456)
segmentID := int64(789)
channelName := fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", collectionID)
rows := 4
bfs := pkoracle.NewBloomFilterSet()
k := metautil.JoinIDPath(collectionID, partitionID, segmentID)
basePath := path.Join(common.SegmentInsertLogPath, k)
manifestPath := packed.MarshalManifestPath(basePath, packed.ManifestEarliest)
seg := metacache.NewSegmentInfo(&datapb.SegmentInfo{ManifestPath: manifestPath}, bfs, nil)
metacache.UpdateNumOfRows(1000)(seg)
mc := metacache.NewMockMetaCache(s.T())
mc.EXPECT().Collection().Return(collectionID).Maybe()
mc.EXPECT().GetSchema(mock.Anything).Return(s.schema).Maybe()
mc.EXPECT().GetSegmentByID(segmentID).Return(seg, true).Maybe()
mc.EXPECT().GetSegmentsBy(mock.Anything, mock.Anything).Return([]*metacache.SegmentInfo{seg}).Maybe()
mc.EXPECT().UpdateSegments(mock.Anything, mock.Anything).
Run(func(action metacache.SegmentAction, filters ...metacache.SegmentFilter) { action(seg) }).
Return().Maybe()
pack := new(SyncPack).
WithCollectionID(collectionID).WithPartitionID(partitionID).WithSegmentID(segmentID).
WithChannelName(channelName).WithInsertData(genInsertData(rows, s.schema))
_, baseVer, err := packed.UnmarshalManifestPath(manifestPath)
s.Require().NoError(err)
var calls int32
var origin func(string, int64, *indexpb.StorageConfig, *packed.ManifestUpdates) (string, error)
patched := mockey.Mock(packed.CommitManifestUpdates).
To(func(basePath string, baseVersion int64, cfg *indexpb.StorageConfig, updates *packed.ManifestUpdates) (string, error) {
if atomic.AddInt32(&calls, 1) == 1 {
return "", packed.ErrLoonTransient
}
return origin(basePath, baseVersion, cfg, updates)
}).
Origin(&origin).
Build()
defer patched.UnPatch()
bw := NewBulkPackWriterV3(mc, s.schema, s.cm, s.logIDAlloc,
packed.DefaultWriteBufferSize, 0, s.storageConfig, s.currentSplit, manifestPath)
_, _, _, _, newManifestPath, _, err := bw.Write(context.Background(), pack)
s.Require().NoError(err)
_, newVer, err := packed.UnmarshalManifestPath(newManifestPath)
s.Require().NoError(err)
s.Equal(baseVer+1, newVer, "retry must not produce extra version bumps")
s.Equal(int32(2), atomic.LoadInt32(&calls), "exactly one retry expected")
}
internal/storage/binlog_record_writer.go
+80 -75
View File
@@ -66,6 +66,10 @@ type packedBinlogRecordWriterBase struct {
columnGroups []storagecommon.ColumnGroup
storageConfig *indexpb.StorageConfig
storagePluginContext *indexcgopb.StoragePluginContext
// basePath is the segment data root, populated by initWriters. The
// underlying packed batch writers do not return a manifest path; the
// caller builds the manifest update against this base path.
basePath string
pkCollector *PkStatsCollector
bm25Collector *Bm25StatsCollector
@@ -376,7 +380,7 @@ var _ BinlogRecordWriter = (*PackedManifestRecordWriter)(nil)
type PackedManifestRecordWriter struct {
packedBinlogRecordWriterBase
// writer and stats generated at runtime
writer *packedRecordManifestWriter
writer *packedRecordBatchWriter
}
func (pw *PackedManifestRecordWriter) Write(r Record) error {
@@ -424,8 +428,8 @@ func (pw *PackedManifestRecordWriter) initWriters(r Record) error {
var err error
k := metautil.JoinIDPath(pw.collectionID, pw.partitionID, pw.segmentID)
basePath := path.Join(pw.storageConfig.GetRootPath(), common.SegmentInsertLogPath, k)
pw.writer, err = NewPackedRecordManifestWriter(basePath, packed.ManifestEarliest, pw.schema, pw.bufferSize, pw.multiPartUploadSize, pw.columnGroups, pw.storageConfig, pw.storagePluginContext)
pw.basePath = path.Join(pw.storageConfig.GetRootPath(), common.SegmentInsertLogPath, k)
pw.writer, err = NewPackedRecordBatchWriter(pw.basePath, pw.schema, pw.bufferSize, pw.multiPartUploadSize, pw.columnGroups, pw.storageConfig, pw.storagePluginContext)
if err != nil {
return merr.WrapErrServiceInternal(fmt.Sprintf("can not new packed record writer %s", err.Error()))
}
@@ -459,38 +463,45 @@ func (pw *PackedManifestRecordWriter) finalizeBinlogs() {
FieldNullCounts: pw.getFieldNullCountsForColumnGroup(columnGroup),
})
}
pw.manifest = pw.writer.GetWrittenManifest()
}
// Close finalizes the V3 segment. It closes the underlying packed writer to
// get column groups, serializes bloom filter / BM25 stat blobs, writes
// those blobs to storage, then performs a single packed.CommitManifestUpdates
// that registers inserts + all stats atomically.
func (pw *PackedManifestRecordWriter) Close() error {
if pw.writer != nil {
if err := pw.writer.Close(); err != nil {
return err
}
if pw.writer == nil {
return nil
}
out, err := pw.writer.Close()
if err != nil {
return err
}
if out != nil {
defer out.Destroy()
}
pw.finalizeBinlogs()
// For V3 (manifest-based) segments, stats are stored under basePath/_stats/
// and registered in the manifest. If pw.manifest is empty, no data was
// written, so there are no stats to write.
if pw.manifest != "" {
if err := pw.writeStatsV3(); err != nil {
return err
}
if out == nil {
return nil
}
updates := &packed.ManifestUpdates{NewFiles: out}
if err := pw.appendV3Stats(updates); err != nil {
return err
}
newManifest, err := packed.CommitManifestUpdates(pw.basePath, packed.ManifestEarliest, pw.storageConfig, updates)
if err != nil {
return fmt.Errorf("PackedManifestRecordWriter.Close commit: %w", err)
}
pw.manifest = newManifest
return nil
}
// writeStatsV3 writes bloom filter and BM25 stats for V3 (manifest-based)
// segments. Stats blobs are written to basePath/_stats/ and registered in the
// manifest via a transaction, leaving pw.statsLog and pw.bm25StatsLog nil so
// callers know stats are embedded in the manifest.
func (pw *PackedManifestRecordWriter) writeStatsV3() error {
basePath, _, err := packed.UnmarshalManifestPath(pw.manifest)
if err != nil {
return fmt.Errorf("writeStatsV3: failed to parse manifest path: %w", err)
}
// --- Bloom filter (PK) stats ---
// appendV3Stats serializes bloom filter / BM25 stat blobs, writes them to
// storage, and appends StatEntry records onto updates so the surrounding
// commit registers inserts + stats atomically. Leaves pw.statsLog and
// pw.bm25StatsLog nil so callers know stats are embedded in the manifest.
func (pw *PackedManifestRecordWriter) appendV3Stats(updates *packed.ManifestUpdates) error {
statsBlob, pkFieldID, err := pw.pkCollector.SerializeBlob(pw.rowNum)
if err != nil {
return err
@@ -500,55 +511,36 @@ func (pw *PackedManifestRecordWriter) writeStatsV3() error {
if err != nil {
return err
}
fullPath := path.Join(basePath, fmt.Sprintf("_stats/bloom_filter.%d/%d", pkFieldID, id))
fullPath := path.Join(pw.basePath, fmt.Sprintf("_stats/bloom_filter.%d/%d", pkFieldID, id))
if err := packed.WriteFile(pw.storageConfig, fullPath, statsBlob.Value); err != nil {
return fmt.Errorf("writeStatsV3: failed to write bloom filter stats: %w", err)
return fmt.Errorf("appendV3Stats: failed to write bloom filter stats: %w", err)
}
newManifest, err := packed.AddStatsToManifest(pw.manifest, pw.storageConfig, []packed.StatEntry{
{
Key: fmt.Sprintf("bloom_filter.%d", pkFieldID),
Files: []string{fullPath},
Metadata: map[string]string{"memory_size": fmt.Sprintf("%d", int64(len(statsBlob.Value)))},
},
updates.Stats = append(updates.Stats, packed.StatEntry{
Key: fmt.Sprintf("bloom_filter.%d", pkFieldID),
Files: []string{fullPath},
Metadata: map[string]string{"memory_size": fmt.Sprintf("%d", int64(len(statsBlob.Value)))},
})
if err != nil {
return fmt.Errorf("writeStatsV3: failed to add bloom filter stats to manifest: %w", err)
}
pw.manifest = newManifest
// Stats are stored in the manifest; leave pw.statsLog nil.
}
// --- BM25 stats ---
bm25Blobs, err := pw.bm25Collector.SerializeBlobs()
if err != nil {
return err
}
if len(bm25Blobs) > 0 {
var statEntries []packed.StatEntry
for fieldID, blob := range bm25Blobs {
id, err := pw.allocator.AllocOne()
if err != nil {
return err
}
fullPath := path.Join(basePath, fmt.Sprintf("_stats/bm25.%d/%d", fieldID, id))
if err := packed.WriteFile(pw.storageConfig, fullPath, blob.Value); err != nil {
return fmt.Errorf("writeStatsV3: failed to write bm25 stats: %w", err)
}
statEntries = append(statEntries, packed.StatEntry{
Key: fmt.Sprintf("bm25.%d", fieldID),
Files: []string{fullPath},
Metadata: map[string]string{"memory_size": fmt.Sprintf("%d", blob.MemorySize)},
})
}
newManifest, err := packed.AddStatsToManifest(pw.manifest, pw.storageConfig, statEntries)
for fieldID, blob := range bm25Blobs {
id, err := pw.allocator.AllocOne()
if err != nil {
return fmt.Errorf("writeStatsV3: failed to add bm25 stats to manifest: %w", err)
return err
}
pw.manifest = newManifest
// Stats are stored in the manifest; leave pw.bm25StatsLog nil.
fullPath := path.Join(pw.basePath, fmt.Sprintf("_stats/bm25.%d/%d", fieldID, id))
if err := packed.WriteFile(pw.storageConfig, fullPath, blob.Value); err != nil {
return fmt.Errorf("appendV3Stats: failed to write bm25 stats: %w", err)
}
updates.Stats = append(updates.Stats, packed.StatEntry{
Key: fmt.Sprintf("bm25.%d", fieldID),
Files: []string{fullPath},
Metadata: map[string]string{"memory_size": fmt.Sprintf("%d", blob.MemorySize)},
})
}
return nil
}
@@ -601,11 +593,11 @@ func newPackedManifestRecordWriter(collectionID, partitionID, segmentID UniqueID
var _ BinlogRecordWriter = (*PackedTextManifestRecordWriter)(nil)
// PackedTextManifestRecordWriter wraps packedTextManifestWriter for TEXT column compaction.
// PackedTextManifestRecordWriter wraps packedTextBatchWriter for TEXT column compaction.
// this writer is used during compaction when TEXT columns need REWRITE_ALL strategy.
type PackedTextManifestRecordWriter struct {
packedBinlogRecordWriterBase
writer *packedTextManifestWriter
writer *packedTextBatchWriter
textColumnConfigs []packed.TextColumnConfig
}
@@ -654,8 +646,8 @@ func (pw *PackedTextManifestRecordWriter) initWriters(r Record) error {
var err error
k := metautil.JoinIDPath(pw.collectionID, pw.partitionID, pw.segmentID)
basePath := path.Join(pw.storageConfig.GetRootPath(), common.SegmentInsertLogPath, k)
pw.writer, err = NewPackedTextManifestWriter(pw.storageConfig.GetBucketName(), basePath, -1, pw.schema, pw.bufferSize, pw.multiPartUploadSize, pw.columnGroups, pw.storageConfig, pw.textColumnConfigs)
pw.basePath = path.Join(pw.storageConfig.GetRootPath(), common.SegmentInsertLogPath, k)
pw.writer, err = NewPackedTextBatchWriter(pw.storageConfig.GetBucketName(), pw.basePath, pw.schema, pw.bufferSize, pw.multiPartUploadSize, pw.columnGroups, pw.storageConfig, pw.textColumnConfigs)
if err != nil {
return merr.WrapErrServiceInternal(fmt.Sprintf("can not new packed text writer %s", err.Error()))
}
@@ -689,20 +681,33 @@ func (pw *PackedTextManifestRecordWriter) finalizeBinlogs() {
FieldNullCounts: pw.getFieldNullCountsForColumnGroup(columnGroup),
})
}
pw.manifest = pw.writer.GetWrittenManifest()
}
// Close finalizes the text-column segment using the same do-then-commit
// pattern as PackedManifestRecordWriter.Close.
func (pw *PackedTextManifestRecordWriter) Close() error {
if pw.writer != nil {
if err := pw.writer.Close(); err != nil {
return err
}
if pw.writer == nil {
return pw.writeStats()
}
pw.finalizeBinlogs()
if err := pw.writeStats(); err != nil {
out, err := pw.writer.Close()
if err != nil {
return err
}
return nil
if out != nil {
defer out.Destroy()
}
pw.finalizeBinlogs()
if out == nil {
return pw.writeStats()
}
newManifest, err := packed.CommitManifestUpdates(pw.basePath, packed.ManifestEarliest, pw.storageConfig,
&packed.ManifestUpdates{NewFiles: out})
if err != nil {
return fmt.Errorf("PackedTextManifestRecordWriter.Close commit: %w", err)
}
pw.manifest = newManifest
return pw.writeStats()
}
// NewPackedTextManifestRecordWriter creates a new BinlogRecordWriter for TEXT column compaction.
internal/storage/binlog_record_writer_v3_test.go
+77
View File
@@ -0,0 +1,77 @@
// 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 storage
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/allocator"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
)
// TestPackedManifestRecordWriter_CloseWithoutWrite verifies the no-data
// short-circuit: Close on a freshly-constructed PackedManifestRecordWriter
// must not call into the FFI commit path and must not produce a manifest.
func TestPackedManifestRecordWriter_CloseWithoutWrite(t *testing.T) {
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
{FieldID: common.TimeStampField, DataType: schemapb.DataType_Int64},
{FieldID: common.RowIDField, DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
}}
dir := t.TempDir()
cfg := &indexpb.StorageConfig{StorageType: "local", RootPath: dir}
w, err := newPackedManifestRecordWriter(1, 2, 3, schema,
ChunkedBlobsWriter(func(_ []*Blob) error { return nil }),
allocator.NewLocalAllocator(1, 1<<20),
1024, 0, 0, nil, cfg, nil)
require.NoError(t, err)
// No Write before Close. The internal `writer` field stays nil so
// Close must return immediately without invoking the FFI commit.
require.NoError(t, w.Close())
_, statsLog, _, manifestPath, _ := w.GetLogs()
assert.Empty(t, manifestPath, "no Write means no manifest should be produced")
assert.Nil(t, statsLog, "no Write means no statsLog should be produced")
}
// TestPackedTextManifestRecordWriter_CloseWithoutWrite is the parallel
// short-circuit test for the text writer: Close with no Writes must
// still succeed without touching the FFI segment writer.
func TestPackedTextManifestRecordWriter_CloseWithoutWrite(t *testing.T) {
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
{FieldID: common.TimeStampField, DataType: schemapb.DataType_Int64},
{FieldID: common.RowIDField, DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, DataType: schemapb.DataType_Text, Name: "doc"},
}}
dir := t.TempDir()
cfg := &indexpb.StorageConfig{StorageType: "local", RootPath: dir}
w, err := NewPackedTextManifestRecordWriter(1, 2, 3, schema,
ChunkedBlobsWriter(func(_ []*Blob) error { return nil }),
allocator.NewLocalAllocator(1, 1<<20),
1024, 0, 0, nil, cfg, nil)
require.NoError(t, err)
// No Write before Close. The text writer's nil-handling path must
// fall through to writeStats() and exit cleanly.
require.NoError(t, w.Close())
_, _, _, manifestPath, _ := w.GetLogs()
assert.Empty(t, manifestPath, "no Write means no manifest should be produced")
}
internal/storage/record_writer.go
+59 -61
View File
@@ -180,7 +180,7 @@ func NewPackedRecordWriter(
}, nil
}
type packedRecordManifestWriter struct {
type packedRecordBatchWriter struct {
writer *packed.FFIPackedWriter
bufferSize int64
columnGroups []storagecommon.ColumnGroup
@@ -192,11 +192,10 @@ type packedRecordManifestWriter struct {
writtenUncompressed uint64
columnGroupUncompressed map[typeutil.UniqueID]uint64
columnGroupCompressed map[typeutil.UniqueID]uint64
outputManifest string
storageConfig *indexpb.StorageConfig
}
func (pw *packedRecordManifestWriter) Write(r Record) error {
func (pw *packedRecordBatchWriter) Write(r Record) error {
var rec arrow.Record
sar, ok := r.(*simpleArrowRecord)
if !ok {
@@ -225,63 +224,63 @@ func (pw *packedRecordManifestWriter) Write(r Record) error {
return pw.writer.WriteRecordBatch(rec)
}
func (pw *packedRecordManifestWriter) GetWrittenUncompressed() uint64 {
func (pw *packedRecordBatchWriter) GetWrittenUncompressed() uint64 {
return pw.writtenUncompressed
}
func (pw *packedRecordManifestWriter) GetColumnGroupWrittenUncompressed(columnGroup typeutil.UniqueID) uint64 {
func (pw *packedRecordBatchWriter) GetColumnGroupWrittenUncompressed(columnGroup typeutil.UniqueID) uint64 {
if size, ok := pw.columnGroupUncompressed[columnGroup]; ok {
return size
}
return 0
}
func (pw *packedRecordManifestWriter) GetColumnGroupWrittenCompressed(columnGroup typeutil.UniqueID) uint64 {
func (pw *packedRecordBatchWriter) GetColumnGroupWrittenCompressed(columnGroup typeutil.UniqueID) uint64 {
if size, ok := pw.columnGroupCompressed[columnGroup]; ok {
return size
}
return 0
}
func (pw *packedRecordManifestWriter) GetWrittenPaths(columnGroup typeutil.UniqueID) string {
func (pw *packedRecordBatchWriter) GetWrittenPaths(columnGroup typeutil.UniqueID) string {
if path, ok := pw.pathsMap[columnGroup]; ok {
return path
}
return ""
}
func (pw *packedRecordManifestWriter) GetWrittenManifest() string {
return pw.outputManifest
}
func (pw *packedRecordManifestWriter) GetWrittenRowNum() int64 {
func (pw *packedRecordBatchWriter) GetWrittenRowNum() int64 {
return pw.rowNum
}
func (pw *packedRecordManifestWriter) Close() error {
if pw.writer != nil {
manifest, err := pw.writer.Close()
if err != nil {
return err
}
pw.outputManifest = manifest
for id := range pw.pathsMap {
pw.columnGroupCompressed[id] = uint64(0)
}
// Close closes the underlying FFI writer and returns the resulting
// column-groups payload. The writer never touches the manifest — the
// caller passes the returned handle to packed.CommitManifestUpdates and
// calls Destroy on it when done.
func (pw *packedRecordBatchWriter) Close() (packed.WriterOutput, error) {
if pw.writer == nil {
return nil, nil
}
return nil
out, err := pw.writer.Close()
if err != nil {
return nil, err
}
pw.writer = nil
for id := range pw.pathsMap {
pw.columnGroupCompressed[id] = uint64(0)
}
return out, nil
}
func NewPackedRecordManifestWriter(
func NewPackedRecordBatchWriter(
basePath string,
baseVersion int64,
schema *schemapb.CollectionSchema,
bufferSize int64,
multiPartUploadSize int64,
columnGroups []storagecommon.ColumnGroup,
storageConfig *indexpb.StorageConfig,
storagePluginContext *indexcgopb.StoragePluginContext,
) (*packedRecordManifestWriter, error) {
) (*packedRecordBatchWriter, error) {
// Validate PK field exists before proceeding
_, err := typeutil.GetPrimaryFieldSchema(schema)
if err != nil {
@@ -294,7 +293,7 @@ func NewPackedRecordManifestWriter(
fmt.Sprintf("can not convert collection schema %s to arrow schema: %s", schema.Name, err.Error()))
}
writer, err := packed.NewFFIPackedWriter(basePath, baseVersion, arrowSchema, columnGroups, storageConfig, storagePluginContext)
writer, err := packed.NewFFIPackedWriter(basePath, arrowSchema, columnGroups, storageConfig, storagePluginContext)
if err != nil {
return nil, merr.WrapErrServiceInternal(
fmt.Sprintf("can not new packed record writer %s", err.Error()))
@@ -312,7 +311,7 @@ func NewPackedRecordManifestWriter(
pathsMap[columnGroup.GroupID] = path.Join(basePath, strconv.FormatInt(columnGroup.GroupID, 10), strconv.FormatInt(start, 10))
}
return &packedRecordManifestWriter{
return &packedRecordBatchWriter{
writer: writer,
schema: schema,
arrowSchema: arrowSchema,
@@ -339,9 +338,9 @@ func NewPackedSerializeWriter(bucketName string, paths []string, schema *schemap
}, batchSize), nil
}
// packedTextManifestWriter wraps FFISegmentWriter for TEXT column support during compaction.
// packedTextBatchWriter wraps FFISegmentWriter for TEXT column support during compaction.
// it handles TEXT column rewriting with LOB file management in REWRITE_ALL mode.
type packedTextManifestWriter struct {
type packedTextBatchWriter struct {
writer *packed.FFISegmentWriter
bufferSize int64
columnGroups []storagecommon.ColumnGroup
@@ -353,11 +352,10 @@ type packedTextManifestWriter struct {
writtenUncompressed uint64
columnGroupUncompressed map[typeutil.UniqueID]uint64
columnGroupCompressed map[typeutil.UniqueID]uint64
outputManifest string
storageConfig *indexpb.StorageConfig
}
func (pw *packedTextManifestWriter) Write(r Record) error {
func (pw *packedTextBatchWriter) Write(r Record) error {
var rec arrow.Record
sar, ok := r.(*simpleArrowRecord)
if !ok {
@@ -386,69 +384,71 @@ func (pw *packedTextManifestWriter) Write(r Record) error {
return pw.writer.Write(rec)
}
func (pw *packedTextManifestWriter) GetWrittenUncompressed() uint64 {
func (pw *packedTextBatchWriter) GetWrittenUncompressed() uint64 {
return pw.writtenUncompressed
}
func (pw *packedTextManifestWriter) GetColumnGroupWrittenUncompressed(columnGroup typeutil.UniqueID) uint64 {
func (pw *packedTextBatchWriter) GetColumnGroupWrittenUncompressed(columnGroup typeutil.UniqueID) uint64 {
if size, ok := pw.columnGroupUncompressed[columnGroup]; ok {
return size
}
return 0
}
func (pw *packedTextManifestWriter) GetColumnGroupWrittenCompressed(columnGroup typeutil.UniqueID) uint64 {
func (pw *packedTextBatchWriter) GetColumnGroupWrittenCompressed(columnGroup typeutil.UniqueID) uint64 {
if size, ok := pw.columnGroupCompressed[columnGroup]; ok {
return size
}
return 0
}
func (pw *packedTextManifestWriter) GetWrittenPaths(columnGroup typeutil.UniqueID) string {
func (pw *packedTextBatchWriter) GetWrittenPaths(columnGroup typeutil.UniqueID) string {
if path, ok := pw.pathsMap[columnGroup]; ok {
return path
}
return ""
}
func (pw *packedTextManifestWriter) GetWrittenManifest() string {
return pw.outputManifest
}
func (pw *packedTextManifestWriter) GetWrittenRowNum() int64 {
func (pw *packedTextBatchWriter) GetWrittenRowNum() int64 {
return pw.rowNum
}
func (pw *packedTextManifestWriter) Close() error {
if pw.writer != nil {
defer pw.writer.Destroy()
result, err := pw.writer.Close()
if err != nil {
return err
}
pw.outputManifest = result.ManifestPath
pw.rowNum = result.RowsWritten
for id := range pw.pathsMap {
pw.columnGroupCompressed[id] = uint64(0)
}
// Close closes the underlying segment writer and returns the resulting
// column-groups + LOB payload. The writer never touches the manifest —
// the caller passes the returned handle to packed.CommitManifestUpdates
// and calls Destroy on it when done. FFISegmentWriter.Close releases its
// own handle and properties, so no extra cleanup is needed here.
func (pw *packedTextBatchWriter) Close() (packed.WriterOutput, error) {
if pw.writer == nil {
return nil, nil
}
return nil
out, err := pw.writer.Close()
pw.writer = nil
if err != nil {
return nil, err
}
if so, ok := out.(*packed.SegmentOutput); ok {
pw.rowNum = so.RowsWritten()
}
for id := range pw.pathsMap {
pw.columnGroupCompressed[id] = uint64(0)
}
return out, nil
}
// NewPackedTextManifestWriter creates a new writer that uses FFISegmentWriter with TEXT column support.
// NewPackedTextBatchWriter creates a new writer that uses FFISegmentWriter with TEXT column support.
// this writer is used during compaction when TEXT columns need REWRITE_ALL strategy.
// textColumnConfigs: TEXT column configurations for REWRITE_ALL fields (nil if no TEXT fields need rewriting)
func NewPackedTextManifestWriter(
func NewPackedTextBatchWriter(
bucketName string,
basePath string,
baseVersion int64,
schema *schemapb.CollectionSchema,
bufferSize int64,
multiPartUploadSize int64,
columnGroups []storagecommon.ColumnGroup,
storageConfig *indexpb.StorageConfig,
textColumnConfigs []packed.TextColumnConfig,
) (*packedTextManifestWriter, error) {
) (*packedTextBatchWriter, error) {
// validate PK field exists before proceeding
_, err := typeutil.GetPrimaryFieldSchema(schema)
if err != nil {
@@ -477,8 +477,6 @@ func NewPackedTextManifestWriter(
// build segment writer config
config := &packed.SegmentWriterConfig{
SegmentPath: basePath,
ReadVersion: baseVersion,
RetryLimit: 3,
TextColumns: textColumnConfigs,
}
@@ -501,7 +499,7 @@ func NewPackedTextManifestWriter(
pathsMap[columnGroup.GroupID] = path.Join(basePath, strconv.FormatInt(columnGroup.GroupID, 10), strconv.FormatInt(start, 10))
}
return &packedTextManifestWriter{
return &packedTextBatchWriter{
writer: writer,
schema: schema,
arrowSchema: arrowSchema,
internal/storagev2/packed/manifest_commit.go
+158
View File
@@ -0,0 +1,158 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
package packed
/*
#cgo pkg-config: milvus_core milvus-storage
#include <stdlib.h>
#include "milvus-storage/ffi_c.h"
*/
import "C"
import (
"fmt"
"unsafe"
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
)
// WriterOutput is the data carrier returned by an FFI writer's Close. It
// owns C memory and must be released via Destroy after the surrounding
// CommitManifestUpdates call returns (success or failure).
//
// Concrete implementations live alongside the writers that produce them:
// - *ColumnGroups (FFIPackedWriter.Close)
// - *SegmentOutput (FFISegmentWriter.Close)
//
// Each implementation knows how to stage its payload onto a loon
// transaction handle via the package-internal applyTo method.
type WriterOutput interface {
// Destroy releases the underlying C resources. Idempotent.
Destroy()
// applyTo stages the output onto a loon transaction handle. Called by
// applyManifestUpdates as part of CommitManifestUpdates.
applyTo(handle C.LoonTransactionHandle) error
}
// ManifestUpdates bundles every data-file-level change a single caller
// wants to apply atomically to a manifest. The slow file writes (parquet,
// deltalog, stat blobs) happen before this payload is assembled;
// CommitManifestUpdates then opens a transaction, applies everything in
// one shot, and commits.
//
// NewFiles holds C memory produced by an FFI writer and MUST be released
// by the caller via Destroy after CommitManifestUpdates returns (success
// or failure).
type ManifestUpdates struct {
// NewFiles is the column-groups / LOB payload returned by an FFI
// writer's Close. nil if no insert files were written.
NewFiles WriterOutput
// DeltaLogs is the list of delta-log entries to register.
DeltaLogs []DeltaLogEntry
// Stats is the list of stat entries (bloom filter, bm25, etc.) to
// register. Each entry's Files / Metadata replace any existing entry
// with the same Key (loon overwrite semantics).
Stats []StatEntry
}
// isEmpty short-circuits CommitManifestUpdates when the caller assembled
// no work. A non-nil NewFiles is treated as real work; callers must not
// pass an already-destroyed payload here (the C transaction would then
// commit with zero staged ops and the loon side would return
// "Cannot commit: no updates recorded").
func (u *ManifestUpdates) isEmpty() bool {
if u == nil {
return true
}
if u.NewFiles != nil {
return false
}
return len(u.DeltaLogs) == 0 && len(u.Stats) == 0
}
// CommitManifestUpdates opens a loon transaction at (basePath, baseVersion),
// applies every change in updates, commits, and returns the new manifest
// path. With no effective changes it returns the unchanged manifest path
// without opening a transaction.
//
// The caller retains ownership of any WriterOutput referenced by updates
// and is responsible for calling Destroy on it after CommitManifestUpdates
// returns.
func CommitManifestUpdates(basePath string, baseVersion int64,
storageConfig *indexpb.StorageConfig, updates *ManifestUpdates,
) (string, error) {
if updates.isEmpty() {
return MarshalManifestPath(basePath, baseVersion), nil
}
cProperties, err := MakePropertiesFromStorageConfig(storageConfig, nil)
if err != nil {
return "", fmt.Errorf("commit manifest: %w", err)
}
defer C.loon_properties_free(cProperties)
cBasePath := C.CString(basePath)
defer C.free(unsafe.Pointer(cBasePath))
// Use OVERWRITE for the whole bundle: stats updates rely on
// overwrite-on-key semantics (loon_transaction_update_stat replaces
// existing entries with the same key), while column-group appends and
// delta-log adds don't collide on a key in the first place. A single
// resolve mode keeps the API simple and matches the prior
// AddStatsToManifest behavior.
var handle C.LoonTransactionHandle
res := C.loon_transaction_begin(cBasePath, cProperties,
C.int64_t(baseVersion),
C.LOON_TRANSACTION_RESOLVE_OVERWRITE,
getRetryLimit(), &handle)
if err := HandleLoonFFIResult(res); err != nil {
return "", fmt.Errorf("commit manifest begin: %w", err)
}
defer C.loon_transaction_destroy(handle)
if err := applyManifestUpdates(handle, updates); err != nil {
return "", err
}
var commitVersion C.int64_t
res = C.loon_transaction_commit(handle, &commitVersion)
if err := HandleLoonFFIResult(res); err != nil {
return "", fmt.Errorf("commit manifest commit: %w", err)
}
return MarshalManifestPath(basePath, int64(commitVersion)), nil
}
// applyManifestUpdates stages every operation in updates onto the loon
// transaction handle.
func applyManifestUpdates(handle C.LoonTransactionHandle, updates *ManifestUpdates) error {
if updates.NewFiles != nil {
if err := updates.NewFiles.applyTo(handle); err != nil {
return err
}
}
for _, entry := range updates.DeltaLogs {
cPath := C.CString(entry.Path)
err := HandleLoonFFIResult(C.loon_transaction_add_delta_log(handle, cPath, C.int64_t(entry.NumEntries)))
C.free(unsafe.Pointer(cPath))
if err != nil {
return fmt.Errorf("commit manifest add_delta_log: %w", err)
}
}
for _, entry := range updates.Stats {
if err := UpdateTransactionStat(handle, entry.Key, entry.Files, entry.Metadata); err != nil {
return fmt.Errorf("commit manifest update_stat: %w", err)
}
}
return nil
}
internal/storagev2/packed/manifest_commit_test.go
+295
View File
@@ -0,0 +1,295 @@
// 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")
}
internal/storagev2/packed/packed_reader_ffi_test.go
+40 -25
View File
@@ -93,25 +93,29 @@ func TestPackedFFIReader(t *testing.T) {
{Columns: []int{0, 1}, GroupID: storagecommon.DefaultShortColumnGroupID},
}
storageConfig := &indexpb.StorageConfig{
RootPath: dir,
StorageType: "local",
}
// Create FFI packed writer and write data
pw, err := NewFFIPackedWriter(basePath, version, schema, columnGroups, nil, nil)
pw, err := NewFFIPackedWriter(basePath, schema, columnGroups, storageConfig, nil)
require.NoError(t, err)
err = pw.WriteRecordBatch(rec)
require.NoError(t, err)
manifest, err := pw.Close()
out, err := pw.Close()
require.NoError(t, err)
defer out.Destroy()
manifest, err := CommitManifestUpdates(basePath, version, storageConfig,
&ManifestUpdates{NewFiles: out})
require.NoError(t, err)
require.NotEmpty(t, manifest)
t.Logf("Successfully wrote %d rows with %d-dim float vectors, manifest: %s", numRows, dim, manifest)
// Create storage config for reader
storageConfig := &indexpb.StorageConfig{
RootPath: dir,
StorageType: "local",
}
// Create FFI packed reader
neededColumns := []string{"pk", "vector"}
reader, err := NewFFIPackedReader(manifest, schema, neededColumns, 8192, storageConfig, nil)
@@ -238,22 +242,27 @@ func TestPackedFFIReaderPartialColumns(t *testing.T) {
{Columns: []int{0, 1, 2}, GroupID: storagecommon.DefaultShortColumnGroupID},
}
// Write data
pw, err := NewFFIPackedWriter(basePath, version, schema, columnGroups, nil, nil)
require.NoError(t, err)
err = pw.WriteRecordBatch(rec)
require.NoError(t, err)
manifest, err := pw.Close()
require.NoError(t, err)
// Create storage config
storageConfig := &indexpb.StorageConfig{
RootPath: dir,
StorageType: "local",
}
// Write data
pw, err := NewFFIPackedWriter(basePath, schema, columnGroups, storageConfig, nil)
require.NoError(t, err)
err = pw.WriteRecordBatch(rec)
require.NoError(t, err)
out, err := pw.Close()
require.NoError(t, err)
defer out.Destroy()
manifest, err := CommitManifestUpdates(basePath, version, storageConfig,
&ManifestUpdates{NewFiles: out})
require.NoError(t, err)
// Read only pk and score columns (skip vector)
neededColumns := []string{"pk", "score"}
partialSchema := arrow.NewSchema([]arrow.Field{
@@ -343,6 +352,11 @@ func TestPackedFFIReaderMultipleBatches(t *testing.T) {
var manifest string
totalWrittenRows := 0
storageConfig := &indexpb.StorageConfig{
RootPath: dir,
StorageType: "local",
}
// Write multiple batches
for batch := 0; batch < numWrites; batch++ {
b := array.NewRecordBuilder(memory.DefaultAllocator, schema)
@@ -364,13 +378,18 @@ func TestPackedFFIReaderMultipleBatches(t *testing.T) {
rec := b.NewRecord()
pw, err := NewFFIPackedWriter(basePath, version, schema, columnGroups, nil, nil)
pw, err := NewFFIPackedWriter(basePath, schema, columnGroups, storageConfig, nil)
require.NoError(t, err)
err = pw.WriteRecordBatch(rec)
require.NoError(t, err)
manifest, err = pw.Close()
out, err := pw.Close()
require.NoError(t, err)
manifest, err = CommitManifestUpdates(basePath, version, storageConfig,
&ManifestUpdates{NewFiles: out})
out.Destroy()
require.NoError(t, err)
_, version, err = UnmarshalManifestPath(manifest)
@@ -382,11 +401,7 @@ func TestPackedFFIReaderMultipleBatches(t *testing.T) {
rec.Release()
}
// Read all data
storageConfig := &indexpb.StorageConfig{
RootPath: dir,
StorageType: "local",
}
// Read all data using the same storageConfig
neededColumns := []string{"pk", "vector"}
reader, err := NewFFIPackedReader(manifest, schema, neededColumns, 8192, storageConfig, nil)
internal/storagev2/packed/packed_writer_ffi.go
+77 -47
View File
@@ -28,16 +28,15 @@ package packed
import "C"
import (
"fmt"
"strings"
"unsafe"
"github.com/apache/arrow/go/v17/arrow"
"github.com/apache/arrow/go/v17/arrow/cdata"
"github.com/samber/lo"
"go.uber.org/zap"
"github.com/milvus-io/milvus/internal/storagecommon"
"github.com/milvus-io/milvus/pkg/v3/log"
"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/util/paramtable"
@@ -76,7 +75,12 @@ func CreateStorageConfig() *indexpb.StorageConfig {
return storageConfig
}
func NewFFIPackedWriter(basePath string, baseVersion int64, schema *arrow.Schema, columnGroups []storagecommon.ColumnGroup, storageConfig *indexpb.StorageConfig, storagePluginContext *indexcgopb.StoragePluginContext) (*FFIPackedWriter, error) {
// NewFFIPackedWriter creates a writer that produces parquet files under
// basePath. The writer knows nothing about manifests or versions — its
// only job is to write data files. Close returns the resulting column
// groups, which the caller passes to packed.CommitManifestUpdates to
// register them with a manifest version.
func NewFFIPackedWriter(basePath string, schema *arrow.Schema, columnGroups []storagecommon.ColumnGroup, storageConfig *indexpb.StorageConfig, storagePluginContext *indexcgopb.StoragePluginContext) (*FFIPackedWriter, error) {
cBasePath := C.CString(basePath)
defer C.free(unsafe.Pointer(cBasePath))
@@ -145,16 +149,16 @@ func NewFFIPackedWriter(basePath string, baseVersion int64, schema *arrow.Schema
return &FFIPackedWriter{
basePath: basePath,
baseVersion: baseVersion,
cWriterHandle: writerHandle,
cProperties: cProperties,
}, nil
}
// AsNewColumnGroups marks this writer so that Close() calls loon_transaction_add_column_group
// for each written column group instead of loon_transaction_append_files.
// Call this when the written columns are brand-new additions to an existing manifest
// (e.g. function-field backfill), not when appending more files to the same column structure.
// AsNewColumnGroups marks this writer so that the column groups returned
// by Close should be staged via loon_transaction_add_column_group instead
// of loon_transaction_append_files when later passed to
// CommitManifestUpdates. Use true when adding columns that do not yet
// exist in the manifest (e.g. function-field backfill).
func (pw *FFIPackedWriter) AsNewColumnGroups() *FFIPackedWriter {
pw.addNewColumnGroups = true
return pw
@@ -176,52 +180,78 @@ func (pw *FFIPackedWriter) WriteRecordBatch(recordBatch arrow.Record) error {
return HandleLoonFFIResult(result)
}
func (pw *FFIPackedWriter) Close() (string, error) {
var cColumnGroups *C.LoonColumnGroups
// ColumnGroups is the data carrier returned by FFIPackedWriter.Close. It
// holds the column-groups payload produced by the C writer and owns C
// memory; the caller MUST call Destroy after passing the handle to
// CommitManifestUpdates (success or failure). Destroy is idempotent;
// a nil cColumnGroups indicates the handle has already been released.
type ColumnGroups struct {
cColumnGroups *C.LoonColumnGroups
addNewColumnGroups bool
}
result := C.loon_writer_close(pw.cWriterHandle, nil, nil, 0, &cColumnGroups)
if err := HandleLoonFFIResult(result); err != nil {
return "", err
// Destroy releases C memory. Safe to call multiple times.
func (f *ColumnGroups) Destroy() {
if f == nil || f.cColumnGroups == nil {
return
}
C.loon_column_groups_destroy(f.cColumnGroups)
f.cColumnGroups = nil
}
cBasePath := C.CString(pw.basePath)
defer C.free(unsafe.Pointer(cBasePath))
var transationHandle C.LoonTransactionHandle
result = C.loon_transaction_begin(cBasePath, pw.cProperties, C.int64_t(pw.baseVersion), C.LOON_TRANSACTION_RESOLVE_OVERWRITE /* resolve_id */, getRetryLimit() /* retry_limit */, &transationHandle)
if err := HandleLoonFFIResult(result); err != nil {
return "", err
// applyTo stages the column groups onto a loon transaction.
//
// When addNewColumnGroups is true the groups are added one-by-one via
// loon_transaction_add_column_group (function-backfill case where the
// schema is being extended). Otherwise they are appended in one call via
// loon_transaction_append_files (normal multi-batch write case).
func (f *ColumnGroups) applyTo(handle C.LoonTransactionHandle) error {
if f == nil || f.cColumnGroups == nil {
return nil
}
defer C.loon_transaction_destroy(transationHandle)
if pw.addNewColumnGroups {
// Each written group is a brand-new column group: add them one-by-one.
// loon_transaction_append_files requires the count to match existing groups,
// which would fail when extending the schema (e.g. backfill adds sparse vector
// to a manifest that already has 2 groups of different columns).
numGroups := int(cColumnGroups.num_of_column_groups)
colGroupSlice := (*[1 << 20]C.LoonColumnGroup)(unsafe.Pointer(cColumnGroups.column_group_array))[:numGroups:numGroups]
for i := range colGroupSlice {
result = C.loon_transaction_add_column_group(transationHandle, &colGroupSlice[i])
if err := HandleLoonFFIResult(result); err != nil {
return "", err
if f.addNewColumnGroups {
num := int(f.cColumnGroups.num_of_column_groups)
slice := unsafe.Slice(f.cColumnGroups.column_group_array, num)
for i := range slice {
if err := HandleLoonFFIResult(C.loon_transaction_add_column_group(handle, &slice[i])); err != nil {
return fmt.Errorf("commit manifest add_column_group: %w", err)
}
}
} else {
result = C.loon_transaction_append_files(transationHandle, cColumnGroups)
if err := HandleLoonFFIResult(result); err != nil {
return "", err
}
return nil
}
var cCommitVersion C.int64_t
result = C.loon_transaction_commit(transationHandle, &cCommitVersion)
if err := HandleLoonFFIResult(result); err != nil {
return "", err
if err := HandleLoonFFIResult(C.loon_transaction_append_files(handle, f.cColumnGroups)); err != nil {
return fmt.Errorf("commit manifest append_files: %w", err)
}
return nil
}
log.Info("FFI writer closed", zap.Int64("version", int64(cCommitVersion)))
defer C.loon_properties_free(pw.cProperties)
return MarshalManifestPath(pw.basePath, int64(cCommitVersion)), nil
// Close closes the underlying loon writer and returns the column-groups
// payload. The writer never touches the manifest — the caller is
// responsible for passing the returned handle to CommitManifestUpdates
// and calling Destroy when done.
//
// Close releases the writer's C resources (loon writer handle and
// cProperties) in a defer, so even when loon_writer_close fails those
// resources are reclaimed. After Close the writer is exhausted; further
// Close or Write calls fail.
func (pw *FFIPackedWriter) Close() (WriterOutput, error) {
if pw.closed {
return nil, fmt.Errorf("FFIPackedWriter already closed")
}
pw.closed = true
defer func() {
if pw.cProperties != nil {
C.loon_properties_free(pw.cProperties)
pw.cProperties = nil
}
}()
var cColumnGroups *C.LoonColumnGroups
result := C.loon_writer_close(pw.cWriterHandle, nil, nil, 0, &cColumnGroups)
if err := HandleLoonFFIResult(result); err != nil {
return nil, err
}
return &ColumnGroups{
cColumnGroups: cColumnGroups,
addNewColumnGroups: pw.addNewColumnGroups,
}, nil
}
internal/storagev2/packed/packed_writer_ffi_test.go
+62 -3
View File
@@ -93,15 +93,21 @@ func TestPackedFFIWriter(t *testing.T) {
}
// Create FFI packed writer
pw, err := NewFFIPackedWriter(basePath, version, schema, columnGroups, nil, nil)
cfg := CreateStorageConfig()
pw, err := NewFFIPackedWriter(basePath, schema, columnGroups, cfg, nil)
require.NoError(t, err)
// Write record batch
err = pw.WriteRecordBatch(rec)
require.NoError(t, err)
// Close writer and get manifest
manifest, err := pw.Close()
// Close writer to obtain column groups, commit via manifest update.
out, err := pw.Close()
require.NoError(t, err)
manifest, err := CommitManifestUpdates(basePath, version, cfg,
&ManifestUpdates{NewFiles: out})
out.Destroy()
require.NoError(t, err)
require.NotEmpty(t, manifest)
@@ -114,3 +120,56 @@ func TestPackedFFIWriter(t *testing.T) {
t.Logf("Successfully wrote %d rows with %d-dim float vectors, manifest: %s", numRows, dim, manifest)
}
}
func TestFFIPackedWriter_CloseThenCommitUpdates(t *testing.T) {
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)
})
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},
}
basePath := "files/close_commit_test/1"
cfg := CreateStorageConfig()
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)
require.NotNil(t, out)
defer out.Destroy()
mfPath, err := CommitManifestUpdates(basePath, ManifestEarliest, cfg,
&ManifestUpdates{NewFiles: out})
require.NoError(t, err)
_, v, err := UnmarshalManifestPath(mfPath)
require.NoError(t, err)
require.Equal(t, int64(1), v, "exactly one version bump expected")
}
internal/storagev2/packed/segment_writer_ffi.go
+83 -82
View File
@@ -44,29 +44,20 @@ type TextColumnConfig struct {
RewriteMode bool // true = input is LOB references, decode & rewrite during compaction
}
// SegmentWriterConfig represents configuration for SegmentWriter.
// ReadVersion and RetryLimit are used by the Go-layer transaction logic,
// SegmentWriterConfig represents configuration for SegmentWriter. The
// writer is concerned only with file output; manifest-level concerns
// (version, retry) live in CommitManifestUpdates.
type SegmentWriterConfig struct {
SegmentPath string
ReadVersion int64 // manifest version for transaction (Go-layer only)
RetryLimit uint32 // transaction retry limit (Go-layer only)
TextColumns []TextColumnConfig
}
// SegmentWriterResult contains the result of closing a SegmentWriter.
type SegmentWriterResult struct {
ManifestPath string
CommittedVersion int64
RowsWritten int64
}
// FFISegmentWriter wraps the C SegmentWriter handle for incremental writes.
type FFISegmentWriter struct {
handle C.LoonSegmentWriterHandle
cProperties *C.LoonProperties
schema *arrow.Schema
basePath string // segment base path for transaction
readVersion int64 // manifest read version for transaction
closed bool
}
// NewFFISegmentWriter creates a new segment writer via FFI.
@@ -107,8 +98,6 @@ func NewFFISegmentWriter(
handle: writerHandle,
cProperties: cProperties,
schema: schema,
basePath: config.SegmentPath,
readVersion: config.ReadVersion,
}, nil
}
@@ -127,85 +116,97 @@ func (w *FFISegmentWriter) Write(record arrow.Record) error {
return HandleLoonFFIResult(result)
}
// Flush flushes buffered data to storage.
func (w *FFISegmentWriter) Flush() error {
// SyncBuffered syncs buffered data to storage without closing the writer.
// Used for mid-stream flushes; Close detaches output and finalizes the writer.
func (w *FFISegmentWriter) SyncBuffered() error {
result := C.loon_segment_writer_flush(w.handle)
return HandleLoonFFIResult(result)
}
// Close closes the writer and commits the manifest via Transaction.
// Returns SegmentWriterResult with the committed manifest path.
// Pattern matches FFIPackedWriter.Close(): C++ writer returns ColumnGroups + LobFiles,
// Go layer handles Transaction begin → append_files → add_lob_files → commit.
func (w *FFISegmentWriter) Close() (*SegmentWriterResult, error) {
var cOutput C.LoonSegmentWriteOutput
// SegmentOutput is the data carrier returned by FFISegmentWriter.Close.
// It holds the column-groups + LOB payload produced by the C writer and
// owns C memory; the caller MUST call Destroy after passing the handle to
// CommitManifestUpdates (success or failure). Destroy is idempotent; a
// nil column_groups pointer indicates the handle has already been released.
type SegmentOutput struct {
cOutput C.LoonSegmentWriteOutput
rowsWritten int64
}
// RowsWritten returns the number of rows the segment writer reported.
func (f *SegmentOutput) RowsWritten() int64 { return f.rowsWritten }
// Destroy releases the C output (LOB file strings + array). Safe to call
// multiple times — the nil column_groups pointer left behind serves as
// the already-released marker.
func (f *SegmentOutput) Destroy() {
if f == nil || f.cOutput.column_groups == nil {
return
}
C.loon_segment_write_output_free(&f.cOutput)
f.cOutput.column_groups = nil
f.cOutput.lob_files = nil
f.cOutput.num_lob_files = 0
}
// applyTo stages the column-groups + LOB payload onto a loon transaction.
// Column groups (when present) are appended via loon_transaction_append_files,
// and each LOB file is registered via loon_transaction_add_lob_file.
func (f *SegmentOutput) applyTo(handle C.LoonTransactionHandle) error {
if f == nil {
return nil
}
if f.cOutput.column_groups != nil {
if err := HandleLoonFFIResult(C.loon_transaction_append_files(handle, f.cOutput.column_groups)); err != nil {
return fmt.Errorf("commit manifest append_files (segment): %w", err)
}
}
if f.cOutput.num_lob_files > 0 && f.cOutput.lob_files != nil {
lob := unsafe.Slice(f.cOutput.lob_files, f.cOutput.num_lob_files)
for i := range lob {
if err := HandleLoonFFIResult(C.loon_transaction_add_lob_file(handle, &lob[i])); err != nil {
return fmt.Errorf("commit manifest add_lob_file: %w", err)
}
}
}
return nil
}
// Close closes the underlying segment writer and returns the column-groups
// + LOB payload. The writer never touches the manifest — the caller is
// responsible for passing the returned handle to CommitManifestUpdates
// and calling Destroy on the returned WriterOutput when done.
//
// Close releases the writer's C resources (segment-writer handle and
// cProperties) in a defer, so even when loon_segment_writer_close fails
// those resources are reclaimed. After Close the writer is exhausted;
// further Close or Write calls fail.
func (w *FFISegmentWriter) Close() (WriterOutput, error) {
if w.closed {
return nil, fmt.Errorf("FFISegmentWriter already closed")
}
w.closed = true
defer func() {
if w.handle != 0 {
C.loon_segment_writer_destroy(w.handle)
w.handle = 0
}
if w.cProperties != nil {
C.loon_properties_free(w.cProperties)
w.cProperties = nil
}
}()
var cOutput C.LoonSegmentWriteOutput
result := C.loon_segment_writer_close(w.handle, &cOutput)
if err := HandleLoonFFIResult(result); err != nil {
return nil, err
}
rowsWritten := int64(cOutput.rows_written)
cBasePath := C.CString(w.basePath)
defer C.free(unsafe.Pointer(cBasePath))
var transactionHandle C.LoonTransactionHandle
result = C.loon_transaction_begin(cBasePath, w.cProperties,
C.int64_t(w.readVersion), C.int32_t(0), getRetryLimit(), &transactionHandle)
if err := HandleLoonFFIResult(result); err != nil {
return nil, err
}
defer C.loon_transaction_destroy(transactionHandle)
// append column groups
if cOutput.column_groups != nil {
result = C.loon_transaction_append_files(transactionHandle, cOutput.column_groups)
if err := HandleLoonFFIResult(result); err != nil {
return nil, err
}
}
// add LOB files
if cOutput.num_lob_files > 0 && cOutput.lob_files != nil {
cLobSlice := unsafe.Slice(cOutput.lob_files, cOutput.num_lob_files)
for _, cLob := range cLobSlice {
result = C.loon_transaction_add_lob_file(transactionHandle, &cLob)
if err := HandleLoonFFIResult(result); err != nil {
return nil, err
}
}
}
// commit
var cCommitVersion C.int64_t
result = C.loon_transaction_commit(transactionHandle, &cCommitVersion)
if err := HandleLoonFFIResult(result); err != nil {
return nil, err
}
// free C output (LOB file strings + array)
C.loon_segment_write_output_free(&cOutput)
return &SegmentWriterResult{
ManifestPath: MarshalManifestPath(w.basePath, int64(cCommitVersion)),
CommittedVersion: int64(cCommitVersion),
RowsWritten: rowsWritten,
return &SegmentOutput{
cOutput: cOutput,
rowsWritten: int64(cOutput.rows_written),
}, nil
}
// Destroy destroys the writer and releases resources.
func (w *FFISegmentWriter) Destroy() {
if w.handle != 0 {
C.loon_segment_writer_destroy(w.handle)
w.handle = 0
}
if w.cProperties != nil {
C.loon_properties_free(w.cProperties)
w.cProperties = nil
}
}
// buildCSegmentWriterConfig converts Go config to C config.
func buildCSegmentWriterConfig(config *SegmentWriterConfig) *C.LoonSegmentWriterConfig {
cConfig := (*C.LoonSegmentWriterConfig)(C.malloc(C.sizeof_LoonSegmentWriterConfig))
internal/storagev2/packed/transaction_test.go
+7 -2
View File
@@ -62,7 +62,7 @@ func createBaseManifest(t *testing.T, basePath string, storageConfig *indexpb.St
{Columns: []int{0, 1}, GroupID: storagecommon.DefaultShortColumnGroupID},
}
pw, err := NewFFIPackedWriter(basePath, 0, schema, columnGroups, storageConfig, nil)
pw, err := NewFFIPackedWriter(basePath, schema, columnGroups, storageConfig, nil)
require.NoError(t, err)
b := array.NewRecordBuilder(memory.DefaultAllocator, schema)
@@ -75,7 +75,12 @@ func createBaseManifest(t *testing.T, basePath string, storageConfig *indexpb.St
err = pw.WriteRecordBatch(rec)
require.NoError(t, err)
manifestPath, err := pw.Close()
out, err := pw.Close()
require.NoError(t, err)
defer out.Destroy()
manifestPath, err := CommitManifestUpdates(basePath, ManifestEarliest, storageConfig,
&ManifestUpdates{NewFiles: out})
require.NoError(t, err)
return manifestPath
}
internal/storagev2/packed/type.go
+7 -6
View File
@@ -37,15 +37,16 @@ type PackedWriter struct {
type FFIPackedWriter struct {
basePath string
baseVersion int64
cWriterHandle C.LoonWriterHandle
cProperties *C.LoonProperties
// addNewColumnGroups controls whether Close() uses loon_transaction_add_column_group
// (true) or loon_transaction_append_files (false).
//
// Use true when adding columns that do not yet exist in the manifest (e.g. backfill).
// Use false (default) when writing more files into the same column structure (e.g. multi-batch write).
// addNewColumnGroups is propagated into the ColumnGroups output of
// Close so CommitManifestUpdates can choose between
// loon_transaction_add_column_group (true) and
// loon_transaction_append_files (false). Use true when adding columns
// that do not yet exist in the manifest (e.g. function backfill).
addNewColumnGroups bool
// closed is set once Close has consumed the underlying C writer handle.
closed bool
}
type PackedReader struct {