github.com/vescale/zgraph@v0.0.0-20230410094002-959c02d50f95/expression/function.go (about)

     1  // Copyright 2023 zGraph Authors. All rights reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package expression
    16  
    17  import (
    18  	"fmt"
    19  	"strings"
    20  
    21  	"github.com/vescale/zgraph/datum"
    22  	"github.com/vescale/zgraph/stmtctx"
    23  	"github.com/vescale/zgraph/types"
    24  )
    25  
    26  var _ Expression = &FuncExpr{}
    27  
    28  type FuncExpr struct {
    29  	Name string
    30  	Args []Expression
    31  	Fn   Function
    32  }
    33  
    34  func (expr *FuncExpr) String() string {
    35  	sb := &strings.Builder{}
    36  	sb.WriteString(expr.Name)
    37  	sb.WriteByte('(')
    38  	sb.WriteString(expr.Args[0].String())
    39  	for i := 1; i < len(expr.Args); i++ {
    40  		sb.WriteString(", ")
    41  		sb.WriteString(expr.Args[i].String())
    42  	}
    43  	sb.WriteByte(')')
    44  	return sb.String()
    45  }
    46  
    47  func (expr *FuncExpr) ReturnType() types.T {
    48  	argTypes := make([]types.T, 0, len(expr.Args))
    49  	for _, arg := range expr.Args {
    50  		argTypes = append(argTypes, arg.ReturnType())
    51  	}
    52  	return expr.Fn.InferReturnType(argTypes)
    53  }
    54  
    55  func (expr *FuncExpr) Eval(stmtCtx *stmtctx.Context, input datum.Row) (datum.Datum, error) {
    56  	funcArgs, nullResult, err := expr.evalFuncArgs(stmtCtx, input)
    57  	if err != nil {
    58  		return nil, err
    59  	}
    60  	if nullResult {
    61  		return datum.Null, nil
    62  	}
    63  	return expr.Fn.Eval(stmtCtx, funcArgs)
    64  }
    65  
    66  func (expr *FuncExpr) evalFuncArgs(stmtCtx *stmtctx.Context, input datum.Row) (_ []datum.Datum, propagateNulls bool, _ error) {
    67  	var funcArgs []datum.Datum
    68  	for _, arg := range expr.Args {
    69  		d, err := arg.Eval(stmtCtx, input)
    70  		if err != nil {
    71  			return nil, false, err
    72  		}
    73  		if d == datum.Null && !expr.Fn.CallOnNullInput() {
    74  			return nil, true, nil
    75  		}
    76  		funcArgs = append(funcArgs, d)
    77  	}
    78  	return funcArgs, false, nil
    79  }
    80  
    81  func NewFuncExpr(name string, args ...Expression) (*FuncExpr, error) {
    82  	fn, ok := builtinFuncs[name]
    83  	if !ok {
    84  		return nil, fmt.Errorf("function %s not found", name)
    85  	}
    86  	if fn.NumArgs() != len(args) {
    87  		return nil, fmt.Errorf("invalid arguments count to call function %s", name)
    88  	}
    89  	return &FuncExpr{
    90  		Name: name,
    91  		Args: args,
    92  		Fn:   fn,
    93  	}, nil
    94  }
    95  
    96  type Function interface {
    97  	NumArgs() int
    98  	InferReturnType(argTypes []types.T) types.T
    99  	CallOnNullInput() bool
   100  	Eval(stmtCtx *stmtctx.Context, args []datum.Datum) (datum.Datum, error)
   101  }
   102  
   103  var builtinFuncs = map[string]Function{
   104  	"id": newBuiltIDFunc(),
   105  }
   106  
   107  type baseBuiltinFunc struct {
   108  	numArgs         int
   109  	callOnNullInput bool
   110  }
   111  
   112  func (b baseBuiltinFunc) NumArgs() int {
   113  	return b.numArgs
   114  }
   115  
   116  func (b baseBuiltinFunc) CallOnNullInput() bool {
   117  	return b.callOnNullInput
   118  }
   119  
   120  func newBaseBuiltinFunc(numArgs int, callOnNullInput bool) baseBuiltinFunc {
   121  	return baseBuiltinFunc{numArgs: numArgs, callOnNullInput: callOnNullInput}
   122  }
   123  
   124  type builtinIDFunc struct {
   125  	baseBuiltinFunc
   126  }
   127  
   128  func (b builtinIDFunc) InferReturnType(_ []types.T) types.T {
   129  	return types.Int
   130  }
   131  
   132  func (b builtinIDFunc) Eval(_ *stmtctx.Context, args []datum.Datum) (datum.Datum, error) {
   133  	switch x := args[0].(type) {
   134  	case *datum.Vertex:
   135  		return datum.NewInt(x.ID), nil
   136  	case *datum.Edge:
   137  		// TODO: Edge should have a unique id.
   138  		return datum.NewInt(0), nil
   139  	default:
   140  		return nil, fmt.Errorf("cannot get id from data type %s", args[0].Type())
   141  	}
   142  }
   143  
   144  func newBuiltIDFunc() Function {
   145  	return builtinIDFunc{newBaseBuiltinFunc(1, false)}
   146  }