Merge #151597

151597: sql: add LTREE key encoding r=paulniziolek a=paulniziolek

#### sql: add LTREE keyside encoding 

This adds basic keyside encoding for LTREE, which is encoded similarly to a string array: ltrees have their own terminator and their labels are encoded via escaped bytes.

Informs: #44657
Epic: CRDB-148

Release note: None


Co-authored-by: Paul Niziolek <[email protected]>

Author: craig[bot]

pkg/sql/catalog/colinfo/col_type_info.go

@@ -148,18 +148,14 @@ func ColumnTypeIsIndexable(t *types.T) bool {
 	}
 
 	switch t.Family() {
-	case types.TupleFamily, types.RefCursorFamily, types.JsonpathFamily, types.LTreeFamily:
-		// TODO(paulniziolek): LTreeFamily should be supported in keyside encoding.
-		// Temporarily, we disallow it, until implemented.
+	case types.TupleFamily, types.RefCursorFamily, types.JsonpathFamily:
 		return false
 	}
 
 	// If the type is an array, check its content type as well.
 	if unwrapped := t.ArrayContents(); unwrapped != nil {
 		switch unwrapped.Family() {
-		case types.TupleFamily, types.RefCursorFamily, types.JsonpathFamily, types.LTreeFamily:
-			// TODO(paulniziolek): LTreeFamily should be supported in keyside encoding.
-			// Temporarily, we disallow it, until implemented.
+		case types.TupleFamily, types.RefCursorFamily, types.JsonpathFamily:
 			return false
 		}
 	}
@@ -175,9 +171,7 @@ func ColumnTypeIsInvertedIndexable(t *types.T) bool {
 	switch t.Family() {
 	case types.ArrayFamily:
 		switch t.ArrayContents().Family() {
-		case types.RefCursorFamily, types.JsonpathFamily, types.LTreeFamily:
-			// TODO(paulniziolek): LTreeFamily should be supported in keyside encoding.
-			// Temporarily, we disallow it, until implemented.
+		case types.RefCursorFamily, types.JsonpathFamily:
 			return false
 		default:
 			return true
@@ -230,10 +224,6 @@ func MustBeValueEncoded(semanticType *types.T) bool {
 		return true
 	case types.PGVectorFamily:
 		return true
-	case types.LTreeFamily:
-		// TODO(paulniziolek): LTreeFamily should be supported in keyside encoding.
-		// Temporarily, we disallow it, until implemented.
-		return true
 		// NB: if you're adding a new type here, you probably also want to
 		// include it into rowenc.mustUseValueEncodingForFingerprinting.
 	}

pkg/sql/rowenc/encoded_datum.go

@@ -331,9 +331,6 @@ func mustUseValueEncodingForFingerprinting(t *types.T) bool {
 	// behavior can result in incorrect results in mixed version clusters).
 	case types.JsonFamily, types.TSQueryFamily, types.TSVectorFamily, types.PGVectorFamily:
 		return true
-	case types.LTreeFamily:
-		// TODO(paulniziolek): remove this once key encoding is added.
-		return true
 	case types.ArrayFamily:
 		// Note that at time of this writing we don't support arrays of JSON
 		// (tracked via #23468) nor of TSQuery / TSVector / PGVector types (tracked by

pkg/sql/rowenc/encoded_datum_test.go

@@ -218,10 +218,6 @@ func TestEncDatumCompare(t *testing.T) {
 		case types.AnyFamily, types.UnknownFamily, types.ArrayFamily, types.JsonFamily, types.TupleFamily, types.VoidFamily,
 			types.TSQueryFamily, types.TSVectorFamily, types.PGVectorFamily, types.TriggerFamily, types.JsonpathFamily:
 			continue
-		case types.LTreeFamily:
-			// TODO(paulniziolek): Temporarily skip LTrees as they are
-			// currently missing keyside indexing support.
-			continue
 		case types.CollatedStringFamily:
 			typ = types.MakeCollatedString(types.String, *randgen.RandCollationLocale(rng))
 		}

pkg/sql/rowenc/index_encoding_test.go

@@ -587,11 +587,6 @@ func TestEncodeContainingArrayInvertedIndexSpans(t *testing.T) {
 		if typ.ArrayContents().Family() == types.JsonpathFamily {
 			continue
 		}
-		// TODO(paulniziolek): Temporarily skip arrays with LTREEs as they are
-		// currently missing keyside indexing support.
-		if typ.ArrayContents().Family() == types.LTreeFamily {
-			continue
-		}
 
 		// Generate two random arrays and evaluate the result of `left @> right`.
 		left := randgen.RandArray(rng, typ, 0 /* nullChance */)
@@ -734,12 +729,6 @@ func TestEncodeContainedArrayInvertedIndexSpans(t *testing.T) {
 	for i := 0; i < 100; i++ {
 		typ := randgen.RandArrayType(rng)
 
-		// TODO(paulniziolek): Temporarily skip arrays with LTREEs as they are
-		// currently missing keyside indexing support.
-		if typ.ArrayContents().Family() == types.LTreeFamily {
-			continue
-		}
-
 		// Generate two random arrays and evaluate the result of `left <@ right`.
 		left := randgen.RandArray(rng, typ, 0 /* nullChance */)
 		right := randgen.RandArray(rng, typ, 0 /* nullChance */)
@@ -984,11 +973,6 @@ func TestEncodeOverlapsArrayInvertedIndexSpans(t *testing.T) {
 		if typ.ArrayContents().Family() == types.JsonpathFamily {
 			continue
 		}
-		// TODO(paulniziolek): Temporarily skip arrays with LTREEs as they are
-		// currently missing keyside indexing support.
-		if typ.ArrayContents().Family() == types.LTreeFamily {
-			continue
-		}
 
 		// Generate two random arrays and evaluate the result of `left && right`.
 		// Using 1/9th as the Null Chance to generate arrays with a small

pkg/sql/rowenc/keyside/BUILD.bazel

@@ -23,6 +23,7 @@ go_library(
         "//pkg/util/encoding",
         "//pkg/util/ipaddr",
         "//pkg/util/json",
+        "//pkg/util/ltree",
         "//pkg/util/timetz",
         "//pkg/util/timeutil/pgdate",
         "//pkg/util/uuid",

pkg/sql/rowenc/keyside/decode.go

@@ -20,6 +20,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/encoding"
 	"github.com/cockroachdb/cockroach/pkg/util/ipaddr"
 	"github.com/cockroachdb/cockroach/pkg/util/json"
+	"github.com/cockroachdb/cockroach/pkg/util/ltree"
 	"github.com/cockroachdb/cockroach/pkg/util/timetz"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil/pgdate"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
@@ -273,6 +274,17 @@ func Decode(
 			rkey, i, err = encoding.DecodeVarintDescending(key)
 		}
 		return a.NewDOid(tree.MakeDOid(oid.Oid(i), valType)), rkey, err
+	case types.LTreeFamily:
+		var l ltree.T
+		if dir == encoding.Ascending {
+			rkey, l, err = encoding.DecodeLTreeAscending(key)
+		} else {
+			rkey, l, err = encoding.DecodeLTreeDescending(key)
+		}
+		if err != nil {
+			return nil, nil, err
+		}
+		return tree.NewDLTree(l), rkey, err
 	case types.EnumFamily:
 		var r []byte
 		if dir == encoding.Ascending {

pkg/sql/rowenc/keyside/encode.go

@@ -181,13 +181,13 @@ func Encode(b []byte, val tree.Datum, dir encoding.Direction) ([]byte, error) {
 		return append(b, []byte(*t)...), nil
 	case *tree.DJSON:
 		return encodeJSONKey(b, t, dir)
+	case *tree.DLTree:
+		if dir == encoding.Ascending {
+			return encoding.EncodeLTreeAscending(b, t.LTree), nil
+		}
+		return encoding.EncodeLTreeDescending(b, t.LTree), nil
 	}
 	if buildutil.CrdbTestBuild {
-		if _, isLTree := val.(*tree.DLTree); isLTree {
-			// TODO(paulniziolek): remove this exception once key encoding is
-			// added.
-			return nil, errors.Newf("LTREE key encoding is not implemented yet")
-		}
 		return nil, errors.AssertionFailedf("unable to encode table key: %T", val)
 	}
 	return nil, errors.Errorf("unable to encode table key: %T", val)

pkg/util/encoding/BUILD.bazel

@@ -61,6 +61,7 @@ go_test(
         "//pkg/util/ipaddr",
         "//pkg/util/json",
         "//pkg/util/log",
+        "//pkg/util/ltree",
         "//pkg/util/randutil",
         "//pkg/util/timeofday",
         "//pkg/util/timetz",

pkg/util/encoding/encoding.go

@@ -136,10 +136,20 @@ const (
 	jsonArrayKeyDescendingMarker      = jsonTrueKeyDescendingMarker - 1
 	jsonObjectKeyDescendingMarker     = jsonArrayKeyDescendingMarker - 1
 
+	// LTREE key encoding markers
+	ltreeKeyMarker           = jsonEmptyArrayKeyDescendingMarker + 1
+	ltreeKeyDescendingMarker = ltreeKeyMarker + 1
+
 	// Terminators for JSON Key encoding.
 	jsonKeyTerminator           byte = 0x00
 	jsonKeyDescendingTerminator byte = 0xFF
 
+	// Terminators for LTREE Key encoding.
+	ltreeKeyTerminator                byte = 0x00
+	ltreeKeyDescendingTerminator      byte = 0xFF
+	ltreeLabelKeyTerminator           byte = 0x01
+	ltreeLabelKeyDescendingTerminator byte = 0xFE
+
 	// IntMin is chosen such that the range of int tags does not overlap the
 	// ascii character set that is frequently used in testing.
 	IntMin      = 0x80 // 128
@@ -863,6 +873,20 @@ func getBytesLength(b []byte, e escapes) (int, error) {
 	}
 }
 
+// getLTreeLength finds the length of a ltree encoding.
+func getLTreeLength(b []byte, dir Direction) (int, error) {
+	var i int
+	if dir == Ascending {
+		i = bytes.IndexByte(b, ltreeKeyTerminator)
+	} else {
+		i = bytes.IndexByte(b, ltreeKeyDescendingTerminator)
+	}
+	if i == -1 {
+		return 0, errors.Errorf("did not find terminator")
+	}
+	return i + 1, nil
+}
+
 // prettyPrintInvertedIndexKey returns a string representation of the path part of a JSON inverted
 // index.
 func prettyPrintInvertedIndexKey(b []byte) (string, []byte, error) {
@@ -1726,6 +1750,101 @@ func DecodeBitArrayDescending(b []byte) ([]byte, bitarray.BitArray, error) {
 	return b, ba, err
 }
 
+// EncodeLTreeAscending encodes a ltree.T value, appends it to the
+// supplied buffer, and returns the final buffer. The encoding is guaranteed to
+// be ordered such that if t1 < t2 then bytes.Compare will order them the same
+// way after encoding.
+//
+// The encoding is in the below format:
+// [ ltreeMarker ]
+// for each label:
+//
+//	[ label raw bytes ] [ ltreeLabelKeyTerminator ]
+//
+// [ ltreeKeyTerminator ]
+func EncodeLTreeAscending(b []byte, d ltree.T) []byte {
+	b = append(b, ltreeKeyMarker)
+	d.ForEachLabel(func(i int, label string) {
+		b = append(b, []byte(label)...)
+		b = append(b, ltreeLabelKeyTerminator)
+	})
+	b = append(b, ltreeKeyTerminator)
+	return b
+}
+
+// EncodeLTreeDescending is the descending version of EncodeLTreeAscending.
+func EncodeLTreeDescending(b []byte, d ltree.T) []byte {
+	b = append(b, ltreeKeyDescendingMarker)
+	d.ForEachLabel(func(i int, label string) {
+		n := len(b)
+		b = append(b, []byte(label)...)
+		onesComplement(b[n:])
+		b = append(b, ltreeLabelKeyDescendingTerminator)
+	})
+	b = append(b, ltreeKeyDescendingTerminator)
+	return b
+}
+
+// DecodeLTreeAscending decodes a ltree.T value which was encoded using
+// EncodeLTreeAscending. The remainder of the input buffer and the
+// decoded ltree.T are returned.
+func DecodeLTreeAscending(b []byte) ([]byte, ltree.T, error) {
+	if PeekType(b) != LTree {
+		return nil, ltree.Empty, errors.Errorf("did not find marker %#x", b)
+	}
+	b = b[1:]
+
+	var labels []string
+	for {
+		if len(b) != 0 && b[0] == ltreeKeyTerminator {
+			b = b[1:]
+			break
+		}
+		i := bytes.IndexByte(b, ltreeLabelKeyTerminator)
+		if i == -1 {
+			return nil, ltree.Empty, errors.Errorf("malformed ltree encoding")
+		}
+		labels = append(labels, string(b[:i]))
+		b = b[i+1:]
+	}
+	l, err := ltree.ParseLTreeFromLabels(labels)
+	if err != nil {
+		return nil, ltree.Empty, err
+	}
+	return b, l, nil
+}
+
+// DecodeLTreeDescending is the descending version of DecodeLTreeAscending.
+func DecodeLTreeDescending(b []byte) ([]byte, ltree.T, error) {
+	if PeekType(b) != LTreeDesc {
+		return nil, ltree.Empty, errors.Errorf("did not find marker %#x", b)
+	}
+	b = b[1:]
+
+	var labels []string
+	for {
+		if len(b) != 0 && b[0] == ltreeKeyDescendingTerminator {
+			b = b[1:]
+			break
+		}
+		i := bytes.IndexByte(b, ltreeLabelKeyDescendingTerminator)
+		if i == -1 {
+			return nil, ltree.Empty, errors.Errorf("malformed ltree encoding")
+		}
+		// Deep copying here is necessary to avoid modifying the input buffer slice.
+		var label []byte
+		label = append(label, b[:i]...)
+		onesComplement(label)
+		labels = append(labels, string(label))
+		b = b[i+1:]
+	}
+	l, err := ltree.ParseLTreeFromLabels(labels)
+	if err != nil {
+		return nil, ltree.Empty, err
+	}
+	return b, l, nil
+}
+
 // Type represents the type of a value encoded by
 // Encode{Null,NotNull,Varint,Uvarint,Float,Bytes}.
 //
@@ -1788,6 +1907,7 @@ const (
 	JsonEmptyArrayDesc Type = 43
 	PGVector           Type = 44
 	LTree              Type = 45
+	LTreeDesc          Type = 46
 )
 
 // typMap maps an encoded type byte to a decoded Type. It's got 256 slots, one
@@ -1890,6 +2010,10 @@ func slowPeekType(b []byte) Type {
 			return Decimal
 		case m == voidMarker:
 			return Void
+		case m == ltreeKeyMarker:
+			return LTree
+		case m == ltreeKeyDescendingMarker:
+			return LTreeDesc
 		}
 	}
 	return Unknown
@@ -2086,6 +2210,10 @@ func PeekLength(b []byte) (int, error) {
 			return 0, errors.Errorf("slice too short for float (%d)", len(b))
 		}
 		return 9, nil
+	case ltreeKeyMarker:
+		return getLTreeLength(b, Ascending)
+	case ltreeKeyDescendingMarker:
+		return getLTreeLength(b, Descending)
 	}
 	if m >= IntMin && m <= IntMax {
 		return getVarintLen(b)
@@ -2384,6 +2512,26 @@ func prettyPrintFirstValue(dir Direction, b []byte) ([]byte, string, error) {
 			return b, "", err
 		}
 		return b, d.StringNanos(), nil
+	case LTree:
+		if dir == Descending {
+			return b, "", errors.Errorf("ascending ltree column dir but descending ltree encoding")
+		}
+		var l ltree.T
+		b, l, err = DecodeLTreeAscending(b)
+		if err != nil {
+			return b, "", err
+		}
+		return b, l.String(), nil
+	case LTreeDesc:
+		if dir == Ascending {
+			return b, "", errors.Errorf("descending ltree column dir but ascending ltree encoding")
+		}
+		var l ltree.T
+		b, l, err = DecodeLTreeDescending(b)
+		if err != nil {
+			return b, "", err
+		}
+		return b, l.String(), nil
 	default:
 		if len(b) >= 1 {
 			switch b[0] {