dolthub--doltgresql/server/expression/in_subquery.go

// 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/hash"
	"github.com/dolthub/go-mysql-server/sql/plan"
	"github.com/dolthub/go-mysql-server/sql/types"
	vitess "github.com/dolthub/vitess/go/vt/sqlparser"

	"github.com/dolthub/doltgresql/server/functions/framework"
	pgtypes "github.com/dolthub/doltgresql/server/types"
)

// InSubquery represents a VALUE IN (SELECT ...) expression.
type InSubquery struct {
	leftExpr  sql.Expression
	rightExpr *plan.Subquery

	// These variables are used so that we can resolve the comparison functions once and reuse them as we iterate over rows.
	// These are assigned in WithChildren, so refer there for more information.
	leftLiteral   *expression.Literal
	rightLiterals []*expression.Literal
	compFuncs     []framework.Function
}

var _ vitess.Injectable = (*InSubquery)(nil)
var _ sql.Expression = (*InSubquery)(nil)
var _ expression.BinaryExpression = (*InSubquery)(nil)

// nilKey is the hash of a row with a single nil value.
var nilKey, _ = hash.HashOf(nil, nil, sql.NewRow(nil))

// NewInSubquery returns a new *InSubquery.
func NewInSubquery() *InSubquery {
	return &InSubquery{}
}

// Children implements the sql.Expression interface.
func (in *InSubquery) Children() []sql.Expression {
	return []sql.Expression{in.leftExpr, in.rightExpr}
}

// Eval implements the sql.Expression interface.
func (in *InSubquery) Eval(ctx *sql.Context, row sql.Row) (any, error) {
	if len(in.compFuncs) == 0 {
		return nil, errors.Errorf("%T: cannot Eval as it has not been fully resolved", in)
	}

	left, err := in.leftExpr.Eval(ctx, row)
	if err != nil {
		return nil, err
	}

	// The NULL handling for IN expressions is tricky. According to
	// https://www.postgresql.org/docs/16/functions-comparisons.html#FUNCTIONS-COMPARISONS-IN-SCALAR:
	// To comply with the SQL standard, IN() returns NULL not only if the expression on the left hand side is NULL, but
	// also if no match is found in the list and one of the expressions in the list is NULL.
	leftNull := left == nil

	if types.NumColumns(in.Left().Type(ctx)) != types.NumColumns(in.Right().Type(ctx)) {
		return nil, sql.ErrInvalidOperandColumns.New(types.NumColumns(in.Left().Type(ctx)), types.NumColumns(in.Right().Type(ctx)))
	}

	right := in.rightExpr

	// TODO: does this work for all pg values?
	values, err := right.HashMultiple(ctx, row)
	if err != nil {
		return nil, err
	}

	// NULL IN (list) returns NULL. NULL IN (empty list) returns 0
	if leftNull {
		if values.Size() == 0 {
			return false, nil
		}
		return nil, nil
	}

	// TODO: it might be possible for the left value to hash to a different value than the right even though they pass
	//  an equality check. We need to perform a type conversion here to catch this case.
	key, err := hash.HashOf(ctx, nil, sql.NewRow(left))
	if err != nil {
		return nil, err
	}

	// If the hashed values don't contain the left value hash, we know it's not there.
	// If we do find the hash of the left value, we still need to check for equality,
	// since non-equal values could have the same hash in some cases.
	val, notFoundErr := values.Get(key)
	if notFoundErr != nil {
		if _, nilValNotFoundErr := values.Get(nilKey); nilValNotFoundErr == nil {
			return nil, nil
		}
		return false, nil
	}

	var r sql.Row
	rowVal, ok := val.([]any)
	if !ok {
		r = sql.Row{val}
	} else {
		r = sql.NewRow(rowVal...)
	}

	return in.valuesEqual(ctx, left, r)
}

// valuesEqual returns true if the left value is equal to the row provided using the equality functions previously
// assigned to |compFuncs| during analysis. If the left value is a single scalar, then |row| has a single value as
// well. Otherwise, (left is a tuple), |row| has a matching number of values.
func (in *InSubquery) valuesEqual(ctx *sql.Context, left interface{}, row sql.Row) (bool, error) {
	// Note that we have to edit the literals in place, since the comparison functions reference them directly.
	in.leftLiteral.Val = left
	for i, v := range row {
		in.rightLiterals[i].Val = v
	}

	for _, compFunc := range in.compFuncs {
		result, err := compFunc.Eval(ctx, nil)
		if err != nil {
			return false, err
		}
		if !result.(bool) {
			return false, nil
		}
	}
	return true, nil
}

// IsNullable implements the sql.Expression interface.
func (in *InSubquery) IsNullable(ctx *sql.Context) bool {
	return true
}

// Resolved implements the sql.Expression interface.
func (in *InSubquery) Resolved() bool {
	if in.leftExpr == nil || !in.leftExpr.Resolved() || in.rightExpr == nil || !in.rightExpr.Resolved() || len(in.compFuncs) == 0 {
		return false
	}
	for _, compFunc := range in.compFuncs {
		if !compFunc.Resolved() {
			return false
		}
	}
	return true
}

// String implements the sql.Expression interface.
func (in *InSubquery) String() string {
	if in.leftExpr == nil || in.rightExpr == nil {
		return "? IN ?"
	}
	return fmt.Sprintf("%s IN %s", in.leftExpr.String(), in.rightExpr.String())
}

// Type implements the sql.Expression interface.
func (in *InSubquery) Type(ctx *sql.Context) sql.Type {
	return pgtypes.Bool
}

// WithChildren implements the sql.Expression interface.
func (in *InSubquery) WithChildren(ctx *sql.Context, children ...sql.Expression) (sql.Expression, error) {
	if len(children) != 2 {
		return nil, sql.ErrInvalidChildrenNumber.New(in, len(children), 2)
	}
	sq, ok := children[1].(*plan.Subquery)
	if !ok {
		return nil, errors.Errorf("%T: expected right child to be `%T` but has type `%T`", in, &plan.Subquery{}, children[1])
	}
	// We'll only resolve the comparison functions once we have all Doltgres types.
	// We may see GMS types during some analyzer steps, so we should wait until those are done.
	if leftType, ok := children[0].Type(ctx).(*pgtypes.DoltgresType); ok {
		// Rather than finding and resolving a comparison function every time we call Eval, we resolve them once and
		// reuse the functions. We also want to avoid re-assigning the parameters of the comparison functions since that
		// will also cause the functions to resolve again. To do this, we store expressions within our struct that the
		// functions reference, so we can freely switch the values within the literals without changing anything
		// regarding the comparison functions. This is usually unsafe, but since we're verifying the types returned by
		// the parameters, and assigning the values to our own literals, we do not have to worry. This offers a
		// significant speedup as function resolution is very expensive, so we want to do it as few times as possible
		// (preferably once).

		// We need a comparison function for each type in the query result
		sch := sq.Query.Schema(ctx)
		leftLiteral := expression.NewLiteral(nil, leftType)
		rightLiterals := make([]*expression.Literal, len(sch))
		compFuncs := make([]framework.Function, len(sch))
		allValidChildren := true
		for i, rightCol := range sch {
			rightType, ok := rightCol.Type.(*pgtypes.DoltgresType)
			if !ok {
				allValidChildren = false
				break
			}
			rightLiterals[i] = expression.NewLiteral(nil, rightType)
			compFuncs[i] = framework.GetBinaryFunction(framework.Operator_BinaryEqual).Compile(ctx, "internal_in_comparison", leftLiteral, rightLiterals[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", in)
			}
		}
		if allValidChildren {
			return &InSubquery{
				leftExpr:      children[0],
				rightExpr:     sq,
				leftLiteral:   leftLiteral,
				rightLiterals: rightLiterals,
				compFuncs:     compFuncs,
			}, nil
		}
	}
	return &InSubquery{
		leftExpr:  children[0],
		rightExpr: sq,
	}, nil
}

// WithResolvedChildren implements the vitess.InjectableExpression interface.
func (in *InSubquery) 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].(*plan.Subquery)
	if !ok {
		return nil, errors.Errorf("expected vitess child to be a *plan.Subquery but has type `%T`", children[1])
	}
	return in.WithChildren(ctx.(*sql.Context), left, right)
}

// Left implements the expression.BinaryExpression interface.
func (in *InSubquery) Left() sql.Expression {
	return in.leftExpr
}

// Right implements the expression.BinaryExpression interface.
func (in *InSubquery) Right() sql.Expression {
	return in.rightExpr
}