expression: Add vector functions (#55021)

ref #54245

pkg/expression/aggregation/aggregation.go

@@ -239,6 +239,9 @@ func CheckAggPushDown(ctx expression.EvalContext, aggFunc *AggFuncDesc, storeTyp
 	if aggFunc.Name == ast.AggFuncApproxPercentile {
 		return false
 	}
+	if !checkVectorAggPushDown(ctx, aggFunc) {
+		return false
+	}
 	ret := true
 	switch storeType {
 	case kv.TiFlash:
@@ -253,6 +256,22 @@ func CheckAggPushDown(ctx expression.EvalContext, aggFunc *AggFuncDesc, storeTyp
 	return ret
 }
 
+// checkVectorAggPushDown returns false if this aggregate function is not supported to push down.
+// - The aggregate function is not calculated over a Vector column (returns true)
+// - The aggregate function is calculated over a Vector column and the function is supported (returns true)
+// - The aggregate function is calculated over a Vector column and the function is not supported (returns false)
+func checkVectorAggPushDown(ctx expression.EvalContext, aggFunc *AggFuncDesc) bool {
+	switch aggFunc.Name {
+	case ast.AggFuncCount, ast.AggFuncMin, ast.AggFuncMax, ast.AggFuncFirstRow:
+		return true
+	default:
+		if aggFunc.Args[0].GetType(ctx).GetType() == mysql.TypeTiDBVectorFloat32 {
+			return false
+		}
+	}
+	return true
+}
+
 // CheckAggPushFlash checks whether an agg function can be pushed to flash storage.
 func CheckAggPushFlash(ctx expression.EvalContext, aggFunc *AggFuncDesc) bool {
 	for _, arg := range aggFunc.Args {

pkg/expression/builtin.go

@@ -929,9 +929,14 @@ var funcs = map[string]functionClass{
 	ast.JSONLength:        &jsonLengthFunctionClass{baseFunctionClass{ast.JSONLength, 1, 2}},
 
 	// vector functions (TiDB extension)
-	ast.VecDims:     &vecDimsFunctionClass{baseFunctionClass{ast.VecDims, 1, 1}},
-	ast.VecFromText: &vecFromTextFunctionClass{baseFunctionClass{ast.VecFromText, 1, 1}},
-	ast.VecAsText:   &vecAsTextFunctionClass{baseFunctionClass{ast.VecAsText, 1, 1}},
+	ast.VecDims:                 &vecDimsFunctionClass{baseFunctionClass{ast.VecDims, 1, 1}},
+	ast.VecL1Distance:           &vecL1DistanceFunctionClass{baseFunctionClass{ast.VecL1Distance, 2, 2}},
+	ast.VecL2Distance:           &vecL2DistanceFunctionClass{baseFunctionClass{ast.VecL2Distance, 2, 2}},
+	ast.VecNegativeInnerProduct: &vecNegativeInnerProductFunctionClass{baseFunctionClass{ast.VecNegativeInnerProduct, 2, 2}},
+	ast.VecCosineDistance:       &vecCosineDistanceFunctionClass{baseFunctionClass{ast.VecCosineDistance, 2, 2}},
+	ast.VecL2Norm:               &vecL2NormFunctionClass{baseFunctionClass{ast.VecL2Norm, 1, 1}},
+	ast.VecFromText:             &vecFromTextFunctionClass{baseFunctionClass{ast.VecFromText, 1, 1}},
+	ast.VecAsText:               &vecAsTextFunctionClass{baseFunctionClass{ast.VecAsText, 1, 1}},
 
 	// TiDB internal function.
 	ast.TiDBDecodeKey: &tidbDecodeKeyFunctionClass{baseFunctionClass{ast.TiDBDecodeKey, 1, 1}},

pkg/expression/builtin_arithmetic.go

@@ -57,6 +57,10 @@ var (
 	_ builtinFunc = &builtinArithmeticModIntSignedSignedSig{}
 	_ builtinFunc = &builtinArithmeticModRealSig{}
 	_ builtinFunc = &builtinArithmeticModDecimalSig{}
+
+	_ builtinFunc = &builtinArithmeticPlusVectorFloat32Sig{}
+	_ builtinFunc = &builtinArithmeticMinusVectorFloat32Sig{}
+	_ builtinFunc = &builtinArithmeticMultiplyVectorFloat32Sig{}
 )
 
 // isConstantBinaryLiteral return true if expr is constant binary literal
@@ -167,6 +171,15 @@ func (c *arithmeticPlusFunctionClass) getFunction(ctx BuildContext, args []Expre
 	if err := c.verifyArgs(args); err != nil {
 		return nil, err
 	}
+	if args[0].GetType(ctx.GetEvalCtx()).EvalType().IsVectorKind() || args[1].GetType(ctx.GetEvalCtx()).EvalType().IsVectorKind() {
+		bf, err := newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETVectorFloat32, types.ETVectorFloat32, types.ETVectorFloat32)
+		if err != nil {
+			return nil, err
+		}
+		sig := &builtinArithmeticPlusVectorFloat32Sig{bf}
+		// sig.setPbCode(tipb.ScalarFuncSig_PlusVectorFloat32)
+		return sig, nil
+	}
 	lhsEvalTp, rhsEvalTp := numericContextResultType(ctx.GetEvalCtx(), args[0]), numericContextResultType(ctx.GetEvalCtx(), args[1])
 	if lhsEvalTp == types.ETReal || rhsEvalTp == types.ETReal {
 		bf, err := newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETReal, types.ETReal, types.ETReal)
@@ -317,6 +330,15 @@ func (c *arithmeticMinusFunctionClass) getFunction(ctx BuildContext, args []Expr
 	if err := c.verifyArgs(args); err != nil {
 		return nil, err
 	}
+	if args[0].GetType(ctx.GetEvalCtx()).EvalType().IsVectorKind() || args[1].GetType(ctx.GetEvalCtx()).EvalType().IsVectorKind() {
+		bf, err := newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETVectorFloat32, types.ETVectorFloat32, types.ETVectorFloat32)
+		if err != nil {
+			return nil, err
+		}
+		sig := &builtinArithmeticMinusVectorFloat32Sig{bf}
+		// sig.setPbCode(tipb.ScalarFuncSig_PlusVectorFloat32)
+		return sig, nil
+	}
 	lhsEvalTp, rhsEvalTp := numericContextResultType(ctx.GetEvalCtx(), args[0]), numericContextResultType(ctx.GetEvalCtx(), args[1])
 	if lhsEvalTp == types.ETReal || rhsEvalTp == types.ETReal {
 		bf, err := newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETReal, types.ETReal, types.ETReal)
@@ -500,6 +522,15 @@ func (c *arithmeticMultiplyFunctionClass) getFunction(ctx BuildContext, args []E
 	if err := c.verifyArgs(args); err != nil {
 		return nil, err
 	}
+	if args[0].GetType(ctx.GetEvalCtx()).EvalType().IsVectorKind() || args[1].GetType(ctx.GetEvalCtx()).EvalType().IsVectorKind() {
+		bf, err := newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETVectorFloat32, types.ETVectorFloat32, types.ETVectorFloat32)
+		if err != nil {
+			return nil, err
+		}
+		sig := &builtinArithmeticMultiplyVectorFloat32Sig{bf}
+		// sig.setPbCode(tipb.ScalarFuncSig_PlusVectorFloat32)
+		return sig, nil
+	}
 	lhsTp, rhsTp := args[0].GetType(ctx.GetEvalCtx()), args[1].GetType(ctx.GetEvalCtx())
 	lhsEvalTp, rhsEvalTp := numericContextResultType(ctx.GetEvalCtx(), args[0]), numericContextResultType(ctx.GetEvalCtx(), args[1])
 	if lhsEvalTp == types.ETReal || rhsEvalTp == types.ETReal {
@@ -1157,3 +1188,90 @@ func (s *builtinArithmeticModIntSignedSignedSig) evalInt(ctx EvalContext, row ch
 
 	return a % b, false, nil
 }
+
+type builtinArithmeticPlusVectorFloat32Sig struct {
+	baseBuiltinFunc
+}
+
+func (s *builtinArithmeticPlusVectorFloat32Sig) Clone() builtinFunc {
+	newSig := &builtinArithmeticPlusVectorFloat32Sig{}
+	newSig.cloneFrom(&s.baseBuiltinFunc)
+	return newSig
+}
+
+func (s *builtinArithmeticPlusVectorFloat32Sig) evalVectorFloat32(ctx EvalContext, row chunk.Row) (types.VectorFloat32, bool, error) {
+	a, isLHSNull, err := s.args[0].EvalVectorFloat32(ctx, row)
+	if err != nil {
+		return types.ZeroVectorFloat32, isLHSNull, err
+	}
+	b, isRHSNull, err := s.args[1].EvalVectorFloat32(ctx, row)
+	if err != nil {
+		return types.ZeroVectorFloat32, isRHSNull, err
+	}
+	if isLHSNull || isRHSNull {
+		return types.ZeroVectorFloat32, true, nil
+	}
+	v, err := a.Add(b)
+	if err != nil {
+		return types.ZeroVectorFloat32, true, err
+	}
+	return v, false, nil
+}
+
+type builtinArithmeticMinusVectorFloat32Sig struct {
+	baseBuiltinFunc
+}
+
+func (s *builtinArithmeticMinusVectorFloat32Sig) Clone() builtinFunc {
+	newSig := &builtinArithmeticMinusVectorFloat32Sig{}
+	newSig.cloneFrom(&s.baseBuiltinFunc)
+	return newSig
+}
+
+func (s *builtinArithmeticMinusVectorFloat32Sig) evalVectorFloat32(ctx EvalContext, row chunk.Row) (types.VectorFloat32, bool, error) {
+	a, isLHSNull, err := s.args[0].EvalVectorFloat32(ctx, row)
+	if err != nil {
+		return types.ZeroVectorFloat32, isLHSNull, err
+	}
+	b, isRHSNull, err := s.args[1].EvalVectorFloat32(ctx, row)
+	if err != nil {
+		return types.ZeroVectorFloat32, isRHSNull, err
+	}
+	if isLHSNull || isRHSNull {
+		return types.ZeroVectorFloat32, true, nil
+	}
+	v, err := a.Sub(b)
+	if err != nil {
+		return types.ZeroVectorFloat32, true, err
+	}
+	return v, false, nil
+}
+
+type builtinArithmeticMultiplyVectorFloat32Sig struct {
+	baseBuiltinFunc
+}
+
+func (s *builtinArithmeticMultiplyVectorFloat32Sig) Clone() builtinFunc {
+	newSig := &builtinArithmeticMultiplyVectorFloat32Sig{}
+	newSig.cloneFrom(&s.baseBuiltinFunc)
+	return newSig
+}
+
+func (s *builtinArithmeticMultiplyVectorFloat32Sig) evalVectorFloat32(ctx EvalContext, row chunk.Row) (types.VectorFloat32, bool, error) {
+	a, isLHSNull, err := s.args[0].EvalVectorFloat32(ctx, row)
+	if err != nil {
+		return types.ZeroVectorFloat32, isLHSNull, err
+	}
+	b, isRHSNull, err := s.args[1].EvalVectorFloat32(ctx, row)
+	if err != nil {
+		return types.ZeroVectorFloat32, isRHSNull, err
+	}
+	if isLHSNull || isRHSNull {
+		return types.ZeroVectorFloat32, true, nil
+	}
+	v, err := a.Mul(b)
+	if err != nil {
+		return types.ZeroVectorFloat32, true, err
+	}
+	return v, false, nil
+}

pkg/expression/builtin_cast.go

@@ -691,7 +691,7 @@ func (c *castAsVectorFloat32FunctionClass) getFunction(ctx BuildContext, args []
 		sig.setPbCode(tipb.ScalarFuncSig_CastVectorFloat32AsVectorFloat32)
 	case types.ETString:
 		sig = &builtinCastStringAsVectorFloat32Sig{bf}
-		sig.setPbCode(tipb.ScalarFuncSig_CastStringAsVectorFloat32)
+		// sig.setPbCode(tipb.ScalarFuncSig_CastStringAsVectorFloat32)
 	default:
 		return nil, errors.Errorf("cannot cast from %s to %s", argTp, "VectorFloat32")
 	}

pkg/expression/builtin_compare.go

@@ -46,6 +46,7 @@ var (
 	_ builtinFunc = &builtinCoalesceStringSig{}
 	_ builtinFunc = &builtinCoalesceTimeSig{}
 	_ builtinFunc = &builtinCoalesceDurationSig{}
+	_ builtinFunc = &builtinCoalesceVectorFloat32Sig{}
 
 	_ builtinFunc = &builtinGreatestIntSig{}
 	_ builtinFunc = &builtinGreatestRealSig{}
@@ -54,13 +55,15 @@ var (
 	_ builtinFunc = &builtinGreatestDurationSig{}
 	_ builtinFunc = &builtinGreatestTimeSig{}
 	_ builtinFunc = &builtinGreatestCmpStringAsTimeSig{}
+	_ builtinFunc = &builtinGreatestVectorFloat32Sig{}
 	_ builtinFunc = &builtinLeastIntSig{}
 	_ builtinFunc = &builtinLeastRealSig{}
 	_ builtinFunc = &builtinLeastDecimalSig{}
 	_ builtinFunc = &builtinLeastStringSig{}
 	_ builtinFunc = &builtinLeastTimeSig{}
 	_ builtinFunc = &builtinLeastDurationSig{}
 	_ builtinFunc = &builtinLeastCmpStringAsTimeSig{}
+	_ builtinFunc = &builtinLeastVectorFloat32Sig{}
 	_ builtinFunc = &builtinIntervalIntSig{}
 	_ builtinFunc = &builtinIntervalRealSig{}
 
@@ -167,6 +170,9 @@ func (c *coalesceFunctionClass) getFunction(ctx BuildContext, args []Expression)
 	case types.ETJson:
 		sig = &builtinCoalesceJSONSig{bf}
 		sig.setPbCode(tipb.ScalarFuncSig_CoalesceJson)
+	case types.ETVectorFloat32:
+		sig = &builtinCoalesceVectorFloat32Sig{bf}
+		// sig.setPbCode(tipb.ScalarFuncSig_CoalesceVectorFloat32)
 	default:
 		return nil, errors.Errorf("%s is not supported for COALESCE()", retEvalTp)
 	}
@@ -331,6 +337,28 @@ func (b *builtinCoalesceJSONSig) evalJSON(ctx EvalContext, row chunk.Row) (res t
 	return res, isNull, err
 }
 
+// builtinCoalesceVectorFloat32Sig is builtin function coalesce signature which return type vector float32.
+// See http://dev.mysql.com/doc/refman/5.7/en/comparison-operators.html#function_coalesce
+type builtinCoalesceVectorFloat32Sig struct {
+	baseBuiltinFunc
+}
+
+func (b *builtinCoalesceVectorFloat32Sig) Clone() builtinFunc {
+	newSig := &builtinCoalesceVectorFloat32Sig{}
+	newSig.cloneFrom(&b.baseBuiltinFunc)
+	return newSig
+}
+
+func (b *builtinCoalesceVectorFloat32Sig) evalVectorFloat32(ctx EvalContext, row chunk.Row) (res types.VectorFloat32, isNull bool, err error) {
+	for _, a := range b.getArgs() {
+		res, isNull, err = a.EvalVectorFloat32(ctx, row)
+		if err != nil || !isNull {
+			break
+		}
+	}
+	return res, isNull, err
+}
+
 func aggregateType(ctx EvalContext, args []Expression) *types.FieldType {
 	fieldTypes := make([]*types.FieldType, len(args))
 	for i := range fieldTypes {
@@ -499,6 +527,9 @@ func (c *greatestFunctionClass) getFunction(ctx BuildContext, args []Expression)
 			sig = &builtinGreatestTimeSig{bf, false}
 			sig.setPbCode(tipb.ScalarFuncSig_GreatestTime)
 		}
+	case types.ETVectorFloat32:
+		sig = &builtinGreatestVectorFloat32Sig{bf}
+		// sig.setPbCode(tipb.ScalarFuncSig_GreatestVectorFloat32)
 	default:
 		return nil, errors.Errorf("unsupported type %s during evaluation", argTp)
 	}
@@ -754,6 +785,29 @@ func (b *builtinGreatestDurationSig) evalDuration(ctx EvalContext, row chunk.Row
 	return res, false, nil
 }
 
+type builtinGreatestVectorFloat32Sig struct {
+	baseBuiltinFunc
+}
+
+func (b *builtinGreatestVectorFloat32Sig) Clone() builtinFunc {
+	newSig := &builtinGreatestVectorFloat32Sig{}
+	newSig.cloneFrom(&b.baseBuiltinFunc)
+	return newSig
+}
+
+func (b *builtinGreatestVectorFloat32Sig) evalVectorFloat32(ctx EvalContext, row chunk.Row) (res types.VectorFloat32, isNull bool, err error) {
+	for i := 0; i < len(b.args); i++ {
+		v, isNull, err := b.args[i].EvalVectorFloat32(ctx, row)
+		if isNull || err != nil {
+			return types.VectorFloat32{}, true, err
+		}
+		if i == 0 || v.Compare(res) > 0 {
+			res = v
+		}
+	}
+	return res, false, nil
+}
+
 type leastFunctionClass struct {
 	baseFunctionClass
 }
@@ -814,6 +868,9 @@ func (c *leastFunctionClass) getFunction(ctx BuildContext, args []Expression) (s
 			sig = &builtinLeastTimeSig{bf, false}
 			sig.setPbCode(tipb.ScalarFuncSig_LeastTime)
 		}
+	case types.ETVectorFloat32:
+		sig = &builtinLeastVectorFloat32Sig{bf}
+		// sig.setPbCode(tipb.ScalarFuncSig_LeastVectorFloat32)
 	default:
 		return nil, errors.Errorf("unsupported type %s during evaluation", argTp)
 	}
@@ -1039,6 +1096,29 @@ func (b *builtinLeastDurationSig) evalDuration(ctx EvalContext, row chunk.Row) (
 	return res, false, nil
 }
 
+type builtinLeastVectorFloat32Sig struct {
+	baseBuiltinFunc
+}
+
+func (b *builtinLeastVectorFloat32Sig) Clone() builtinFunc {
+	newSig := &builtinLeastVectorFloat32Sig{}
+	newSig.cloneFrom(&b.baseBuiltinFunc)
+	return newSig
+}
+
+func (b *builtinLeastVectorFloat32Sig) evalVectorFloat32(ctx EvalContext, row chunk.Row) (res types.VectorFloat32, isNull bool, err error) {
+	for i := 0; i < len(b.args); i++ {
+		v, isNull, err := b.args[i].EvalVectorFloat32(ctx, row)
+		if isNull || err != nil {
+			return types.VectorFloat32{}, true, err
+		}
+		if i == 0 || v.Compare(res) < 0 {
+			res = v
+		}
+	}
+	return res, false, nil
+}
+
 type intervalFunctionClass struct {
 	baseFunctionClass
 }