[Go SDK] Add the standard timestamp logical type with millisecond, microsecond and nanosecond precision

CHANGES.md

@@ -69,10 +69,12 @@
 ## New Features / Improvements
 
 * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
+* (Go) Added `schema.RegisterLogicalTypeConversion` and `schema.RegisterLogicalTypeCoder` to register portable logical types in one call, and Go types for the standard `beam:logical_type:date:v1`, `beam:logical_type:micros_instant:v1`, `beam:logical_type:millis_instant:v1`, `beam:logical_type:decimal:v1`, `beam:logical_type:fixed_char:v1`, `beam:logical_type:var_char:v1`, `beam:logical_type:fixed_bytes:v1`, `beam:logical_type:var_bytes:v1` and `beam:logical_type:timestamp:v1` (precisions 3, 6 and 9) logical types, so schemas from other SDKs with these types can be used in cross-language pipelines ([#39684](https://github.com/apache/beam/issues/39684)).
 
 ## Breaking Changes
 
 * X behavior was changed ([#X](https://github.com/apache/beam/issues/X)).
+* (Go) The row coder now encodes `int16` and `uint16` struct fields as 2 byte big endian INT16 values, matching the Java and Python SDKs. Rows with such fields written by earlier Go SDK versions decode differently with this version ([#40151](https://github.com/apache/beam/issues/40151)).
 
 ## Deprecations
 

sdks/go/pkg/beam/core/graph/coder/int.go

@@ -69,3 +69,34 @@ func DecodeInt32(r io.Reader) (int32, error) {
 	}
 	return int32(ret), nil
 }
+
+// EncodeUint16 encodes an uint16 in big endian format.
+func EncodeUint16(value uint16, w io.Writer) error {
+	var data [2]byte
+	binary.BigEndian.PutUint16(data[:], value)
+	_, err := ioutilx.WriteUnsafe(w, data[:])
+	return err
+}
+
+// DecodeUint16 decodes an uint16 in big endian format.
+func DecodeUint16(r io.Reader) (uint16, error) {
+	var data [2]byte
+	if err := ioutilx.ReadNBufUnsafe(r, data[:]); err != nil {
+		return 0, err
+	}
+	return binary.BigEndian.Uint16(data[:]), nil
+}
+
+// EncodeInt16 encodes an int16 in big endian format.
+func EncodeInt16(value int16, w io.Writer) error {
+	return EncodeUint16(uint16(value), w)
+}
+
+// DecodeInt16 decodes an int16 in big endian format.
+func DecodeInt16(r io.Reader) (int16, error) {
+	ret, err := DecodeUint16(r)
+	if err != nil {
+		return 0, err
+	}
+	return int16(ret), nil
+}

sdks/go/pkg/beam/core/graph/coder/int_test.go

@@ -86,3 +86,35 @@ func TestEncodeDecodeInt32(t *testing.T) {
 		}
 	}
 }
+
+func TestEncodeDecodeInt16(t *testing.T) {
+	tests := []struct {
+		value int16
+		want  []byte
+	}{
+		{-32768, []byte{0x80, 0x00}},
+		{-2, []byte{0xff, 0xfe}},
+		{-1, []byte{0xff, 0xff}},
+		{0, []byte{0x00, 0x00}},
+		{1, []byte{0x00, 0x01}},
+		{999, []byte{0x03, 0xe7}},
+		{32767, []byte{0x7f, 0xff}},
+	}
+
+	for _, test := range tests {
+		var buf bytes.Buffer
+		if err := EncodeInt16(test.value, &buf); err != nil {
+			t.Fatalf("EncodeInt16(%v) failed: %v", test.value, err)
+		}
+		if got := buf.Bytes(); !bytes.Equal(got, test.want) {
+			t.Errorf("EncodeInt16(%v) = %v, want %v", test.value, got, test.want)
+		}
+		actual, err := DecodeInt16(&buf)
+		if err != nil {
+			t.Fatalf("DecodeInt16(<%v>) failed: %v", test.value, err)
+		}
+		if actual != test.value {
+			t.Errorf("DecodeInt16(<%v>) = %v, want %v", test.value, actual, test.value)
+		}
+	}
+}

sdks/go/pkg/beam/core/graph/coder/row.go

@@ -226,3 +226,19 @@ func ReadSimpleRowHeader(fields int, r io.Reader) error {
 	}
 	return nil
 }
+
+// RowFieldEncoderForType returns an encoding function for values of the given
+// type, using the encoding of that type as a field of a beam schema row.
+//
+// Returns an error if the given type cannot be a field of a beam schema row.
+func RowFieldEncoderForType(rt reflect.Type) (func(any, io.Writer) error, error) {
+	return defaultEnc.fieldEncoderForType(rt)
+}
+
+// RowFieldDecoderForType returns a decoding function for values of the given
+// type, using the encoding of that type as a field of a beam schema row.
+//
+// Returns an error if the given type cannot be a field of a beam schema row.
+func RowFieldDecoderForType(rt reflect.Type) (func(io.Reader) (any, error), error) {
+	return defaultDec.fieldDecoderForType(rt)
+}

sdks/go/pkg/beam/core/graph/coder/row_decoder.go

@@ -241,6 +241,24 @@ func reflectDecodeUint(rv reflect.Value, r io.Reader) error {
 	return nil
 }
 
+func reflectDecodeInt16(rv reflect.Value, r io.Reader) error {
+	v, err := DecodeInt16(r)
+	if err != nil {
+		return errors.Wrap(err, "error decoding int16 field")
+	}
+	rv.SetInt(int64(v))
+	return nil
+}
+
+func reflectDecodeUint16(rv reflect.Value, r io.Reader) error {
+	v, err := DecodeUint16(r)
+	if err != nil {
+		return errors.Wrap(err, "error decoding uint16 field")
+	}
+	rv.SetUint(uint64(v))
+	return nil
+}
+
 func reflectDecodeSinglePrecisionFloat(rv reflect.Value, r io.Reader) error {
 	v, err := DecodeSinglePrecisionFloat(r)
 	if err != nil {
@@ -341,10 +359,14 @@ func (b *RowDecoderBuilder) decoderForSingleTypeReflect(t reflect.Type) (typeDec
 		return typeDecoderFieldReflect{decode: reflectDecodeByte}, nil
 	case reflect.String:
 		return typeDecoderFieldReflect{decode: reflectDecodeString}, nil
-	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
+	case reflect.Int, reflect.Int8, reflect.Int32, reflect.Int64:
 		return typeDecoderFieldReflect{decode: reflectDecodeInt}, nil
-	case reflect.Uint, reflect.Uint64, reflect.Uint32, reflect.Uint16:
+	case reflect.Int16:
+		return typeDecoderFieldReflect{decode: reflectDecodeInt16}, nil
+	case reflect.Uint, reflect.Uint64, reflect.Uint32:
 		return typeDecoderFieldReflect{decode: reflectDecodeUint}, nil
+	case reflect.Uint16:
+		return typeDecoderFieldReflect{decode: reflectDecodeUint16}, nil
 	case reflect.Float32:
 		return typeDecoderFieldReflect{decode: reflectDecodeSinglePrecisionFloat}, nil
 	case reflect.Float64:
@@ -413,3 +435,23 @@ type typeDecoderFieldReflect struct {
 	// of the field value (i.e. &foo.bar) and not the field value (i.e. foo.bar).
 	addr bool
 }
+
+// fieldDecoderForType returns a decoder for values of t, using the encoding
+// of t as a row field.
+func (b *RowDecoderBuilder) fieldDecoderForType(t reflect.Type) (func(io.Reader) (any, error), error) {
+	decf, err := b.decoderForSingleTypeReflect(t)
+	if err != nil {
+		return nil, err
+	}
+	return func(r io.Reader) (any, error) {
+		rv := reflect.New(t)
+		v := rv.Elem()
+		if decf.addr {
+			v = rv
+		}
+		if err := decf.decode(v, r); err != nil {
+			return nil, err
+		}
+		return rv.Elem().Interface(), nil
+	}, nil
+}

sdks/go/pkg/beam/core/graph/coder/row_encoder.go

@@ -200,14 +200,24 @@ func (b *RowEncoderBuilder) encoderForSingleTypeReflect(t reflect.Type) (typeEnc
 		return typeEncoderFieldReflect{encode: func(rv reflect.Value, w io.Writer) error {
 			return EncodeStringUTF8(rv.String(), w)
 		}}, nil
-	case reflect.Int, reflect.Int64, reflect.Int16, reflect.Int32, reflect.Int8:
+	case reflect.Int, reflect.Int64, reflect.Int32, reflect.Int8:
 		return typeEncoderFieldReflect{encode: func(rv reflect.Value, w io.Writer) error {
 			return EncodeVarInt(rv.Int(), w)
 		}}, nil
-	case reflect.Uint, reflect.Uint64, reflect.Uint32, reflect.Uint16:
+	case reflect.Int16:
+		// INT16 schema fields are big endian, matching the Java and Python row coders.
+		return typeEncoderFieldReflect{encode: func(rv reflect.Value, w io.Writer) error {
+			return EncodeInt16(int16(rv.Int()), w)
+		}}, nil
+	case reflect.Uint, reflect.Uint64, reflect.Uint32:
 		return typeEncoderFieldReflect{encode: func(rv reflect.Value, w io.Writer) error {
 			return EncodeVarUint64(rv.Uint(), w)
 		}}, nil
+	case reflect.Uint16:
+		// uint16 is stored as an INT16 schema field.
+		return typeEncoderFieldReflect{encode: func(rv reflect.Value, w io.Writer) error {
+			return EncodeUint16(uint16(rv.Uint()), w)
+		}}, nil
 	case reflect.Float32:
 		return typeEncoderFieldReflect{encode: func(rv reflect.Value, w io.Writer) error {
 			return EncodeSinglePrecisionFloat(float32(rv.Float()), w)
@@ -353,3 +363,21 @@ func (b *RowEncoderBuilder) encoderForStructReflect(t reflect.Type) (func(reflec
 		return nil
 	}, nil
 }
+
+// fieldEncoderForType returns an encoder for values of t, using the encoding
+// of t as a row field.
+func (b *RowEncoderBuilder) fieldEncoderForType(t reflect.Type) (func(any, io.Writer) error, error) {
+	encf, err := b.encoderForSingleTypeReflect(t)
+	if err != nil {
+		return nil, err
+	}
+	return func(v any, w io.Writer) error {
+		rv := reflect.ValueOf(v)
+		if encf.addr {
+			p := reflect.New(t)
+			p.Elem().Set(rv)
+			rv = p
+		}
+		return encf.encode(rv, w)
+	}, nil
+}

sdks/go/pkg/beam/core/graph/coder/row_test.go

@@ -905,3 +905,96 @@ func TestRowHeader_TrailingZeroBytes(t *testing.T) {
 	}
 
 }
+
+// TestReflectionRowCoderGeneration_Int16BigEndian checks that int16 and uint16
+// fields use the 2 byte big endian encoding of the INT16 schema type, which is
+// what the Java and Python row coders produce.
+func TestReflectionRowCoderGeneration_Int16BigEndian(t *testing.T) {
+	type row struct {
+		A int16
+		B uint16
+		C int16
+	}
+	want := row{A: 999, B: 65535, C: -2}
+	wantBytes := []byte{
+		0x03, 0x00, // 3 fields, no nils.
+		0x03, 0xe7, // A
+		0xff, 0xff, // B
+		0xff, 0xfe, // C
+	}
+	rt := reflect.TypeOf(want)
+	enc, err := RowEncoderForStruct(rt)
+	if err != nil {
+		t.Fatalf("RowEncoderForStruct(%v) = %v, want nil error", rt, err)
+	}
+	var buf bytes.Buffer
+	if err := enc(want, &buf); err != nil {
+		t.Fatalf("enc(%v) = %v, want nil error", want, err)
+	}
+	if got := buf.Bytes(); !bytes.Equal(got, wantBytes) {
+		t.Fatalf("enc(%v) = %v, want %v", want, got, wantBytes)
+	}
+	dec, err := RowDecoderForStruct(rt)
+	if err != nil {
+		t.Fatalf("RowDecoderForStruct(%v) = %v, want nil error", rt, err)
+	}
+	got, err := dec(bytes.NewBuffer(wantBytes))
+	if err != nil {
+		t.Fatalf("dec(%v) = %v, want nil error", wantBytes, err)
+	}
+	if d := cmp.Diff(want, got); d != "" {
+		t.Fatalf("dec(enc(%v)) = %v\ndiff (-want, +got): %v", want, got, d)
+	}
+}
+
+func TestRowFieldCoderForType(t *testing.T) {
+	type pair struct {
+		A int64
+		B string
+	}
+	tests := []struct {
+		value     any
+		wantBytes []byte
+	}{
+		{value: int64(999), wantBytes: []byte{0xe7, 0x07}},
+		{value: int16(999), wantBytes: []byte{0x03, 0xe7}},
+		{value: "abc", wantBytes: []byte{0x03, 'a', 'b', 'c'}},
+		{value: []byte{1, 2}, wantBytes: []byte{0x02, 1, 2}},
+		{value: true, wantBytes: []byte{0x01}},
+		{value: []int32{1, 2}, wantBytes: []byte{0x00, 0x00, 0x00, 0x02, 0x01, 0x02}},
+		{value: pair{A: 7, B: "x"}, wantBytes: []byte{0x02, 0x00, 0x07, 0x01, 'x'}},
+	}
+	for _, test := range tests {
+		rt := reflect.TypeOf(test.value)
+		t.Run(rt.String(), func(t *testing.T) {
+			enc, err := RowFieldEncoderForType(rt)
+			if err != nil {
+				t.Fatalf("RowFieldEncoderForType(%v) = %v, want nil error", rt, err)
+			}
+			var buf bytes.Buffer
+			if err := enc(test.value, &buf); err != nil {
+				t.Fatalf("enc(%v) = %v, want nil error", test.value, err)
+			}
+			if got := buf.Bytes(); !bytes.Equal(got, test.wantBytes) {
+				t.Fatalf("enc(%v) = %v, want %v", test.value, got, test.wantBytes)
+			}
+			dec, err := RowFieldDecoderForType(rt)
+			if err != nil {
+				t.Fatalf("RowFieldDecoderForType(%v) = %v, want nil error", rt, err)
+			}
+			got, err := dec(bytes.NewBuffer(test.wantBytes))
+			if err != nil {
+				t.Fatalf("dec(%v) = %v, want nil error", test.wantBytes, err)
+			}
+			if d := cmp.Diff(test.value, got); d != "" {
+				t.Fatalf("dec(enc(%v)) = %v\ndiff (-want, +got): %v", test.value, got, d)
+			}
+		})
+	}
+	if _, err := RowFieldEncoderForType(reflect.TypeOf(func() {})); err == nil {
+		t.Errorf("RowFieldEncoderForType(func) = nil error, want error")
+	}
+	if _, err := RowFieldDecoderForType(reflect.TypeOf(func() {})); err == nil {
+		t.Errorf("RowFieldDecoderForType(func) = nil error, want error")
+	}
+}