Implement temporal comparisons (#17826)

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

go/vt/vtgate/evalengine/compiler.go

@@ -335,7 +335,7 @@ func (c *compiler) compileToDecimal(ct ctype, offset int) ctype {
 		c.asm.Convert_id(offset)
 	case sqltypes.Uint64:
 		c.asm.Convert_ud(offset)
-	case sqltypes.Datetime, sqltypes.Time:
+	case sqltypes.Datetime, sqltypes.Time, sqltypes.Timestamp:
 		scale = ct.Size
 		size = ct.Size + decimalSizeBase
 		fallthrough
@@ -345,6 +345,28 @@ func (c *compiler) compileToDecimal(ct ctype, offset int) ctype {
 	return ctype{Type: sqltypes.Decimal, Flag: ct.Flag, Col: collationNumeric, Scale: scale, Size: size}
 }
 
+func (c *compiler) compileToTemporal(doct ctype, typ sqltypes.Type, offset, prec int) ctype {
+	switch doct.Type {
+	case typ:
+		if int(doct.Size) == prec {
+			return doct
+		}
+		fallthrough
+	default:
+		switch typ {
+		case sqltypes.Date:
+			c.asm.Convert_xD(offset, c.sqlmode.AllowZeroDate())
+		case sqltypes.Datetime:
+			c.asm.Convert_xDT(offset, prec, c.sqlmode.AllowZeroDate())
+		case sqltypes.Timestamp:
+			c.asm.Convert_xDTs(offset, prec, c.sqlmode.AllowZeroDate())
+		case sqltypes.Time:
+			c.asm.Convert_xT(offset, prec)
+		}
+	}
+	return ctype{Type: typ, Col: collationBinary, Flag: flagNullable}
+}
+
 func (c *compiler) compileToDate(doct ctype, offset int) ctype {
 	switch doct.Type {
 	case sqltypes.Date:
@@ -366,6 +388,17 @@ func (c *compiler) compileToDateTime(doct ctype, offset, prec int) ctype {
 	return ctype{Type: sqltypes.Datetime, Size: int32(prec), Col: collationBinary, Flag: flagNullable}
 }
 
+func (c *compiler) compileToTimestamp(doct ctype, offset, prec int) ctype {
+	switch doct.Type {
+	case sqltypes.Timestamp:
+		c.asm.Convert_tp(offset, prec)
+		return doct
+	default:
+		c.asm.Convert_xDTs(offset, prec, c.sqlmode.AllowZeroDate())
+	}
+	return ctype{Type: sqltypes.Timestamp, Size: int32(prec), Col: collationBinary, Flag: flagNullable}
+}
+
 func (c *compiler) compileToTime(doct ctype, offset, prec int) ctype {
 	switch doct.Type {
 	case sqltypes.Time:

go/vt/vtgate/evalengine/compiler_asm.go

@@ -767,11 +767,11 @@ func (asm *assembler) CmpDates() {
 	}, "CMP DATE(SP-2), DATE(SP-1)")
 }
 
-func (asm *assembler) Collate(col collations.ID) {
+func (asm *assembler) Collate(col collations.TypedCollation) {
 	asm.emit(func(env *ExpressionEnv) int {
 		a := env.vm.stack[env.vm.sp-1].(*evalBytes)
 		a.tt = int16(sqltypes.VarChar)
-		a.col.Collation = col
+		a.col = col
 		return 1
 	}, "COLLATE VARCHAR(SP-1), %d", col)
 }
@@ -1170,6 +1170,21 @@ func (asm *assembler) Convert_xDT(offset, prec int, allowZero bool) {
 	}, "CONV (SP-%d), DATETIME", offset)
 }
 
+func (asm *assembler) Convert_xDTs(offset, prec int, allowZero bool) {
+	asm.emit(func(env *ExpressionEnv) int {
+		// Need to explicitly check here or we otherwise
+		// store a nil wrapper in an interface vs. a direct
+		// nil.
+		dt := evalToTimestamp(env.vm.stack[env.vm.sp-offset], prec, env.now, allowZero)
+		if dt == nil {
+			env.vm.stack[env.vm.sp-offset] = nil
+		} else {
+			env.vm.stack[env.vm.sp-offset] = dt
+		}
+		return 1
+	}, "CONV (SP-%d), TIMESTAMP", offset)
+}
+
 func (asm *assembler) Convert_xT(offset, prec int) {
 	asm.emit(func(env *ExpressionEnv) int {
 		t := evalToTime(env.vm.stack[env.vm.sp-offset], prec)
@@ -2670,6 +2685,40 @@ func (asm *assembler) Fn_MULTICMP_u(args int, lessThan bool) {
 	}, "FN MULTICMP UINT64(SP-%d)...UINT64(SP-1)", args)
 }
 
+func (asm *assembler) Fn_MULTICMP_temporal(args int, lessThan bool) {
+	asm.adjustStack(-(args - 1))
+
+	asm.emit(func(env *ExpressionEnv) int {
+		var x *evalTemporal
+		x, _ = env.vm.stack[env.vm.sp-args].(*evalTemporal)
+		for sp := env.vm.sp - args + 1; sp < env.vm.sp; sp++ {
+			if env.vm.stack[sp] == nil {
+				if lessThan {
+					x = nil
+				}
+				continue
+			}
+			y := env.vm.stack[sp].(*evalTemporal)
+			if lessThan == (y.compare(x) < 0) {
+				x = y
+			}
+		}
+		env.vm.stack[env.vm.sp-args] = x
+		env.vm.sp -= args - 1
+		return 1
+	}, "FN MULTICMP TEMPORAL(SP-%d)...TEMPORAL(SP-1)", args)
+}
+
+func (asm *assembler) Fn_MULTICMP_temporal_fallback(f multiComparisonFunc, args int, cmp, prec int) {
+	asm.adjustStack(-(args - 1))
+
+	asm.emit(func(env *ExpressionEnv) int {
+		env.vm.stack[env.vm.sp-args], env.vm.err = f(env, env.vm.stack[env.vm.sp-args:env.vm.sp], cmp, prec)
+		env.vm.sp -= args - 1
+		return 1
+	}, "FN MULTICMP_FALLBACK TEMPORAL(SP-%d)...TEMPORAL(SP-1)", args)
+}
+
 func (asm *assembler) Fn_REPEAT(base sqltypes.Type, fallback sqltypes.Type) {
 	asm.adjustStack(-1)
 

go/vt/vtgate/evalengine/compiler_asm_push.go

@@ -362,6 +362,23 @@ func (asm *assembler) PushColumn_datetime(offset int) {
 	}, "PUSH DATETIME(:%d)", offset)
 }
 
+func push_timestamp(env *ExpressionEnv, raw []byte) int {
+	env.vm.stack[env.vm.sp], env.vm.err = parseTimestamp(raw)
+	env.vm.sp++
+	return 1
+}
+
+func (asm *assembler) PushColumn_timestamp(offset int) {
+	asm.adjustStack(1)
+	asm.emit(func(env *ExpressionEnv) int {
+		col := env.Row[offset]
+		if col.IsNull() {
+			return push_null(env)
+		}
+		return push_timestamp(env, col.Raw())
+	}, "PUSH TIMESTAMP(:%d)", offset)
+}
+
 func (asm *assembler) PushBVar_datetime(key string) {
 	asm.adjustStack(1)
 	asm.emit(func(env *ExpressionEnv) int {
@@ -374,6 +391,18 @@ func (asm *assembler) PushBVar_datetime(key string) {
 	}, "PUSH DATETIME(:%q)", key)
 }
 
+func (asm *assembler) PushBVar_timestamp(key string) {
+	asm.adjustStack(1)
+	asm.emit(func(env *ExpressionEnv) int {
+		var bvar *querypb.BindVariable
+		bvar, env.vm.err = env.lookupBindVar(key)
+		if env.vm.err != nil {
+			return 0
+		}
+		return push_timestamp(env, bvar.Value)
+	}, "PUSH TIMESTAMP(:%q)", key)
+}
+
 func push_date(env *ExpressionEnv, raw []byte) int {
 	env.vm.stack[env.vm.sp], env.vm.err = parseDate(raw)
 	env.vm.sp++

go/vt/vtgate/evalengine/compiler_test.go

@@ -24,12 +24,12 @@ import (
 	"testing"
 	"time"
 
-	"github.com/stretchr/testify/assert"
-
 	"github.com/olekukonko/tablewriter"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"vitess.io/vitess/go/mysql/collations"
+	"vitess.io/vitess/go/mysql/collations/colldata"
 	"vitess.io/vitess/go/sqltypes"
 	querypb "vitess.io/vitess/go/vt/proto/query"
 	"vitess.io/vitess/go/vt/sqlparser"
@@ -119,7 +119,7 @@ func TestCompilerReference(t *testing.T) {
 			var supported, total int
 			env := evalengine.EmptyExpressionEnv(venv)
 
-			tc.Run(func(query string, row []sqltypes.Value) {
+			tc.Run(func(query string, row []sqltypes.Value, _ bool) {
 				env.Row = row
 				total++
 				testCompilerCase(t, query, venv, tc.Schema, env)
@@ -171,6 +171,7 @@ func testCompilerCase(t *testing.T, query string, venv *vtenv.Environment, schem
 		eval := expected.String()
 		comp := res.String()
 		assert.Equalf(t, eval, comp, "bad evaluation from compiler:\nSQL:  %s\nEval: %s\nComp: %s", query, eval, comp)
+		assert.Equalf(t, expected.Collation(), res.Collation(), "bad collation from compiler:\nSQL:  %s\nEval: %s\nComp: %s", query, colldata.Lookup(expected.Collation()).Name(), colldata.Lookup(res.Collation()).Name())
 	case vmErr == nil:
 		t.Errorf("failed evaluation from evalengine:\nSQL:  %s\nError: %s", query, evalErr)
 	case evalErr == nil: