sql: add schema support for vector indexing in CREATE TABLE

Support usage of VECTOR INDEX in a CREATE TABLE statement, e.g.:

  CREATE TABLE simple (
    a INT PRIMARY KEY,
    vec1 VECTOR(3),
    VECTOR INDEX (vec1)
  )

Create the corresponding table and index schema objects for the vector
index. Check various error conditions, e.g. that only a column of type
VECTOR can be the last column in the index. Add unit and logic tests.
CREATE VECTOR INDEX support will come in a future PR.

Epic: CRDB-42943

Release note: None

Co-authored-by: Drew Kimball <[email protected]>

pkg/ccl/logictestccl/testdata/logic_test/refcursor

@@ -573,29 +573,29 @@ statement ok
 CREATE TABLE t112099_no_index (x INT NOT NULL, y TEXT NOT NULL, r REFCURSOR NOT NULL);
 
 # REFCURSOR is not allowed in a primary key.
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor and thus is not indexable
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor, which is not indexable
 CREATE TABLE t112099 (r REFCURSOR PRIMARY KEY);
 
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor and thus is not indexable
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor, which is not indexable
 CREATE TABLE t112099 (x INT, y TEXT, r REFCURSOR, PRIMARY KEY (x, r, y));
 
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor and thus is not indexable
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor, which is not indexable
 ALTER TABLE t112099_no_index ADD PRIMARY KEY (r);
 
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor and thus is not indexable
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor, which is not indexable
 ALTER TABLE t112099_no_index ADD PRIMARY KEY (x, r, y);
 
 # REFCURSOR is not allowed in a secondary index.
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor and thus is not indexable
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor, which is not indexable
 CREATE TABLE t112099 (r REFCURSOR, INDEX (r));
 
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor and thus is not indexable
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor, which is not indexable
 CREATE TABLE t112099 (x INT, y TEXT, r REFCURSOR, INDEX (x, r, y));
 
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor and thus is not indexable
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor, which is not indexable
 CREATE INDEX ON t112099_no_index (r);
 
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor and thus is not indexable
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor, which is not indexable
 CREATE INDEX ON t112099_no_index (x, r, y);
 
 # Regression test for #115701 - REFCURSOR[] is not a valid index column.
@@ -605,40 +605,35 @@ statement ok
 CREATE TABLE t115701_no_index (x INT NOT NULL, y TEXT NOT NULL, r REFCURSOR[] NOT NULL);
 
 # REFCURSOR[] is not allowed in a primary key.
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor\[\] and thus is not indexable
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor\[\], which is not indexable
 CREATE TABLE t115701 (r REFCURSOR[] PRIMARY KEY);
 
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor\[\] and thus is not indexable
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor\[\], which is not indexable
 CREATE TABLE t115701 (x INT, y TEXT, r REFCURSOR[], PRIMARY KEY (x, r, y));
 
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor\[\] and thus is not indexable
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor\[\], which is not indexable
 ALTER TABLE t115701_no_index ADD PRIMARY KEY (r);
 
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor\[\] and thus is not indexable
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor\[\], which is not indexable
 ALTER TABLE t115701_no_index ADD PRIMARY KEY (x, r, y);
 
 # REFCURSOR[] is not allowed in a secondary index.
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor\[\] and thus is not indexable
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor\[\], which is not indexable
 CREATE TABLE t115701 (r REFCURSOR[], INDEX (r));
 
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor\[\] and thus is not indexable
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor\[\], which is not indexable
 CREATE TABLE t115701 (x INT, y TEXT, r REFCURSOR[], INDEX (x, r, y));
 
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor\[\] and thus is not indexable
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor\[\], which is not indexable
 CREATE INDEX ON t115701_no_index (r);
 
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor\[\] and thus is not indexable
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor\[\], which is not indexable
 CREATE INDEX ON t115701_no_index (x, r, y);
 
-statement error pgcode 0A000 pq: column r of type refcursor\[\] is not allowed as the last column in an inverted index
+statement error pgcode 0A000 pq: column r has type refcursor\[\], which is not allowed as the last column in an inverted index
 CREATE INDEX ON t115701_no_index USING GIN (r)
 
-skipif config local-legacy-schema-changer
-statement error pgcode 0A000 pq: unimplemented: column r is of type refcursor\[\] and thus is not indexable
-CREATE INDEX ON t115701_no_index USING GIN (x, r, y gin_trgm_ops)
-
-onlyif config local-legacy-schema-changer
-statement error pgcode 0A000  pq: column r of type refcursor\[\] is only allowed as the last column in an inverted index
+statement error pgcode 0A000 pq: unimplemented: column r has type refcursor\[\], which is not indexable
 CREATE INDEX ON t115701_no_index USING GIN (x, r, y gin_trgm_ops)
 
 # REFCURSOR is allowed as a STORING column.

pkg/ccl/logictestccl/testdata/logic_test/vector

@@ -12,7 +12,7 @@ CREATE TABLE v (v vector(0))
 statement error pgcode 42601 dimensions for type vector cannot exceed 16000
 CREATE TABLE v (v vector(16001))
 
-statement error column v is of type vector and thus is not indexable
+statement error column v has type vector, which is not indexable in a non-vector index
 CREATE TABLE v (v vector(2) PRIMARY KEY)
 
 statement ok
@@ -189,7 +189,4 @@ CREATE INDEX ON t_vec USING CSPANN (v);
 statement ok
 DROP TABLE t_vec;
 
-statement error pgcode 0A000 pq: unimplemented: VECTOR indexes are not yet supported
-CREATE TABLE t_vec (k INT PRIMARY KEY, v VECTOR(128), VECTOR INDEX (v));
-
 subtest end

pkg/sql/catalog/BUILD.bazel

@@ -46,6 +46,7 @@ go_library(
         "//pkg/sql/sessiondatapb",
         "//pkg/sql/sqlclustersettings",
         "//pkg/sql/types",
+        "//pkg/sql/vecindex/vecpb",
         "//pkg/util",
         "//pkg/util/hlc",
         "//pkg/util/intsets",

pkg/sql/catalog/catformat/index.go

@@ -102,8 +102,13 @@ func indexForDisplay(
 	if index.Unique {
 		f.WriteString("UNIQUE ")
 	}
-	if !f.HasFlags(tree.FmtPGCatalog) && index.Type == idxtype.INVERTED {
-		f.WriteString("INVERTED ")
+	if !f.HasFlags(tree.FmtPGCatalog) {
+		switch index.Type {
+		case idxtype.INVERTED:
+			f.WriteString("INVERTED ")
+		case idxtype.VECTOR:
+			f.WriteString("VECTOR ")
+		}
 	}
 	f.WriteString("INDEX ")
 	f.FormatNameP(&index.Name)
@@ -114,9 +119,12 @@ func indexForDisplay(
 
 	if f.HasFlags(tree.FmtPGCatalog) {
 		f.WriteString(" USING")
-		if index.Type == idxtype.INVERTED {
+		switch index.Type {
+		case idxtype.INVERTED:
 			f.WriteString(" gin")
-		} else {
+		case idxtype.VECTOR:
+			f.WriteString(" cspann")
+		default:
 			f.WriteString(" btree")
 		}
 	}
@@ -240,17 +248,20 @@ func FormatIndexElements(
 		} else {
 			f.FormatNameP(&index.KeyColumnNames[i])
 		}
+		// TODO(drewk): we might need to print something like "vector_l2_ops" for
+		// vector indexes.
 		if index.Type == idxtype.INVERTED &&
 			col.GetID() == index.InvertedColumnID() && len(index.InvertedColumnKinds) > 0 {
 			switch index.InvertedColumnKinds[0] {
 			case catpb.InvertedIndexColumnKind_TRIGRAM:
 				f.WriteString(" gin_trgm_ops")
 			}
 		}
-		// The last column of an inverted index cannot have a DESC direction.
-		// Since the default direction is ASC, we omit the direction entirely
-		// for inverted index columns.
-		if i < n-1 || index.Type != idxtype.INVERTED {
+		// The last column of an inverted or vector index cannot have a DESC
+		// direction because it does not have a linear ordering. Since the default
+		// direction is ASC, we omit the direction entirely for inverted/vector
+		// index columns.
+		if i < n-1 || index.Type.HasLinearOrdering() {
 			f.WriteByte(' ')
 			f.WriteString(index.KeyColumnDirections[i].String())
 		}

pkg/sql/catalog/catformat/index_test.go

@@ -112,6 +112,12 @@ func TestIndexForDisplay(t *testing.T) {
 		ColumnNames:  []string{"a"},
 	}
 
+	// VECTOR INDEX baz (a)
+	vectorIndex := baseIndex
+	vectorIndex.Type = idxtype.VECTOR
+	vectorIndex.KeyColumnNames = []string{"a"}
+	vectorIndex.KeyColumnIDs = descpb.ColumnIDs{1}
+
 	testData := []struct {
 		index       descpb.IndexDescriptor
 		tableName   tree.TableName
@@ -266,6 +272,22 @@ func TestIndexForDisplay(t *testing.T) {
 			expected:    "CREATE INDEX baz ON foo.public.bar (a DESC) USING HASH WITH (bucket_count=8)",
 			pgExpected:  "CREATE INDEX baz ON foo.public.bar USING btree (a DESC) USING HASH WITH (bucket_count=8)",
 		},
+		{
+			index:       vectorIndex,
+			tableName:   descpb.AnonymousTable,
+			partition:   "",
+			displayMode: IndexDisplayDefOnly,
+			expected:    "VECTOR INDEX baz (a)",
+			pgExpected:  "INDEX baz USING cspann (a)",
+		},
+		{
+			index:       vectorIndex,
+			tableName:   tableName,
+			partition:   "",
+			displayMode: IndexDisplayShowCreate,
+			expected:    "CREATE VECTOR INDEX baz ON foo.public.bar (a)",
+			pgExpected:  "CREATE INDEX baz ON foo.public.bar USING cspann (a)",
+		},
 	}
 
 	sd := &sessiondata.SessionData{}

pkg/sql/catalog/colinfo/BUILD.bazel

@@ -16,6 +16,7 @@ go_library(
     importpath = "github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/docs",
         "//pkg/settings/cluster",
         "//pkg/sql/catalog",
         "//pkg/sql/catalog/catpb",
@@ -24,7 +25,9 @@ go_library(
         "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/sem/catconstants",
         "//pkg/sql/sem/eval",
+        "//pkg/sql/sem/idxtype",
         "//pkg/sql/sem/tree",
+        "//pkg/sql/sqlerrors",
         "//pkg/sql/types",
         "//pkg/util/encoding",
         "//pkg/util/errorutil/unimplemented",

pkg/sql/catalog/colinfo/col_type_info.go

@@ -9,10 +9,13 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/cockroachdb/cockroach/pkg/docs"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/idxtype"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/errorutil/unimplemented"
 	"github.com/cockroachdb/errors"
@@ -172,6 +175,12 @@ func ColumnTypeIsOnlyInvertedIndexable(t *types.T) bool {
 	return true
 }
 
+// ColumnTypeIsVectorIndexable returns true if the type t can be indexed using a
+// vector index.
+func ColumnTypeIsVectorIndexable(t *types.T) bool {
+	return t.Family() == types.PGVectorFamily
+}
+
 // MustBeValueEncoded returns true if columns of the given kind can only be value
 // encoded.
 func MustBeValueEncoded(semanticType *types.T) bool {
@@ -192,3 +201,72 @@ func MustBeValueEncoded(semanticType *types.T) bool {
 	}
 	return false
 }
+
+// ValidateColumnForIndex checks that the given column type is allowed in the
+// given index. If "isLastCol" is true, then it is the last column in the index.
+// "colDesc" can either be a column name or an expression in the form (a + b).
+func ValidateColumnForIndex(
+	indexType idxtype.T, colDesc string, colType *types.T, isLastCol bool,
+) error {
+	// Inverted and vector indexes only allow certain types for the last column.
+	if isLastCol {
+		switch indexType {
+		case idxtype.INVERTED:
+			if !ColumnTypeIsInvertedIndexable(colType) {
+				return sqlerrors.NewInvalidLastColumnError(colDesc, colType.Name(), indexType)
+			}
+			return nil
+
+		case idxtype.VECTOR:
+			if !ColumnTypeIsVectorIndexable(colType) {
+				return sqlerrors.NewInvalidLastColumnError(colDesc, colType.Name(), indexType)
+			} else if colType.Width() <= 0 {
+				return errors.WithDetail(
+					pgerror.Newf(
+						pgcode.FeatureNotSupported,
+						"%s column %s does not have a fixed number of dimensions, so it cannot be indexed",
+						colType.Name(), colDesc,
+					),
+					"specify the number of dimensions in the type, like VECTOR(128) for 128 dimensions",
+				)
+			}
+			return nil
+		}
+	}
+
+	if !ColumnTypeIsIndexable(colType) {
+		// If the column is indexable as the last column in the right kind of
+		// index, then use a more descriptive error message.
+		if ColumnTypeIsInvertedIndexable(colType) {
+			if indexType == idxtype.INVERTED {
+				// Column type is allowed, but only as the last column.
+				return sqlerrors.NewColumnOnlyIndexableError(colDesc, colType.Name(), idxtype.INVERTED)
+			}
+
+			// Column type is allowed in an inverted index.
+			return errors.WithHint(pgerror.Newf(
+				pgcode.FeatureNotSupported,
+				"column %s has type %s, which is not indexable in a non-inverted index",
+				colDesc, colType),
+				"you may want to create an inverted index instead. See the documentation for inverted indexes: "+docs.URL("inverted-indexes.html"))
+		}
+
+		if ColumnTypeIsVectorIndexable(colType) {
+			if indexType == idxtype.VECTOR {
+				// Column type is allowed, but only as the last column.
+				return sqlerrors.NewColumnOnlyIndexableError(colDesc, colType.Name(), idxtype.VECTOR)
+			}
+
+			// Column type is allowed in a vector index.
+			return errors.WithHint(pgerror.Newf(
+				pgcode.FeatureNotSupported,
+				"column %s has type %s, which is not indexable in a non-vector index",
+				colDesc, colType),
+				"you may want to create a vector index instead")
+		}
+
+		return sqlerrors.NewColumnNotIndexableError(colDesc, colType.Name(), colType.DebugString())
+	}
+
+	return nil
+}

pkg/sql/catalog/descpb/index.go

@@ -97,7 +97,7 @@ func (desc *IndexDescriptor) GetName() string {
 }
 
 // InvertedColumnID returns the ColumnID of the inverted column of the inverted
-// index. This is always the last column in ColumnIDs. Panics if the index is
+// index. This is always the last column in KeyColumnIDs. Panics if the index is
 // not inverted.
 func (desc *IndexDescriptor) InvertedColumnID() ColumnID {
 	if desc.Type != idxtype.INVERTED {
@@ -126,3 +126,23 @@ func (desc *IndexDescriptor) InvertedColumnKeyType() *types.T {
 	}
 	return types.EncodedKey
 }
+
+// VectorColumnID returns the ColumnID of the vector column of the vector index.
+// This is always the last column in KeyColumnIDs. Panics if the index is not a
+// vector index.
+func (desc *IndexDescriptor) VectorColumnID() ColumnID {
+	if desc.Type != idxtype.VECTOR {
+		panic(errors.AssertionFailedf("index is not a vector index"))
+	}
+	return desc.KeyColumnIDs[len(desc.KeyColumnIDs)-1]
+}
+
+// VectorColumnName returns the name of the vector column of the vector index.
+// This is always the last column in KeyColumnNames. Panics if the index is
+// not a vector index.
+func (desc *IndexDescriptor) VectorColumnName() string {
+	if desc.Type != idxtype.VECTOR {
+		panic(errors.AssertionFailedf("index is not a vector index"))
+	}
+	return desc.KeyColumnNames[len(desc.KeyColumnNames)-1]
+}