tree: add extended vector friendly version of ParseAndRequireString

In order to build coldata.Vec's from string data a new version of
ParseAndRequireString is provided which pulls out the special types
supported by the vector engine and delegates to ParseAndRequireString
for anything else.

Informs: cockroachdb#91831

Release note: None

pkg/col/coldataext/BUILD.bazel

@@ -6,6 +6,7 @@ go_library(
     srcs = [
         "datum_vec.go",
         "extended_column_factory.go",
+        "vec_handler.go",
     ],
     importpath = "github.com/cockroachdb/cockroach/pkg/col/coldataext",
     visibility = ["//visibility:public"],
@@ -19,6 +20,10 @@ go_library(
         "//pkg/sql/sem/eval",
         "//pkg/sql/sem/tree",
         "//pkg/sql/types",
+        "//pkg/util/duration",
+        "//pkg/util/json",
+        "//pkg/util/timeutil/pgdate",
+        "@com_github_cockroachdb_apd_v3//:apd",
         "@com_github_cockroachdb_errors//:errors",
     ],
 )

pkg/sql/sem/tree/datum.go

@@ -1041,10 +1041,14 @@ func ParseDDecimal(s string) (*DDecimal, error) {
 // SetString sets d to s. Any non-standard NaN values are converted to a
 // normal NaN. Any negative zero is converted to positive.
 func (d *DDecimal) SetString(s string) error {
+	return setDecimalString(s, &d.Decimal)
+}
+
+func setDecimalString(s string, d *apd.Decimal) error {
 	// ExactCtx should be able to handle any decimal, but if there is any rounding
 	// or other inexact conversion, it will result in an error.
 	//_, res, err := HighPrecisionCtx.SetString(&d.Decimal, s)
-	_, res, err := ExactCtx.SetString(&d.Decimal, s)
+	_, res, err := ExactCtx.SetString(d, s)
 	if res != 0 || err != nil {
 		return MakeParseError(s, types.Decimal, err)
 	}
@@ -2854,13 +2858,20 @@ type DTimestampTZ struct {
 	time.Time
 }
 
-// MakeDTimestampTZ creates a DTimestampTZ with specified precision.
-func MakeDTimestampTZ(t time.Time, precision time.Duration) (*DTimestampTZ, error) {
+func checkTimeBounds(t time.Time, precision time.Duration) (time.Time, error) {
 	ret := t.Round(precision)
 	if ret.After(MaxSupportedTime) || ret.Before(MinSupportedTime) {
-		return nil, NewTimestampExceedsBoundsError(ret)
+		return time.Time{}, NewTimestampExceedsBoundsError(ret)
+	}
+	return ret, nil
+}
+
+// MakeDTimestampTZ creates a DTimestampTZ with specified precision.
+func MakeDTimestampTZ(t time.Time, precision time.Duration) (_ *DTimestampTZ, err error) {
+	if t, err = checkTimeBounds(t, precision); err != nil {
+		return nil, err
 	}
-	return &DTimestampTZ{Time: ret}, nil
+	return &DTimestampTZ{Time: t}, nil
 }
 
 // MustMakeDTimestampTZ wraps MakeDTimestampTZ but panics if there is an error.
@@ -2891,6 +2902,9 @@ func MakeDTimestampTZFromDate(loc *time.Location, d *DDate) (*DTimestampTZ, erro
 //
 // The dependsOnContext return value indicates if we had to consult the
 // ParseTimeContext (either for the time or the local timezone).
+//
+// Parts of this function are inlined into ParseAndRequireEx, if this changes materially
+// ParseAndRequireEx may need to change too.
 func ParseDTimestampTZ(
 	ctx ParseTimeContext, s string, precision time.Duration,
 ) (_ *DTimestampTZ, dependsOnContext bool, _ error) {

pkg/sql/sem/tree/parse_string.go

@@ -13,9 +13,17 @@ package tree
 import (
 	"strconv"
 	"strings"
+	"time"
 
+	"github.com/cockroachdb/apd/v3"
+	"github.com/cockroachdb/cockroach/pkg/sql/lex"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util"
+	"github.com/cockroachdb/cockroach/pkg/util/duration"
+	"github.com/cockroachdb/cockroach/pkg/util/json"
+	"github.com/cockroachdb/cockroach/pkg/util/timeutil/pgdate"
 	"github.com/cockroachdb/errors"
 	"github.com/lib/pq/oid"
 )
@@ -72,12 +80,7 @@ func ParseAndRequireString(
 			d, err = ParseDOidAsInt(s)
 		}
 	case types.StringFamily:
-		// If the string type specifies a limit we truncate to that limit:
-		//   'hello'::CHAR(2) -> 'he'
-		// This is true of all the string type variants.
-		if t.Width() > 0 {
-			s = util.TruncateString(s, int(t.Width()))
-		}
+		s = truncateString(s, t)
 		return NewDString(s), false, nil
 	case types.TimeFamily:
 		d, dependsOnContext, err = ParseDTime(ctx, s, TimeFamilyPrecisionToRoundDuration(t.Precision()))
@@ -113,6 +116,16 @@ func ParseAndRequireString(
 	return d, dependsOnContext, err
 }
 
+func truncateString(s string, t *types.T) string {
+	// If the string type specifies a limit we truncate to that limit:
+	//   'hello'::CHAR(2) -> 'he'
+	// This is true of all the string type variants.
+	if t.Width() > 0 {
+		s = util.TruncateString(s, int(t.Width()))
+	}
+	return s
+}
+
 // ParseDOidAsInt parses the input and returns it as an OID. If the input
 // is not formatted as an int, an error is returned.
 func ParseDOidAsInt(s string) (*DOid, error) {
@@ -141,3 +154,87 @@ func FormatBitArrayToType(d *DBitArray, t *types.T) *DBitArray {
 	}
 	return &DBitArray{a}
 }
+
+// ValueHandler is an interface to allow raw types to extracted from strings.
+type ValueHandler interface {
+	Date(d pgdate.Date)
+	Datum(d Datum)
+	Bool(b bool)
+	Bytes(b []byte)
+	Decimal() *apd.Decimal
+	Float(f float64)
+	Int(i int64)
+	Duration(d duration.Duration)
+	JSON(j json.JSON)
+	String(s string)
+	TimestampTZ(t time.Time)
+}
+
+func ParseAndRequireStringEx(t *types.T, s string, ctx ParseTimeContext, vh ValueHandler, ph *pgdate.ParseHelper) (err error) {
+	switch t.Family() {
+	case types.BoolFamily:
+		var b bool
+		if b, err = ParseBool(strings.TrimSpace(s)); err == nil {
+			vh.Bool(b)
+		}
+	case types.BytesFamily:
+		var res []byte
+		if res, err = lex.DecodeRawBytesToByteArrayAuto([]byte(s)); err != nil {
+			vh.Bytes(res)
+		} else {
+			err = MakeParseError(s, types.Bytes, err)
+		}
+	case types.DateFamily:
+		now := relativeParseTime(ctx)
+		var t pgdate.Date
+		if t, _, err = pgdate.ParseDate(now, dateStyle(ctx), s, ph); err == nil {
+			vh.Date(t)
+		}
+	case types.DecimalFamily:
+		dec := vh.Decimal()
+		if err = setDecimalString(s, dec); err != nil {
+			// Erase any invalid results.
+			*dec = apd.Decimal{}
+			err = MakeParseError(s, types.Decimal, err)
+		}
+	case types.FloatFamily:
+		var f float64
+		if f, err = strconv.ParseFloat(s, 64); err == nil {
+			vh.Float(f)
+		} else {
+			err = MakeParseError(s, types.Float, err)
+		}
+	case types.IntFamily:
+		var i int64
+		if i, err = strconv.ParseInt(s, 0, 64); err == nil {
+			vh.Int(i)
+		} else {
+			err = MakeParseError(s, types.Int, err)
+		}
+	case types.JsonFamily:
+		var j json.JSON
+		if j, err = json.ParseJSON(s); err == nil {
+			vh.JSON(j)
+		} else {
+			err = pgerror.Wrapf(err, pgcode.Syntax, "could not parse JSON")
+		}
+	case types.StringFamily:
+		s = truncateString(s, t)
+		vh.String(s)
+	case types.TimestampTZFamily:
+		now := relativeParseTime(ctx)
+		var ts time.Time
+		if ts, _, err = pgdate.ParseTimestamp(now, dateStyle(ctx), s); err == nil {
+			// Always normalize time to the current location.
+			if ts, err = checkTimeBounds(ts, TimeFamilyPrecisionToRoundDuration(t.Precision())); err == nil {
+				vh.TimestampTZ(ts)
+			}
+		}
+	default:
+		var d Datum
+		if d, _, err = ParseAndRequireString(t, s, ctx); err == nil {
+			vh.Datum(d)
+		}
+	}
+	return err
+}