Make user side coders more convenient.

Author: lostluck

sdks/go/pkg/beam/coder.go

@@ -18,10 +18,12 @@ package beam
 import (
 	"encoding/json"
 	"fmt"
+	"io"
 	"reflect"
 
 	"github.com/apache/beam/sdks/go/pkg/beam/core/graph/coder"
 	"github.com/apache/beam/sdks/go/pkg/beam/core/runtime/coderx"
+	"github.com/apache/beam/sdks/go/pkg/beam/core/runtime/exec"
 	"github.com/apache/beam/sdks/go/pkg/beam/core/typex"
 	"github.com/apache/beam/sdks/go/pkg/beam/core/util/reflectx"
 	"github.com/golang/protobuf/proto"
@@ -68,20 +70,56 @@ func (c Coder) String() string {
 	return c.coder.String()
 }
 
-// TODO(herohde) 4/4/2017: for convenience, we use the magic json coding
-// everywhere. To be replaced by Coder registry, sharing, etc.
+// NewElementEncoder returns a new encoding function for the given type.
+func NewElementEncoder(t reflect.Type) ElementEncoder {
+	c, err := inferCoder(typex.New(t))
+	if err != nil {
+		panic(err)
+	}
+	return &execEncoder{enc: exec.MakeElementEncoder(c)}
+}
 
-// TODO: select optimal coder based on type, notably handling int, string, etc.
+// execEncoder wraps an exec.ElementEncoder to implement the ElementDecoder interface
+// in this package.
+type execEncoder struct {
+	enc   exec.ElementEncoder
+	coder *coder.Coder
+}
 
-// TODO(herohde) 7/11/2017: figure out best way to let transformation use
-// coders (like passert). For now, we just allow them to grab in the internal
-// coder. Maybe it's cleaner to pull Encode/Decode into beam instead, if
-// adequate. The issue is that we would need non-windowed coding. Maybe focus on
-// coder registry and construction: then type -> coder might be adequate.
+func (e *execEncoder) Encode(element interface{}, w io.Writer) error {
+	return e.enc.Encode(exec.FullValue{Elm: element}, w)
+}
 
-// UnwrapCoder returns the internal coder.
-func UnwrapCoder(c Coder) *coder.Coder {
-	return c.coder
+func (e *execEncoder) String() string {
+	return e.coder.String()
+}
+
+// NewElementDecoder returns an ElementDecoder the given type.
+func NewElementDecoder(t reflect.Type) ElementDecoder {
+	c, err := inferCoder(typex.New(t))
+	if err != nil {
+		panic(err)
+	}
+	return &execDecoder{dec: exec.MakeElementDecoder(c)}
+}
+
+// execDecoder wraps an exec.ElementDecoder to implement the ElementDecoder interface
+// in this package.
+type execDecoder struct {
+	dec   exec.ElementDecoder
+	coder *coder.Coder
+}
+
+func (d *execDecoder) Decode(r io.Reader) (interface{}, error) {
+	fv, err := d.dec.Decode(r)
+	if err != nil {
+		return nil, err
+	}
+	return fv.Elm, nil
+}
+
+func (d *execDecoder) String() string {
+	return d.coder.String()
 }
 
 // NewCoder infers a Coder for any bound full type.
@@ -179,11 +217,6 @@ func inferCoders(list []FullType) ([]*coder.Coder, error) {
 	return ret, nil
 }
 
-// TODO(herohde) 4/5/2017: decide whether we want an Encoded form. For now,
-// we'll use exploded form coders only using typex.T. We might also need a
-// form that doesn't require LengthPrefix'ing to cut up the bytestream from
-// the FnHarness.
-
 // protoEnc marshals the supplied proto.Message.
 func protoEnc(in T) ([]byte, error) {
 	return proto.Marshal(in.(proto.Message))

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

@@ -19,6 +19,7 @@ package coder
 
 import (
 	"fmt"
+	"io"
 	"reflect"
 	"strings"
 
@@ -92,6 +93,16 @@ var (
 		OptReturn: []reflect.Type{reflectx.Error}}
 )
 
+// ElementEncoder encapsulates being able to encode an element into a writer.
+type ElementEncoder interface {
+	Encode(element interface{}, w io.Writer) error
+}
+
+// ElementDecoder encapsulates being able to decode an element from a reader.
+type ElementDecoder interface {
+	Decode(r io.Reader) (interface{}, error)
+}
+
 func validateEncoder(t reflect.Type, encode interface{}) error {
 	// Check if it uses the real type in question.
 	if err := funcx.Satisfy(encode, funcx.Replace(encodeSig, typex.TType, t)); err != nil {

sdks/go/pkg/beam/create.go

@@ -19,13 +19,8 @@ import (
 	"bytes"
 	"fmt"
 	"reflect"
-
-	"github.com/apache/beam/sdks/go/pkg/beam/core/runtime/exec"
-	"github.com/apache/beam/sdks/go/pkg/beam/core/typex"
 )
 
-// TODO(herohde) 7/11/2017: add variants that use coder encoding.
-
 // Create inserts a fixed set of values into the pipeline. The values must
 // be of the same type 'A' and the returned PCollection is of type A.
 //
@@ -58,15 +53,15 @@ func TryCreate(s Scope, values ...interface{}) (PCollection, error) {
 	}
 
 	t := reflect.ValueOf(values[0]).Type()
-	coder := NewCoder(typex.New(t))
-	fn := &createFn{Coder: EncodedCoder{Coder: coder}}
-	en := exec.MakeElementEncoder(UnwrapCoder(coder))
+	fn := &createFn{Type: EncodedType{T: t}}
+	enc := NewElementEncoder(t)
+
 	for i, value := range values {
 		if other := reflect.ValueOf(value).Type(); other != t {
 			return PCollection{}, fmt.Errorf("value %v at index %v has type %v, want %v", value, i, other, t)
 		}
 		var buf bytes.Buffer
-		if err := en.Encode(exec.FullValue{Elm: value}, &buf); err != nil {
+		if err := enc.Encode(value, &buf); err != nil {
 			return PCollection{}, fmt.Errorf("marshalling of %v failed: %v", value, err)
 		}
 		fn.Values = append(fn.Values, buf.Bytes())
@@ -84,18 +79,18 @@ func TryCreate(s Scope, values ...interface{}) (PCollection, error) {
 // TODO(herohde) 6/26/2017: make 'create' a SDF once supported. See BEAM-2421.
 
 type createFn struct {
-	Values [][]byte     `json:"values"`
-	Coder  EncodedCoder `json:"coder"`
+	Values [][]byte    `json:"values"`
+	Type   EncodedType `json:"type"`
 }
 
 func (c *createFn) ProcessElement(_ []byte, emit func(T)) error {
-	dec := exec.MakeElementDecoder(UnwrapCoder(c.Coder.Coder))
+	dec := NewElementDecoder(c.Type.T)
 	for _, val := range c.Values {
-		fv, err := dec.Decode(bytes.NewBuffer(val))
+		element, err := dec.Decode(bytes.NewBuffer(val))
 		if err != nil {
 			return err
 		}
-		emit(fv.Elm)
+		emit(element)
 	}
 	return nil
 }

sdks/go/pkg/beam/forward.go

@@ -92,14 +92,19 @@ func RegisterInit(hook func()) {
 //  func(reflect.Type, []byte) (T, error)
 //
 // where T is the matching user type.
-//
 func RegisterCoder(t reflect.Type, encoder, decoder interface{}) {
 	runtime.RegisterType(t)
 	runtime.RegisterFunction(encoder)
 	runtime.RegisterFunction(decoder)
 	coder.RegisterCoder(t, encoder, decoder)
 }
 
+// ElementEncoder encapsulates being able to encode an element into a writer.
+type ElementEncoder = coder.ElementEncoder
+
+// ElementDecoder encapsulates being able to decode an element from a reader.
+type ElementDecoder = coder.ElementDecoder
+
 // Init is the hook that all user code must call after flags processing and
 // other static initialization, for now.
 func Init() {

sdks/go/pkg/beam/testing/passert/passert.go

@@ -23,8 +23,6 @@ import (
 	"fmt"
 
 	"github.com/apache/beam/sdks/go/pkg/beam"
-	"github.com/apache/beam/sdks/go/pkg/beam/core/graph/coder"
-	"github.com/apache/beam/sdks/go/pkg/beam/core/runtime/exec"
 	"github.com/apache/beam/sdks/go/pkg/beam/core/typex"
 	"github.com/apache/beam/sdks/go/pkg/beam/transforms/filter"
 )
@@ -61,26 +59,25 @@ func equals(s beam.Scope, actual, expected beam.PCollection) beam.PCollection {
 // because all values are held in memory at the same time.
 func Diff(s beam.Scope, a, b beam.PCollection) (left, both, right beam.PCollection) {
 	imp := beam.Impulse(s)
-	return beam.ParDo3(s, &diffFn{Coder: beam.EncodedCoder{Coder: a.Coder()}}, imp, beam.SideInput{Input: a}, beam.SideInput{Input: b})
-}
 
-// TODO(herohde) 7/11/2017: should there be a first-class way to obtain the coder,
-// such a a specially-typed parameter?
+	t := beam.ValidateNonCompositeType(a)
+	beam.ValidateNonCompositeType(b)
+	return beam.ParDo3(s, &diffFn{Type: beam.EncodedType{T: t.Type()}}, imp, beam.SideInput{Input: a}, beam.SideInput{Input: b})
+}
 
 // diffFn computes the symmetrical multi-set difference of 2 collections, under
 // coder equality. The Go values returned may be any of the coder-equal ones.
 type diffFn struct {
-	Coder beam.EncodedCoder `json:"coder"`
+	Type beam.EncodedType `json:"type"`
 }
 
 func (f *diffFn) ProcessElement(_ []byte, ls, rs func(*beam.T) bool, left, both, right func(t beam.T)) error {
-	c := beam.UnwrapCoder(f.Coder.Coder)
-
-	indexL, err := index(c, ls)
+	enc := beam.NewElementEncoder(f.Type.T)
+	indexL, err := index(enc, ls)
 	if err != nil {
 		return err
 	}
-	indexR, err := index(c, rs)
+	indexR, err := index(enc, rs)
 	if err != nil {
 		return err
 	}
@@ -130,15 +127,14 @@ type indexEntry struct {
 	value beam.T
 }
 
-func index(c *coder.Coder, iter func(*beam.T) bool) (map[string]indexEntry, error) {
+func index(enc beam.ElementEncoder, iter func(*beam.T) bool) (map[string]indexEntry, error) {
 	ret := make(map[string]indexEntry)
-	enc := exec.MakeElementEncoder(c)
 
 	var val beam.T
 	for iter(&val) {
 		var buf bytes.Buffer
-		if err := enc.Encode(exec.FullValue{Elm: val}, &buf); err != nil {
-			return nil, fmt.Errorf("value %v not encodable by %v", val, c)
+		if err := enc.Encode(val, &buf); err != nil {
+			return nil, fmt.Errorf("value %v not encodable with %v", val, enc)
 		}
 		encoded := buf.String()
 
@@ -148,9 +144,6 @@ func index(c *coder.Coder, iter func(*beam.T) bool) (map[string]indexEntry, erro
 	return ret, nil
 }
 
-// TODO(herohde) 7/11/2017: perhaps extract the coder helpers as more
-// general and polished utilities for working with coders in user code.
-
 // True asserts that all elements satisfy the given predicate.
 func True(s beam.Scope, col beam.PCollection, fn interface{}) beam.PCollection {
 	fail(s, filter.Exclude(s, col, fn), "predicate(%v) = false, want true")