Fix ordering of query results on non-unique indexes

The query results for non-unique indexes were sorted properly by the
secondary key, but not by primary due to the substition scheme not
properly retaining the order. Fix this by using the following scheme:

  0x00 => 0x0101
  0x01 => 0x0102

This ensures that the original ordering is retained. Additionally also
encode the secondary key using this scheme to make sure 0x00 in the
secondary key does not mess the results.

The TestDB_Quick was extended to validate that the query results for
non-unique indexes are properly sorted by both the primary and secondary
keys.

Signed-off-by: Jussi Maki <[email protected]>

Author: joamaki

any_table.go

@@ -69,8 +69,7 @@ func (t AnyTable) Get(txn ReadTxn, index string, key string) (any, Revision, boo
 		if !ok {
 			break
 		}
-		secondary, _ := decodeNonUniqueKey(k)
-		if len(secondary) == len(rawKey) {
+		if nonUniqueKey(k).secondaryLen() == len(rawKey) {
 			return obj.data, obj.revision, true, nil
 		}
 	}

db_test.go

@@ -1150,31 +1150,26 @@ func TestWriteJSON(t *testing.T) {
 func Test_nonUniqueKey(t *testing.T) {
 	// empty keys
 	key := encodeNonUniqueKey(nil, nil)
-	secondary, _ := decodeNonUniqueKey(key)
-	assert.Len(t, secondary, 0)
+	nuk := nonUniqueKey(key)
+	assert.Equal(t, 0, nuk.secondaryLen())
 
 	// empty primary
 	key = encodeNonUniqueKey(nil, []byte("foo"))
-	secondary, _ = decodeNonUniqueKey(key)
-	assert.Equal(t, string(secondary), "foo")
+	nuk = nonUniqueKey(key)
+	assert.Zero(t, nuk.primaryLen())
+	assert.Equal(t, 3, nuk.secondaryLen())
 
 	// empty secondary
 	key = encodeNonUniqueKey([]byte("quux"), []byte{})
-	secondary, _ = decodeNonUniqueKey(key)
-	assert.Len(t, secondary, 0)
+	nuk = nonUniqueKey(key)
+	assert.Zero(t, nuk.secondaryLen())
 
 	// non-empty
 	key = encodeNonUniqueKey([]byte("foo"), []byte("quux"))
-	secondary, primary := decodeNonUniqueKey(key)
-	assert.EqualValues(t, secondary, "quux")
-	assert.EqualValues(t, primary, "foo")
-
-	// non-empty, primary with substitutions:
-	// 0x0 => 0xfe, 0xfe => 0xfd01, 0xfd => 0xfd00
-	key = encodeNonUniqueKey([]byte{0x0, 0xfd, 0xfe}, []byte("quux"))
-	secondary, primary = decodeNonUniqueKey(key)
-	assert.EqualValues(t, secondary, "quux")
-	assert.EqualValues(t, primary, []byte{0xfe, 0xfd, 0x01, 0xfd, 0x00})
+	nuk = nonUniqueKey(key)
+	assert.Equal(t, 4, nuk.secondaryLen())
+	assert.Equal(t, 3, nuk.primaryLen())
+	assert.EqualValues(t, "foo", nuk.encodedPrimary())
 }
 
 func Test_validateTableName(t *testing.T) {

http.go

@@ -124,6 +124,9 @@ type QueryResponse struct {
 
 func runQuery(indexTxn indexReadTxn, lowerbound bool, queryKey []byte, onObject func(object) error) {
 	var iter *part.Iterator[object]
+	if !indexTxn.unique {
+		queryKey = encodeNonUniqueBytes(queryKey)
+	}
 	if lowerbound {
 		iter = indexTxn.LowerBound(queryKey)
 	} else {
@@ -137,8 +140,7 @@ func runQuery(indexTxn indexReadTxn, lowerbound bool, queryKey []byte, onObject
 		match = func(k []byte) bool { return len(k) == len(queryKey) }
 	default:
 		match = func(k []byte) bool {
-			secondary, _ := decodeNonUniqueKey(k)
-			return len(secondary) == len(queryKey)
+			return nonUniqueKey(k).secondaryLen() == len(queryKey)
 		}
 	}
 	for key, obj, ok := iter.Next(); ok; key, obj, ok = iter.Next() {

iterator.go

@@ -113,18 +113,21 @@ func nonUniqueSeq[Obj any](iter *part.Iterator[object], prefixSearch bool, searc
 				break
 			}
 
-			secondary, primary := decodeNonUniqueKey(key)
+			nuk := nonUniqueKey(key)
+			secondaryLen := nuk.secondaryLen()
 
 			switch {
-			case !prefixSearch && len(secondary) != len(searchKey):
+			case !prefixSearch && secondaryLen != len(searchKey):
 				// This a List(), thus secondary key must match length exactly.
 				continue
-			case prefixSearch && len(secondary) < len(searchKey):
+			case prefixSearch && secondaryLen < len(searchKey):
 				// This is Prefix(), thus key must be equal or longer to search key.
 				continue
 			}
 
 			if prefixSearch {
+				primary := nuk.encodedPrimary()
+
 				// When doing a prefix search on a non-unique index we may see the
 				// same object multiple times since multiple keys may point it.
 				// Skip if we've already seen this object.
@@ -157,8 +160,10 @@ func nonUniqueLowerBoundSeq[Obj any](iter *part.Iterator[object], searchKey []by
 			// With a non-unique index we have a composite key <secondary><primary><secondary len>.
 			// This means we need to check every key that it's larger or equal to the search key.
 			// Just seeking to the first one isn't enough as the secondary key length may vary.
-			secondary, primary := decodeNonUniqueKey(key)
+			nuk := nonUniqueKey(key)
+			secondary := nuk.encodedSecondary()
 			if bytes.Compare(secondary, searchKey) >= 0 {
+				primary := nuk.encodedPrimary()
 				if _, found := visited[string(primary)]; found {
 					continue
 				}

quick_test.go

@@ -1,8 +1,12 @@
 package statedb
 
 import (
+	"bytes"
 	"cmp"
+	"fmt"
 	"iter"
+	"math/rand"
+	"reflect"
 	"strings"
 	"testing"
 	"testing/quick"
@@ -13,6 +17,18 @@ import (
 
 var quickConfig = &quick.Config{
 	MaxCount: 5000,
+
+	// Make 1-8 byte long strings as input data. Keep the strings shorter
+	// than the default quick value generation to hit the more interesting cases
+	// often.
+	Values: func(args []reflect.Value, rand *rand.Rand) {
+		for i := range args {
+			numBytes := 1 + rand.Intn(8)
+			bs := make([]byte, numBytes)
+			rand.Read(bs)
+			args[i] = reflect.ValueOf(string(bs))
+		}
+	},
 }
 
 // Use an object with strings for both primary and secondary
@@ -23,11 +39,8 @@ type quickObj struct {
 	A, B string
 }
 
-func (q quickObj) getA() string {
-	return q.A
-}
-func (q quickObj) getB() string {
-	return q.B
+func (q quickObj) String() string {
+	return fmt.Sprintf("%x %x", []byte(q.A), []byte(q.B))
 }
 
 var (
@@ -52,13 +65,27 @@ var (
 	}
 )
 
-func isOrdered[A cmp.Ordered, B any](t *testing.T, it iter.Seq2[A, B]) bool {
-	var prev A
+func isOrdered(t *testing.T, aFirst bool, it iter.Seq2[quickObj, Revision]) bool {
+	var prev quickObj
 	for a := range it {
-		if cmp.Compare(a, prev) < 0 {
-			t.Logf("isOrdered: %#v < %#v!", a, prev)
-			return false
+		if aFirst {
+			if ret := cmp.Compare(a.A, prev.A); ret < 0 {
+				t.Logf("isOrdered(A): %s < %s!", a, prev)
+				return false
+			} else if ret == 0 && cmp.Compare(a.B, prev.B) < 0 {
+				t.Logf("isOrdered(B): %s < %s!", a, prev)
+				return false
+			}
+		} else {
+			if ret := cmp.Compare(a.B, prev.B); ret < 0 {
+				t.Logf("isOrdered(B): %s < %s!", a, prev)
+				return false
+			} else if ret == 0 && cmp.Compare(a.A, prev.A) < 0 {
+				t.Logf("isOrdered(A): %s < %s!", a, prev)
+				return false
+			}
 		}
+
 		prev = a
 	}
 	return true
@@ -106,15 +133,15 @@ func TestDB_Quick(t *testing.T) {
 		// Check queries against the primary index
 		//
 
-		if numExpected != seqLen(Map(table.All(rtxn), quickObj.getA)) {
+		if numExpected != seqLen(table.All(rtxn)) {
 			t.Logf("All() via aIndex wrong length")
 			return false
 		}
-		if numExpected != seqLen(Map(table.Prefix(rtxn, aIndex.Query("")), quickObj.getA)) {
+		if numExpected != seqLen(table.Prefix(rtxn, aIndex.Query(""))) {
 			t.Logf("Prefix() via aIndex wrong length")
 			return false
 		}
-		if numExpected != seqLen(Map(table.LowerBound(rtxn, aIndex.Query("")), quickObj.getA)) {
+		if numExpected != seqLen(table.LowerBound(rtxn, aIndex.Query(""))) {
 			t.Logf("LowerBound() via aIndex wrong length")
 			return false
 		}
@@ -151,20 +178,20 @@ func TestDB_Quick(t *testing.T) {
 		for anyObj := range anyObjs {
 			obj := anyObj.(quickObj)
 			if cmp.Compare(obj.A, a) < 0 {
-				t.Logf("AnyTable.LowerBound() order wrong")
+				t.Logf("AnyTable.LowerBound(%x) order wrong: %x < %x", []byte(a), []byte(obj.A), []byte(a))
 				return false
 			}
 		}
 
-		if !isOrdered(t, Map(table.All(rtxn), quickObj.getA)) {
+		if !isOrdered(t, true, table.All(rtxn)) {
 			t.Logf("All() wrong order")
 			return false
 		}
-		if !isOrdered(t, Map(table.Prefix(rtxn, aIndex.Query("")), quickObj.getA)) {
+		if !isOrdered(t, true, table.Prefix(rtxn, aIndex.Query(""))) {
 			t.Logf("Prefix() via aIndex wrong order")
 			return false
 		}
-		if !isOrdered(t, Map(table.LowerBound(rtxn, aIndex.Query("")), quickObj.getA)) {
+		if !isOrdered(t, true, table.LowerBound(rtxn, aIndex.Query(""))) {
 			t.Logf("LowerBound() via aIndex wrong order")
 			return false
 		}
@@ -178,6 +205,7 @@ func TestDB_Quick(t *testing.T) {
 			t.Logf("Prefix() via bIndex wrong length")
 			return false
 		}
+
 		if numExpected != seqLen(table.LowerBound(rtxn, bIndex.Query(""))) {
 			t.Logf("LowerBound() via bIndex wrong length")
 			return false
@@ -187,7 +215,7 @@ func TestDB_Quick(t *testing.T) {
 		// not be the one that we just inserted.
 		obj, _, found = table.Get(rtxn, bIndex.Query(b))
 		if !found || obj.B != b {
-			t.Logf("Get() via bIndex not found or wrong B")
+			t.Logf("Get(%q) via bIndex not found (%v) or wrong B (%q vs %q)", b, found, obj.B, b)
 			return false
 		}
 
@@ -224,7 +252,7 @@ func TestDB_Quick(t *testing.T) {
 		for anyObj := range anyObjs {
 			obj := anyObj.(quickObj)
 			if !strings.HasPrefix(obj.B, b) {
-				t.Logf("AnyTable.Prefix() via bIndex has wrong prefix")
+				t.Logf("AnyTable.Prefix() via bIndex has wrong prefix: %q vs %q", obj.B, b)
 				return false
 			}
 		}
@@ -253,17 +281,75 @@ func TestDB_Quick(t *testing.T) {
 		}
 
 		// Iterating over the secondary index returns the objects in order
-		// defined by the "B" key.
-		if !isOrdered(t, Map(table.Prefix(rtxn, bIndex.Query("")), quickObj.getB)) {
-			t.Logf("Prefix() via bIndex has wrong order")
+		// defined by the "B" key first and then by the "A" key.
+		if !isOrdered(t, false, table.Prefix(rtxn, bIndex.Query(""))) {
+			t.Logf("Prefix() via bIndex wrong order")
+			rtxn.getTxn().mustIndexReadTxn(table, table.indexPos("b")).PrintTree()
 			return false
 		}
-		if !isOrdered(t, Map(table.LowerBound(rtxn, bIndex.Query("")), quickObj.getB)) {
-			t.Logf("LowerBound() via bIndex has wrong order")
+		if !isOrdered(t, false, table.LowerBound(rtxn, bIndex.Query(""))) {
+			t.Logf("Prefix() via bIndex wrong order")
 			return false
 		}
 		return true
 	}
 
 	require.NoError(t, quick.Check(check, quickConfig))
 }
+
+func Test_Quick_nonUniqueKey(t *testing.T) {
+	check := func(p1, s1, p2, s2 []byte) bool {
+		key1 := encodeNonUniqueKey(p1, s1)
+		expectedLen := encodedLength(p1) + 1 + encodedLength(s1) + 2
+		minLen := len(p1) + 1 + len(s1) + 2
+		if expectedLen < minLen {
+			t.Logf("expected length too short (%d), must be >= %d", expectedLen, minLen)
+			return false
+		}
+		if len(key1) != expectedLen {
+			t.Logf("length mismatch, expected %d, got %d", expectedLen, len(key1))
+			return false
+		}
+
+		nuk1 := nonUniqueKey(key1)
+		if len(nuk1.encodedPrimary()) < len(p1) {
+			t.Logf("encodedPrimary() length (%d) shorter than original (%d)", len(nuk1.encodedPrimary()), len(p1))
+			return false
+		}
+		if len(nuk1.encodedPrimary()) != encodedLength(p1) {
+			t.Logf("encodedPrimary() length (%d) does not match encodedLength() (%d)", len(nuk1.encodedPrimary()), len(p1))
+			return false
+		}
+		if len(nuk1.encodedSecondary()) < len(s1) {
+			t.Logf("encodedSecondary() length (%d) shorter than original (%d)", len(nuk1.encodedSecondary()), len(s1))
+		}
+		if len(nuk1.encodedSecondary()) != encodedLength(s1) {
+			t.Logf("encodedSecondary() length (%d) does not match encodedLength() (%d)", len(nuk1.encodedSecondary()), len(s1))
+			return false
+		}
+
+		// Do another key and check that ordering is preserved.
+		key2 := encodeNonUniqueKey(p2, s2)
+		scmp := bytes.Compare(s1, s2)
+		pcmp := bytes.Compare(p1, p2)
+		kcmp := bytes.Compare(key1, key2)
+
+		if (scmp == 0 && pcmp != kcmp) /* secondary key matches, primary key determines order */ ||
+			(scmp != 0 && scmp != kcmp) /* secondary key determines order */ {
+			t.Logf("ordering not preserved: key1=%v key2=%v, p1=%v, s1=%v, p2=%v, s2=%v", key1, key2, p1, s1, p2, s2)
+			return false
+		}
+		return true
+	}
+	require.NoError(t, quick.Check(check, &quick.Config{
+		MaxCount: 50000,
+		Values: func(args []reflect.Value, rand *rand.Rand) {
+			for i := range args {
+				numBytes := 1 + rand.Intn(8)
+				bs := make([]byte, numBytes)
+				rand.Read(bs)
+				args[i] = reflect.ValueOf(bs)
+			}
+		},
+	}))
+}

table.go

@@ -321,8 +321,7 @@ func (t *genTable[Obj]) GetWatch(txn ReadTxn, q Query[Obj]) (obj Obj, revision u
 		}
 
 		// Check that we have a full match on the key
-		secondary, _ := decodeNonUniqueKey(key)
-		if len(secondary) == len(q.key) {
+		if nonUniqueKey(key).secondaryLen() == len(q.key) {
 			break
 		}
 	}