Merge branch 'main' into james/implicit

Author: James Cor

enginetest/queries/index_query_plans.go

@@ -7620,6 +7620,20 @@ where
 					{3, "c", float64(4)},
 				},
 			},
+			{
+				Query: "alter table t modify column e enum('asdf', 'a', 'b', 'c', 'd');",
+				Expected: []sql.Row{
+					{types.NewOkResult(0)},
+				},
+			},
+			{
+				Query: "select i, e, e + 0 from t;",
+				Expected: []sql.Row{
+					{1, "a", float64(2)},
+					{2, "b", float64(3)},
+					{3, "c", float64(4)},
+				},
+			},
 			{
 				Query:       "alter table t modify column e enum('abc');",
 				ExpectedErr: types.ErrConvertingToEnum,

enginetest/queries/script_queries.go

@@ -56,6 +56,17 @@ var VectorIndexQueries = []ScriptTest{
 				},
 				ExpectedIndexes: []string{"v_idx"},
 			},
+			{
+				// Use the index even when there's a projection involved.
+				Query: "select `id`+1 from vectors order by VEC_DISTANCE('[0.0,0.0]', v) limit 4",
+				Expected: []sql.Row{
+					{3},
+					{4},
+					{5},
+					{2},
+				},
+				ExpectedIndexes: []string{"v_idx"},
+			},
 			{
 				// Only queries with a limit can use a vector index.
 				Query: "select * from vectors order by VEC_DISTANCE('[0.0,0.0]', v)",
@@ -88,6 +99,21 @@ var VectorIndexQueries = []ScriptTest{
 				},
 				ExpectedIndexes: []string{},
 			},
+			{
+				// Modify the index after creation.
+				Query: "insert into vectors values (5, '[1.0,0.0]')",
+			},
+			{
+				Query: "select * from vectors order by VEC_DISTANCE('[0.0,0.0]', v)",
+				Expected: []sql.Row{
+					{2, types.MustJSON(`[0.0, 0.0]`)},
+					{5, types.MustJSON(`[1.0, 0.0]`)},
+					{3, types.MustJSON(`[-1.0, 1.0]`)},
+					{4, types.MustJSON(`[0.0, -2.0]`)},
+					{1, types.MustJSON(`[4.0, 3.0]`)},
+				},
+				ExpectedIndexes: []string{},
+			},
 		},
 	},
 }

enginetest/queries/vector_index_queries.go

@@ -95,4 +95,28 @@ create table pref_index_t3 (v1 varchar(10), v2 varchar(10), unique index (v1(3),
 
 exec
 create table pref_index_t4 (i int primary key, v1 varchar(10), v2 varchar(10), unique index (v1(3),v2(5)));
+----
+
+exec
+CREATE TABLE comp_vector_index_t0 (pk BIGINT PRIMARY KEY, v1 BIGINT, v2 JSON);
+----
+
+exec
+INSERT INTO comp_vector_index_t0 VALUES (0,0,"[3,16]"),(1,2,"[65,9]"),(2,3,"[38,37]"),(3,3,"[99,99]"),(4,5,"[17,42]"),(5,6,"[6,76]"),(6,6,"[81,33]"),
+(7,7,"[33,51]"),(8,7,"[37,42]"),(9,8,"[9,21]"),(10,8,"[37,90]"),(11,9,"[39,20]"),(12,9,"[71,82]"),(13,10,"[16,21]"),(14,10,"[32,46]"),(15,10,"[47,36]"),
+(16,12,"[44,84]"),(17,12,"[66,40]"),(18,13,"[47,30]"),(19,13,"[56,41]"),(20,14,"[38,24]"),(21,14,"[91,1]"),(22,15,"[2,69]"),(23,16,"[40,36]"),
+(24,20,"[29,93]"),(25,21,"[9,89]"),(26,21,"[42,76]"),(27,23,"[13,53]"),(28,23,"[28,68]"),(29,23,"[28,90]"),(30,23,"[30,44]"),(31,24,"[20,8]"),
+(32,25,"[49,88]"),(33,26,"[15,28]"),(34,27,"[35,12]"),(35,28,"[39,84]"),(36,29,"[7,38]"),(37,29,"[21,74]"),(38,29,"[27,48]"),(39,29,"[77,46]"),
+(40,31,"[47,21]"),(41,31,"[47,91]"),(42,32,"[40,76]"),(43,33,"[70,50]"),(44,34,"[27,58]"),(45,35,"[32,36]"),(46,36,"[4,36]"),(47,36,"[84,75]"),
+(48,37,"[27,32]"),(49,38,"[88,68]"),(50,41,"[17,68]"),(51,41,"[77,26]"),(52,42,"[80,85]"),(53,45,"[1,57]"),(54,46,"[58,8]"),(55,49,"[26,11]"),
+(56,50,"[49,20]"),(57,50,"[86,6]"),(58,54,"[13,78]"),(59,54,"[57,83]"),(60,55,"[45,46]"),(61,55,"[81,80]"),(62,56,"[0,97]"),(63,56,"[8,78]"),
+(64,56,"[58,4]"),(65,56,"[66,33]"),(66,57,"[7,52]"),(67,59,"[77,53]"),(68,60,"[8,70]"),(69,61,"[11,25]"),(70,63,"[85,23]"),(71,65,"[17,9]"),
+(72,66,"[46,46]"),(73,66,"[73,4]"),(74,67,"[55,27]"),(75,70,"[8,54]"),(76,70,"[58,33]"),(77,71,"[39,15]"),(78,72,"[65,64]"),(79,74,"[78,26]"),
+(80,75,"[91,35]"),(81,76,"[40,52]"),(82,76,"[44,87]"),(83,81,"[32,4]"),(84,82,"[11,6]"),(85,82,"[46,32]"),(86,84,"[40,8]"),(87,84,"[93,37]"),
+(88,85,"[53,50]"),(89,86,"[63,79]"),(90,87,"[22,34]"),(91,87,"[57,62]"),(92,88,"[88,42]"),(93,90,"[30,67]"),(94,91,"[15,15]"),(95,93,"[7,26]"),
+(96,94,"[92,38]"),(97,95,"[89,66]"),(98,97,"[63,19]"),(99,98,"[31,21]"),(100,98,"[42,22]")
+----
+
+exec
+create VECTOR INDEX v_idx on comp_vector_index_t0 (v2)
 ----

enginetest/scriptgen/setup/scripts/comp_index_tables

@@ -550,6 +550,9 @@ func GetDeferredProjections(iter sql.RowIter) (sql.RowIter, []sql.Expression) {
 // resultForMax1RowIter ensures that an empty iterator returns at most one row
 func resultForMax1RowIter(ctx *sql.Context, schema sql.Schema, iter sql.RowIter, resultFields []*querypb.Field, buf *sql.ByteBuffer) (*sqltypes.Result, error) {
 	defer trace.StartRegion(ctx, "Handler.resultForMax1RowIter").End()
+
+	defer iter.Close(ctx)
+
 	row, err := iter.Next(ctx)
 	if err == io.EOF {
 		return &sqltypes.Result{Fields: resultFields}, nil
@@ -560,9 +563,7 @@ func resultForMax1RowIter(ctx *sql.Context, schema sql.Schema, iter sql.RowIter,
 	if _, err = iter.Next(ctx); err != io.EOF {
 		return nil, fmt.Errorf("result max1Row iterator returned more than one row")
 	}
-	if err := iter.Close(ctx); err != nil {
-		return nil, err
-	}
+
 	outputRow, err := RowToSQL(ctx, schema, row, nil, buf)
 	if err != nil {
 		return nil, err

enginetest/scriptgen/setup/setup_data.sg.go

@@ -10,15 +10,20 @@ import (
 
 // replaceIdxSort applies an IndexAccess when there is an `OrderBy` over a prefix of any columns with Indexes
 func replaceIdxOrderByDistance(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.Scope, sel RuleSelector, qFlags *sql.QueryFlags) (sql.Node, transform.TreeIdentity, error) {
-	return replaceIdxOrderByDistanceHelper(ctx, scope, n, nil)
+	return replaceIdxOrderByDistanceHelper(ctx, scope, n, nil, nil)
 }
 
-func replaceIdxOrderByDistanceHelper(ctx *sql.Context, scope *plan.Scope, node sql.Node, sortNode *plan.TopN) (sql.Node, transform.TreeIdentity, error) {
+func replaceIdxOrderByDistanceHelper(ctx *sql.Context, scope *plan.Scope, node sql.Node, sortNode plan.Sortable, limit sql.Expression) (sql.Node, transform.TreeIdentity, error) {
 	switch n := node.(type) {
 	case *plan.TopN:
 		sortNode = n // lowest parent sort node
+		limit = n.Limit
+	case plan.Sortable:
+		sortNode = n
+	case *plan.Limit:
+		limit = n.Limit
 	case *plan.ResolvedTable:
-		if sortNode == nil {
+		if sortNode == nil || limit == nil {
 			return n, transform.SameTree, nil
 		}
 
@@ -44,7 +49,7 @@ func replaceIdxOrderByDistanceHelper(ctx *sql.Context, scope *plan.Scope, node s
 
 		// Column references have not been assigned their final indexes yet, so do that for the ORDER BY expression now.
 		// We can safely do this because an expression that references other tables won't pass `isSortFieldsValidPrefix` below.
-		sortNode = offsetAssignIndexes(sortNode).(*plan.TopN)
+		sortNode = offsetAssignIndexes(sortNode).(plan.Sortable)
 
 		sfExprs := normalizeExpressions(tableAliases, sortNode.GetSortFields().ToExpressions()...)
 		sfAliases := aliasedExpressionsInNode(sortNode)
@@ -100,8 +105,6 @@ func replaceIdxOrderByDistanceHelper(ctx *sql.Context, scope *plan.Scope, node s
 			return n, transform.SameTree, nil
 		}
 
-		limit := sortNode.Limit
-
 		lookup := sql.IndexLookup{
 			Index:  idx,
 			Ranges: sql.MySQLRangeCollection{},
@@ -125,7 +128,7 @@ func replaceIdxOrderByDistanceHelper(ctx *sql.Context, scope *plan.Scope, node s
 		same := transform.SameTree
 		switch c := child.(type) {
 		case *plan.Project, *plan.TableAlias, *plan.ResolvedTable, *plan.Filter, *plan.Limit, *plan.TopN, *plan.Offset, *plan.Sort, *plan.IndexedTableAccess:
-			newChildren[i], same, err = replaceIdxOrderByDistanceHelper(ctx, scope, child, sortNode)
+			newChildren[i], same, err = replaceIdxOrderByDistanceHelper(ctx, scope, child, sortNode, limit)
 		default:
 			newChildren[i] = c
 		}

server/handler.go

@@ -114,7 +114,7 @@ func TestVectorIndex(t *testing.T) {
 
 	for _, testCase := range vectorIndexTestCases(t, db, vectorIndexTable) {
 		t.Run(testCase.name, func(t *testing.T) {
-			res, same, err := replaceIdxOrderByDistanceHelper(nil, nil, testCase.inputPlan, nil)
+			res, same, err := replaceIdxOrderByDistanceHelper(nil, nil, testCase.inputPlan, nil, nil)
 			require.NoError(t, err)
 			require.Equal(t, testCase.usesVectorIndex, !bool(same))
 			res = offsetAssignIndexes(res)
@@ -218,7 +218,6 @@ func (i vectorIndexTable) SkipIndexCosting() bool {
 }
 
 func (i vectorIndexTable) IndexWithPrefix(ctx *sql.Context, expressions []string) (sql.Index, error) {
-	//TODO implement me
 	panic("implement me")
 }
 

sql/analyzer/replace_order_by_distance.go

@@ -525,7 +525,7 @@ func (i *modifyColumnIter) rewriteTable(ctx *sql.Context, rwt sql.RewritableTabl
 
 	oldEnum, isOldEnum := oldCol.Type.(sql.EnumType)
 	newEnum, isNewEnum := newCol.Type.(sql.EnumType)
-	if isOldEnum && isNewEnum && !oldEnum.Equals(newEnum) {
+	if isOldEnum && isNewEnum && !oldEnum.IsSubsetOf(newEnum) {
 		rewriteRequired = true
 	}
 

sql/analyzer/vector_index_test.go

@@ -186,6 +186,9 @@ type EnumType interface {
 	Collation() CollationID
 	// IndexOf returns the index of the given string. If the string was not found, then this returns -1.
 	IndexOf(v string) int
+	// IsSubsetOf returns whether every element in this is also in |otherType|, with the same indexes.
+	// |otherType| may contain additional elements not in this.
+	IsSubsetOf(otherType EnumType) bool
 	// NumberOfElements returns the number of enumerations.
 	NumberOfElements() uint16
 	// Values returns the elements, in order, of every enumeration.