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:
XuanYang-cn
2026-10-09 10:39:44 +08:00
committed by GitHub
parent 5f0342f680
commit 2ca0398246
16 changed files with 1783 additions and 88 deletions
@@ -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 := &paramtable.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))
}
+95 -7
View File
@@ -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),
+209
View File
@@ -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)
}
+7
View File
@@ -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() {
+5
View File
@@ -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) {
+17 -2
View File
@@ -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)
+43 -32
View File
@@ -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)
+169 -2
View File
@@ -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)
})
}
+5 -3
View File
@@ -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,