btree insert root

pkg/bptree/bptree.go

@@ -36,8 +36,7 @@ func (t *BPTree) root() (*BPTreeNode, pointer.MemoryPointer, error) {
 type TraversalRecord struct {
 	node  *BPTreeNode
 	index int
-	// the offset is useful so we know which page to free when we split
-	ptr pointer.MemoryPointer
+	ptr   pointer.MemoryPointer
 }
 
 type TraversalIterator struct {
@@ -274,6 +273,7 @@ func (t *BPTree) Insert(key pointer.ReferencedValue, value pointer.MemoryPointer
 		if err != nil {
 			return err
 		}
+
 		return t.MetaPage.SetRoot(pointer.MemoryPointer{Offset: uint64(offset), Length: uint32(len(buf))})
 	}
 

pkg/bptree/bptree_test.go

@@ -82,6 +82,7 @@ func TestBPTree(t *testing.T) {
 		if err := tree.Insert(pointer.ReferencedValue{Value: []byte("hello")}, pointer.MemoryPointer{Offset: 1, Length: 5}); err != nil {
 			t.Fatal(err)
 		}
+
 		k, v, err := tree.Find(pointer.ReferencedValue{Value: []byte("hello")})
 		if err != nil {
 			t.Fatal(err)

pkg/bptree/node.go

@@ -43,7 +43,7 @@ func (n *BPTreeNode) NumPointers() int {
 func (n *BPTreeNode) Size() int64 {
 	size := 4 // number of keys
 	for _, k := range n.Keys {
-		o := encoding.SizeVarint(uint64(k.DataPointer.Offset))
+		o := encoding.SizeVarint(k.DataPointer.Offset)
 		l := encoding.SizeVarint(uint64(k.DataPointer.Length))
 		size += l + o
 

pkg/btree/btree.go

@@ -1,10 +1,13 @@
 package btree
 
 import (
+	"fmt"
+	"github.com/kevmo314/appendable/pkg/hnsw"
 	"github.com/kevmo314/appendable/pkg/metapage"
 	"github.com/kevmo314/appendable/pkg/pagefile"
 	"github.com/kevmo314/appendable/pkg/pointer"
 	"io"
+	"slices"
 )
 
 type BTree struct {
@@ -18,7 +21,7 @@ type BTree struct {
 
 func (t *BTree) root() (*BTreeNode, pointer.MemoryPointer, error) {
 	mp, err := t.MetaPage.Root()
-	if err != nil {
+	if err != nil || mp.Length == 0 {
 		return nil, mp, err
 	}
 
@@ -48,3 +51,159 @@ func (t *BTree) readNode(offset uint64) (*BTreeNode, error) {
 
 	return node, nil
 }
+
+func (t *BTree) Insert(key pointer.ReferencedId, vector hnsw.Point) error {
+	root, rootOffset, err := t.root()
+	if err != nil {
+		return fmt.Errorf("read root node: %w", err)
+	}
+
+	if root == nil {
+		node := &BTreeNode{Width: t.Width, VectorDim: t.VectorDim}
+		node.Ids = []pointer.ReferencedId{key}
+		node.Vectors = []hnsw.Point{vector}
+		node.Offsets = make([]uint64, 0)
+
+		buf, err := node.MarshalBinary()
+		if err != nil {
+			return err
+		}
+		offset, err := t.PageFile.NewPage(buf)
+		if err != nil {
+			return err
+		}
+		return t.MetaPage.SetRoot(pointer.MemoryPointer{Offset: uint64(offset), Length: uint32(len(buf))})
+	}
+
+	parent, parentOffset := root, rootOffset.Offset
+	for !parent.Leaf() {
+		index, found := slices.BinarySearchFunc(parent.Ids, key, pointer.CompareReferencedIds)
+
+		if found {
+			panic("cannot insert duplicate key")
+		}
+
+		loffset := parent.Offsets[index]
+		child, err := t.readNode(loffset)
+		if err != nil {
+			return err
+		}
+
+		if int(child.Size()) > t.PageFile.PageSize() {
+			// split node here
+			mid := len(child.Ids) / 2
+			midKey := child.Ids[mid]
+
+			rightChild := &BTreeNode{Width: t.Width, VectorDim: t.VectorDim}
+			if !child.Leaf() {
+				rightChild.Offsets = child.Offsets[mid+1:]
+				child.Offsets = child.Offsets[:mid]
+			}
+			rightChild.Vectors = child.Vectors[mid+1:]
+			rightChild.Ids = child.Ids[mid+1:]
+
+			rbuf, err := rightChild.MarshalBinary()
+			if err != nil {
+				return err
+			}
+			roffset, err := t.PageFile.NewPage(rbuf)
+			if err != nil {
+				return err
+			}
+
+			// shrink left child (child)
+			child.Ids = child.Ids[:mid]
+			child.Vectors = child.Vectors[:mid]
+			if _, err := t.PageFile.Seek(int64(loffset), io.SeekStart); err != nil {
+				return err
+			}
+
+			if _, err := child.WriteTo(t.PageFile); err != nil {
+				return err
+			}
+
+			// update parent to include new key and store left right offsets
+			if index == len(parent.Ids) {
+				parent.Ids = append(parent.Ids, midKey)
+			} else {
+				parent.Ids = append(parent.Ids[:index+1], parent.Ids[index:]...)
+				parent.Ids[index] = midKey
+			}
+
+			parent.Offsets = append(parent.Offsets[:index+2], parent.Offsets[:index+1]...)
+			parent.Offsets[index+1] = uint64(roffset)
+			if _, err := t.PageFile.Seek(int64(parentOffset), io.SeekStart); err != nil {
+				return err
+			}
+			if _, err := parent.WriteTo(t.PageFile); err != nil {
+				return err
+			}
+
+			if pointer.CompareReferencedIds(midKey, key) == 1 {
+				parent, parentOffset = child, loffset
+			} else {
+				parent, parentOffset = rightChild, uint64(roffset)
+			}
+		} else {
+			if _, err := t.PageFile.Seek(int64(parentOffset), io.SeekStart); err != nil {
+				return err
+			}
+			if _, err := parent.WriteTo(t.PageFile); err != nil {
+				return err
+			}
+			parent, parentOffset = child, loffset
+		}
+	}
+
+	index, found := slices.BinarySearchFunc(parent.Ids, key, pointer.CompareReferencedIds)
+	if found {
+		panic("cannot insert duplicate key")
+	}
+
+	parent.Ids = append(parent.Ids[:index+1], parent.Ids[index:]...)
+	parent.Ids[index] = key
+
+	parent.Vectors = append(parent.Vectors[:index+1], parent.Vectors[index:]...)
+	parent.Vectors[index] = vector
+
+	if _, err := t.PageFile.Seek(int64(parentOffset), io.SeekStart); err != nil {
+		return err
+	}
+	if _, err := parent.WriteTo(t.PageFile); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (t *BTree) Find(key pointer.ReferencedId) (pointer.ReferencedId, pointer.MemoryPointer, error) {
+	node, _, err := t.root()
+	if err != nil {
+		return pointer.ReferencedId{}, pointer.MemoryPointer{}, err
+	}
+
+	for {
+		if node == nil {
+			return pointer.ReferencedId{}, pointer.MemoryPointer{}, nil
+		}
+
+		index, found := slices.BinarySearchFunc(node.Ids, key, pointer.CompareReferencedIds)
+
+		if found {
+			return node.Ids[index], pointer.MemoryPointer{Offset: node.Ids[index].DataPointer.Offset}, nil
+		}
+
+		// no key found
+		if node.Leaf() {
+			return pointer.ReferencedId{}, pointer.MemoryPointer{}, nil
+		}
+
+		newOffset := node.Offsets[index]
+		newNode, err := t.readNode(newOffset)
+		if err != nil {
+			return pointer.ReferencedId{}, pointer.MemoryPointer{}, err
+		}
+
+		node = newNode
+	}
+}

pkg/btree/btree_test.go

@@ -1 +1,93 @@
 package btree
+
+import (
+	"encoding/binary"
+	"github.com/kevmo314/appendable/pkg/buftest"
+	"github.com/kevmo314/appendable/pkg/hnsw"
+	"github.com/kevmo314/appendable/pkg/pagefile"
+	"github.com/kevmo314/appendable/pkg/pointer"
+	"io"
+	"testing"
+)
+
+type testMetaPage struct {
+	pf   *pagefile.PageFile
+	root pointer.MemoryPointer
+}
+
+func (m *testMetaPage) SetRoot(mp pointer.MemoryPointer) error {
+	m.root = mp
+	return m.write()
+}
+
+func (m *testMetaPage) Root() (pointer.MemoryPointer, error) {
+	return m.root, nil
+}
+
+func (m *testMetaPage) write() error {
+	buf := make([]byte, 8)
+	binary.LittleEndian.PutUint64(buf, m.root.Offset)
+	if _, err := m.pf.Seek(4096, io.SeekStart); err != nil {
+		return err
+	}
+	if _, err := m.pf.Write(buf); err != nil {
+		return err
+	}
+	return nil
+}
+
+func newTestMetaPage(t *testing.T, pf *pagefile.PageFile) *testMetaPage {
+	meta := &testMetaPage{pf: pf}
+	offset, err := pf.NewPage([]byte{0, 0, 0, 0, 0, 0, 0, 0})
+	if err != nil {
+		t.Fatal(err)
+	}
+	// first page is garbage collection
+	if offset != 4096 {
+		t.Fatalf("expected offset 0, got %d", offset)
+	}
+	return meta
+}
+
+func TestBTree(t *testing.T) {
+	t.Run("empty tree", func(t *testing.T) {
+		b := buftest.NewSeekableBuffer()
+		p, err := pagefile.NewPageFile(b)
+		if err != nil {
+			t.Fatal(err)
+		}
+		tree := &BTree{PageFile: p, MetaPage: newTestMetaPage(t, p)}
+		// find a key that doesn't exist
+		k, _, err := tree.Find(pointer.ReferencedId{Value: hnsw.Id(0)})
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if k.Value != hnsw.Id(0) {
+			t.Fatalf("expected id 0, got %d", k)
+		}
+
+	})
+
+	t.Run("insert creates a root", func(t *testing.T) {
+		b := buftest.NewSeekableBuffer()
+		p, err := pagefile.NewPageFile(b)
+		if err != nil {
+			t.Fatal(err)
+		}
+		tree := &BTree{PageFile: p, MetaPage: newTestMetaPage(t, p), Width: uint16(0), VectorDim: 1}
+		if err := tree.Insert(pointer.ReferencedId{Value: 1}, hnsw.Point{1}); err != nil {
+			t.Fatal(err)
+		}
+		k, _, err := tree.Find(pointer.ReferencedId{Value: 1})
+
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if k.Value != hnsw.Id(1) {
+			t.Fatalf("expected id 1, got %d", k)
+		}
+	})
+
+}

pkg/btree/node.go

@@ -39,7 +39,7 @@ func (n *BTreeNode) Size() int64 {
 }
 
 func (n *BTreeNode) Leaf() bool {
-	return n.Offsets == nil || len(n.Offsets) == 0
+	return len(n.Offsets) == 0
 }
 
 func (n *BTreeNode) MarshalBinary() ([]byte, error) {

pkg/pointer/referenced_value_test.go

@@ -0,0 +1,49 @@
+package pointer
+
+import (
+	"slices"
+	"testing"
+)
+
+func TestReferencedValue(t *testing.T) {
+	t.Run("compare referenced value", func(t *testing.T) {
+		keys := []ReferencedId{
+			{
+				Value: 1,
+				DataPointer: MemoryPointer{
+					Offset: 100,
+					Length: 0,
+				},
+			},
+			{
+				Value: 2,
+				DataPointer: MemoryPointer{
+					Offset: 200,
+					Length: 0,
+				},
+			},
+			{
+				Value: 3,
+				DataPointer: MemoryPointer{
+					Offset: 300,
+					Length: 0,
+				},
+			},
+		}
+
+		index, found := slices.BinarySearchFunc(keys, ReferencedId{
+			DataPointer: MemoryPointer{},
+			Value:       1,
+		}, CompareReferencedIds)
+
+		if !found {
+			t.Fatal("expected to find key 1")
+		}
+
+		index++
+		if index != 1 {
+			t.Fatalf("expected index to be 1, got: %v", index)
+		}
+	})
+
+}