// DatumToEncDatumWithInferredType initializes an EncDatum with the given
// Datum, setting its Type automatically. This does not work if the base
// Datum's type is DNull, in which case an error is returned.
func DatumToEncDatumWithInferredType(datum parser.Datum) (EncDatum, error) {
dType, ok := ColumnType_Kind_value[strings.ToUpper(datum.ResolvedType().String())]
if !ok {
return EncDatum{}, errors.Errorf(
"Unknown type %s, could not convert to EncDatum", datum.ResolvedType())
}
return DatumToEncDatum(ColumnType_Kind(dType), datum), nil
}
func TestEncDatumCompare(t *testing.T) {
defer leaktest.AfterTest(t)()
a := &DatumAlloc{}
rng, _ := randutil.NewPseudoRand()
for kind := range ColumnType_Kind_name {
kind := ColumnType_Kind(kind)
// TODO(cuongdo,eisen): we don't support persistence for arrays or collated
// strings yet
if kind == ColumnType_COLLATEDSTRING || kind == ColumnType_INT_ARRAY {
continue
}
typ := ColumnType{Kind: kind}
// Generate two datums d1 < d2
var d1, d2 parser.Datum
for {
d1 = RandDatum(rng, typ, false)
d2 = RandDatum(rng, typ, false)
if cmp := d1.Compare(d2); cmp < 0 {
break
}
}
v1 := DatumToEncDatum(typ, d1)
v2 := DatumToEncDatum(typ, d2)
if val, err := v1.Compare(a, &v2); err != nil {
t.Fatal(err)
} else if val != -1 {
t.Errorf("compare(1, 2) = %d", val)
}
asc := DatumEncoding_ASCENDING_KEY
desc := DatumEncoding_DESCENDING_KEY
noncmp := DatumEncoding_VALUE
checkEncDatumCmp(t, a, &v1, &v2, asc, asc, -1, false)
checkEncDatumCmp(t, a, &v2, &v1, asc, asc, +1, false)
checkEncDatumCmp(t, a, &v1, &v1, asc, asc, 0, false)
checkEncDatumCmp(t, a, &v2, &v2, asc, asc, 0, false)
checkEncDatumCmp(t, a, &v1, &v2, desc, desc, -1, false)
checkEncDatumCmp(t, a, &v2, &v1, desc, desc, +1, false)
checkEncDatumCmp(t, a, &v1, &v1, desc, desc, 0, false)
checkEncDatumCmp(t, a, &v2, &v2, desc, desc, 0, false)
checkEncDatumCmp(t, a, &v1, &v2, noncmp, noncmp, -1, true)
checkEncDatumCmp(t, a, &v2, &v1, desc, noncmp, +1, true)
checkEncDatumCmp(t, a, &v1, &v1, asc, desc, 0, true)
checkEncDatumCmp(t, a, &v2, &v2, desc, asc, 0, true)
}
}
// SetDatum initializes the EncDatum with the given Datum.
func (ed *EncDatum) SetDatum(typ ColumnType_Kind, d parser.Datum) {
if d == nil {
panic("nil datum given")
}
if d != parser.DNull && !typ.ToDatumType().Equal(d.ResolvedType()) {
panic(fmt.Sprintf("invalid datum type given: %s, expected %s",
d.ResolvedType(), typ.ToDatumType()))
}
ed.Type = typ
ed.encoded = nil
ed.Datum = d
}
// DatumToEncDatum initializes an EncDatum with the given Datum.
func DatumToEncDatum(ctyp ColumnType, d parser.Datum) EncDatum {
if d == nil {
panic("Cannot convert nil datum to EncDatum")
}
if ptyp := ctyp.ToDatumType(); d != parser.DNull && !ptyp.Equal(d.ResolvedType()) {
panic(fmt.Sprintf("invalid datum type given: %s, expected %s",
d.ResolvedType(), ptyp))
}
return EncDatum{
Type: ctyp,
Datum: d,
}
}
func TestEncDatumCompare(t *testing.T) {
a := &DatumAlloc{}
rng, _ := randutil.NewPseudoRand()
for typ := range ColumnType_Kind_name {
typ := ColumnType_Kind(typ)
// TODO(cuongdo): we don't support persistence for arrays yet
if typ == ColumnType_INT_ARRAY {
continue
}
// Generate two datums d1 < d2
var d1, d2 parser.Datum
for {
d1 = RandDatum(rng, typ, false)
d2 = RandDatum(rng, typ, false)
if cmp := d1.Compare(d2); cmp < 0 {
break
}
}
v1 := &EncDatum{}
v1.SetDatum(typ, d1)
v2 := &EncDatum{}
v2.SetDatum(typ, d2)
if val, err := v1.Compare(a, v2); err != nil {
t.Fatal(err)
} else if val != -1 {
t.Errorf("compare(1, 2) = %d", val)
}
asc := DatumEncoding_ASCENDING_KEY
desc := DatumEncoding_DESCENDING_KEY
noncmp := DatumEncoding_VALUE
checkEncDatumCmp(t, a, v1, v2, asc, asc, -1, false)
checkEncDatumCmp(t, a, v2, v1, asc, asc, +1, false)
checkEncDatumCmp(t, a, v1, v1, asc, asc, 0, false)
checkEncDatumCmp(t, a, v2, v2, asc, asc, 0, false)
checkEncDatumCmp(t, a, v1, v2, desc, desc, -1, false)
checkEncDatumCmp(t, a, v2, v1, desc, desc, +1, false)
checkEncDatumCmp(t, a, v1, v1, desc, desc, 0, false)
checkEncDatumCmp(t, a, v2, v2, desc, desc, 0, false)
checkEncDatumCmp(t, a, v1, v2, noncmp, noncmp, -1, true)
checkEncDatumCmp(t, a, v2, v1, desc, noncmp, +1, true)
checkEncDatumCmp(t, a, v1, v1, asc, desc, 0, true)
checkEncDatumCmp(t, a, v2, v2, desc, asc, 0, true)
}
}
func TestEncDatumCompare(t *testing.T) {
a := &DatumAlloc{}
rng, _ := randutil.NewPseudoRand()
for typ := ColumnType_Kind(0); int(typ) < len(ColumnType_Kind_value); typ++ {
// Generate two datums d1 < d2
var d1, d2 parser.Datum
for {
d1 = RandDatum(rng, typ, false)
d2 = RandDatum(rng, typ, false)
if cmp := d1.Compare(d2); cmp < 0 {
break
}
}
v1 := &EncDatum{}
v1.SetDatum(typ, d1)
v2 := &EncDatum{}
v2.SetDatum(typ, d2)
if val, err := v1.Compare(a, v2); err != nil {
t.Fatal(err)
} else if val != -1 {
t.Errorf("compare(1, 2) = %d", val)
}
asc := DatumEncoding_ASCENDING_KEY
desc := DatumEncoding_DESCENDING_KEY
noncmp := DatumEncoding_VALUE
checkEncDatumCmp(t, a, v1, v2, asc, asc, -1, false)
checkEncDatumCmp(t, a, v2, v1, asc, asc, +1, false)
checkEncDatumCmp(t, a, v1, v1, asc, asc, 0, false)
checkEncDatumCmp(t, a, v2, v2, asc, asc, 0, false)
checkEncDatumCmp(t, a, v1, v2, desc, desc, -1, false)
checkEncDatumCmp(t, a, v2, v1, desc, desc, +1, false)
checkEncDatumCmp(t, a, v1, v1, desc, desc, 0, false)
checkEncDatumCmp(t, a, v2, v2, desc, desc, 0, false)
checkEncDatumCmp(t, a, v1, v2, noncmp, noncmp, -1, true)
checkEncDatumCmp(t, a, v2, v1, desc, noncmp, +1, true)
checkEncDatumCmp(t, a, v1, v1, asc, desc, 0, true)
checkEncDatumCmp(t, a, v2, v2, desc, asc, 0, true)
}
}
// Encodes datum at the end of key, using direction `dir` for the encoding.
// It takes in an inclusive key and returns an inclusive key if
// isLastEndConstraint is not set, and an exclusive key otherwise (the idea is
// that, for inclusive constraints, the value for the last column in the
// constraint needs to be adapted to an exclusive span.EndKey).
func encodeInclusiveEndValue(
key roachpb.Key, datum parser.Datum, dir encoding.Direction, isLastEndConstraint bool,
) roachpb.Key {
// Since the end of a span is exclusive, if the last constraint is an
// inclusive one, we might need to make the key exclusive by applying a
// PrefixEnd(). We normally avoid doing this by transforming "a = x" to
// "a = x±1" for the last end constraint, depending on the encoding direction
// (since this keeps the key nice and pretty-printable).
// However, we might not be able to do the ±1.
needExclusiveKey := false
if isLastEndConstraint {
if dir == encoding.Ascending {
if datum.IsMax() {
needExclusiveKey = true
} else {
nextVal, hasNext := datum.Next()
if !hasNext {
needExclusiveKey = true
} else {
datum = nextVal
}
}
} else {
if datum.IsMin() {
needExclusiveKey = true
} else {
prevVal, hasPrev := datum.Prev()
if !hasPrev {
needExclusiveKey = true
} else {
datum = prevVal
}
}
}
}
key, err := sqlbase.EncodeTableKey(key, datum, dir)
if err != nil {
panic(err)
}
if needExclusiveKey {
key = key.PrefixEnd()
}
return key
}
// MarshalColumnValue returns a Go primitive value equivalent of val, of the
// type expected by col. If val's type is incompatible with col, or if
// col's type is not yet implemented, an error is returned.
func MarshalColumnValue(col ColumnDescriptor, val parser.Datum) (roachpb.Value, error) {
var r roachpb.Value
if val == parser.DNull {
return r, nil
}
switch col.Type.Kind {
case ColumnType_BOOL:
if v, ok := val.(*parser.DBool); ok {
r.SetBool(bool(*v))
return r, nil
}
case ColumnType_INT:
if v, ok := val.(*parser.DInt); ok {
r.SetInt(int64(*v))
return r, nil
}
case ColumnType_FLOAT:
if v, ok := val.(*parser.DFloat); ok {
r.SetFloat(float64(*v))
return r, nil
}
case ColumnType_DECIMAL:
if v, ok := val.(*parser.DDecimal); ok {
err := r.SetDecimal(&v.Dec)
return r, err
}
case ColumnType_STRING:
if v, ok := val.(*parser.DString); ok {
r.SetString(string(*v))
return r, nil
}
case ColumnType_BYTES:
if v, ok := val.(*parser.DBytes); ok {
r.SetString(string(*v))
return r, nil
}
if v, ok := val.(*parser.DString); ok {
r.SetString(string(*v))
return r, nil
}
case ColumnType_DATE:
if v, ok := val.(*parser.DDate); ok {
r.SetInt(int64(*v))
return r, nil
}
case ColumnType_TIMESTAMP:
if v, ok := val.(*parser.DTimestamp); ok {
r.SetTime(v.Time)
return r, nil
}
case ColumnType_TIMESTAMPTZ:
if v, ok := val.(*parser.DTimestampTZ); ok {
r.SetTime(v.Time)
return r, nil
}
case ColumnType_INTERVAL:
if v, ok := val.(*parser.DInterval); ok {
err := r.SetDuration(v.Duration)
return r, err
}
default:
return r, errors.Errorf("unsupported column type: %s", col.Type.Kind)
}
return r, fmt.Errorf("value type %s doesn't match type %s of column %q",
val.ResolvedType(), col.Type.Kind, col.Name)
}
func (b *writeBuffer) writeTextDatum(d parser.Datum, sessionLoc *time.Location) {
if log.V(2) {
log.Infof(context.TODO(), "pgwire writing TEXT datum of type: %T, %#v", d, d)
}
if d == parser.DNull {
// NULL is encoded as -1; all other values have a length prefix.
b.putInt32(-1)
return
}
switch v := d.(type) {
case *parser.DBool:
b.putInt32(1)
if *v {
b.writeByte('t')
} else {
b.writeByte('f')
}
case *parser.DInt:
// Start at offset 4 because `putInt32` clobbers the first 4 bytes.
s := strconv.AppendInt(b.putbuf[4:4], int64(*v), 10)
b.putInt32(int32(len(s)))
b.write(s)
case *parser.DFloat:
// Start at offset 4 because `putInt32` clobbers the first 4 bytes.
s := strconv.AppendFloat(b.putbuf[4:4], float64(*v), 'f', -1, 64)
b.putInt32(int32(len(s)))
b.write(s)
case *parser.DDecimal:
b.writeLengthPrefixedDatum(v)
case *parser.DBytes:
// http://www.postgresql.org/docs/current/static/datatype-binary.html#AEN5667
// Code cribbed from github.com/lib/pq.
result := make([]byte, 2+hex.EncodedLen(len(*v)))
result[0] = '\\'
result[1] = 'x'
hex.Encode(result[2:], []byte(*v))
b.putInt32(int32(len(result)))
b.write(result)
case *parser.DString:
b.writeLengthPrefixedString(string(*v))
case *parser.DCollatedString:
b.writeLengthPrefixedString(v.Contents)
case *parser.DDate:
t := time.Unix(int64(*v)*secondsInDay, 0)
// Start at offset 4 because `putInt32` clobbers the first 4 bytes.
s := formatTs(t, nil, b.putbuf[4:4])
b.putInt32(int32(len(s)))
b.write(s)
case *parser.DTimestamp:
// Start at offset 4 because `putInt32` clobbers the first 4 bytes.
s := formatTs(v.Time, nil, b.putbuf[4:4])
b.putInt32(int32(len(s)))
b.write(s)
case *parser.DTimestampTZ:
// Start at offset 4 because `putInt32` clobbers the first 4 bytes.
s := formatTs(v.Time, sessionLoc, b.putbuf[4:4])
b.putInt32(int32(len(s)))
b.write(s)
case *parser.DInterval:
b.writeLengthPrefixedString(v.ValueAsString())
case *parser.DTuple:
b.variablePutbuf.WriteString("(")
for i, d := range *v {
if i > 0 {
b.variablePutbuf.WriteString(",")
}
if d == parser.DNull {
// Emit nothing on NULL.
continue
}
d.Format(&b.variablePutbuf, parser.FmtSimple)
}
b.variablePutbuf.WriteString(")")
b.writeLengthPrefixedVariablePutbuf()
case *parser.DArray:
b.variablePutbuf.WriteString("{")
for i, d := range v.Array {
if i > 0 {
b.variablePutbuf.WriteString(",")
}
d.Format(&b.variablePutbuf, parser.FmtSimple)
}
b.variablePutbuf.WriteString("}")
b.writeLengthPrefixedVariablePutbuf()
default:
b.setError(errors.Errorf("unsupported type %T", d))
}
}
// writeLengthPrefixedDatum writes a length-prefixed Datum in its
// string representation. The length is encoded as an int32.
func (b *writeBuffer) writeLengthPrefixedDatum(d parser.Datum) {
d.Format(&b.variablePutbuf, parser.FmtSimple)
b.writeLengthPrefixedVariablePutbuf()
}
