exec,coldata: support timestamp type

This commit adds support for the TIMESTAMP type to the coldata and exec
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"
@@ -93,6 +94,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{}
@@ -172,6 +175,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))
 	}
@@ -221,6 +226,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 AppendArgs) {

pkg/col/colserde/arrowbatchconverter_test.go

@@ -31,7 +31,7 @@ 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 {
+		if typ == coltypes.Decimal || typ == coltypes.Timestamp {
 			continue
 		}
 		availableTyps = append(availableTyps, typ)
@@ -117,7 +117,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

@@ -196,7 +196,7 @@ func TestRecordBatchSerializerSerializeDeserializeRandom(t *testing.T) {
 
 	// We do not support decimals 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

@@ -12,6 +12,7 @@ package coltypes
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/cockroachdb/apd"
 )
@@ -40,6 +41,8 @@ const (
 	Float32
 	// 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
@@ -78,6 +81,7 @@ func init() {
 	CompatibleTypes[Int64] = append(CompatibleTypes[Int64], NumberTypes...)
 	CompatibleTypes[Float32] = append(CompatibleTypes[Float32], 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
@@ -104,6 +108,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))
 	}
@@ -130,6 +136,8 @@ func (t T) GoTypeName() string {
 		return "float32"
 	case Float64:
 		return "float64"
+	case Timestamp:
+		return "time.Time"
 	default:
 		panic(fmt.Sprintf("unhandled type %d", t))
 	}

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"
@@ -84,6 +86,10 @@ func decodeUntaggedDatumToCol(vec coldata.Vec, idx uint16, t *types.T, buf []byt
 			// We map these to 64-bit INT now. See #34161.
 			vec.Int64()[idx] = i
 		}
+	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/exec/any_not_null_agg_tmpl.go

@@ -20,6 +20,8 @@
 package exec
 
 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/exec/const_tmpl.go

@@ -21,6 +21,7 @@ package exec
 
 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/exec/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/exec/execgen/cmd/execgen/overloads.go

@@ -343,6 +343,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}
@@ -807,11 +810,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/exec/execgen/cmd/execgen/overloads_test_utils_gen.go

@@ -23,6 +23,7 @@ import (
 	"math"
 
 	"github.com/cockroachdb/apd"
+	"time"
 	"github.com/cockroachdb/cockroach/pkg/sql/exec/execerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 )

pkg/sql/exec/execgen/cmd/execgen/projection_ops_gen.go

@@ -21,9 +21,10 @@ const projTemplate = `
 package exec
 
 import (
-	"bytes"
+  "bytes"
   "context"
   "math"
+  "time"
 
 	"github.com/cockroachdb/apd"
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"

pkg/sql/exec/execgen/cmd/execgen/selection_ops_gen.go

@@ -21,9 +21,10 @@ const selTemplate = `
 package exec
 
 import (
-	"bytes"
+  "bytes"
   "context"
   "math"
+  "time"
 
 	"github.com/cockroachdb/apd"
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"

pkg/sql/exec/mem_estimation.go

@@ -61,6 +61,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/exec/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/exec/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/exec/min_max_agg_tmpl.go

@@ -22,6 +22,7 @@ package exec
 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/exec/random_testutils.go

@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
 	"github.com/cockroachdb/cockroach/pkg/col/coltypes"
 	"github.com/cockroachdb/cockroach/pkg/sql/exec/execerror"
+	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 )
 
 // maxVarLen specifies a length limit for variable length types (e.g. byte slices).
@@ -94,6 +95,11 @@ func RandomVec(rng *rand.Rand, typ coltypes.T, vec coldata.Vec, n int, nullProba
 		for i := 0; i < n; i++ {
 			floats[i] = rng.Float64()
 		}
+	case coltypes.Timestamp:
+		timestamps := vec.Timestamp()
+		for i := 0; i < n; i++ {
+			timestamps[i] = timeutil.Unix(rng.Int63n(1000000), rng.Int63n(1000000))
+		}
 	default:
 		execerror.VectorizedInternalPanic(fmt.Sprintf("unhandled type %s", typ))
 	}