You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by al...@apache.org on 2019/02/13 22:46:32 UTC
[beam] branch master updated: Make user side coders more convenient.
This is an automated email from the ASF dual-hosted git repository.
altay pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new ad05c34 Make user side coders more convenient.
new b41c6b8 Merge pull request #7834 from lostluck/surface
ad05c34 is described below
commit ad05c34c0e2b3abb903e4e85dbd86efc6ceaa33d
Author: Robert Burke <ro...@frantil.com>
AuthorDate: Wed Feb 13 21:44:12 2019 +0000
Make user side coders more convenient.
---
sdks/go/pkg/beam/coder.go | 65 ++++++++++++++++++++++-------
sdks/go/pkg/beam/core/graph/coder/coder.go | 11 +++++
sdks/go/pkg/beam/create.go | 23 ++++------
sdks/go/pkg/beam/forward.go | 7 +++-
sdks/go/pkg/beam/testing/passert/passert.go | 29 +++++--------
sdks/go/pkg/beam/transforms/top/top.go | 58 ++++++++++++++-----------
sdks/go/pkg/beam/transforms/top/top_test.go | 8 ++--
7 files changed, 123 insertions(+), 78 deletions(-)
diff --git a/sdks/go/pkg/beam/coder.go b/sdks/go/pkg/beam/coder.go
index 83b934f..07c8e56 100644
--- a/sdks/go/pkg/beam/coder.go
+++ b/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))
diff --git a/sdks/go/pkg/beam/core/graph/coder/coder.go b/sdks/go/pkg/beam/core/graph/coder/coder.go
index 7432c10..a28f8ca 100644
--- a/sdks/go/pkg/beam/core/graph/coder/coder.go
+++ b/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 {
diff --git a/sdks/go/pkg/beam/create.go b/sdks/go/pkg/beam/create.go
index 25ea19b..93c11d0 100644
--- a/sdks/go/pkg/beam/create.go
+++ b/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
}
diff --git a/sdks/go/pkg/beam/forward.go b/sdks/go/pkg/beam/forward.go
index af722d5..37ffefc 100644
--- a/sdks/go/pkg/beam/forward.go
+++ b/sdks/go/pkg/beam/forward.go
@@ -92,7 +92,6 @@ 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)
@@ -100,6 +99,12 @@ func RegisterCoder(t reflect.Type, encoder, decoder interface{}) {
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() {
diff --git a/sdks/go/pkg/beam/testing/passert/passert.go b/sdks/go/pkg/beam/testing/passert/passert.go
index 0a00074..c267960 100644
--- a/sdks/go/pkg/beam/testing/passert/passert.go
+++ b/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")
diff --git a/sdks/go/pkg/beam/transforms/top/top.go b/sdks/go/pkg/beam/transforms/top/top.go
index 8b7edfd..a7a0208 100644
--- a/sdks/go/pkg/beam/transforms/top/top.go
+++ b/sdks/go/pkg/beam/transforms/top/top.go
@@ -21,11 +21,11 @@ import (
"bytes"
"encoding/json"
"fmt"
+ "reflect"
"sort"
"github.com/apache/beam/sdks/go/pkg/beam"
"github.com/apache/beam/sdks/go/pkg/beam/core/funcx"
- "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"
)
@@ -52,7 +52,7 @@ func Largest(s beam.Scope, col beam.PCollection, n int, less interface{}) beam.P
t := beam.ValidateNonCompositeType(col)
validate(t, n, less)
- return beam.Combine(s, newCombineFn(less, n, t, false), col)
+ return beam.Combine(s, newCombineFn(less, n, t.Type(), false), col)
}
// LargestPerKey returns the largest N values for each key of a PCollection<KV<K,T>>.
@@ -65,7 +65,7 @@ func LargestPerKey(s beam.Scope, col beam.PCollection, n int, less interface{})
_, t := beam.ValidateKVType(col)
validate(t, n, less)
- return beam.CombinePerKey(s, newCombineFn(less, n, t, false), col)
+ return beam.CombinePerKey(s, newCombineFn(less, n, t.Type(), false), col)
}
// Smallest returns the smallest N elements of a PCollection<T>. The order is
@@ -83,7 +83,7 @@ func Smallest(s beam.Scope, col beam.PCollection, n int, less interface{}) beam.
t := beam.ValidateNonCompositeType(col)
validate(t, n, less)
- return beam.Combine(s, newCombineFn(less, n, t, true), col)
+ return beam.Combine(s, newCombineFn(less, n, t.Type(), true), col)
}
// SmallestPerKey returns the smallest N values for each key of a PCollection<KV<K,T>>.
@@ -96,7 +96,7 @@ func SmallestPerKey(s beam.Scope, col beam.PCollection, n int, less interface{})
_, t := beam.ValidateKVType(col)
validate(t, n, less)
- return beam.Combine(s, newCombineFn(less, n, t, true), col)
+ return beam.Combine(s, newCombineFn(less, n, t.Type(), true), col)
}
func validate(t typex.FullType, n int, less interface{}) {
@@ -106,16 +106,21 @@ func validate(t typex.FullType, n int, less interface{}) {
funcx.MustSatisfy(less, funcx.Replace(sig, beam.TType, t.Type()))
}
-func newCombineFn(less interface{}, n int, t typex.FullType, reversed bool) *combineFn {
- coder := beam.NewCoder(t)
- return &combineFn{Less: beam.EncodedFunc{Fn: reflectx.MakeFunc(less)}, N: n, Coder: beam.EncodedCoder{Coder: coder}, Reversed: reversed}
+func newCombineFn(less interface{}, n int, t reflect.Type, reversed bool) *combineFn {
+ fn := &combineFn{Less: beam.EncodedFunc{Fn: reflectx.MakeFunc(less)}, N: n, Type: beam.EncodedType{T: t}, Reversed: reversed}
+ // Running SetupFn at pipeline construction helps validate the
+ // combineFn, and simplify testing.
+ fn.SetupFn()
+ return fn
}
// TODO(herohde) 5/25/2017: use a heap instead of a sorted slice.
type accum struct {
- coder beam.Coder
- data [][]byte
+ enc beam.ElementEncoder
+ dec beam.ElementDecoder
+
+ data [][]byte
// list stores the elements of type A in order. It has at most size N.
list []interface{}
}
@@ -131,28 +136,26 @@ func (a *accum) unmarshal() error {
if a.data == nil {
return nil
}
- dec := exec.MakeElementDecoder(beam.UnwrapCoder(a.coder))
for _, val := range a.data {
- fv, err := dec.Decode(bytes.NewBuffer(val))
+ element, err := a.dec.Decode(bytes.NewBuffer(val))
if err != nil {
return fmt.Errorf("top.accum: error unmarshal: %v", err)
}
- a.list = append(a.list, fv.Elm)
+ a.list = append(a.list, element)
}
a.data = nil
return nil
}
-// MarshalJSON uses the hook into the JSON encoder library to
+// MarshalJSON uses the hook into the JSON encoder library to encode the accumulator.
func (a accum) MarshalJSON() ([]byte, error) {
- if !a.coder.IsValid() {
- return nil, fmt.Errorf("top.accum: element coder unspecified")
+ if a.enc == nil {
+ return nil, fmt.Errorf("top.accum: element encoder unspecified")
}
- enc := exec.MakeElementEncoder(beam.UnwrapCoder(a.coder))
var values [][]byte
for _, value := range a.list {
var buf bytes.Buffer
- if err := enc.Encode(exec.FullValue{Elm: value}, &buf); err != nil {
+ if err := a.enc.Encode(value, &buf); err != nil {
return nil, fmt.Errorf("top.accum: marshalling of %v failed: %v", value, err)
}
values = append(values, buf.Bytes())
@@ -171,14 +174,21 @@ type combineFn struct {
Reversed bool `json:"reversed"`
// N is the number of elements to keep.
N int `json:"n"`
- // Coder is the element coder for the underlying type, A.
- Coder beam.EncodedCoder `json:"coder"`
+ // Type is the element type A
+ Type beam.EncodedType `json:"type"`
+ enc beam.ElementEncoder
+ dec beam.ElementDecoder
less reflectx.Func2x1
}
+func (f *combineFn) SetupFn() {
+ f.enc = beam.NewElementEncoder(f.Type.T)
+ f.dec = beam.NewElementDecoder(f.Type.T)
+}
+
func (f *combineFn) CreateAccumulator() accum {
- return accum{coder: f.Coder.Coder}
+ return accum{enc: f.enc, dec: f.dec}
}
func (f *combineFn) AddInput(a accum, val beam.T) accum {
@@ -187,8 +197,8 @@ func (f *combineFn) AddInput(a accum, val beam.T) accum {
}
func (f *combineFn) MergeAccumulators(a, b accum) accum {
- a.coder = f.Coder.Coder
- b.coder = f.Coder.Coder
+ a.enc, a.dec = f.enc, f.dec
+ b.enc, b.dec = f.enc, f.dec
if err := a.unmarshal(); err != nil {
panic(err)
}
@@ -225,5 +235,5 @@ func (f *combineFn) trim(ret []interface{}) accum {
if len(ret) > f.N {
ret = ret[:f.N]
}
- return accum{coder: f.Coder.Coder, list: ret}
+ return accum{enc: f.enc, dec: f.dec, list: ret}
}
diff --git a/sdks/go/pkg/beam/transforms/top/top_test.go b/sdks/go/pkg/beam/transforms/top/top_test.go
index 76130ec..a6a38b9 100644
--- a/sdks/go/pkg/beam/transforms/top/top_test.go
+++ b/sdks/go/pkg/beam/transforms/top/top_test.go
@@ -20,8 +20,6 @@ import (
"reflect"
"testing"
- "github.com/apache/beam/sdks/go/pkg/beam/core/typex"
-
"github.com/apache/beam/sdks/go/pkg/beam/core/util/reflectx"
)
@@ -31,7 +29,7 @@ func TestCombineFn3String(t *testing.T) {
less := func(a, b string) bool {
return len(a) < len(b)
}
- fn := newCombineFn(less, 3, typex.New(reflectx.String), false)
+ fn := newCombineFn(less, 3, reflectx.String, false)
tests := []struct {
Elms []string
@@ -59,7 +57,7 @@ func TestCombineFn3RevString(t *testing.T) {
less := func(a, b string) bool {
return len(a) < len(b)
}
- fn := newCombineFn(less, 3, typex.New(reflectx.String), true)
+ fn := newCombineFn(less, 3, reflectx.String, true)
tests := []struct {
Elms []string
@@ -86,7 +84,7 @@ func TestCombineFnMerge(t *testing.T) {
less := func(a, b string) bool {
return len(a) < len(b)
}
- fn := newCombineFn(less, 3, typex.New(reflectx.String), false)
+ fn := newCombineFn(less, 3, reflectx.String, false)
tests := []struct {
Elms [][]string
Expected []string