rowexec: fix inverted join with prefix columns with DESC direction

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