[Go SDK] Add portable logical type registration and standard date and micros_instant types

CHANGES.md

@@ -74,6 +74,7 @@
 * (Go) Pipeline-level `--resource_hints` are now forwarded to expansion services, so they apply to cross-language transforms as they do in the Java and Python SDKs. The `max_active_bundles_per_worker` hint name is now also accepted by `--resource_hints` ([#23893](https://github.com/apache/beam/issues/23893)).
 * ReadFromBigQuery now supports Lakehouse runtime catalog (BigLake metastore) tables with `method=DIRECT_READ`, using 4-part `project.catalog.namespace.table` identifiers. Previously `project:catalog.namespace.table` was silently mis-parsed and tables that report no `numBytes` failed to split (Python) ([#39597](https://github.com/apache/beam/issues/39597)).
 * (Python) Added `AwsOptions` with `--aws_*` pipeline options (e.g. `--aws_access_key_id`, `--aws_region_name`, `--aws_endpoint_url`) for configuring AWS service clients such as S3 ([#40528](https://github.com/apache/beam/issues/40528)).
+* (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` and `beam:logical_type:micros_instant:v1` 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
 

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

@@ -95,6 +95,23 @@ func (b *RowDecoderBuilder) decoderForType(t reflect.Type) (func(io.Reader) (any
 	if t.Kind() == reflect.Ptr {
 		isPtr = true
 		t = t.Elem()
+		// A provider for the value type also handles a top level pointer.
+		f, addr, err := b.customFunc(t)
+		if err != nil {
+			return nil, err
+		}
+		if f != nil && !addr {
+			vt := t
+			return func(r io.Reader) (any, error) {
+				v, err := f(r)
+				if err != nil {
+					return nil, err
+				}
+				rv := reflect.New(vt)
+				rv.Elem().Set(reflect.ValueOf(v))
+				return rv.Interface(), nil
+			}, nil
+		}
 	}
 	dec, err := b.decoderForStructReflect(t)
 	if err != nil {
@@ -435,3 +452,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

@@ -363,3 +363,24 @@ 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 !rv.IsValid() || (rv.Kind() == reflect.Ptr && rv.IsNil()) {
+			return errors.Errorf("cannot encode a nil value as a row field of type %v", t)
+		}
+		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

@@ -945,3 +945,108 @@ func TestReflectionRowCoderGeneration_Int16BigEndian(t *testing.T) {
 		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")
+	}
+	enc, err := RowFieldEncoderForType(reflect.TypeOf(int64(0)))
+	if err != nil {
+		t.Fatalf("RowFieldEncoderForType(int64) = %v, want nil error", err)
+	}
+	if err := enc(nil, &bytes.Buffer{}); err == nil {
+		t.Errorf("enc(nil) = nil error, want error")
+	}
+	ptrEnc, err := RowFieldEncoderForType(reflect.TypeOf((*int64)(nil)))
+	if err != nil {
+		t.Fatalf("RowFieldEncoderForType(*int64) = %v, want nil error", err)
+	}
+	if err := ptrEnc((*int64)(nil), &bytes.Buffer{}); err == nil {
+		t.Errorf("enc((*int64)(nil)) = nil error, want error")
+	}
+	if _, err := RowFieldDecoderForType(reflect.TypeOf(func() {})); err == nil {
+		t.Errorf("RowFieldDecoderForType(func) = nil error, want error")
+	}
+}
+
+// A provider for a value type also encodes and decodes a top level pointer to
+// that type.
+func TestRowCoder_TopLevelPointerWithValueProvider(t *testing.T) {
+	type tenths struct{ V int64 }
+	rt := reflect.TypeOf(tenths{})
+	var encB RowEncoderBuilder
+	encB.Register(rt, func(reflect.Type) (func(any, io.Writer) error, error) {
+		return func(v any, w io.Writer) error { return EncodeVarInt(v.(tenths).V*10, w) }, nil
+	})
+	var decB RowDecoderBuilder
+	decB.Register(rt, func(reflect.Type) (func(io.Reader) (any, error), error) {
+		return func(r io.Reader) (any, error) {
+			v, err := DecodeVarInt(r)
+			return tenths{V: v / 10}, err
+		}, nil
+	})
+	prt := reflect.PointerTo(rt)
+	enc, err := encB.Build(prt)
+	if err != nil {
+		t.Fatalf("RowEncoderBuilder.Build(%v) = %v, want nil error", prt, err)
+	}
+	dec, err := decB.Build(prt)
+	if err != nil {
+		t.Fatalf("RowDecoderBuilder.Build(%v) = %v, want nil error", prt, err)
+	}
+	want := &tenths{V: 42}
+	var buf bytes.Buffer
+	if err := enc(want, &buf); err != nil {
+		t.Fatalf("enc(%v) = %v, want nil error", want, err)
+	}
+	got, err := dec(&buf)
+	if err != nil {
+		t.Fatalf("dec = %v, want nil error", err)
+	}
+	if d := cmp.Diff(want, got); d != "" {
+		t.Errorf("dec(enc(%v)) diff (-want, +got): %v", want, d)
+	}
+}

sdks/go/pkg/beam/core/runtime/graphx/schema/logicaltypes.go

@@ -17,9 +17,11 @@ package schema
 
 import (
 	"fmt"
+	"io"
 	"reflect"
 	"sync"
 
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx"
 	pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
 )
@@ -224,4 +226,115 @@ func preRegLogicalTypes(r *Registry) {
 
 func init() {
 	preRegLogicalTypes(defaultRegistry)
+	registerStandardLogicalTypes(defaultRegistry)
+}
+
+// RegisterLogicalTypeConversion registers a portable logical type identified
+// by urn. Values of GoT are stored in schema rows as values of StorageT,
+// converted with toStorage and fromStorage, and encoded with the row field
+// encoding of StorageT. StorageT must have a schema representation: an atomic
+// type, a struct of such types, or a registered logical type. GoT must be a
+// concrete type. A struct kind GoT with a non row StorageT can be used as a
+// field of a row, but not as a top level schema type.
+//
+// RegisterLogicalTypeConversion must be called before beam.Init(), and
+// conventionally is called in a package init() function.
+func RegisterLogicalTypeConversion[GoT, StorageT any](urn string, toStorage func(GoT) (StorageT, error), fromStorage func(StorageT) (GoT, error)) {
+	registerLogicalTypeConversion(defaultRegistry, ToLogicalType(urn, typeOf[GoT](), typeOf[StorageT]()), toStorage, fromStorage)
+}
+
+// RegisterParameterizedLogicalTypeConversion is RegisterLogicalTypeConversion
+// for a logical type with an argument, such as a timestamp and its precision.
+// Each argument value is registered with its own Go type, and a schema field
+// with the urn and that argument maps to GoT. See ToLogicalTypeWithArgument
+// for the supported argument types.
+func RegisterParameterizedLogicalTypeConversion[GoT, StorageT any](urn string, argument any, toStorage func(GoT) (StorageT, error), fromStorage func(StorageT) (GoT, error)) {
+	registerLogicalTypeConversion(defaultRegistry, ToLogicalTypeWithArgument(urn, typeOf[GoT](), typeOf[StorageT](), argument), toStorage, fromStorage)
+}
+
+// RegisterLogicalTypeCoder registers a portable logical type identified by
+// urn whose wire encoding differs from the row field encoding of StorageT.
+// StorageT is only used as the schema representation of the logical type.
+// enc and dec produce and consume the encoded bytes of a value directly, in
+// the wire encoding that the urn defines. GoT must be a concrete type.
+// StorageT cannot be inferred, so both type parameters are given explicitly,
+// for example
+// RegisterLogicalTypeCoder[Money, []byte]("example:money:v1", encodeMoney, decodeMoney).
+//
+// RegisterLogicalTypeCoder must be called before beam.Init(), and
+// conventionally is called in a package init() function.
+func RegisterLogicalTypeCoder[GoT, StorageT any](urn string, enc func(GoT, io.Writer) error, dec func(io.Reader) (GoT, error)) {
+	registerLogicalTypeCoder(defaultRegistry, ToLogicalType(urn, typeOf[GoT](), typeOf[StorageT]()), enc, dec)
+}
+
+func typeOf[T any]() reflect.Type {
+	return reflect.TypeOf((*T)(nil)).Elem()
+}
+
+// registerLogicalTypeConversion registers lt with the registry and registers
+// row coder providers for its Go type, built from the row field coder of the
+// storage type and the conversion functions.
+func registerLogicalTypeConversion[GoT, StorageT any](r *Registry, lt LogicalType, toStorage func(GoT) (StorageT, error), fromStorage func(StorageT) (GoT, error)) {
+	st := lt.StorageType()
+	checkLogicalTypeKinds(lt)
+	r.RegisterLogicalType(lt)
+	coder.RegisterSchemaProviders(lt.GoType(),
+		func(reflect.Type) (func(any, io.Writer) error, error) {
+			enc, err := coder.RowFieldEncoderForType(st)
+			if err != nil {
+				return nil, err
+			}
+			return func(v any, w io.Writer) error {
+				s, err := toStorage(v.(GoT))
+				if err != nil {
+					return err
+				}
+				return enc(s, w)
+			}, nil
+		},
+		func(reflect.Type) (func(io.Reader) (any, error), error) {
+			dec, err := coder.RowFieldDecoderForType(st)
+			if err != nil {
+				return nil, err
+			}
+			return func(rd io.Reader) (any, error) {
+				s, err := dec(rd)
+				if err != nil {
+					return nil, err
+				}
+				return fromStorage(s.(StorageT))
+			}, nil
+		})
+}
+
+// registerLogicalTypeCoder registers lt with the registry and registers enc
+// and dec as the row coder providers for its Go type.
+func registerLogicalTypeCoder[GoT any](r *Registry, lt LogicalType, enc func(GoT, io.Writer) error, dec func(io.Reader) (GoT, error)) {
+	checkLogicalTypeKinds(lt)
+	r.RegisterLogicalType(lt)
+	coder.RegisterSchemaProviders(lt.GoType(),
+		func(reflect.Type) (func(any, io.Writer) error, error) {
+			return func(v any, w io.Writer) error {
+				return enc(v.(GoT), w)
+			}, nil
+		},
+		func(reflect.Type) (func(io.Reader) (any, error), error) {
+			return func(rd io.Reader) (any, error) {
+				return dec(rd)
+			}, nil
+		})
+}
+
+// checkLogicalTypeKinds panics when the Go type of lt is an interface or its
+// storage type is a pointer. The row coder matches an interface by all of its
+// implementations, and the schema matches only the interface itself. The row
+// coders decide nullability from the Go type of a field, so a pointer storage
+// type behind a non pointer Go type would fail at encoding time.
+func checkLogicalTypeKinds(lt LogicalType) {
+	if lt.GoType().Kind() == reflect.Interface {
+		panic(fmt.Sprintf("LogicalType[%v] has the interface Go type %v, Go types must be concrete", lt.ID(), lt.GoType()))
+	}
+	if lt.StorageType().Kind() == reflect.Ptr {
+		panic(fmt.Sprintf("LogicalType[%v] has the pointer storage type %v, storage types must not be pointers", lt.ID(), lt.StorageType()))
+	}
 }