Repository: polarsignals/frostdb
Pull request #608 (2023-12-04 18:29:50)

Split arrow part and parquet part
Now that Part is an interface we can split the Parquet buffer and Arrow record implementations into two distinct implementations 


Patch:
diff --git a/index/lsm.go b/index/lsm.go
index 3f824e5..d8a8369 100644
--- a/index/lsm.go
+++ b/index/lsm.go
@@ -167,7 +167,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), l.schema, parts.WithCompactionLevel(int(L0))))
+	l.levels.Prepend(parts.NewArrowPart(tx, record, uint64(size), l.schema, 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 {
diff --git a/parts/arrow.go b/parts/arrow.go
new file mode 100644
index 0000000..52763de
--- /dev/null
+++ b/parts/arrow.go
@@ -0,0 +1,140 @@
+package parts
+
+import (
+	"bytes"
+
+	"github.com/apache/arrow/go/v14/arrow"
+	"github.com/parquet-go/parquet-go"
+
+	"github.com/polarsignals/frostdb/dynparquet"
+	"github.com/polarsignals/frostdb/pqarrow"
+)
+
+// arrow implments the Part interface backed by an Arrow record.
+type arrowPart struct {
+	basePart
+
+	schema *dynparquet.Schema
+	record arrow.Record
+	size   uint64
+}
+
+// NewArrowPart returns a new Arrow part.
+func NewArrowPart(tx uint64, record arrow.Record, size uint64, schema *dynparquet.Schema, options ...Option) Part {
+	p := &arrowPart{
+		basePart: basePart{
+			tx: tx,
+		},
+		schema: schema,
+		record: record,
+		size:   size,
+	}
+
+	for _, option := range options {
+		option(&p.basePart)
+	}
+
+	return p
+}
+
+func (p *arrowPart) Record() arrow.Record {
+	return p.record
+}
+
+func (p *arrowPart) SerializeBuffer(schema *dynparquet.Schema, w dynparquet.ParquetWriter) error {
+	return pqarrow.RecordToFile(schema, w, p.record)
+}
+
+func (p *arrowPart) AsSerializedBuffer(schema *dynparquet.Schema) (*dynparquet.SerializedBuffer, error) {
+	// If this is a Arrow record part, convert the record into a serialized buffer
+	b := &bytes.Buffer{}
+
+	w, err := schema.GetWriter(b, pqarrow.RecordDynamicCols(p.record), false)
+	if err != nil {
+		return nil, err
+	}
+	defer schema.PutWriter(w)
+	if err := p.SerializeBuffer(schema, w.ParquetWriter); err != nil {
+		return nil, err
+	}
+
+	f, err := parquet.OpenFile(bytes.NewReader(b.Bytes()), int64(b.Len()))
+	if err != nil {
+		return nil, err
+	}
+
+	buf, err := dynparquet.NewSerializedBuffer(f)
+	if err != nil {
+		return nil, err
+	}
+
+	return buf, nil
+}
+
+func (p *arrowPart) NumRows() int64 {
+	return p.record.NumRows()
+}
+
+func (p *arrowPart) Size() int64 {
+	return int64(p.size)
+}
+
+// Least returns the least row  in the part.
+func (p *arrowPart) Least() (*dynparquet.DynamicRow, error) {
+	if p.minRow != nil {
+		return p.minRow, nil
+	}
+
+	dynCols := pqarrow.RecordDynamicCols(p.record)
+	pooledSchema, err := p.schema.GetDynamicParquetSchema(dynCols)
+	if err != nil {
+		return nil, err
+	}
+	defer p.schema.PutPooledParquetSchema(pooledSchema)
+	p.minRow, err = pqarrow.RecordToDynamicRow(p.schema, pooledSchema.Schema, p.record, dynCols, 0)
+	if err != nil {
+		return nil, err
+	}
+
+	return p.minRow, nil
+}
+
+func (p *arrowPart) Most() (*dynparquet.DynamicRow, error) {
+	if p.maxRow != nil {
+		return p.maxRow, nil
+	}
+
+	dynCols := pqarrow.RecordDynamicCols(p.record)
+	pooledSchema, err := p.schema.GetDynamicParquetSchema(dynCols)
+	if err != nil {
+		return nil, err
+	}
+	defer p.schema.PutPooledParquetSchema(pooledSchema)
+	p.maxRow, err = pqarrow.RecordToDynamicRow(p.schema, pooledSchema.Schema, p.record, dynCols, int(p.record.NumRows()-1))
+	if err != nil {
+		return nil, err
+	}
+
+	return p.maxRow, nil
+}
+
+func (p *arrowPart) OverlapsWith(schema *dynparquet.Schema, otherPart Part) (bool, error) {
+	a, err := p.Least()
+	if err != nil {
+		return false, err
+	}
+	b, err := p.Most()
+	if err != nil {
+		return false, err
+	}
+	c, err := otherPart.Least()
+	if err != nil {
+		return false, err
+	}
+	d, err := otherPart.Most()
+	if err != nil {
+		return false, err
+	}
+
+	return schema.Cmp(a, d) <= 0 && schema.Cmp(c, b) <= 0, nil
+}
diff --git a/parts/parquet.go b/parts/parquet.go
new file mode 100644
index 0000000..3262692
--- /dev/null
+++ b/parts/parquet.go
@@ -0,0 +1,121 @@
+package parts
+
+import (
+	"fmt"
+
+	"github.com/apache/arrow/go/v14/arrow"
+	"github.com/parquet-go/parquet-go"
+
+	"github.com/polarsignals/frostdb/dynparquet"
+)
+
+// This file contains the implementation of the Part interface backed by a Parquet Buffer.
+type parquetPart struct {
+	basePart
+
+	buf *dynparquet.SerializedBuffer
+}
+
+func (p *parquetPart) Record() arrow.Record {
+	return nil
+}
+
+func (p *parquetPart) SerializeBuffer(_ *dynparquet.Schema, _ dynparquet.ParquetWriter) error {
+	return fmt.Errorf("not a record part")
+}
+
+func (p *parquetPart) AsSerializedBuffer(_ *dynparquet.Schema) (*dynparquet.SerializedBuffer, error) {
+	return p.buf, nil
+}
+
+func NewParquetPart(tx uint64, buf *dynparquet.SerializedBuffer, options ...Option) Part {
+	p := &parquetPart{
+		basePart: basePart{
+			tx: tx,
+		},
+		buf: buf,
+	}
+
+	for _, opt := range options {
+		opt(&p.basePart)
+	}
+
+	return p
+}
+
+func (p *parquetPart) NumRows() int64 {
+	return p.buf.NumRows()
+}
+
+func (p *parquetPart) Size() int64 {
+	return p.buf.ParquetFile().Size()
+}
+
+// Least returns the least row  in the part.
+func (p *parquetPart) Least() (*dynparquet.DynamicRow, error) {
+	if p.minRow != nil {
+		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)
+	}
+
+	// Copy here so that this reference does not prevent the decompressed page
+	// from being GCed.
+	p.minRow = rowBuf.GetCopy(0)
+	return p.minRow, nil
+}
+
+func (p *parquetPart) Most() (*dynparquet.DynamicRow, error) {
+	if p.maxRow != nil {
+		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)
+	}
+
+	// Copy here so that this reference does not prevent the decompressed page
+	// from being GCed.
+	p.maxRow = rowBuf.GetCopy(0)
+	return p.maxRow, nil
+}
+
+func (p *parquetPart) OverlapsWith(schema *dynparquet.Schema, otherPart Part) (bool, error) {
+	a, err := p.Least()
+	if err != nil {
+		return false, err
+	}
+	b, err := p.Most()
+	if err != nil {
+		return false, err
+	}
+	c, err := otherPart.Least()
+	if err != nil {
+		return false, err
+	}
+	d, err := otherPart.Most()
+	if err != nil {
+		return false, err
+	}
+
+	return schema.Cmp(a, d) <= 0 && schema.Cmp(c, b) <= 0, nil
+}
diff --git a/parts/part.go b/parts/part.go
index 65b66fc..b6ef96e 100644
--- a/parts/part.go
+++ b/parts/part.go
@@ -1,16 +1,11 @@
 package parts
 
 import (
-	"bytes"
-	"fmt"
-	"math"
 	"sort"
 
 	"github.com/apache/arrow/go/v14/arrow"
-	"github.com/parquet-go/parquet-go"
 
 	"github.com/polarsignals/frostdb/dynparquet"
-	"github.com/polarsignals/frostdb/pqarrow"
 )
 
 type Part interface {
@@ -28,237 +23,27 @@ type Part interface {
 	OverlapsWith(schema *dynparquet.Schema, otherPart Part) (bool, error)
 }
 
-type part struct {
-	buf    *dynparquet.SerializedBuffer
-	record arrow.Record
-	// record relative size is how many bytes are roughly attributed to this record.
-	// It's considered rough because it may be sharing underlying dictionaries with other records.
-	recordRelativeSize int
-	schema             *dynparquet.Schema
-
-	// tx is the id of the transaction that created this part.
-	tx uint64
-
+type basePart struct {
+	tx              uint64
 	compactionLevel int
-
-	minRow *dynparquet.DynamicRow
-	maxRow *dynparquet.DynamicRow
+	minRow          *dynparquet.DynamicRow
+	maxRow          *dynparquet.DynamicRow
 }
 
-func (p *part) Record() arrow.Record {
-	return p.record
-}
-
-func (p *part) SerializeBuffer(schema *dynparquet.Schema, w dynparquet.ParquetWriter) error {
-	if p.record == nil {
-		return fmt.Errorf("not a record part")
-	}
-
-	return pqarrow.RecordToFile(schema, w, p.record)
+func (p *basePart) CompactionLevel() int {
+	return p.compactionLevel
 }
 
-func (p *part) AsSerializedBuffer(schema *dynparquet.Schema) (*dynparquet.SerializedBuffer, error) {
-	if p.buf != nil {
-		return p.buf, nil
-	}
+func (p *basePart) TX() uint64 { return p.tx }
 
-	// If this is a Arrow record part, convert the record into a serialized buffer
-	b := &bytes.Buffer{}
-
-	w, err := schema.GetWriter(b, pqarrow.RecordDynamicCols(p.record), false)
-	if err != nil {
-		return nil, err
-	}
-	defer schema.PutWriter(w)
-	if err := p.SerializeBuffer(schema, w.ParquetWriter); err != nil {
-		return nil, err
-	}
-
-	f, err := parquet.OpenFile(bytes.NewReader(b.Bytes()), int64(b.Len()))
-	if err != nil {
-		return nil, err
-	}
-
-	buf, err := dynparquet.NewSerializedBuffer(f)
-	if err != nil {
-		return nil, err
-	}
-
-	return buf, nil
-}
-
-type Option func(*part)
+type Option func(*basePart)
 
 func WithCompactionLevel(level int) Option {
-	return func(p *part) {
+	return func(p *basePart) {
 		p.compactionLevel = level
 	}
 }
 
-// NewArrowPart returns a new Arrow part.
-func NewArrowPart(tx uint64, record arrow.Record, size int, schema *dynparquet.Schema, options ...Option) Part {
-	p := &part{
-		tx:                 tx,
-		record:             record,
-		recordRelativeSize: size,
-		schema:             schema,
-	}
-
-	for _, opt := range options {
-		opt(p)
-	}
-
-	return p
-}
-
-func NewPart(tx uint64, buf *dynparquet.SerializedBuffer, options ...Option) Part {
-	p := &part{
-		tx:  tx,
-		buf: buf,
-	}
-
-	for _, opt := range options {
-		opt(p)
-	}
-
-	return p
-}
-
-func (p *part) NumRows() int64 {
-	if p.buf != nil {
-		return p.buf.NumRows()
-	}
-
-	return p.record.NumRows()
-}
-
-func (p *part) Size() int64 {
-	if p.buf != nil {
-		return p.buf.ParquetFile().Size()
-	}
-
-	return int64(p.recordRelativeSize)
-}
-
-func (p *part) CompactionLevel() int {
-	return p.compactionLevel
-}
-
-// TX returns the transaction id for the part.
-func (p *part) TX() uint64 { return p.tx }
-
-// Least returns the least row  in the part.
-func (p *part) Least() (*dynparquet.DynamicRow, error) {
-	if p.minRow != nil {
-		return p.minRow, nil
-	}
-
-	if p.record != nil {
-		dynCols := pqarrow.RecordDynamicCols(p.record)
-		pooledSchema, err := p.schema.GetDynamicParquetSchema(dynCols)
-		if err != nil {
-			return nil, err
-		}
-		defer p.schema.PutPooledParquetSchema(pooledSchema)
-		p.minRow, err = pqarrow.RecordToDynamicRow(p.schema, pooledSchema.Schema, p.record, dynCols, 0)
-		if err != nil {
-			return nil, err
-		}
-
-		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)
-	}
-
-	// Copy here so that this reference does not prevent the decompressed page
-	// from being GCed.
-	p.minRow = rowBuf.GetCopy(0)
-	return p.minRow, nil
-}
-
-func (p *part) Most() (*dynparquet.DynamicRow, error) {
-	if p.maxRow != nil {
-		return p.maxRow, nil
-	}
-
-	if p.record != nil {
-		dynCols := pqarrow.RecordDynamicCols(p.record)
-		pooledSchema, err := p.schema.GetDynamicParquetSchema(dynCols)
-		if err != nil {
-			return nil, err
-		}
-		defer p.schema.PutPooledParquetSchema(pooledSchema)
-		p.maxRow, err = pqarrow.RecordToDynamicRow(p.schema, pooledSchema.Schema, p.record, dynCols, int(p.record.NumRows()-1))
-		if err != nil {
-			return nil, err
-		}
-
-		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)
-	}
-
-	// Copy here so that this reference does not prevent the decompressed page
-	// from being GCed.
-	p.maxRow = rowBuf.GetCopy(0)
-	return p.maxRow, nil
-}
-
-func (p *part) OverlapsWith(schema *dynparquet.Schema, otherPart Part) (bool, error) {
-	a, err := p.Least()
-	if err != nil {
-		return false, err
-	}
-	b, err := p.Most()
-	if err != nil {
-		return false, err
-	}
-	c, err := otherPart.Least()
-	if err != nil {
-		return false, err
-	}
-	d, err := otherPart.Most()
-	if err != nil {
-		return false, err
-	}
-
-	return schema.Cmp(a, d) <= 0 && schema.Cmp(c, b) <= 0, nil
-}
-
-// Tombstone marks all the parts with the max tx id to ensure they aren't
-// included in reads. Tombstoned parts will be eventually be dropped from the
-// database during compaction.
-func Tombstone(parts []*part) {
-	for _, part := range parts {
-		part.tx = math.MaxUint64
-	}
-}
-
-func (p *part) HasTombstone() bool {
-	return p.tx == math.MaxUint64
-}
-
 type PartSorter struct {
 	schema *dynparquet.Schema
 	parts  []Part
diff --git a/snapshot.go b/snapshot.go
index aaf4e45..a5ae91c 100644
--- a/snapshot.go
+++ b/snapshot.go
@@ -632,7 +632,7 @@ func loadSnapshot(ctx context.Context, db *DB, r io.ReaderAt, size int64) ([]byt
 						if err != nil {
 							return err
 						}
-						resultParts = append(resultParts, parts.NewPart(partMeta.Tx, serBuf, partOptions))
+						resultParts = append(resultParts, parts.NewParquetPart(partMeta.Tx, serBuf, partOptions))
 					case snapshotpb.Part_ENCODING_ARROW:
 						if err := func() error {
 							arrowReader, err := ipc.NewReader(bytes.NewReader(partBytes))
@@ -649,7 +649,7 @@ func loadSnapshot(ctx context.Context, db *DB, r io.ReaderAt, size int64) ([]byt
 							record.Retain()
 							resultParts = append(
 								resultParts,
-								parts.NewArrowPart(partMeta.Tx, record, int(util.TotalRecordSize(record)), table.schema, partOptions),
+								parts.NewArrowPart(partMeta.Tx, record, uint64(util.TotalRecordSize(record)), table.schema, partOptions),
 							)
 							return nil
 						}(); err != nil {
diff --git a/table.go b/table.go
index de56f98..21b977e 100644
--- a/table.go
+++ b/table.go
@@ -1144,7 +1144,7 @@ func (t *Table) parquetCompaction(compact []parts.Part, options ...parts.Option)
 		postCompactionSize = buf.ParquetFile().Size()
 	}
 
-	return []parts.Part{parts.NewPart(0, buf, options...)}, preCompactionSize, postCompactionSize, nil
+	return []parts.Part{parts.NewParquetPart(0, buf, options...)}, preCompactionSize, postCompactionSize, 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 fb138b2..92d0970 100644
--- a/index/lsm_test.go
+++ b/index/lsm_test.go
@@ -27,7 +27,7 @@ func parquetCompaction(compact []parts.Part, _ ...parts.Option) ([]parts.Part, i
 	if err != nil {
 		return nil, 0, 0, err
 	}
-	return []parts.Part{parts.NewPart(0, buf)}, size, int64(b.Len()), nil
+	return []parts.Part{parts.NewParquetPart(0, buf)}, size, int64(b.Len()), nil
 }
 
 func compactParts(w io.Writer, compact []parts.Part) (int64, error) {
diff --git a/parts/part_test.go b/parts/part_test.go
index 916929b..53bda21 100644
--- a/parts/part_test.go
+++ b/parts/part_test.go
@@ -92,7 +92,7 @@ func TestFindMaximumNonOverlappingSet(t *testing.T) {
 				require.NoError(t, testSchema.SerializeBuffer(&b, buf))
 				serBuf, err := dynparquet.ReaderFromBytes(b.Bytes())
 				require.NoError(t, err)
-				parts[i] = NewPart(0, serBuf)
+				parts[i] = NewParquetPart(0, serBuf)
 			}
 			nonOverlapping, overlapping, err := FindMaximumNonOverlappingSet(testSchema, parts)
 			require.NoError(t, err)