mirror of
https://github.com/milvus-io/milvus.git
synced 2026-10-12 02:19:10 +00:00
fix: batch clustering compaction output across input records (#53678)
Clustering compaction converts each input row to a Value and flushes each bucket at input Record boundaries. When rows are spread across buckets, this produces many small Arrow Records and amplifies allocation and retained-object overhead in packed writers. Accumulate surviving rows in a RecordBuilder owned by each bucket across input boundaries. Batch at half the binlog budget using allocated Arrow buffer bytes, including null slots, validity/offset buffers and spare capacity; use the same accounting for pending-bucket pressure. Keep payload-size accounting for existing sorting callers and halve the V2 writer buffer option. Preserve filtering, routing, nullable/default values and output statistics, including the Double and Timestamptz default fixes required by the RecordBuilder path. On an append or writer failure, discard the bucket batch under the bucket lock, reject further writes to that bucket, cancel sibling mapping workers and join them before cleanup while preserving the original error. Reject malformed timestamp column types with a typed internal error instead of panicking. Always close writers after flush failures and during task cleanup. Keep batching for V1/V2/V3; project records onto the writer schema, normalize all-null TEXT columns to its binary/UTF8 representation, and return an error for incompatible non-null columns before constructing the Arrow Record. V3 clustering batches use the writer’s physical field names, including ExternalField overrides, so matching batches reuse their Arrow Record. The previous projection path did not copy column payloads; this only removes a record wrapper and column retain/release work per matching batch. Validation: Targeted failing tests reproduced a panic for malformed or absent timestamp columns, continued writes after an injected writer failure, and the V3 builder/writer schema mismatch. The focused internal/storage and internal/datanode/compactor regressions pass after the fixes, including TEXT normalization, V1/V2 persisted readback, the V3 identity fast path, lazy-writer close, and concurrent cancellation. A controlled rebuild of the Milvus native libraries and Go binary passed against the pinned native storage dependency. The full repository static check passed. Memory scope: RecordBuilder.GetMemorySize() counts only the pending batch and resets after Build; V2/V3 writer-retained Arrow/native memory remains outside that counter. Their GetBufferUncompressed() reports the configured multipart buffer size rather than actual retained bytes, and FlushChunk does not force native release. As a result, getBufferTotalUsedMemorySize() and flushLargestBuffers() can undercount or mis-rank writer memory; the high-watermark is not a bound on process RSS. A focused V2 test forces an 8 MiB writer buffer to verify FlushChunk retention under that fixture, not production peak memory. Accurate native-writer accounting and release remain separate follow-up work; native writer implementation and file layout are unchanged. issue: #53432 --------- Signed-off-by: yangxuan <[email protected]>
This commit is contained in:
@@ -28,6 +28,8 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/apache/arrow/go/v17/arrow"
|
||||
"github.com/apache/arrow/go/v17/arrow/array"
|
||||
"github.com/samber/lo"
|
||||
"go.opentelemetry.io/otel"
|
||||
"go.uber.org/atomic"
|
||||
@@ -112,33 +114,96 @@ type clusteringCompactionTask struct {
|
||||
type ClusterBuffer struct {
|
||||
id int
|
||||
writer *MultiSegmentWriter
|
||||
builder *storage.RecordBuilder
|
||||
appendErr error
|
||||
closed bool
|
||||
clusteringKeyFieldStats *storage.FieldStats
|
||||
|
||||
lock sync.RWMutex
|
||||
}
|
||||
|
||||
func (b *ClusterBuffer) Write(v *storage.Value) error {
|
||||
// WriteRecord copies a row into the bucket's batch, independently of the input
|
||||
// Record's lifetime. The writer receives Records sized by bytes, not input rows.
|
||||
func (b *ClusterBuffer) WriteRecord(r storage.Record, row int) error {
|
||||
b.lock.Lock()
|
||||
defer b.lock.Unlock()
|
||||
return b.writer.WriteValue(v)
|
||||
if b.appendErr != nil {
|
||||
return b.appendErr
|
||||
}
|
||||
if b.closed {
|
||||
return merr.WrapErrServiceInternalMsg("cannot write to closed clustering buffer")
|
||||
}
|
||||
if b.builder == nil {
|
||||
if b.writer.storageVersion == storage.StorageV3 {
|
||||
b.builder = storage.NewRecordBuilderWithFieldID(b.writer.schema)
|
||||
} else {
|
||||
b.builder = storage.NewRecordBuilder(b.writer.schema)
|
||||
}
|
||||
}
|
||||
if err := b.builder.Append(r, row, row+1); err != nil {
|
||||
// Append can leave columns at different lengths. Poison this bucket
|
||||
// before another mapping worker can acquire it, and discard the batch.
|
||||
b.appendErr = err
|
||||
b.releaseBuilder()
|
||||
return err
|
||||
}
|
||||
// Leave room for buffering in the underlying writer.
|
||||
if b.builder.GetMemorySize() >= max(uint64(1), b.writer.binLogMaxSize/2) {
|
||||
return b.writeRecord()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (b *ClusterBuffer) Flush() error {
|
||||
b.lock.Lock()
|
||||
defer b.lock.Unlock()
|
||||
return b.writer.Flush()
|
||||
// writeRecord requires b.lock. Submitting a Record may transfer its buffers to
|
||||
// the writer; it does not imply that the underlying writer released memory.
|
||||
func (b *ClusterBuffer) writeRecord() error {
|
||||
if b.appendErr != nil {
|
||||
return b.appendErr
|
||||
}
|
||||
if b.builder == nil || b.builder.GetRowNum() == 0 {
|
||||
return nil
|
||||
}
|
||||
record := b.builder.Build()
|
||||
defer record.Release()
|
||||
if err := b.writer.Write(record); err != nil {
|
||||
// Another mapping worker may reach this bucket before cancellation.
|
||||
// Keep the first write failure and reject every later row.
|
||||
b.appendErr = err
|
||||
b.releaseBuilder()
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (b *ClusterBuffer) FlushChunk() error {
|
||||
b.lock.Lock()
|
||||
defer b.lock.Unlock()
|
||||
if err := b.writeRecord(); err != nil {
|
||||
return err
|
||||
}
|
||||
return b.writer.FlushChunk()
|
||||
}
|
||||
|
||||
func (b *ClusterBuffer) Close() error {
|
||||
b.lock.Lock()
|
||||
defer b.lock.Unlock()
|
||||
return b.writer.Close()
|
||||
defer b.releaseBuilder()
|
||||
if b.closed {
|
||||
return nil
|
||||
}
|
||||
defer func() { b.closed = true }()
|
||||
err := b.writeRecord()
|
||||
closeErr := b.writer.Close()
|
||||
// Keep the write/append failure as the primary cause while retaining any
|
||||
// close failure. Closing must release the writer even after a failed flush.
|
||||
return merr.Combine(closeErr, err)
|
||||
}
|
||||
|
||||
func (b *ClusterBuffer) releaseBuilder() {
|
||||
if b.builder != nil {
|
||||
b.builder.Release()
|
||||
b.builder = nil
|
||||
}
|
||||
}
|
||||
|
||||
func (b *ClusterBuffer) GetCompactionSegments() []*datapb.CompactionSegment {
|
||||
@@ -150,7 +215,11 @@ func (b *ClusterBuffer) GetCompactionSegments() []*datapb.CompactionSegment {
|
||||
func (b *ClusterBuffer) GetBufferSize() uint64 {
|
||||
b.lock.RLock()
|
||||
defer b.lock.RUnlock()
|
||||
return b.writer.GetBufferUncompressed()
|
||||
size := b.writer.GetBufferUncompressed()
|
||||
if b.builder != nil {
|
||||
size += b.builder.GetMemorySize()
|
||||
}
|
||||
return size
|
||||
}
|
||||
|
||||
func newClusterBuffer(id int, writer *MultiSegmentWriter, clusteringKeyFieldStats *storage.FieldStats) *ClusterBuffer {
|
||||
@@ -494,16 +563,31 @@ func (t *clusteringCompactionTask) mapping(ctx context.Context,
|
||||
ctx, span := otel.Tracer(typeutil.DataNodeRole).Start(ctx, fmt.Sprintf("mapping-%d", t.GetPlanID()))
|
||||
defer span.End()
|
||||
inputSegments := t.plan.GetSegmentBinlogs()
|
||||
ctx, cancel := context.WithCancelCause(ctx)
|
||||
defer cancel(nil)
|
||||
mapStart := time.Now()
|
||||
futures := make([]*conc.Future[any], 0, len(inputSegments))
|
||||
for _, segment := range inputSegments {
|
||||
future := t.mappingPool.Submit(func() (any, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
err := t.mappingSegment(ctx, segment)
|
||||
if err != nil {
|
||||
cancel(err)
|
||||
}
|
||||
return struct{}{}, err
|
||||
})
|
||||
futures = append(futures, future)
|
||||
}
|
||||
if err := conc.AwaitAll(futures...); err != nil {
|
||||
// Join every worker before Compact's deferred cleanup releases shared buffers.
|
||||
if err := conc.BlockOnAll(futures...); err != nil {
|
||||
// A sibling may return context.Canceled before the failed future is
|
||||
// visited. Preserve the failure that triggered cancellation. If several
|
||||
// workers fail independently, the first cancellation cause wins.
|
||||
if cause := context.Cause(ctx); cause != nil {
|
||||
return nil, nil, cause
|
||||
}
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
@@ -612,9 +696,19 @@ func (t *clusteringCompactionTask) mappingSegment(
|
||||
defer rr.Close()
|
||||
|
||||
hasTTLField := t.ttlFieldID >= common.StartOfUserFieldID
|
||||
var ttlField *schemapb.FieldSchema
|
||||
if hasTTLField {
|
||||
ttlField = typeutil.GetField(t.plan.Schema, t.ttlFieldID)
|
||||
}
|
||||
|
||||
offset := int64(-1)
|
||||
done := ctx.Done()
|
||||
for {
|
||||
select {
|
||||
case <-done:
|
||||
return ctx.Err()
|
||||
default:
|
||||
}
|
||||
r, err := rr.Next()
|
||||
if err != nil {
|
||||
if err == sio.EOF {
|
||||
@@ -624,40 +718,48 @@ func (t *clusteringCompactionTask) mappingSegment(
|
||||
return err
|
||||
}
|
||||
|
||||
vs := make([]*storage.Value, r.Len())
|
||||
if err = storage.ValueDeserializerWithSchema(r, vs, t.plan.Schema, true); err != nil {
|
||||
mlog.Warn(context.TODO(), "compact wrong, failed to deserialize data", mlog.Err(err))
|
||||
return err
|
||||
pkColumn := r.Column(t.primaryKeyField.GetFieldID())
|
||||
ts, exists := storage.TryRecordColumn(r, common.TimeStampField)
|
||||
tsColumn, ok := ts.(*array.Int64)
|
||||
if !exists || !ok {
|
||||
return merr.WrapErrServiceInternalMsg("timestamp field is not an int64 column in clustering compaction record")
|
||||
}
|
||||
|
||||
for _, v := range vs {
|
||||
for row := 0; row < r.Len(); row++ {
|
||||
select {
|
||||
case <-done:
|
||||
return ctx.Err()
|
||||
default:
|
||||
}
|
||||
offset++
|
||||
|
||||
row, ok := v.Value.(map[typeutil.UniqueID]interface{})
|
||||
if !ok {
|
||||
mlog.Warn(context.TODO(), "convert interface to map wrong")
|
||||
return merr.WrapErrServiceInternalMsg("unexpected error")
|
||||
pk, err := clusteringScalarValue(pkColumn, t.primaryKeyField, row)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
expireTs := int64(-1)
|
||||
if hasTTLField {
|
||||
if val, exists := row[t.ttlFieldID]; exists {
|
||||
if v, ok := val.(int64); ok {
|
||||
expireTs = v
|
||||
}
|
||||
value, err := clusteringScalarValue(r.Column(t.ttlFieldID), ttlField, row)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if timestamp, ok := value.(int64); ok {
|
||||
expireTs = timestamp
|
||||
}
|
||||
}
|
||||
if entityFilter.Filtered(v.PK.GetValue(), uint64(v.Timestamp), expireTs) {
|
||||
if entityFilter.Filtered(pk, uint64(tsColumn.Value(row)), expireTs) {
|
||||
continue
|
||||
}
|
||||
|
||||
clusteringKey := row[t.clusteringKeyField.FieldID]
|
||||
var clusterBuffer *ClusterBuffer
|
||||
if t.isVectorClusteringKey {
|
||||
clusterBuffer = t.offsetToBufferFunc(offset, mappingStats.GetCentroidIdMapping())
|
||||
} else {
|
||||
clusterBuffer = t.keyToBufferFunc(clusteringKey)
|
||||
key, err := clusteringScalarValue(r.Column(t.clusteringKeyField.FieldID), t.clusteringKeyField, row)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
clusterBuffer = t.keyToBufferFunc(key)
|
||||
}
|
||||
if err := clusterBuffer.Write(v); err != nil {
|
||||
if err := clusterBuffer.WriteRecord(r, row); err != nil {
|
||||
return err
|
||||
}
|
||||
t.writtenRowNum.Inc()
|
||||
@@ -675,11 +777,6 @@ func (t *clusteringCompactionTask) mappingSegment(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// all cluster buffers are flushed for a certain record, since the values read from the same record are references instead of copies
|
||||
for _, buffer := range t.clusterBuffers {
|
||||
buffer.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
missing := entityFilter.GetMissingDeleteCount()
|
||||
@@ -698,6 +795,37 @@ func (t *clusteringCompactionTask) mappingSegment(
|
||||
return nil
|
||||
}
|
||||
|
||||
// clusteringScalarValue reads only the routing/filter columns. Keep null and
|
||||
// default-value semantics consistent with the scalar analysis deserializer.
|
||||
func clusteringScalarValue(column arrow.Array, field *schemapb.FieldSchema, row int) (any, error) {
|
||||
if column.IsNull(row) {
|
||||
if field.GetDefaultValue() != nil {
|
||||
return storage.GetDefaultValue(field), nil
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
switch column := column.(type) {
|
||||
case *array.Boolean:
|
||||
return column.Value(row), nil
|
||||
case *array.Int8:
|
||||
return column.Value(row), nil
|
||||
case *array.Int16:
|
||||
return column.Value(row), nil
|
||||
case *array.Int32:
|
||||
return column.Value(row), nil
|
||||
case *array.Int64:
|
||||
return column.Value(row), nil
|
||||
case *array.Float32:
|
||||
return column.Value(row), nil
|
||||
case *array.Float64:
|
||||
return column.Value(row), nil
|
||||
case *array.String:
|
||||
return column.Value(row), nil
|
||||
default:
|
||||
return nil, merr.WrapErrServiceInternalMsg("unsupported clustering scalar column %T", column)
|
||||
}
|
||||
}
|
||||
|
||||
func (t *clusteringCompactionTask) getWorkerPoolSize() int {
|
||||
return int(math.Max(float64(paramtable.Get().DataNodeCfg.ClusteringCompactionWorkerPoolSize.GetAsInt()), 1.0))
|
||||
}
|
||||
@@ -759,7 +887,7 @@ func (t *clusteringCompactionTask) flushLargestBuffers(ctx context.Context) erro
|
||||
break
|
||||
}
|
||||
}
|
||||
if err := conc.AwaitAll(futures...); err != nil {
|
||||
if err := conc.BlockOnAll(futures...); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -780,7 +908,7 @@ func (t *clusteringCompactionTask) flushAll() error {
|
||||
})
|
||||
futures = append(futures, future)
|
||||
}
|
||||
if err := conc.AwaitAll(futures...); err != nil {
|
||||
if err := conc.BlockOnAll(futures...); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
@@ -809,6 +937,17 @@ func (t *clusteringCompactionTask) uploadPartitionStats(ctx context.Context, col
|
||||
|
||||
// cleanUp try best to clean all temp datas
|
||||
func (t *clusteringCompactionTask) cleanUp(ctx context.Context) {
|
||||
for _, buffer := range t.clusterBuffers {
|
||||
buffer.lock.Lock()
|
||||
buffer.releaseBuilder()
|
||||
// Mapping has joined all workers. Discard pending rows on failure and
|
||||
// close any writer that flushAll did not reach.
|
||||
err := buffer.writer.Close()
|
||||
buffer.lock.Unlock()
|
||||
if err != nil {
|
||||
mlog.RatedWarn(ctx, 1.0, "failed to close clustering compaction writer during cleanup", mlog.Int("bufferID", buffer.id), mlog.Err(err))
|
||||
}
|
||||
}
|
||||
if t.mappingPool != nil {
|
||||
t.mappingPool.Release()
|
||||
}
|
||||
@@ -1033,7 +1172,7 @@ func (t *clusteringCompactionTask) GetStorageConfig() *indexpb.StorageConfig {
|
||||
// Includes TEXT column configs when lobContext requires REWRITE_ALL.
|
||||
func (t *clusteringCompactionTask) getWriterOpts() []storage.RwOption {
|
||||
opts := []storage.RwOption{
|
||||
storage.WithBufferSize(t.bufferSize),
|
||||
storage.WithBufferSize(max(int64(1), t.bufferSize/2)),
|
||||
storage.WithStorageConfig(t.compactionParams.StorageConfig),
|
||||
storage.WithUseLoonFFI(t.compactionParams.UseLoonFFI),
|
||||
storage.WithWriterFormat(t.compactionParams.GetStorageFormat()),
|
||||
|
||||
@@ -18,14 +18,20 @@ package compactor
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/apache/arrow/go/v17/arrow/array"
|
||||
"github.com/apache/arrow/go/v17/arrow/memory"
|
||||
"github.com/samber/lo"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/stretchr/testify/suite"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
@@ -39,6 +45,7 @@ import (
|
||||
"github.com/milvus-io/milvus/pkg/v3/objectstorage"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/conc"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
|
||||
@@ -228,3 +235,66 @@ func genInsertData(size int, seed int64, schema *schemapb.CollectionSchema) []*s
|
||||
}
|
||||
return []*storage.InsertData{buf}
|
||||
}
|
||||
|
||||
// Force the native writer to retain submitted Records so this test isolates
|
||||
// FlushChunk behavior. This fixture does not measure production peak memory.
|
||||
func TestClusteringFlushChunkRetainsNativeWriterBuffers(t *testing.T) {
|
||||
schema := clusteringWideSchema()
|
||||
task := &clusteringCompactionTask{memoryLimit: 1 << 20, flushPool: conc.NewPool[any](1)}
|
||||
defer task.flushPool.Release()
|
||||
alloc := memory.NewCheckedAllocator(memory.DefaultAllocator)
|
||||
for i := 0; i < 4; i++ {
|
||||
buffer, _ := newClusteringTestBuffer(t, schema, 64<<10, storage.StorageV2, storage.WithBufferSize(8<<20))
|
||||
buffer.id = i
|
||||
// Install tracking only while constructing the output builders. They
|
||||
// retain the allocator; input records and other writers are untracked.
|
||||
original := memory.DefaultAllocator
|
||||
memory.DefaultAllocator = alloc
|
||||
buffer.builder = storage.NewRecordBuilder(schema)
|
||||
memory.DefaultAllocator = original
|
||||
task.clusterBuffers = append(task.clusterBuffers, buffer)
|
||||
}
|
||||
for bucket, buffer := range task.clusterBuffers {
|
||||
record := clusteringTestRecord(t, schema, bucket*64, 64)
|
||||
for row := 0; row < record.Len(); row++ {
|
||||
require.NoError(t, buffer.WriteRecord(record, row))
|
||||
}
|
||||
record.Release()
|
||||
require.NoError(t, buffer.FlushChunk())
|
||||
require.Zero(t, buffer.builder.GetRowNum())
|
||||
}
|
||||
before := alloc.CurrentAlloc()
|
||||
require.Greater(t, before, int(task.memoryLimit), "tracked output buffers alone exceed the task budget")
|
||||
require.NoError(t, task.flushLargestBuffers(context.Background()))
|
||||
require.Equal(t, before, alloc.CurrentAlloc(), "V2 FlushChunk has not freed the submitted Arrow buffers")
|
||||
t.Logf("forced writer buffering: limit=%d retained Arrow bytes=%d buckets=%d builderBytes=%d writerBytes=%d",
|
||||
task.memoryLimit, before, len(task.clusterBuffers), task.clusterBuffers[0].writer.binLogMaxSize/2, 8<<20)
|
||||
for bucket, buffer := range task.clusterBuffers {
|
||||
require.NoError(t, buffer.Close())
|
||||
segments := buffer.GetCompactionSegments()
|
||||
require.Len(t, segments, 1)
|
||||
require.EqualValues(t, 64, segments[0].NumOfRows)
|
||||
reader, err := storage.NewBinlogRecordReader(context.Background(), segments[0].InsertLogs, schema,
|
||||
storage.WithVersion(storage.StorageV2), storage.WithStorageConfig(buffer.writer.params.StorageConfig))
|
||||
require.NoError(t, err)
|
||||
rowID := bucket * 64
|
||||
for {
|
||||
record, err := reader.Next()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
require.NoError(t, err)
|
||||
for row := 0; row < record.Len(); row++ {
|
||||
require.EqualValues(t, rowID, record.Column(100).(*array.Int64).Value(row))
|
||||
vector := record.Column(103).(*array.FixedSizeBinary).Value(row)
|
||||
for dim := 0; dim < 2048; dim++ {
|
||||
require.Equal(t, float32(rowID+dim), math.Float32frombits(binary.LittleEndian.Uint32(vector[dim*4:])))
|
||||
}
|
||||
rowID++
|
||||
}
|
||||
}
|
||||
reader.Close()
|
||||
require.Equal(t, (bucket+1)*64, rowID)
|
||||
}
|
||||
alloc.AssertSize(t, 0)
|
||||
}
|
||||
|
||||
@@ -19,20 +19,24 @@ package compactor
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"path"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/apache/arrow/go/v17/arrow/array"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/stretchr/testify/suite"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
||||
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
"github.com/milvus-io/milvus/internal/allocator"
|
||||
"github.com/milvus-io/milvus/internal/compaction"
|
||||
binlogio "github.com/milvus-io/milvus/internal/flushcommon/io"
|
||||
"github.com/milvus-io/milvus/internal/flushcommon/metacache"
|
||||
"github.com/milvus-io/milvus/internal/flushcommon/metacache/pkoracle"
|
||||
"github.com/milvus-io/milvus/internal/flushcommon/syncmgr"
|
||||
@@ -564,3 +568,92 @@ func genTextLOBInsertData(rows int, seed int64, schema *schemapb.CollectionSchem
|
||||
}
|
||||
return []*storage.InsertData{buf}
|
||||
}
|
||||
|
||||
func newClusteringTextTask(fixture *bumpFixture) *clusteringCompactionTask {
|
||||
plan := proto.Clone(fixture.task.plan).(*datapb.CompactionPlan)
|
||||
plan.Type = datapb.CompactionType_ClusteringCompaction
|
||||
plan.ClusteringKeyField = bumpFxPKField
|
||||
plan.PreferSegmentRows, plan.MaxSegmentRows = 128, 128
|
||||
plan.AnalyzeResultPath = fixture.cfg.RootPath + "/analyze_stats/999"
|
||||
return NewClusteringCompactionTask(context.Background(), binlogio.NewBinlogIO(fixture.task.chunkManager), plan, fixture.task.compactionParams)
|
||||
}
|
||||
|
||||
func TestClusteringTextRewriteAcrossPartitionNamespaces(t *testing.T) {
|
||||
setupBumpUTEnv(t)
|
||||
const textID = int64(105)
|
||||
const rows = 6
|
||||
fixture := buildBumpFixture(t, withRows(rows),
|
||||
withSourceFields(&schemapb.FieldSchema{FieldID: textID, Name: "text_lob", DataType: schemapb.DataType_Text}),
|
||||
withFillValue(func(i int, _ uint64, values map[int64]any) { values[textID] = bumpFxLobText(i) }),
|
||||
withTextLOBSource(textID), withLegacySourceNamespace())
|
||||
task := newClusteringTextTask(fixture)
|
||||
result, err := task.Compact()
|
||||
require.NoError(t, err)
|
||||
require.Len(t, result.Segments, 1)
|
||||
segment := result.Segments[0]
|
||||
require.EqualValues(t, rows, segment.NumOfRows)
|
||||
base, _, err := packed.UnmarshalManifestPath(segment.Manifest)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, storage.SegmentManifestBasePath(fixture.cfg.RootPath, CollectionID, PartitionID, segment.SegmentID), base)
|
||||
lobFiles, err := packed.GetManifestLobFiles(segment.Manifest, fixture.cfg)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, lobFiles)
|
||||
require.NotNil(t, task.lobContext)
|
||||
require.True(t, task.lobContext.DecodeTextFromSource)
|
||||
configs, err := task.lobContext.GetSourceTextColumnConfigs(segment.Manifest)
|
||||
require.NoError(t, err)
|
||||
reader, err := storage.NewTextDecodedManifestRecordReader(context.Background(), segment.Manifest, task.plan.Schema, configs,
|
||||
storage.WithVersion(storage.StorageV3), storage.WithStorageConfig(fixture.cfg))
|
||||
require.NoError(t, err)
|
||||
defer reader.Close()
|
||||
readRows := 0
|
||||
for {
|
||||
record, err := reader.Next()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
require.NoError(t, err)
|
||||
for row := 0; row < record.Len(); row++ {
|
||||
require.EqualValues(t, readRows, record.Column(bumpFxPKField).(*array.Int64).Value(row))
|
||||
require.Equal(t, bumpFxLobText(readRows), record.Column(textID).(*array.String).Value(row))
|
||||
readRows++
|
||||
}
|
||||
}
|
||||
require.Equal(t, rows, readRows)
|
||||
}
|
||||
|
||||
func TestClusteringTextAddedAfterAllSourceSegments(t *testing.T) {
|
||||
setupBumpUTEnv(t)
|
||||
const textID = int64(105)
|
||||
const rows = 6
|
||||
fixture := buildBumpFixture(t, withRows(rows), withLegacySourceNamespace(),
|
||||
withTargetAddedField(&schemapb.FieldSchema{FieldID: textID, Name: "added_text", DataType: schemapb.DataType_Text, Nullable: true}))
|
||||
task := newClusteringTextTask(fixture)
|
||||
task.compactionParams.BinLogMaxSize = 1 // Each source row produces a separate all-null TEXT batch.
|
||||
result, err := task.Compact()
|
||||
require.NoError(t, err)
|
||||
require.True(t, task.lobContext.DecodeTextFromSource)
|
||||
require.Len(t, result.Segments, 1)
|
||||
segment := result.Segments[0]
|
||||
require.EqualValues(t, rows, segment.NumOfRows)
|
||||
configs, err := task.lobContext.GetSourceTextColumnConfigs(segment.Manifest)
|
||||
require.NoError(t, err)
|
||||
reader, err := storage.NewTextDecodedManifestRecordReader(context.Background(), segment.Manifest, task.plan.Schema, configs,
|
||||
storage.WithVersion(storage.StorageV3), storage.WithStorageConfig(fixture.cfg))
|
||||
require.NoError(t, err)
|
||||
defer reader.Close()
|
||||
readRows := 0
|
||||
for {
|
||||
record, err := reader.Next()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, record.Len(), record.Column(textID).NullN())
|
||||
for row := 0; row < record.Len(); row++ {
|
||||
require.EqualValues(t, readRows, record.Column(bumpFxPKField).(*array.Int64).Value(row))
|
||||
readRows++
|
||||
}
|
||||
}
|
||||
require.Equal(t, rows, readRows)
|
||||
}
|
||||
|
||||
@@ -18,25 +18,43 @@ package compactor
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"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/bytedance/mockey"
|
||||
"github.com/cockroachdb/errors"
|
||||
"github.com/samber/lo"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/stretchr/testify/suite"
|
||||
"go.uber.org/atomic"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
||||
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
"github.com/milvus-io/milvus/internal/allocator"
|
||||
"github.com/milvus-io/milvus/internal/compaction"
|
||||
binlogio "github.com/milvus-io/milvus/internal/flushcommon/io"
|
||||
"github.com/milvus-io/milvus/internal/mocks/flushcommon/mock_util"
|
||||
"github.com/milvus-io/milvus/internal/storage"
|
||||
"github.com/milvus-io/milvus/internal/storagecommon"
|
||||
"github.com/milvus-io/milvus/internal/storagev2/packed"
|
||||
"github.com/milvus-io/milvus/internal/util/initcore"
|
||||
"github.com/milvus-io/milvus/pkg/v3/common"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/clusteringpb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/conc"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/merr"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
|
||||
@@ -702,3 +720,855 @@ func genRowWithBM25(magic int64) map[int64]interface{} {
|
||||
102: typeutil.CreateAndSortSparseFloatRow(map[uint32]float32{1: 1}),
|
||||
}
|
||||
}
|
||||
|
||||
// Observe the Record interface while retaining the real binlog writer below it.
|
||||
type clusteringRecordObserver struct {
|
||||
storage.BinlogRecordWriter
|
||||
rows []int
|
||||
values []*storage.Value
|
||||
captureValues bool
|
||||
writeError error
|
||||
blobs map[string][]byte
|
||||
}
|
||||
|
||||
func (w *clusteringRecordObserver) Write(r storage.Record) error {
|
||||
if w.writeError != nil {
|
||||
return w.writeError
|
||||
}
|
||||
if w.captureValues {
|
||||
values := make([]*storage.Value, r.Len())
|
||||
if err := storage.ValueDeserializerWithSchema(r, values, w.Schema(), true); err != nil {
|
||||
return err
|
||||
}
|
||||
w.values = append(w.values, values...)
|
||||
}
|
||||
if err := w.BinlogRecordWriter.Write(r); err != nil {
|
||||
return err
|
||||
}
|
||||
w.rows = append(w.rows, r.Len())
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *ClusteringCompactionTaskSuite) TestMappingAccumulatesRecords() {
|
||||
s.checkMappingRecords(false, -1)
|
||||
}
|
||||
|
||||
func (s *ClusteringCompactionTaskSuite) TestMappingVectorOffsetsAfterDelete() {
|
||||
s.checkMappingRecords(true, -1)
|
||||
}
|
||||
|
||||
func (s *ClusteringCompactionTaskSuite) TestMappingRecordTTL() {
|
||||
s.checkMappingRecords(false, 511)
|
||||
}
|
||||
|
||||
type clusteringOffsetReader struct {
|
||||
binlogio.BinlogIO
|
||||
offsets []byte
|
||||
}
|
||||
|
||||
func (r *clusteringOffsetReader) Download(ctx context.Context, paths []string) ([][]byte, error) {
|
||||
if len(paths) == 1 && paths[0] == "offsets" {
|
||||
return [][]byte{r.offsets}, nil
|
||||
}
|
||||
return r.BinlogIO.Download(ctx, paths)
|
||||
}
|
||||
|
||||
func (s *ClusteringCompactionTaskSuite) checkMappingRecords(vector bool, ttlCutoff int64) {
|
||||
s.preparScalarCompactionNormalTask()
|
||||
s.Require().NoError(s.task.init())
|
||||
defer s.task.cleanUp(context.Background())
|
||||
writer, err := NewMultiSegmentWriter(context.Background(), s.mockBinlogIO,
|
||||
NewCompactionAllocator(s.task.segIDAlloc, s.task.logIDAlloc),
|
||||
s.plan.MaxSize, s.plan.Schema, s.task.compactionParams, s.plan.MaxSegmentRows,
|
||||
PartitionID, CollectionID, "", 100, storage.WithStorageConfig(s.task.compactionParams.StorageConfig))
|
||||
s.Require().NoError(err)
|
||||
s.Require().NoError(writer.rotateWriter())
|
||||
observer := &clusteringRecordObserver{BinlogRecordWriter: writer.writer.BinlogRecordWriter, captureValues: true}
|
||||
writer.writer = storage.NewBinlogValueWriter(observer, 100)
|
||||
buffer := newClusterBuffer(0, writer, nil)
|
||||
s.task.clusterBuffers = []*ClusterBuffer{buffer}
|
||||
s.task.keyToBufferFunc = func(any) *ClusterBuffer { return buffer }
|
||||
s.task.memoryLimit = 1 << 30
|
||||
if ttlCutoff >= 0 {
|
||||
// Reuse the fixture's int64 PK column as expiration timestamps. Rows
|
||||
// through ttlCutoff expire, while the existing delete still filters 100.
|
||||
s.task.ttlFieldID = 100
|
||||
s.task.currentTime = time.UnixMicro(ttlCutoff)
|
||||
}
|
||||
if vector {
|
||||
ids := make([]uint32, 10240)
|
||||
for i := range ids {
|
||||
ids[i] = uint32(i % 7)
|
||||
}
|
||||
encoded, err := proto.Marshal(&clusteringpb.ClusteringCentroidIdMappingStats{CentroidIdMapping: ids})
|
||||
s.Require().NoError(err)
|
||||
s.task.binlogIO = &clusteringOffsetReader{BinlogIO: s.mockBinlogIO, offsets: encoded}
|
||||
s.task.isVectorClusteringKey = true
|
||||
s.task.segmentIDOffsetMapping = map[int64]string{s.plan.SegmentBinlogs[0].SegmentID: "offsets"}
|
||||
nextOffset := int64(0)
|
||||
s.task.offsetToBufferFunc = func(offset int64, mapping []uint32) *ClusterBuffer {
|
||||
if nextOffset == 100 {
|
||||
nextOffset++
|
||||
}
|
||||
s.Equal(nextOffset, offset)
|
||||
s.Equal(uint32(offset%7), mapping[offset])
|
||||
nextOffset++
|
||||
return buffer
|
||||
}
|
||||
}
|
||||
defer buffer.Close()
|
||||
|
||||
s.Require().NoError(s.task.mappingSegment(context.Background(), s.plan.SegmentBinlogs[0]))
|
||||
s.Empty(observer.rows, "input Record boundaries must not submit underfilled output batches")
|
||||
s.Require().NoError(buffer.Close())
|
||||
expectedRows := 10239
|
||||
if ttlCutoff >= 100 {
|
||||
expectedRows = 10240 - int(ttlCutoff) - 1
|
||||
}
|
||||
s.Equal([]int{expectedRows}, observer.rows)
|
||||
s.Require().Len(observer.values, expectedRows)
|
||||
for i, v := range observer.values {
|
||||
id := int64(i)
|
||||
if ttlCutoff >= 100 {
|
||||
id += ttlCutoff + 1
|
||||
}
|
||||
if ttlCutoff < 100 && id >= 100 {
|
||||
id++
|
||||
} // deleted by the fixture's deltalog
|
||||
s.Equal(id, v.PK.GetValue())
|
||||
s.Equal(genRow(id), v.Value)
|
||||
}
|
||||
}
|
||||
|
||||
// Use both a wide vector and variable-length nullable columns: row count alone
|
||||
// must not determine when a bucket submits a Record.
|
||||
func clusteringWideSchema() *schemapb.CollectionSchema {
|
||||
schema := genCollectionSchema()
|
||||
schema.Fields[4].Nullable = true
|
||||
schema.Fields[5].TypeParams = []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "2048"}}
|
||||
for i := int64(104); i < 112; i++ {
|
||||
schema.Fields = append(schema.Fields, &schemapb.FieldSchema{
|
||||
FieldID: i, Name: fmt.Sprint(i), DataType: schemapb.DataType_VarChar,
|
||||
Nullable: true, TypeParams: []*commonpb.KeyValuePair{{Key: common.MaxLengthKey, Value: "1024"}},
|
||||
})
|
||||
}
|
||||
return schema
|
||||
}
|
||||
|
||||
func clusteringTestRecord(t *testing.T, schema *schemapb.CollectionSchema, first, rows int) storage.Record {
|
||||
t.Helper()
|
||||
arrowSchema, err := storage.ConvertToArrowSchema(schema, false)
|
||||
require.NoError(t, err)
|
||||
builder := array.NewRecordBuilder(memory.DefaultAllocator, arrowSchema)
|
||||
defer builder.Release()
|
||||
field2Col := make(map[int64]int)
|
||||
for i, field := range schema.Fields {
|
||||
field2Col[field.FieldID] = i
|
||||
for row := first; row < first+rows; row++ {
|
||||
b := builder.Field(i)
|
||||
if field.Nullable && row%3 == 0 {
|
||||
b.AppendNull()
|
||||
continue
|
||||
}
|
||||
switch b := b.(type) {
|
||||
case *array.Int64Builder:
|
||||
b.Append(int64(row))
|
||||
case *array.Int32Builder:
|
||||
b.Append(int32(row))
|
||||
case *array.Float64Builder:
|
||||
b.Append(float64(row) + 0.25)
|
||||
case *array.StringBuilder:
|
||||
b.Append(strings.Repeat("x", 128+row%17) + fmt.Sprint(row))
|
||||
case *array.FixedSizeBinaryBuilder:
|
||||
vector := make([]byte, b.Type().(*arrow.FixedSizeBinaryType).ByteWidth)
|
||||
for dim := 0; dim < len(vector)/4; dim++ {
|
||||
binary.LittleEndian.PutUint32(vector[dim*4:], math.Float32bits(float32(row+dim)))
|
||||
}
|
||||
b.Append(vector)
|
||||
default:
|
||||
t.Fatalf("unsupported fixture field %v", field)
|
||||
}
|
||||
}
|
||||
}
|
||||
return storage.NewSimpleArrowRecord(builder.NewRecord(), field2Col)
|
||||
}
|
||||
|
||||
func newClusteringTestBuffer(t *testing.T, schema *schemapb.CollectionSchema, batchBytes uint64, version int64, overrides ...storage.RwOption) (*ClusterBuffer, *clusteringRecordObserver) {
|
||||
t.Helper()
|
||||
paramtable.Init()
|
||||
params := compaction.GenParams()
|
||||
params.BinLogMaxSize = batchBytes
|
||||
params.StorageVersion = version
|
||||
params.StorageConfig = &indexpb.StorageConfig{StorageType: "local", RootPath: t.TempDir()}
|
||||
binlogIO := mock_util.NewMockBinlogIO(t)
|
||||
blobs := make(map[string][]byte)
|
||||
binlogIO.EXPECT().Upload(mock.Anything, mock.Anything).RunAndReturn(func(_ context.Context, data map[string][]byte) error {
|
||||
for key, value := range data {
|
||||
blobs[key] = append([]byte(nil), value...)
|
||||
}
|
||||
return nil
|
||||
}).Maybe()
|
||||
task := &clusteringCompactionTask{bufferSize: int64(batchBytes), compactionParams: params}
|
||||
opts := append(task.getWriterOpts(), storage.WithColumnGroups([]storagecommon.ColumnGroup{
|
||||
{GroupID: 0, Fields: []int64{0, 1, 100, 101, 102}, Columns: []int{0, 1, 2, 3, 4}},
|
||||
{GroupID: 103, Fields: []int64{103}, Columns: []int{5}},
|
||||
{GroupID: 104, Fields: []int64{104, 105, 106, 107, 108, 109, 110, 111}, Columns: []int{6, 7, 8, 9, 10, 11, 12, 13}},
|
||||
}))
|
||||
opts = append(opts, overrides...)
|
||||
writer, err := NewMultiSegmentWriter(context.Background(), binlogIO,
|
||||
NewCompactionAllocator(allocator.NewLocalAllocator(1, 100), allocator.NewLocalAllocator(100, 10000)),
|
||||
1<<30, schema, params, 10000, PartitionID, CollectionID, "", 100, opts...)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, writer.rotateWriter())
|
||||
observer := &clusteringRecordObserver{BinlogRecordWriter: writer.writer.BinlogRecordWriter, blobs: blobs}
|
||||
writer.writer = storage.NewBinlogValueWriter(observer, 100)
|
||||
buffer := newClusterBuffer(0, writer, nil)
|
||||
t.Cleanup(func() { require.NoError(t, buffer.Close()) })
|
||||
return buffer, observer
|
||||
}
|
||||
|
||||
func TestClusteringMappingRejectsInvalidTimestamp(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
missing bool
|
||||
wrapped bool
|
||||
}{
|
||||
{name: "wrong type"},
|
||||
{name: "missing", missing: true},
|
||||
{name: "missing materialized", missing: true, wrapped: true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
schema := genCollectionSchema()
|
||||
malformed := proto.Clone(schema).(*schemapb.CollectionSchema)
|
||||
if tc.missing {
|
||||
malformed.Fields = append(malformed.Fields[:1], malformed.Fields[2:]...)
|
||||
} else {
|
||||
malformed.Fields[1].DataType = schemapb.DataType_VarChar
|
||||
}
|
||||
record := clusteringTestRecord(t, malformed, 0, 1)
|
||||
defer record.Release()
|
||||
var readerRecord storage.Record
|
||||
readerRecord = record
|
||||
if tc.wrapped {
|
||||
readerRecord = &materializedRecord{base: record}
|
||||
}
|
||||
existingFields := make(map[int64]struct{}, len(malformed.Fields))
|
||||
for _, field := range malformed.Fields {
|
||||
existingFields[field.FieldID] = struct{}{}
|
||||
}
|
||||
patch := mockey.Mock(newTextDecodedCompactionSegmentRecordReader).
|
||||
Return(&mockReader{records: []storage.Record{readerRecord}}, existingFields, nil).Build()
|
||||
defer patch.UnPatch()
|
||||
segment := &datapb.CompactionSegmentBinlogs{CollectionID: CollectionID, SegmentID: 1}
|
||||
task := &clusteringCompactionTask{
|
||||
plan: &datapb.CompactionPlan{Schema: schema, SegmentBinlogs: []*datapb.CompactionSegmentBinlogs{segment}},
|
||||
binlogIO: mock_util.NewMockBinlogIO(t),
|
||||
primaryKeyField: typeutil.GetField(schema, 100),
|
||||
compactionParams: compaction.GenParams(),
|
||||
}
|
||||
err := task.mappingSegment(context.Background(), segment)
|
||||
require.ErrorIs(t, err, merr.ErrServiceInternal)
|
||||
require.ErrorContains(t, err, "timestamp field is not an int64 column")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestClusterBufferV3BatchUsesWriterSchema(t *testing.T) {
|
||||
schema := clusteringWideSchema()
|
||||
schema.Fields[3].ExternalField = "source_field_101"
|
||||
writer := &MultiSegmentWriter{schema: schema, storageVersion: storage.StorageV3, binLogMaxSize: 64 << 20}
|
||||
buffer := newClusterBuffer(0, writer, nil)
|
||||
defer buffer.releaseBuilder()
|
||||
input := clusteringTestRecord(t, schema, 1, 1)
|
||||
require.NoError(t, buffer.WriteRecord(input, 0))
|
||||
input.Release()
|
||||
batch := buffer.builder.Build()
|
||||
defer batch.Release()
|
||||
wanted, err := storage.ConvertToArrowSchema(schema, true)
|
||||
require.NoError(t, err)
|
||||
require.True(t, batch.(interface{ ArrowSchema() *arrow.Schema }).ArrowSchema().Equal(wanted),
|
||||
"V3 batches should match the writer schema and reuse the Arrow record")
|
||||
}
|
||||
|
||||
func TestClusterBufferRecordBatches(t *testing.T) {
|
||||
for _, version := range []int64{storage.StorageV1, storage.StorageV2} {
|
||||
t.Run(fmt.Sprint(version), func(t *testing.T) {
|
||||
schema := clusteringWideSchema()
|
||||
// Half of the 64 KiB binlog budget is for allocated Arrow buffers.
|
||||
// The third wide row grows their capacity past 32 KiB.
|
||||
buffer, observer := newClusteringTestBuffer(t, schema, 64<<10, version)
|
||||
observer.captureValues = true
|
||||
for i := 0; i < 7; i++ {
|
||||
record := clusteringTestRecord(t, schema, i, 1)
|
||||
require.NoError(t, buffer.WriteRecord(record, 0))
|
||||
record.Release() // bucket must own copies across input lifetimes
|
||||
if i == 0 {
|
||||
require.Empty(t, observer.rows)
|
||||
}
|
||||
}
|
||||
require.Equal(t, []int{3, 3}, observer.rows)
|
||||
require.GreaterOrEqual(t, buffer.GetBufferSize(), uint64(8192))
|
||||
require.NoError(t, buffer.Close())
|
||||
require.Equal(t, []int{3, 3, 1}, observer.rows)
|
||||
require.Len(t, observer.values, 7)
|
||||
for i, value := range observer.values {
|
||||
require.Equal(t, int64(i), value.PK.GetValue())
|
||||
row := value.Value.(map[int64]interface{})
|
||||
vector := row[103].([]float32)
|
||||
require.Len(t, vector, 2048)
|
||||
for dim, v := range vector {
|
||||
require.Equal(t, float32(i+dim), v)
|
||||
}
|
||||
if i%3 == 0 {
|
||||
require.Nil(t, row[102])
|
||||
} else {
|
||||
require.Equal(t, strings.Repeat("x", 128+i%17)+fmt.Sprint(i), row[102])
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestClusterBufferCloseFlushesLazyWriter(t *testing.T) {
|
||||
schema := clusteringWideSchema()
|
||||
params := compaction.GenParams()
|
||||
params.StorageVersion = storage.StorageV1
|
||||
params.BinLogMaxSize = 64 << 20
|
||||
params.StorageConfig = &indexpb.StorageConfig{StorageType: "local", RootPath: t.TempDir()}
|
||||
binlogIO := mock_util.NewMockBinlogIO(t)
|
||||
binlogIO.EXPECT().Upload(mock.Anything, mock.Anything).Return(nil).Maybe()
|
||||
writer, err := NewMultiSegmentWriter(context.Background(), binlogIO,
|
||||
NewCompactionAllocator(allocator.NewLocalAllocator(1, 100), allocator.NewLocalAllocator(100, 10000)),
|
||||
1<<30, schema, params, 10000, PartitionID, CollectionID, "", 100,
|
||||
storage.WithStorageConfig(params.StorageConfig))
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, writer.writer)
|
||||
buffer := newClusterBuffer(0, writer, nil)
|
||||
input := clusteringTestRecord(t, schema, 1, 1)
|
||||
require.NoError(t, buffer.WriteRecord(input, 0))
|
||||
input.Release()
|
||||
require.Nil(t, writer.writer, "the writer should still be unopened before final close")
|
||||
require.NoError(t, buffer.Close())
|
||||
require.Len(t, writer.res, 1)
|
||||
require.EqualValues(t, 1, writer.res[0].NumOfRows)
|
||||
require.NoError(t, buffer.Close())
|
||||
}
|
||||
|
||||
func TestClusterBufferOversizedRowAndWriteError(t *testing.T) {
|
||||
buffer, observer := newClusteringTestBuffer(t, clusteringWideSchema(), 1024, storage.StorageV1)
|
||||
record := clusteringTestRecord(t, clusteringWideSchema(), 1, 2)
|
||||
defer record.Release()
|
||||
require.NoError(t, buffer.WriteRecord(record, 0))
|
||||
require.Equal(t, []int{1}, observer.rows)
|
||||
want := errors.New("write record failed")
|
||||
observer.writeError = want
|
||||
require.ErrorIs(t, buffer.WriteRecord(record, 1), want)
|
||||
observer.writeError = nil
|
||||
require.ErrorIs(t, buffer.WriteRecord(record, 0), want, "a failed write must poison the bucket")
|
||||
require.Equal(t, []int{1}, observer.rows, "later rows must not reach the writer")
|
||||
require.ErrorIs(t, buffer.FlushChunk(), want)
|
||||
require.ErrorIs(t, buffer.Close(), want)
|
||||
require.Nil(t, buffer.builder)
|
||||
}
|
||||
|
||||
func TestClusterBufferConcurrentRecords(t *testing.T) {
|
||||
buffer, observer := newClusteringTestBuffer(t, clusteringWideSchema(), 64<<10, storage.StorageV1)
|
||||
observer.captureValues = true
|
||||
var wg sync.WaitGroup
|
||||
for worker := 0; worker < 4; worker++ {
|
||||
record := clusteringTestRecord(t, clusteringWideSchema(), worker*10, 10)
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
defer record.Release()
|
||||
for row := 0; row < record.Len(); row++ {
|
||||
if err := buffer.WriteRecord(record, row); err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
require.NoError(t, buffer.Close())
|
||||
seen := make(map[int64]bool)
|
||||
for _, value := range observer.values {
|
||||
key := value.PK.GetValue().(int64)
|
||||
require.False(t, seen[key])
|
||||
seen[key] = true
|
||||
}
|
||||
require.Len(t, seen, 40)
|
||||
}
|
||||
|
||||
func TestClusteringScalarNullAndDefault(t *testing.T) {
|
||||
for _, dataType := range []schemapb.DataType{schemapb.DataType_Int64, schemapb.DataType_VarChar} {
|
||||
t.Run(dataType.String(), func(t *testing.T) {
|
||||
field := &schemapb.FieldSchema{DataType: dataType, Nullable: true}
|
||||
var builder array.Builder
|
||||
var expected any
|
||||
if dataType == schemapb.DataType_Int64 {
|
||||
b := array.NewInt64Builder(memory.DefaultAllocator)
|
||||
b.Append(42)
|
||||
builder, expected = b, int64(42)
|
||||
} else {
|
||||
b := array.NewStringBuilder(memory.DefaultAllocator)
|
||||
b.Append("key")
|
||||
builder, expected = b, "key"
|
||||
}
|
||||
defer builder.Release()
|
||||
builder.AppendNull()
|
||||
values := builder.NewArray()
|
||||
defer values.Release()
|
||||
value, err := clusteringScalarValue(values, field, 0)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, expected, value)
|
||||
value, err = clusteringScalarValue(values, field, 1)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, value)
|
||||
if dataType == schemapb.DataType_Int64 {
|
||||
field.DefaultValue = &schemapb.ValueField{Data: &schemapb.ValueField_LongData{LongData: 42}}
|
||||
} else {
|
||||
field.DefaultValue = &schemapb.ValueField{Data: &schemapb.ValueField_StringData{StringData: "key"}}
|
||||
}
|
||||
value, err = clusteringScalarValue(values, field, 1)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, expected, value)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestClusterBufferFlushChunkAndCleanup(t *testing.T) {
|
||||
buffer, observer := newClusteringTestBuffer(t, clusteringWideSchema(), 64<<20, storage.StorageV1)
|
||||
record := clusteringTestRecord(t, clusteringWideSchema(), 0, 1)
|
||||
defer record.Release()
|
||||
require.NoError(t, buffer.WriteRecord(record, 0))
|
||||
require.Empty(t, observer.rows)
|
||||
require.NoError(t, buffer.FlushChunk())
|
||||
require.Equal(t, []int{1}, observer.rows)
|
||||
require.Zero(t, buffer.GetBufferSize(), "V1 pressure flush must include the pending Builder")
|
||||
require.NoError(t, buffer.WriteRecord(record, 0))
|
||||
task := &clusteringCompactionTask{clusterBuffers: []*ClusterBuffer{buffer}}
|
||||
task.cleanUp(context.Background())
|
||||
require.Nil(t, buffer.builder, "cleanup discards an unfinished batch on task failure")
|
||||
require.Equal(t, []int{1}, observer.rows)
|
||||
}
|
||||
|
||||
func TestClusterBufferAccountsForNullColumns(t *testing.T) {
|
||||
schema := clusteringWideSchema()
|
||||
for id := int64(112); id < 352; id++ {
|
||||
schema.Fields = append(schema.Fields, &schemapb.FieldSchema{
|
||||
FieldID: id, Name: fmt.Sprint(id), DataType: schemapb.DataType_Double, Nullable: true,
|
||||
})
|
||||
}
|
||||
buffer, observer := newClusteringTestBuffer(t, schema, 2<<20, storage.StorageV1)
|
||||
record := clusteringTestRecord(t, schema, 0, 1)
|
||||
defer record.Release()
|
||||
checked := memory.NewCheckedAllocator(memory.DefaultAllocator)
|
||||
original := memory.DefaultAllocator
|
||||
memory.DefaultAllocator = checked
|
||||
buffer.builder = storage.NewRecordBuilder(schema)
|
||||
memory.DefaultAllocator = original
|
||||
defer func() {
|
||||
require.NoError(t, buffer.Close())
|
||||
checked.AssertSize(t, 0)
|
||||
}()
|
||||
for range 32 {
|
||||
require.NoError(t, buffer.WriteRecord(record, 0))
|
||||
}
|
||||
require.Empty(t, observer.rows)
|
||||
require.EqualValues(t, checked.CurrentAlloc(), buffer.GetBufferSize(), "pressure accounting must include all null column buffers")
|
||||
for range 33 {
|
||||
require.NoError(t, buffer.WriteRecord(record, 0))
|
||||
}
|
||||
require.NotEmpty(t, observer.rows, "allocated capacity must trigger a batch before payload bytes reach the threshold")
|
||||
}
|
||||
|
||||
func TestClusterBufferWideReadback(t *testing.T) {
|
||||
for _, version := range []int64{storage.StorageV1, storage.StorageV2} {
|
||||
t.Run(fmt.Sprint(version), func(t *testing.T) {
|
||||
const rows = 9
|
||||
schema := clusteringWideSchema()
|
||||
schema.Fields[4].DefaultValue = &schemapb.ValueField{Data: &schemapb.ValueField_StringData{StringData: "fallback"}}
|
||||
schema.Fields[6] = &schemapb.FieldSchema{
|
||||
FieldID: 104, Name: "double", DataType: schemapb.DataType_Double, Nullable: true,
|
||||
DefaultValue: &schemapb.ValueField{Data: &schemapb.ValueField_DoubleData{DoubleData: 3.5}},
|
||||
}
|
||||
const defaultTimestamp int64 = 1_700_000_000_000_000
|
||||
schema.Fields[7] = &schemapb.FieldSchema{
|
||||
FieldID: 105, Name: "timestamp", DataType: schemapb.DataType_Timestamptz, Nullable: true,
|
||||
DefaultValue: &schemapb.ValueField{Data: &schemapb.ValueField_TimestamptzData{TimestamptzData: defaultTimestamp}},
|
||||
}
|
||||
buffer, observer := newClusteringTestBuffer(t, schema, 16<<10, version)
|
||||
for row := 0; row < rows; row++ {
|
||||
record := clusteringTestRecord(t, schema, row, 1)
|
||||
err := buffer.WriteRecord(record, 0)
|
||||
record.Release()
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.NoError(t, buffer.Close())
|
||||
segments := buffer.GetCompactionSegments()
|
||||
require.Len(t, segments, 1)
|
||||
segment := segments[0]
|
||||
require.EqualValues(t, rows, segment.NumOfRows)
|
||||
reader, err := storage.NewBinlogRecordReader(context.Background(), segment.InsertLogs, schema,
|
||||
storage.WithVersion(version), storage.WithStorageConfig(buffer.writer.params.StorageConfig),
|
||||
storage.WithDownloader(func(_ context.Context, paths []string) ([][]byte, error) {
|
||||
data := make([][]byte, len(paths))
|
||||
for i, key := range paths {
|
||||
require.Contains(t, observer.blobs, key)
|
||||
data[i] = observer.blobs[key]
|
||||
}
|
||||
return data, nil
|
||||
}))
|
||||
require.NoError(t, err)
|
||||
defer reader.Close()
|
||||
rowID := 0
|
||||
for {
|
||||
record, err := reader.Next()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
require.NoError(t, err)
|
||||
for row := 0; row < record.Len(); row++ {
|
||||
for _, field := range schema.Fields {
|
||||
column := record.Column(field.FieldID)
|
||||
null := field.Nullable && rowID%3 == 0 && field.DefaultValue == nil
|
||||
require.Equal(t, null, column.IsNull(row), "field %d row %d", field.FieldID, rowID)
|
||||
if null {
|
||||
continue
|
||||
}
|
||||
switch column := column.(type) {
|
||||
case *array.Int64:
|
||||
want := int64(rowID)
|
||||
if field.GetDataType() == schemapb.DataType_Timestamptz && rowID%3 == 0 {
|
||||
want = defaultTimestamp
|
||||
}
|
||||
require.Equal(t, want, column.Value(row))
|
||||
case *array.Int32:
|
||||
require.EqualValues(t, rowID, column.Value(row))
|
||||
case *array.Float64:
|
||||
want := float64(rowID) + 0.25
|
||||
if rowID%3 == 0 {
|
||||
want = 3.5
|
||||
}
|
||||
require.Equal(t, want, column.Value(row))
|
||||
case *array.String:
|
||||
want := strings.Repeat("x", 128+rowID%17) + fmt.Sprint(rowID)
|
||||
if rowID%3 == 0 {
|
||||
want = "fallback"
|
||||
}
|
||||
require.Equal(t, want, column.Value(row))
|
||||
case *array.FixedSizeBinary:
|
||||
for dim := 0; dim < 2048; dim++ {
|
||||
require.Equal(t, float32(rowID+dim), math.Float32frombits(binary.LittleEndian.Uint32(column.Value(row)[dim*4:])))
|
||||
}
|
||||
default:
|
||||
t.Fatalf("unexpected column %T", column)
|
||||
}
|
||||
}
|
||||
rowID++
|
||||
}
|
||||
}
|
||||
require.Equal(t, rows, rowID)
|
||||
for _, group := range segment.InsertLogs {
|
||||
var entries int64
|
||||
for _, log := range group.Binlogs {
|
||||
require.EqualValues(t, entries, log.TimestampFrom)
|
||||
entries += log.EntriesNum
|
||||
require.EqualValues(t, entries-1, log.TimestampTo)
|
||||
}
|
||||
require.EqualValues(t, rows, entries)
|
||||
}
|
||||
if version == storage.StorageV2 {
|
||||
for _, field := range schema.Fields {
|
||||
want := int64(0)
|
||||
if field.Nullable && field.DefaultValue == nil {
|
||||
want = 3
|
||||
}
|
||||
require.Contains(t, segment.GetStats().GetNullCounts(), field.FieldID)
|
||||
require.Equal(t, want, segment.GetStats().GetNullCounts()[field.FieldID], "field %d", field.FieldID)
|
||||
}
|
||||
}
|
||||
require.NotEmpty(t, segment.Field2StatslogPaths)
|
||||
for _, field := range segment.Field2StatslogPaths {
|
||||
for _, log := range field.Binlogs {
|
||||
stats, err := storage.DeserializeStats([]*storage.Blob{{Value: observer.blobs[log.LogPath]}})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, stats, 1)
|
||||
require.EqualValues(t, 0, stats[0].MinPk.GetValue())
|
||||
require.EqualValues(t, rows-1, stats[0].MaxPk.GetValue())
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// The second Next releases the input and pauses while the bucket still owns an
|
||||
// underfilled batch. This exercises Compact's cleanup after real mapping work.
|
||||
type pausedClusteringReader struct {
|
||||
ctx context.Context
|
||||
record storage.Record
|
||||
read bool
|
||||
ready chan<- struct{}
|
||||
resume <-chan struct{}
|
||||
closed chan struct{}
|
||||
err error
|
||||
}
|
||||
|
||||
func (r *pausedClusteringReader) Next() (storage.Record, error) {
|
||||
if !r.read {
|
||||
r.read = true
|
||||
return r.record, nil
|
||||
}
|
||||
r.record.Release()
|
||||
r.record = nil
|
||||
r.ready <- struct{}{}
|
||||
select {
|
||||
case <-r.ctx.Done():
|
||||
return nil, r.ctx.Err()
|
||||
case <-r.resume:
|
||||
return nil, r.err
|
||||
}
|
||||
}
|
||||
|
||||
func (r *pausedClusteringReader) Close() error {
|
||||
if r.record != nil {
|
||||
r.record.Release()
|
||||
r.record = nil
|
||||
}
|
||||
close(r.closed)
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestClusteringReleasesPendingBuildersOnFailure(t *testing.T) {
|
||||
for _, cancelTask := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("cancel=%v", cancelTask), func(t *testing.T) {
|
||||
schema := clusteringWideSchema()
|
||||
buffer, observer := newClusteringTestBuffer(t, schema, 64<<20, storage.StorageV1)
|
||||
workers := ¶mtable.Get().DataNodeCfg.ClusteringCompactionWorkerPoolSize
|
||||
previousWorkers := workers.GetValue()
|
||||
require.NoError(t, paramtable.Get().Save(workers.Key, "2"))
|
||||
t.Cleanup(func() { require.NoError(t, paramtable.Get().Save(workers.Key, previousWorkers)) })
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
checked := memory.NewCheckedAllocator(memory.DefaultAllocator)
|
||||
originalAllocator := memory.DefaultAllocator
|
||||
memory.DefaultAllocator = checked
|
||||
defer func() { memory.DefaultAllocator = originalAllocator }()
|
||||
ready := make(chan struct{}, 2)
|
||||
gates := []chan struct{}{make(chan struct{}), make(chan struct{})}
|
||||
readFailure := errors.New("injected read failure after a buffered record")
|
||||
readers := make([]*pausedClusteringReader, 2)
|
||||
for i := range readers {
|
||||
readers[i] = &pausedClusteringReader{
|
||||
ctx: ctx, record: clusteringTestRecord(t, schema, i, 1),
|
||||
ready: ready, resume: gates[i], closed: make(chan struct{}), err: readFailure,
|
||||
}
|
||||
}
|
||||
plan := &datapb.CompactionPlan{
|
||||
Type: datapb.CompactionType_ClusteringCompaction, Schema: schema, ClusteringKeyField: 100,
|
||||
SegmentBinlogs: []*datapb.CompactionSegmentBinlogs{{CollectionID: CollectionID, SegmentID: 0}, {CollectionID: CollectionID, SegmentID: 1}},
|
||||
}
|
||||
task := NewClusteringCompactionTask(ctx, buffer.writer.binlogIO, plan, buffer.writer.params)
|
||||
analyze := mockey.Mock((*clusteringCompactionTask).getScalarAnalyzeResult).To(func(task *clusteringCompactionTask, _ context.Context) error {
|
||||
task.clusterBuffers = []*ClusterBuffer{buffer}
|
||||
task.keyToBufferFunc = func(any) *ClusterBuffer { return buffer }
|
||||
return nil
|
||||
}).Build()
|
||||
defer analyze.UnPatch()
|
||||
readerFactory := mockey.Mock(newTextDecodedCompactionSegmentRecordReader).To(func(_ context.Context, segment *datapb.CompactionSegmentBinlogs,
|
||||
_ *schemapb.CollectionSchema, _ *indexpb.StorageConfig, _ []packed.TextColumnConfig, _ ...storage.RwOption,
|
||||
) (storage.RecordReader, map[int64]struct{}, error) {
|
||||
fields := make(map[int64]struct{})
|
||||
for _, field := range schema.Fields {
|
||||
fields[field.FieldID] = struct{}{}
|
||||
}
|
||||
return readers[segment.SegmentID], fields, nil
|
||||
}).Build()
|
||||
defer readerFactory.UnPatch()
|
||||
finished := make(chan error, 1)
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
_, err := task.Compact()
|
||||
finished <- err
|
||||
}()
|
||||
defer func() {
|
||||
cancel()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Error("mapping workers did not stop during test cleanup")
|
||||
}
|
||||
}()
|
||||
for range readers {
|
||||
select {
|
||||
case <-ready:
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("mapping did not reach the pending batch")
|
||||
}
|
||||
}
|
||||
require.Greater(t, checked.CurrentAlloc(), 0)
|
||||
require.Equal(t, 2, buffer.builder.GetRowNum())
|
||||
require.Empty(t, observer.rows)
|
||||
if cancelTask {
|
||||
cancel()
|
||||
} else {
|
||||
close(gates[0])
|
||||
select {
|
||||
case <-readers[0].closed:
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("failed reader did not close")
|
||||
}
|
||||
select {
|
||||
case err := <-finished:
|
||||
t.Fatalf("Compact returned while another mapping worker still owned a pending batch: %v", err)
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
close(gates[1])
|
||||
}
|
||||
select {
|
||||
case err := <-finished:
|
||||
if cancelTask {
|
||||
require.ErrorIs(t, err, context.Canceled)
|
||||
} else {
|
||||
require.ErrorIs(t, err, readFailure)
|
||||
}
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("Compact did not finish after readers stopped")
|
||||
}
|
||||
require.Nil(t, buffer.builder)
|
||||
checked.AssertSize(t, 0)
|
||||
require.Empty(t, observer.rows)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type clusteringFailingWriter struct {
|
||||
storage.BinlogRecordWriter
|
||||
writeErr error
|
||||
closeErr error
|
||||
writes int
|
||||
closes int
|
||||
}
|
||||
|
||||
func (w *clusteringFailingWriter) GetWrittenUncompressed() uint64 { return 0 }
|
||||
func (w *clusteringFailingWriter) GetBufferUncompressed() uint64 { return 0 }
|
||||
func (w *clusteringFailingWriter) FlushChunk() error { return nil }
|
||||
func (w *clusteringFailingWriter) Write(storage.Record) error {
|
||||
w.writes++
|
||||
return w.writeErr
|
||||
}
|
||||
|
||||
func (w *clusteringFailingWriter) Close() error {
|
||||
w.closes++
|
||||
return w.closeErr
|
||||
}
|
||||
|
||||
func newClusteringFailureBuffer(t *testing.T, schema *schemapb.CollectionSchema, writer *clusteringFailingWriter) *ClusterBuffer {
|
||||
t.Helper()
|
||||
buffer := newClusterBuffer(0, &MultiSegmentWriter{
|
||||
schema: schema, binLogMaxSize: 64 << 20, segmentSize: 1 << 30,
|
||||
allocator: &compactionAlloactor{}, writer: storage.NewBinlogValueWriter(writer, 100),
|
||||
}, nil)
|
||||
t.Cleanup(buffer.releaseBuilder)
|
||||
return buffer
|
||||
}
|
||||
|
||||
func TestClusterBufferRejectsWritesAfterPartialAppend(t *testing.T) {
|
||||
schema := clusteringWideSchema()
|
||||
writer := &clusteringFailingWriter{closeErr: errors.New("injected close error")}
|
||||
buffer := newClusteringFailureBuffer(t, schema, writer)
|
||||
good := clusteringTestRecord(t, schema, 1, 1)
|
||||
defer good.Release()
|
||||
require.NoError(t, buffer.WriteRecord(good, 0))
|
||||
badSchema := proto.Clone(schema).(*schemapb.CollectionSchema)
|
||||
badSchema.Fields[len(badSchema.Fields)-1].DataType = schemapb.DataType_Int64
|
||||
bad := clusteringTestRecord(t, badSchema, 2, 1)
|
||||
defer bad.Release()
|
||||
appendErr := buffer.WriteRecord(bad, 0)
|
||||
require.ErrorContains(t, appendErr, "failed to append value")
|
||||
|
||||
// Simulate workers that were already mapping when another worker failed.
|
||||
var wg sync.WaitGroup
|
||||
for range 4 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
assert.ErrorIs(t, buffer.WriteRecord(good, 0), appendErr)
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
require.ErrorIs(t, buffer.FlushChunk(), appendErr)
|
||||
require.Nil(t, buffer.builder, "the partially appended batch must be discarded")
|
||||
require.Zero(t, writer.writes, "no partial or later batch may reach the writer")
|
||||
}
|
||||
|
||||
func TestClusteringMappingCancelsWorkersOnFailure(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
task := &clusteringCompactionTask{
|
||||
plan: &datapb.CompactionPlan{SegmentBinlogs: []*datapb.CompactionSegmentBinlogs{{SegmentID: 0}, {SegmentID: 1}}},
|
||||
mappingPool: conc.NewPool[any](2),
|
||||
}
|
||||
defer task.mappingPool.Release()
|
||||
started := make(chan struct{})
|
||||
exited := make(chan struct{})
|
||||
want := errors.New("injected append failure")
|
||||
patch := mockey.Mock((*clusteringCompactionTask).mappingSegment).To(func(_ *clusteringCompactionTask, ctx context.Context, segment *datapb.CompactionSegmentBinlogs) error {
|
||||
if segment.SegmentID == 0 {
|
||||
close(started)
|
||||
<-ctx.Done()
|
||||
close(exited)
|
||||
return ctx.Err()
|
||||
}
|
||||
<-started
|
||||
return want
|
||||
}).Build()
|
||||
defer patch.UnPatch()
|
||||
segments, stats, err := task.mapping(ctx)
|
||||
require.ErrorIs(t, err, want, "sibling cancellation must not mask the original failure")
|
||||
require.Nil(t, segments)
|
||||
require.Nil(t, stats)
|
||||
select {
|
||||
case <-exited:
|
||||
default:
|
||||
t.Fatal("mapping returned before the canceled worker exited")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClusterBufferClosesWriterAfterFlushFailure(t *testing.T) {
|
||||
schema := clusteringWideSchema()
|
||||
writeErr := errors.New("injected write error")
|
||||
closeErr := errors.New("injected close error")
|
||||
writer := &clusteringFailingWriter{writeErr: writeErr, closeErr: closeErr}
|
||||
buffer := newClusteringFailureBuffer(t, schema, writer)
|
||||
record := clusteringTestRecord(t, schema, 1, 1)
|
||||
defer record.Release()
|
||||
require.NoError(t, buffer.WriteRecord(record, 0))
|
||||
err := buffer.Close()
|
||||
require.ErrorIs(t, err, writeErr)
|
||||
require.ErrorIs(t, err, closeErr)
|
||||
require.Equal(t, 1, writer.closes)
|
||||
require.Nil(t, buffer.builder)
|
||||
require.Nil(t, buffer.writer.writer)
|
||||
require.NoError(t, buffer.Close())
|
||||
require.Equal(t, 1, writer.closes)
|
||||
}
|
||||
|
||||
func TestClusteringCleanupClosesWriterWithoutFlushingPendingRows(t *testing.T) {
|
||||
schema := clusteringWideSchema()
|
||||
writer := &clusteringFailingWriter{closeErr: errors.New("injected close error")}
|
||||
buffer := newClusteringFailureBuffer(t, schema, writer)
|
||||
record := clusteringTestRecord(t, schema, 1, 1)
|
||||
defer record.Release()
|
||||
require.NoError(t, buffer.WriteRecord(record, 0))
|
||||
task := &clusteringCompactionTask{clusterBuffers: []*ClusterBuffer{buffer}}
|
||||
task.cleanUp(context.Background())
|
||||
require.Zero(t, writer.writes)
|
||||
require.Equal(t, 1, writer.closes)
|
||||
require.Nil(t, buffer.builder)
|
||||
require.Nil(t, buffer.writer.writer)
|
||||
task.cleanUp(context.Background())
|
||||
require.Equal(t, 1, writer.closes)
|
||||
}
|
||||
|
||||
@@ -30,6 +30,7 @@ import (
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/suite"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
||||
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
@@ -319,10 +320,9 @@ func (s *MixCompactionTaskStorageV1Suite) setupTestWithTextField() {
|
||||
s.task = NewMixCompactionTask(context.Background(), s.mockBinlogIO, nil, plan, compaction.GenParams(), []int64{pk.FieldID})
|
||||
}
|
||||
|
||||
// V1 binlogs deserialize TEXT as arrow String while the retained-row rebuilder
|
||||
// allocates a Binary builder for TEXT, so with any filtered row rb.Append fails
|
||||
// deterministically: the task must fail instead of silently dropping the
|
||||
// retained rows of that batch.
|
||||
// Feed an Int64 column where the plan declares TEXT so rb.Append fails on a
|
||||
// real type mismatch. UTF8 TEXT is valid input for RecordBuilder. The task must
|
||||
// propagate the append error instead of silently dropping retained rows.
|
||||
func TestMixCompactionPropagatesRecordBuilderAppendError(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
@@ -341,11 +341,13 @@ func TestMixCompactionPropagatesRecordBuilderAppendError(t *testing.T) {
|
||||
|
||||
segmentID := int64(700)
|
||||
alloc := allocator.NewLocalAllocator(888888, math.MaxInt64)
|
||||
segWriter, err := NewSegmentWriter(s.meta.GetSchema(), 65535, compactionBatchSize, segmentID, PartitionID, CollectionID, []int64{})
|
||||
sourceSchema := proto.Clone(s.meta.GetSchema()).(*schemapb.CollectionSchema)
|
||||
typeutil.GetField(sourceSchema, textFieldForAppendErrorTest).DataType = schemapb.DataType_Int64
|
||||
segWriter, err := NewSegmentWriter(sourceSchema, 65535, compactionBatchSize, segmentID, PartitionID, CollectionID, []int64{})
|
||||
s.Require().NoError(err)
|
||||
for i := int64(0); i < 2; i++ {
|
||||
row := getRow(segmentID+i, 0)
|
||||
row[textFieldForAppendErrorTest] = "text-payload"
|
||||
row[textFieldForAppendErrorTest] = int64(42)
|
||||
err = segWriter.Write(&storage.Value{
|
||||
PK: storage.NewInt64PrimaryKey(segmentID + i),
|
||||
Timestamp: int64(tsoutil.ComposeTSByTime(getMilvusBirthday())),
|
||||
|
||||
@@ -185,6 +185,13 @@ func (r *materializedRecord) Column(fieldID storage.FieldID) arrow.Array {
|
||||
return r.base.Column(fieldID)
|
||||
}
|
||||
|
||||
func (r *materializedRecord) TryColumn(fieldID storage.FieldID) (arrow.Array, bool) {
|
||||
if col, ok := r.computed[fieldID]; ok {
|
||||
return col, true
|
||||
}
|
||||
return storage.TryRecordColumn(r.base, fieldID)
|
||||
}
|
||||
|
||||
func (r *materializedRecord) Len() int {
|
||||
return r.base.Len()
|
||||
}
|
||||
|
||||
@@ -27,6 +27,13 @@ func (r *timestampOverwriteRecord) Column(i storage.FieldID) arrow.Array {
|
||||
return r.inner.Column(i)
|
||||
}
|
||||
|
||||
func (r *timestampOverwriteRecord) TryColumn(i storage.FieldID) (arrow.Array, bool) {
|
||||
if i == common.TimeStampField {
|
||||
return r.tsArray, true
|
||||
}
|
||||
return storage.TryRecordColumn(r.inner, i)
|
||||
}
|
||||
|
||||
func (r *timestampOverwriteRecord) Len() int { return r.inner.Len() }
|
||||
func (r *timestampOverwriteRecord) Release() { r.tsArray.Release(); r.inner.Release() }
|
||||
func (r *timestampOverwriteRecord) Retain() { r.tsArray.Retain(); r.inner.Retain() }
|
||||
|
||||
@@ -71,6 +71,9 @@ func TestOverwriteRecordTimestamps_NonZero(t *testing.T) {
|
||||
assert.Equal(t, 3, out.Len())
|
||||
|
||||
tsCol := out.Column(common.TimeStampField).(*array.Int64)
|
||||
safeColumn, ok := storage.TryRecordColumn(out, common.TimeStampField)
|
||||
assert.True(t, ok)
|
||||
assert.Same(t, tsCol, safeColumn)
|
||||
for i := 0; i < 3; i++ {
|
||||
assert.Equal(t, int64(5000), tsCol.Value(i))
|
||||
}
|
||||
|
||||
@@ -18,6 +18,7 @@ package storage
|
||||
|
||||
import (
|
||||
"strconv"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/apache/arrow/go/v17/arrow"
|
||||
"github.com/apache/arrow/go/v17/arrow/array"
|
||||
@@ -152,7 +153,11 @@ func appendValueAt(builder array.Builder, a arrow.Array, idx int, field *schemap
|
||||
}
|
||||
if ia.IsNull(idx) {
|
||||
if defaultValue != nil {
|
||||
b.Append(defaultValue.GetLongData())
|
||||
if field.GetDataType() == schemapb.DataType_Timestamptz {
|
||||
b.Append(defaultValue.GetTimestamptzData())
|
||||
} else {
|
||||
b.Append(defaultValue.GetLongData())
|
||||
}
|
||||
return 8, nil
|
||||
}
|
||||
b.AppendNull()
|
||||
@@ -192,6 +197,10 @@ func appendValueAt(builder array.Builder, a arrow.Array, idx int, field *schemap
|
||||
return 0, merr.WrapErrServiceInternalMsg("invalid value type %T, expect %T", a.DataType(), builder.Type())
|
||||
}
|
||||
if fa.IsNull(idx) {
|
||||
if defaultValue != nil {
|
||||
b.Append(defaultValue.GetDoubleData())
|
||||
return 8, nil
|
||||
}
|
||||
b.AppendNull()
|
||||
return 0, nil
|
||||
} else {
|
||||
@@ -440,11 +449,37 @@ type RecordBuilder struct {
|
||||
arrowFields []arrow.Field
|
||||
builders []array.Builder
|
||||
defaults []appendValueDefault
|
||||
allocator *recordBuilderAllocator
|
||||
|
||||
nRows int
|
||||
size uint64
|
||||
}
|
||||
|
||||
// Each output batch has its own allocator counter. Records handed to a writer
|
||||
// can outlive the builder and can be released from another goroutine.
|
||||
type recordBuilderAllocator struct {
|
||||
memory.Allocator
|
||||
size atomic.Int64
|
||||
}
|
||||
|
||||
func (a *recordBuilderAllocator) Allocate(size int) []byte {
|
||||
buf := a.Allocator.Allocate(size)
|
||||
a.size.Add(int64(len(buf)))
|
||||
return buf
|
||||
}
|
||||
|
||||
func (a *recordBuilderAllocator) Reallocate(size int, buf []byte) []byte {
|
||||
oldSize := len(buf)
|
||||
buf = a.Allocator.Reallocate(size, buf)
|
||||
a.size.Add(int64(len(buf) - oldSize))
|
||||
return buf
|
||||
}
|
||||
|
||||
func (a *recordBuilderAllocator) Free(buf []byte) {
|
||||
a.Allocator.Free(buf)
|
||||
a.size.Add(-int64(len(buf)))
|
||||
}
|
||||
|
||||
func (b *RecordBuilder) prepareAppendDefaults() error {
|
||||
if b.defaults != nil {
|
||||
return nil
|
||||
@@ -469,6 +504,23 @@ func (b *RecordBuilder) Append(rec Record, start, end int) error {
|
||||
for i, builder := range b.builders {
|
||||
f := b.fields[i]
|
||||
col := rec.Column(f.FieldID)
|
||||
// TEXT may be decoded UTF8 or a binary LOB reference. Preserve the
|
||||
// reader's representation instead of treating decoded text as a ref.
|
||||
if f.GetDataType() == schemapb.DataType_Text &&
|
||||
(col.DataType().ID() == arrow.STRING || col.DataType().ID() == arrow.BINARY) &&
|
||||
col.DataType().ID() != builder.Type().ID() {
|
||||
if builder.NullN() == builder.Len() {
|
||||
// Missing TEXT fields can have been filled with binary nulls.
|
||||
nulls := builder.Len()
|
||||
builder.Release()
|
||||
builder = array.NewBuilder(b.allocator, col.DataType())
|
||||
builder.AppendNulls(nulls)
|
||||
b.builders[i] = builder
|
||||
} else if col.IsNull(offset) {
|
||||
builder.AppendNull()
|
||||
continue
|
||||
}
|
||||
}
|
||||
size, err := appendValueAt(builder, col, offset, f, b.defaults[i])
|
||||
if err != nil {
|
||||
return merr.Wrapf(err, "failed to append value at offset %d for field %s", offset, f.GetName())
|
||||
@@ -488,6 +540,13 @@ func (b *RecordBuilder) GetSize() uint64 {
|
||||
return b.size
|
||||
}
|
||||
|
||||
func (b *RecordBuilder) GetMemorySize() uint64 {
|
||||
// Count allocated Arrow buffers, including null slots, validity/offset
|
||||
// buffers and spare capacity. Payload bytes alone can severely undercount
|
||||
// wide nullable schemas. This excludes batches already handed to writers.
|
||||
return uint64(b.allocator.size.Load())
|
||||
}
|
||||
|
||||
func (b *RecordBuilder) Release() {
|
||||
for _, builder := range b.builders {
|
||||
builder.Release()
|
||||
@@ -515,10 +574,30 @@ func (b *RecordBuilder) Build() Record {
|
||||
}
|
||||
b.nRows = 0
|
||||
b.size = 0
|
||||
// NewArray transfers the buffers but keeps the allocator on each builder.
|
||||
// Start fresh builders/counters so retained output Records cannot affect
|
||||
// the next batch's accounting when their buffers are eventually released.
|
||||
b.allocator = &recordBuilderAllocator{Allocator: b.allocator.Allocator}
|
||||
for i, builder := range b.builders {
|
||||
dataType := builder.Type()
|
||||
builder.Release()
|
||||
b.builders[i] = array.NewBuilder(b.allocator, dataType)
|
||||
}
|
||||
return rec
|
||||
}
|
||||
|
||||
func NewRecordBuilder(schema *schemapb.CollectionSchema) *RecordBuilder {
|
||||
return newRecordBuilder(schema, false)
|
||||
}
|
||||
|
||||
// NewRecordBuilderWithFieldID uses the physical field names expected by V3
|
||||
// writers, including ExternalField overrides, so matching batches can be
|
||||
// written without rebuilding the Arrow record.
|
||||
func NewRecordBuilderWithFieldID(schema *schemapb.CollectionSchema) *RecordBuilder {
|
||||
return newRecordBuilder(schema, true)
|
||||
}
|
||||
|
||||
func newRecordBuilder(schema *schemapb.CollectionSchema, useFieldID bool) *RecordBuilder {
|
||||
// assumes 5 sub fields per StructArrayField
|
||||
fields := make([]*schemapb.FieldSchema, 0, len(schema.Fields)+len(schema.StructArrayFields)*5)
|
||||
fields = append(fields, schema.Fields...)
|
||||
@@ -528,6 +607,7 @@ func NewRecordBuilder(schema *schemapb.CollectionSchema) *RecordBuilder {
|
||||
|
||||
builders := make([]array.Builder, len(fields))
|
||||
arrowFields := make([]arrow.Field, len(fields))
|
||||
allocator := &recordBuilderAllocator{Allocator: memory.DefaultAllocator}
|
||||
for i, field := range fields {
|
||||
dim, _ := typeutil.GetDim(field)
|
||||
|
||||
@@ -536,26 +616,34 @@ func NewRecordBuilder(schema *schemapb.CollectionSchema) *RecordBuilder {
|
||||
elementType = field.GetElementType()
|
||||
}
|
||||
if field.GetNullable() && isNullableDenseVectorArrowType(field.DataType) {
|
||||
builders[i] = array.NewBinaryBuilder(memory.DefaultAllocator, arrow.BinaryTypes.Binary)
|
||||
builders[i] = array.NewBinaryBuilder(allocator, arrow.BinaryTypes.Binary)
|
||||
} else if field.DataType == schemapb.DataType_Text {
|
||||
// TEXT fields are stored as binary (LOB references) in manifest storage,
|
||||
// so the builder must use binary type to match what the reader returns.
|
||||
builders[i] = array.NewBinaryBuilder(memory.DefaultAllocator, arrow.BinaryTypes.Binary)
|
||||
builders[i] = array.NewBinaryBuilder(allocator, arrow.BinaryTypes.Binary)
|
||||
} else {
|
||||
arrowType := serdeMap[field.DataType].arrowType(int(dim), elementType, field.GetElementNullable())
|
||||
builders[i] = array.NewBuilder(memory.DefaultAllocator, arrowType)
|
||||
builders[i] = array.NewBuilder(allocator, arrowType)
|
||||
}
|
||||
arrowFields[i] = newRecordBuilderArrowField(field, builders[i].Type(), dim, elementType)
|
||||
arrowFields[i] = newRecordBuilderArrowField(field, builders[i].Type(), dim, elementType, useFieldID)
|
||||
}
|
||||
|
||||
return &RecordBuilder{
|
||||
fields: fields,
|
||||
arrowFields: arrowFields,
|
||||
builders: builders,
|
||||
allocator: allocator,
|
||||
}
|
||||
}
|
||||
|
||||
func newRecordBuilderArrowField(field *schemapb.FieldSchema, arrowType arrow.DataType, dim int64, elementType schemapb.DataType) arrow.Field {
|
||||
func newRecordBuilderArrowField(field *schemapb.FieldSchema, arrowType arrow.DataType, dim int64, elementType schemapb.DataType, useFieldID bool) arrow.Field {
|
||||
name := field.GetName()
|
||||
if useFieldID {
|
||||
name = strconv.FormatInt(field.GetFieldID(), 10)
|
||||
if field.GetExternalField() != "" {
|
||||
name = field.GetExternalField()
|
||||
}
|
||||
}
|
||||
keys := []string{packed.ArrowFieldIdMetadataKey}
|
||||
values := []string{strconv.Itoa(int(field.GetFieldID()))}
|
||||
|
||||
@@ -570,7 +658,7 @@ func newRecordBuilderArrowField(field *schemapb.FieldSchema, arrowType arrow.Dat
|
||||
}
|
||||
|
||||
return arrow.Field{
|
||||
Name: field.GetName(),
|
||||
Name: name,
|
||||
Type: arrowType,
|
||||
Nullable: field.GetNullable(),
|
||||
Metadata: arrow.NewMetadata(keys, values),
|
||||
|
||||
@@ -17,6 +17,7 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math/rand"
|
||||
"testing"
|
||||
|
||||
@@ -258,6 +259,160 @@ func TestGenerateEmptyArray(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordBuilderPreservesTextRepresentation(t *testing.T) {
|
||||
for _, dataType := range []arrow.DataType{arrow.BinaryTypes.String, arrow.BinaryTypes.Binary} {
|
||||
t.Run(dataType.Name(), func(t *testing.T) {
|
||||
field := &schemapb.FieldSchema{FieldID: 100, Name: "text", DataType: schemapb.DataType_Text, Nullable: true}
|
||||
rb := NewRecordBuilder(&schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{field}})
|
||||
defer rb.Release()
|
||||
for batch := 0; batch < 2; batch++ {
|
||||
for _, value := range []string{"decoded text", "second record"} {
|
||||
input := array.NewBuilder(memory.DefaultAllocator, dataType)
|
||||
switch input := input.(type) {
|
||||
case *array.StringBuilder:
|
||||
input.Append(value)
|
||||
case *array.BinaryBuilder:
|
||||
input.Append([]byte(value))
|
||||
}
|
||||
input.AppendNull()
|
||||
column := input.NewArray()
|
||||
input.Release()
|
||||
record := NewSimpleArrowRecord(array.NewRecord(
|
||||
arrow.NewSchema([]arrow.Field{{Name: field.Name, Type: dataType, Nullable: true}}, nil),
|
||||
[]arrow.Array{column}, 2), map[FieldID]int{field.FieldID: 0})
|
||||
column.Release()
|
||||
err := rb.Append(record, 0, record.Len())
|
||||
record.Release()
|
||||
require.NoError(t, err)
|
||||
}
|
||||
output := rb.Build()
|
||||
column := output.Column(field.FieldID)
|
||||
assert.Equal(t, dataType, column.DataType())
|
||||
assert.Equal(t, 4, column.Len())
|
||||
assert.True(t, column.IsNull(1))
|
||||
assert.True(t, column.IsNull(3))
|
||||
switch column := column.(type) {
|
||||
case *array.String:
|
||||
assert.Equal(t, "decoded text", column.Value(0))
|
||||
assert.Equal(t, "second record", column.Value(2))
|
||||
case *array.Binary:
|
||||
assert.Equal(t, []byte("decoded text"), column.Value(0))
|
||||
assert.Equal(t, []byte("second record"), column.Value(2))
|
||||
}
|
||||
output.Release()
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordBuilderTextNullsAcrossRepresentations(t *testing.T) {
|
||||
field := &schemapb.FieldSchema{FieldID: 100, Name: "text", DataType: schemapb.DataType_Text, Nullable: true}
|
||||
rb := NewRecordBuilder(&schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{field}})
|
||||
defer rb.Release()
|
||||
for _, dataType := range []arrow.DataType{arrow.BinaryTypes.Binary, arrow.BinaryTypes.String, arrow.BinaryTypes.Binary} {
|
||||
input := array.NewBuilder(memory.DefaultAllocator, dataType)
|
||||
if input, ok := input.(*array.StringBuilder); ok {
|
||||
input.Append("decoded")
|
||||
}
|
||||
input.AppendNull()
|
||||
column := input.NewArray()
|
||||
input.Release()
|
||||
record := NewSimpleArrowRecord(array.NewRecord(
|
||||
arrow.NewSchema([]arrow.Field{{Name: field.Name, Type: dataType, Nullable: true}}, nil),
|
||||
[]arrow.Array{column}, int64(column.Len())), map[FieldID]int{field.FieldID: 0})
|
||||
column.Release()
|
||||
err := rb.Append(record, 0, record.Len())
|
||||
record.Release()
|
||||
require.NoError(t, err)
|
||||
}
|
||||
output := rb.Build()
|
||||
defer output.Release()
|
||||
values := output.Column(field.FieldID).(*array.String)
|
||||
require.Equal(t, 4, values.Len())
|
||||
require.Equal(t, 3, values.NullN())
|
||||
require.Equal(t, "decoded", values.Value(1))
|
||||
for _, row := range []int{0, 2, 3} {
|
||||
require.True(t, values.IsNull(row))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordBuilderFillsNullableInt64Defaults(t *testing.T) {
|
||||
const defaultValue int64 = 1_700_000_000_000_000
|
||||
for _, dataType := range []schemapb.DataType{schemapb.DataType_Int64, schemapb.DataType_Timestamptz} {
|
||||
t.Run(dataType.String(), func(t *testing.T) {
|
||||
for _, withDefault := range []bool{false, true} {
|
||||
field := &schemapb.FieldSchema{FieldID: 100, Name: "value", DataType: dataType, Nullable: true}
|
||||
if withDefault {
|
||||
field.DefaultValue = &schemapb.ValueField{Data: &schemapb.ValueField_LongData{LongData: defaultValue}}
|
||||
if dataType == schemapb.DataType_Timestamptz {
|
||||
field.DefaultValue = &schemapb.ValueField{Data: &schemapb.ValueField_TimestamptzData{TimestamptzData: defaultValue}}
|
||||
}
|
||||
}
|
||||
rb := NewRecordBuilder(&schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{field}})
|
||||
defer rb.Release()
|
||||
for batch := 0; batch < 2; batch++ {
|
||||
input := array.NewInt64Builder(memory.DefaultAllocator)
|
||||
input.AppendNull()
|
||||
input.Append(0)
|
||||
input.Append(123)
|
||||
column := input.NewArray()
|
||||
input.Release()
|
||||
record := NewSimpleArrowRecord(array.NewRecord(
|
||||
arrow.NewSchema([]arrow.Field{{Name: field.Name, Type: column.DataType(), Nullable: true}}, nil),
|
||||
[]arrow.Array{column}, 3), map[FieldID]int{field.FieldID: 0})
|
||||
column.Release()
|
||||
err := rb.Append(record, 0, record.Len())
|
||||
record.Release()
|
||||
require.NoError(t, err)
|
||||
output := rb.Build()
|
||||
values := output.Column(field.FieldID).(*array.Int64)
|
||||
assert.Equal(t, !withDefault, values.IsNull(0))
|
||||
if withDefault {
|
||||
assert.Equal(t, defaultValue, values.Value(0))
|
||||
}
|
||||
assert.EqualValues(t, 0, values.Value(1))
|
||||
assert.EqualValues(t, 123, values.Value(2))
|
||||
output.Release()
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordBuilderFillsNullableDoubleDefault(t *testing.T) {
|
||||
for _, withDefault := range []bool{false, true} {
|
||||
field := &schemapb.FieldSchema{FieldID: 100, Name: "double", DataType: schemapb.DataType_Double, Nullable: true}
|
||||
if withDefault {
|
||||
field.DefaultValue = &schemapb.ValueField{Data: &schemapb.ValueField_DoubleData{DoubleData: 3.5}}
|
||||
}
|
||||
rb := NewRecordBuilder(&schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{field}})
|
||||
for batch := 0; batch < 2; batch++ {
|
||||
input := array.NewFloat64Builder(memory.DefaultAllocator)
|
||||
input.AppendNull()
|
||||
input.Append(0)
|
||||
input.Append(1.25)
|
||||
column := input.NewArray()
|
||||
input.Release()
|
||||
record := NewSimpleArrowRecord(array.NewRecord(
|
||||
arrow.NewSchema([]arrow.Field{{Name: field.Name, Type: column.DataType(), Nullable: true}}, nil),
|
||||
[]arrow.Array{column}, 3), map[FieldID]int{field.FieldID: 0})
|
||||
column.Release()
|
||||
require.NoError(t, rb.Append(record, 0, record.Len()))
|
||||
record.Release()
|
||||
output := rb.Build()
|
||||
values := output.Column(field.FieldID).(*array.Float64)
|
||||
assert.Equal(t, !withDefault, values.IsNull(0))
|
||||
if withDefault {
|
||||
assert.Equal(t, 3.5, values.Value(0))
|
||||
}
|
||||
assert.Equal(t, float64(0), values.Value(1))
|
||||
assert.Equal(t, 1.25, values.Value(2))
|
||||
output.Release()
|
||||
}
|
||||
rb.Release()
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordBuilderFillsNullableGeometryDefault(t *testing.T) {
|
||||
const defaultWKT = "POINT (1 2)"
|
||||
defaultWKB, err := common.ConvertWKTToWKB(defaultWKT)
|
||||
@@ -600,3 +755,57 @@ func TestRecordBuilderBuildReleasesCreatorRefs(t *testing.T) {
|
||||
}()
|
||||
alloc.AssertSize(t, 0)
|
||||
}
|
||||
|
||||
func TestRecordBuilderAccountsForNullBuffers(t *testing.T) {
|
||||
fields := make([]*schemapb.FieldSchema, 243)
|
||||
arrowFields := make([]arrow.Field, len(fields))
|
||||
field2Col := make(map[FieldID]int, len(fields))
|
||||
for i := range fields {
|
||||
fields[i] = &schemapb.FieldSchema{FieldID: int64(i), Name: fmt.Sprint(i), DataType: schemapb.DataType_Double, Nullable: true}
|
||||
arrowFields[i] = arrow.Field{Name: fields[i].Name, Type: arrow.PrimitiveTypes.Float64, Nullable: true}
|
||||
if i < 3 {
|
||||
fields[i].DataType = schemapb.DataType_Int64
|
||||
arrowFields[i].Type = arrow.PrimitiveTypes.Int64
|
||||
}
|
||||
field2Col[int64(i)] = i
|
||||
}
|
||||
input := array.NewRecordBuilder(memory.DefaultAllocator, arrow.NewSchema(arrowFields, nil))
|
||||
for i := range fields {
|
||||
if i < 3 {
|
||||
for row := range 4096 {
|
||||
input.Field(i).(*array.Int64Builder).Append(int64(row))
|
||||
}
|
||||
} else {
|
||||
input.Field(i).AppendNulls(4096)
|
||||
}
|
||||
}
|
||||
source := NewSimpleArrowRecord(input.NewRecord(), field2Col)
|
||||
input.Release()
|
||||
defer source.Release()
|
||||
|
||||
checked := memory.NewCheckedAllocator(memory.DefaultAllocator)
|
||||
original := memory.DefaultAllocator
|
||||
memory.DefaultAllocator = checked
|
||||
rb := NewRecordBuilder(&schemapb.CollectionSchema{Fields: fields})
|
||||
memory.DefaultAllocator = original
|
||||
defer func() { rb.Release(); checked.AssertSize(t, 0) }()
|
||||
require.NoError(t, rb.Append(source, 0, source.Len()))
|
||||
require.EqualValues(t, checked.CurrentAlloc(), rb.GetMemorySize(), "batching must include null slots, validity buffers and capacity growth")
|
||||
first := rb.Build()
|
||||
defer func() {
|
||||
if first != nil {
|
||||
first.Release()
|
||||
}
|
||||
}()
|
||||
require.Zero(t, rb.GetMemorySize(), "transferred records are no longer pending builder memory")
|
||||
retained := checked.CurrentAlloc()
|
||||
require.NoError(t, rb.Append(source, 0, source.Len()))
|
||||
require.EqualValues(t, checked.CurrentAlloc()-retained, rb.GetMemorySize())
|
||||
first.Release()
|
||||
first = nil
|
||||
require.EqualValues(t, checked.CurrentAlloc(), rb.GetMemorySize(), "late releases must not change the next batch's counter")
|
||||
second := rb.Build()
|
||||
require.Zero(t, rb.GetMemorySize())
|
||||
second.Release()
|
||||
checked.AssertSize(t, 0)
|
||||
}
|
||||
|
||||
@@ -667,6 +667,13 @@ func (r *absentFilledRecord) Column(i FieldID) arrow.Array {
|
||||
return r.base.Column(i)
|
||||
}
|
||||
|
||||
func (r *absentFilledRecord) TryColumn(i FieldID) (arrow.Array, bool) {
|
||||
if col, ok := r.computed[i]; ok {
|
||||
return col, true
|
||||
}
|
||||
return TryRecordColumn(r.base, i)
|
||||
}
|
||||
|
||||
func (r *absentFilledRecord) Len() int { return r.base.Len() }
|
||||
|
||||
func (r *absentFilledRecord) Retain() {
|
||||
|
||||
@@ -169,6 +169,11 @@ func TestAbsentFieldFillReaderFillsDefaultAndNull(t *testing.T) {
|
||||
require.Equal(t, []int64{42, 42, 42}, []int64{def.Value(0), def.Value(1), def.Value(2)}) // default materialized
|
||||
require.Equal(t, 3, rec.Column(nullField).(*array.Int64).NullN()) // no default -> null
|
||||
require.Equal(t, int64(2), rec.Column(idField).(*array.Int64).Value(1)) // present passthrough
|
||||
column, ok := TryRecordColumn(rec, defField)
|
||||
require.True(t, ok)
|
||||
require.Same(t, def, column)
|
||||
_, ok = TryRecordColumn(rec, 999)
|
||||
require.False(t, ok)
|
||||
}
|
||||
|
||||
func TestAbsentFieldFillReaderRejectsNonNullableAbsent(t *testing.T) {
|
||||
|
||||
@@ -50,7 +50,7 @@ func RecordToInsertData(
|
||||
|
||||
for _, field := range typeutil.GetAllFieldSchemas(schema) {
|
||||
fieldID := field.GetFieldID()
|
||||
col, ok := recordColumn(rec, fieldID)
|
||||
col, ok := TryRecordColumn(rec, fieldID)
|
||||
if !ok {
|
||||
if requiredFields != nil && requiredFields.Contain(fieldID) {
|
||||
return nil, merr.WrapErrParameterInvalidMsg("required field %s (ID=%d) not found in record",
|
||||
@@ -113,8 +113,18 @@ func RecordToInsertData(
|
||||
return insertData, nil
|
||||
}
|
||||
|
||||
func recordColumn(rec Record, fieldID FieldID) (col arrow.Array, ok bool) {
|
||||
// TryRecordColumn checks field presence for built-in Records and wrappers that
|
||||
// implement TryColumn. Custom Records should return nil for absent fields or
|
||||
// implement TryColumn to support a non-panicking lookup.
|
||||
func TryRecordColumn(rec Record, fieldID FieldID) (col arrow.Array, ok bool) {
|
||||
if rec == nil {
|
||||
return nil, false
|
||||
}
|
||||
switch r := rec.(type) {
|
||||
case interface {
|
||||
TryColumn(FieldID) (arrow.Array, bool)
|
||||
}:
|
||||
return r.TryColumn(fieldID)
|
||||
case *simpleArrowRecord:
|
||||
colIdx, ok := r.field2Col[fieldID]
|
||||
if !ok {
|
||||
@@ -127,6 +137,11 @@ func recordColumn(rec Record, fieldID FieldID) (col arrow.Array, ok bool) {
|
||||
case *compositeRecord:
|
||||
col := r.Column(fieldID)
|
||||
return col, col != nil
|
||||
case *selectiveRecord:
|
||||
if r.fieldId != fieldID {
|
||||
return nil, false
|
||||
}
|
||||
return TryRecordColumn(r.r, fieldID)
|
||||
}
|
||||
|
||||
col = rec.Column(fieldID)
|
||||
|
||||
@@ -25,6 +25,7 @@ import (
|
||||
|
||||
"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/samber/lo"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
@@ -45,18 +46,43 @@ import (
|
||||
// addition to the appended field) to exactly the writer's column set. Passing
|
||||
// such an over-wide record straight to the FFI writer misaligns column count
|
||||
// against the writer's manifest and crashes loon_writer_write.
|
||||
func toWriterRecord(r Record, schema *schemapb.CollectionSchema, arrowSchema *arrow.Schema) (arrow.Record, func()) {
|
||||
func toWriterRecord(r Record, schema *schemapb.CollectionSchema, arrowSchema *arrow.Schema) (arrow.Record, func(), error) {
|
||||
if sar, ok := r.(*simpleArrowRecord); ok && sar.r.Schema().Equal(arrowSchema) {
|
||||
return sar.r, func() {}
|
||||
return sar.r, func() {}, nil
|
||||
}
|
||||
// Include struct sub-fields, matching the writer's arrow schema layout.
|
||||
allFields := typeutil.GetAllFieldSchemas(schema)
|
||||
arrays := make([]arrow.Array, len(allFields))
|
||||
var normalized []arrow.Array
|
||||
defer func() {
|
||||
for _, column := range normalized {
|
||||
column.Release()
|
||||
}
|
||||
}()
|
||||
for i, field := range allFields {
|
||||
arrays[i] = r.Column(field.GetFieldID())
|
||||
column := r.Column(field.GetFieldID())
|
||||
expected := arrowSchema.Field(i).Type
|
||||
if column == nil || column.Len() < r.Len() {
|
||||
return nil, nil, merr.WrapErrServiceInternalMsg("missing or short writer column for field %s", field.GetName())
|
||||
}
|
||||
if !arrow.TypeEqual(column.DataType(), expected) {
|
||||
// Absent TEXT fields are filled with binary nulls, including for
|
||||
// decoded UTF8 writers. Nulls have no LOB/string payload to convert;
|
||||
// normalize them to the writer's fixed schema on every batch.
|
||||
sourceType, targetType := column.DataType().ID(), expected.ID()
|
||||
if field.GetDataType() == schemapb.DataType_Text && column.NullN() == column.Len() &&
|
||||
(sourceType == arrow.STRING || sourceType == arrow.BINARY) &&
|
||||
(targetType == arrow.STRING || targetType == arrow.BINARY) {
|
||||
column = array.MakeArrayOfNull(memory.DefaultAllocator, expected, r.Len())
|
||||
normalized = append(normalized, column)
|
||||
} else {
|
||||
return nil, nil, merr.WrapErrServiceInternalMsg("writer column type mismatch for field %s: got %s, expected %s",
|
||||
field.GetName(), column.DataType(), expected)
|
||||
}
|
||||
}
|
||||
arrays[i] = column
|
||||
}
|
||||
rec := array.NewRecord(arrowSchema, arrays, int64(r.Len()))
|
||||
return rec, rec.Release
|
||||
return rec, rec.Release, nil
|
||||
}
|
||||
|
||||
var _ RecordWriter = (*packedRecordWriter)(nil)
|
||||
@@ -78,20 +104,11 @@ type packedRecordWriter struct {
|
||||
}
|
||||
|
||||
func (pw *packedRecordWriter) Write(r Record) error {
|
||||
var rec arrow.Record
|
||||
sar, ok := r.(*simpleArrowRecord)
|
||||
if !ok {
|
||||
// Get all fields including struct sub-fields
|
||||
allFields := typeutil.GetAllFieldSchemas(pw.schema)
|
||||
arrays := make([]arrow.Array, len(allFields))
|
||||
for i, field := range allFields {
|
||||
arrays[i] = r.Column(field.FieldID)
|
||||
}
|
||||
rec = array.NewRecord(pw.arrowSchema, arrays, int64(r.Len()))
|
||||
defer rec.Release()
|
||||
} else {
|
||||
rec = sar.r
|
||||
rec, release, err := toWriterRecord(r, pw.schema, pw.arrowSchema)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer release()
|
||||
pw.rowNum += int64(r.Len())
|
||||
for col, arr := range rec.Columns() {
|
||||
size := calculateActualDataSize(arr)
|
||||
@@ -218,7 +235,10 @@ type packedRecordBatchWriter struct {
|
||||
}
|
||||
|
||||
func (pw *packedRecordBatchWriter) Write(r Record) error {
|
||||
rec, release := toWriterRecord(r, pw.schema, pw.arrowSchema)
|
||||
rec, release, err := toWriterRecord(r, pw.schema, pw.arrowSchema)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer release()
|
||||
pw.rowNum += int64(r.Len())
|
||||
for col, arr := range rec.Columns() {
|
||||
@@ -466,20 +486,11 @@ type packedTextBatchWriter struct {
|
||||
}
|
||||
|
||||
func (pw *packedTextBatchWriter) Write(r Record) error {
|
||||
var rec arrow.Record
|
||||
sar, ok := r.(*simpleArrowRecord)
|
||||
if !ok {
|
||||
// get all fields including struct sub-fields
|
||||
allFields := typeutil.GetAllFieldSchemas(pw.schema)
|
||||
arrays := make([]arrow.Array, len(allFields))
|
||||
for i, field := range allFields {
|
||||
arrays[i] = r.Column(field.FieldID)
|
||||
}
|
||||
rec = array.NewRecord(pw.arrowSchema, arrays, int64(r.Len()))
|
||||
defer rec.Release()
|
||||
} else {
|
||||
rec = sar.r
|
||||
rec, release, err := toWriterRecord(r, pw.schema, pw.arrowSchema)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer release()
|
||||
pw.rowNum += int64(r.Len())
|
||||
for col, arr := range rec.Columns() {
|
||||
size := calculateActualDataSize(arr)
|
||||
|
||||
@@ -22,9 +22,12 @@ import (
|
||||
"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/bytedance/mockey"
|
||||
"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/merr"
|
||||
)
|
||||
|
||||
// TestToWriterRecordProjectsOverWideRecord pins the schema-bump additive
|
||||
@@ -62,7 +65,8 @@ func TestToWriterRecordProjectsOverWideRecord(t *testing.T) {
|
||||
defer wideRec.Release()
|
||||
overWide := NewSimpleArrowRecord(wideRec, map[FieldID]int{0: 0, 103: 1})
|
||||
|
||||
rec, release := toWriterRecord(overWide, writerSchema, writerArrow)
|
||||
rec, release, err := toWriterRecord(overWide, writerSchema, writerArrow)
|
||||
require.NoError(t, err)
|
||||
defer release()
|
||||
|
||||
require.EqualValues(t, 1, rec.NumCols(), "record must be narrowed to the writer's single column")
|
||||
@@ -92,8 +96,171 @@ func TestToWriterRecordFastPathOnMatchingSchema(t *testing.T) {
|
||||
defer matchRec.Release()
|
||||
sar := NewSimpleArrowRecord(matchRec, map[FieldID]int{103: 0})
|
||||
|
||||
rec, release := toWriterRecord(sar, writerSchema, writerArrow)
|
||||
rec, release, err := toWriterRecord(sar, writerSchema, writerArrow)
|
||||
require.NoError(t, err)
|
||||
defer release()
|
||||
|
||||
require.Same(t, matchRec, rec, "matching schema must take the fast path and return the backing record")
|
||||
}
|
||||
|
||||
func TestToWriterRecordV3BuilderFastPath(t *testing.T) {
|
||||
field := &schemapb.FieldSchema{FieldID: 103, Name: "added", ExternalField: "source_added", DataType: schemapb.DataType_VarChar}
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{field}}
|
||||
inputSchema, err := ConvertToArrowSchema(schema, false)
|
||||
require.NoError(t, err)
|
||||
inputBuilder := array.NewStringBuilder(memory.DefaultAllocator)
|
||||
inputBuilder.Append("value")
|
||||
column := inputBuilder.NewArray()
|
||||
inputBuilder.Release()
|
||||
input := NewSimpleArrowRecord(array.NewRecord(inputSchema, []arrow.Array{column}, 1), map[FieldID]int{103: 0})
|
||||
column.Release()
|
||||
builder := NewRecordBuilderWithFieldID(schema)
|
||||
defer builder.Release()
|
||||
require.NoError(t, builder.Append(input, 0, 1))
|
||||
input.Release()
|
||||
batch := builder.Build()
|
||||
defer batch.Release()
|
||||
writerSchema, err := ConvertToArrowSchema(schema, true)
|
||||
require.NoError(t, err)
|
||||
converted, release, err := toWriterRecord(batch, schema, writerSchema)
|
||||
require.NoError(t, err)
|
||||
defer release()
|
||||
require.Same(t, batch.(*simpleArrowRecord).r, converted)
|
||||
}
|
||||
|
||||
func TestToWriterRecordV3BuilderNormalizesNullText(t *testing.T) {
|
||||
field := &schemapb.FieldSchema{FieldID: 103, Name: "text", DataType: schemapb.DataType_Text, Nullable: true}
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{field}}
|
||||
column := array.MakeArrayOfNull(memory.DefaultAllocator, arrow.BinaryTypes.Binary, 1)
|
||||
input := NewSimpleArrowRecord(array.NewRecord(arrow.NewSchema([]arrow.Field{
|
||||
{Name: field.Name, Type: arrow.BinaryTypes.Binary, Nullable: true},
|
||||
}, nil), []arrow.Array{column}, 1), map[FieldID]int{103: 0})
|
||||
column.Release()
|
||||
builder := NewRecordBuilderWithFieldID(schema)
|
||||
defer builder.Release()
|
||||
require.NoError(t, builder.Append(input, 0, 1))
|
||||
input.Release()
|
||||
batch := builder.Build()
|
||||
defer batch.Release()
|
||||
writerSchema, err := ConvertToArrowSchema(schema, true)
|
||||
require.NoError(t, err)
|
||||
converted, release, err := toWriterRecord(batch, schema, writerSchema)
|
||||
require.NoError(t, err)
|
||||
defer release()
|
||||
require.NotSame(t, batch.(*simpleArrowRecord).r, converted)
|
||||
require.True(t, writerSchema.Equal(converted.Schema()))
|
||||
require.Equal(t, 1, converted.Column(0).NullN())
|
||||
}
|
||||
|
||||
func TestToWriterRecordNormalizesNullText(t *testing.T) {
|
||||
field := &schemapb.FieldSchema{FieldID: 100, Name: "text", DataType: schemapb.DataType_Text, Nullable: true}
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{field}}
|
||||
for _, target := range []arrow.DataType{arrow.BinaryTypes.String, arrow.BinaryTypes.Binary} {
|
||||
t.Run(target.Name(), func(t *testing.T) {
|
||||
source := arrow.BinaryTypes.Binary
|
||||
if target.ID() == arrow.BINARY {
|
||||
source = arrow.BinaryTypes.String
|
||||
}
|
||||
column := array.MakeArrayOfNull(memory.DefaultAllocator, source, 3)
|
||||
defer column.Release()
|
||||
record := NewSimpleArrowRecord(array.NewRecord(arrow.NewSchema([]arrow.Field{
|
||||
{Name: field.Name, Type: source, Nullable: true},
|
||||
}, nil), []arrow.Array{column}, 3), map[FieldID]int{100: 0})
|
||||
defer record.Release()
|
||||
writerSchema := arrow.NewSchema([]arrow.Field{{Name: field.Name, Type: target, Nullable: true}}, nil)
|
||||
require.NotPanics(t, func() {
|
||||
converted, release, err := toWriterRecord(record, schema, writerSchema)
|
||||
require.NoError(t, err)
|
||||
defer release()
|
||||
require.True(t, writerSchema.Equal(converted.Schema()))
|
||||
require.Equal(t, target, converted.Column(0).DataType())
|
||||
require.Equal(t, 3, converted.Column(0).NullN())
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPackedTextWriterUsesWriterSchemaAcrossBatches(t *testing.T) {
|
||||
field := &schemapb.FieldSchema{FieldID: 100, Name: "text", DataType: schemapb.DataType_Text, Nullable: true}
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{field}}
|
||||
for _, target := range []arrow.DataType{arrow.BinaryTypes.String, arrow.BinaryTypes.Binary} {
|
||||
t.Run(target.Name(), func(t *testing.T) {
|
||||
writerSchema := arrow.NewSchema([]arrow.Field{{Name: field.Name, Type: target, Nullable: true}}, nil)
|
||||
writer := &packedTextBatchWriter{schema: schema, arrowSchema: writerSchema, writer: &packed.FFISegmentWriter{}}
|
||||
writes := 0
|
||||
patch := mockey.Mock((*packed.FFISegmentWriter).Write).To(func(_ *packed.FFISegmentWriter, record arrow.Record) error {
|
||||
require.True(t, writerSchema.Equal(record.Schema()))
|
||||
require.Equal(t, target, record.Column(0).DataType())
|
||||
if writes != 1 {
|
||||
require.Equal(t, 1, record.Column(0).NullN())
|
||||
} else {
|
||||
switch column := record.Column(0).(type) {
|
||||
case *array.String:
|
||||
require.Equal(t, "decoded text", column.Value(0))
|
||||
case *array.Binary:
|
||||
require.Equal(t, []byte("encoded ref"), column.Value(0))
|
||||
}
|
||||
}
|
||||
writes++
|
||||
return nil
|
||||
}).Build()
|
||||
defer patch.UnPatch()
|
||||
rb := NewRecordBuilder(schema)
|
||||
defer rb.Release()
|
||||
for batch := range 3 {
|
||||
var sourceType arrow.DataType = arrow.BinaryTypes.Binary
|
||||
if target.ID() == arrow.BINARY {
|
||||
sourceType = arrow.BinaryTypes.String
|
||||
}
|
||||
if batch == 1 {
|
||||
sourceType = target
|
||||
}
|
||||
input := array.NewBuilder(memory.DefaultAllocator, sourceType)
|
||||
if batch == 1 {
|
||||
switch input := input.(type) {
|
||||
case *array.StringBuilder:
|
||||
input.Append("decoded text")
|
||||
case *array.BinaryBuilder:
|
||||
input.Append([]byte("encoded ref"))
|
||||
}
|
||||
} else {
|
||||
input.AppendNull()
|
||||
}
|
||||
column := input.NewArray()
|
||||
input.Release()
|
||||
record := NewSimpleArrowRecord(array.NewRecord(arrow.NewSchema([]arrow.Field{
|
||||
{Name: field.Name, Type: sourceType, Nullable: true},
|
||||
}, nil), []arrow.Array{column}, 1), map[FieldID]int{100: 0})
|
||||
column.Release()
|
||||
require.NoError(t, rb.Append(record, 0, 1))
|
||||
record.Release()
|
||||
out := rb.Build()
|
||||
err := writer.Write(out)
|
||||
out.Release()
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.Equal(t, 3, writes)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestToWriterRecordRejectsNonNullTextTypeMismatch(t *testing.T) {
|
||||
field := &schemapb.FieldSchema{FieldID: 100, Name: "text", DataType: schemapb.DataType_Text, Nullable: true}
|
||||
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{field}}
|
||||
builder := array.NewBinaryBuilder(memory.DefaultAllocator, arrow.BinaryTypes.Binary)
|
||||
builder.Append([]byte("encoded LOB reference"))
|
||||
column := builder.NewArray()
|
||||
builder.Release()
|
||||
defer column.Release()
|
||||
record := NewSimpleArrowRecord(array.NewRecord(arrow.NewSchema([]arrow.Field{
|
||||
{Name: field.Name, Type: arrow.BinaryTypes.Binary, Nullable: true},
|
||||
}, nil), []arrow.Array{column}, 1), map[FieldID]int{100: 0})
|
||||
defer record.Release()
|
||||
writerSchema := arrow.NewSchema([]arrow.Field{{Name: field.Name, Type: arrow.BinaryTypes.String, Nullable: true}}, nil)
|
||||
require.NotPanics(t, func() {
|
||||
converted, release, err := toWriterRecord(record, schema, writerSchema)
|
||||
require.ErrorIs(t, err, merr.ErrServiceInternal)
|
||||
require.Nil(t, converted)
|
||||
require.Nil(t, release)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -413,8 +413,10 @@ func TestMergeSort(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestMergeSortReturnsRecordBuilderAppendError(t *testing.T) {
|
||||
textBuilder := array.NewStringBuilder(memory.DefaultAllocator)
|
||||
textBuilder.Append("not-a-lob-ref")
|
||||
// UTF8 is valid decoded TEXT. An integer column violates both supported
|
||||
// TEXT representations and must still propagate the Append error.
|
||||
textBuilder := array.NewInt64Builder(memory.DefaultAllocator)
|
||||
textBuilder.Append(42)
|
||||
textColumn := textBuilder.NewArray()
|
||||
defer textColumn.Release()
|
||||
textBuilder.Release()
|
||||
@@ -428,7 +430,7 @@ func TestMergeSortReturnsRecordBuilderAppendError(t *testing.T) {
|
||||
rec := NewSimpleArrowRecord(array.NewRecord(
|
||||
arrow.NewSchema([]arrow.Field{
|
||||
{Name: "pk", Type: arrow.PrimitiveTypes.Int64},
|
||||
{Name: "text", Type: arrow.BinaryTypes.String},
|
||||
{Name: "text", Type: arrow.PrimitiveTypes.Int64},
|
||||
}, nil),
|
||||
[]arrow.Array{pkColumn, textColumn},
|
||||
1,
|
||||
|
||||
Reference in New Issue
Block a user