Repository: polarsignals/frostdb
Pull request #719 (2024-02-01 14:36:09)

Fix missing compaction tx
Fixes compaction being able to write files without the `compaction_tx` metadata.

Also no longer treats a missing `compaction_tx` as an error, and simply downgrades the part tx.

asubiotto: Is there some more integration testing we can do here that would've caught this issue?

Patch:
diff --git a/index/compaction.go b/index/compaction.go
index 6895c5c..61ff0f6 100644
--- a/index/compaction.go
+++ b/index/compaction.go
@@ -248,14 +248,17 @@ func (f *FileCompaction) recover(options ...parts.Option) ([]parts.Part, error)
 					return err
 				}
 
+				var tx int
 				txstr, ok := buf.ParquetFile().Lookup(ParquetCompactionTXKey)
 				if !ok {
-					return fmt.Errorf("failed to find compaction_tx metadata: %s", file.Name())
-				}
-
-				tx, err := strconv.Atoi(txstr)
-				if err != nil {
-					return fmt.Errorf("failed to parse compaction_tx metadata: %v", err)
+					level.Warn(f.logger).Log("msg", "failed to find compaction_tx metadata", "file", file.Name())
+					tx = 0 // Downgrade the compaction tx so that all future reads will be able to read this part.
+				} else {
+					tx, err = strconv.Atoi(txstr)
+					if err != nil {
+						level.Warn(f.logger).Log("msg", "failed to parse compaction_tx metadata", "file", file.Name(), "err", err)
+						tx = 0 // Downgrade the compaction tx so that all future reads will be able to read this part.
+					}
 				}
 
 				f.parts.Add(1)
diff --git a/table.go b/table.go
index 83f0f71..c972bab 100644
--- a/table.go
+++ b/table.go
@@ -1300,7 +1300,7 @@ func (t *Table) compactParts(w io.Writer, compact []parts.Part, options ...parqu
 		}
 	}
 
-	bufs, err := t.buffersForCompaction(w, compact)
+	bufs, err := t.buffersForCompaction(w, compact, options...)
 	if err != nil {
 		return 0, err
 	}
@@ -1362,7 +1362,7 @@ func (t *Table) compactParts(w io.Writer, compact []parts.Part, options ...parqu
 // the minimum slice of dynamic row groups to be merged together for compaction.
 // If nil, nil is returned, the resulting serialized buffer is written directly
 // to w as an optimization.
-func (t *Table) buffersForCompaction(w io.Writer, inputParts []parts.Part) ([]dynparquet.DynamicRowGroup, error) {
+func (t *Table) buffersForCompaction(w io.Writer, inputParts []parts.Part, options ...parquet.WriterOption) ([]dynparquet.DynamicRowGroup, error) {
 	nonOverlappingParts, overlappingParts, err := parts.FindMaximumNonOverlappingSet(t.schema, inputParts)
 	if err != nil {
 		return nil, err
@@ -1424,7 +1424,7 @@ func (t *Table) buffersForCompaction(w io.Writer, inputParts []parts.Part) ([]dy
 		records = append(records, p.Record())
 	}
 
-	if err := t.writeRecordsToParquet(w, records, false); err != nil {
+	if err := t.writeRecordsToParquet(w, records, false, options...); err != nil {
 		return nil, err
 	}
 


Test patch:
diff --git a/db_test.go b/db_test.go
index 3f8c624..4974b79 100644
--- a/db_test.go
+++ b/db_test.go
@@ -2796,3 +2796,87 @@ func Test_DB_PersistentDiskCompaction_BlockRotation(t *testing.T) {
 
 	validateRows(2400)
 }
+
+func Test_DB_PersistentDiskCompaction_NonOverlappingCompaction(t *testing.T) {
+	t.Parallel()
+	config := NewTableConfig(
+		dynparquet.SampleDefinition(),
+	)
+	logger := newTestLogger(t)
+
+	dir := t.TempDir()
+	cfg := []*index.LevelConfig{
+		{Level: index.L0, MaxSize: 336, Type: index.CompactionTypeParquetDisk},
+		{Level: index.L1, MaxSize: 1024 * 1024 * 128, Type: index.CompactionTypeParquetDisk},
+		{Level: index.L2, MaxSize: 1024 * 1024 * 512},
+	}
+	c, err := New(
+		WithLogger(logger),
+		WithIndexConfig(cfg),
+		WithStoragePath(dir),
+		WithWAL(),
+		WithManualBlockRotation(),
+	)
+	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)
+
+	ctx := context.Background()
+	for i := 0; i < 3; i++ { // Create non-overlapping samples for compaction
+		samples := dynparquet.GenerateTestSamples(1)
+		r, err := samples.ToRecord()
+		require.NoError(t, err)
+		_, err = table.InsertRecord(ctx, r)
+		require.NoError(t, err)
+	}
+	require.Eventually(t, func() bool {
+		return table.active.index.LevelSize(index.L1) != 0
+	}, time.Second, time.Millisecond*10)
+
+	// Snapshot the file
+	success := false
+	table.db.snapshot(context.Background(), false, func() {
+		success = true
+	})
+	require.True(t, success)
+
+	// Close the database
+	require.NoError(t, c.Close())
+
+	// Recover from the snapshot
+	c, err = New(
+		WithLogger(logger),
+		WithIndexConfig(cfg),
+		WithStoragePath(dir),
+		WithWAL(),
+	)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		require.NoError(t, c.Close())
+	})
+	db, err = c.DB(context.Background(), "test")
+	require.NoError(t, err)
+	table, err = db.Table("test", config)
+	require.NoError(t, err)
+
+	validateRows := func(expected int64) {
+		pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
+		defer pool.AssertSize(t, 0)
+		rows := int64(0)
+		engine := query.NewEngine(pool, db.TableProvider())
+		err = engine.ScanTable("test").
+			Execute(context.Background(), func(ctx context.Context, r arrow.Record) error {
+				rows += r.NumRows()
+				return nil
+			})
+		require.NoError(t, err)
+		require.Equal(t, expected, rows)
+	}
+
+	validateRows(3)
+}