sdks/go/pkg/beam/example_schema_test.go - beam - Git at Google

 // Licensed to the Apache Software Foundation (ASF) under one or more
 // contributor license agreements.  See the NOTICE file distributed with
 // this work for additional information regarding copyright ownership.
 // The ASF licenses this file to You under the Apache License, Version 2.0
 // (the "License"); you may not use this file except in compliance with
 // the License.  You may obtain a copy of the License at
 //
 //    http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 package beam_test

 import (
 	"bytes"
 	"fmt"
 	"io"
 	"reflect"

 	"github.com/apache/beam/sdks/go/pkg/beam"
 	"github.com/apache/beam/sdks/go/pkg/beam/core/graph/coder"
 	"github.com/google/go-cmp/cmp"
 )

 // RegisterSchemaProvider must be called before beam.Init, and conventionally in a package init block.
 func init() {
 	beam.RegisterSchemaProvider(reflect.TypeOf((*Alphabet)(nil)).Elem(), &AlphabetProvider{})
 	// TODO(BEAM-9615): Registerying a self encoding type causes a cycle. Needs resolving.
 	// beam.RegisterType(reflect.TypeOf((*Cyrillic)(nil)))
 	beam.RegisterType(reflect.TypeOf((*Latin)(nil)))
 	beam.RegisterType(reflect.TypeOf((*Ελληνικά)(nil)))
 }

 type Alphabet interface {
 	alphabet() string
 }

 type Cyrillic struct {
 	A, B int
 }

 func (*Cyrillic) alphabet() string {
 	return "Cyrillic"
 }

 type Latin struct {
 	// Unexported fields are not serializable by beam schemas by default
 	// so we need to handle this ourselves.
 	c uint64
 	d *float32
 }

 func (*Latin) alphabet() string {
 	return "Latin"
 }

 type Ελληνικά struct {
 	q string
 	G func() string
 }

 func (*Ελληνικά) alphabet() string {
 	return "Ελληνικά"
 }

 // AlphabetProvider provides encodings for types that implement the Alphabet interface.
 type AlphabetProvider struct {
 	enc *coder.RowEncoderBuilder
 	dec *coder.RowDecoderBuilder
 }

 var (
 	typeCyrillic = reflect.TypeOf((*Cyrillic)(nil))
 	typeLatin    = reflect.TypeOf((*Latin)(nil))
 	typeΕλληνικά = reflect.TypeOf((*Ελληνικά)(nil))
 )

 func (p *AlphabetProvider) FromLogicalType(rt reflect.Type) (reflect.Type, error) {
 	// FromLogicalType produces schema representative types, which match the encoders
 	// and decoders that this function generates for this type.
 	// While this example uses statically assigned schema representative types, it's
 	// possible to generate the returned reflect.Type dynamically instead, using the
 	// reflect package.
 	switch rt {
 	// The Cyrillic type is able to be encoded by default, so we simply use it directly
 	// as it's own representative type.
 	case typeCyrillic:
 		return typeCyrillic, nil
 	case typeLatin:
 		// The Latin type only has unexported fields, so we need to make the equivalent
 		// have exported fields.
 		return reflect.TypeOf((*struct {
 			C uint64
 			D *float32
 		})(nil)).Elem(), nil
 	case typeΕλληνικά:
 		return reflect.TypeOf((*struct{ Q string })(nil)).Elem(), nil
 	}
 	return nil, fmt.Errorf("Unknown Alphabet: %v", rt)
 }

 // BuildEncoder returns beam schema encoder functions for types with the Alphabet interface.
 func (p *AlphabetProvider) BuildEncoder(rt reflect.Type) (func(interface{}, io.Writer) error, error) {
 	switch rt {
 	case typeCyrillic:
 		if p.enc == nil {
 			p.enc = &coder.RowEncoderBuilder{}
 		}
 		// Since Cyrillic is by default encodable, defer to the standard schema row decoder for the type.
 		return p.enc.Build(rt)
 	case typeLatin:
 		return func(iface interface{}, w io.Writer) error {
 			v := iface.(*Latin)
 			// Beam Schema Rows have a header that indicates which fields if any, are nil.
 			if err := coder.WriteRowHeader(2, func(i int) bool {
 				if i == 1 {
 					return v.d == nil
 				}
 				return false
 			}, w); err != nil {
 				return err
 			}
 			// Afterwards, each field is encoded using the appropriate helper.
 			if err := coder.EncodeVarUint64(v.c, w); err != nil {
 				return err
 			}
 			// Nil fields have nothing written for them other than the header.
 			if v.d != nil {
 				if err := coder.EncodeDouble(float64(*v.d), w); err != nil {
 					return err
 				}
 			}
 			return nil
 		}, nil
 	case typeΕλληνικά:
 		return func(iface interface{}, w io.Writer) error {
 			// Since the representation for Ελληνικά never has nil fields
 			// we can use the simple header helper.
 			if err := coder.WriteSimpleRowHeader(1, w); err != nil {
 				return err
 			}
 			v := iface.(*Ελληνικά)
 			if err := coder.EncodeStringUTF8(v.q, w); err != nil {
 				return fmt.Errorf("decoding string field A: %v", err)
 			}
 			return nil
 		}, nil
 	}
 	return nil, fmt.Errorf("Unknown Alphabet: %v", rt)
 }

 // BuildDecoder returns beam schema decoder functions for types with the Alphabet interface.
 func (p *AlphabetProvider) BuildDecoder(rt reflect.Type) (func(io.Reader) (interface{}, error), error) {
 	switch rt {
 	case typeCyrillic:
 		if p.dec == nil {
 			p.dec = &coder.RowDecoderBuilder{}
 		}
 		// Since Cyrillic is by default encodable, defer to the standard schema row decoder for the type.
 		return p.dec.Build(rt)
 	case typeLatin:
 		return func(r io.Reader) (interface{}, error) {
 			// Since the d field can be nil, we use the header get the nil bits.
 			n, nils, err := coder.ReadRowHeader(r)
 			if err != nil {
 				return nil, err
 			}
 			// Header returns the number of fields, so we check if it has what we
 			// expect. This allows schemas to evolve if necessary.
 			if n != 2 {
 				return nil, fmt.Errorf("expected 2 fields, but got %v", n)
 			}
 			c, err := coder.DecodeVarUint64(r)
 			if err != nil {
 				return nil, err
 			}
 			// Check if the field is nil before trying to decode a value for it.
 			var d *float32
 			if !coder.IsFieldNil(nils, 1) {
 				f, err := coder.DecodeDouble(r)
 				if err != nil {
 					return nil, err
 				}
 				f32 := float32(f)
 				d = &f32
 			}
 			return &Latin{
 				c: c,
 				d: d,
 			}, nil
 		}, nil
 	case typeΕλληνικά:
 		return func(r io.Reader) (interface{}, error) {
 			// Since the representation for Ελληνικά never has nil fields
 			// we can use the simple header helper. Returns an error if
 			// something unexpected occurs.
 			if err := coder.ReadSimpleRowHeader(1, r); err != nil {
 				return nil, err
 			}
 			q, err := coder.DecodeStringUTF8(r)
 			if err != nil {
 				return nil, fmt.Errorf("decoding string field A: %v", err)
 			}
 			return &Ελληνικά{
 				q: q,
 			}, nil
 		}, nil
 	}
 	return nil, nil
 }

 // Schema providers work on fields of schema encoded types.
 type translation struct {
 	C *Cyrillic
 	L *Latin
 	E *Ελληνικά
 }

 func ExampleRegisterSchemaProvider() {
 	f := float32(42.789)
 	want := translation{
 		C: &Cyrillic{A: 123, B: 456},
 		L: &Latin{c: 789, d: &f},
 		E: &Ελληνικά{q: "testing"},
 	}
 	rt := reflect.TypeOf((*translation)(nil)).Elem()
 	enc, err := coder.RowEncoderForStruct(rt)
 	if err != nil {
 		panic(err)
 	}
 	dec, err := coder.RowDecoderForStruct(rt)
 	if err != nil {
 		panic(err)
 	}
 	var buf bytes.Buffer
 	if err := enc(want, &buf); err != nil {
 		panic(err)
 	}
 	got, err := dec(&buf)
 	if err != nil {
 		panic(err)
 	}
 	if d := cmp.Diff(want, got,
 		cmp.AllowUnexported(Latin{}, Ελληνικά{})); d != "" {
 		fmt.Printf("diff in schema encoding translation: (-want,+got)\n%v\n", d)
 	} else {
 		fmt.Println("No diffs!")
 	}
 	// Output: No diffs!
 }
	// Licensed to the Apache Software Foundation (ASF) under one or more
	// contributor license agreements. See the NOTICE file distributed with
	// this work for additional information regarding copyright ownership.
	// The ASF licenses this file to You under the Apache License, Version 2.0
	// (the "License"); you may not use this file except in compliance with
	// the License. You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	package beam_test

	import (
	"bytes"
	"fmt"
	"io"
	"reflect"

	"github.com/apache/beam/sdks/go/pkg/beam"
	"github.com/apache/beam/sdks/go/pkg/beam/core/graph/coder"
	"github.com/google/go-cmp/cmp"
	)

	// RegisterSchemaProvider must be called before beam.Init, and conventionally in a package init block.
	func init() {
	beam.RegisterSchemaProvider(reflect.TypeOf((*Alphabet)(nil)).Elem(), &AlphabetProvider{})
	// TODO(BEAM-9615): Registerying a self encoding type causes a cycle. Needs resolving.
	// beam.RegisterType(reflect.TypeOf((*Cyrillic)(nil)))
	beam.RegisterType(reflect.TypeOf((*Latin)(nil)))
	beam.RegisterType(reflect.TypeOf((*Ελληνικά)(nil)))
	}

	type Alphabet interface {
	alphabet() string
	}

	type Cyrillic struct {
	A, B int
	}

	func (*Cyrillic) alphabet() string {
	return "Cyrillic"
	}

	type Latin struct {
	// Unexported fields are not serializable by beam schemas by default
	// so we need to handle this ourselves.
	c uint64
	d *float32
	}

	func (*Latin) alphabet() string {
	return "Latin"
	}

	type Ελληνικά struct {
	q string
	G func() string
	}

	func (*Ελληνικά) alphabet() string {
	return "Ελληνικά"
	}

	// AlphabetProvider provides encodings for types that implement the Alphabet interface.
	type AlphabetProvider struct {
	enc *coder.RowEncoderBuilder
	dec *coder.RowDecoderBuilder
	}

	var (
	typeCyrillic = reflect.TypeOf((*Cyrillic)(nil))
	typeLatin = reflect.TypeOf((*Latin)(nil))
	typeΕλληνικά = reflect.TypeOf((*Ελληνικά)(nil))
	)

	func (p *AlphabetProvider) FromLogicalType(rt reflect.Type) (reflect.Type, error) {
	// FromLogicalType produces schema representative types, which match the encoders
	// and decoders that this function generates for this type.
	// While this example uses statically assigned schema representative types, it's
	// possible to generate the returned reflect.Type dynamically instead, using the
	// reflect package.
	switch rt {
	// The Cyrillic type is able to be encoded by default, so we simply use it directly
	// as it's own representative type.
	case typeCyrillic:
	return typeCyrillic, nil
	case typeLatin:
	// The Latin type only has unexported fields, so we need to make the equivalent
	// have exported fields.
	return reflect.TypeOf((*struct {
	C uint64
	D *float32
	})(nil)).Elem(), nil
	case typeΕλληνικά:
	return reflect.TypeOf((*struct{ Q string })(nil)).Elem(), nil
	}
	return nil, fmt.Errorf("Unknown Alphabet: %v", rt)
	}

	// BuildEncoder returns beam schema encoder functions for types with the Alphabet interface.
	func (p *AlphabetProvider) BuildEncoder(rt reflect.Type) (func(interface{}, io.Writer) error, error) {
	switch rt {
	case typeCyrillic:
	if p.enc == nil {
	p.enc = &coder.RowEncoderBuilder{}
	}
	// Since Cyrillic is by default encodable, defer to the standard schema row decoder for the type.
	return p.enc.Build(rt)
	case typeLatin:
	return func(iface interface{}, w io.Writer) error {
	v := iface.(*Latin)
	// Beam Schema Rows have a header that indicates which fields if any, are nil.
	if err := coder.WriteRowHeader(2, func(i int) bool {
	if i == 1 {
	return v.d == nil
	}
	return false
	}, w); err != nil {
	return err
	}
	// Afterwards, each field is encoded using the appropriate helper.
	if err := coder.EncodeVarUint64(v.c, w); err != nil {
	return err
	}
	// Nil fields have nothing written for them other than the header.
	if v.d != nil {
	if err := coder.EncodeDouble(float64(*v.d), w); err != nil {
	return err
	}
	}
	return nil
	}, nil
	case typeΕλληνικά:
	return func(iface interface{}, w io.Writer) error {
	// Since the representation for Ελληνικά never has nil fields
	// we can use the simple header helper.
	if err := coder.WriteSimpleRowHeader(1, w); err != nil {
	return err
	}
	v := iface.(*Ελληνικά)
	if err := coder.EncodeStringUTF8(v.q, w); err != nil {
	return fmt.Errorf("decoding string field A: %v", err)
	}
	return nil
	}, nil
	}
	return nil, fmt.Errorf("Unknown Alphabet: %v", rt)
	}

	// BuildDecoder returns beam schema decoder functions for types with the Alphabet interface.
	func (p *AlphabetProvider) BuildDecoder(rt reflect.Type) (func(io.Reader) (interface{}, error), error) {
	switch rt {
	case typeCyrillic:
	if p.dec == nil {
	p.dec = &coder.RowDecoderBuilder{}
	}
	// Since Cyrillic is by default encodable, defer to the standard schema row decoder for the type.
	return p.dec.Build(rt)
	case typeLatin:
	return func(r io.Reader) (interface{}, error) {
	// Since the d field can be nil, we use the header get the nil bits.
	n, nils, err := coder.ReadRowHeader(r)
	if err != nil {
	return nil, err
	}
	// Header returns the number of fields, so we check if it has what we
	// expect. This allows schemas to evolve if necessary.
	if n != 2 {
	return nil, fmt.Errorf("expected 2 fields, but got %v", n)
	}
	c, err := coder.DecodeVarUint64(r)
	if err != nil {
	return nil, err
	}
	// Check if the field is nil before trying to decode a value for it.
	var d *float32
	if !coder.IsFieldNil(nils, 1) {
	f, err := coder.DecodeDouble(r)
	if err != nil {
	return nil, err
	}
	f32 := float32(f)
	d = &f32
	}
	return &Latin{
	c: c,
	d: d,
	}, nil
	}, nil
	case typeΕλληνικά:
	return func(r io.Reader) (interface{}, error) {
	// Since the representation for Ελληνικά never has nil fields
	// we can use the simple header helper. Returns an error if
	// something unexpected occurs.
	if err := coder.ReadSimpleRowHeader(1, r); err != nil {
	return nil, err
	}
	q, err := coder.DecodeStringUTF8(r)
	if err != nil {
	return nil, fmt.Errorf("decoding string field A: %v", err)
	}
	return &Ελληνικά{
	q: q,
	}, nil
	}, nil
	}
	return nil, nil
	}

	// Schema providers work on fields of schema encoded types.
	type translation struct {
	C *Cyrillic
	L *Latin
	E *Ελληνικά
	}

	func ExampleRegisterSchemaProvider() {
	f := float32(42.789)
	want := translation{
	C: &Cyrillic{A: 123, B: 456},
	L: &Latin{c: 789, d: &f},
	E: &Ελληνικά{q: "testing"},
	}
	rt := reflect.TypeOf((*translation)(nil)).Elem()
	enc, err := coder.RowEncoderForStruct(rt)
	if err != nil {
	panic(err)
	}
	dec, err := coder.RowDecoderForStruct(rt)
	if err != nil {
	panic(err)
	}
	var buf bytes.Buffer
	if err := enc(want, &buf); err != nil {
	panic(err)
	}
	got, err := dec(&buf)
	if err != nil {
	panic(err)
	}
	if d := cmp.Diff(want, got,
	cmp.AllowUnexported(Latin{}, Ελληνικά{})); d != "" {
	fmt.Printf("diff in schema encoding translation: (-want,+got)\n%v\n", d)
	} else {
	fmt.Println("No diffs!")
	}
	// Output: No diffs!
	}