Repository: polarsignals/frostdb (go)

Issue:
L1+ disk compaction 
It would be useful if we implemented a new compaction function that compacts parts into on disk parts, to reduce memory usage.


Discussion:

Patch:
diff --git a/db.go b/db.go
--- a/db.go
+++ b/db.go
@@ -65,7 +65,7 @@ type ColumnStore struct {
 	// splitSize is the number of new granules that are created when granules are split (default =2)
 	splitSize int
 	// indexConfig is the configuration settings for the lsm index
-	indexConfig []*index.LevelConfig
+	indexConfig []*IndexConfig
 
 	sources []DataSource
 	sinks   []DataSink
@@ -227,7 +227,7 @@ func WithStoragePath(path string) Option {
 	}
 }
 
-func WithIndexConfig(indexConfig []*index.LevelConfig) Option {
+func WithIndexConfig(indexConfig []*IndexConfig) Option {
 	return func(s *ColumnStore) error {
 		s.indexConfig = indexConfig
 		return nil
diff --git a/index/lsm.go b/index/lsm.go
--- a/index/lsm.go
+++ b/index/lsm.go
@@ -421,9 +421,7 @@ func (l *LSM) merge(level SentinelType, externalWriter func([]parts.Part) (parts
 	l.Lock()
 	defer l.Unlock()
 	for _, part := range mergeList {
-		if r := part.Record(); r != nil {
-			r.Release()
-		}
+		part.Release()
 	}
 
 	return nil
diff --git a/parts/arrow.go b/parts/arrow.go
--- a/parts/arrow.go
+++ b/parts/arrow.go
@@ -41,6 +41,8 @@ func (p *arrowPart) Record() arrow.Record {
 	return p.record
 }
 
+func (p *arrowPart) Release() { p.record.Release() }
+
 func (p *arrowPart) SerializeBuffer(schema *dynparquet.Schema, w dynparquet.ParquetWriter) error {
 	return pqarrow.RecordToFile(schema, w, p.record)
 }
diff --git a/parts/parquet.go b/parts/parquet.go
--- a/parts/parquet.go
+++ b/parts/parquet.go
@@ -20,6 +20,12 @@ func (p *parquetPart) Record() arrow.Record {
 	return nil
 }
 
+func (p *parquetPart) Release() {
+	if p.release != nil {
+		p.release()
+	}
+}
+
 func (p *parquetPart) SerializeBuffer(_ *dynparquet.Schema, _ dynparquet.ParquetWriter) error {
 	return fmt.Errorf("not a record part")
 }
@@ -57,19 +63,12 @@ func (p *parquetPart) Least() (*dynparquet.DynamicRow, error) {
 		return p.minRow, nil
 	}
 
-	rowBuf := &dynparquet.DynamicRows{Rows: make([]parquet.Row, 1)}
-	reader := p.buf.DynamicRowGroup(0).DynamicRows()
-	defer reader.Close()
-
-	if n, err := reader.ReadRows(rowBuf); err != nil {
-		return nil, fmt.Errorf("read first row of part: %w", err)
-	} else if n != 1 {
-		return nil, fmt.Errorf("expected to read exactly 1 row, but read %d", n)
+	minRow, err := min(p.buf)
+	if err != nil {
+		return nil, err
 	}
 
-	// Copy here so that this reference does not prevent the decompressed page
-	// from being GCed.
-	p.minRow = rowBuf.GetCopy(0)
+	p.minRow = minRow
 	return p.minRow, nil
 }
 
@@ -78,24 +77,11 @@ func (p *parquetPart) Most() (*dynparquet.DynamicRow, error) {
 		return p.maxRow, nil
 	}
 
-	rowBuf := &dynparquet.DynamicRows{Rows: make([]parquet.Row, 1)}
-	rg := p.buf.DynamicRowGroup(p.buf.NumRowGroups() - 1)
-	reader := rg.DynamicRows()
-	defer reader.Close()
-
-	if err := reader.SeekToRow(rg.NumRows() - 1); err != nil {
-		return nil, fmt.Errorf("seek to last row of part: %w", err)
-	}
-
-	if n, err := reader.ReadRows(rowBuf); err != nil {
-		return nil, fmt.Errorf("read last row of part: %w", err)
-	} else if n != 1 {
-		return nil, fmt.Errorf("expected to read exactly 1 row, but read %d", n)
+	maxRow, err := max(p.buf)
+	if err != nil {
+		return nil, err
 	}
-
-	// Copy here so that this reference does not prevent the decompressed page
-	// from being GCed.
-	p.maxRow = rowBuf.GetCopy(0)
+	p.maxRow = maxRow
 	return p.maxRow, nil
 }
 
@@ -119,3 +105,40 @@ func (p *parquetPart) OverlapsWith(schema *dynparquet.Schema, otherPart Part) (b
 
 	return schema.Cmp(a, d) <= 0 && schema.Cmp(c, b) <= 0, nil
 }
+
+func max(buf *dynparquet.SerializedBuffer) (*dynparquet.DynamicRow, error) {
+	rowBuf := &dynparquet.DynamicRows{Rows: make([]parquet.Row, 1)}
+	rg := buf.DynamicRowGroup(buf.NumRowGroups() - 1)
+	reader := rg.DynamicRows()
+	defer reader.Close()
+
+	if err := reader.SeekToRow(rg.NumRows() - 1); err != nil {
+		return nil, fmt.Errorf("seek to last row of part: %w", err)
+	}
+
+	if n, err := reader.ReadRows(rowBuf); err != nil {
+		return nil, fmt.Errorf("read last row of part: %w", err)
+	} else if n != 1 {
+		return nil, fmt.Errorf("expected to read exactly 1 row, but read %d", n)
+	}
+
+	// Copy here so that this reference does not prevent the decompressed page
+	// from being GCed.
+	return rowBuf.GetCopy(0), nil
+}
+
+func min(buf *dynparquet.SerializedBuffer) (*dynparquet.DynamicRow, error) {
+	rowBuf := &dynparquet.DynamicRows{Rows: make([]parquet.Row, 1)}
+	reader := buf.DynamicRowGroup(0).DynamicRows()
+	defer reader.Close()
+
+	if n, err := reader.ReadRows(rowBuf); err != nil {
+		return nil, fmt.Errorf("read first row of part: %w", err)
+	} else if n != 1 {
+		return nil, fmt.Errorf("expected to read exactly 1 row, but read %d", n)
+	}
+
+	// Copy here so that this reference does not prevent the decompressed page
+	// from being GCed.
+	return rowBuf.GetCopy(0), nil
+}
diff --git a/parts/part.go b/parts/part.go
--- a/parts/part.go
+++ b/parts/part.go
@@ -12,6 +12,7 @@ type Part interface {
 	// Record returns the Arrow record for the part. If the part is not an Arrow
 	// record part, nil is returned.
 	Record() arrow.Record
+	Release()
 	SerializeBuffer(schema *dynparquet.Schema, w dynparquet.ParquetWriter) error
 	AsSerializedBuffer(schema *dynparquet.Schema) (*dynparquet.SerializedBuffer, error)
 	NumRows() int64
@@ -28,6 +29,7 @@ type basePart struct {
 	compactionLevel int
 	minRow          *dynparquet.DynamicRow
 	maxRow          *dynparquet.DynamicRow
+	release         func()
 }
 
 func (p *basePart) CompactionLevel() int {
@@ -44,6 +46,12 @@ func WithCompactionLevel(level int) Option {
 	}
 }
 
+func WithRelease(release func()) Option {
+	return func(p *basePart) {
+		p.release = release
+	}
+}
+
 type PartSorter struct {
 	schema *dynparquet.Schema
 	parts  []Part
diff --git a/table.go b/table.go
--- a/table.go
+++ b/table.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"math/rand"
+	"os"
 	"path/filepath"
 	"runtime"
 	"sort"
@@ -52,11 +53,11 @@ var (
 
 // DefaultIndexConfig returns the default level configs used. This is a function
 // So that any modifications to the result will not affect the default config.
-func DefaultIndexConfig() []*index.LevelConfig {
-	return []*index.LevelConfig{
-		{Level: index.L0, MaxSize: 1024 * 1024 * 15},  // Compact to Parquet after 15MiB of data
-		{Level: index.L1, MaxSize: 1024 * 1024 * 128}, // Compact to a single Parquet after 128MiB of Parquet files
-		{Level: index.L2, MaxSize: 1024 * 1024 * 512}, // Final level. Rotate after 512MiB of Parquet files
+func DefaultIndexConfig() []*IndexConfig {
+	return []*IndexConfig{
+		{Level: int(index.L0), MaxSize: 1024 * 1024 * 15, Type: CompactionTypeParquet},  // Compact to in-memory Parquet buffer after 15MiB of data
+		{Level: int(index.L1), MaxSize: 1024 * 1024 * 128, Type: CompactionTypeParquet}, // Compact to a single in-memory Parquet buffer after 128MiB of Parquet files
+		{Level: int(index.L2), MaxSize: 1024 * 1024 * 512},                              // Final level. Rotate after 512MiB of Parquet files
 	}
 }
 
@@ -169,6 +170,7 @@ type Table struct {
 
 	wal     WAL
 	closing bool
+	closers []io.Closer
 }
 
 type WAL interface {
@@ -1092,23 +1094,52 @@ func (t *Table) close() {
 	t.active.pendingWritersWg.Wait()
 	t.closing = true
 	t.active.index.WaitForPendingCompactions()
+
+	for _, closer := range t.closers {
+		if err := closer.Close(); err != nil {
+			level.Error(t.logger).Log("msg", "table closer", "err", err)
+		}
+	}
+}
+
+type CompactionType int
+
+const (
+	CompactionTypeUnknown CompactionType = iota
+	CompactionTypeParquet
+	CompactionTypeParquetDisk
+)
+
+type IndexConfig struct {
+	Level   int
+	MaxSize int64
+	Type    CompactionType
 }
 
 // configureLSMLevels configures the level configs for this table.
-func (t *Table) configureLSMLevels(levels []*index.LevelConfig) []*index.LevelConfig {
+func (t *Table) configureLSMLevels(levels []*IndexConfig) []*index.LevelConfig {
 	config := make([]*index.LevelConfig, 0, len(levels))
 
-	for _, level := range levels {
-		config = append(config, &index.LevelConfig{
-			Level:   level.Level,
+	for i, level := range levels {
+		cfg := &index.LevelConfig{
+			Level:   index.SentinelType(level.Level),
 			MaxSize: level.MaxSize,
-			Compact: t.parquetCompaction, // NOTE: right now we only support parquet compaction but in the future this could compact parts into larger arrow records.
-		})
+		}
+		switch level.Type {
+		case CompactionTypeParquet:
+			cfg.Compact = t.parquetCompaction
+		case CompactionTypeParquetDisk:
+			fileCompaction := t.parquetFileCompaction(i + 1)
+			t.closers = append(t.closers, fileCompaction) // Append to closers so that the underlying files are closed on table close.
+			cfg.Compact = fileCompaction.writeRecordsToParquetFile
+		default:
+			if i != len(levels)-1 { // Compaction type should not be set for last level
+				panic(fmt.Sprintf("unknown compaction type: %v", level.Type))
+			}
+		}
+		config = append(config, cfg)
 	}
 
-	// Don't set the compaction functino for the last level. As the LSM doesn't have a level after the last one.
-	config[len(levels)-1].Compact = nil
-
 	return config
 }
 
@@ -1379,3 +1410,86 @@ func (t *Table) distinctRecordsForCompaction(compact []parts.Part) ([]arrow.Reco
 	}
 	return newRecords, nil
 }
+
+type fileCompaction struct {
+	t      *Table
+	file   *os.File
+	offset int64 // Writing offsets into the file
+	ref    int64 // Number of references to file.
+}
+
+func (f *fileCompaction) Close() error {
+	return f.file.Close()
+}
+
+func (t *Table) parquetFileCompaction(lvl int) *fileCompaction {
+	f := &fileCompaction{
+		t: t,
+	}
+
+	file, err := os.CreateTemp("", fmt.Sprintf("L%v-*.parquet", lvl))
+	if err != nil {
+		panic(err)
+	}
+	level.Info(t.logger).Log("msg", "created temp file for level", "level", lvl, "file", file.Name())
+	f.file = file
+
+	return f
+}
+
+// accountingWriter is a writer that accounts for the number of bytes written.
+type accountingWriter struct {
+	w io.Writer
+	n int64
+}
+
+func (a *accountingWriter) Write(p []byte) (int, error) {
+	n, err := a.w.Write(p)
+	a.n += int64(n)
+	return n, err
+}
+
+// writeRecordsToParquetFile will compact the given parts into a Parquet file written to the next level file.
+func (f *fileCompaction) writeRecordsToParquetFile(compact []parts.Part, options ...parts.Option) ([]parts.Part, int64, int64, error) {
+	accountant := &accountingWriter{w: f.file}
+	preCompactionSize, err := f.t.compactParts(accountant, compact) // compact into the next level
+	if err != nil {
+		return nil, 0, 0, err
+	}
+
+	// Record the writing offset into the file.
+	prevOffset := f.offset
+	f.offset += accountant.n
+	f.ref++
+
+	pf, err := parquet.OpenFile(io.NewSectionReader(f.file, prevOffset, accountant.n), accountant.n)
+	if err != nil {
+		return nil, 0, 0, err
+	}
+
+	buf, err := dynparquet.NewSerializedBuffer(pf)
+	if err != nil {
+		return nil, 0, 0, err
+	}
+
+	return []parts.Part{parts.NewParquetPart(0, buf, append(options, parts.WithRelease(f.release()))...)}, preCompactionSize, accountant.n, nil
+}
+
+// release will account for all the Parts currently pointing to this file. Once the last one has been released it will truncate the file.
+// release is not safe to call concurrently.
+func (f *fileCompaction) release() func() {
+	return func() {
+		f.ref--
+		if f.ref == 0 {
+			level.Info(f.t.logger).Log("msg", "truncating file", "file", f.file.Name())
+			if err := os.Truncate(f.file.Name(), 0); err != nil {
+				panic(fmt.Errorf("failed to truncate the level file: %v", err))
+			}
+
+			f.offset = 0
+			if _, err := f.file.Seek(0, io.SeekStart); err != nil {
+				panic(fmt.Errorf("failed to seek the level file: %v", err))
+			}
+		}
+	}
+}


Test patch:
diff --git a/db_test.go b/db_test.go
--- a/db_test.go
+++ b/db_test.go
@@ -2250,3 +2250,56 @@ func TestDBConcurrentOpen(t *testing.T) {
 	}
 	require.NoError(t, errg.Wait())
 }
+
+func Test_DB_WithParquetDiskCompaction(t *testing.T) {
+	config := NewTableConfig(
+		dynparquet.SampleDefinition(),
+	)
+
+	logger := newTestLogger(t)
+
+	cfg := DefaultIndexConfig()
+	cfg[0].Type = CompactionTypeParquetDisk // Create disk compaction
+	cfg[1].Type = CompactionTypeParquetDisk
+	c, err := New(
+		WithLogger(logger),
+		WithIndexConfig(cfg),
+	)
+	t.Cleanup(func() {
+		require.NoError(t, c.Close())
+	})
+	require.NoError(t, err)
+	db, err := c.DB(context.Background(), "test")
+	require.NoError(t, err)
+	table, err := db.Table("test", config)
+	require.NoError(t, err)
+
+	samples := dynparquet.NewTestSamples()
+
+	ctx := context.Background()
+	for i := 0; i < 100; i++ {
+		r, err := samples.ToRecord()
+		require.NoError(t, err)
+		_, err = table.InsertRecord(ctx, r)
+		require.NoError(t, err)
+	}
+	require.NoError(t, table.EnsureCompaction())
+
+	// Ensure that disk compacted data can be recovered
+	pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
+	defer pool.AssertSize(t, 0)
+	rows := int64(0)
+	err = table.View(ctx, func(ctx context.Context, tx uint64) error {
+		return table.Iterator(
+			ctx,
+			tx,
+			pool,
+			[]logicalplan.Callback{func(ctx context.Context, ar arrow.Record) error {
+				rows += ar.NumRows()
+				return nil
+			}},
+		)
+	})
+	require.NoError(t, err)
+	require.Equal(t, int64(300), rows)
+}
diff --git a/table_test.go b/table_test.go
--- a/table_test.go
+++ b/table_test.go
@@ -451,9 +451,9 @@ func Test_Table_NewTableValidSplitSize(t *testing.T) {
 
 func Test_Table_Bloomfilter(t *testing.T) {
 	c, table := basicTable(t, WithIndexConfig(
-		[]*index.LevelConfig{
-			{Level: index.L0, MaxSize: 452}, // NOTE: 452 is the current size of the 3 records that are inserted
-			{Level: index.L1, MaxSize: 100000},
+		[]*IndexConfig{
+			{Level: int(index.L0), MaxSize: 452, Type: CompactionTypeParquet}, // NOTE: 452 is the current size of the 3 records that are inserted
+			{Level: int(index.L1), MaxSize: 100000},
 		},
 	))
 	defer c.Close()
@@ -1136,9 +1136,9 @@ func Test_Table_DynamicColumnNotDefined(t *testing.T) {
 }
 
 func TestTableUniquePrimaryIndex(t *testing.T) {
-	c, err := New(WithGranuleSizeBytes(1), WithIndexConfig([]*index.LevelConfig{
-		{Level: index.L0, MaxSize: 180},
-		{Level: index.L1, MaxSize: 1 * TiB},
+	c, err := New(WithGranuleSizeBytes(1), WithIndexConfig([]*IndexConfig{
+		{Level: int(index.L0), MaxSize: 180, Type: CompactionTypeParquet},
+		{Level: int(index.L1), MaxSize: 1 * TiB},
 	}))
 	require.NoError(t, err)
 	defer c.Close()