[BEAM-3612] Type specialize stats.Sum

apache · Nov 10, 2018 · 86c75eb · 86c75eb
1 parent d1c8010
commit 86c75eb
Show file tree

Hide file tree

Showing 4 changed files with 376 additions and 23 deletions.
diff --git a/sdks/go/pkg/beam/transforms/stats/sum.go b/sdks/go/pkg/beam/transforms/stats/sum.go
@@ -21,6 +21,7 @@ import (
 
 //go:generate specialize --input=sum_switch.tmpl --x=integers,floats
 //go:generate gofmt -w sum_switch.go
+//go:generate go generate sum_switch.go
 
 // Sum returns the sum of the elements in a PCollection<A> as a singleton
 // PCollection<A>. It can only be used for numbers, such as int, uint16,

diff --git a/sdks/go/pkg/beam/transforms/stats/sum_switch.go b/sdks/go/pkg/beam/transforms/stats/sum_switch.go
@@ -20,24 +20,9 @@ package stats
 import (
 	"fmt"
 	"reflect"
-
-	"github.com/apache/beam/sdks/go/pkg/beam"
 )
 
-func init() {
-	beam.RegisterFunction(sumIntFn)
-	beam.RegisterFunction(sumInt8Fn)
-	beam.RegisterFunction(sumInt16Fn)
-	beam.RegisterFunction(sumInt32Fn)
-	beam.RegisterFunction(sumInt64Fn)
-	beam.RegisterFunction(sumUintFn)
-	beam.RegisterFunction(sumUint8Fn)
-	beam.RegisterFunction(sumUint16Fn)
-	beam.RegisterFunction(sumUint32Fn)
-	beam.RegisterFunction(sumUint64Fn)
-	beam.RegisterFunction(sumFloat32Fn)
-	beam.RegisterFunction(sumFloat64Fn)
-}
+//go:generate starcgen --inputs=sum_switch.go --output=sum_switch.shims.go --identifiers=sumIntFn,sumInt8Fn,sumInt16Fn,sumInt32Fn,sumInt64Fn,sumUintFn,sumUint8Fn,sumUint16Fn,sumUint32Fn,sumUint64Fn,sumFloat32Fn,sumFloat64Fn
 
 func findSumFn(t reflect.Type) interface{} {
 	switch t.String() {