Add Match and MatchAt API to FunctionVariant interface

extensions/extension_mgr.go

@@ -9,6 +9,7 @@
 	"path"
 	"sort"
 
+	"github.com/creasty/defaults"
 	"github.com/goccy/go-yaml"
 	substraitgo "github.com/substrait-io/substrait-go"
 	"github.com/substrait-io/substrait-go/proto/extensions"
@@ -179,14 +180,23 @@
 	simpleNames := make(map[string]string)
 
 	for _, f := range file.ScalarFunctions {
+		if err := defaults.Set(&f); err != nil {
+			return fmt.Errorf("failure setting defaults for scalar functions: %w", err)
+		}
 		addToMaps[*ScalarFunctionVariant](id, &f, c.scalarMap, simpleNames)
 	}
 
 	for _, f := range file.AggregateFunctions {
+		if err := defaults.Set(&f); err != nil {
+			return fmt.Errorf("failure setting defaults for aggregate functions: %w", err)
+		}
 		addToMaps[*AggregateFunctionVariant](id, &f, c.aggregateMap, simpleNames)
 	}
 
 	for _, f := range file.WindowFunctions {
+		if err := defaults.Set(&f); err != nil {
+			return fmt.Errorf("failure setting defaults for window functions: %w", err)
+		}
 		addToMaps[*WindowFunctionVariant](id, &f, c.windowMap, simpleNames)
 	}
 

extensions/extension_mgr_test.go

@@ -328,6 +328,8 @@ func TestCollection_GetAllScalarFunctions(t *testing.T) {
 				sf, ok := c.GetScalarFunc(extensions.ID{URI: tt.uri, Name: tt.signature})
 				assert.True(t, ok)
 				assert.Contains(t, scalarFunctions, sf)
+				// verify that default nullability is set to MIRROR
+				assert.Equal(t, extensions.MirrorNullability, sf.Nullability())
 			}
 			if tt.isAggregate {
 				af, ok := c.GetAggregateFunc(extensions.ID{URI: tt.uri, Name: tt.signature})

extensions/simple_extension.go

@@ -188,7 +188,7 @@ const (
 type VariadicBehavior struct {
 	Min                  int                  `yaml:",omitempty"`
 	Max                  int                  `yaml:",omitempty"`
-	ParameterConsistency ParameterConsistency `yaml:"parameterConsistency,omitempty"`
+	ParameterConsistency ParameterConsistency `yaml:"parameterConsistency,omitempty" default:"CONSISTENT"`
 }
 
 func (v *VariadicBehavior) IsValidArgumentCount(count int) bool {
@@ -205,7 +205,7 @@ type ScalarFunctionImpl struct {
 	Variadic         *VariadicBehavior     `yaml:",omitempty"`
 	SessionDependent bool                  `yaml:"sessionDependent,omitempty"`
 	Deterministic    bool                  `yaml:",omitempty"`
-	Nullability      NullabilityHandling   `yaml:",omitempty"`
+	Nullability      NullabilityHandling   `yaml:",omitempty" default:"MIRROR"`
 	Return           parser.TypeExpression `yaml:",omitempty"`
 	Implementation   map[string]string     `yaml:",omitempty"`
 }

extensions/variants.go

@@ -21,6 +21,19 @@
 	URI() string
 	ResolveType(argTypes []types.Type) (types.Type, error)
 	Variadic() *VariadicBehavior
+	// Match this function matches input arguments against this functions parameter list
+	// returns (true, nil) if all input argument can type replace the function definition argument
+	// returns (false, err) for invalid input argument. For e.g. if input argument nullability is not correctly
+	// set this function will return error
+	// returns (false, nil) valid input arguments and argument list type replace parameter list
+	Match(argumentTypes []types.Type) (bool, error)
+	// MatchAt this function matches input argument at position against definition of this
+	// functions argument at same position
+	// returns (true, nil) if input argument can type replace the function definition argument
+	// returns (false, err) for invalid input argument. For e.g. if input arg position is negative or
+	// argument nullability is not correctly set this function will return error
+	// returns (false, nil) valid input argument type and argument can't type replace parameter at argPos
+	MatchAt(typ types.Type, pos int) (bool, error)
 }
 
 func EvaluateTypeExpression(nullHandling NullabilityHandling, expr parser.TypeExpression, paramTypeList ArgumentList, actualTypes []types.Type) (types.Type, error) {
@@ -84,6 +97,109 @@
 	return outType, nil
 }
 
+func matchArguments(nullability NullabilityHandling, paramTypeList ArgumentList, variadicBehavior *VariadicBehavior, actualTypes []types.Type) (bool, error) {
+	if variadicBehavior == nil && len(actualTypes) != len(paramTypeList) {
+		return false, nil
+	} else if variadicBehavior != nil && !validateVariadicBehaviorForMatch(variadicBehavior, actualTypes) {
+		return false, nil
+	}
+	funcDefArgList, err := getFuncDefFromArgList(paramTypeList)
+	if err != nil {
+		return false, nil
+	}
+	// loop over actualTypes and not params since actualTypes can be more than params
+	// considering variadic type
+	for argPos := range actualTypes {
+		match, err1 := matchArgumentAtCommon(actualTypes[argPos], argPos, nullability, funcDefArgList, variadicBehavior)
+		if err1 != nil {
+			return false, err1
+		}
+		if !match {
+			return false, nil
+		}
+	}
+	return true, nil
+}
+
+func matchArgumentAt(actualType types.Type, argPos int, nullability NullabilityHandling, paramTypeList ArgumentList, variadicBehavior *VariadicBehavior) (bool, error) {
+	if argPos < 0 {
+		return false, fmt.Errorf("non-zero argument position")
+	}
+	funcDefArgList, err := getFuncDefFromArgList(paramTypeList)
+	if err != nil {
+		return false, nil
+	}
+	return matchArgumentAtCommon(actualType, argPos, nullability, funcDefArgList, variadicBehavior)
+}
+
+func matchArgumentAtCommon(actualType types.Type, argPos int, nullability NullabilityHandling, funcDefArgList []types.FuncDefArgType, variadicBehavior *VariadicBehavior) (bool, error) {
+	// check if argument out of range
+	if variadicBehavior == nil && argPos >= len(funcDefArgList) {
+		return false, nil
+	} else if variadicBehavior != nil && !variadicBehavior.IsValidArgumentCount(argPos+1) {
+		// this argument position can't be more than the max allowed by the variadic behavior
+		return false, nil
+	}
+
+	if HasSyncParams(funcDefArgList) {
+		return false, fmt.Errorf("%w: function has sync params", substraitgo.ErrNotImplemented)
+	}
+	// if argPos is >= len(funcDefArgList) than last funcDefArg type should be considered for type match
+	// already checked for parameter in range above (considering variadic) so no need to check again for variadic
+	var funcDefArg types.FuncDefArgType
+	if argPos < len(funcDefArgList) {
+		funcDefArg = funcDefArgList[argPos]
+	} else {
+		funcDefArg = funcDefArgList[len(funcDefArgList)-1]
+	}
+	switch nullability {
+	case DiscreteNullability:
+		return funcDefArg.MatchWithNullability(actualType), nil
+	case MirrorNullability, DeclaredOutputNullability:
+		return funcDefArg.MatchWithoutNullability(actualType), nil
+	}
+	// unreachable case
+	return false, fmt.Errorf("invalid nullability type: %s", nullability)
+}
+
+func validateVariadicBehaviorForMatch(variadicBehavior *VariadicBehavior, actualTypes []types.Type) bool {
+	if !variadicBehavior.IsValidArgumentCount(len(actualTypes)) {
+		return false
+	}
+	// verify consistency of variadic behavior
+	if variadicBehavior.ParameterConsistency == ConsistentParams {
+		// all concrete types must be equal for all variable arguments
+		firstVariadicArgIdx := variadicBehavior.Min - 1
+		for i := firstVariadicArgIdx; i < len(actualTypes)-1; i++ {
+			if !actualTypes[i].Equals(actualTypes[i+1]) {
+				return false
+			}
+		}
+	}
+	return true
+}
+
+func getFuncDefFromArgList(paramTypeList ArgumentList) ([]types.FuncDefArgType, error) {
+	var out []types.FuncDefArgType
+	for argPos, param := range paramTypeList {
+		switch paramType := param.(type) {
+		case ValueArg:
+			funcDefArgType, err := paramType.Value.Expr.(*parser.Type).ArgType()
+			if err != nil {
+				return nil, err
+			}
+			out = append(out, funcDefArgType)
+		case EnumArg:
+			return nil, fmt.Errorf("%w: invalid argument at position %d for match operation", substraitgo.ErrInvalidType, argPos)
+		case TypeArg:
+			return nil, fmt.Errorf("%w: invalid argument at position %d for match operation", substraitgo.ErrInvalidType, argPos)
+		default:
+			return nil, fmt.Errorf("%w: invalid argument at position %d for match operation", substraitgo.ErrInvalidType, argPos)
+		}
+	}
+	return out, nil
+}
+
 func parseFuncName(compoundName string) (name string, args ArgumentList) {
 	name, argsStr, _ := strings.Cut(compoundName, ":")
 	if len(argsStr) == 0 {
@@ -160,6 +276,14 @@
 	return ID{URI: s.uri, Name: s.CompoundName()}
 }
 
+func (s *ScalarFunctionVariant) Match(argumentTypes []types.Type) (bool, error) {
+	return matchArguments(s.Nullability(), s.impl.Args, s.impl.Variadic, argumentTypes)
+}
+
+func (s *ScalarFunctionVariant) MatchAt(typ types.Type, pos int) (bool, error) {
+	return matchArgumentAt(typ, pos, s.Nullability(), s.impl.Args, s.impl.Variadic)
+}
+
 // NewAggFuncVariant constructs a variant with the provided name and uri
 // and uses the defaults for everything else.
 //
@@ -268,6 +392,12 @@
 }
 func (s *AggregateFunctionVariant) Ordered() bool { return s.impl.Ordered }
 func (s *AggregateFunctionVariant) MaxSet() int   { return s.impl.MaxSet }
+func (s *AggregateFunctionVariant) Match(argumentTypes []types.Type) (bool, error) {
+	return matchArguments(s.Nullability(), s.impl.Args, s.impl.Variadic, argumentTypes)
+}
+func (s *AggregateFunctionVariant) MatchAt(typ types.Type, pos int) (bool, error) {
+	return matchArgumentAt(typ, pos, s.Nullability(), s.impl.Args, s.impl.Variadic)
+}
 
 type WindowFunctionVariant struct {
 	name        string
@@ -376,6 +506,12 @@
 func (s *WindowFunctionVariant) Ordered() bool          { return s.impl.Ordered }
 func (s *WindowFunctionVariant) MaxSet() int            { return s.impl.MaxSet }
 func (s *WindowFunctionVariant) WindowType() WindowType { return s.impl.WindowType }
+func (s *WindowFunctionVariant) Match(argumentTypes []types.Type) (bool, error) {
+	return matchArguments(s.Nullability(), s.impl.Args, s.impl.Variadic, argumentTypes)
+}
+func (s *WindowFunctionVariant) MatchAt(typ types.Type, pos int) (bool, error) {
+	return matchArgumentAt(typ, pos, s.Nullability(), s.impl.Args, s.impl.Variadic)
+}
 
 // HasSyncParams This API returns if params share a leaf param name
 func HasSyncParams(params []types.FuncDefArgType) bool {