Repository: polarsignals/frostdb
Pull request #739 (2024-02-14 12:30:46)

Set OptInt32Builder data to nil
The reason we want to do this is so that when pooling the builder as part of the `MultiColSorter` we don't reuse the builder's data for a concurrent existing array.

asubiotto: Can you add a regression test? You don't even need concurrency. I would just reuse a builder and check that a previously returned array is intact.

I'm a little concerned by how this will affect scan performance though since we will introduce more allocations. Could you run `BenchmarkQuery` on a prod db?
asubiotto: I would also change this for all builders. I think this is a "contract-level" change. i.e. the caller of `NewArray` retains ownership of this array and the builder loses it, so any builder reuse should not affect existing arrays.
metalmatze: Here are the results for the `BenchmarkQuery`:
```
name             old time/op  new time/op  delta
Query/Types-24    754µs ± 1%   773µs ± 2%  +2.53%  (p=0.032 n=5+5)
Query/Labels-24   292µs ± 1%   263µs ±10%  -9.79%  (p=0.016 n=4+5)
Query/Values-24   592µs ±20%   628µs ±16%    ~     (p=0.548 n=5+5)
Query/Merge-24    560ms ± 7%   568ms ±10%    ~     (p=1.000 n=5+5)
Query/Range-24   85.3ms ± 5%  87.0ms ±11%    ~     (p=0.310 n=5+5)
```
It doesn't seem too bad, is it?

Patch:
diff --git a/db.go b/db.go
index 4bf531c..b2c63f1 100644
--- a/db.go
+++ b/db.go
@@ -984,8 +984,22 @@ func (db *DB) Close(options ...CloseOption) error {
 	for _, opt := range options {
 		opt(opts)
 	}
-	level.Info(db.logger).Log("msg", "closing DB")
+
 	shouldPersist := len(db.sinks) > 0 && !db.columnStore.manualBlockRotation
+	if !shouldPersist && db.columnStore.snapshotTriggerSize != 0 && !opts.clearStorage {
+		start := time.Now()
+		db.snapshot(context.Background(), false, func() {
+			level.Info(db.logger).Log("msg", "snapshot on close completed", "duration", time.Since(start))
+			if err := db.reclaimDiskSpace(context.Background(), nil); err != nil {
+				level.Error(db.logger).Log(
+					"msg", "failed to reclaim disk space after snapshot",
+					"err", err,
+				)
+			}
+		})
+	}
+
+	level.Info(db.logger).Log("msg", "closing DB")
 	for _, table := range db.tables {
 		table.close()
 		if shouldPersist {
@@ -1000,19 +1014,6 @@ func (db *DB) Close(options ...CloseOption) error {
 	}
 	level.Info(db.logger).Log("msg", "closed all tables")
 
-	if !shouldPersist && db.columnStore.snapshotTriggerSize != 0 && !opts.clearStorage {
-		start := time.Now()
-		db.snapshot(context.Background(), false, func() {
-			level.Info(db.logger).Log("msg", "snapshot on close completed", "duration", time.Since(start))
-			if err := db.reclaimDiskSpace(context.Background(), nil); err != nil {
-				level.Error(db.logger).Log(
-					"msg", "failed to reclaim disk space after snapshot",
-					"err", err,
-				)
-			}
-		})
-	}
-
 	if err := db.closeInternal(); err != nil {
 		return err
 	}
diff --git a/pqarrow/builder/optbuilders.go b/pqarrow/builder/optbuilders.go
index bde4d54..9187640 100644
--- a/pqarrow/builder/optbuilders.go
+++ b/pqarrow/builder/optbuilders.go
@@ -179,7 +179,7 @@ func (b *OptBinaryBuilder) NewArray() arrow.Array {
 	)
 	b.reset()
 	b.offsets = b.offsets[:0]
-	b.data = b.data[:0]
+	b.data = nil
 
 	return array.NewBinaryData(data)
 }
@@ -340,7 +340,7 @@ func (b *OptInt64Builder) NewArray() arrow.Array {
 		0,
 	)
 	b.reset()
-	b.data = b.data[:0]
+	b.data = nil
 	return array.NewInt64Data(data)
 }
 
@@ -459,7 +459,7 @@ func (b *OptBooleanBuilder) NewArray() arrow.Array {
 		0,
 	)
 	b.reset()
-	b.data = b.data[:0]
+	b.data = nil
 	array := array.NewBooleanData(data)
 	return array
 }
@@ -584,7 +584,7 @@ func (b *OptInt32Builder) NewArray() arrow.Array {
 		0,
 	)
 	b.reset()
-	b.data = b.data[:0]
+	b.data = nil
 	return array.NewInt32Data(data)
 }
 
diff --git a/snapshot.go b/snapshot.go
index c4a2d3d..20374a3 100644
--- a/snapshot.go
+++ b/snapshot.go
@@ -346,18 +346,18 @@ func (w *offsetWriter) checksum() uint32 {
 
 func (db *DB) snapshotWriter(tx uint64) func(context.Context, io.Writer) error {
 	return func(ctx context.Context, w io.Writer) error {
-		return WriteSnapshot(ctx, tx, db, w, false)
+		return WriteSnapshot(ctx, tx, db, w)
 	}
 }
 
 // offlineSnapshotWriter is used when a database is closing after all the tables have closed.
 func (db *DB) offlineSnapshotWriter(tx uint64) func(context.Context, io.Writer) error {
 	return func(ctx context.Context, w io.Writer) error {
-		return WriteSnapshot(ctx, tx, db, w, true)
+		return WriteSnapshot(ctx, tx, db, w)
 	}
 }
 
-func WriteSnapshot(ctx context.Context, tx uint64, db *DB, w io.Writer, offline bool) error {
+func WriteSnapshot(ctx context.Context, tx uint64, db *DB, w io.Writer) error {
 	offW := newOffsetWriter(w)
 	w = offW
 	var tables []*Table
@@ -376,18 +376,11 @@ func WriteSnapshot(ctx context.Context, tx uint64, db *DB, w io.Writer, offline
 		if err := func() error {
 			// Obtain a write block to prevent racing with
 			// compaction/persistence.
-			var block *TableBlock
-			if offline {
-				block = t.ActiveBlock()
-			} else {
-				var done func()
-				var err error
-				block, done, err = t.ActiveWriteBlock()
-				if err != nil {
-					return err
-				}
-				defer done()
+			block, done, err := t.ActiveWriteBlock()
+			if err != nil {
+				return err
 			}
+			defer done()
 			blockUlid, err := block.ulid.MarshalBinary()
 			if err != nil {
 				return err


Test patch:
diff --git a/bench_test.go b/bench_test.go
index 81b3934..df0adc2 100644
--- a/bench_test.go
+++ b/bench_test.go
@@ -342,7 +342,8 @@ func BenchmarkQuery(b *testing.B) {
 	// BenchmarkFilter benchmarks performing a simple filter. This is useful to get an idea of our scan speed (minus the
 	// execution engine).
 	b.Run("Filter", func(b *testing.B) {
-		table, err := db.Table(tableName, nil)
+		schema := dynparquet.SampleDefinition()
+		table, err := db.Table(tableName, NewTableConfig(schema))
 		require.NoError(b, err)
 		size := table.ActiveBlock().Size()
 		for i := 0; i < b.N; i++ {
@@ -409,7 +410,7 @@ func BenchmarkSnapshot(b *testing.B) {
 	bytesWritten := 0
 	for i := 0; i < b.N; i++ {
 		wc := &writeCounter{Writer: io.Discard}
-		require.NoError(b, WriteSnapshot(ctx, db.HighWatermark(), db, wc, false))
+		require.NoError(b, WriteSnapshot(ctx, db.HighWatermark(), db, wc))
 		bytesWritten += wc.count
 	}
 	b.ReportMetric(float64(bytesWritten)/float64(b.N), "size/op")
diff --git a/pqarrow/arrowutils/sort_test.go b/pqarrow/arrowutils/sort_test.go
index 4fa8c4a..f08ae66 100644
--- a/pqarrow/arrowutils/sort_test.go
+++ b/pqarrow/arrowutils/sort_test.go
@@ -1,7 +1,8 @@
-package arrowutils_test
+package arrowutils
 
 import (
 	"context"
+	"sort"
 	"testing"
 
 	"github.com/apache/arrow/go/v14/arrow"
@@ -9,8 +10,6 @@ import (
 	"github.com/apache/arrow/go/v14/arrow/compute"
 	"github.com/apache/arrow/go/v14/arrow/memory"
 	"github.com/stretchr/testify/require"
-
-	"github.com/polarsignals/frostdb/pqarrow/arrowutils"
 )
 
 func TestSortRecord(t *testing.T) {
@@ -30,14 +29,14 @@ func TestSortRecord(t *testing.T) {
 		{
 			Name:    "No Nows",
 			Samples: Samples{},
-			Columns: []arrowutils.SortingColumn{{Index: 0}},
+			Columns: []SortingColumn{{Index: 0}},
 		},
 		{
 			Name: "One Row",
 			Samples: Samples{
 				{},
 			},
-			Columns: []arrowutils.SortingColumn{
+			Columns: []SortingColumn{
 				{
 					Index: 0,
 				},
@@ -51,7 +50,7 @@ func TestSortRecord(t *testing.T) {
 				{Int: 2},
 				{Int: 1},
 			},
-			Columns: []arrowutils.SortingColumn{
+			Columns: []SortingColumn{
 				{Index: 0},
 			},
 			Indices: []int32{2, 1, 0},
@@ -64,8 +63,8 @@ func TestSortRecord(t *testing.T) {
 				{Int: 3},
 			},
 
-			Columns: []arrowutils.SortingColumn{
-				{Index: 0, Direction: arrowutils.Descending},
+			Columns: []SortingColumn{
+				{Index: 0, Direction: Descending},
 			},
 			Indices: []int32{2, 1, 0},
 		},
@@ -76,7 +75,7 @@ func TestSortRecord(t *testing.T) {
 				{Double: 2},
 				{Double: 1},
 			},
-			Columns: []arrowutils.SortingColumn{{Index: 1}},
+			Columns: []SortingColumn{{Index: 1}},
 			Indices: []int32{2, 1, 0},
 		},
 		{
@@ -86,7 +85,7 @@ func TestSortRecord(t *testing.T) {
 				{Double: 2},
 				{Double: 3},
 			},
-			Columns: []arrowutils.SortingColumn{{Index: 1, Direction: arrowutils.Descending}},
+			Columns: []SortingColumn{{Index: 1, Direction: Descending}},
 			Indices: []int32{2, 1, 0},
 		},
 		{
@@ -96,7 +95,7 @@ func TestSortRecord(t *testing.T) {
 				{String: "2"},
 				{String: "1"},
 			},
-			Columns: []arrowutils.SortingColumn{{Index: 2}},
+			Columns: []SortingColumn{{Index: 2}},
 			Indices: []int32{2, 1, 0},
 		},
 		{
@@ -106,7 +105,7 @@ func TestSortRecord(t *testing.T) {
 				{String: "2"},
 				{String: "3"},
 			},
-			Columns: []arrowutils.SortingColumn{{Index: 2, Direction: arrowutils.Descending}},
+			Columns: []SortingColumn{{Index: 2, Direction: Descending}},
 			Indices: []int32{2, 1, 0},
 		},
 		{
@@ -116,7 +115,7 @@ func TestSortRecord(t *testing.T) {
 				{Dict: "2"},
 				{Dict: "1"},
 			},
-			Columns: []arrowutils.SortingColumn{{Index: 3}},
+			Columns: []SortingColumn{{Index: 3}},
 			Indices: []int32{2, 1, 0},
 		},
 		{
@@ -126,7 +125,7 @@ func TestSortRecord(t *testing.T) {
 				{Dict: "2"},
 				{Dict: "3"},
 			},
-			Columns: []arrowutils.SortingColumn{{Index: 3, Direction: arrowutils.Descending}},
+			Columns: []SortingColumn{{Index: 3, Direction: Descending}},
 			Indices: []int32{2, 1, 0},
 		},
 		{
@@ -136,7 +135,7 @@ func TestSortRecord(t *testing.T) {
 				{},
 				{Nullable: null(1)},
 			},
-			Columns: []arrowutils.SortingColumn{{Index: 4}},
+			Columns: []SortingColumn{{Index: 4}},
 			Indices: []int32{2, 0, 1},
 		},
 		{
@@ -146,7 +145,7 @@ func TestSortRecord(t *testing.T) {
 				{},
 				{Nullable: null(1)},
 			},
-			Columns: []arrowutils.SortingColumn{{Index: 4, NullsFirst: true}},
+			Columns: []SortingColumn{{Index: 4, NullsFirst: true}},
 			Indices: []int32{0, 1, 2},
 		},
 		{
@@ -156,7 +155,7 @@ func TestSortRecord(t *testing.T) {
 				{},
 				{Nullable: null(1)},
 			},
-			Columns: []arrowutils.SortingColumn{{Index: 4, Direction: arrowutils.Descending}},
+			Columns: []SortingColumn{{Index: 4, Direction: Descending}},
 			Indices: []int32{2, 0, 1},
 		},
 		{
@@ -166,7 +165,7 @@ func TestSortRecord(t *testing.T) {
 				{},
 				{Nullable: null(1)},
 			},
-			Columns: []arrowutils.SortingColumn{{Index: 4, Direction: arrowutils.Descending, NullsFirst: true}},
+			Columns: []SortingColumn{{Index: 4, Direction: Descending, NullsFirst: true}},
 			Indices: []int32{0, 1, 2},
 		},
 		{
@@ -177,7 +176,7 @@ func TestSortRecord(t *testing.T) {
 				{String: "3", Int: 2},
 				{String: "4", Int: 1},
 			},
-			Columns: []arrowutils.SortingColumn{
+			Columns: []SortingColumn{
 				{Index: 0},
 				{Index: 2},
 			},
@@ -191,9 +190,9 @@ func TestSortRecord(t *testing.T) {
 				{String: "3", Int: 2},
 				{String: "4", Int: 1},
 			},
-			Columns: []arrowutils.SortingColumn{
-				{Index: 0, Direction: arrowutils.Ascending},
-				{Index: 2, Direction: arrowutils.Descending},
+			Columns: []SortingColumn{
+				{Index: 0, Direction: Ascending},
+				{Index: 2, Direction: Descending},
 			},
 			Indices: []int32{3, 2, 1, 0},
 		},
@@ -206,6 +205,40 @@ func TestSortRecord(t *testing.T) {
 	}
 }
 
+func TestSortRecordBuilderReuse(t *testing.T) {
+	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+
+	schema := arrow.NewSchema([]arrow.Field{{Name: "int64", Type: arrow.PrimitiveTypes.Int64}}, nil)
+
+	b1 := array.NewInt64Builder(mem)
+	b1.AppendValues([]int64{3, 2, 1}, nil)
+	arr1 := b1.NewArray()
+	r1 := array.NewRecord(schema, []arrow.Array{arr1}, 3)
+
+	ms, err := newMultiColSorter(r1, []SortingColumn{{Index: 0}})
+	require.Nil(t, err)
+	sort.Sort(ms)
+	sortedArr1 := ms.indices.NewArray().(*array.Int32)
+	require.Equal(t, []int32{2, 1, 0}, sortedArr1.Int32Values())
+	ms.Release() // usually defer
+
+	b2 := array.NewInt64Builder(mem)
+	b2.AppendValues([]int64{2, 1}, nil)
+	arr2 := b2.NewArray()
+	r2 := array.NewRecord(schema, []arrow.Array{arr2}, 2)
+
+	ms, err = newMultiColSorter(r2, []SortingColumn{{Index: 0}})
+	require.Nil(t, err)
+	sort.Sort(ms)
+	sortedArr2 := ms.indices.NewArray().(*array.Int32)
+	require.Equal(t, []int32{1, 0}, sortedArr2.Int32Values())
+	ms.Release() // usually defer
+
+	// This failed before the fix because the builder's data was reused.
+	require.Equal(t, []int32{2, 1, 0}, sortedArr1.Int32Values())
+	require.Equal(t, []int32{1, 0}, sortedArr2.Int32Values())
+}
+
 func TestReorderRecord(t *testing.T) {
 	t.Run("Without dictionary field", func(t *testing.T) {
 		mem := memory.NewGoAllocator()
@@ -225,7 +258,7 @@ func TestReorderRecord(t *testing.T) {
 		indices := array.NewInt32Builder(mem)
 		indices.AppendValues([]int32{2, 1, 0}, nil)
 		by := indices.NewInt32Array()
-		result, err := arrowutils.Take(
+		result, err := Take(
 			compute.WithAllocator(context.Background(), mem), r, by)
 		require.Nil(t, err)
 		defer result.Release()
@@ -258,7 +291,7 @@ func TestReorderRecord(t *testing.T) {
 		indices := array.NewInt32Builder(mem)
 		indices.AppendValues([]int32{2, 1, 0}, nil)
 		by := indices.NewInt32Array()
-		result, err := arrowutils.Take(
+		result, err := Take(
 			compute.WithAllocator(context.Background(), mem), r, by)
 		require.Nil(t, err)
 		defer result.Release()
@@ -337,7 +370,7 @@ func (s Samples) Record() arrow.Record {
 type SortCase struct {
 	Name    string
 	Samples Samples
-	Columns []arrowutils.SortingColumn
+	Columns []SortingColumn
 	Indices []int32
 	Error   string
 }
@@ -345,7 +378,7 @@ type SortCase struct {
 func sortAndCompare(t *testing.T, kase SortCase) {
 	t.Helper()
 
-	got, err := arrowutils.SortRecord(kase.Samples.Record(), kase.Columns)
+	got, err := SortRecord(kase.Samples.Record(), kase.Columns)
 	if kase.Error != "" {
 		require.NotNil(t, err, kase.Error)
 		return