github.com/matrixorigin/matrixone@v0.7.0/pkg/sql/plan/function/builtin/multi/json_extract.go (about)

     1  // Copyright 2022 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 multi
    16  
    17  import (
    18  	"github.com/matrixorigin/matrixone/pkg/container/bytejson"
    19  	"github.com/matrixorigin/matrixone/pkg/container/types"
    20  	"github.com/matrixorigin/matrixone/pkg/container/vector"
    21  	"github.com/matrixorigin/matrixone/pkg/vm/process"
    22  )
    23  
    24  type computeFn func([]byte, []*bytejson.Path) (*bytejson.ByteJson, error)
    25  
    26  func computeJson(json []byte, paths []*bytejson.Path) (*bytejson.ByteJson, error) {
    27  	bj := types.DecodeJson(json)
    28  	return bj.Query(paths), nil
    29  }
    30  func computeString(json []byte, paths []*bytejson.Path) (*bytejson.ByteJson, error) {
    31  	bj, err := types.ParseSliceToByteJson(json)
    32  	if err != nil {
    33  		return nil, err
    34  	}
    35  	return bj.Query(paths), nil
    36  }
    37  
    38  func JsonExtract(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int) error {
    39  	jsonVec := parameters[0]
    40  	var fn computeFn
    41  	switch jsonVec.Typ.Oid {
    42  	case types.T_json:
    43  		fn = computeJson
    44  	default:
    45  		fn = computeString
    46  	}
    47  	jsonWrapper := vector.GenerateFunctionStrParameter(jsonVec)
    48  	pathWrapers := make([]vector.FunctionParameterWrapper[types.Varlena], len(parameters)-1)
    49  	rs := vector.MustFunctionResult[types.Varlena](result)
    50  	paths := make([]*bytejson.Path, len(parameters)-1)
    51  	for i := 0; i < len(parameters)-1; i++ {
    52  		pathWrapers[i] = vector.GenerateFunctionStrParameter(parameters[i+1])
    53  	}
    54  	for i := uint64(0); i < uint64(length); i++ {
    55  		jsonBytes, jIsNull := jsonWrapper.GetStrValue(i)
    56  		if jIsNull {
    57  			err := rs.AppendStr(nil, true)
    58  			if err != nil {
    59  				return err
    60  			}
    61  			continue
    62  		}
    63  		skip := false
    64  		for j := 0; j < len(parameters)-1; j++ {
    65  			pathBytes, pIsNull := pathWrapers[j].GetStrValue(i)
    66  			if pIsNull {
    67  				skip = true
    68  				break
    69  			}
    70  			p, err := types.ParseStringToPath(string(pathBytes))
    71  			if err != nil {
    72  				return err
    73  			}
    74  			paths[j] = &p
    75  		}
    76  		if skip {
    77  			err := rs.AppendStr(nil, true)
    78  			if err != nil {
    79  				return err
    80  			}
    81  			continue
    82  		}
    83  		out, err := fn(jsonBytes, paths)
    84  		if err != nil {
    85  			return err
    86  		}
    87  		if out.IsNull() {
    88  			err := rs.AppendStr(nil, true)
    89  			if err != nil {
    90  				return err
    91  			}
    92  			continue
    93  		}
    94  		dt, _ := out.Marshal()
    95  		err = rs.AppendStr(dt, false)
    96  		if err != nil {
    97  			return err
    98  		}
    99  	}
   100  	return nil
   101  }