Repository: polarsignals/frostdb
Pull request #572 (2023-10-25 10:18:30)

table: run distinct on arrow records during compaction
This applies specifically to tables with a unique primary index. This reduces
the amount of data that is converted to parquet and also avoids a merge step
for L0 compactions.

```
pkg: github.com/polarsignals/frostdb
                         │ benchmain_labels │         bench_distinct_sw          │
                         │      sec/op      │   sec/op     vs base               │
Replay//labels_view/0-12         7.104 ± 2%   2.793 ± 15%  -60.68% (p=0.002 n=6)

                         │ benchmain_labels │          bench_distinct_sw          │
                         │       B/op       │     B/op      vs base               │
Replay//labels_view/0-12      10.180Gi ± 2%   5.513Gi ± 3%  -45.84% (p=0.002 n=6)

                         │ benchmain_labels │         bench_distinct_sw          │
                         │    allocs/op     │  allocs/op   vs base               │
Replay//labels_view/0-12        52.25M ± 2%   20.84M ± 0%  -60.11% (p=0.002 n=6)
```


Patch:
diff --git a/cmd/parquet-reencode/main.go b/cmd/parquet-reencode/main.go
index 8a975f0..00c6d99 100644
--- a/cmd/parquet-reencode/main.go
+++ b/cmd/parquet-reencode/main.go
@@ -52,7 +52,7 @@ func main() {
 		log.Fatal(fmt.Errorf("create output file: %w", err))
 	}
 
-	w, err := newSchema.GetWriter(outf, serBuf.DynamicColumns())
+	w, err := newSchema.GetWriter(outf, serBuf.DynamicColumns(), false)
 	if err != nil {
 		log.Fatal(fmt.Errorf("get writer: %w", err))
 	}
diff --git a/dynparquet/schema.go b/dynparquet/schema.go
index 095baac..e828a86 100644
--- a/dynparquet/schema.go
+++ b/dynparquet/schema.go
@@ -1185,7 +1185,7 @@ func (s *Schema) NewBufferV2(dynamicColumns ...*schemav2pb.Node) (*Buffer, error
 }
 
 func (s *Schema) SerializeBuffer(w io.Writer, buffer *Buffer) error {
-	pw, err := s.GetWriter(w, buffer.DynamicColumns())
+	pw, err := s.GetWriter(w, buffer.DynamicColumns(), false)
 	if err != nil {
 		return fmt.Errorf("create writer: %w", err)
 	}
@@ -1211,7 +1211,7 @@ const bloomFilterBitsPerValue = 10
 
 // NewWriter returns a new parquet writer with a concrete parquet schema
 // generated using the given concrete dynamic column names.
-func (s *Schema) NewWriter(w io.Writer, dynamicColumns map[string][]string) (ParquetWriter, error) {
+func (s *Schema) NewWriter(w io.Writer, dynamicColumns map[string][]string, sorting bool) (ParquetWriter, error) {
 	ps, err := s.GetDynamicParquetSchema(dynamicColumns)
 	if err != nil {
 		return nil, err
@@ -1248,6 +1248,9 @@ func (s *Schema) NewWriter(w io.Writer, dynamicColumns map[string][]string) (Par
 			parquet.SortingColumns(cols...),
 		),
 	}
+	if sorting {
+		return parquet.NewSortingWriter[any](w, 32*1024, writerOptions...), nil
+	}
 	return parquet.NewGenericWriter[any](w, writerOptions...), nil
 }
 
@@ -1265,12 +1268,12 @@ type PooledWriter struct {
 	ParquetWriter
 }
 
-func (s *Schema) GetWriter(w io.Writer, dynamicColumns map[string][]string) (*PooledWriter, error) {
+func (s *Schema) GetWriter(w io.Writer, dynamicColumns map[string][]string, sorting bool) (*PooledWriter, error) {
 	key := serializeDynamicColumns(dynamicColumns)
-	pool, _ := s.writers.LoadOrStore(key, &sync.Pool{})
+	pool, _ := s.writers.LoadOrStore(fmt.Sprintf("%s,sorting=%t", key, sorting), &sync.Pool{})
 	pooled := pool.(*sync.Pool).Get()
 	if pooled == nil {
-		new, err := s.NewWriter(w, dynamicColumns)
+		new, err := s.NewWriter(w, dynamicColumns, sorting)
 		if err != nil {
 			return nil, err
 		}
diff --git a/parts/part.go b/parts/part.go
index bb017fe..9c19b3e 100644
--- a/parts/part.go
+++ b/parts/part.go
@@ -50,7 +50,7 @@ func (p *Part) AsSerializedBuffer(schema *dynparquet.Schema) (*dynparquet.Serial
 	// 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))
+	w, err := schema.GetWriter(b, pqarrow.RecordDynamicCols(p.record), false)
 	if err != nil {
 		return nil, err
 	}
diff --git a/table.go b/table.go
index 886198d..f09dd28 100644
--- a/table.go
+++ b/table.go
@@ -39,6 +39,7 @@ import (
 	"github.com/polarsignals/frostdb/parts"
 	"github.com/polarsignals/frostdb/pqarrow"
 	"github.com/polarsignals/frostdb/query/logicalplan"
+	"github.com/polarsignals/frostdb/query/physicalplan"
 	"github.com/polarsignals/frostdb/recovery"
 	"github.com/polarsignals/frostdb/wal"
 	walpkg "github.com/polarsignals/frostdb/wal"
@@ -935,7 +936,7 @@ type parquetRowWriterOption func(p *parquetRowWriter)
 
 // rowWriter returns a new Parquet row writer with the given dynamic columns.
 func (t *TableBlock) rowWriter(writer io.Writer, dynCols map[string][]string, options ...parquetRowWriterOption) (*parquetRowWriter, error) {
-	w, err := t.table.schema.NewWriter(writer, dynCols)
+	w, err := t.table.schema.NewWriter(writer, dynCols, false)
 	if err != nil {
 		return nil, err
 	}
@@ -1162,6 +1163,24 @@ func (t *Table) compactParts(w io.Writer, compact []*parts.Part) (int64, error)
 		preCompactionSize += p.Size()
 	}
 
+	if t.schema.UniquePrimaryIndex {
+		distinctRecords, err := t.distinctRecordsForCompaction(compact)
+		if err != nil {
+			return 0, err
+		}
+		if distinctRecords != nil {
+			// Arrow distinction was successful.
+			defer func() {
+				for _, r := range distinctRecords {
+					r.Release()
+				}
+			}()
+			// Note that the records must be sorted (sortInput=true) because
+			// there is no guarantee that order is maintained.
+			return preCompactionSize, t.writeRecordsToParquet(w, distinctRecords, true)
+		}
+	}
+
 	bufs, err := t.buffersForCompaction(w, compact)
 	if err != nil {
 		return 0, err
@@ -1268,21 +1287,11 @@ func (t *Table) buffersForCompaction(w io.Writer, inputParts []*parts.Part) ([]d
 	}
 
 	records := make([]arrow.Record, 0, len(nonOverlappingParts))
-	dynColSets := make([]map[string][]string, 0, len(nonOverlappingParts))
 	for _, p := range nonOverlappingParts {
-		r := p.Record()
-		dynColSets = append(dynColSets, pqarrow.RecordDynamicCols(r))
-		records = append(records, r)
+		records = append(records, p.Record())
 	}
-	dynCols := dynparquet.MergeDynamicColumnSets(dynColSets)
 
-	pw, err := t.schema.GetWriter(w, dynCols)
-	if err != nil {
-		return nil, err
-	}
-	defer t.schema.PutWriter(pw)
-
-	if err := pqarrow.RecordsToFile(t.schema, pw, records); err != nil {
+	if err := t.writeRecordsToParquet(w, records, false); err != nil {
 		return nil, err
 	}
 
@@ -1298,3 +1307,69 @@ func (t *Table) buffersForCompaction(w io.Writer, inputParts []*parts.Part) ([]d
 	result = append(result, buf.MultiDynamicRowGroup())
 	return result, nil
 }
+
+func (t *Table) writeRecordsToParquet(w io.Writer, records []arrow.Record, sortInput bool) error {
+	dynColSets := make([]map[string][]string, 0, len(records))
+	for _, r := range records {
+		dynColSets = append(dynColSets, pqarrow.RecordDynamicCols(r))
+	}
+	dynCols := dynparquet.MergeDynamicColumnSets(dynColSets)
+	pw, err := t.schema.GetWriter(w, dynCols, sortInput)
+	if err != nil {
+		return err
+	}
+	defer t.schema.PutWriter(pw)
+
+	return pqarrow.RecordsToFile(t.schema, pw, records)
+}
+
+// distinctRecordsForCompaction performs a distinct on the given parts. If at
+// least one non-arrow part is found, nil, nil is returned in which case, the
+// caller should fall back to normal compaction. On success, the caller is
+// responsible for releasing the returned records.
+func (t *Table) distinctRecordsForCompaction(compact []*parts.Part) ([]arrow.Record, error) {
+	sortingCols := t.schema.SortingColumns()
+	columnExprs := make([]logicalplan.Expr, 0, len(sortingCols))
+	for _, col := range sortingCols {
+		var expr logicalplan.Expr
+		if col.Dynamic {
+			expr = logicalplan.DynCol(col.Name)
+		} else {
+			expr = logicalplan.Col(col.Name)
+		}
+		columnExprs = append(columnExprs, expr)
+	}
+
+	d := physicalplan.Distinct(memory.NewGoAllocator(), t.tracer, columnExprs)
+	output := physicalplan.OutputPlan{}
+	newRecords := make([]arrow.Record, 0)
+	output.SetNextCallback(func(ctx context.Context, r arrow.Record) error {
+		r.Retain()
+		newRecords = append(newRecords, r)
+		return nil
+	})
+	d.SetNext(&output)
+
+	if ok, err := func() (bool, error) {
+		ctx := context.TODO()
+		for _, p := range compact {
+			if p.Record() == nil {
+				// Caller should fall back to parquet distinction.
+				return false, nil
+			}
+			if err := d.Callback(ctx, p.Record()); err != nil {
+				return false, err
+			}
+		}
+		if err := d.Finish(ctx); err != nil {
+			return false, err
+		}
+		return true, nil
+	}(); !ok || err != nil {
+		for _, r := range newRecords {
+			r.Release()
+		}
+		return nil, err
+	}
+	return newRecords, nil
+}


Test patch:
diff --git a/dynparquet/reader_test.go b/dynparquet/reader_test.go
index 3356863..58edc47 100644
--- a/dynparquet/reader_test.go
+++ b/dynparquet/reader_test.go
@@ -17,7 +17,7 @@ func TestReader(t *testing.T) {
 	b := bytes.NewBuffer(nil)
 	w, err := schema.NewWriter(b, map[string][]string{
 		"labels": samples.SampleLabelNames(),
-	})
+	}, false)
 	require.NoError(t, err)
 
 	_, err = parquet.CopyRows(w, buf.Rows())
diff --git a/index/lsm_test.go b/index/lsm_test.go
index cf6008e..262b9b7 100644
--- a/index/lsm_test.go
+++ b/index/lsm_test.go
@@ -47,7 +47,7 @@ func compactParts(w io.Writer, compact []*parts.Part) (int64, error) {
 		return 0, err
 	}
 	err = func() error {
-		writer, err := schema.GetWriter(w, merged.DynamicColumns())
+		writer, err := schema.GetWriter(w, merged.DynamicColumns(), false)
 		if err != nil {
 			return err
 		}