Repository: polarsignals/frostdb
Pull request #555 (2023-10-09 17:21:38)

lsm: store the part compaction level
This is useful when replaying a snapshot to replay the part into the correct level of the lsm


Patch:
diff --git a/index/lsm.go b/index/lsm.go
index 2ed79eb..e9eaba1 100644
--- a/index/lsm.go
+++ b/index/lsm.go
@@ -53,7 +53,7 @@ type LSMMetrics struct {
 type LevelConfig struct {
 	Level   SentinelType
 	MaxSize int64
-	Compact func([]*parts.Part) ([]*parts.Part, int64, int64, error)
+	Compact func([]*parts.Part, ...parts.Option) ([]*parts.Part, int64, int64, error)
 }
 
 type LSMOption func(*LSM)
@@ -158,7 +158,7 @@ func (l *LSM) MaxLevel() SentinelType {
 func (l *LSM) Add(tx uint64, record arrow.Record) {
 	record.Retain()
 	size := util.TotalRecordSize(record)
-	l.levels.Prepend(parts.NewArrowPart(tx, record, int(size), nil))
+	l.levels.Prepend(parts.NewArrowPart(tx, record, int(size), nil, parts.WithCompactionLevel(int(L0))))
 	l0 := l.sizes[L0].Add(int64(size))
 	l.metrics.LevelSize.WithLabelValues(L0.String()).Set(float64(l0))
 	if l0 >= l.configs[L0].MaxSize {
@@ -360,7 +360,7 @@ func (l *LSM) merge(level SentinelType, externalWriter func([]*parts.Part) (*par
 			s.next.Store(next)
 		}
 	} else {
-		compacted, size, compactedSize, err = l.configs[level].Compact(mergeList)
+		compacted, size, compactedSize, err = l.configs[level].Compact(mergeList, parts.WithCompactionLevel(int(level)+1))
 		if err != nil {
 			return err
 		}
diff --git a/parts/part.go b/parts/part.go
index adf61f5..6df7e30 100644
--- a/parts/part.go
+++ b/parts/part.go
@@ -13,18 +13,6 @@ import (
 	"github.com/polarsignals/frostdb/pqarrow"
 )
 
-type CompactionLevel uint8
-
-const (
-	// CompactionLevel0 is the default compaction level for new Parts. This
-	// means that the Part contains multiple variable-length row groups.
-	CompactionLevel0 CompactionLevel = iota
-	// CompactionLevel1 is the compaction level for Parts that are the result of
-	// a compaction. Parts with this compaction level contain multiple row
-	// groups with the row group size specified on table creation.
-	CompactionLevel1
-)
-
 type Part struct {
 	buf    *dynparquet.SerializedBuffer
 	record arrow.Record
@@ -36,7 +24,7 @@ type Part struct {
 	// tx is the id of the transaction that created this part.
 	tx uint64
 
-	compactionLevel CompactionLevel
+	compactionLevel int
 
 	minRow *dynparquet.DynamicRow
 	maxRow *dynparquet.DynamicRow
@@ -86,7 +74,7 @@ func (p *Part) AsSerializedBuffer(schema *dynparquet.Schema) (*dynparquet.Serial
 
 type Option func(*Part)
 
-func WithCompactionLevel(level CompactionLevel) Option {
+func WithCompactionLevel(level int) Option {
 	return func(p *Part) {
 		p.compactionLevel = level
 	}
@@ -137,7 +125,7 @@ func (p *Part) Size() int64 {
 	return int64(p.recordRelativeSize)
 }
 
-func (p *Part) CompactionLevel() CompactionLevel {
+func (p *Part) CompactionLevel() int {
 	return p.compactionLevel
 }
 
diff --git a/snapshot.go b/snapshot.go
index 7a4e559..853c961 100644
--- a/snapshot.go
+++ b/snapshot.go
@@ -621,7 +621,7 @@ func loadSnapshot(ctx context.Context, db *DB, r io.ReaderAt, size int64) ([]byt
 					if _, err := r.ReadAt(partBytes, startOffset); err != nil {
 						return err
 					}
-					partOptions := parts.WithCompactionLevel(parts.CompactionLevel(partMeta.CompactionLevel))
+					partOptions := parts.WithCompactionLevel(int(partMeta.CompactionLevel))
 					switch partMeta.Encoding {
 					case snapshotpb.Part_ENCODING_PARQUET:
 						serBuf, err := dynparquet.ReaderFromBytes(partBytes)
@@ -647,11 +647,7 @@ func loadSnapshot(ctx context.Context, db *DB, r io.ReaderAt, size int64) ([]byt
 				}
 
 				for _, part := range resultParts {
-					if part.Record() != nil { // TODO: eventually store the level in the part metadata
-						newIdx.InsertPart(index.L0, part)
-					} else {
-						newIdx.InsertPart(index.L1, part)
-					}
+					newIdx.InsertPart(index.SentinelType(part.CompactionLevel()), part)
 				}
 			}
 
diff --git a/table.go b/table.go
index 1bbc866..a241fda 100644
--- a/table.go
+++ b/table.go
@@ -1085,7 +1085,7 @@ func (t *Table) configureLSMLevels(levels []*index.LevelConfig) []*index.LevelCo
 	return config
 }
 
-func (t *Table) parquetCompaction(compact []*parts.Part) ([]*parts.Part, int64, int64, error) {
+func (t *Table) parquetCompaction(compact []*parts.Part, options ...parts.Option) ([]*parts.Part, int64, int64, error) {
 	b := &bytes.Buffer{}
 	size, err := t.compactParts(b, compact)
 	if err != nil {
@@ -1096,7 +1096,7 @@ func (t *Table) parquetCompaction(compact []*parts.Part) ([]*parts.Part, int64,
 	if err != nil {
 		return nil, 0, 0, err
 	}
-	return []*parts.Part{parts.NewPart(0, buf)}, size, int64(b.Len()), nil
+	return []*parts.Part{parts.NewPart(0, buf, options...)}, size, int64(b.Len()), nil
 }
 
 func (t *Table) externalParquetCompaction(writer io.Writer) func(compact []*parts.Part) (*parts.Part, int64, int64, error) {


Test patch:
diff --git a/index/lsm_test.go b/index/lsm_test.go
index a3c84da..b855859 100644
--- a/index/lsm_test.go
+++ b/index/lsm_test.go
@@ -16,7 +16,7 @@ import (
 	"github.com/polarsignals/frostdb/parts"
 )
 
-func parquetCompaction(compact []*parts.Part) ([]*parts.Part, int64, int64, error) {
+func parquetCompaction(compact []*parts.Part, options ...parts.Option) ([]*parts.Part, int64, int64, error) {
 	b := &bytes.Buffer{}
 	size, err := compactParts(b, compact)
 	if err != nil {