From d5c94660536bb0c66394f79074210811e08db588 Mon Sep 17 00:00:00 2001
From: Radu Berinde <radu@cockroachlabs.com>
Date: Wed, 19 Apr 2017 16:43:15 -0400
Subject: [PATCH] sql: fix EncodePartialIndexKey function

The partial version of EncodeIndexKey wasn't stopping correctly once we encoded
the given columns. Added a testcase that was crashing with the old function.
---
 pkg/sql/sqlbase/table.go           | 19 ++++++++++++++++---
 pkg/sql/testdata/logic_test/ranges | 25 ++++++++++++++++++++++---
 2 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/pkg/sql/sqlbase/table.go b/pkg/sql/sqlbase/table.go
index c47bbad61d29..eee59e63d9ac 100644
--- a/pkg/sql/sqlbase/table.go
+++ b/pkg/sql/sqlbase/table.go
@@ -245,7 +245,7 @@ func EncodePartialIndexKey(
 ) (key []byte, containsNull bool, err error) {
 	colIDs := index.ColumnIDs[:numCols]
 	key = keyPrefix
-	dirs := directions(index.ColumnDirections)
+	dirs := directions(index.ColumnDirections)[:numCols]
 
 	if len(index.Interleave.Ancestors) > 0 {
 		for i, ancestor := range index.Interleave.Ancestors {
@@ -255,15 +255,25 @@ func EncodePartialIndexKey(
 				key = encoding.EncodeUvarintAscending(key, uint64(ancestor.IndexID))
 			}
 
+			partial := false
 			length := int(ancestor.SharedPrefixLen)
+			if length > len(colIDs) {
+				length = len(colIDs)
+				partial = true
+			}
 			var n bool
 			key, n, err = EncodeColumns(colIDs[:length], dirs[:length], colMap, values, key)
 			if err != nil {
 				return key, containsNull, err
 			}
-			colIDs, dirs = colIDs[length:], dirs[length:]
 			containsNull = containsNull || n
-
+			if partial {
+				// Early stop. Note that if we had exactly SharedPrefixLen columns
+				// remaining, we want to append the next tableID/indexID pair because
+				// that results in a more specific key.
+				return key, containsNull, nil
+			}
+			colIDs, dirs = colIDs[length:], dirs[length:]
 			// We reuse NotNullDescending (0xfe) as the interleave sentinel.
 			key = encoding.EncodeNotNullDescending(key)
 		}
@@ -296,6 +306,9 @@ func EncodeColumns(
 	values []parser.Datum,
 	keyPrefix []byte,
 ) (key []byte, containsNull bool, err error) {
+	if len(columnIDs) == 0 {
+		return keyPrefix, false, nil
+	}
 	// We know we will append to the key which will cause the capacity to grow
 	// so make it bigger from the get-go.
 	key = make([]byte, len(keyPrefix), 2*len(keyPrefix))
diff --git a/pkg/sql/testdata/logic_test/ranges b/pkg/sql/testdata/logic_test/ranges
index bcae8f759231..c1af1b121267 100644
--- a/pkg/sql/testdata/logic_test/ranges
+++ b/pkg/sql/testdata/logic_test/ranges
@@ -155,6 +155,22 @@ NULL           /1             {1}       1
 /7/8/#/52/1/9  /10            {1,2,4}   4
 /10            NULL           {1}       1
 
+statement ok
+ALTER TABLE t0 SPLIT AT VALUES (11)
+
+query TTTI colnames
+SHOW TESTING_RANGES FROM TABLE t0
+----
+Start Key      End Key        Replicas  Lease Holder
+NULL           /1             {1}       1
+/1             /5/1           {3,4}     3
+/5/1           /5/2           {1,2,3}   1
+/5/2           /5/3           {2,3,5}   5
+/5/3           /7/8/#/52/1/9  {1,2,4}   4
+/7/8/#/52/1/9  /10            {1,2,4}   4
+/10            /11            {1}       1
+/11            NULL           {1}       1
+
 query TTTI colnames
 SHOW TESTING_RANGES FROM TABLE t
 ----
@@ -165,7 +181,8 @@ NULL           /1             {1}       1
 /5/2           /5/3           {2,3,5}   5
 /5/3           /7/8/#/52/1/9  {1,2,4}   4
 /7/8/#/52/1/9  /10            {1,2,4}   4
-/10            NULL           {1}       1
+/10            /11            {1}       1
+/11            NULL           {1}       1
 
 statement ok
 CREATE TABLE t1 (k INT PRIMARY KEY, v1 INT, v2 INT, v3 INT)
@@ -184,7 +201,8 @@ NULL           /1             {1}       1
 /5/2           /5/3           {2,3,5}   5
 /5/3           /7/8/#/52/1/9  {1,2,4}   4
 /7/8/#/52/1/9  /10            {1,2,4}   4
-/10            NULL           {1}       1
+/10            /11            {1}       1
+/11            NULL           {1}       1
 
 statement ok
 ALTER INDEX t1@idx SPLIT AT VALUES (15,16)
@@ -199,7 +217,8 @@ NULL           /1             {1}       1
 /5/2           /5/3           {2,3,5}   5
 /5/3           /7/8/#/52/1/9  {1,2,4}   4
 /7/8/#/52/1/9  /10            {1,2,4}   4
-/10            /15/16/#/53/2  {1}       1
+/10            /11            {1}       1
+/11            /15/16/#/53/2  {1}       1
 /15/16/#/53/2  NULL           {1}       1
 
 statement error too many columns in SPLIT AT data