colexec,coldata: support timestamp type

This commit adds support for the TIMESTAMP type to the coldata and colexec
packages.

Release note: None

pkg/col/coldata/vec.go

@@ -12,6 +12,7 @@ package coldata
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/cockroachdb/apd"
 	"github.com/cockroachdb/cockroach/pkg/col/coltypes"
@@ -76,6 +77,8 @@ type Vec interface {
 	// TODO(jordan): should this be [][]byte?
 	// Decimal returns an apd.Decimal slice.
 	Decimal() []apd.Decimal
+	// Timestamp returns a time.Time slice.
+	Timestamp() []time.Time
 
 	// Col returns the raw, typeless backing storage for this Vec.
 	Col() interface{}
@@ -151,6 +154,8 @@ func NewMemColumn(t coltypes.T, n int) Vec {
 		return &memColumn{t: t, col: make([]float64, n), nulls: nulls}
 	case coltypes.Decimal:
 		return &memColumn{t: t, col: make([]apd.Decimal, n), nulls: nulls}
+	case coltypes.Timestamp:
+		return &memColumn{t: t, col: make([]time.Time, n), nulls: nulls}
 	default:
 		panic(fmt.Sprintf("unhandled type %s", t))
 	}
@@ -192,6 +197,10 @@ func (m *memColumn) Decimal() []apd.Decimal {
 	return m.col.([]apd.Decimal)
 }
 
+func (m *memColumn) Timestamp() []time.Time {
+	return m.col.([]time.Time)
+}
+
 func (m *memColumn) Col() interface{} {
 	return m.col
 }

pkg/col/coldata/vec_tmpl.go

@@ -21,6 +21,7 @@ package coldata
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/cockroachdb/apd"
 	"github.com/cockroachdb/cockroach/pkg/col/coltypes"
@@ -45,6 +46,9 @@ type _GOTYPESLICE interface{}
 // Dummy import to pull in "apd" package.
 var _ apd.Decimal
 
+// Dummy import to pull in "time" package.
+var _ time.Time
+
 // */}}
 
 func (m *memColumn) Append(args SliceArgs) {

pkg/col/colserde/arrowbatchconverter_test.go

@@ -30,8 +30,8 @@ func randomBatch() ([]coltypes.T, coldata.Batch) {
 
 	availableTyps := make([]coltypes.T, 0, len(coltypes.AllTypes))
 	for _, typ := range coltypes.AllTypes {
-		// TODO(asubiotto): We do not support decimal conversion yet.
-		if typ == coltypes.Decimal {
+		// TODO(asubiotto,jordan): We do not support decimal, timestamp conversion yet.
+		if typ == coltypes.Decimal || typ == coltypes.Timestamp {
 			continue
 		}
 		availableTyps = append(availableTyps, typ)
@@ -119,7 +119,7 @@ func assertEqualBatches(t *testing.T, expected, actual coldata.Batch) {
 func TestArrowBatchConverterRejectsUnsupportedTypes(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 
-	typs := []coltypes.T{coltypes.Decimal}
+	typs := []coltypes.T{coltypes.Decimal, coltypes.Timestamp}
 	_, err := NewArrowBatchConverter(typs)
 	require.Error(t, err)
 }

pkg/col/colserde/record_batch_test.go

@@ -180,9 +180,9 @@ func TestRecordBatchSerializerSerializeDeserializeRandom(t *testing.T) {
 		buf             = bytes.Buffer{}
 	)
 
-	// We do not support decimals yet.
+	// We do not support decimals or timestamps yet.
 	for _, t := range coltypes.AllTypes {
-		if t == coltypes.Decimal {
+		if t == coltypes.Decimal || t == coltypes.Timestamp {
 			continue
 		}
 		supportedTypes = append(supportedTypes, t)

pkg/col/coltypes/types.go

@@ -14,6 +14,7 @@ import (
 	"fmt"
 	"strings"
 	"text/template"
+	"time"
 
 	"github.com/cockroachdb/apd"
 )
@@ -38,6 +39,8 @@ const (
 	Int64
 	// Float64 is a column of type float64
 	Float64
+	// Timestamp is a column of type time.Time
+	Timestamp
 
 	// Unhandled is a temporary value that represents an unhandled type.
 	// TODO(jordan): this should be replaced by a panic once all types are
@@ -74,6 +77,7 @@ func init() {
 	CompatibleTypes[Int32] = append(CompatibleTypes[Int32], NumberTypes...)
 	CompatibleTypes[Int64] = append(CompatibleTypes[Int64], NumberTypes...)
 	CompatibleTypes[Float64] = append(CompatibleTypes[Float64], NumberTypes...)
+	CompatibleTypes[Timestamp] = append(CompatibleTypes[Timestamp], Timestamp)
 }
 
 // FromGoType returns the type for a Go value, if applicable. Shouldn't be used at
@@ -96,6 +100,8 @@ func FromGoType(v interface{}) T {
 		return Bytes
 	case apd.Decimal:
 		return Decimal
+	case time.Time:
+		return Timestamp
 	default:
 		panic(fmt.Sprintf("type %T not supported yet", t))
 	}
@@ -118,6 +124,8 @@ func (t T) GoTypeName() string {
 		return "int64"
 	case Float64:
 		return "float64"
+	case Timestamp:
+		return "time.Time"
 	default:
 		panic(fmt.Sprintf("unhandled type %d", t))
 	}

pkg/sql/colencoding/key_encoding.go

@@ -11,6 +11,8 @@
 package colencoding
 
 import (
+	"time"
+
 	"github.com/cockroachdb/apd"
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -225,6 +227,14 @@ func decodeTableKeyToCol(
 			rkey, t, err = encoding.DecodeVarintDescending(key)
 		}
 		vec.Int64()[idx] = t
+	case types.TimestampFamily:
+		var t time.Time
+		if dir == sqlbase.IndexDescriptor_ASC {
+			rkey, t, err = encoding.DecodeTimeAscending(key)
+		} else {
+			rkey, t, err = encoding.DecodeTimeDescending(key)
+		}
+		vec.Timestamp()[idx] = t
 	default:
 		return rkey, errors.AssertionFailedf("unsupported type %+v", log.Safe(valType))
 	}

pkg/sql/colencoding/value_encoding.go

@@ -11,6 +11,8 @@
 package colencoding
 
 import (
+	"time"
+
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/encoding"
@@ -91,6 +93,10 @@ func decodeUntaggedDatumToCol(vec coldata.Vec, idx uint16, t *types.T, buf []byt
 		if err == nil {
 			vec.Bytes().Set(int(idx), data.GetBytes())
 		}
+	case types.TimestampFamily:
+		var t time.Time
+		buf, t, err = encoding.DecodeUntaggedTimeValue(buf)
+		vec.Timestamp()[idx] = t
 	default:
 		return buf, errors.AssertionFailedf(
 			"couldn't decode type: %s", log.Safe(t))

pkg/sql/colexec/any_not_null_agg_tmpl.go

@@ -20,6 +20,8 @@
 package colexec
 
 import (
+	"time"
+
 	"github.com/cockroachdb/apd"
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
 	"github.com/cockroachdb/cockroach/pkg/col/coltypes"
@@ -45,6 +47,9 @@ func newAnyNotNullAgg(t coltypes.T) (aggregateFunc, error) {
 // Dummy import to pull in "apd" package.
 var _ apd.Decimal
 
+// Dummy import to pull in "time" package.
+var _ time.Time
+
 // _GOTYPESLICE is the template Go type slice variable for this operator. It
 // will be replaced by the Go slice representation for each type in coltypes.T, for
 // example []int64 for coltypes.Int64.

pkg/sql/colexec/const_tmpl.go

@@ -21,6 +21,7 @@ package colexec
 
 import (
 	"context"
+	"time"
 
 	"github.com/cockroachdb/apd"
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
@@ -36,6 +37,9 @@ import (
 // Dummy import to pull in "apd" package.
 var _ apd.Decimal
 
+// Dummy import to pull in "time" package.
+var _ time.Time
+
 // _TYPES_T is the template type variable for coltypes.T. It will be replaced by
 // coltypes.Foo for each type Foo in the coltypes.T type.
 const _TYPES_T = coltypes.Unhandled

pkg/sql/colexec/distinct_tmpl.go

@@ -23,6 +23,7 @@ import (
 	"bytes"
 	"context"
 	"math"
+	"time"
 
 	"github.com/cockroachdb/apd"
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
@@ -104,6 +105,9 @@ var _ bytes.Buffer
 // Dummy import to pull in "apd" package.
 var _ apd.Decimal
 
+// Dummy import to pull in "time" package.
+var _ time.Time
+
 // Dummy import to pull in "tree" package.
 var _ tree.Datum
 

pkg/sql/colexec/execgen/cmd/execgen/overloads.go

@@ -537,6 +537,9 @@ type floatIntCustomizer struct{}
 // side and a float right-hand side.
 type intFloatCustomizer struct{}
 
+// timestampCustomizer is necessary since time.Time doesn't have infix operators.
+type timestampCustomizer struct{}
+
 func (boolCustomizer) getCmpOpCompareFunc() compareFunc {
 	return func(target, l, r string) string {
 		args := map[string]string{"Target": target, "Left": l, "Right": r}
@@ -997,11 +1000,43 @@ func (c intFloatCustomizer) getCmpOpCompareFunc() compareFunc {
 	return getFloatCmpOpCompareFunc(false /* checkLeftNan */, true /* checkRightNan */)
 }
 
+func (c timestampCustomizer) getCmpOpCompareFunc() compareFunc {
+	return func(target, l, r string) string {
+		args := map[string]string{"Target": target, "Left": l, "Right": r}
+		buf := strings.Builder{}
+		// Inline the code from tree.compareTimestamps.
+		t := template.Must(template.New("").Parse(`
+      if {{.Left}}.Before({{.Right}}) {
+				{{.Target}} = -1
+			} else if {{.Right}}.Before({{.Left}}) {
+				{{.Target}} = 1
+			} else { 
+        {{.Target}} = 0
+      }`))
+
+		if err := t.Execute(&buf, args); err != nil {
+			execerror.VectorizedInternalPanic(err)
+		}
+		return buf.String()
+	}
+}
+
+func (timestampCustomizer) getHashAssignFunc() assignFunc {
+	return func(op overload, target, v, _ string) string {
+		return fmt.Sprintf(`
+		  s, ns := %[2]s.Second(), %[2]s.Nanosecond()
+		  %[1]s = memhash64(noescape(unsafe.Pointer(&s)), %[1]s)
+		  %[1]s = memhash64(noescape(unsafe.Pointer(&ns)), %[1]s)
+		`, target, v)
+	}
+}
+
 func registerTypeCustomizers() {
 	typeCustomizers = make(map[coltypePair]typeCustomizer)
 	registerTypeCustomizer(coltypePair{coltypes.Bool, coltypes.Bool}, boolCustomizer{})
 	registerTypeCustomizer(coltypePair{coltypes.Bytes, coltypes.Bytes}, bytesCustomizer{})
 	registerTypeCustomizer(coltypePair{coltypes.Decimal, coltypes.Decimal}, decimalCustomizer{})
+	registerTypeCustomizer(coltypePair{coltypes.Timestamp, coltypes.Timestamp}, timestampCustomizer{})
 	for _, leftFloatType := range coltypes.FloatTypes {
 		for _, rightFloatType := range coltypes.FloatTypes {
 			registerTypeCustomizer(coltypePair{leftFloatType, rightFloatType}, floatCustomizer{width: 64})

pkg/sql/colexec/execgen/cmd/execgen/overloads_test_utils_gen.go

@@ -21,6 +21,7 @@ package colexec
 import (
 	"bytes"
 	"math"
+	"time"
 
 	"github.com/cockroachdb/apd"
 	"github.com/cockroachdb/cockroach/pkg/sql/colexec/execerror"

pkg/sql/colexec/mem_estimation.go

@@ -55,6 +55,9 @@ func EstimateBatchSizeBytes(vecTypes []coltypes.T, batchLength int) int {
 			// Similar to byte arrays, we can't tell how much space is used
 			// to hold the arbitrary precision decimal objects.
 			acc += 50
+		case coltypes.Timestamp:
+			// time.Time has 2 int64s and a pointer.
+			acc += sizeOfInt64 * 3
 		default:
 			execerror.VectorizedInternalPanic(fmt.Sprintf("unhandled type %s", t))
 		}

pkg/sql/colexec/mergejoinbase_tmpl.go

@@ -23,6 +23,7 @@ import (
 	"bytes"
 	"fmt"
 	"math"
+	"time"
 
 	"github.com/cockroachdb/apd"
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
@@ -44,6 +45,9 @@ var _ tree.Datum
 // Dummy import to pull in "apd" package.
 var _ apd.Decimal
 
+// Dummy import to pull in "time" package.
+var _ time.Time
+
 // Dummy import to pull in "math" package.
 var _ = math.MaxInt64
 

pkg/sql/colexec/mergejoiner_tmpl.go

@@ -24,6 +24,7 @@ import (
 	"context"
 	"fmt"
 	"math"
+	"time"
 
 	"github.com/cockroachdb/apd"
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
@@ -46,6 +47,9 @@ var _ tree.Datum
 // Dummy import to pull in "apd" package.
 var _ apd.Decimal
 
+// Dummy import to pull in "time" package.
+var _ time.Time
+
 // Dummy import to pull in "math" package.
 var _ = math.MaxInt64
 

pkg/sql/colexec/min_max_agg_tmpl.go

@@ -22,6 +22,7 @@ package colexec
 import (
 	"bytes"
 	"math"
+	"time"
 
 	"github.com/cockroachdb/apd"
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
@@ -43,6 +44,9 @@ var _ bytes.Buffer
 // Dummy import to pull in "apd" package.
 var _ apd.Decimal
 
+// Dummy import to pull in "time" package.
+var _ time.Time
+
 // Dummy import to pull in "tree" package.
 var _ tree.Datum
 

pkg/sql/colexec/proj_const_ops_tmpl.go

@@ -23,6 +23,7 @@ import (
 	"bytes"
 	"context"
 	"math"
+	"time"
 
 	"github.com/cockroachdb/apd"
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
@@ -52,6 +53,9 @@ var _ tree.Datum
 // Dummy import to pull in "math" package.
 var _ = math.MaxInt64
 
+// Dummy import to pull in "time" package.
+var _ = time.Day
+
 // _ASSIGN is the template function for assigning the first input to the result
 // of computation an operation on the second and the third inputs.
 func _ASSIGN(_, _, _ interface{}) {