Merge #154914

154914: rowexec: fix inverted join with prefix columns with DESC direction r=yuzefovich a=yuzefovich

This commit fixes a bug that was added long time ago when we added support for inverted joins on multi-column inverted indexes (i.e. such indexes that have a prefix of "forward" columns that we can constrain using an equality filter) in d3f9a8b. In particular, in that change we added tricky logic which prepended the prefix to the encoded inverted key; however, we then didn't change how we use the `span.Builder`. That guy assumes that it receives datums for each index column (both prefix and inverted) whereas we only provided the inverted bytes (with the prepended prefix). This led to us being confused about which key column is being encoded during span construction.

Consider an example where we have one prefix key column like `i INT DESC` (followed by the inverted column). In this case, when calling `MakeKeyFromEncDatums` we'd try to encode the single EncDatum using the DESC key direction. The EncDatum internally stores the prefix + inverted key using ASC key direction, so we'd try to decode it and fail - because of the prepended prefix. After this patch, we'll correctly only use a single column in the call to `MakeKeyFromEncDatums`.

Note that if prefix columns had only ASC directions, then we would simply get lucky. Namely, since we have EncDatumRow with a single EncDatum AND that EncDatum already has encoded value of ASC direction stored, we'd just reuse it and exit.

The code here is quite difficult to follow and perhaps it could be refactored (e.g. the comment on `SpansFromInvertedSpans` says that multi-column inverted indexes are expected to use non-nil constraint, yet in the inverted joiner we always pass nil), but I chose a more targetted fix that seems reasonable.

Additionally, to simplify the change a bit, I inlined the code for `rowenc.MakeSpanFromEncDatums` into two call sites (which also allows us to avoid constructing the unnecessary end key for the inverted spans case).

Fixes: #153499.
Release note (bug fix): Previously, CockroachDB would hit an internal error when performing an inverted join using an inverted index in which the first prefix column had DESC direction. The bug has been present since 21.1 and is now fixed.

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

Author: craig[bot]

pkg/sql/distsql_spec_exec_factory.go

@@ -282,7 +282,7 @@ func (e *distSQLSpecExecFactory) ConstructScan(
 	var spans roachpb.Spans
 	var err error
 	if params.InvertedConstraint != nil {
-		spans, err = sb.SpansFromInvertedSpans(e.ctx, params.InvertedConstraint, params.IndexConstraint, nil /* scratch */)
+		spans, err = sb.SpansFromInvertedSpans(e.ctx, params.InvertedConstraint, params.IndexConstraint, false /* prefixIncludedInKeys */, nil /* scratch */)
 	} else {
 		var splitter span.Splitter
 		if params.Locking.MustLockAllRequestedColumnFamilies() {

pkg/sql/logictest/testdata/logic_test/inverted_join_multi_column

@@ -61,7 +61,7 @@ CREATE TABLE j2 (
   i INT,
   j JSON,
   s STRING,
-  INVERTED INDEX ij_idx (i, j),
+  INVERTED INDEX ij_idx (i DESC, j),
   INVERTED INDEX isj_idx (i, s, j)
 )
 

pkg/sql/opt/exec/execbuilder/testdata/inverted_index_multi_column

@@ -7,15 +7,15 @@ CREATE TABLE t (
     s STRING,
     j JSON,
     FAMILY (k, i, s, j),
-    INVERTED INDEX (i, j),
+    INVERTED INDEX (i DESC, j),
     INVERTED INDEX (i, s, j)
 )
 
 query T kvtrace
 SELECT k FROM t WHERE i = 10 AND j @> '1'
 ----
 Scan /Table/20/1/10{6-7}
-Scan /Table/106/2/10/1{-/PrefixEnd}, /Table/106/2/10/Arr/1{-/PrefixEnd}
+Scan /Table/106/2/-11/1{-/PrefixEnd}, /Table/106/2/-11/Arr/1{-/PrefixEnd}
 
 query T kvtrace
 SELECT k FROM t WHERE i = 10 AND s = 'foo' AND j @> '1'
@@ -39,7 +39,7 @@ Put /Table/106/3/333/"foo"/"a"/"b"/1/0 -> /BYTES/
 query T kvtrace
 SELECT * FROM t WHERE i = 333 AND j @> '{"a": "b"}'
 ----
-Scan /Table/106/2/333/"a"/"b"{-/PrefixEnd}
+Scan /Table/106/2/-334/"a"/"b"{-/PrefixEnd}
 Scan /Table/106/1/1/0
 
 # Don't insert duplicate values.
@@ -103,7 +103,7 @@ UPDATE t SET j = NULL WHERE k = 4
 ----
 Scan /Table/106/1/4/0 lock Exclusive (Block, Unreplicated)
 Put /Table/106/1/4/0 -> /TUPLE/2:2:Int/333/1:3:Bytes/foo
-Del /Table/106/2/333/Arr/1/4/0
+Del /Table/106/2/-334/Arr/1/4/0
 Del /Table/106/3/333/"foo"/Arr/1/4/0
 
 # Deleting a NULL shouldn't remove anything from the inv idx.
@@ -138,7 +138,7 @@ UPDATE t SET i = NULL WHERE k = 5
 ----
 Scan /Table/106/1/5/0 lock Exclusive (Block, Unreplicated)
 Put /Table/106/1/5/0 -> /TUPLE/3:3:Bytes/foo/1:4:SentinelType/{"a": "b"}
-Del /Table/106/2/333/"a"/"b"/5/0
+Del /Table/106/2/-334/"a"/"b"/5/0
 Put /Table/106/2/NULL/"a"/"b"/5/0 -> /BYTES/
 Del /Table/106/3/333/"foo"/"a"/"b"/5/0
 Put /Table/106/3/NULL/"foo"/"a"/"b"/5/0 -> /BYTES/

pkg/sql/opt_exec_factory.go

@@ -179,7 +179,7 @@ func generateScanSpans(
 	var sb span.Builder
 	sb.InitAllowingExternalRowData(evalCtx, codec, tabDesc, index)
 	if params.InvertedConstraint != nil {
-		return sb.SpansFromInvertedSpans(ctx, params.InvertedConstraint, params.IndexConstraint, nil /* scratch */)
+		return sb.SpansFromInvertedSpans(ctx, params.InvertedConstraint, params.IndexConstraint, false /* prefixIncludedInKeys */, nil /* scratch */)
 	}
 	var splitter span.Splitter
 	if params.Locking.MustLockAllRequestedColumnFamilies() {

pkg/sql/rowenc/index_encoding.go

@@ -141,26 +141,6 @@ func (d Directions) Get(i int) (encoding.Direction, error) {
 	return encoding.Ascending, nil
 }
 
-// MakeSpanFromEncDatums creates a minimal index key span on the input
-// values. A minimal index key span is a span that includes the fewest possible
-// keys after the start key generated by the input values.
-//
-// The start key is generated by concatenating keyPrefix with the encodings of
-// the given EncDatum values. The values, types, and dirs parameters should be
-// specified in the same order as the index key columns and may be a prefix.
-func MakeSpanFromEncDatums(
-	values EncDatumRow,
-	keyCols []fetchpb.IndexFetchSpec_KeyColumn,
-	alloc *tree.DatumAlloc,
-	keyPrefix []byte,
-) (_ roachpb.Span, containsNull bool, _ error) {
-	startKey, containsNull, err := MakeKeyFromEncDatums(values, keyCols, alloc, keyPrefix)
-	if err != nil {
-		return roachpb.Span{}, false, err
-	}
-	return roachpb.Span{Key: startKey, EndKey: startKey.PrefixEnd()}, containsNull, nil
-}
-
 // NeededColumnFamilyIDs returns the minimal set of column families required to
 // retrieve neededCols for the specified table and index. The returned
 // descpb.FamilyIDs are in sorted order. If forSideEffect is true, column

pkg/sql/rowexec/inverted_joiner.go

@@ -499,7 +499,8 @@ func (ij *invertedJoiner) readInput() (invertedJoinerState, *execinfrapb.Produce
 	}
 	// NB: spans is already sorted, and that sorting is preserved when
 	// generating ij.indexSpans.
-	ij.indexSpans, err = ij.spanBuilder.SpansFromInvertedSpans(ij.Ctx(), spans, nil /* constraint */, ij.indexSpans)
+	prefixIncludedInKeys := len(ij.prefixEqualityCols) > 0
+	ij.indexSpans, err = ij.spanBuilder.SpansFromInvertedSpans(ij.Ctx(), spans, nil /* constraint */, prefixIncludedInKeys, ij.indexSpans)
 	if err != nil {
 		ij.MoveToDraining(err)
 		return ijStateUnknown, ij.DrainHelper()

pkg/sql/span/span_builder.go

@@ -109,7 +109,11 @@ func (s *Builder) IsPreEncoded(values rowenc.EncDatumRow) bool {
 func (s *Builder) SpanFromEncDatums(
 	values rowenc.EncDatumRow,
 ) (_ roachpb.Span, containsNull bool, _ error) {
-	return rowenc.MakeSpanFromEncDatums(values, s.keyAndPrefixCols, &s.alloc, s.KeyPrefix)
+	startKey, containsNull, err := rowenc.MakeKeyFromEncDatums(values, s.keyAndPrefixCols, &s.alloc, s.KeyPrefix)
+	if err != nil {
+		return roachpb.Span{}, false, err
+	}
+	return roachpb.Span{Key: startKey, EndKey: startKey.PrefixEnd()}, containsNull, nil
 }
 
 // SpanFromEncDatumsWithRange encodes a range span. The inequality is assumed to
@@ -408,16 +412,31 @@ var _ InvertedSpans = inverted.SpanExpressionProtoSpans{}
 // If the index is a multi-column inverted index, c should constrain the
 // non-inverted prefix columns of the index. Each span in c must have a single
 // key. The resulting roachpb.Spans are created by performing a cross product of
-// keys in c and the invertedSpan keys.
+// keys in c and the invertedSpan keys. Cannot be used in combination with
+// prefixIncludedInKeys.
+//
+// prefixIncludedInKeys, if set, indicates that we have a multi-column inverted
+// index AND the non-inverted prefix columns have already been prepended into
+// the inverted spans. Cannot be used in combination with non-nil constraint.
+// TODO(yuzefovich): consider splitting this method into two: one which accepts
+// the constraint and constructs a cross-product of spans, and another which
+// already has the inverted spans fully encoded.
 //
 // scratch can be an optional roachpb.Spans slice that will be reused to
 // populate the result.
 func (s *Builder) SpansFromInvertedSpans(
-	ctx context.Context, invertedSpans InvertedSpans, c *constraint.Constraint, scratch roachpb.Spans,
+	ctx context.Context,
+	invertedSpans InvertedSpans,
+	c *constraint.Constraint,
+	prefixIncludedInKeys bool,
+	scratch roachpb.Spans,
 ) (roachpb.Spans, error) {
 	if invertedSpans == nil {
 		return nil, errors.AssertionFailedf("invertedSpans cannot be nil")
 	}
+	if c != nil && prefixIncludedInKeys {
+		return nil, errors.AssertionFailedf("constraint and prefixIncludedInKeys cannot be used at the same time")
+	}
 
 	var scratchRows []rowenc.EncDatumRow
 	if c != nil {
@@ -453,10 +472,10 @@ func (s *Builder) SpansFromInvertedSpans(
 		for j, n := 0, invertedSpans.Len(); j < n; j++ {
 			var indexSpan roachpb.Span
 			var err error
-			if indexSpan.Key, err = s.generateInvertedSpanKey(invertedSpans.Start(j), scratchRows[i]); err != nil {
+			if indexSpan.Key, err = s.generateInvertedSpanKey(invertedSpans.Start(j), scratchRows[i], prefixIncludedInKeys); err != nil {
 				return nil, err
 			}
-			if indexSpan.EndKey, err = s.generateInvertedSpanKey(invertedSpans.End(j), scratchRows[i]); err != nil {
+			if indexSpan.EndKey, err = s.generateInvertedSpanKey(invertedSpans.End(j), scratchRows[i], prefixIncludedInKeys); err != nil {
 				return nil, err
 			}
 			scratch = append(scratch, indexSpan)
@@ -471,7 +490,7 @@ func (s *Builder) SpansFromInvertedSpans(
 // of scratchRow is greater than one, the EncDatums that precede the last slot
 // are encoded as prefix keys of enc.
 func (s *Builder) generateInvertedSpanKey(
-	enc []byte, scratchRow rowenc.EncDatumRow,
+	enc []byte, scratchRow rowenc.EncDatumRow, prefixIncludedInKeys bool,
 ) (roachpb.Key, error) {
 	keyLen := len(scratchRow) - 1
 	scratchRow = scratchRow[:keyLen]
@@ -489,8 +508,14 @@ func (s *Builder) generateInvertedSpanKey(
 	// generate a span, of which we will only use Span.Key. Span.EndKey is
 	// generated by the caller in the second call, with RKeyMax.
 
-	span, _, err := s.SpanFromEncDatums(scratchRow[:keyLen])
-	return span.Key, err
+	keyCols := s.keyAndPrefixCols
+	if prefixIncludedInKeys {
+		// When prefix is already included in keys, then we only have a single
+		// "value" which is prefix + encoded inverted value.
+		keyCols = keyCols[len(keyCols)-1:]
+	}
+	startKey, _, err := rowenc.MakeKeyFromEncDatums(scratchRow[:keyLen], keyCols, &s.alloc, s.KeyPrefix)
+	return startKey, err
 }
 
 // KeysFromVectorPrefixConstraint extracts the encoded prefix keys from a