evalengine: Fix JSON weight string computation (#13669)

Signed-off-by: Dirkjan Bussink <d.bussink@gmail.com>

go/mysql/json/helpers.go

@@ -100,6 +100,12 @@ func NewFromSQL(v sqltypes.Value) (*Value, error) {
 		return NewString(v.RawStr()), nil
 	case v.IsBinary():
 		return NewBlob(v.RawStr()), nil
+	case v.IsDateTime(), v.IsTimestamp():
+		return NewDateTime(v.RawStr()), nil
+	case v.IsDate():
+		return NewDate(v.RawStr()), nil
+	case v.IsTime():
+		return NewTime(v.RawStr()), nil
 	default:
 		return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "cannot coerce %v as a JSON type", v)
 	}

go/sqltypes/value.go

@@ -486,6 +486,16 @@ func (v Value) IsDateTime() bool {
 	return v.typ == querypb.Type_DATETIME
 }
 
+// IsTimestamp returns true if Value is date.
+func (v Value) IsTimestamp() bool {
+	return v.typ == querypb.Type_TIMESTAMP
+}
+
+// IsDate returns true if Value is date.
+func (v Value) IsDate() bool {
+	return v.typ == querypb.Type_DATE
+}
+
 // IsTime returns true if Value is time.
 func (v Value) IsTime() bool {
 	return v.typ == querypb.Type_TIME

go/vt/vtgate/evalengine/compiler_asm.go

@@ -734,10 +734,10 @@ func (asm *assembler) CmpTupleNullsafe() {
 		l := env.vm.stack[env.vm.sp-2].(*evalTuple)
 		r := env.vm.stack[env.vm.sp-1].(*evalTuple)
 
-		var equals bool
+		var equals int
 		equals, env.vm.err = evalCompareTuplesNullSafe(l.t, r.t)
 
-		env.vm.stack[env.vm.sp-2] = env.vm.arena.newEvalBool(equals)
+		env.vm.stack[env.vm.sp-2] = env.vm.arena.newEvalBool(equals == 0)
 		env.vm.sp -= 1
 		return 1
 	}, "CMP NULLSAFE TUPLE(SP-2), TUPLE(SP-1)")
@@ -2732,15 +2732,15 @@ func (asm *assembler) Fn_TO_BASE64(t sqltypes.Type, col collations.TypedCollatio
 	}, "FN TO_BASE64 VARCHAR(SP-1)")
 }
 
-func (asm *assembler) Fn_WEIGHT_STRING(length int) {
+func (asm *assembler) Fn_WEIGHT_STRING(typ sqltypes.Type, length int) {
 	asm.emit(func(env *ExpressionEnv) int {
 		input := env.vm.stack[env.vm.sp-1]
 		w, _, err := evalWeightString(nil, input, length, 0)
 		if err != nil {
 			env.vm.err = err
 			return 1
 		}
-		env.vm.stack[env.vm.sp-1] = env.vm.arena.newEvalBinary(w)
+		env.vm.stack[env.vm.sp-1] = env.vm.arena.newEvalRaw(w, typ, collationBinary)
 		return 1
 	}, "FN WEIGHT_STRING (SP-1)")
 }

go/vt/vtgate/evalengine/expr_compare.go

@@ -110,7 +110,7 @@ func (compareGE) compare(left, right eval) (boolean, error) {
 func (compareNullSafeEQ) String() string { return "<=>" }
 func (compareNullSafeEQ) compare(left, right eval) (boolean, error) {
 	cmp, err := evalCompareNullSafe(left, right)
-	return makeboolean(cmp), err
+	return makeboolean(cmp == 0), err
 }
 
 func typeIsTextual(tt sqltypes.Type) bool {
@@ -164,15 +164,21 @@ func compareAsJSON(l, r sqltypes.Type) bool {
 	return l == sqltypes.TypeJSON || r == sqltypes.TypeJSON
 }
 
-func evalCompareNullSafe(lVal, rVal eval) (bool, error) {
-	if lVal == nil || rVal == nil {
-		return lVal == rVal, nil
+func evalCompareNullSafe(lVal, rVal eval) (int, error) {
+	if lVal == nil {
+		if rVal == nil {
+			return 0, nil
+		}
+		return -1, nil
+	}
+	if rVal == nil {
+		return 1, nil
 	}
 	if left, right, ok := compareAsTuples(lVal, rVal); ok {
 		return evalCompareTuplesNullSafe(left.t, right.t)
 	}
 	n, err := evalCompare(lVal, rVal)
-	return n == 0, err
+	return n, err
 }
 
 func evalCompareMany(left, right []eval, fulleq bool) (int, bool, error) {
@@ -263,20 +269,20 @@ func fallbackBinary(t sqltypes.Type) bool {
 	return false
 }
 
-func evalCompareTuplesNullSafe(left, right []eval) (bool, error) {
+func evalCompareTuplesNullSafe(left, right []eval) (int, error) {
 	if len(left) != len(right) {
 		panic("did not typecheck cardinality")
 	}
 	for idx, lResult := range left {
 		res, err := evalCompareNullSafe(lResult, right[idx])
 		if err != nil {
-			return false, err
+			return 0, err
 		}
-		if !res {
-			return false, nil
+		if res != 0 {
+			return res, nil
 		}
 	}
-	return true, nil
+	return 0, nil
 }
 
 // eval implements the Expr interface

go/vt/vtgate/evalengine/fn_string.go

@@ -459,7 +459,11 @@ func (c *builtinWeightString) callable() []Expr {
 }
 
 func (c *builtinWeightString) typeof(env *ExpressionEnv, fields []*querypb.Field) (sqltypes.Type, typeFlag) {
-	_, f := c.Expr.typeof(env, fields)
+	tt, f := c.Expr.typeof(env, fields)
+	switch tt {
+	case sqltypes.Blob, sqltypes.Text, sqltypes.TypeJSON:
+		return sqltypes.Blob, f
+	}
 	return sqltypes.VarBinary, f
 }
 
@@ -471,18 +475,44 @@ func (c *builtinWeightString) eval(env *ExpressionEnv) (eval, error) {
 		return nil, err
 	}
 
+	typ := sqltypes.VarBinary
+
 	if c.Cast == "binary" {
+		switch input.SQLType() {
+		case sqltypes.Blob, sqltypes.Text, sqltypes.TypeJSON:
+			typ = sqltypes.Blob
+		}
+
 		weights, _, err = evalWeightString(weights, evalToBinary(input), c.Len, 0)
 		if err != nil {
 			return nil, err
 		}
-		return newEvalBinary(weights), nil
+		return newEvalRaw(typ, weights, collationBinary), nil
 	}
 
 	switch val := input.(type) {
 	case *evalInt64, *evalUint64, *evalTemporal:
 		weights, _, err = evalWeightString(weights, val, 0, 0)
+	case *evalJSON:
+		// JSON doesn't actually use a sortable weight string for this function, but
+		// returns the weight string directly for the string based representation. This
+		// means that ordering etc. is not correct for JSON values, but that's how MySQL
+		// works here for this function. We still have the internal weight string logic
+		// that can order these correctly.
+		out, err := evalToVarchar(val, collationJSON.Collation, false)
+		if err != nil {
+			return nil, err
+		}
+		weights, _, err = evalWeightString(weights, out, 0, 0)
+		if err != nil {
+			return nil, err
+		}
+		typ = sqltypes.Blob
 	case *evalBytes:
+		switch val.SQLType() {
+		case sqltypes.Blob, sqltypes.Text:
+			typ = sqltypes.Blob
+		}
 		if val.isBinary() {
 			weights, _, err = evalWeightString(weights, val, 0, 0)
 		} else {
@@ -500,7 +530,7 @@ func (c *builtinWeightString) eval(env *ExpressionEnv) (eval, error) {
 		return nil, err
 	}
 
-	return newEvalBinary(weights), nil
+	return newEvalRaw(typ, weights, collationBinary), nil
 }
 
 func (call *builtinWeightString) compile(c *compiler) (ctype, error) {
@@ -514,34 +544,49 @@ func (call *builtinWeightString) compile(c *compiler) (ctype, error) {
 		flag = flag | flagNullable
 	}
 
+	typ := sqltypes.VarBinary
 	skip := c.compileNullCheck1(str)
 	if call.Cast == "binary" {
 		if !sqltypes.IsBinary(str.Type) {
 			c.asm.Convert_xb(1, sqltypes.VarBinary, 0, false)
 		}
-		c.asm.Fn_WEIGHT_STRING(call.Len)
+		switch str.Type {
+		case sqltypes.Blob, sqltypes.Text, sqltypes.TypeJSON:
+			typ = sqltypes.Blob
+		}
+
+		c.asm.Fn_WEIGHT_STRING(typ, call.Len)
 		c.asm.jumpDestination(skip)
 		return ctype{Type: sqltypes.VarBinary, Flag: flagNullable | flagNull, Col: collationBinary}, nil
 	}
 
 	switch str.Type {
 	case sqltypes.Int64, sqltypes.Uint64, sqltypes.Date, sqltypes.Datetime, sqltypes.Timestamp, sqltypes.Time, sqltypes.VarBinary, sqltypes.Binary, sqltypes.Blob:
-		c.asm.Fn_WEIGHT_STRING(0)
-
+		if str.Type == sqltypes.Blob {
+			typ = sqltypes.Blob
+		}
+		c.asm.Fn_WEIGHT_STRING(typ, 0)
+	case sqltypes.TypeJSON:
+		typ = sqltypes.Blob
+		c.asm.Convert_xce(1, sqltypes.VarChar, collationJSON.Collation)
+		c.asm.Fn_WEIGHT_STRING(typ, 0)
 	case sqltypes.VarChar, sqltypes.Char, sqltypes.Text:
+		if str.Type == sqltypes.Text {
+			typ = sqltypes.Blob
+		}
 		var strLen int
 		if call.Cast == "char" {
 			strLen = call.Len
 		}
-		c.asm.Fn_WEIGHT_STRING(strLen)
+		c.asm.Fn_WEIGHT_STRING(typ, strLen)
 
 	default:
 		c.asm.SetNull(1)
 		flag = flag | flagNull | flagNullable
 	}
 
 	c.asm.jumpDestination(skip)
-	return ctype{Type: sqltypes.VarBinary, Flag: flag, Col: collationBinary}, nil
+	return ctype{Type: typ, Flag: flag, Col: collationBinary}, nil
 }
 
 func (call builtinLeftRight) eval(env *ExpressionEnv) (eval, error) {

go/vt/vtgate/evalengine/testcases/cases.go

@@ -847,6 +847,29 @@ func WeightString(yield Query) {
 		`time'-11:22:33'`, `time'-11:22:33' as char(12)`, `time'-11:22:33' as char(2)`, `time'-11:22:33' as binary(12)`, `time'-11:22:33' as binary(2)`,
 		`time'11:22:33'`, `time'11:22:33' as char(12)`, `time'11:22:33' as char(2)`, `time'11:22:33' as binary(12)`, `time'11:22:33' as binary(2)`,
 		`time'101:22:33'`, `time'101:22:33' as char(12)`, `time'101:22:33' as char(2)`, `time'101:22:33' as binary(12)`, `time'101:22:33' as binary(2)`,
+		"cast(0 as json)", "cast(1 as json)",
+		"cast(true as json)", "cast(false as json)",
+		"cast('{}' as json)", "cast('[]' as json)",
+		"cast('null' as json)", "cast('true' as json)", "cast('false' as json)",
+		"cast('1' as json)", "cast('2' as json)", "cast('1.1' as json)", "cast('-1.1' as json)",
+		"cast('9223372036854775807' as json)", "cast('18446744073709551615' as json)",
+		// JSON strings
+		"cast('\"foo\"' as json)", "cast('\"bar\"' as json)", "cast('invalid' as json)",
+		// JSON binary values
+		"cast(_binary' \"foo\"' as json)", "cast(_binary '\"bar\"' as json)",
+		"cast(0xFF666F6F626172FF as json)", "cast(0x666F6F626172FF as json)",
+		"cast(0b01 as json)", "cast(0b001 as json)",
+		// JSON arrays
+		"cast('[\"a\"]' as json)", "cast('[\"ab\"]' as json)",
+		"cast('[\"ab\", \"cd\", \"ef\"]' as json)", "cast('[\"ab\", \"ef\"]' as json)",
+		// JSON objects
+		"cast('{\"a\": 1, \"b\": 2}' as json)", "cast('{\"b\": 2, \"a\": 1}' as json)",
+		"cast('{\"c\": 1, \"b\": 2}' as json)", "cast('{\"b\": 2, \"c\": 1}' as json)",
+		"cast(' \"b\": 2}' as json)", "cast('\"a\": 1' as json)",
+		// JSON date, datetime & time
+		"cast(date '2000-01-01' as json)", "cast(date '2000-01-02' as json)",
+		"cast(timestamp '2000-01-01 12:34:58' as json)",
+		"cast(time '12:34:56' as json)", "cast(time '12:34:58' as json)", "cast(time '5 12:34:58' as json)",
 	}
 
 	for _, i := range inputs {

go/vt/vtgate/evalengine/weights.go

@@ -169,6 +169,8 @@ func evalWeightString(dst []byte, e eval, length, precision int) ([]byte, bool,
 		return coll.WeightString(dst, b, length), false, nil
 	case *evalTemporal:
 		return e.dt.WeightString(dst), true, nil
+	case *evalJSON:
+		return e.WeightString(dst), false, nil
 	}
 
 	return dst, false, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "unexpected type %v", e.SQLType())