diff --git a/go/store/prolly/tuple_map.go b/go/store/prolly/tuple_map.go index ad611e19f4..5607c006cf 100644 --- a/go/store/prolly/tuple_map.go +++ b/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) diff --git a/go/store/val/tuple.go b/go/store/val/tuple.go index 93de5a5a04..63f8e9d938 100644 --- a/go/store/val/tuple.go +++ b/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 diff --git a/go/store/val/tuple_test.go b/go/store/val/tuple_test.go index f6f7700e6a..376c623181 100644 --- a/go/store/val/tuple_test.go +++ b/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 -}