diff --git a/expression/builtin_string_vec.go b/expression/builtin_string_vec.go index 3f48e01fee73f..2263b8495054a 100644 --- a/expression/builtin_string_vec.go +++ b/expression/builtin_string_vec.go @@ -1175,11 +1175,62 @@ func (b *builtinLocateBinary2ArgsSig) vecEvalInt(input *chunk.Chunk, result *chu } func (b *builtinLocateBinary3ArgsSig) vectorized() bool { - return false + return true } +// vecEvalInt evals LOCATE(substr,str,pos), case-sensitive. +// See https://dev.mysql.com/doc/refman/5.7/en/string-functions.html#function_locate func (b *builtinLocateBinary3ArgsSig) vecEvalInt(input *chunk.Chunk, result *chunk.Column) error { - return errors.Errorf("not implemented") + n := input.NumRows() + buf0, err := b.bufAllocator.get(types.ETString, n) + if err != nil { + return err + } + defer b.bufAllocator.put(buf0) + if err := b.args[0].VecEvalString(b.ctx, input, buf0); err != nil { + return err + } + buf1, err := b.bufAllocator.get(types.ETString, n) + if err != nil { + return err + } + defer b.bufAllocator.put(buf1) + if err := b.args[1].VecEvalString(b.ctx, input, buf1); err != nil { + return err + } + // store positions in result + if err := b.args[2].VecEvalInt(b.ctx, input, result); err != nil { + return err + } + + result.MergeNulls(buf0, buf1) + i64s := result.Int64s() + for i := 0; i < n; i++ { + if result.IsNull(i) { + continue + } + pos := i64s[i] + // Transfer the argument which starts from 1 to real index which starts from 0. + pos-- + subStr := buf0.GetString(i) + str := buf1.GetString(i) + subStrLen := len(subStr) + if pos < 0 || pos > int64(len(str)-subStrLen) { + i64s[i] = 0 + continue + } else if subStrLen == 0 { + i64s[i] = pos + 1 + continue + } + slice := str[pos:] + idx := strings.Index(slice, subStr) + if idx != -1 { + i64s[i] = pos + int64(idx) + 1 + continue + } + i64s[i] = 0 + } + return nil } func (b *builtinExportSet4ArgSig) vectorized() bool { diff --git a/expression/builtin_string_vec_test.go b/expression/builtin_string_vec_test.go index 75a1faaece68e..192328e8aa7f4 100644 --- a/expression/builtin_string_vec_test.go +++ b/expression/builtin_string_vec_test.go @@ -191,6 +191,18 @@ var vecBuiltinStringCases = map[string][]vecExprBenchCase{ childrenFieldTypes: []*types.FieldType{{Tp: mysql.TypeString, Flag: mysql.BinaryFlag, Collate: charset.CollationBin}, {Tp: mysql.TypeString, Flag: mysql.BinaryFlag, Collate: charset.CollationBin}}, geners: []dataGenerator{&selectStringGener{[]string{"01", "10", "001", "110", "0001", "1110"}}, &selectStringGener{[]string{"010010001000010", "101101110111101"}}}, }, + { + retEvalType: types.ETInt, + childrenTypes: []types.EvalType{types.ETString, types.ETString, types.ETInt}, + childrenFieldTypes: []*types.FieldType{{Tp: mysql.TypeString, Flag: mysql.BinaryFlag, Collate: charset.CollationBin}, {Tp: mysql.TypeString, Flag: mysql.BinaryFlag, Collate: charset.CollationBin}, {Tp: mysql.TypeInt24}}, + geners: []dataGenerator{&randLenStrGener{0, 10}, &randLenStrGener{0, 20}, &rangeInt64Gener{-10, 20}}, + }, + { + retEvalType: types.ETInt, + childrenTypes: []types.EvalType{types.ETString, types.ETString, types.ETInt}, + childrenFieldTypes: []*types.FieldType{{Tp: mysql.TypeString, Flag: mysql.BinaryFlag, Collate: charset.CollationBin}, {Tp: mysql.TypeString, Flag: mysql.BinaryFlag, Collate: charset.CollationBin}, {Tp: mysql.TypeInt24}}, + geners: []dataGenerator{&selectStringGener{[]string{"01", "10", "001", "110", "0001", "1110"}}, &selectStringGener{[]string{"010010001000010", "101101110111101"}}, &rangeInt64Gener{-10, 20}}, + }, }, ast.Hex: { {retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&randHexStrGener{10, 100}}},