chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,374 @@
|
||||
// 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 expression
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/cockroachdb/errors"
|
||||
"github.com/dolthub/go-mysql-server/sql"
|
||||
"github.com/dolthub/go-mysql-server/sql/expression"
|
||||
"github.com/dolthub/go-mysql-server/sql/plan"
|
||||
|
||||
"github.com/dolthub/doltgresql/server/functions/framework"
|
||||
pgtypes "github.com/dolthub/doltgresql/server/types"
|
||||
)
|
||||
|
||||
// AnyExpr represents the ANY/SOME expression.
|
||||
type AnyExpr struct {
|
||||
leftExpr sql.Expression
|
||||
rightExpr sql.Expression
|
||||
subOperator string
|
||||
name string // ANY or SOME
|
||||
|
||||
subqueryAnyExpr *subqueryAnyExpr
|
||||
expressionAnyExpr *expressionAnyExpr
|
||||
}
|
||||
|
||||
// subqueryAnyExpr represents the resolved comparison functions for a plan.Subquery.
|
||||
type subqueryAnyExpr struct {
|
||||
rightSub *plan.Subquery
|
||||
staticLiteral *expression.Literal
|
||||
arrayLiterals []*expression.Literal
|
||||
compFuncs []framework.Function
|
||||
}
|
||||
|
||||
// expressionAnyExpr represents the resolved comparison function for a sql.Expression.
|
||||
type expressionAnyExpr struct {
|
||||
rightExpr sql.Expression
|
||||
staticLiteral *expression.Literal
|
||||
arrayLiteral *expression.Literal
|
||||
compFunc framework.Function
|
||||
}
|
||||
|
||||
// NewAnyExpr creates a new AnyExpr expression.
|
||||
func NewAnyExpr(subOperator string) *AnyExpr {
|
||||
return &AnyExpr{
|
||||
leftExpr: nil,
|
||||
rightExpr: nil,
|
||||
subOperator: subOperator,
|
||||
name: "ANY",
|
||||
}
|
||||
}
|
||||
|
||||
// Children implements the Expression interface.
|
||||
func (a *AnyExpr) Children() []sql.Expression {
|
||||
return []sql.Expression{a.leftExpr, a.rightExpr}
|
||||
}
|
||||
|
||||
// Resolved implements the Expression interface.
|
||||
func (a *AnyExpr) Resolved() bool {
|
||||
if a.leftExpr == nil || !a.leftExpr.Resolved() || a.rightExpr == nil || !a.rightExpr.Resolved() {
|
||||
return false
|
||||
}
|
||||
if a.subqueryAnyExpr != nil {
|
||||
return a.subqueryAnyExpr.resolved()
|
||||
}
|
||||
if a.expressionAnyExpr != nil {
|
||||
return a.expressionAnyExpr.resolved()
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// IsNullable implements the Expression interface.
|
||||
func (a *AnyExpr) IsNullable(ctx *sql.Context) bool {
|
||||
return a.leftExpr.IsNullable(ctx) || a.rightExpr.IsNullable(ctx)
|
||||
}
|
||||
|
||||
// Type implements the Expression interface.
|
||||
func (a *AnyExpr) Type(ctx *sql.Context) sql.Type {
|
||||
return pgtypes.Bool
|
||||
}
|
||||
|
||||
// resolved checks if the comparison functions for subqueryAnyExpr is resolved.
|
||||
func (a *subqueryAnyExpr) resolved() bool {
|
||||
if len(a.compFuncs) == 0 {
|
||||
return false
|
||||
}
|
||||
for _, compFunc := range a.compFuncs {
|
||||
if !compFunc.Resolved() {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// eval evaluates the comparison functions for subqueryAnyExpr.
|
||||
func (a *subqueryAnyExpr) eval(ctx *sql.Context, subOperator string, row sql.Row, left interface{}) (interface{}, error) {
|
||||
if len(a.compFuncs) == 0 {
|
||||
return nil, errors.Errorf("%T: cannot Eval as it has not been fully resolved", a)
|
||||
}
|
||||
|
||||
// TODO: This sometimes panics in `evalMultiple` for subqueries that return
|
||||
// more than one row, when len(row) > len(iter.Next())
|
||||
rightValues, err := a.rightSub.EvalMultiple(ctx, row)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if len(rightValues) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// TODO: This is a workaround some subqueries where the schema length does not
|
||||
// match the row length
|
||||
if len(a.arrayLiterals) == 1 && len(rightValues) != 1 {
|
||||
op, err := framework.GetOperatorFromString(subOperator)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
for i := len(a.arrayLiterals); i < len(rightValues); i++ {
|
||||
arrayLiteral := expression.NewLiteral(nil, a.arrayLiterals[0].Type(ctx))
|
||||
a.arrayLiterals = append(a.arrayLiterals, arrayLiteral)
|
||||
compFunc := framework.GetBinaryFunction(op).Compile(ctx, "internal_any_comparison", a.staticLiteral, a.arrayLiterals[i])
|
||||
a.compFuncs = append(a.compFuncs, compFunc)
|
||||
}
|
||||
}
|
||||
|
||||
if len(a.arrayLiterals) != len(rightValues) {
|
||||
return nil, errors.Errorf("%T: expected right child to return `%d` values but returned `%d`", a, len(a.arrayLiterals), len(rightValues))
|
||||
}
|
||||
|
||||
// Next we'll assign our evaluated values to the expressions that the comparison functions reference
|
||||
// Note that the compiled function has a reference to the staticLiteral and arrayLiterals, so we must alter them in place
|
||||
a.staticLiteral.Val = left
|
||||
for i, rightValue := range rightValues {
|
||||
a.arrayLiterals[i].Val = rightValue
|
||||
}
|
||||
// Now we can loop over all comparison functions, as they'll reference their respective values
|
||||
for _, compFunc := range a.compFuncs {
|
||||
result, err := compFunc.Eval(ctx, row)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if result.(bool) {
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// resolved checks if the comparison function for expressionAnyExpr is resolved.
|
||||
func (a *expressionAnyExpr) resolved() bool {
|
||||
if a.compFunc == nil || !a.compFunc.Resolved() {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// eval evaluates the comparison function for expressionAnyExpr.
|
||||
func (a *expressionAnyExpr) eval(ctx *sql.Context, row sql.Row, left interface{}) (interface{}, error) {
|
||||
if a.compFunc == nil {
|
||||
return nil, errors.Errorf("%T: cannot Eval as it has not been fully resolved", a)
|
||||
}
|
||||
|
||||
rightInterface, err := a.rightExpr.Eval(ctx, row)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if rightInterface == nil {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
rightValues, ok := rightInterface.([]any)
|
||||
if !ok {
|
||||
return nil, errors.Errorf("%T: expected right child to return `%T` but returned `%T`", a, []any{}, rightInterface)
|
||||
}
|
||||
if len(rightValues) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Next we'll assign our evaluated values to the expressions that the comparison function reference
|
||||
// Note that the compiled function has a reference to the staticLiteral and arrayLiteral, so we must alter them in place
|
||||
a.staticLiteral.Val = left
|
||||
for _, rightValue := range rightValues {
|
||||
a.arrayLiteral.Val = rightValue
|
||||
result, err := a.compFunc.Eval(ctx, row)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if result == nil {
|
||||
return nil, nil
|
||||
}
|
||||
if result.(bool) {
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// Eval implements the Expression interface.
|
||||
func (a *AnyExpr) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
|
||||
left, err := a.leftExpr.Eval(ctx, row)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if a.subqueryAnyExpr != nil {
|
||||
return a.subqueryAnyExpr.eval(ctx, a.subOperator, row, left)
|
||||
}
|
||||
|
||||
if a.expressionAnyExpr != nil {
|
||||
return a.expressionAnyExpr.eval(ctx, row, left)
|
||||
}
|
||||
|
||||
return nil, errors.Errorf("%T: cannot Eval as it has not been fully resolved", a)
|
||||
}
|
||||
|
||||
// WithChildren implements the Expression interface.
|
||||
func (a *AnyExpr) WithChildren(ctx *sql.Context, children ...sql.Expression) (sql.Expression, error) {
|
||||
if len(children) != 2 {
|
||||
return nil, sql.ErrInvalidChildrenNumber.New(a, len(children), 2)
|
||||
}
|
||||
|
||||
leftExpr := children[0]
|
||||
rightExpr := children[1]
|
||||
// Unmodified BindVars use deferred type resolution, so we replace the deference with the left's type in array form
|
||||
if bv, ok := rightExpr.(*expression.BindVar); ok {
|
||||
if _, ok = bv.Typ.(*pgtypes.DoltgresType); !ok {
|
||||
if leftType, ok := leftExpr.Type(ctx).(*pgtypes.DoltgresType); ok {
|
||||
bv.Typ = leftType.ToArrayType()
|
||||
}
|
||||
}
|
||||
}
|
||||
anyExpr := &AnyExpr{
|
||||
leftExpr: leftExpr,
|
||||
rightExpr: rightExpr,
|
||||
subOperator: a.subOperator,
|
||||
name: a.name,
|
||||
}
|
||||
|
||||
if sub, ok := children[1].(*plan.Subquery); ok {
|
||||
return anySubqueryWithChildren(ctx, anyExpr, sub)
|
||||
}
|
||||
|
||||
return anyExpressionWithChildren(ctx, anyExpr)
|
||||
}
|
||||
|
||||
// WithResolvedChildren implements the Expression interface.
|
||||
func (a *AnyExpr) WithResolvedChildren(ctx context.Context, children []any) (any, error) {
|
||||
if len(children) != 2 {
|
||||
return nil, errors.Errorf("invalid vitess child count, expected `2` but got `%d`", len(children))
|
||||
}
|
||||
left, ok := children[0].(sql.Expression)
|
||||
if !ok {
|
||||
return nil, errors.Errorf("expected vitess child to be an expression but has type `%T`", children[0])
|
||||
}
|
||||
right, ok := children[1].(sql.Expression)
|
||||
if !ok {
|
||||
return nil, errors.Errorf("expected vitess child to be an expression but has type `%T`", children[1])
|
||||
}
|
||||
return a.WithChildren(ctx.(*sql.Context), left, right)
|
||||
}
|
||||
|
||||
// String implements the fmt.Stringer interface.
|
||||
func (a *AnyExpr) String() string {
|
||||
if a.leftExpr == nil || a.rightExpr == nil {
|
||||
return fmt.Sprintf("? %s (?)", a.name)
|
||||
}
|
||||
return fmt.Sprintf("%s = %s (%s)", a.leftExpr, a.name, a.rightExpr)
|
||||
}
|
||||
|
||||
// DebugString implements the Expression interface.
|
||||
func (a *AnyExpr) DebugString(ctx *sql.Context) string {
|
||||
return fmt.Sprintf("%s %s (%s)", sql.DebugString(ctx, a.leftExpr), a.name, sql.DebugString(ctx, a.rightExpr))
|
||||
}
|
||||
|
||||
// anySubqueryWithChildren resolves the comparison functions for a plan.Subquery.
|
||||
func anySubqueryWithChildren(ctx *sql.Context, anyExpr *AnyExpr, sub *plan.Subquery) (sql.Expression, error) {
|
||||
schema := sub.Query.Schema(ctx)
|
||||
subTypes := make([]*pgtypes.DoltgresType, len(schema))
|
||||
for i, col := range schema {
|
||||
dgType, ok := col.Type.(*pgtypes.DoltgresType)
|
||||
if !ok {
|
||||
return nil, errors.Errorf("expected right child to be a DoltgresType but got `%T`", sub)
|
||||
}
|
||||
subTypes[i] = dgType
|
||||
}
|
||||
|
||||
op, err := framework.GetOperatorFromString(anyExpr.subOperator)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if leftType, ok := anyExpr.leftExpr.Type(ctx).(*pgtypes.DoltgresType); ok {
|
||||
// Resolve comparison functions once and reuse the functions in Eval.
|
||||
staticLiteral := expression.NewLiteral(nil, leftType)
|
||||
arrayLiterals := make([]*expression.Literal, len(subTypes))
|
||||
// Each expression may be a different type (which is valid), so we need a comparison function for each expression.
|
||||
compFuncs := make([]framework.Function, len(subTypes))
|
||||
for i, rightType := range subTypes {
|
||||
arrayLiterals[i] = expression.NewLiteral(nil, rightType)
|
||||
compFuncs[i] = framework.GetBinaryFunction(op).Compile(ctx, "internal_any_comparison", staticLiteral, arrayLiterals[i])
|
||||
if compFuncs[i] == nil {
|
||||
return nil, errors.Errorf("operator does not exist: %s = %s", leftType.String(), rightType.String())
|
||||
}
|
||||
if compFuncs[i].Type(ctx).(*pgtypes.DoltgresType).ID != pgtypes.Bool.ID {
|
||||
// This should never happen, but this is just to be safe
|
||||
return nil, errors.Errorf("%T: found equality comparison that does not return a bool", anyExpr)
|
||||
}
|
||||
}
|
||||
|
||||
anyExpr.subqueryAnyExpr = &subqueryAnyExpr{
|
||||
rightSub: sub,
|
||||
staticLiteral: staticLiteral,
|
||||
arrayLiterals: arrayLiterals,
|
||||
compFuncs: compFuncs,
|
||||
}
|
||||
}
|
||||
|
||||
return anyExpr, nil
|
||||
}
|
||||
|
||||
// anyExpressionWithChildren resolves the comparison functions for a sql.Expression.
|
||||
func anyExpressionWithChildren(ctx *sql.Context, anyExpr *AnyExpr) (sql.Expression, error) {
|
||||
arrType, ok := anyExpr.rightExpr.Type(ctx).(*pgtypes.DoltgresType)
|
||||
if !ok {
|
||||
return nil, errors.Errorf("expected right child to be a DoltgresType but got `%T`", anyExpr.rightExpr)
|
||||
}
|
||||
rightType := arrType.ArrayBaseType()
|
||||
op, err := framework.GetOperatorFromString(anyExpr.subOperator)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if leftType, ok := anyExpr.leftExpr.Type(ctx).(*pgtypes.DoltgresType); ok {
|
||||
// Resolve comparison function once and reuse the function in Eval.
|
||||
staticLiteral := expression.NewLiteral(nil, leftType)
|
||||
arrayLiteral := expression.NewLiteral(nil, rightType)
|
||||
compFunc := framework.GetBinaryFunction(op).Compile(ctx, "internal_any_comparison", staticLiteral, arrayLiteral)
|
||||
if compFunc == nil || compFunc.StashedError() != nil {
|
||||
return nil, errors.Errorf("operator does not exist: %s = %s", leftType.String(), rightType.String())
|
||||
}
|
||||
compFuncType := compFunc.Type(ctx)
|
||||
if compFuncType.(*pgtypes.DoltgresType).ID != pgtypes.Bool.ID {
|
||||
// This should never happen, but this is just to be safe
|
||||
return nil, errors.Errorf("%T: found equality comparison that does not return a bool", anyExpr)
|
||||
}
|
||||
anyExpr.expressionAnyExpr = &expressionAnyExpr{
|
||||
rightExpr: anyExpr.rightExpr,
|
||||
staticLiteral: staticLiteral,
|
||||
arrayLiteral: arrayLiteral,
|
||||
compFunc: compFunc,
|
||||
}
|
||||
}
|
||||
|
||||
return anyExpr, nil
|
||||
}
|
||||
Reference in New Issue
Block a user