github.com/matrixorigin/matrixone@v1.2.0/pkg/sql/plan/function/python_udf.go (about)

     1  // Copyright 2023 Matrix Origin
     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 function
    16  
    17  import (
    18  	"encoding/json"
    19  
    20  	"github.com/matrixorigin/matrixone/pkg/container/types"
    21  	"github.com/matrixorigin/matrixone/pkg/container/vector"
    22  	"github.com/matrixorigin/matrixone/pkg/udf"
    23  	"github.com/matrixorigin/matrixone/pkg/vm/process"
    24  )
    25  
    26  // param inputs has four parts:
    27  //  1. inputs[0]: udf, function self
    28  //  2. inputs[1 : size+1]: receivedArgs, args which function received
    29  //  3. inputs[size+1 : 2*size+1]: requiredArgs, args which function required
    30  //  4. inputs[2*size+1]: ret, function ret
    31  //     which size = (len(inputs) - 2) / 2
    32  func checkPythonUdf(overloads []overload, inputs []types.Type) checkResult {
    33  
    34  	if len(inputs)%2 == 1 {
    35  		return newCheckResultWithFailure(failedFunctionParametersWrong)
    36  	}
    37  	if len(inputs) == 2 {
    38  		return newCheckResultWithSuccess(0)
    39  	}
    40  	size := (len(inputs) - 2) / 2
    41  	receivedArgs := inputs[1 : size+1]
    42  	requiredArgs := inputs[size+1 : 2*size+1]
    43  	needCast := false
    44  	for i := 0; i < size; i++ {
    45  		if receivedArgs[i].Oid != requiredArgs[i].Oid {
    46  			canCast, _ := fixedImplicitTypeCast(receivedArgs[i], requiredArgs[i].Oid)
    47  			if !canCast {
    48  				return newCheckResultWithFailure(failedFunctionParametersWrong)
    49  			}
    50  			needCast = true
    51  		}
    52  	}
    53  	if needCast {
    54  		castType := make([]types.Type, size+2)
    55  		castType[0] = inputs[0]
    56  		for i, typ := range requiredArgs {
    57  			castType[i+1] = typ
    58  		}
    59  		castType[size+1] = inputs[2*size+1]
    60  		return newCheckResultWithCast(0, castType)
    61  	}
    62  	return newCheckResultWithSuccess(0)
    63  }
    64  
    65  // param parameters is same with param inputs in function checkPythonUdf
    66  func pythonUdfRetType(parameters []types.Type) types.Type {
    67  	return parameters[len(parameters)-1]
    68  }
    69  
    70  // param parameters has two parts:
    71  //  1. parameters[0]: const vector udf
    72  //  2. parameters[1:]: data vectors
    73  func runPythonUdf(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int) error {
    74  	// udf with context
    75  	u := &UdfWithContext{}
    76  	bytes, _ := vector.GenerateFunctionStrParameter(parameters[0]).GetStrValue(0)
    77  	err := json.Unmarshal(bytes, u)
    78  	if err != nil {
    79  		return err
    80  	}
    81  
    82  	// request
    83  	body := &NonSqlUdfBody{}
    84  	err = json.Unmarshal([]byte(u.Body), body)
    85  	if err != nil {
    86  		return err
    87  	}
    88  	request := &udf.Request{
    89  		Udf: &udf.Udf{
    90  			Handler:      body.Handler,
    91  			IsImport:     body.Import,
    92  			Body:         body.Body,
    93  			RetType:      t2DataType[u.GetRetType().Oid],
    94  			Language:     udf.LanguagePython,
    95  			Db:           u.Db,
    96  			ModifiedTime: u.ModifiedTime,
    97  		},
    98  		Vectors: make([]*udf.DataVector, len(parameters)-1),
    99  		Length:  int64(length),
   100  		Type:    udf.RequestType_DataRequest,
   101  		Context: u.Context,
   102  	}
   103  	for i := 1; i < len(parameters); i++ {
   104  		dataVector, _ := vector2DataVector(parameters[i])
   105  		request.Vectors[i-1] = dataVector
   106  	}
   107  
   108  	// getPkg
   109  	reader := &DefaultPkgReader{
   110  		Proc: proc,
   111  	}
   112  
   113  	// run
   114  	response, err := proc.UdfService.Run(proc.Ctx, request, reader)
   115  	if err != nil {
   116  		return err
   117  	}
   118  
   119  	// response
   120  	err = writeResponse(response, result)
   121  	if err != nil {
   122  		return err
   123  	}
   124  
   125  	return nil
   126  }