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

     1  // Copyright 2021 - 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/common/moerr"
    19  	"github.com/matrixorigin/matrixone/pkg/container/nulls"
    20  	"github.com/matrixorigin/matrixone/pkg/container/types"
    21  	"github.com/matrixorigin/matrixone/pkg/container/vector"
    22  	"github.com/matrixorigin/matrixone/pkg/vectorize/split_part"
    23  	"github.com/matrixorigin/matrixone/pkg/vm/process"
    24  )
    25  
    26  var fnMap = []func(s1, s2 []string, s3 []uint32, nsps []*nulls.Nulls, rs []string, rnsp *nulls.Nulls){
    27  	1: split_part.SplitPart1,
    28  	split_part.SplitPart2,
    29  	split_part.SplitPart3,
    30  	split_part.SplitPart4,
    31  	split_part.SplitPart5,
    32  	split_part.SplitPart6,
    33  	split_part.SplitPart7,
    34  }
    35  
    36  func SplitPart(vectors []*vector.Vector, proc *process.Process) (vec *vector.Vector, err error) {
    37  	defer func() {
    38  		if err != nil && vec != nil {
    39  			vec.Free(proc.Mp())
    40  		}
    41  	}()
    42  	v1, v2, v3 := vectors[0], vectors[1], vectors[2]
    43  	resultType := types.T_varchar.ToType()
    44  	maxLen := findMaxLen(vectors)
    45  	if v1.IsScalarNull() || v2.IsScalarNull() || v3.IsScalarNull() {
    46  		vec = proc.AllocConstNullVector(resultType, maxLen)
    47  		return
    48  	}
    49  	if !validCount(v3) {
    50  		err = moerr.NewInvalidInput(proc.Ctx, "split_part: field contains non-positive integer")
    51  		return
    52  	}
    53  	s1, s2, s3 := vector.MustStrCols(v1), vector.MustStrCols(v2), vector.MustTCols[uint32](v3)
    54  	if v1.IsScalar() && v2.IsScalar() && v3.IsScalar() {
    55  		vec = proc.AllocScalarVector(resultType)
    56  		ret, isNull := split_part.SplitSingle(s1[0], s2[0], s3[0])
    57  		if isNull {
    58  			vec.Nsp.Set(0)
    59  			return
    60  		}
    61  		err = vec.Append([]byte(ret), false, proc.Mp())
    62  		return
    63  	}
    64  	vec, err = proc.AllocVectorOfRows(resultType, int64(maxLen), nil)
    65  	if err != nil {
    66  		return
    67  	}
    68  	fnIdx := determineFn(vectors)
    69  
    70  	rs := make([]string, maxLen)
    71  	fnMap[fnIdx](s1, s2, s3, []*nulls.Nulls{v1.Nsp, v2.Nsp, v3.Nsp}, rs, vec.Nsp)
    72  	for i, r := range rs {
    73  		if vec.Nsp.Contains(uint64(i)) {
    74  			continue
    75  		}
    76  		err = vector.SetStringAt(vec, i, r, proc.Mp())
    77  		if err != nil {
    78  			return
    79  		}
    80  	}
    81  	return vec, err
    82  }
    83  
    84  func findMaxLen(vecs []*vector.Vector) int {
    85  	ret := 0
    86  	for _, vec := range vecs {
    87  		if vec.Length() > ret {
    88  			ret = vec.Length()
    89  		}
    90  	}
    91  	return ret
    92  }
    93  
    94  func determineFn(vecs []*vector.Vector) int {
    95  	ret := 0
    96  	vecCnt := 3
    97  	for i, vec := range vecs {
    98  		if !vec.IsScalar() {
    99  			ret |= 1 << (vecCnt - i - 1)
   100  		}
   101  	}
   102  	return ret
   103  }
   104  
   105  func validCount(v *vector.Vector) bool {
   106  	s := vector.MustTCols[uint32](v)
   107  	for i, x := range s {
   108  		if v.Nsp.Contains(uint64(i)) {
   109  			continue
   110  		}
   111  		if x == 0 {
   112  			return false
   113  		}
   114  	}
   115  	return true
   116  }