From 1d42f9f238d89bb8c6189d45d0ce6d023ad205ea Mon Sep 17 00:00:00 2001
From: Rohan Yadav <rohany@alumni.cmu.edu>
Date: Thu, 22 Aug 2019 17:25:21 -0400
Subject: [PATCH] sql: Bugfixes and increased tests for IndexSkipTableReader

Fixes #39659.

Release note: None
---
 pkg/sql/distsqlrun/index_skip_table_reader.go | 19 +++-----
 .../index_skip_table_reader_test.go           | 45 +++++++++++++++++++
 2 files changed, 50 insertions(+), 14 deletions(-)

diff --git a/pkg/sql/distsqlrun/index_skip_table_reader.go b/pkg/sql/distsqlrun/index_skip_table_reader.go
index 40edd775be46..131e657360f5 100644
--- a/pkg/sql/distsqlrun/index_skip_table_reader.go
+++ b/pkg/sql/distsqlrun/index_skip_table_reader.go
@@ -181,7 +181,7 @@ func (t *indexSkipTableReader) Next() (sqlbase.EncDatumRow, *distsqlpb.ProducerM
 			}
 		}
 
-		var key []byte
+		var key roachpb.Key
 		{
 			// Fetch the PartialKey in a separate scope so that it cannot be reused
 			// outside this scope.
@@ -212,19 +212,10 @@ func (t *indexSkipTableReader) Next() (sqlbase.EncDatumRow, *distsqlpb.ProducerM
 		}
 
 		if !t.reverse {
-			// 0xff is the largest prefix marker for any encoded key. To ensure that
-			// our new key is larger than any value with the same prefix, we place
-			// 0xff at all other index column values, and one more to guard against
-			// 0xff present as a value in the table (0xff encodes a type of null).
-			// TODO(asubiotto): We should delegate to PrefixEnd here (it exists for
-			//  this purpose and is widely used), but we still need to handle maximal
-			//  keys. Maybe we should just exit early (t.currentSpan++) when we
-			//  encounter one. We also need a test to ensure that we behave properly
-			//  when we encounter maximal (i.e. all nulls) prefix keys.
-			for i := 0; i < (t.indexLen - t.keyPrefixLen + 1); i++ {
-				key = append(key, 0xff)
-			}
-			t.spans[t.currentSpan].Key = key
+			// We set the new key to be the largest key with the prefix that we have
+			// so that we skip all values with the same prefix, and "skip" to the
+			// next distinct value.
+			t.spans[t.currentSpan].Key = key.PrefixEnd()
 		} else {
 			// In the case of reverse, this is much easier. The reverse fetcher
 			// returns the key retrieved, in this case the first key smaller
diff --git a/pkg/sql/distsqlrun/index_skip_table_reader_test.go b/pkg/sql/distsqlrun/index_skip_table_reader_test.go
index f0f7abb23efa..d3aa732f2217 100644
--- a/pkg/sql/distsqlrun/index_skip_table_reader_test.go
+++ b/pkg/sql/distsqlrun/index_skip_table_reader_test.go
@@ -141,12 +141,29 @@ func TestIndexSkipTableReader(t *testing.T) {
 		99,
 		sqlutils.ToRowFn(xFnt6, yFnt6))
 
+	// create a table t7 where each row is:
+	//
+	// |     x     |   y   |  z  |
+	// |-------------------------|
+	// | rowId%10  | NULL | NULL|
+	xFnt7 := func(row int) tree.Datum {
+		return tree.NewDInt(tree.DInt(row % 10))
+	}
+	nullt7 := func(_ int) tree.Datum {
+		return tree.DNull
+	}
+	sqlutils.CreateTable(t, sqlDB, "t7",
+		"x INT, y INT, z INT, PRIMARY KEY (x), INDEX i1 (x, y DESC, z DESC), INDEX i2 (y DESC, z DESC)",
+		10,
+		sqlutils.ToRowFn(xFnt7, nullt7, nullt7))
+
 	td1 := sqlbase.GetTableDescriptor(kvDB, "test", "t1")
 	td2 := sqlbase.GetTableDescriptor(kvDB, "test", "t2")
 	td3 := sqlbase.GetTableDescriptor(kvDB, "test", "t3")
 	td4 := sqlbase.GetTableDescriptor(kvDB, "test", "t4")
 	td5 := sqlbase.GetTableDescriptor(kvDB, "test", "t5")
 	td6 := sqlbase.GetTableDescriptor(kvDB, "test", "t6")
+	td7 := sqlbase.GetTableDescriptor(kvDB, "test", "t7")
 
 	makeIndexSpan := func(td *sqlbase.TableDescriptor, start, end int) distsqlpb.TableReaderSpan {
 		var span roachpb.Span
@@ -369,6 +386,34 @@ func TestIndexSkipTableReader(t *testing.T) {
 			},
 			expected: "[[10] [9] [8] [7] [6] [5] [4] [3] [2] [1]]",
 		},
+		{
+			// Distinct scan on index with multiple null values
+			desc:      "IndexMultipleNulls",
+			tableDesc: td7,
+			spec: distsqlpb.IndexSkipTableReaderSpec{
+				Spans:    []distsqlpb.TableReaderSpan{{Span: td7.IndexSpan(2)}},
+				IndexIdx: 1,
+			},
+			post: distsqlpb.PostProcessSpec{
+				Projection:    true,
+				OutputColumns: []uint32{0},
+			},
+			expected: "[[0] [1] [2] [3] [4] [5] [6] [7] [8] [9]]",
+		},
+		{
+			// Distinct scan on index with only null values
+			desc:      "IndexAllNulls",
+			tableDesc: td7,
+			spec: distsqlpb.IndexSkipTableReaderSpec{
+				Spans:    []distsqlpb.TableReaderSpan{{Span: td7.IndexSpan(3)}},
+				IndexIdx: 2,
+			},
+			post: distsqlpb.PostProcessSpec{
+				Projection:    true,
+				OutputColumns: []uint32{1},
+			},
+			expected: "[[NULL]]",
+		},
 	}
 
 	for _, c := range testCases {