Add support for Poltgres-compatible indexes

go.mod

@@ -8,10 +8,10 @@ require (
 	github.com/PuerkitoBio/goquery v1.8.1
 	github.com/cockroachdb/apd/v2 v2.0.3-0.20200518165714-d020e156310a
 	github.com/cockroachdb/errors v1.7.5
-	github.com/dolthub/dolt/go v0.40.5-0.20240815155659-80f6ce76a5a7
-	github.com/dolthub/dolt/go/gen/proto/dolt/services/eventsapi v0.0.0-20240529071237-4a099b896ce8
+	github.com/dolthub/dolt/go v0.40.5-0.20240827111219-e4bb9ca3442d
+	github.com/dolthub/dolt/go/gen/proto/dolt/services/eventsapi v0.0.0-20240827111219-e4bb9ca3442d
 	github.com/dolthub/flatbuffers/v23 v23.3.3-dh.2
-	github.com/dolthub/go-mysql-server v0.18.2-0.20240815142344-761713e36043
+	github.com/dolthub/go-mysql-server v0.18.2-0.20240827100900-3bf086dd5c18
 	github.com/dolthub/sqllogictest/go v0.0.0-20240618184124-ca47f9354216
 	github.com/dolthub/vitess v0.0.0-20240807181005-71d735078e24
 	github.com/fatih/color v1.13.0

go.sum

@@ -214,18 +214,18 @@ github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZm
 github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw=
 github.com/docker/go-connections v0.4.0/go.mod h1:Gbd7IOopHjR8Iph03tsViu4nIes5XhDvyHbTtUxmeec=
 github.com/docker/go-units v0.4.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
-github.com/dolthub/dolt/go v0.40.5-0.20240815155659-80f6ce76a5a7 h1:ktioGEKbWiZm01COmKFjariwrUVKJv376wiN9n2sNj8=
-github.com/dolthub/dolt/go v0.40.5-0.20240815155659-80f6ce76a5a7/go.mod h1:Dawp8iKBWNJZi6qXNhKwvvvDtjqmmyseD5VZwhOsQi8=
-github.com/dolthub/dolt/go/gen/proto/dolt/services/eventsapi v0.0.0-20240529071237-4a099b896ce8 h1:izuogF6KRc6Pr5g5KevRtn8JK/KwyEGjbpqWJIORbQo=
-github.com/dolthub/dolt/go/gen/proto/dolt/services/eventsapi v0.0.0-20240529071237-4a099b896ce8/go.mod h1:L5RDYZbC9BBWmoU2+TjTekeqqhFXX5EqH9ln00O0stY=
+github.com/dolthub/dolt/go v0.40.5-0.20240827111219-e4bb9ca3442d h1:LNssHEPGHIbeyJ0nxS4jerKRZTXYF6msLxAAE4iwqRE=
+github.com/dolthub/dolt/go v0.40.5-0.20240827111219-e4bb9ca3442d/go.mod h1:Hg+BBTRR74OcOXP+WrVcHyyzRD+WtnLHtMcy1wtvS2E=
+github.com/dolthub/dolt/go/gen/proto/dolt/services/eventsapi v0.0.0-20240827111219-e4bb9ca3442d h1:RZkQeYOrDrOWzCxaP2ttkvg4E2TM9n8lnEsIBLKjqkM=
+github.com/dolthub/dolt/go/gen/proto/dolt/services/eventsapi v0.0.0-20240827111219-e4bb9ca3442d/go.mod h1:L5RDYZbC9BBWmoU2+TjTekeqqhFXX5EqH9ln00O0stY=
 github.com/dolthub/flatbuffers/v23 v23.3.3-dh.2 h1:u3PMzfF8RkKd3lB9pZ2bfn0qEG+1Gms9599cr0REMww=
 github.com/dolthub/flatbuffers/v23 v23.3.3-dh.2/go.mod h1:mIEZOHnFx4ZMQeawhw9rhsj+0zwQj7adVsnBX7t+eKY=
 github.com/dolthub/fslock v0.0.3 h1:iLMpUIvJKMKm92+N1fmHVdxJP5NdyDK5bK7z7Ba2s2U=
 github.com/dolthub/fslock v0.0.3/go.mod h1:QWql+P17oAAMLnL4HGB5tiovtDuAjdDTPbuqx7bYfa0=
 github.com/dolthub/go-icu-regex v0.0.0-20230524105445-af7e7991c97e h1:kPsT4a47cw1+y/N5SSCkma7FhAPw7KeGmD6c9PBZW9Y=
 github.com/dolthub/go-icu-regex v0.0.0-20230524105445-af7e7991c97e/go.mod h1:KPUcpx070QOfJK1gNe0zx4pA5sicIK1GMikIGLKC168=
-github.com/dolthub/go-mysql-server v0.18.2-0.20240815142344-761713e36043 h1:KgrDVE4o4Y04XLnAs5BGv6I6z+Rd82FWntCbQEmbTKs=
-github.com/dolthub/go-mysql-server v0.18.2-0.20240815142344-761713e36043/go.mod h1:nbdOzd0ceWONE80vbfwoRBjut7z3CIj69ZgDF/cKuaA=
+github.com/dolthub/go-mysql-server v0.18.2-0.20240827100900-3bf086dd5c18 h1:1lgwZvnecrjoc9v0iqxjdKBvaasAPiQzty40uTKOHsE=
+github.com/dolthub/go-mysql-server v0.18.2-0.20240827100900-3bf086dd5c18/go.mod h1:nbdOzd0ceWONE80vbfwoRBjut7z3CIj69ZgDF/cKuaA=
 github.com/dolthub/gozstd v0.0.0-20240423170813-23a2903bca63 h1:OAsXLAPL4du6tfbBgK0xXHZkOlos63RdKYS3Sgw/dfI=
 github.com/dolthub/gozstd v0.0.0-20240423170813-23a2903bca63/go.mod h1:lV7lUeuDhH5thVGDCKXbatwKy2KW80L4rMT46n+Y2/Q=
 github.com/dolthub/ishell v0.0.0-20240701202509-2b217167d718 h1:lT7hE5k+0nkBdj/1UOSFwjWpNxf+LCApbRHgnCA17XE=

postgres/parser/sem/tree/name_part.go

@@ -231,3 +231,35 @@ func (u *UnresolvedName) ToUnresolvedObjectName(idx AnnotationIdx) (*UnresolvedO
 		idx,
 	)
 }
+
+// GetUnresolvedObjectName pulls the UnresolvedObjectName from the UnresolvedName. It is possible that the
+// UnresolvedName does not reference an upper object, therefore the UnresolvedObjectName may be empty.
+func (u *UnresolvedName) GetUnresolvedObjectName() *UnresolvedObjectName {
+	switch u.NumParts {
+	case 1:
+		return &UnresolvedObjectName{
+			NumParts: 0,
+			Parts:    [3]string{},
+		}
+	case 2:
+		return &UnresolvedObjectName{
+			NumParts: 1,
+			Parts:    [3]string{u.Parts[1], "", ""},
+		}
+	case 3:
+		return &UnresolvedObjectName{
+			NumParts: 2,
+			Parts:    [3]string{u.Parts[1], u.Parts[2], ""},
+		}
+	case 4:
+		return &UnresolvedObjectName{
+			NumParts: 3,
+			Parts:    [3]string{u.Parts[1], u.Parts[2], u.Parts[3]},
+		}
+	default:
+		return &UnresolvedObjectName{
+			NumParts: 0,
+			Parts:    [3]string{},
+		}
+	}
+}

server/analyzer/init.go

@@ -24,6 +24,7 @@ const (
 	ruleId_TypeSanitizer analyzer.RuleId = iota + 1000
 	ruleId_AssignInsertCasts
 	ruleId_AssignUpdateCasts
+	ruleId_ReplaceIndexedTables
 	ruleId_ReplaceSerial
 	ruleId_InsertContextRootFinalizer
 )
@@ -36,6 +37,7 @@ func Init() {
 		getAnalyzerRule(analyzer.OnceBeforeDefault, analyzer.ValidateColumnDefaultsId),
 		analyzer.Rule{Id: ruleId_AssignInsertCasts, Apply: AssignInsertCasts},
 		analyzer.Rule{Id: ruleId_AssignUpdateCasts, Apply: AssignUpdateCasts},
+		analyzer.Rule{Id: ruleId_ReplaceIndexedTables, Apply: ReplaceIndexedTables},
 	)
 
 	// Column default validation was moved to occur after type sanitization, so we'll remove it from its original place
@@ -52,9 +54,6 @@ func Init() {
 	// The auto-commit rule writes the contents of the context, so we need to insert our finalizer before that
 	analyzer.OnceAfterAll = insertAnalyzerRules(analyzer.OnceAfterAll, analyzer.AutocommitId, true,
 		analyzer.Rule{Id: ruleId_InsertContextRootFinalizer, Apply: InsertContextRootFinalizer})
-
-	// Handle the function overrides
-	analyzer.IndexLeafChildren = IndexLeafChildren
 }
 
 // getAnalyzerRule returns the rule matching the given ID.

server/analyzer/replace_indexed_tables.go

@@ -0,0 +1,67 @@
+// Copyright 2024 Dolthub, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package analyzer
+
+import (
+	"github.com/dolthub/dolt/go/libraries/doltcore/sqle"
+	"github.com/dolthub/go-mysql-server/sql"
+	"github.com/dolthub/go-mysql-server/sql/analyzer"
+	"github.com/dolthub/go-mysql-server/sql/plan"
+	"github.com/dolthub/go-mysql-server/sql/transform"
+
+	"github.com/dolthub/doltgresql/server/index"
+)
+
+// ReplaceIndexedTables replaces Dolt tables with Doltgres tables that can properly handle indexed access.
+func ReplaceIndexedTables(ctx *sql.Context, a *analyzer.Analyzer, node sql.Node, scope *plan.Scope, selector analyzer.RuleSelector, qFlags *sql.QueryFlags) (sql.Node, transform.TreeIdentity, error) {
+	return transform.Node(node, func(n sql.Node) (sql.Node, transform.TreeIdentity, error) {
+		if filter, ok := n.(*plan.Filter); ok {
+			return transform.Node(filter, replaceIndexedTablesFilter)
+		}
+		return n, transform.SameTree, nil
+	})
+}
+
+// replaceIndexedTablesFilter is the transform function for ReplaceIndexedTables that handles the resolved table.
+func replaceIndexedTablesFilter(n sql.Node) (sql.Node, transform.TreeIdentity, error) {
+	switch n := n.(type) {
+	case *plan.ResolvedTable:
+		if newTable, ok, err := doltTableToIndexedTable(n.UnderlyingTable()); err != nil {
+			return n, transform.SameTree, err
+		} else if ok {
+			nt, err := n.WithTable(newTable)
+			return nt, transform.NewTree, err
+		}
+		return n, transform.SameTree, nil
+	default:
+		return n, transform.SameTree, nil
+	}
+}
+
+// doltTableToIndexedTable replaces Dolt tables with Doltgres' indexed tables.
+func doltTableToIndexedTable(table sql.Table) (sql.Table, bool, error) {
+	switch table := table.(type) {
+	case *sqle.AlterableDoltTable:
+		return &index.WritableDoltgresTable{WritableDoltTable: &table.WritableDoltTable}, true, nil
+	case *sqle.WritableDoltTable:
+		return &index.WritableDoltgresTable{WritableDoltTable: table}, true, nil
+	case *sqle.DoltTable:
+		return &index.DoltgresTable{DoltTable: table}, true, nil
+	case *index.DoltgresTable:
+		return table, false, nil
+	default:
+		return table, false, nil
+	}
+}

server/ast/expr.go

@@ -327,7 +327,7 @@ func nodeExpr(node tree.Expr) (vitess.Expr, error) {
 			switch right := right.(type) {
 			case vitess.ValTuple:
 				return vitess.InjectedExpr{
-					Expression: pgexprs.NewInTuple(),
+					Expression: pgexprs.NewInTupleDecaying(),
 					Children:   vitess.Exprs{left, right},
 				}, nil
 			case *vitess.Subquery:

server/ast/select_clause.go

@@ -32,6 +32,87 @@ func nodeSelectClause(node *tree.SelectClause) (*vitess.Select, error) {
 	if err != nil {
 		return nil, err
 	}
+	// Multiple tables in the FROM column with an "equals" filter for some columns within each table should be treated
+	// as a join. The analyzer should catch this, however GMS processes this form of a join differently than a standard
+	// join, which is currently incompatible with Doltgres expressions. As a workaround, we rewrite the tree so that we
+	// pass along a join node.
+	// TODO: handle more than two tables, also make this more robust with handling more node types
+	if len(node.From.Tables) == 2 && node.Where != nil {
+		// We don't yet handle AS OF for rewrites
+		if node.From.AsOf.Expr != nil {
+			goto PostJoinRewrite
+		}
+		tableNames := make(map[tree.TableName]int)
+		tableAliases := make(map[tree.TableName]int)
+		// First we need to get the table names and aliases, since they'll be referenced by the filters
+		for i := range node.From.Tables {
+			switch table := node.From.Tables[i].(type) {
+			case *tree.AliasedTableExpr:
+				if tableName, ok := table.Expr.(*tree.TableName); ok {
+					tableNames[*tableName] = i
+				} else {
+					goto PostJoinRewrite
+				}
+				tableAliases[tree.MakeUnqualifiedTableName(table.As.Alias)] = i
+			case *tree.TableName:
+				tableNames[*table] = i
+			case *tree.UnresolvedObjectName:
+				tableNames[table.ToTableName()] = i
+			default:
+				goto PostJoinRewrite
+			}
+		}
+		// For now, we'll check if the entire filter should be moved into the join condition. Eventually, this should
+		// move only the needed expressions into the join condition.
+		var delveExprs func(expr tree.Expr) bool
+		delveExprs = func(expr tree.Expr) bool {
+			switch expr := expr.(type) {
+			case *tree.AndExpr:
+				return delveExprs(expr.Left) && delveExprs(expr.Right)
+			case *tree.OrExpr:
+				return delveExprs(expr.Left) && delveExprs(expr.Right)
+			case *tree.ComparisonExpr:
+				if expr.Operator != tree.EQ {
+					return false
+				}
+				var refTables [2]int
+				for argIndex, arg := range []tree.Expr{expr.Left, expr.Right} {
+					switch arg := arg.(type) {
+					case *tree.UnresolvedName:
+						refTable := arg.GetUnresolvedObjectName().ToTableName()
+						if aliasIndex, ok := tableAliases[refTable]; ok {
+							refTables[argIndex] = aliasIndex
+						} else if tableIndex, ok := tableNames[refTable]; ok {
+							refTables[argIndex] = tableIndex
+						} else {
+							return false
+						}
+					default:
+						return false
+					}
+				}
+				// In this case, the expression does not reference multiple tables, so it's not a join condition
+				if refTables[0] == refTables[1] {
+					return false
+				}
+				return true
+			default:
+				return false
+			}
+		}
+		if !delveExprs(node.Where.Expr) {
+			goto PostJoinRewrite
+		}
+		// The filter condition represents a join, so we need to rewrite our FROM node to be a join node
+		node.From.Tables = tree.TableExprs{&tree.JoinTableExpr{
+			JoinType: "",
+			Left:     node.From.Tables[0],
+			Right:    node.From.Tables[1],
+			Cond:     &tree.OnJoinCond{Expr: node.Where.Expr},
+		}}
+		node.Where = nil
+	}
+PostJoinRewrite:
 	from, err := nodeFrom(node.From)
 	if err != nil {
 		return nil, err

server/expression/array.go

@@ -77,7 +77,16 @@ func (array *Array) Eval(ctx *sql.Context, row sql.Row) (any, error) {
 		castFunc := framework.GetImplicitCast(doltgresType.BaseID(), resultTyp.BaseID())
 		if castFunc == nil {
 			if doltgresType.BaseID() == pgtypes.DoltgresTypeBaseID_Unknown {
-				castFunc = framework.CastFromUnknownType
+				castFunc = func(ctx *sql.Context, val any, targetType pgtypes.DoltgresType) (any, error) {
+					if val == nil {
+						return nil, nil
+					}
+					str, err := pgtypes.Unknown.IoOutput(ctx, val)
+					if err != nil {
+						return nil, err
+					}
+					return targetType.IoInput(ctx, str)
+				}
 			} else {
 				return nil, fmt.Errorf("cannot find cast function from %s to %s", doltgresType.String(), resultTyp.String())
 			}

server/expression/assignment_cast.go

@@ -55,7 +55,16 @@ func (ac *AssignmentCast) Eval(ctx *sql.Context, row sql.Row) (any, error) {
 	castFunc := framework.GetAssignmentCast(ac.fromType.BaseID(), ac.toType.BaseID())
 	if castFunc == nil {
 		if ac.fromType.BaseID() == pgtypes.DoltgresTypeBaseID_Unknown {
-			castFunc = framework.CastFromUnknownType
+			castFunc = func(ctx *sql.Context, val any, targetType pgtypes.DoltgresType) (any, error) {
+				if val == nil {
+					return nil, nil
+				}
+				str, err := pgtypes.Unknown.IoOutput(ctx, val)
+				if err != nil {
+					return nil, err
+				}
+				return targetType.IoInput(ctx, str)
+			}
 		} else {
 			return nil, fmt.Errorf("ASSIGNMENT_CAST: target is of type %s but expression is of type %s: %s",
 				ac.toType.String(), ac.fromType.String(), ac.expr.String())