Merge #152353 #152523

152353: opt: index accelerate LTREE with ancestry operators, @> and <@ r=paulniziolek a=paulniziolek

#### opt: index accelerate `@>` operator for ltree

The `@>` operator is accelerated by restricting the spans to be the union of
equality spans for all sub-ltrees rooted at ''. For example, a query with
predicate `WHERE a `@>` 'A.B'` would create spans [/'' - /''], [/'A' - /'A'],
[/'A.B' - /'A.B'].

Informs: #44657
Epic: CRDB-148

Release note (performance improvement): LTREE is now index accelerated with the
``@>`` operator.

#### opt: index accelerate <@ operator for ltree

The <@ is index accelerated by restricting the span of key-encoded ltrees to be
between a given ltree and the ltree with an incremented last label. For
example, a query with predicate `WHERE a <@ 'A.B'` would create the span
[/'A.B' - /'A.C'].

Informs: #44657
Epic: CRDB-148

Release note (performance improvement): LTREE is now index accelerated with the
`<`@`` operator.


152523: sql: harden recent fix to deleteRange r=yuzefovich a=yuzefovich

In recently merged 59a28ec we fixed how we count "rows affected" by the deleteRangeNode, which was done by maintaining "cur row prefix" across BatchRequests. AI-generated code review pointed out that we can alias the memory of now-old BatchRequest while processing the response to the new one. Although I don't think it can lead to problems (since we shouldn't be modifying the BatchRequest's or BatchResponse's keys - which is verified via `GRPCTransportFactory` "race" variant), it seems prudent that we make a copy of the row prefix (among other benefits, this might allow for the old keys to be GCed sooner).

Epic: None
Release note: None

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

Author: 3 people

pkg/sql/delete_range.go

@@ -111,6 +111,7 @@ func (d *deleteRangeNode) startExec(params runParams) error {
 
 	ctx := params.ctx
 	log.VEvent(ctx, 2, "fast delete: skipping scan")
+	// TODO(yuzefovich): why are we making a copy of spans?
 	spans := make([]roachpb.Span, len(d.spans))
 	copy(spans, d.spans)
 	if !d.autoCommitEnabled {
@@ -198,6 +199,18 @@ func (d *deleteRangeNode) deleteSpans(params runParams, b *kv.Batch, spans roach
 func (d *deleteRangeNode) processResults(
 	results []kv.Result, resumeSpans []roachpb.Span,
 ) (roachpb.Spans, error) {
+	if !d.autoCommitEnabled {
+		defer func() {
+			// Make a copy of curRowPrefix to avoid referencing the memory from
+			// the now-old BatchRequest.
+			//
+			// When auto-commit is enabled, we expect to see any resume spans,
+			// so we won't need to access d.curRowPrefix later.
+			curRowPrefix := make([]byte, len(d.curRowPrefix))
+			copy(curRowPrefix, d.curRowPrefix)
+			d.curRowPrefix = curRowPrefix
+		}()
+	}
 	for _, r := range results {
 		// TODO(yuzefovich): when the table has 1 column family, we don't need
 		// to compare the key prefixes since each deleted key corresponds to a

pkg/sql/opt/exec/execbuilder/testdata/select_index

@@ -1904,3 +1904,36 @@ vectorized: true
   missing stats
   table: t4@t4_pkey
   spans: [/1/1 - /1/1] [/1/5 - /1/5] [/5/1 - /5/1] [/5/5 - /5/5]
+
+# ------------------------------------------------------------------------------
+# Tests for index acceleration on ancestry operators (@> and <@) using
+# LTREE columns
+# ------------------------------------------------------------------------------
+
+statement ok
+CREATE TABLE t5 (a LTREE)
+
+statement ok
+CREATE INDEX t5_a_idx ON t5 (a)
+
+query T
+EXPLAIN SELECT a FROM t5 WHERE a @> 'A.B.C'
+----
+distribution: local
+vectorized: true
+·
+• scan
+  missing stats
+  table: t5@t5_a_idx
+  spans: [/'' - /''] [/'A' - /'A'] [/'A.B' - /'A.B'] [/'A.B.C' - /'A.B.C']
+
+query T
+EXPLAIN SELECT a FROM t5 WHERE a <@ 'A.B.C'
+----
+distribution: local
+vectorized: true
+·
+• scan
+  missing stats
+  table: t5@t5_a_idx
+  spans: [/'A.B.C' - /'A.B.D')

pkg/sql/opt/idxconstraint/BUILD.bazel

@@ -16,6 +16,7 @@ go_library(
         "//pkg/sql/sem/tree",
         "//pkg/sql/types",
         "//pkg/util",
+        "//pkg/util/ltree",
         "@com_github_cockroachdb_errors//:errors",
     ],
 )

pkg/sql/opt/idxconstraint/index_constraints.go

@@ -20,6 +20,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util"
+	"github.com/cockroachdb/cockroach/pkg/util/ltree"
 	"github.com/cockroachdb/errors"
 )
 
@@ -338,6 +339,56 @@ func (c *indexConstraintCtx) makeSpansForSingleColumnDatum(
 				return complete
 			}
 		}
+
+	case opt.ContainsOp:
+		if l, ok := datum.(*tree.DLTree); ok {
+			var spans constraint.Spans
+			// We need to create an equality span for each subtree of the LTree that
+			// is rooted from the root, including the empty ltree.
+			spans.Alloc(l.LTree.Len() + 1)
+			keyCtx := &c.keyCtx[offset]
+			for i := 0; i <= l.LTree.Len(); i++ {
+				var subLTree ltree.T
+				if l.LTree.Compare(ltree.Empty) != 0 {
+					var err error
+					subLTree, err = l.LTree.SubPath(0 /* offset */, i /* length */)
+					if err != nil {
+						panic(err)
+					}
+				} else {
+					// SubPath is not graceful with empty ltree, thus we handle it here.
+					subLTree = ltree.Empty
+				}
+				key := constraint.MakeKey(tree.NewDLTree(subLTree))
+				var sp constraint.Span
+				sp.Init(key, includeBoundary, key, includeBoundary)
+				spans.Append(&sp)
+			}
+			// Given how we've constructed the spans, they already are ordered and
+			// unique, but we choose to call SortAndMerge for symmetry with other
+			// expressions and as a sanity check (the function exits quickly if the
+			// ordering is already correct).
+			spans.SortAndMerge(keyCtx)
+			out.Init(keyCtx, &spans)
+			return true
+		}
+
+	case opt.ContainedByOp:
+		if l, ok := datum.(*tree.DLTree); ok {
+			if l.LTree.Compare(ltree.Empty) == 0 {
+				// Empty LTree shouldn't be constrained.
+				// TODO: This could be constrained by excluding NULLs.
+				break
+			}
+			startKey := constraint.MakeKey(l)
+			endKey := constraint.MakeKey(tree.NewDLTree(l.LTree.NextSibling()))
+			c.singleSpan(
+				offset, startKey, includeBoundary, endKey, excludeBoundary,
+				c.columns[offset].Descending(),
+				out,
+			)
+			return true
+		}
 	}
 	c.unconstrained(offset, out)
 	return false

pkg/sql/opt/idxconstraint/testdata/ltree

@@ -0,0 +1,79 @@
+index-constraints vars=(a ltree) index=a
+a @> '-'
+----
+[/'' - /'']
+[/'-' - /'-']
+
+index-constraints vars=(a ltree) index=a
+a @> 'A.B.C'
+----
+[/'' - /'']
+[/'A' - /'A']
+[/'A.B' - /'A.B']
+[/'A.B.C' - /'A.B.C']
+
+index-constraints vars=(a ltree) index=a
+a @> ''
+----
+[/'' - /'']
+
+index-constraints vars=(a ltree) index=a
+a @> NULL
+----
+
+
+index-constraints vars=(a ltree) index=a
+a <@ 'A.B.C'
+----
+[/'A.B.C' - /'A.B.D')
+
+index-constraints vars=(a ltree) index=a
+a <@ 'z'
+----
+[/'z' - /'z-')
+
+index-constraints vars=(a ltree) index=a
+a <@ ''
+----
+[ - ]
+Remaining filter: a <@ ''
+
+index-constraints vars=(a ltree) index=a
+a <@ NULL
+----
+
+index-constraints vars=(a ltree) index=a
+a @> '-' OR a @> 'foo.bar'
+----
+[/'' - /'']
+[/'-' - /'-']
+[/'foo' - /'foo']
+[/'foo.bar' - /'foo.bar']
+
+index-constraints vars=(a ltree) index=a
+a <@ '-' OR a <@ 'foo.bar'
+----
+[/'-' - /'0')
+[/'foo.bar' - /'foo.bas')
+
+index-constraints vars=(a ltree) index=a
+a <@ 'A.B.C' OR a @> 'A.B'
+----
+[/'' - /'']
+[/'A' - /'A']
+[/'A.B' - /'A.B']
+[/'A.B.C' - /'A.B.D')
+
+index-constraints vars=(a ltree) index=a
+a @> 'A.B.C' OR a <@ 'A.B'
+----
+[/'' - /'']
+[/'A' - /'A']
+[/'A.B' - /'A.C')
+
+index-constraints vars=(a ltree) index=a
+a <@ 'A.B.C' AND a @> 'A.B.C.D.E'
+----
+[/'A.B.C' - /'A.B.C']
+[/'A.B.C.D' - /'A.B.C.D']
+[/'A.B.C.D.E' - /'A.B.C.D.E']

pkg/util/ltree/ltree.go

@@ -295,3 +295,46 @@ func LCA(ltrees []T) (_ T, isNull bool) { // lint: uppercase function OK
 
 	return T{path: ltrees[0].path[:i:i]}, false
 }
+
+// NextSibling returns a LTree with a lexicographically incremented last label.
+// This is different from the next lexicographic LTree. This is mainly used for
+// defining a key-encoded span for ancestry operators.
+// Note that this could produce a LTREE with more labels than the maximum
+// allowed.
+// Example: 'A.B' -> 'A.C'
+func (lt T) NextSibling() T {
+	if lt.Len() == 0 {
+		return Empty
+	}
+	lastLabel := lt.path[len(lt.path)-1]
+	nextLabel := incrementLabel(lastLabel)
+
+	newLTree := lt.Copy()
+	newLTree.path[newLTree.Len()-1] = nextLabel
+	return newLTree
+}
+
+var nextCharMap = map[byte]byte{
+	'-': '0',
+	'9': 'A',
+	'Z': '_',
+	'_': 'a',
+	'z': 0,
+}
+
+func incrementLabel(label string) string {
+	nextLabel := []byte(label)
+	nextChar, ok := nextCharMap[nextLabel[len(nextLabel)-1]]
+
+	if ok && nextChar == 0 {
+		// Technically, this could mean exceeding the length of the label.
+		return string(append(nextLabel, '-'))
+	}
+
+	if !ok {
+		nextChar = nextLabel[len(nextLabel)-1] + 1
+	}
+
+	nextLabel[len(nextLabel)-1] = nextChar
+	return string(nextLabel)
+}

pkg/util/ltree/ltree_test.go

@@ -152,3 +152,29 @@ func TestCompare(t *testing.T) {
 		}
 	}
 }
+
+func TestNext(t *testing.T) {
+	tests := []struct {
+		input    string
+		expected string
+	}{
+		{"a", "b"},
+		{"a.b", "a.c"},
+		{"a.z", "a.z-"},
+		{"-", "0"},
+		{"9", "A"},
+		{"Z", "_"},
+		{"_", "a"},
+	}
+
+	for _, tc := range tests {
+		a, err := ParseLTree(tc.input)
+		if err != nil {
+			t.Fatalf("unexpected error parsing %q: %v", tc.input, err)
+		}
+		next := a.NextSibling()
+		if next.String() != tc.expected {
+			t.Errorf("expected next of %q to be %q, got %q", tc.input, tc.expected, next.String())
+		}
+	}
+}