Repository: polarsignals/frostdb
Pull request #704 (2024-01-23 22:26:38)

Snapshot index rotation
Added a design doc in the designs directory with more detailed information.

TL;DR this fixes some data loss scenarios with the persistent lsm index files that were present with snapshots.

Previously if a snapshot occurred that mean the only copy of those now truncated writes that were in the WAL were now only found in an index file. If that index file was compacted into/out of and the node shutdown during this period that compaction could cause file corruption. This meant the only copy of the data was lost, and data loss occurred. 

Instead of relying solely on the index files we now have the snapshot hard link the index files into the snapshot directory. This is still much faster than copying the data into the snapshot file, and can be restored on startup by hard linking the snapshotted files back into the index.

This also adds a couple of QoL changes around the LSM where now the LSM configures its own levels from the settings instead of having to call configureLSM before calling NewLSM 


Patch:
diff --git a/compaction.go b/compaction.go
deleted file mode 100644
index 37141a3..0000000
--- a/compaction.go
+++ /dev/null
@@ -1,209 +0,0 @@
-package frostdb
-
-import (
-	"encoding/binary"
-	"fmt"
-	"io"
-	"os"
-	"path/filepath"
-	"strconv"
-
-	"github.com/go-kit/log/level"
-	"github.com/oklog/ulid"
-	"github.com/parquet-go/parquet-go"
-
-	"github.com/polarsignals/frostdb/dynparquet"
-	"github.com/polarsignals/frostdb/parts"
-)
-
-const (
-	ParquetCompactionTXKey = "compaction_tx"
-)
-
-var ErrCompactionRecoveryFailed = fmt.Errorf("failed to recover compacted level")
-
-type FileCompaction struct {
-	t      *Table
-	file   *os.File
-	offset int64 // Writing offsets into the file
-	ref    int64 // Number of references to file.
-}
-
-func (f *FileCompaction) Close(cleanup bool) error {
-	err := f.file.Close()
-	if cleanup {
-		if err := os.Remove(f.file.Name()); err != nil {
-			return fmt.Errorf("failed to remove file: %v", err)
-		}
-	}
-	return err
-}
-
-func NewFileCompaction(t *Table, block ulid.ULID, lvl int) *FileCompaction {
-	f := &FileCompaction{
-		t: t,
-	}
-
-	if t.db.storagePath == "" { // No storage path provided, use temp files.
-		file, err := os.CreateTemp("", fmt.Sprintf("L%v-*.parquet", lvl))
-		if err != nil {
-			panic(err)
-		}
-		level.Info(t.logger).Log("msg", "created temp file for level", "level", lvl, "file", file.Name())
-		f.file = file
-		return f
-	}
-
-	indexDir := filepath.Join(t.db.indexDir(), t.name)
-	if err := os.MkdirAll(indexDir, dirPerms); err != nil {
-		panic(err)
-	}
-
-	file, err := os.OpenFile(filepath.Join(indexDir, fmt.Sprintf("L%v-%s.parquet", lvl, block.String())), os.O_CREATE|os.O_RDWR, filePerms)
-	if err != nil {
-		panic(err)
-	}
-	f.file = file
-	return f
-}
-
-// accountingWriter is a writer that accounts for the number of bytes written.
-type accountingWriter struct {
-	w io.Writer
-	n int64
-}
-
-func (a *accountingWriter) Write(p []byte) (int, error) {
-	n, err := a.w.Write(p)
-	a.n += int64(n)
-	return n, err
-}
-
-// writeRecordsToParquetFile will compact the given parts into a Parquet file written to the next level file.
-func (f *FileCompaction) writeRecordsToParquetFile(compact []parts.Part, options ...parts.Option) ([]parts.Part, int64, int64, error) {
-	if len(compact) == 0 {
-		return nil, 0, 0, fmt.Errorf("no parts to compact")
-	}
-
-	accountant := &accountingWriter{w: f.file}
-	preCompactionSize, err := f.t.compactParts(accountant, compact,
-		parquet.KeyValueMetadata(
-			ParquetCompactionTXKey, // Compacting up through this transaction.
-			fmt.Sprintf("%v", compact[0].TX()),
-		),
-	) // compact into the next level
-	if err != nil {
-		return nil, 0, 0, err
-	}
-
-	// Record the writing offset into the file.
-	prevOffset := f.offset
-	f.offset += accountant.n + 8
-	f.ref++
-
-	// Record the file size for recovery.
-	size := make([]byte, 8)
-	binary.LittleEndian.PutUint64(size, uint64(accountant.n))
-	if n, err := f.file.Write(size); n != 8 {
-		return nil, 0, 0, fmt.Errorf("failed to write size to file: %v", err)
-	}
-
-	pf, err := parquet.OpenFile(io.NewSectionReader(f.file, prevOffset, accountant.n), accountant.n)
-	if err != nil {
-		return nil, 0, 0, err
-	}
-
-	buf, err := dynparquet.NewSerializedBuffer(pf)
-	if err != nil {
-		return nil, 0, 0, err
-	}
-
-	return []parts.Part{parts.NewParquetPart(compact[0].TX(), buf, append(options, parts.WithRelease(f.release()))...)}, preCompactionSize, accountant.n, nil
-}
-
-// Truncate will truncate the file to 0 bytes. This is used when a compaction recovery fails.
-func (f *FileCompaction) Truncate() error {
-	return f.file.Truncate(0)
-}
-
-// release will account for all the Parts currently pointing to this file. Once the last one has been released it will truncate the file.
-// release is not safe to call concurrently.
-func (f *FileCompaction) release() func() {
-	return func() {
-		f.ref--
-		if f.ref == 0 {
-			level.Info(f.t.logger).Log("msg", "truncating file", "file", f.file.Name())
-			if err := os.Truncate(f.file.Name(), 0); err != nil {
-				panic(fmt.Errorf("failed to truncate the level file: %v", err))
-			}
-
-			f.offset = 0
-			if _, err := f.file.Seek(0, io.SeekStart); err != nil {
-				panic(fmt.Errorf("failed to seek the level file: %v", err))
-			}
-		}
-	}
-}
-
-func (f *FileCompaction) recover(options ...parts.Option) ([]parts.Part, error) {
-	recovered, err := func() ([]parts.Part, error) {
-		info, err := os.Stat(f.file.Name())
-		if err != nil {
-			return nil, err
-		}
-
-		if info.Size() == 0 { // file was truncated, nothing to recover.
-			return nil, nil
-		}
-
-		recovered := []parts.Part{}
-
-		// Recover all parts from file.
-		for offset := info.Size(); offset > 0; {
-			offset -= 8
-			size := make([]byte, 8)
-			if n, err := f.file.ReadAt(size, offset); n != 8 {
-				return nil, fmt.Errorf("failed to read size from file: %v", err)
-			}
-			parquetSize := int64(binary.LittleEndian.Uint64(size))
-			offset -= parquetSize
-
-			pf, err := parquet.OpenFile(io.NewSectionReader(f.file, offset, parquetSize), parquetSize)
-			if err != nil {
-				return nil, err
-			}
-
-			buf, err := dynparquet.NewSerializedBuffer(pf)
-			if err != nil {
-				return nil, err
-			}
-
-			txstr, ok := buf.ParquetFile().Lookup(ParquetCompactionTXKey)
-			if !ok {
-				return nil, fmt.Errorf("failed to find compaction_tx metadata")
-			}
-
-			tx, err := strconv.Atoi(txstr)
-			if err != nil {
-				return nil, fmt.Errorf("failed to parse compaction_tx metadata: %v", err)
-			}
-
-			recovered = append(recovered, parts.NewParquetPart(uint64(tx), buf, append(options, parts.WithRelease(f.release()))...))
-		}
-
-		return recovered, nil
-	}()
-	if err != nil {
-		level.Error(f.t.logger).Log("msg", "truncating file after failed recovery", "err", err, "file", f.file.Name())
-		if err := os.Truncate(f.file.Name(), 0); err != nil {
-			return nil, fmt.Errorf("failed to truncate the level file %s: %v", f.file.Name(), err)
-		}
-
-		return nil, ErrCompactionRecoveryFailed
-	}
-
-	return recovered, err
-}
-
-// Sync calls Sync on the underlying file.
-func (f *FileCompaction) Sync() error { return f.file.Sync() }
diff --git a/db.go b/db.go
index a860510..3715bc0 100644
--- a/db.go
+++ b/db.go
@@ -66,7 +66,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 []*IndexConfig
+	indexConfig []*index.LevelConfig
 
 	sources []DataSource
 	sinks   []DataSink
@@ -129,9 +129,9 @@ func New(
 	}
 
 	for _, cfg := range s.indexConfig {
-		if cfg.Type == CompactionTypeParquetDisk {
-			if !s.enableWAL && s.storagePath != "" {
-				return nil, fmt.Errorf("persistent disk compaction requires WAL to be enabled")
+		if cfg.Type == index.CompactionTypeParquetDisk {
+			if !s.enableWAL || s.storagePath == "" {
+				return nil, fmt.Errorf("persistent disk compaction requires WAL and storage path to be enabled")
 			}
 		}
 	}
@@ -228,7 +228,7 @@ func WithStoragePath(path string) Option {
 	}
 }
 
-func WithIndexConfig(indexConfig []*IndexConfig) Option {
+func WithIndexConfig(indexConfig []*index.LevelConfig) Option {
 	return func(s *ColumnStore) error {
 		s.indexConfig = indexConfig
 		return nil
@@ -740,25 +740,17 @@ func (db *DB) recover(ctx context.Context, wal WAL) error {
 				// already been persisted.
 				db.mtx.Lock()
 				if table, ok := db.tables[e.TableBlockPersisted.TableName]; ok {
-					cfg, recovered, err := table.ActiveBlock().configureLSMLevels(db.columnStore.indexConfig)
-					if err != nil {
-						return fmt.Errorf("configure lsm levels: %w", err)
-					}
-
 					table.ActiveBlock().index, err = index.NewLSM(
-						table.name,
+						filepath.Join(table.db.indexDir(), table.name, table.ActiveBlock().ulid.String()), // Any index files are found at <db.indexDir>/<table.name>/<block.id>
 						table.schema,
-						cfg,
+						table.IndexConfig(),
 						db.Wait,
 						index.LSMWithMetrics(table.metrics.indexMetrics),
+						index.LSMWithLogger(table.logger),
 					)
 					if err != nil {
 						return fmt.Errorf("create new lsm index: %w", err)
 					}
-
-					for _, p := range recovered {
-						table.ActiveBlock().index.InsertPart(p)
-					}
 				}
 				db.mtx.Unlock()
 			}
@@ -1301,18 +1293,6 @@ func (db *DB) resetToTxn(txn uint64, wal WAL) {
 	db.tx.Store(txn)
 	db.highWatermark.Store(txn)
 	if wal != nil {
-		// Before resetting the wal make sure any pending writes to the index are flushed.
-		errg := errgroup.Group{}
-		for _, table := range db.tables {
-			for _, syncer := range table.ActiveBlock().syncers {
-				errg.Go(syncer.Sync)
-			}
-		}
-
-		if err := errg.Wait(); err != nil {
-			level.Error(db.logger).Log("msg", "failed to sync index before resetting WAL; dataloss may occur", "err", err)
-		}
-
 		// This call resets the WAL to a zero state so that new records can be
 		// logged.
 		if err := wal.Reset(txn + 1); err != nil {
diff --git a/design/index_files.md b/design/index_files.md
new file mode 100644
index 0000000..8c67795
--- /dev/null
+++ b/design/index_files.md
@@ -0,0 +1,64 @@
+# On-Disk persistent index files
+
+Authored: Jan 22, 2024
+Author: Thor Hansen (@thorfour)
+
+## Abstract
+
+FrostDB currently uses temporary index files to store Parquet data. These are lost on shutdown, and need to be recreated from snapshots and WAL.
+Furthermore snapshots are less effecient than necessary because snapshotting copies Parquet data that is already on disk, and writes it to a different file (snapshot file)
+when we could just reference the index file instead.
+
+## Proposal
+
+Instead FrostDB should use persistent index files that don't requrie snapshots to copy data but instead hard link the files in the snapshot directory for recovery.
+
+Index files have the format of:
+
+|--------------------|
+|   <ParquetFile>    |
+|   <uint64 size>    |
+|   <ParquetFile>    |
+|   <uint64 size>    |
+|   <ParquetFile>    |
+|   <uint64 size>    |
+|--------------------|
+
+Where Parquet files are appending togehter with a uint64 that represents the number of bytes in the previous ParquetFile.
+Each Parquet file represents a single Part in the LSM index. The Part's transaction id is stored in the Parquet file with a metadata key of 'compaction_tx'
+
+Index files will be stored in a directory with the prefix of:
+databases/<database>/index/<table>/<block>/<index_level>/
+
+With the filename format of 000000000001.parquet with a 20 byte representation of the ordering of the index file.
+
+### Index File Lifecycle
+
+There are 4 events that influence the lifecycle of an index file.
+- Writes
+- Snapshots
+- Compaction
+- Block Rotation
+
+Writes: 
+    These are compaction from a lower level that write into the index file. These writes shall write into the currently active index file in the level,
+    appending the newly compacted Parquet file into the index file with the uint64 at the end of the number of bytes written.
+    This file shall be fsync'd after a completed write.
+
+Snapshots: 
+    When a snapshot occurs the index files in every level of the index shall be rotated. This means a new index file is created the old one is fsync'd
+    and all future writes to the level are writen into that new file leaving the previous file immutable. The old files are then hard linked into the snapshot directory,
+    such that when the index deletes those files they can still be recovered from the snapshot.
+
+Compaction: 
+    This is when a level is compacted into a new Parquet file and written to the next level. All the files in that compacted level shall be deleted, as they are now
+    covered by the latest write into the next level.
+
+Rotation:
+    When an entire table block is rotated out, after rotation is completed the block directory that holds all levels of index files can be deleted.
+
+
+### Replay
+
+When a table block is created it shall walk through the table blocks directory. Upon discovering index files it will open and read all parts of that file.
+Returning a list of parts that were recovered with the block was created. The Table shall than replay those recovered parts into the LSM.
diff --git a/index/compaction.go b/index/compaction.go
new file mode 100644
index 0000000..6f495c2
--- /dev/null
+++ b/index/compaction.go
@@ -0,0 +1,317 @@
+package index
+
+import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"strconv"
+	"sync"
+
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+	"github.com/parquet-go/parquet-go"
+
+	"github.com/polarsignals/frostdb/dynparquet"
+	"github.com/polarsignals/frostdb/parts"
+)
+
+const (
+	IndexFileExtension     = ".idx"
+	ParquetCompactionTXKey = "compaction_tx"
+	dirPerms               = os.FileMode(0o755)
+	filePerms              = os.FileMode(0o640)
+)
+
+type Compaction func(w io.Writer, compact []parts.Part, options ...parquet.WriterOption) (int64, error)
+
+type FileCompaction struct {
+	// settings
+	dir     string
+	compact Compaction
+	maxSize int64
+
+	// internal data
+	indexFiles []*os.File
+	offset     int64          // Writing offsets into the file
+	parts      sync.WaitGroup // Wait group for parts that are currently reference in this level.
+
+	// Options
+	logger log.Logger
+}
+
+func NewFileCompaction(dir string, maxSize int64, compact Compaction, logger log.Logger) (*FileCompaction, error) {
+	f := &FileCompaction{
+		dir:     dir,
+		compact: compact,
+		maxSize: maxSize,
+		logger:  logger,
+	}
+
+	if err := os.MkdirAll(dir, dirPerms); err != nil {
+		return nil, err
+	}
+
+	return f, nil
+}
+
+func (f *FileCompaction) MaxSize() int64 { return f.maxSize }
+
+// Snapshot will create a new index file and all future compacitons are written to that index file.
+func (f *FileCompaction) Snapshot(dir string) error {
+	if len(f.indexFiles) > 0 && f.offset == 0 { // If we are at the beginning of the file, no need to create a new file.
+		return nil
+	}
+
+	// First sync the current file if there is one.
+	if len(f.indexFiles) > 0 {
+		if err := f.Sync(); err != nil {
+			return err
+		}
+
+		for _, file := range f.indexFiles {
+			// Hard link the file into the snapshot directory.
+			if err := os.Link(file.Name(), filepath.Join(dir, filepath.Base(file.Name()))); err != nil {
+				return err
+			}
+		}
+	}
+
+	_, err := f.createIndexFile(len(f.indexFiles))
+	return err
+}
+
+func (f *FileCompaction) createIndexFile(id int) (*os.File, error) {
+	file, err := os.OpenFile(filepath.Join(f.dir, fmt.Sprintf("%020d%s", id, IndexFileExtension)), os.O_CREATE|os.O_RDWR, filePerms)
+	if err != nil {
+		return nil, err
+	}
+
+	f.offset = 0
+	f.indexFiles = append(f.indexFiles, file)
+	return file, nil
+}
+
+func (f *FileCompaction) openIndexFile(path string) (*os.File, error) {
+	file, err := os.Open(path)
+	if err != nil {
+		return nil, err
+	}
+
+	f.indexFiles = append(f.indexFiles, file)
+	return file, nil
+}
+
+// file returns the currently active index file.
+func (f *FileCompaction) file() *os.File {
+	return f.indexFiles[len(f.indexFiles)-1]
+}
+
+// accountingWriter is a writer that accounts for the number of bytes written.
+type accountingWriter struct {
+	w io.Writer
+	n int64
+}
+
+func (a *accountingWriter) Write(p []byte) (int, error) {
+	n, err := a.w.Write(p)
+	a.n += int64(n)
+	return n, err
+}
+
+// Compact will compact the given parts into a Parquet file written to the next level file.
+func (f *FileCompaction) Compact(compact []parts.Part, options ...parts.Option) ([]parts.Part, int64, int64, error) {
+	if len(compact) == 0 {
+		return nil, 0, 0, fmt.Errorf("no parts to compact")
+	}
+
+	accountant := &accountingWriter{w: f.file()}
+	preCompactionSize, err := f.compact(accountant, compact,
+		parquet.KeyValueMetadata(
+			ParquetCompactionTXKey, // Compacting up through this transaction.
+			fmt.Sprintf("%v", compact[0].TX()),
+		),
+	) // compact into the next level
+	if err != nil {
+		return nil, 0, 0, err
+	}
+
+	// Record the writing offset into the file.
+	prevOffset := f.offset
+
+	// Record the file size for recovery.
+	size := make([]byte, 8)
+	binary.LittleEndian.PutUint64(size, uint64(accountant.n))
+	if n, err := f.file().Write(size); n != 8 {
+		return nil, 0, 0, fmt.Errorf("failed to write size to file: %v", err)
+	}
+	f.offset += accountant.n + 8
+
+	// Sync file after writing.
+	if err := f.Sync(); err != nil {
+		return nil, 0, 0, fmt.Errorf("failed to sync file: %v", err)
+	}
+
+	pf, err := parquet.OpenFile(io.NewSectionReader(f.file(), prevOffset, accountant.n), accountant.n)
+	if err != nil {
+		return nil, 0, 0, fmt.Errorf("failed to open file after compaction: %w", err)
+	}
+
+	buf, err := dynparquet.NewSerializedBuffer(pf)
+	if err != nil {
+		return nil, 0, 0, err
+	}
+
+	f.parts.Add(1)
+	return []parts.Part{parts.NewParquetPart(compact[0].TX(), buf, append(options, parts.WithRelease(f.parts.Done))...)}, preCompactionSize, accountant.n, nil
+}
+
+// Reset is called when the level no longer has active parts in it at the end of a compaction.
+func (f *FileCompaction) Reset() {
+	f.parts.Wait() // Wait for all parts to be released.
+	for _, file := range f.indexFiles {
+		if err := file.Close(); err != nil {
+			level.Error(f.logger).Log("msg", "failed to close level file", "err", err)
+		}
+	}
+
+	// Delete all the files in the directory level. And open a new file.
+	if err := os.RemoveAll(f.dir); err != nil {
+		level.Error(f.logger).Log("msg", "failed to remove level directory", "err", err)
+	}
+
+	if err := os.MkdirAll(f.dir, dirPerms); err != nil {
+		level.Error(f.logger).Log("msg", "failed to create level directory", "err", err)
+	}
+
+	f.indexFiles = nil
+	_, err := f.createIndexFile(len(f.indexFiles))
+	if err != nil {
+		level.Error(f.logger).Log("msg", "failed to create new level file", "err", err)
+	}
+}
+
+// recovery the level from the given directory.
+func (f *FileCompaction) recover(options ...parts.Option) ([]parts.Part, error) {
+	defer func() {
+		_, err := f.createIndexFile(len(f.indexFiles))
+		if err != nil {
+			level.Error(f.logger).Log("msg", "failed to create new level file", "err", err)
+		}
+	}()
+	recovered := []parts.Part{}
+	err := filepath.WalkDir(f.dir, func(path string, d os.DirEntry, err error) error {
+		if err != nil {
+			return err
+		}
+
+		if filepath.Ext(path) != IndexFileExtension {
+			return nil
+		}
+
+		info, err := d.Info()
+		if err != nil {
+			return fmt.Errorf("failed to get file info: %v", err)
+		}
+
+		if info.Size() == 0 { // file empty, nothing to recover.
+			return nil
+		}
+
+		file, err := f.openIndexFile(path)
+		if err != nil {
+			return fmt.Errorf("failed to open file: %v", err)
+		}
+
+		// Recover all parts from file.
+		fileParts := []parts.Part{}
+		if err := func() error {
+			for offset := info.Size(); offset > 0; {
+				offset -= 8
+				size := make([]byte, 8)
+				if n, err := file.ReadAt(size, offset); n != 8 {
+					return fmt.Errorf("failed to read size from file: %v", err)
+				}
+				parquetSize := int64(binary.LittleEndian.Uint64(size))
+				offset -= parquetSize
+
+				pf, err := parquet.OpenFile(io.NewSectionReader(file, offset, parquetSize), parquetSize)
+				if err != nil {
+					return err
+				}
+
+				buf, err := dynparquet.NewSerializedBuffer(pf)
+				if err != nil {
+					return err
+				}
+
+				txstr, ok := buf.ParquetFile().Lookup(ParquetCompactionTXKey)
+				if !ok {
+					return fmt.Errorf("failed to find compaction_tx metadata")
+				}
+
+				tx, err := strconv.Atoi(txstr)
+				if err != nil {
+					return fmt.Errorf("failed to parse compaction_tx metadata: %v", err)
+				}
+
+				f.parts.Add(1)
+				fileParts = append(fileParts, parts.NewParquetPart(uint64(tx), buf, append(options, parts.WithRelease(f.parts.Done))...))
+			}
+
+			return nil
+		}(); err != nil {
+			for _, part := range fileParts {
+				part.Release()
+			}
+
+			// If we failed to recover the file, remove it.
+			if err := f.file().Close(); err != nil {
+				level.Error(f.logger).Log("msg", "failed to close level file after failed recovery", "err", err)
+			}
+			f.indexFiles = f.indexFiles[:len(f.indexFiles)-1] // Remove the file from the list of files.
+			return err
+		}
+
+		recovered = append(recovered, fileParts...)
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return recovered, nil
+}
+
+// Sync calls Sync on the underlying file.
+func (f *FileCompaction) Sync() error { return f.file().Sync() }
+
+type inMemoryLevel struct {
+	compact Compaction
+	maxSize int64
+}
+
+func (l *inMemoryLevel) MaxSize() int64          { return l.maxSize }
+func (l *inMemoryLevel) Snapshot(_ string) error { return nil }
+func (l *inMemoryLevel) Reset()                  {}
+func (l *inMemoryLevel) Compact(toCompact []parts.Part, options ...parts.Option) ([]parts.Part, int64, int64, error) {
+	if len(toCompact) == 0 {
+		return nil, 0, 0, fmt.Errorf("no parts to compact")
+	}
+
+	var b bytes.Buffer
+	preCompactionSize, err := l.compact(&b, toCompact)
+	if err != nil {
+		return nil, 0, 0, err
+	}
+
+	buf, err := dynparquet.ReaderFromBytes(b.Bytes())
+	if err != nil {
+		return nil, 0, 0, err
+	}
+
+	postCompactionSize := int64(b.Len())
+	return []parts.Part{parts.NewParquetPart(0, buf, options...)}, preCompactionSize, postCompactionSize, nil
+}
diff --git a/index/lsm.go b/index/lsm.go
index 18004c4..30af97f 100644
--- a/index/lsm.go
+++ b/index/lsm.go
@@ -3,6 +3,8 @@ package index
 import (
 	"context"
 	"fmt"
+	"os"
+	"path/filepath"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -20,6 +22,18 @@ import (
 	"github.com/polarsignals/frostdb/query/logicalplan"
 )
 
+type CompactionType int
+
+const (
+	CompactionTypeUnknown CompactionType = iota
+
+	// CompactionTypeParquetDisk is a compaction type that will compact the parts into a Parquet file on disk.
+	CompactionTypeParquetDisk
+
+	// CompactionTypeParquetMemory is a compaction type that will compact the parts into a Parquet file in memory.
+	CompactionTypeParquetMemory
+)
+
 // LSM is a log-structured merge-tree like index. It is implemented as a single linked list of parts.
 //
 // Arrow records are always added to the L0 list. When a list reaches it's configured max size it is compacted
@@ -28,16 +42,16 @@ import (
 // [L0]->[record]->[record]->[L1]->[record/parquet]->[record/parquet] etc.
 type LSM struct {
 	sync.RWMutex
-	compacting   *atomic.Bool
+	compacting   sync.Mutex
 	compactionWg sync.WaitGroup
 
 	schema *dynparquet.Schema
 
-	prefix        string
+	dir           string
 	maxTXRecoverd []uint64
-	levels        *Node
+	levels        []Level
+	partList      *Node
 	sizes         []atomic.Int64
-	configs       []*LevelConfig
 
 	// Options
 	logger  log.Logger
@@ -52,14 +66,20 @@ type LSMMetrics struct {
 	CompactionDuration prometheus.Histogram
 }
 
-// LevelConfig is the configuration for a level in the LSM tree.
-// The Level is the sentinel node that represents the level.
-// The MaxSize is the maximum size in bytes that the level can reach before it triggers compaction into the next level.
-// The Compact function is called when the level reaches it's max size. NOTE: that this is not yet implemtened and the database will rotate the full block as normal.
+// LevelConfig is the configuration for a level in the LSM.
+// The MaxSize is the maximum size of the level in bytes before it triggers a compaction into the next level.
 type LevelConfig struct {
 	Level   SentinelType
 	MaxSize int64
-	Compact func([]parts.Part, ...parts.Option) ([]parts.Part, int64, int64, error)
+	Type    CompactionType
+	Compact Compaction
+}
+
+type Level interface {
+	Compact(parts []parts.Part, options ...parts.Option) ([]parts.Part, int64, int64, error)
+	MaxSize() int64
+	Snapshot(dir string) error
+	Reset()
 }
 
 type LSMOption func(*LSM)
@@ -99,19 +119,18 @@ func NewLSMMetrics(reg prometheus.Registerer) *LSMMetrics {
 // NewLSM returns an LSM-like index of len(levels) levels.
 // wait is a function that will block until the given transaction has been committed; this is used only during compaction to ensure
 // that all the tx in the level up to the compaction tx have been committed before compacting.
-func NewLSM(prefix string, schema *dynparquet.Schema, levels []*LevelConfig, wait func(uint64), options ...LSMOption) (*LSM, error) {
+func NewLSM(dir string, schema *dynparquet.Schema, levels []*LevelConfig, wait func(uint64), options ...LSMOption) (*LSM, error) {
 	if err := validateLevels(levels); err != nil {
 		return nil, err
 	}
 
 	lsm := &LSM{
 		schema:        schema,
-		prefix:        prefix,
+		dir:           dir,
 		maxTXRecoverd: make([]uint64, len(levels)),
-		levels:        NewList(L0),
+		partList:      NewList(L0),
 		sizes:         make([]atomic.Int64, len(levels)),
-		configs:       levels,
-		compacting:    &atomic.Bool{},
+		compacting:    sync.Mutex{},
 		logger:        log.NewNopLogger(),
 		wait:          wait,
 	}
@@ -120,18 +139,98 @@ func NewLSM(prefix string, schema *dynparquet.Schema, levels []*LevelConfig, wai
 		opt(lsm)
 	}
 
+	// Configure the LSM levels.
+	settings, recovered, err := configureLSMLevels(dir, levels, lsm.logger)
+	if err != nil {
+		return nil, err
+	}
+	defer func() {
+		for _, part := range recovered {
+			part.Release()
+		}
+	}()
+	lsm.levels = settings
+
 	// Reverse iterate (due to prepend) to create the chain of sentinel nodes.
 	for i := len(levels) - 1; i > 0; i-- {
-		lsm.levels.Sentinel(levels[i].Level)
+		lsm.partList.Sentinel(levels[i].Level)
 	}
 
 	if lsm.metrics == nil {
 		lsm.metrics = NewLSMMetrics(prometheus.NewRegistry())
 	}
 
+	// Replay the recovered parts
+	for _, part := range recovered {
+		lsm.InsertPart(part)
+	}
+
 	return lsm, nil
 }
 
+func (l *LSM) Close() error {
+	l.Lock()
+	defer l.Unlock()
+
+	// Release all the parts to free up any underlying resources.
+	l.partList.Iterate(func(node *Node) bool {
+		if node.part != nil {
+			node.part.Release()
+		}
+		return true
+	})
+
+	// Reset the levels to ensure that none of the parts are still being referenced.
+	for i := range l.levels {
+		if l.levels[i] != nil {
+			l.levels[i].Reset()
+		}
+	}
+
+	// Remove the index directory
+	if err := os.RemoveAll(l.dir); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// configureLSMLevels will configure the LSM levels. It will recover the levels from disk and return the recovered parts.
+func configureLSMLevels(dir string, levels []*LevelConfig, logger log.Logger) ([]Level, []parts.Part, error) {
+	settings := make([]Level, len(levels))
+	recovered := []parts.Part{}
+
+	// Recover in reverse order so that the highest level is recovered first.
+	// This allows us to throw away parts that were compacted into a higher level but for some reason weren't successfully removed.
+	for i := len(levels) - 1; i >= 0; i-- {
+		lvl := levels[i]
+		switch lvl.Type {
+		case CompactionTypeParquetMemory:
+			settings[i] = &inMemoryLevel{
+				maxSize: lvl.MaxSize,
+				compact: lvl.Compact,
+			}
+		case CompactionTypeParquetDisk:
+			fileCompaction, err := NewFileCompaction(filepath.Join(dir, fmt.Sprintf("L%v", i+1)), lvl.MaxSize, lvl.Compact, logger) // TODO: it would be nice to not need to inject the compact function here.
+			if err != nil {
+				return nil, nil, err
+			}
+			parts, err := fileCompaction.recover(parts.WithCompactionLevel(i + 1))
+			if err != nil {
+				level.Error(logger).Log("msg", "failed to recover parts", "err", err, "level", i+1)
+			}
+			recovered = append(recovered, parts...)
+			settings[i] = fileCompaction
+		default:
+			if i != len(levels)-1 { // Compaction type should not be set for last level
+				panic(fmt.Sprintf("unknown compaction type: %v", lvl.Type))
+			}
+		}
+	}
+
+	return settings, recovered, nil
+}
+
 // Size returns the total size of the index in bytes.
 func (l *LSM) Size() int64 {
 	var size int64
@@ -146,6 +245,30 @@ func (l *LSM) LevelSize(t SentinelType) int64 {
 	return l.sizes[t].Load()
 }
 
+// Snapshot will create a snapshot of the LSM. This will block until all compactions have completed.
+func (l *LSM) Snapshot(dir string) error {
+	l.compacting.Lock()
+	defer l.compacting.Unlock()
+
+	// call snapshot on all the levels.
+	for i, lvl := range l.levels {
+		if lvl == nil {
+			continue
+		}
+
+		lvldir := filepath.Join(dir, fmt.Sprintf("L%v", i+1))
+		if err := os.MkdirAll(lvldir, dirPerms); err != nil {
+			return err
+		}
+
+		if err := lvl.Snapshot(lvldir); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 func validateLevels(levels []*LevelConfig) error {
 	for i, l := range levels {
 		if int(l.Level) != i {
@@ -168,19 +291,20 @@ func validateLevels(levels []*LevelConfig) error {
 }
 
 func (l *LSM) MaxLevel() SentinelType {
-	return SentinelType(len(l.configs) - 1)
+	return SentinelType(len(l.levels) - 1)
 }
 
 func (l *LSM) Add(tx uint64, record arrow.Record) {
 	record.Retain()
 	size := util.TotalRecordSize(record)
-	l.levels.Insert(parts.NewArrowPart(tx, record, uint64(size), l.schema, parts.WithCompactionLevel(int(L0))))
+	l.partList.Insert(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 {
-		if l.compacting.CompareAndSwap(false, true) {
+	if l0 >= l.levels[L0].MaxSize() {
+		if l.compacting.TryLock() {
 			l.compactionWg.Add(1)
 			go func() {
+				defer l.compacting.Unlock()
 				defer l.compactionWg.Done()
 				_ = l.compact(false)
 			}()
@@ -206,9 +330,8 @@ func (l *LSM) InsertPart(part parts.Part) {
 		}
 	}
 
-	if r := part.Record(); r != nil {
-		r.Retain()
-	}
+	// Retain the part
+	part.Retain()
 
 	if tx := part.TX(); tx > l.maxTXRecoverd[level] {
 		l.maxTXRecoverd[level] = tx
@@ -226,18 +349,18 @@ func (l *LSM) String() string {
 		s += fmt.Sprintf("L%v: %d ", i, l.sizes[i].Load())
 	}
 	s += "\n"
-	s += l.levels.String()
+	s += l.partList.String()
 	return s
 }
 
 func (l *LSM) Prefixes(_ context.Context, _ string) ([]string, error) {
-	return []string{l.prefix}, nil
+	return []string{}, nil
 }
 
 func (l *LSM) Iterate(iter func(node *Node) bool) {
 	l.RLock()
 	defer l.RUnlock()
-	l.levels.Iterate(iter)
+	l.partList.Iterate(iter)
 }
 
 func (l *LSM) Scan(ctx context.Context, _ string, _ *dynparquet.Schema, filter logicalplan.Expr, tx uint64, callback func(context.Context, any) error) error {
@@ -249,7 +372,7 @@ func (l *LSM) Scan(ctx context.Context, _ string, _ *dynparquet.Schema, filter l
 		return fmt.Errorf("boolean expr: %w", err)
 	}
 	var iterError error
-	l.levels.Iterate(func(node *Node) bool {
+	l.partList.Iterate(func(node *Node) bool {
 		if node.part == nil { // encountered a sentinel node; continue on
 			return true
 		}
@@ -282,7 +405,8 @@ func (l *LSM) Scan(ctx context.Context, _ string, _ *dynparquet.Schema, filter l
 			}
 
 			if mayContainUsefulData {
-				if err := callback(ctx, rg); err != nil {
+				node.part.Retain() // Create another reference to this part
+				if err := callback(ctx, &releaseableRowGroup{DynamicRowGroup: rg, release: node.part.Release}); err != nil {
 					iterError = err
 					return false
 				}
@@ -293,10 +417,24 @@ func (l *LSM) Scan(ctx context.Context, _ string, _ *dynparquet.Schema, filter l
 	return iterError
 }
 
+type releaseableRowGroup struct {
+	dynparquet.DynamicRowGroup
+	release func()
+}
+
+func (r *releaseableRowGroup) Release() {
+	r.release()
+}
+
+type ReleaseableRowGroup interface {
+	dynparquet.DynamicRowGroup
+	Release()
+}
+
 // TODO: this should be changed to just retain the sentinel nodes in the lsm struct to do an O(1) lookup.
 func (l *LSM) findLevel(level SentinelType) *Node {
 	var list *Node
-	l.levels.Iterate(func(node *Node) bool {
+	l.partList.Iterate(func(node *Node) bool {
 		if node.part == nil && node.sentinel == level {
 			list = node
 			return false
@@ -310,7 +448,7 @@ func (l *LSM) findLevel(level SentinelType) *Node {
 // findNode returns the node that points to node.
 func (l *LSM) findNode(node *Node) *Node {
 	var list *Node
-	l.levels.Iterate(func(n *Node) bool {
+	l.partList.Iterate(func(n *Node) bool {
 		if n.next.Load() == node {
 			list = n
 			return false
@@ -324,40 +462,43 @@ func (l *LSM) findNode(node *Node) *Node {
 // EnsureCompaction forces a compaction of all levels, regardless of whether the
 // levels are below the target size.
 func (l *LSM) EnsureCompaction() error {
-	for !l.compacting.CompareAndSwap(false, true) { // TODO: should backoff retry this probably
-		// Satisfy linter with a statement.
-		continue
-	}
+	l.compacting.Lock()
+	defer l.compacting.Unlock()
 	return l.compact(true /* ignoreSizes */)
 }
 
-func (l *LSM) Rotate(level SentinelType, externalWriter func([]parts.Part) (parts.Part, int64, int64, error)) error {
-	for !l.compacting.CompareAndSwap(false, true) { // TODO: should backoff retry this probably
-		// Satisfy linter with a statement.
-		continue
-	}
-	defer l.compacting.Store(false)
+// Rotate will write all parts in the LSM into the external writer. No changes are made to the LSM.
+func (l *LSM) Rotate(externalWriter func([]parts.Part) (parts.Part, int64, int64, error)) error {
+	l.compacting.Lock()
+	defer l.compacting.Unlock()
 	start := time.Now()
 	defer func() {
 		l.metrics.CompactionDuration.Observe(time.Since(start).Seconds())
 	}()
 
-	for i := 0; i < len(l.configs)-1; i++ {
-		if err := l.merge(SentinelType(i), nil); err != nil {
-			return err
+	// Write all the parts to the external writer
+	compact := []parts.Part{}
+	l.partList.Iterate(func(node *Node) bool {
+		if node.part == nil {
+			return true
 		}
-	}
-	return l.merge(level, externalWriter)
+
+		compact = append(compact, node.part)
+		return true
+	})
+
+	_, _, _, err := externalWriter(compact)
+	return err
 }
 
 // Merge will merge the given level into an arrow record for the next level using the configured Compact function for the given level.
 // If this is the max level of the LSM an external writer must be provided to write the merged part elsewhere.
-func (l *LSM) merge(level SentinelType, externalWriter func([]parts.Part) (parts.Part, int64, int64, error)) error {
-	if int(level) > len(l.configs) {
+func (l *LSM) merge(level SentinelType) error {
+	if int(level) > len(l.levels) {
 		return fmt.Errorf("level %d does not exist", level)
 	}
-	if int(level) == len(l.configs)-1 && externalWriter == nil {
-		return fmt.Errorf("cannot merge the last level without an external writer")
+	if int(level) == len(l.levels)-1 {
+		return fmt.Errorf("cannot merge the last level")
 	}
 	l.metrics.Compactions.WithLabelValues(level.String()).Inc()
 
@@ -411,43 +552,32 @@ func (l *LSM) merge(level SentinelType, externalWriter func([]parts.Part) (parts
 	s := &Node{
 		sentinel: level + 1,
 	}
-	if externalWriter != nil {
-		_, size, _, err = externalWriter(mergeList)
-		if err != nil {
-			return err
-		}
-		// Drop compacted files from list
-		if next != nil {
-			s.next.Store(next)
-		}
-	} else {
-		compacted, size, compactedSize, err = l.configs[level].Compact(mergeList, parts.WithCompactionLevel(int(level)+1))
-		if err != nil {
-			return err
-		}
+	compacted, size, compactedSize, err = l.levels[level].Compact(mergeList, parts.WithCompactionLevel(int(level)+1))
+	if err != nil {
+		return err
+	}
 
-		// Create new list for the compacted parts.
-		compactedList := &Node{
-			part: compacted[0],
-		}
-		node := compactedList
-		for _, p := range compacted[1:] {
-			node.next.Store(&Node{
-				part: p,
-			})
-			node = node.next.Load()
-		}
-		s.next.Store(compactedList)
-		if next != nil {
-			node.next.Store(next)
-		}
-		l.sizes[level+1].Add(int64(compactedSize))
-		l.metrics.LevelSize.WithLabelValues(SentinelType(level + 1).String()).Set(float64(l.sizes[level+1].Load()))
+	// Create new list for the compacted parts.
+	compactedList := &Node{
+		part: compacted[0],
+	}
+	node := compactedList
+	for _, p := range compacted[1:] {
+		node.next.Store(&Node{
+			part: p,
+		})
+		node = node.next.Load()
+	}
+	s.next.Store(compactedList)
+	if next != nil {
+		node.next.Store(next)
 	}
+	l.sizes[level+1].Add(int64(compactedSize))
+	l.metrics.LevelSize.WithLabelValues(SentinelType(level + 1).String()).Set(float64(l.sizes[level+1].Load()))
 
 	// Replace the compacted list with the new list
 	// find the node that points to the first node in our compacted list.
-	node := l.findNode(nodeList[0])
+	node = l.findNode(nodeList[0])
 	for !node.next.CompareAndSwap(nodeList[0], s) {
 		// This can happen at most once in the scenario where a new part is added to the L0 list while we are trying to replace it.
 		node = l.findNode(nodeList[0])
@@ -462,21 +592,25 @@ func (l *LSM) merge(level SentinelType, externalWriter func([]parts.Part) (parts
 		part.Release()
 	}
 
+	// Reset the level that was just compacted
+	if level != L0 {
+		l.levels[level-1].Reset()
+	}
+
 	return nil
 }
 
 // compact is a cascading compaction routine. It will start at the lowest level and compact until the next level is either the max level or the next level does not exceed the max size.
 // compact can not be run concurrently.
 func (l *LSM) compact(ignoreSizes bool) error {
-	defer l.compacting.Store(false)
 	start := time.Now()
 	defer func() {
 		l.metrics.CompactionDuration.Observe(time.Since(start).Seconds())
 	}()
 
-	for i := 0; i < len(l.configs)-1; i++ {
-		if ignoreSizes || l.sizes[i].Load() >= l.configs[i].MaxSize {
-			if err := l.merge(SentinelType(i), nil); err != nil {
+	for i := 0; i < len(l.levels)-1; i++ {
+		if ignoreSizes || l.sizes[i].Load() >= l.levels[i].MaxSize() {
+			if err := l.merge(SentinelType(i)); err != nil {
 				level.Error(l.logger).Log("msg", "failed to merge level", "level", i, "err", err)
 				return err
 			}
diff --git a/parts/arrow.go b/parts/arrow.go
index c94ce1d..e29cbbe 100644
--- a/parts/arrow.go
+++ b/parts/arrow.go
@@ -34,6 +34,8 @@ func NewArrowPart(tx uint64, record arrow.Record, size uint64, schema *dynparque
 	return p
 }
 
+func (p *arrowPart) Retain() { p.record.Retain() }
+
 func (p *arrowPart) Record() arrow.Record {
 	return p.record
 }
diff --git a/parts/parquet.go b/parts/parquet.go
index 278d95f..a498f65 100644
--- a/parts/parquet.go
+++ b/parts/parquet.go
@@ -2,6 +2,7 @@ package parts
 
 import (
 	"fmt"
+	"sync/atomic"
 
 	"github.com/apache/arrow/go/v14/arrow"
 	"github.com/parquet-go/parquet-go"
@@ -13,6 +14,7 @@ import (
 type parquetPart struct {
 	basePart
 
+	ref *atomic.Int32
 	buf *dynparquet.SerializedBuffer
 }
 
@@ -20,8 +22,11 @@ func (p *parquetPart) Record() arrow.Record {
 	return nil
 }
 
+func (p *parquetPart) Retain() { p.ref.Add(1) }
+
 func (p *parquetPart) Release() {
-	if p.release != nil {
+	ref := p.ref.Add(-1)
+	if ref <= 0 && p.release != nil {
 		p.release()
 	}
 }
@@ -39,6 +44,7 @@ func NewParquetPart(tx uint64, buf *dynparquet.SerializedBuffer, options ...Opti
 		basePart: basePart{
 			tx: tx,
 		},
+		ref: &atomic.Int32{},
 		buf: buf,
 	}
 
@@ -46,6 +52,7 @@ func NewParquetPart(tx uint64, buf *dynparquet.SerializedBuffer, options ...Opti
 		opt(&p.basePart)
 	}
 
+	p.ref.Add(1)
 	return p
 }
 
diff --git a/parts/part.go b/parts/part.go
index 6f812ce..f253976 100644
--- a/parts/part.go
+++ b/parts/part.go
@@ -13,6 +13,7 @@ type Part interface {
 	// record part, nil is returned.
 	Record() arrow.Record
 	Release()
+	Retain()
 	SerializeBuffer(schema *dynparquet.Schema, w dynparquet.ParquetWriter) error
 	AsSerializedBuffer(schema *dynparquet.Schema) (*dynparquet.SerializedBuffer, error)
 	NumRows() int64
diff --git a/snapshot.go b/snapshot.go
index f118b7f..6c57297 100644
--- a/snapshot.go
+++ b/snapshot.go
@@ -187,7 +187,7 @@ func (db *DB) snapshotAtTX(ctx context.Context, tx uint64, writeSnapshot func(co
 	var fileSize int64
 	start := time.Now()
 	if err := func() error {
-		snapshotsDir := db.snapshotsDir()
+		snapshotsDir := snapshotDir(db, tx)
 		if err := os.MkdirAll(snapshotsDir, dirPerms); err != nil {
 			return err
 		}
@@ -252,16 +252,16 @@ func (db *DB) loadLatestSnapshotFromDir(ctx context.Context, dir string) (uint64
 	// No error should be returned from snapshotsDo.
 	_ = db.snapshotsDo(ctx, dir, func(parsedTx uint64, entry os.DirEntry) (bool, error) {
 		if err := func() error {
-			f, err := os.Open(filepath.Join(dir, entry.Name()))
+			f, err := os.Open(filepath.Join(dir, entry.Name(), snapshotFileName(parsedTx)))
 			if err != nil {
 				return err
 			}
 			defer f.Close()
-			info, err := entry.Info()
+			info, err := f.Stat()
 			if err != nil {
 				return err
 			}
-			watermark, err := LoadSnapshot(ctx, db, parsedTx, f, info.Size(), false)
+			watermark, err := LoadSnapshot(ctx, db, parsedTx, f, info.Size(), filepath.Join(dir, entry.Name()), false)
 			if err != nil {
 				return err
 			}
@@ -291,8 +291,8 @@ func (db *DB) loadLatestSnapshotFromDir(ctx context.Context, dir string) (uint64
 	return 0, fmt.Errorf("%s", errString)
 }
 
-func LoadSnapshot(ctx context.Context, db *DB, tx uint64, r io.ReaderAt, size int64, truncateWAL bool) (uint64, error) {
-	if err := loadSnapshot(ctx, db, r, size); err != nil {
+func LoadSnapshot(ctx context.Context, db *DB, tx uint64, r io.ReaderAt, size int64, dir string, truncateWAL bool) (uint64, error) {
+	if err := loadSnapshot(ctx, db, r, size, dir); err != nil {
 		return 0, err
 	}
 	watermark := tx
@@ -310,12 +310,12 @@ func (db *DB) getLatestValidSnapshotTxn(ctx context.Context) (uint64, error) {
 	// No error should be returned from snapshotsDo.
 	_ = db.snapshotsDo(ctx, dir, func(parsedTx uint64, entry os.DirEntry) (bool, error) {
 		if err := func() error {
-			f, err := os.Open(filepath.Join(dir, entry.Name()))
+			f, err := os.Open(filepath.Join(dir, entry.Name(), snapshotFileName(parsedTx)))
 			if err != nil {
 				return err
 			}
 			defer f.Close()
-			info, err := entry.Info()
+			info, err := f.Stat()
 			if err != nil {
 				return err
 			}
@@ -382,7 +382,7 @@ func (db *DB) offlineSnapshotWriter(tx uint64) func(context.Context, io.Writer)
 	}
 }
 
-func WriteSnapshot(ctx context.Context, _ uint64, db *DB, w io.Writer, offline bool) error {
+func WriteSnapshot(ctx context.Context, tx uint64, db *DB, w io.Writer, offline bool) error {
 	offW := newOffsetWriter(w)
 	w = offW
 	var tables []*Table
@@ -417,6 +417,7 @@ func WriteSnapshot(ctx context.Context, _ uint64, db *DB, w io.Writer, offline b
 			if err != nil {
 				return err
 			}
+
 			tableMeta := &snapshotpb.Table{
 				Name:   t.name,
 				Config: t.config.Load(),
@@ -458,7 +459,7 @@ func WriteSnapshot(ctx context.Context, _ uint64, db *DB, w io.Writer, offline b
 
 					// If file compaction is enabled we do not need to snapshot the parquet parts.
 					for _, cfg := range db.columnStore.indexConfig {
-						if cfg.Type == CompactionTypeParquetDisk {
+						if cfg.Type == index.CompactionTypeParquetDisk {
 							return ErrSkipPart
 						}
 					}
@@ -491,6 +492,13 @@ func WriteSnapshot(ctx context.Context, _ uint64, db *DB, w io.Writer, offline b
 				tableMeta.GranuleMetadata = append(tableMeta.GranuleMetadata, granuleMeta) // TODO: we have one part per granule now
 				return true
 			})
+
+			// Snapshot the table index. Perform this after the parts have been written to the snapshot to prevent a scenario where
+			// we snapshot the index, but before we can write the L0 parts they get compacted to the new index files.
+			if err := block.index.Snapshot(snapshotIndexDir(db, tx, t.name, block.ulid.String())); err != nil {
+				return fmt.Errorf("failed to snapshot table %s index: %w", t.name, err)
+			}
+
 			metadata.TableMetadata = append(metadata.TableMetadata, tableMeta)
 			return ascendErr
 		}(); err != nil {
@@ -574,7 +582,7 @@ func readFooter(r io.ReaderAt, size int64) (*snapshotpb.FooterData, error) {
 // loadSnapshot loads a snapshot from the given io.ReaderAt and returns the
 // txnMetadata (if any) the snapshot was created with and an error if any
 // occurred.
-func loadSnapshot(ctx context.Context, db *DB, r io.ReaderAt, size int64) error {
+func loadSnapshot(ctx context.Context, db *DB, r io.ReaderAt, size int64, dir string) error {
 	footer, err := readFooter(r, size)
 	if err != nil {
 		return err
@@ -609,6 +617,11 @@ func loadSnapshot(ctx context.Context, db *DB, r io.ReaderAt, size int64) error
 				return err
 			}
 
+			// Restore the table index from tx snapshot dir
+			if err := restoreIndexFilesFromSnapshot(db, tableMeta.Name, dir, blockUlid.String()); err != nil {
+				return err
+			}
+
 			table, err := db.table(tableMeta.Name, tableConfig, blockUlid)
 			if err != nil {
 				return err
@@ -703,14 +716,14 @@ func (db *DB) cleanupSnapshotDir(ctx context.Context, tx uint64) error {
 			// Continue.
 			return true, nil
 		}
-		if err := os.Remove(filepath.Join(dir, entry.Name())); err != nil {
+		if err := os.RemoveAll(filepath.Join(dir, entry.Name())); err != nil {
 			return false, err
 		}
 		return true, nil
 	})
 }
 
-// snapshotsDo executes the given callback with the filename of each snapshot
+// snapshotsDo executes the given callback with the directory of each snapshot
 // in dir in reverse lexicographical order (most recent snapshot first). If
 // false or an error is returned by the callback, the iteration is aborted and
 // the error returned.
@@ -725,7 +738,7 @@ func (db *DB) snapshotsDo(ctx context.Context, dir string, callback func(tx uint
 			return ctx.Err()
 		}
 		name := entry.Name()
-		if entry.IsDir() || len(name) < 20 {
+		if len(name) < 20 {
 			continue
 		}
 		parsedTx, err := getTxFromSnapshotFileName(name)
@@ -747,3 +760,51 @@ func StoreSnapshot(ctx context.Context, tx uint64, db *DB, snapshot io.Reader) e
 		return err
 	})
 }
+
+// Will restore the index files found in the given directory back to the table's index directory.
+func restoreIndexFilesFromSnapshot(db *DB, table, snapshotDir, blockID string) error {
+	// Remove the current index directory.
+	if err := os.RemoveAll(filepath.Join(db.indexDir(), table)); err != nil {
+		return fmt.Errorf("failed to remove index directory: %w", err)
+	}
+
+	snapshotIndexDir := filepath.Join(snapshotDir, "index", table, blockID)
+
+	// Restore the index files from the snapshot files.
+	return filepath.WalkDir(snapshotIndexDir, func(path string, d os.DirEntry, err error) error {
+		if err != nil {
+			return fmt.Errorf("failed to walk snapshot index directory: %w", err)
+		}
+
+		if d.IsDir() { // Level dirs expected
+			return nil
+		}
+
+		if filepath.Ext(path) != index.IndexFileExtension {
+			return nil // unknown file
+		}
+
+		// Expected file path is .../<level>/<file>
+		filename := filepath.Base(path)
+		lvl := filepath.Base(filepath.Dir(path))
+
+		if err := os.MkdirAll(filepath.Join(db.indexDir(), table, blockID, lvl), dirPerms); err != nil {
+			return err
+		}
+
+		// Hard link the file back into the index directory.
+		if err := os.Link(path, filepath.Join(db.indexDir(), table, blockID, lvl, filename)); err != nil {
+			return err
+		}
+
+		return nil
+	})
+}
+
+func snapshotDir(db *DB, tx uint64) string {
+	return filepath.Join(db.snapshotsDir(), fmt.Sprintf("%020d", tx))
+}
+
+func snapshotIndexDir(db *DB, tx uint64, table, block string) string {
+	return filepath.Join(snapshotDir(db, tx), "index", table, block)
+}
diff --git a/table.go b/table.go
index 96c8525..18996fc 100644
--- a/table.go
+++ b/table.go
@@ -53,11 +53,11 @@ var (
 
 // DefaultIndexConfig returns the default level configs used. This is a function
 // So that any modifications to the result will not affect the default config.
-func DefaultIndexConfig() []*IndexConfig {
-	return []*IndexConfig{
-		{Level: int(index.L0), MaxSize: 1024 * 1024 * 15, Type: CompactionTypeParquet},  // Compact to in-memory Parquet buffer after 15MiB of data
-		{Level: int(index.L1), MaxSize: 1024 * 1024 * 128, Type: CompactionTypeParquet}, // Compact to a single in-memory Parquet buffer after 128MiB of Parquet files
-		{Level: int(index.L2), MaxSize: 1024 * 1024 * 512},                              // Final level. Rotate after 512MiB of Parquet files
+func DefaultIndexConfig() []*index.LevelConfig {
+	return []*index.LevelConfig{
+		{Level: index.L0, MaxSize: 1024 * 1024 * 15, Type: index.CompactionTypeParquetMemory},  // Compact to in-memory Parquet buffer after 15MiB of data
+		{Level: index.L1, MaxSize: 1024 * 1024 * 128, Type: index.CompactionTypeParquetMemory}, // Compact to a single in-memory Parquet buffer after 128MiB of Parquet files
+		{Level: index.L2, MaxSize: 1024 * 1024 * 512},                                          // Final level. Rotate after 512MiB of Parquet files
 	}
 }
 
@@ -318,9 +318,7 @@ type TableBlock struct {
 	// snapshot was last triggered.
 	lastSnapshotSize atomic.Int64
 
-	index   *index.LSM
-	closers []Closer
-	syncers []Sync
+	index *index.LSM
 
 	pendingWritersWg sync.WaitGroup
 	pendingReadersWg sync.WaitGroup
@@ -493,7 +491,9 @@ func (t *Table) dropPendingBlock(block *TableBlock) {
 	block.pendingReadersWg.Wait()
 	block.pendingWritersWg.Wait()
 
-	// TODO thor call Release on all records in the block...
+	if err := block.index.Close(); err != nil {
+		level.Error(t.logger).Log("msg", "failed to close index", "err", err)
+	}
 }
 
 func (t *Table) writeBlock(block *TableBlock, skipPersist, snapshotDB bool) {
@@ -626,12 +626,7 @@ func (t *Table) RotateBlock(_ context.Context, block *TableBlock, skipPersist bo
 	// We don't check t.db.columnStore.manualBlockRotation here because this is
 	// the entry point for users to trigger a manual block rotation and they
 	// will specify through skipPersist if they want the block to be persisted.
-	go func() {
-		t.writeBlock(block, skipPersist, true)
-		if err := block.Close(true); err != nil {
-			level.Error(t.logger).Log("msg", "closing block", "err", err)
-		}
-	}()
+	go t.writeBlock(block, skipPersist, true)
 
 	return nil
 }
@@ -815,6 +810,26 @@ func (t *Table) Iterator(
 						if err != nil {
 							return err
 						}
+					case index.ReleaseableRowGroup:
+						defer t.Release()
+						if err := converter.Convert(ctx, t); err != nil {
+							return fmt.Errorf("failed to convert row group to arrow record: %v", err)
+						}
+						// This RowGroup had no relevant data. Ignore it.
+						if len(converter.Fields()) == 0 {
+							continue
+						}
+						if converter.NumRows() >= bufferSize {
+							err := func() error {
+								r := converter.NewRecord()
+								defer r.Release()
+								converter.Reset() // Reset the converter to drop any dictionaries that were built.
+								return callback(ctx, r)
+							}()
+							if err != nil {
+								return err
+							}
+						}
 					case dynparquet.DynamicRowGroup:
 						if err := converter.Convert(ctx, t); err != nil {
 							return fmt.Errorf("failed to convert row group to arrow record: %v", err)
@@ -912,6 +927,25 @@ func (t *Table) SchemaIterator(
 						if err != nil {
 							return err
 						}
+					case index.ReleaseableRowGroup:
+						if rg == nil {
+							return errors.New("received nil rowGroup") // shouldn't happen, but anyway
+						}
+						defer t.Release()
+						parquetFields := t.Schema().Fields()
+						fieldNames := make([]string, 0, len(parquetFields))
+						for _, f := range parquetFields {
+							fieldNames = append(fieldNames, f.Name())
+						}
+
+						b.Field(0).(*array.StringBuilder).AppendValues(fieldNames, nil)
+
+						record := b.NewRecord()
+						if err := callback(ctx, record); err != nil {
+							return err
+						}
+						record.Release()
+						b.Release()
 					case dynparquet.DynamicRowGroup:
 						if rg == nil {
 							return errors.New("received nil rowGroup") // shouldn't happen, but anyway
@@ -966,27 +1000,19 @@ func newTableBlock(table *Table, prevTx, tx uint64, id ulid.ULID) (*TableBlock,
 		prevTx: prevTx,
 	}
 
-	cfg, recovered, err := tb.configureLSMLevels(table.db.columnStore.indexConfig)
-	if err != nil {
-		return nil, fmt.Errorf("configure LSM levels: %w", err)
-	}
-
+	var err error
 	tb.index, err = index.NewLSM(
-		table.name,
+		filepath.Join(table.db.indexDir(), table.name, id.String()), // Any index files are found at <db.indexDir>/<table.name>/<block.id>
 		table.schema,
-		cfg,
+		table.IndexConfig(),
 		table.db.Wait,
 		index.LSMWithMetrics(table.metrics.indexMetrics),
+		index.LSMWithLogger(table.logger),
 	)
 	if err != nil {
 		return nil, fmt.Errorf("new LSM: %w", err)
 	}
 
-	// Inserts parts into the index that were recovered.
-	for _, p := range recovered {
-		tb.index.InsertPart(p)
-	}
-
 	return tb, nil
 }
 
@@ -1026,17 +1052,7 @@ func (t *TableBlock) Index() *index.LSM {
 
 // Serialize the table block into a single Parquet file.
 func (t *TableBlock) Serialize(writer io.Writer) error {
-	return t.index.Rotate(t.index.MaxLevel(), t.table.externalParquetCompaction(writer))
-}
-
-// Close the block and release all resources.
-func (t *TableBlock) Close(cleanup bool) error {
-	for _, closer := range t.closers {
-		if err := closer.Close(cleanup); err != nil {
-			return fmt.Errorf("closer: %w", err)
-		}
-	}
-	return nil
+	return t.index.Rotate(t.table.externalParquetCompaction(writer))
 }
 
 type ParquetWriter interface {
@@ -1214,105 +1230,6 @@ func (t *Table) close() {
 	t.active.pendingWritersWg.Wait()
 	t.closing = true
 	t.active.index.WaitForPendingCompactions()
-
-	if err := t.active.Close(false); err != nil {
-		level.Error(t.logger).Log("msg", "closing table block", "err", err)
-	}
-}
-
-type CompactionType int
-
-const (
-	CompactionTypeUnknown CompactionType = iota
-	CompactionTypeParquet
-	CompactionTypeParquetDisk
-)
-
-type IndexConfig struct {
-	Level   int
-	MaxSize int64
-	Type    CompactionType
-}
-
-// configureLSMLevels configures the level configs for this table block.
-func (t *TableBlock) configureLSMLevels(levels []*IndexConfig) ([]*index.LevelConfig, []parts.Part, error) {
-	config := make([]*index.LevelConfig, len(levels))
-	recovered := []parts.Part{}
-	abortRecovery := false
-
-	// Recover in reverse order so that the highest level is recovered first.
-	// This allows us to throw away parts that were compacted into a higher level but for some reason weren't successfully removed.
-	for i := len(levels) - 1; i >= 0; i-- {
-		lvl := levels[i]
-		cfg := &index.LevelConfig{
-			Level:   index.SentinelType(lvl.Level),
-			MaxSize: lvl.MaxSize,
-		}
-		switch lvl.Type {
-		case CompactionTypeParquet:
-			cfg.Compact = t.table.parquetCompaction
-		case CompactionTypeParquetDisk:
-			fileCompaction := NewFileCompaction(t.table, t.ulid, i+1)
-			if abortRecovery { // If we failed to recover an LSM file, we need to abort recovery of all lower levels as well, and let the WAL recover.
-				if err := fileCompaction.Truncate(); err != nil {
-					return nil, nil, err
-				}
-			} else {
-				parts, err := fileCompaction.recover(parts.WithCompactionLevel(i + 1))
-				if err != nil {
-					if errors.Is(ErrCompactionRecoveryFailed, err) {
-						abortRecovery = true
-					} else {
-						return nil, nil, err
-					}
-				}
-				recovered = append(recovered, parts...)
-			}
-			t.closers = append(t.closers, fileCompaction) // Append to closers so that the underlying files are closed on table close.
-			t.syncers = append(t.syncers, fileCompaction) // Append to syncers so that the underlying files are synced on wal truncation.
-			cfg.Compact = fileCompaction.writeRecordsToParquetFile
-		default:
-			if i != len(levels)-1 { // Compaction type should not be set for last level
-				panic(fmt.Sprintf("unknown compaction type: %v", lvl.Type))
-			}
-		}
-		config[i] = cfg
-	}
-
-	return config, recovered, nil
-}
-
-func (t *Table) parquetCompaction(compact []parts.Part, options ...parts.Option) ([]parts.Part, int64, int64, error) {
-	var (
-		buf                                   *dynparquet.SerializedBuffer
-		preCompactionSize, postCompactionSize int64
-		err                                   error
-	)
-	if len(compact) > 1 {
-		var b bytes.Buffer
-		preCompactionSize, err = t.compactParts(&b, compact)
-		if err != nil {
-			return nil, 0, 0, err
-		}
-		buf, err = dynparquet.ReaderFromBytes(b.Bytes())
-		if err != nil {
-			return nil, 0, 0, err
-		}
-		postCompactionSize = int64(b.Len())
-	} else if len(compact) == 1 {
-		// It's more efficient to skip compactParts if there's only one part.
-		// The only thing we want to ensure is that this part is converted to
-		// parquet if it is an arrow part.
-		singlePart := compact[0]
-		preCompactionSize = singlePart.Size()
-		buf, err = compact[0].AsSerializedBuffer(t.schema)
-		if err != nil {
-			return nil, 0, 0, err
-		}
-		postCompactionSize = buf.ParquetFile().Size()
-	}
-
-	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) {
@@ -1329,6 +1246,10 @@ func (t *Table) externalParquetCompaction(writer io.Writer) func(compact []parts
 // compactParts will compact the given parts into a Parquet file written to w.
 // It returns the size in bytes of the compacted parts.
 func (t *Table) compactParts(w io.Writer, compact []parts.Part, options ...parquet.WriterOption) (int64, error) {
+	if len(compact) == 0 {
+		return 0, nil
+	}
+
 	preCompactionSize := int64(0)
 	for _, p := range compact {
 		preCompactionSize += p.Size()
@@ -1558,3 +1479,23 @@ func (t *Table) distinctRecordsForCompaction(compact []parts.Part) ([]arrow.Reco
 	}
 	return newRecords, nil
 }
+
+// IndexConfig returns the index configuration for the table. It makes a copy of the column store index config and injects it's compactParts method.
+func (t *Table) IndexConfig() []*index.LevelConfig {
+	config := make([]*index.LevelConfig, 0, len(t.db.columnStore.indexConfig))
+	for i, c := range t.db.columnStore.indexConfig {
+		compactFunc := t.compactParts
+		if i == len(t.db.columnStore.indexConfig)-1 {
+			// The last level is the in-memory level, which is never compacted.
+			compactFunc = nil
+		}
+		config = append(config, &index.LevelConfig{
+			Level:   c.Level,
+			MaxSize: c.MaxSize,
+			Type:    c.Type,
+			Compact: compactFunc, // TODO: this is bad and it should feel bad. We shouldn't need the table object to define how parts are compacted. Refactor needed.
+		})
+	}
+
+	return config
+}


Test patch:
diff --git a/db_test.go b/db_test.go
index 81fc1ec..38e0623 100644
--- a/db_test.go
+++ b/db_test.go
@@ -2228,16 +2228,18 @@ func Test_DB_WithParquetDiskCompaction(t *testing.T) {
 	logger := newTestLogger(t)
 
 	cfg := DefaultIndexConfig()
-	cfg[0].Type = CompactionTypeParquetDisk // Create disk compaction
-	cfg[1].Type = CompactionTypeParquetDisk
+	cfg[0].Type = index.CompactionTypeParquetDisk // Create disk compaction
+	cfg[1].Type = index.CompactionTypeParquetDisk
 	c, err := New(
 		WithLogger(logger),
 		WithIndexConfig(cfg),
+		WithWAL(),
+		WithStoragePath(t.TempDir()),
 	)
+	require.NoError(t, err)
 	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)
@@ -2280,12 +2282,37 @@ func Test_DB_PersistentDiskCompaction(t *testing.T) {
 	)
 	logger := newTestLogger(t)
 
+	type write struct {
+		n                int
+		ensureCompaction bool
+		op               func(table *Table, dir string)
+	}
+
+	defaultWrites := []write{
+		{
+			n:                100,
+			ensureCompaction: true,
+		},
+		{
+			n:                100,
+			ensureCompaction: true,
+		},
+		{
+			n:                100,
+			ensureCompaction: true,
+		},
+		{
+			n: 100,
+		},
+	}
+
 	tests := map[string]struct {
 		beforeReplay func(table *Table, dir string)
 		beforeClose  func(table *Table, dir string)
 		lvl2Parts    int
 		lvl1Parts    int
 		finalRows    int64
+		writes       []write
 	}{
 		"happy path": {
 			beforeReplay: func(_ *Table, _ string) {},
@@ -2293,15 +2320,18 @@ func Test_DB_PersistentDiskCompaction(t *testing.T) {
 			lvl2Parts:    3,
 			lvl1Parts:    1,
 			finalRows:    1200,
+			writes:       defaultWrites,
 		},
 		"corrupted L1": {
 			beforeClose: func(table *Table, dir string) {},
 			lvl2Parts:   3,
 			lvl1Parts:   1,
 			finalRows:   1200,
+			writes:      defaultWrites,
 			beforeReplay: func(table *Table, dir string) {
 				// Corrupt the LSM file; this should trigger a recovery from the WAL
-				levelFile := filepath.Join(dir, "databases", "test", "index", "test", fmt.Sprintf("L1-%s.parquet", table.active.ulid.String()))
+				l1Dir := filepath.Join(table.db.indexDir(), "test", table.active.ulid.String(), "L1")
+				levelFile := filepath.Join(l1Dir, fmt.Sprintf("00000000000000000000%s", index.IndexFileExtension))
 				info, err := os.Stat(levelFile)
 				require.NoError(t, err)
 				require.NoError(t, os.Truncate(levelFile, info.Size()-1)) // truncate the last byte to pretend the write didn't finish
@@ -2309,46 +2339,210 @@ func Test_DB_PersistentDiskCompaction(t *testing.T) {
 		},
 		"corrupted L2": {
 			beforeClose: func(table *Table, dir string) {},
-			lvl2Parts:   3,
-			lvl1Parts:   1,
+			lvl2Parts:   4,
+			lvl1Parts:   0,
 			finalRows:   1200,
 			beforeReplay: func(table *Table, dir string) {
 				// Corrupt the LSM file; this should trigger a recovery from the WAL
-				levelFile := filepath.Join(dir, "databases", "test", "index", "test", fmt.Sprintf("L2-%s.parquet", table.active.ulid.String()))
+				l2Dir := filepath.Join(table.db.indexDir(), "test", table.active.ulid.String(), "L2")
+				levelFile := filepath.Join(l2Dir, fmt.Sprintf("00000000000000000000%s", index.IndexFileExtension))
 				info, err := os.Stat(levelFile)
 				require.NoError(t, err)
 				require.NoError(t, os.Truncate(levelFile, info.Size()-1)) // truncate the last byte to pretend the write didn't finish
 			},
+			writes: []write{
+				{
+					n:                100,
+					ensureCompaction: true,
+				},
+				{
+					n:                100,
+					ensureCompaction: true,
+				},
+				{
+					n:                100,
+					ensureCompaction: true,
+				},
+				{
+					n:                100,
+					ensureCompaction: true,
+				},
+			},
 		},
-		"untruncated L1": { // Test that if we fail to truncate the L1 file, we can still recover
-			lvl2Parts:    4,
-			lvl1Parts:    1,
-			finalRows:    1200,
-			beforeReplay: func(_ *Table, _ string) {},
-			beforeClose: func(table *Table, dir string) { // trigger a compaction, but save off the L1 file and move it back. This mimics a failure to truncate the L1 file
-				// Save the L1 file
-				levelFile := filepath.Join(dir, "databases", "test", "index", "test", fmt.Sprintf("L1-%s.parquet", table.active.ulid.String()))
-				src, err := os.Open(levelFile)
+		"L1 cleanup failure with corruption": { // Compaction from L1->L2 happens, but L1 is not cleaned up and L2 gets corrupted
+			beforeClose: func(table *Table, dir string) {},
+			lvl2Parts:   1,
+			lvl1Parts:   0,
+			finalRows:   300,
+			beforeReplay: func(table *Table, dir string) {
+				// Copy the L2 file back to L1 to simulate that it was not cleaned up
+				l2Dir := filepath.Join(table.db.indexDir(), "test", table.active.ulid.String(), "L2")
+				l2File := filepath.Join(l2Dir, fmt.Sprintf("00000000000000000000%s", index.IndexFileExtension))
+				l2, err := os.Open(l2File)
+				require.NoError(t, err)
+
+				l1Dir := filepath.Join(table.db.indexDir(), "test", table.active.ulid.String(), "L1")
+				l1File := filepath.Join(l1Dir, fmt.Sprintf("00000000000000000000%s", index.IndexFileExtension))
+				l1, err := os.Create(l1File)
 				require.NoError(t, err)
 
-				saveFile := filepath.Join(dir, "databases", "test", "index", "test", fmt.Sprintf("L1-%s.parquet.sav", table.active.ulid.String()))
-				dst, err := os.Create(saveFile)
+				_, err = io.Copy(l1, l2)
 				require.NoError(t, err)
+				require.NoError(t, l2.Close())
+				require.NoError(t, l1.Close())
 
-				_, err = io.Copy(dst, src)
+				// Corrupt the L2 file
+				info, err := os.Stat(l2File)
+				require.NoError(t, err)
+				require.NoError(t, os.Truncate(l2File, info.Size()-1)) // truncate the last byte to pretend the write didn't finish
+			},
+			writes: []write{
+				{
+					n:                100,
+					ensureCompaction: true,
+				},
+			},
+		},
+		"L1 cleanup failure": { // Compaction from L1->L2 happens, but L1 is not cleaned up
+			beforeClose: func(table *Table, dir string) {},
+			lvl2Parts:   1,
+			lvl1Parts:   0,
+			finalRows:   300,
+			beforeReplay: func(table *Table, dir string) {
+				// Copy the L2 file back to L1 to simulate that it was not cleaned up
+				l2Dir := filepath.Join(table.db.indexDir(), "test", table.active.ulid.String(), "L2")
+				l2File := filepath.Join(l2Dir, fmt.Sprintf("00000000000000000000%s", index.IndexFileExtension))
+				l2, err := os.Open(l2File)
 				require.NoError(t, err)
-				dst.Close()
-				src.Close()
 
-				// This truncates the L1 file
-				require.NoError(t, table.EnsureCompaction())
+				l1Dir := filepath.Join(table.db.indexDir(), "test", table.active.ulid.String(), "L1")
+				l1File := filepath.Join(l1Dir, fmt.Sprintf("00000000000000000000%s", index.IndexFileExtension))
+				l1, err := os.Create(l1File)
+				require.NoError(t, err)
 
-				// Move the L1 file back
-				require.NoError(t, os.Rename(saveFile, levelFile))
+				_, err = io.Copy(l1, l2)
+				require.NoError(t, err)
+				require.NoError(t, l2.Close())
+				require.NoError(t, l1.Close())
+			},
+			writes: []write{
+				{
+					n:                100,
+					ensureCompaction: true,
+				},
+			},
+		},
+		"L1 cleanup failure with corruption after snapshot": { // Snapshot happens after L1 compaction; Then compaction happens from L1->L2, but L1 is not cleaned up and L2 gets corrupted
+			beforeClose: func(table *Table, dir string) {},
+			lvl2Parts:   1,
+			lvl1Parts:   0,
+			finalRows:   303,
+			beforeReplay: func(table *Table, dir string) {
+				// Move the save files back into L1 to simulate an unssuccessful cleanup after L1 compaction
+				l1Dir := filepath.Join(table.db.indexDir(), "test", table.active.ulid.String(), "L1")
+				saveDir := filepath.Join(table.db.indexDir(), "save")
+				require.NoError(t, filepath.WalkDir(saveDir, func(path string, info fs.DirEntry, err error) error {
+					if filepath.Ext(path) == index.IndexFileExtension {
+						// Move the save file back to the original file
+						sv, err := os.Open(path)
+						require.NoError(t, err)
+						defer sv.Close()
+
+						lvl, err := os.Create(filepath.Join(l1Dir, filepath.Base(path)))
+						require.NoError(t, err)
+						defer lvl.Close()
+
+						_, err = io.Copy(lvl, sv)
+						require.NoError(t, err)
+					}
+					return nil
+				}))
+
+				// Corrupt the L2 file to simulate that the failure occurred when writing the L2 file
+				l2Dir := filepath.Join(table.db.indexDir(), "test", table.active.ulid.String(), "L2")
+				l2File := filepath.Join(l2Dir, fmt.Sprintf("00000000000000000000%s", index.IndexFileExtension))
+				info, err := os.Stat(l2File)
+				require.NoError(t, err)
+				require.NoError(t, os.Truncate(l2File, info.Size()-1)) // truncate the last byte to pretend the write didn't finish
+			},
+			writes: []write{
+				{ // Get to L1 compaction
+					n: 100,
+					op: func(table *Table, dir string) {
+						require.Eventually(t, func() bool {
+							return table.active.index.LevelSize(index.L1) != 0
+						}, time.Second, time.Millisecond*10)
+					},
+				},
+				{ // trigger a snapshot (leaving the only valid data in L1)
+					op: func(table *Table, dir string) {
+						tx := table.db.highWatermark.Load()
+						require.NoError(t, table.db.snapshotAtTX(context.Background(), tx, table.db.snapshotWriter(tx)))
+					},
+				},
+				{
+					op: func(table *Table, dir string) { // save the L1 files to restore them
+						l1Dir := filepath.Join(table.db.indexDir(), "test", table.active.ulid.String(), "L1")
+						saveDir := filepath.Join(dir, "save")
+						require.NoError(t, os.MkdirAll(saveDir, 0o755)) // Create a directory to save the files
+						require.NoError(t, filepath.WalkDir(l1Dir, func(path string, info fs.DirEntry, err error) error {
+							if filepath.Ext(path) == index.IndexFileExtension {
+								lvl, err := os.Open(path)
+								require.NoError(t, err)
+								defer lvl.Close()
+
+								sv, err := os.Create(filepath.Join(saveDir, filepath.Base(path)))
+								require.NoError(t, err)
+								defer sv.Close()
+
+								_, err = io.Copy(sv, lvl)
+								require.NoError(t, err)
+							}
+							return nil
+						}))
+					},
+				},
+				{ // Compact L1->L2
+					n:                1, // New write to "trigger" the compaction
+					ensureCompaction: true,
+				},
 			},
 		},
 		"snapshot": {
 			beforeReplay: func(_ *Table, _ string) {},
+			beforeClose: func(table *Table, dir string) {
+				// trigger a snapshot
+				success := false
+				table.db.snapshot(context.Background(), false, func() {
+					success = true
+				})
+				require.True(t, success)
+
+				// Write more to the table and trigger a compaction
+				samples := dynparquet.NewTestSamples()
+				for j := 0; j < 100; j++ {
+					r, err := samples.ToRecord()
+					require.NoError(t, err)
+					_, err = table.InsertRecord(context.Background(), r)
+					require.NoError(t, err)
+				}
+				require.Eventually(t, func() bool {
+					return table.active.index.LevelSize(index.L1) != 0
+				}, time.Second, time.Millisecond*10)
+			},
+			lvl2Parts: 3,
+			lvl1Parts: 2,
+			finalRows: 1500,
+			writes:    defaultWrites,
+		},
+		"corruption after snapshot": {
+			beforeReplay: func(table *Table, dir string) {
+				l1Dir := filepath.Join(table.db.indexDir(), "test", table.active.ulid.String(), "L1")
+				levelFile := filepath.Join(l1Dir, fmt.Sprintf("00000000000000000001%v", index.IndexFileExtension))
+				info, err := os.Stat(levelFile)
+				require.NoError(t, err)
+				require.NoError(t, os.Truncate(levelFile, info.Size()-1)) // truncate the last byte to pretend the write didn't finish
+			},
 			beforeClose: func(table *Table, dir string) {
 				// trigger a snapshot
 				tx := table.db.highWatermark.Load()
@@ -2369,16 +2563,19 @@ func Test_DB_PersistentDiskCompaction(t *testing.T) {
 			lvl2Parts: 3,
 			lvl1Parts: 2,
 			finalRows: 1500,
+			writes:    defaultWrites,
 		},
 	}
 
 	for name, test := range tests {
+		test := test
 		t.Run(name, func(t *testing.T) {
+			t.Parallel()
 			dir := t.TempDir()
-			cfg := []*IndexConfig{
-				{Level: int(index.L0), MaxSize: 25700, Type: CompactionTypeParquetDisk},
-				{Level: int(index.L1), MaxSize: 1024 * 1024 * 128, Type: CompactionTypeParquetDisk},
-				{Level: int(index.L2), MaxSize: 1024 * 1024 * 512},
+			cfg := []*index.LevelConfig{
+				{Level: index.L0, MaxSize: 25700, 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),
@@ -2398,71 +2595,51 @@ func Test_DB_PersistentDiskCompaction(t *testing.T) {
 			samples := dynparquet.NewTestSamples()
 
 			ctx := context.Background()
-			compactions := 3 // 3 compacted parts in L2 file
-			for i := 0; i < compactions; i++ {
-				for j := 0; j < 100; j++ {
+			insertedRows := int64(0)
+			for _, w := range test.writes {
+				for j := 0; j < w.n; j++ {
 					r, err := samples.ToRecord()
 					require.NoError(t, err)
 					_, err = table.InsertRecord(ctx, r)
+					insertedRows += r.NumRows()
 					require.NoError(t, err)
 				}
-				require.NoError(t, table.EnsureCompaction())
+				if w.ensureCompaction {
+					require.NoError(t, table.EnsureCompaction())
+				}
+				if w.op != nil {
+					w.op(table, db.indexDir())
+				}
 			}
 
-			for j := 0; j < 100; j++ {
-				r, err := samples.ToRecord()
-				require.NoError(t, err)
-				_, err = table.InsertRecord(ctx, r)
+			if test.lvl1Parts > 0 {
+				require.Eventually(t, func() bool {
+					return table.active.index.LevelSize(index.L1) != 0
+				}, 30*time.Second, time.Millisecond*10)
+			}
+
+			rows := func(db *DB) 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)
+				return rows
 			}
-			require.Eventually(t, func() bool {
-				return table.active.index.LevelSize(index.L1) != 0
-			}, time.Second, time.Millisecond*10)
 
 			// Ensure that disk compacted data can be recovered
-			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, int64((300*compactions)+300), rows)
+			require.Equal(t, insertedRows, rows(db))
 
 			test.beforeClose(table, dir)
-			id := table.active.ulid
 
 			// Close the database
-			wm := db.highWatermark.Load()
 			require.NoError(t, c.Close())
 
-			fc := NewFileCompaction(table, id, 2)
-			p, err := fc.recover()
-			require.NoError(t, err)
-			require.Equal(t, test.lvl2Parts, len(p))
-
-			for i := 0; i < compactions; i++ {
-				buf, err := p[i].AsSerializedBuffer(nil)
-				require.NoError(t, err)
-				tx, ok := buf.ParquetFile().Lookup("compaction_tx")
-				require.True(t, ok)
-				require.Equal(t, fmt.Sprintf("%v", test.lvl2Parts*100-(i*100)+1), tx)
-			}
-
-			fc = NewFileCompaction(table, id, 1)
-			p, err = fc.recover()
-			require.NoError(t, err)
-			require.Equal(t, test.lvl1Parts, len(p))
-
-			buf, err := p[0].AsSerializedBuffer(nil)
-			require.NoError(t, err)
-			tx, ok := buf.ParquetFile().Lookup("compaction_tx")
-			require.True(t, ok)
-			require.Equal(t, fmt.Sprintf("%v", wm), tx)
-
 			// Run the beforeReplay hook
 			test.beforeReplay(table, dir)
 
@@ -2479,16 +2656,30 @@ func Test_DB_PersistentDiskCompaction(t *testing.T) {
 			})
 			db, err = c.DB(context.Background(), "test")
 			require.NoError(t, err)
-
-			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
-				})
+			table, err = db.Table("test", config)
 			require.NoError(t, err)
-			require.Equal(t, test.finalRows, rows)
+
+			// Check the final row count
+			require.Equal(t, test.finalRows, rows(db))
+
+			// Now write more data after a recovery into the levels and ensure that it can be read back.
+
+			// L1 writes
+			l1Before := table.active.index.LevelSize(index.L1)
+			for j := 0; j < 100; j++ {
+				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) != l1Before
+			}, 30*time.Second, time.Millisecond*10)
+			require.Equal(t, test.finalRows+300, rows(db))
+
+			// Move that data to L2
+			require.NoError(t, table.EnsureCompaction())
+			require.Equal(t, test.finalRows+300, rows(db))
 		})
 	}
 }
@@ -2504,10 +2695,10 @@ func Test_DB_PersistentDiskCompaction_BlockRotation(t *testing.T) {
 	sinksource := NewDefaultObjstoreBucket(bucket)
 
 	dir := t.TempDir()
-	cfg := []*IndexConfig{
-		{Level: int(index.L0), MaxSize: 25700, Type: CompactionTypeParquetDisk},
-		{Level: int(index.L1), MaxSize: 1024 * 1024 * 128, Type: CompactionTypeParquetDisk},
-		{Level: int(index.L2), MaxSize: 1024 * 1024 * 512},
+	cfg := []*index.LevelConfig{
+		{Level: index.L0, MaxSize: 25700, 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),
diff --git a/index/lsm_test.go b/index/lsm_test.go
index 378b083..66cecdd 100644
--- a/index/lsm_test.go
+++ b/index/lsm_test.go
@@ -1,7 +1,6 @@
 package index
 
 import (
-	"bytes"
 	"context"
 	"errors"
 	"io"
@@ -20,21 +19,7 @@ import (
 	"github.com/polarsignals/frostdb/parts"
 )
 
-func parquetCompaction(compact []parts.Part, _ ...parts.Option) ([]parts.Part, int64, int64, error) {
-	b := &bytes.Buffer{}
-	size, err := compactParts(b, compact)
-	if err != nil {
-		return nil, 0, 0, err
-	}
-
-	buf, err := dynparquet.ReaderFromBytes(b.Bytes())
-	if err != nil {
-		return nil, 0, 0, err
-	}
-	return []parts.Part{parts.NewParquetPart(0, buf)}, size, int64(b.Len()), nil
-}
-
-func compactParts(w io.Writer, compact []parts.Part) (int64, error) {
+func compactParts(w io.Writer, compact []parts.Part, _ ...parquet.WriterOption) (int64, error) {
 	schema := dynparquet.NewSampleSchema()
 	bufs := []dynparquet.DynamicRowGroup{}
 	var size int64
@@ -80,7 +65,7 @@ func compactParts(w io.Writer, compact []parts.Part) (int64, error) {
 func check(t *testing.T, lsm *LSM, records, buffers int) {
 	t.Helper()
 	seen := map[SentinelType]bool{}
-	lsm.levels.Iterate(func(node *Node) bool {
+	lsm.partList.Iterate(func(node *Node) bool {
 		if node.part == nil {
 			if seen[node.sentinel] {
 				t.Fatal("duplicate sentinel")
@@ -107,8 +92,8 @@ func check(t *testing.T, lsm *LSM, records, buffers int) {
 func Test_LSM_Basic(t *testing.T) {
 	t.Parallel()
 	lsm, err := NewLSM("test", nil, []*LevelConfig{
-		{Level: L0, MaxSize: 1024 * 1024 * 1024, Compact: parquetCompaction},
-		{Level: L1, MaxSize: 1024 * 1024 * 1024, Compact: parquetCompaction},
+		{Level: L0, MaxSize: 1024 * 1024 * 1024, Type: CompactionTypeParquetMemory, Compact: compactParts},
+		{Level: L1, MaxSize: 1024 * 1024 * 1024, Type: CompactionTypeParquetMemory, Compact: compactParts},
 		{Level: L2, MaxSize: 1024 * 1024 * 1024},
 	},
 		func(uint64) {},
@@ -123,27 +108,27 @@ func Test_LSM_Basic(t *testing.T) {
 	lsm.Add(2, r)
 	lsm.Add(3, r)
 	check(t, lsm, 3, 0)
-	require.NoError(t, lsm.merge(L0, nil))
+	require.NoError(t, lsm.merge(L0))
 	check(t, lsm, 0, 1)
 	lsm.Add(4, r)
 	check(t, lsm, 1, 1)
 	lsm.Add(5, r)
 	check(t, lsm, 2, 1)
-	require.NoError(t, lsm.merge(L0, nil))
+	require.NoError(t, lsm.merge(L0))
 	check(t, lsm, 0, 2)
 	lsm.Add(6, r)
 	check(t, lsm, 1, 2)
-	require.NoError(t, lsm.merge(L1, nil))
+	require.NoError(t, lsm.merge(L1))
 	check(t, lsm, 1, 1)
-	require.NoError(t, lsm.merge(L0, nil))
+	require.NoError(t, lsm.merge(L0))
 	check(t, lsm, 0, 2)
 }
 
 func Test_LSM_DuplicateSentinel(t *testing.T) {
 	t.Parallel()
 	lsm, err := NewLSM("test", nil, []*LevelConfig{
-		{Level: L0, MaxSize: 1024 * 1024 * 1024, Compact: parquetCompaction},
-		{Level: L1, MaxSize: 1024 * 1024 * 1024, Compact: parquetCompaction},
+		{Level: L0, MaxSize: 1024 * 1024 * 1024, Type: CompactionTypeParquetMemory, Compact: compactParts},
+		{Level: L1, MaxSize: 1024 * 1024 * 1024, Type: CompactionTypeParquetMemory, Compact: compactParts},
 		{Level: L2, MaxSize: 1024 * 1024 * 1024},
 	},
 		func(uint64) {},
@@ -158,16 +143,16 @@ func Test_LSM_DuplicateSentinel(t *testing.T) {
 	lsm.Add(2, r)
 	lsm.Add(3, r)
 	check(t, lsm, 3, 0)
-	require.NoError(t, lsm.merge(L0, nil))
+	require.NoError(t, lsm.merge(L0))
 	check(t, lsm, 0, 1)
-	require.NoError(t, lsm.merge(L0, nil))
+	require.NoError(t, lsm.merge(L0))
 	check(t, lsm, 0, 1)
 }
 
 func Test_LSM_Compaction(t *testing.T) {
 	t.Parallel()
 	lsm, err := NewLSM("test", nil, []*LevelConfig{
-		{Level: L0, MaxSize: 1, Compact: parquetCompaction},
+		{Level: L0, MaxSize: 1, Type: CompactionTypeParquetMemory, Compact: compactParts},
 		{Level: L1, MaxSize: 1024 * 1024 * 1024},
 	},
 		func(uint64) {},
@@ -187,10 +172,10 @@ func Test_LSM_Compaction(t *testing.T) {
 func Test_LSM_CascadeCompaction(t *testing.T) {
 	t.Parallel()
 	lsm, err := NewLSM("test", nil, []*LevelConfig{
-		{Level: L0, MaxSize: 257, Compact: parquetCompaction},
-		{Level: L1, MaxSize: 2281, Compact: parquetCompaction},
-		{Level: L2, MaxSize: 2281, Compact: parquetCompaction},
-		{Level: 3, MaxSize: 2281, Compact: parquetCompaction},
+		{Level: L0, MaxSize: 257, Type: CompactionTypeParquetMemory, Compact: compactParts},
+		{Level: L1, MaxSize: 2281, Type: CompactionTypeParquetMemory, Compact: compactParts},
+		{Level: L2, MaxSize: 2281, Type: CompactionTypeParquetMemory, Compact: compactParts},
+		{Level: 3, MaxSize: 2281, Type: CompactionTypeParquetMemory, Compact: compactParts},
 		{Level: 4, MaxSize: 2281},
 	},
 		func(uint64) {},
@@ -222,8 +207,8 @@ func Test_LSM_CascadeCompaction(t *testing.T) {
 func Test_LSM_InOrderInsert(t *testing.T) {
 	t.Parallel()
 	lsm, err := NewLSM("test", nil, []*LevelConfig{
-		{Level: L0, MaxSize: 1024 * 1024 * 1024, Compact: parquetCompaction},
-		{Level: L1, MaxSize: 1024 * 1024 * 1024, Compact: parquetCompaction},
+		{Level: L0, MaxSize: 1024 * 1024 * 1024, Type: CompactionTypeParquetMemory, Compact: compactParts},
+		{Level: L1, MaxSize: 1024 * 1024 * 1024, Type: CompactionTypeParquetMemory, Compact: compactParts},
 		{Level: L2, MaxSize: 1024 * 1024 * 1024},
 	},
 		func(uint64) {},
diff --git a/table_test.go b/table_test.go
index 1bfe363..127700f 100644
--- a/table_test.go
+++ b/table_test.go
@@ -438,9 +438,9 @@ func Test_Table_NewTableValidSplitSize(t *testing.T) {
 
 func Test_Table_Bloomfilter(t *testing.T) {
 	c, table := basicTable(t, WithIndexConfig(
-		[]*IndexConfig{
-			{Level: int(index.L0), MaxSize: 452, Type: CompactionTypeParquet}, // NOTE: 452 is the current size of the 3 records that are inserted
-			{Level: int(index.L1), MaxSize: 100000},
+		[]*index.LevelConfig{
+			{Level: index.L0, MaxSize: 452, Type: index.CompactionTypeParquetMemory}, // NOTE: 452 is the current size of the 3 records that are inserted
+			{Level: index.L1, MaxSize: 100000},
 		},
 	))
 	defer c.Close()
@@ -1087,9 +1087,9 @@ func Test_Table_DynamicColumnNotDefined(t *testing.T) {
 }
 
 func TestTableUniquePrimaryIndex(t *testing.T) {
-	c, err := New(WithIndexConfig([]*IndexConfig{
-		{Level: int(index.L0), MaxSize: 180, Type: CompactionTypeParquet},
-		{Level: int(index.L1), MaxSize: 1 * TiB},
+	c, err := New(WithIndexConfig([]*index.LevelConfig{
+		{Level: index.L0, MaxSize: 180, Type: index.CompactionTypeParquetMemory},
+		{Level: index.L1, MaxSize: 1 * TiB},
 	}))
 	require.NoError(t, err)
 	defer c.Close()