go/store/val: add TuplePrefix and TupleSuffix functions

go/store/prolly/tuple_map.go

@@ -252,6 +252,8 @@ func (m Map) GetPrefix(ctx context.Context, key val.Tuple, prefDesc val.TupleDes
 	return m.tuples.GetPrefix(ctx, key, prefDesc, cb)
 }
 
+// todo(andy): iter prefix
+
 // Has returns true is |key| is present in the Map.
 func (m Map) Has(ctx context.Context, key val.Tuple) (ok bool, err error) {
 	return m.tuples.Has(ctx, key)

go/store/val/tuple.go

@@ -113,6 +113,42 @@ func NewTuple(pool pool.BuffPool, values ...[]byte) Tuple {
 	return tup
 }
 
+func TuplePrefix(pool pool.BuffPool, tup Tuple, k int) Tuple {
+	cnt := tup.Count()
+	if k >= cnt {
+		return tup
+	}
+	for k > 0 && tup.FieldIsNull(k-1) {
+		k-- // trim NULL suffix
+	}
+	if k == 0 {
+		return EmptyTuple
+	}
+
+	stop, _ := tup.GetOffset(k)
+	prefix, offs := allocateTuple(pool, ByteSize(stop), k)
+	split := ByteSize(len(tup)) - uint16Size*ByteSize(cnt)
+
+	copy(prefix, tup[:stop])
+	copy(offs, tup[split:])
+	return prefix
+}
+
+func TupleSuffix(pool pool.BuffPool, tup Tuple, k int) Tuple {
+	// todo(andy)
+	cnt := tup.Count()
+	if k == 0 {
+		return EmptyTuple
+	} else if k >= cnt {
+		return tup
+	}
+	fields := make([][]byte, k)
+	for i := range fields {
+		fields[i] = tup.GetField((cnt - k) + i)
+	}
+	return NewTuple(pool, fields...)
+}
+
 func trimNullSuffix(values [][]byte) [][]byte {
 	n := len(values)
 	for i := len(values) - 1; i >= 0; i-- {
@@ -191,12 +227,6 @@ func (tup Tuple) GetField(i int) []byte {
 	return tup[start:stop]
 }
 
-// GetManyFields takes a sorted slice of ordinals |indexes| and returns the requested
-// tuple fields. It populates field data into |slices| to avoid allocating.
-func (tup Tuple) GetManyFields(indexes []int, slices [][]byte) [][]byte {
-	return sliceManyFields(tup, indexes, slices)
-}
-
 func (tup Tuple) FieldIsNull(i int) bool {
 	return tup.GetField(i) == nil
 }
@@ -219,56 +249,6 @@ func writeFieldCount(tup Tuple, count int) {
 	WriteUint16(sl, uint16(count))
 }
 
-func sliceManyFields(tuple Tuple, indexes []int, slices [][]byte) [][]byte {
-	cnt := tuple.Count()
-	sz := ByteSize(len(tuple))
-	split := sz - uint16Size*ByteSize(cnt)
-
-	data := tuple[:split]
-	offs := offsets(tuple[split : sz-countSize])
-
-	// if count is 1, we assume |indexes| is [0]
-	if cnt == 1 {
-		slices[0] = data
-		if len(data) == 0 {
-			slices[0] = nil
-		}
-		return slices
-	}
-
-	subset := slices
-	// we don't have a "stop" offset for the last field
-	n := len(slices)
-	if indexes[n-1] == cnt-1 {
-		o := ReadUint16(offs[len(offs)-2:])
-		slices[n-1] = data[o:]
-		indexes = indexes[:n-1]
-		subset = subset[:n-1]
-	}
-
-	// we don't have a "start" offset for the first field
-	if len(indexes) > 0 && indexes[0] == 0 {
-		o := ReadUint16(offs[:2])
-		slices[0] = data[:o]
-		indexes = indexes[1:]
-		subset = subset[1:]
-	}
-
-	for i, k := range indexes {
-		start := ReadUint16(offs[(k-1)*2 : k*2])
-		stop := ReadUint16(offs[k*2 : (k+1)*2])
-		subset[i] = tuple[start:stop]
-	}
-
-	for i := range slices {
-		if len(slices[i]) == 0 {
-			slices[i] = nil
-		}
-	}
-
-	return slices
-}
-
 type offsets []byte
 
 // offsetsSize returns the number of bytes needed to

go/store/val/tuple_test.go

@@ -16,31 +16,16 @@ package val
 
 import (
 	"math/rand"
-	"sort"
 	"testing"
 
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-
 	"github.com/dolthub/dolt/go/store/pool"
+	"github.com/stretchr/testify/assert"
 )
 
 var testPool = pool.NewBuffPool()
 
-// todo(andy): randomize test seed
-var testRand = rand.New(rand.NewSource(1))
-
 func TestNewTuple(t *testing.T) {
-	t.Run("test tuple round trip", func(t *testing.T) {
-		roundTripTupleFields(t)
-	})
-	t.Run("test tuple get many", func(t *testing.T) {
-		testTupleGetMany(t)
-	})
-}
-
-func roundTripTupleFields(t *testing.T) {
-	for n := 0; n < 100; n++ {
+	for n := 0; n < 1024; n++ {
 		fields := randomByteFields(t)
 		tup := NewTuple(testPool, fields...)
 		for i, field := range fields {
@@ -49,31 +34,25 @@ func roundTripTupleFields(t *testing.T) {
 	}
 }
 
-func testTupleGetMany(t *testing.T) {
-	for n := 0; n < 1000; n++ {
+func TestTuplePrefix(t *testing.T) {
+	for n := 0; n < 1024; n++ {
 		fields := randomByteFields(t)
-		tup := NewTuple(testPool, fields...)
-
-		// GetManyFields must not be called with indexes that are greater than
-		// or equal to the tuple count.
-		indexes := randomFieldIndexes(fields)
-		for i := len(indexes) - 1; i >= 0; i-- {
-			idx := indexes[i]
-			if idx < tup.Count() {
-				break
-			}
-			require.Equal(t, 0, len(fields[idx]))
-			indexes = indexes[:i]
-		}
-		if len(indexes) == 0 {
-			continue
+		full := NewTuple(testPool, fields...)
+		for i := 0; i <= len(fields); i++ {
+			exp := NewTuple(testPool, fields[:i]...)
+			act := TuplePrefix(testPool, full, i)
+			assert.Equal(t, exp, act)
 		}
+	}
+}
 
-		actual := tup.GetManyFields(indexes, make([][]byte, len(indexes)))
-
-		for k, idx := range indexes {
-			exp := fields[idx]
-			act := actual[k]
+func TestTupleSuffix(t *testing.T) {
+	for n := 0; n < 1024; n++ {
+		fields := randomByteFields(t)
+		full := NewTuple(testPool, fields...)
+		for i := 0; i <= full.Count(); i++ {
+			exp := NewTuple(testPool, fields[i:]...)
+			act := TupleSuffix(testPool, full, full.Count()-i)
 			assert.Equal(t, exp, act)
 		}
 	}
@@ -93,23 +72,3 @@ func randomByteFields(t *testing.T) (fields [][]byte) {
 	}
 	return
 }
-
-func randomFieldIndexes(fields [][]byte) []int {
-	indexes := make([]int, len(fields))
-	for i := range indexes {
-		indexes[i] = i
-	}
-
-	k := testRand.Intn(len(indexes))
-	if k == 0 {
-		k++
-	}
-
-	testRand.Shuffle(len(indexes), func(i, j int) {
-		indexes[i], indexes[j] = indexes[j], indexes[i]
-	})
-	indexes = indexes[:k]
-	sort.Ints(indexes)
-
-	return indexes
-}