expression: implement vectorized evaluation for builtinVectorFloat32 (#…

…55254)

ref #54245

pkg/expression/BUILD.bazel

@@ -44,6 +44,7 @@ go_library(
         "builtin_time_vec.go",
         "builtin_time_vec_generated.go",
         "builtin_vec.go",
+        "builtin_vec_vec.go",
         "builtin_vectorized.go",
         "chunk_executor.go",
         "collation.go",
@@ -182,6 +183,7 @@ go_test(
         "builtin_time_test.go",
         "builtin_time_vec_generated_test.go",
         "builtin_time_vec_test.go",
+        "builtin_vec_vec_test.go",
         "builtin_vectorized_test.go",
         "collation_test.go",
         "column_test.go",

pkg/expression/bench_test.go

@@ -387,6 +387,29 @@ func (g *jsonStringGener) gen() any {
 	return j.String()
 }
 
+type vectorFloat32RandGener struct {
+	dimension int
+	randGen   *defaultRandGen
+}
+
+// create a vectorfloat32 randomly with dimension. if dimension = -1, return nil vectorfloat32
+func newVectorFloat32RandGener(dimension int) *vectorFloat32RandGener {
+	return &vectorFloat32RandGener{dimension, newDefaultRandGen()}
+}
+
+func (g *vectorFloat32RandGener) gen() any {
+	if g.dimension == -1 {
+		return nil
+	}
+	var values []float32
+	for i := 0; i < g.dimension; i++ {
+		values = append(values, g.randGen.Float32())
+	}
+	vec := types.InitVectorFloat32(g.dimension)
+	copy(vec.Elements(), values)
+	return vec
+}
+
 type decimalStringGener struct {
 	randGen *defaultRandGen
 }
@@ -1224,6 +1247,8 @@ func fillColumnWithGener(eType types.EvalType, chk *chunk.Chunk, colIdx int, gen
 			col.AppendJSON(v.(types.BinaryJSON))
 		case types.ETString:
 			col.AppendString(v.(string))
+		case types.ETVectorFloat32:
+			col.AppendVectorFloat32(v.(types.VectorFloat32))
 		}
 	}
 }
@@ -1260,6 +1285,8 @@ func eType2FieldType(eType types.EvalType) *types.FieldType {
 		return types.NewFieldType(mysql.TypeJSON)
 	case types.ETString:
 		return types.NewFieldType(mysql.TypeVarString)
+	case types.ETVectorFloat32:
+		return types.NewFieldType(mysql.TypeTiDBVectorFloat32)
 	default:
 		panic(fmt.Sprintf("EvalType=%v is not supported.", eType))
 	}
@@ -1669,6 +1696,21 @@ func testVectorizedBuiltinFunc(t *testing.T, vecExprCases vecExprBenchCases) {
 					}
 					i++
 				}
+			case types.ETVectorFloat32:
+				err := baseFunc.vecEvalVectorFloat32(ctx, input, output)
+				require.NoErrorf(t, err, "func: %v, case: %+v", baseFuncName, testCase)
+				// do not forget to call ResizeXXX/ReserveXXX
+				require.Equal(t, input.NumRows(), getColumnLen(output, testCase.retEvalType))
+				vecWarnCnt = ctx.GetSessionVars().StmtCtx.WarningCount()
+				for row := it.Begin(); row != it.End(); row = it.Next() {
+					val, isNull, err := baseFunc.evalVectorFloat32(ctx, row)
+					require.NoErrorf(t, err, commentf(i))
+					require.Equal(t, output.IsNull(i), isNull, commentf(i))
+					if !isNull {
+						require.Equal(t, output.GetVectorFloat32(i).Compare(val), 0, commentf(i))
+					}
+					i++
+				}
 			default:
 				t.Fatalf("evalType=%v is not supported", testCase.retEvalType)
 			}
@@ -1823,6 +1865,12 @@ func benchmarkVectorizedBuiltinFunc(b *testing.B, vecExprCases vecExprBenchCases
 							b.Fatal(err)
 						}
 					}
+				case types.ETVectorFloat32:
+					for i := 0; i < b.N; i++ {
+						if err := baseFunc.vecEvalVectorFloat32(ctx, input, output); err != nil {
+							b.Fatal(err)
+						}
+					}
 				default:
 					b.Fatalf("evalType=%v is not supported", testCase.retEvalType)
 				}
@@ -1936,6 +1984,21 @@ func benchmarkVectorizedBuiltinFunc(b *testing.B, vecExprCases vecExprBenchCases
 							}
 						}
 					}
+				case types.ETVectorFloat32:
+					for i := 0; i < b.N; i++ {
+						output.Reset(testCase.retEvalType)
+						for row := it.Begin(); row != it.End(); row = it.Next() {
+							v, isNull, err := baseFunc.evalVectorFloat32(ctx, row)
+							if err != nil {
+								b.Fatal(err)
+							}
+							if isNull {
+								output.AppendNull()
+							} else {
+								output.AppendVectorFloat32(v)
+							}
+						}
+					}
 				default:
 					b.Fatalf("evalType=%v is not supported", testCase.retEvalType)
 				}