featurebase/tstore/debug_utils_test.go

package tstore

import (
	"fmt"
	"math/rand"
	"os"
	"sort"
	"testing"

	"github.com/featurebasedb/featurebase/v3/bufferpool"
	"github.com/featurebasedb/featurebase/v3/sql3/parser"
	"github.com/featurebasedb/featurebase/v3/sql3/planner/types"
	"github.com/stretchr/testify/assert"
)

func buildSample(objectID int32, shard int32) (*BTree, error, func()) {
	diskManager := bufferpool.NewTupleStoreDiskManager()

	dataFile := fmt.Sprintf("ts-shard.%04d", shard)

	diskManager.CreateOrOpenShard(objectID, shard, dataFile)

	bufferPool := bufferpool.NewBufferPool(100, diskManager)

	tableSchema := types.Schema{
		&types.PlannerColumn{
			ColumnName: "vtest",
			Type:       parser.NewDataTypeVarchar(50),
		},
	}
	b, e := NewBTree(8, objectID, shard, tableSchema, bufferPool)
	return b, e, func() {
		os.Remove(dataFile)
	}
}

func TestBTree_RangeIterator(t *testing.T) {
	b, err, c := buildSample(1, 0)
	defer c()
	assert.Nil(t, err)

	rowSchema := types.Schema{
		&types.PlannerColumn{
			ColumnName: "_id",
			Type:       parser.NewDataTypeID(),
		},
		&types.PlannerColumn{
			ColumnName: "vtest",
			Type:       parser.NewDataTypeVarchar(50),
		},
	}

	inserts := make([]int, 0)
	for i := 1; i <= 90; i++ {
		inserts = append(inserts, i)
	}
	rand.Shuffle(len(inserts), func(i, j int) { inserts[i], inserts[j] = inserts[j], inserts[i] })

	rr := make(types.Row, 2)
	want1 := make([]string, 0)
	want2 := make([]string, 0)
	want3 := make([]string, 0)
	for _, i := range inserts {
		rr[0] = int64(i)
		rr[1] = fmt.Sprintf("This is a test of things %d", i)

		if i >= 1 && i < 10 { // items in first page
			want1 = append(want1, rr[1].(string))
		}

		if i >= 75 && i < 80 { // items in last page
			want2 = append(want2, rr[1].(string))
		}
		if i >= 35 && i < 75 { // items in both pages
			want3 = append(want3, rr[1].(string))
		}

		tup := &BTreeTuple{
			TupleSchema: rowSchema,
			Tuple:       rr,
		}
		err = b.Insert(tup)
		if err != nil {
			t.Fatal(err)
		}
	}

	// b.Dot(os.Stdout, "first", false)
	root, err := b.fetchNode(b.rootNode)
	assert.Nil(t, err)
	t.Run("first page", func(t *testing.T) {
		itr := NewRangeIterator(b, root, Int(1), Int(10), b.schema)
		defer itr.Dispose()
		got := make([]string, 0)
		for itr.Next() {
			item, _ := itr.Item()
			got = append(got, item.Tuple[0].(string))
		}
		sort.Strings(got)
		sort.Strings(want1)
		assert.Equal(t, want1, got)
	})
	t.Run("last page", func(t *testing.T) {
		itr := NewRangeIterator(b, root, Int(75), Int(80), b.schema)
		defer itr.Dispose()
		got := make([]string, 0)
		for itr.Next() {
			item, _ := itr.Item()
			got = append(got, item.Tuple[0].(string))
		}
		sort.Strings(got)
		sort.Strings(want2)
		assert.Equal(t, want2, got)
	})
	t.Run("cover pages", func(t *testing.T) {
		itr := NewRangeIterator(b, root, Int(35), Int(75), b.schema)
		defer itr.Dispose()
		got := make([]string, 0)
		for itr.Next() {
			item, _ := itr.Item()
			got = append(got, item.Tuple[0].(string))
		}
		sort.Strings(got)
		sort.Strings(want3)
		assert.Equal(t, want3, got)
	})
}

func TestDotOverflow(t *testing.T) {
	t.Skip("need to figure out fail")
	objectId := int32(1)
	shard := int32(0)
	dot, err := OpenBtree(fmt.Sprintf("ts-shard.%04d", shard))
	assert.NotNil(t, err)
	dot.Dot(os.Stdout, "next", true)

	diskManager := bufferpool.NewTupleStoreDiskManager()

	dataFile := fmt.Sprintf("ts-shard.%04d", shard)
	defer os.Remove(dataFile)

	diskManager.CreateOrOpenShard(objectId, shard, dataFile)
	bufferPool := bufferpool.NewBufferPool(100, diskManager)

	tableSchema := make(types.Schema, 0)

	numCols := 100
	numRecs := 10

	// build schema
	for i := 0; i < numCols; i++ {
		tableSchema = append(tableSchema, &types.PlannerColumn{
			ColumnName: fmt.Sprintf("vtest%d", i+1),
			Type:       parser.NewDataTypeVarchar(4),
		})
	}

	b, err := NewBTree(8, objectId, shard, tableSchema, bufferPool)
	_ = b
	_ = numRecs
	if err != nil {
		t.Fatal(err)
	}
	rowSchema := make(types.Schema, 0)

	rowSchema = append(rowSchema, &types.PlannerColumn{
		ColumnName: "_id",
		Type:       parser.NewDataTypeID(),
	})

	for i := 0; i < numCols; i++ {
		rowSchema = append(rowSchema, &types.PlannerColumn{
			ColumnName: fmt.Sprintf("vtest%d", i+1),
			Type:       parser.NewDataTypeVarchar(4),
		})
	}

	inserts := make([]int, 0)

	for i := 1; i <= numRecs; i++ {
		inserts = append(inserts, i)
	}

	rr := make(types.Row, numCols+1)

	for _, i := range inserts {
		rr[0] = int64(i)

		for j := 0; j < numCols; j++ {
			rr[j+1] = fmt.Sprintf("%04d", j)
		}

		tup := &BTreeTuple{
			TupleSchema: rowSchema,
			Tuple:       rr,
		}

		// fmt.Printf("[%d]row key %v\n\n", j, i)

		err = b.Insert(tup)
		if err != nil {
			t.Fatal(err)
		}

	}

	// b.Dump(0)
	b.Dot(os.Stdout, "first", true)
}