Parquet file compaction (#621)

Adds support for file based compaction

Adds Parquet file compaction that will append the Parquet files
into a os temporary file. During compaction of the file part it will
truncate the file after the parts are released so that new parts can be
appended to the file.

Author: thorfour

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

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)
+}

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

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)
 }

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
+}

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