You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by lo...@apache.org on 2022/04/14 16:52:20 UTC
[beam] branch master updated: [BEAM-11105] Basic Watermark Estimation (Wall Clock Observing) (#17267)
This is an automated email from the ASF dual-hosted git repository.
lostluck 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 73c683fe2b0 [BEAM-11105] Basic Watermark Estimation (Wall Clock Observing) (#17267)
73c683fe2b0 is described below
commit 73c683fe2b0a02ae7a5327e0cdf2991079ca3500
Author: Danny McCormick <da...@google.com>
AuthorDate: Thu Apr 14 12:52:15 2022 -0400
[BEAM-11105] Basic Watermark Estimation (Wall Clock Observing) (#17267)
---
sdks/go/pkg/beam/core/graph/fn.go | 104 +++++++++++++++++++--
sdks/go/pkg/beam/core/graph/fn_test.go | 104 +++++++++++++++++++++
sdks/go/pkg/beam/core/runtime/exec/datasource.go | 3 +
.../pkg/beam/core/runtime/exec/datasource_test.go | 13 +++
sdks/go/pkg/beam/core/runtime/exec/pardo.go | 1 +
sdks/go/pkg/beam/core/runtime/exec/plan.go | 3 +
sdks/go/pkg/beam/core/runtime/exec/sdf.go | 37 ++++++++
sdks/go/pkg/beam/core/runtime/exec/sdf_invokers.go | 54 ++++++++++-
.../beam/core/runtime/exec/sdf_invokers_test.go | 31 ++++++
sdks/go/pkg/beam/core/runtime/exec/sdf_test.go | 59 ++++++++++++
sdks/go/pkg/beam/core/runtime/exec/translate.go | 8 +-
sdks/go/pkg/beam/core/runtime/genx/genx.go | 5 +
sdks/go/pkg/beam/core/runtime/genx/genx_test.go | 7 +-
sdks/go/pkg/beam/core/runtime/harness/harness.go | 7 +-
sdks/go/pkg/beam/core/sdf/sdf.go | 10 ++
sdks/go/pkg/beam/core/sdf/watermark_estimator.go | 29 ++++++
16 files changed, 462 insertions(+), 13 deletions(-)
diff --git a/sdks/go/pkg/beam/core/graph/fn.go b/sdks/go/pkg/beam/core/graph/fn.go
index 31bd4ee7057..e0f0c8b9619 100644
--- a/sdks/go/pkg/beam/core/graph/fn.go
+++ b/sdks/go/pkg/beam/core/graph/fn.go
@@ -168,6 +168,8 @@ const (
restrictionSizeName = "RestrictionSize"
createTrackerName = "CreateTracker"
+ createWatermarkEstimatorName = "CreateWatermarkEstimator"
+
createAccumulatorName = "CreateAccumulator"
addInputName = "AddInput"
mergeAccumulatorsName = "MergeAccumulators"
@@ -187,15 +189,20 @@ var doFnNames = []string{
splitRestrictionName,
restrictionSizeName,
createTrackerName,
+ createWatermarkEstimatorName,
}
-var sdfNames = []string{
+var requiredSdfNames = []string{
createInitialRestrictionName,
splitRestrictionName,
restrictionSizeName,
createTrackerName,
}
+var watermarkEstimationNames = []string{
+ createWatermarkEstimatorName,
+}
+
var combineFnNames = []string{
createAccumulatorName,
addInputName,
@@ -301,6 +308,22 @@ func (f *SplittableDoFn) RestrictionT() reflect.Type {
return f.CreateInitialRestrictionFn().Ret[0].T
}
+// IsWatermarkEstimating returns whether the DoFn implements a custom watermark estimator.
+func (f *SplittableDoFn) IsWatermarkEstimating() bool {
+ _, ok := f.methods[createWatermarkEstimatorName]
+ return ok
+}
+
+// createWatermarkEstimatorFn returns the "createWatermarkEstimator" function, if present
+func (f *SplittableDoFn) CreateWatermarkEstimatorFn() *funcx.Fn {
+ return f.methods[createWatermarkEstimatorName]
+}
+
+// WatermarkEstimatorT returns the type of the watermark estimator from the SDF
+func (f *SplittableDoFn) WatermarkEstimatorT() reflect.Type {
+ return f.CreateWatermarkEstimatorFn().Ret[0].T
+}
+
// TODO(herohde) 5/19/2017: we can sometimes detect whether the main input must be
// a KV or not based on the other signatures (unless we're more loose about which
// sideinputs are present). Bind should respect that.
@@ -490,6 +513,18 @@ func AsDoFn(fn *Fn, numMainIn mainInputs) (*DoFn, error) {
}
}
+ isWatermarkEstimating, err := validateIsWatermarkEstimating(fn, isSdf)
+ if err != nil {
+ return nil, addContext(err, fn)
+ }
+
+ if isWatermarkEstimating {
+ err := validateWatermarkSig(fn)
+ if err != nil {
+ return nil, addContext(err, fn)
+ }
+ }
+
return (*DoFn)(fn), nil
}
@@ -689,7 +724,7 @@ func validateSideInputsNumUnknown(processFnInputs []funcx.FnParam, method *funcx
// requirements for either case.
//
// For a Fn to be an SDF it must:
-// * Implement all the SDF methods.
+// * Implement all the required (non-watermark related) SDF methods.
// * Include an RTracker parameter in ProcessElement.
// For a Fn to not be an SDF, it must:
// * Implement none of the SDF methods.
@@ -697,7 +732,7 @@ func validateSideInputsNumUnknown(processFnInputs []funcx.FnParam, method *funcx
func validateIsSdf(fn *Fn) (bool, error) {
// Store missing method names so we can output them to the user if validation fails.
var missing []string
- for _, name := range sdfNames {
+ for _, name := range requiredSdfNames {
_, ok := fn.methods[name]
if !ok {
missing = append(missing, name)
@@ -708,10 +743,10 @@ func validateIsSdf(fn *Fn) (bool, error) {
switch len(missing) {
case 0: // All SDF methods present.
isSdf = true
- case len(sdfNames): // No SDF methods.
+ case len(requiredSdfNames): // No SDF methods.
isSdf = false
default: // Anything else means an invalid # of SDF methods.
- err := errors.Errorf("not all SplittableDoFn methods are present. Missing methods: %v", missing)
+ err := errors.Errorf("not all required SplittableDoFn methods are present. Missing methods: %v", missing)
return false, err
}
@@ -784,7 +819,7 @@ func validateSdfSigNumbers(fn *Fn, num int) error {
}
returnNum := 1 // TODO(BEAM-3301): Enable optional error params in SDF methods.
- for _, name := range sdfNames {
+ for _, name := range requiredSdfNames {
method := fn.methods[name]
if len(method.Param) != paramNums[name] {
err := errors.Errorf("unexpected number of params in method %v. got: %v, want: %v",
@@ -813,7 +848,7 @@ func validateSdfSigTypes(fn *Fn, num int) error {
restrictionT := fn.methods[createInitialRestrictionName].Ret[0].T
rTrackerT := reflect.TypeOf((*sdf.RTracker)(nil)).Elem()
- for _, name := range sdfNames {
+ for _, name := range requiredSdfNames {
method := fn.methods[name]
switch name {
case createInitialRestrictionName:
@@ -913,6 +948,61 @@ func validateSdfElementT(fn *Fn, name string, method *funcx.Fn, num int) error {
return nil
}
+// validateIsWatermarkEstimating returns true if watermark estimator methods are present on the DoFn, returns
+// false if they aren't, and returns an error if they are present but the function isn't an sdf and thus doesn't
+// support watermark estimation
+func validateIsWatermarkEstimating(fn *Fn, isSdf bool) (bool, error) {
+ _, isWatermarkEstimating := fn.methods[createWatermarkEstimatorName]
+ if !isSdf && isWatermarkEstimating {
+ return false, errors.Errorf("watermark estimation method %v is defined on non-splittable DoFn. Watermark"+
+ "estimation is only valid on splittable DoFns", createWatermarkEstimatorName)
+ }
+ return isWatermarkEstimating, nil
+}
+
+// validateWatermarkSig validates that all watermark related functions are valid
+func validateWatermarkSig(fn *Fn) error {
+ paramRange := map[string][]int{
+ createWatermarkEstimatorName: []int{0, 0},
+ }
+ returnNum := 1 // TODO(BEAM-3301): Enable optional error params in SDF methods.
+
+ watermarkEstimatorT := reflect.TypeOf((*sdf.WatermarkEstimator)(nil)).Elem()
+
+ for _, name := range watermarkEstimationNames {
+ if method, ok := fn.methods[name]; ok {
+ if len(method.Param) < paramRange[name][0] || len(method.Param) > paramRange[name][1] {
+ err := errors.Errorf("unexpected number of params in method %v. got: %v, want number in range: %v to %v",
+ name, len(method.Param), paramRange[name][0], paramRange[name][1])
+ return errors.SetTopLevelMsgf(err, "Unexpected number of parameters in method %v. "+
+ "Got: %v, Want number in range: %v to %v. Check that the signature conforms to the expected signature for %v, "+
+ "and that elements in SDF method parameters match elements in %v.",
+ name, len(method.Param), paramRange[name][0], paramRange[name][1], name, processElementName)
+ }
+ if len(method.Ret) != returnNum {
+ err := errors.Errorf("unexpected number of returns in method %v. got: %v, want: %v",
+ name, len(method.Ret), returnNum)
+ return errors.SetTopLevelMsgf(err, "Unexpected number of return values in method %v. "+
+ "Got: %v, Want: %v. Check that the signature conforms to the expected signature for %v.",
+ name, len(method.Ret), returnNum, name)
+ }
+
+ switch name {
+ case createWatermarkEstimatorName:
+ if method.Ret[0].T.Implements(watermarkEstimatorT) == false {
+ err := errors.Errorf("invalid output type in method %v, return %v: %v does not implement sdf.WatermarkEstimator",
+ createWatermarkEstimatorName, 0, method.Ret[0].T)
+ return errors.SetTopLevelMsgf(err, "Invalid output type in method %v, "+
+ "return value at index %v (type: %v). Output of method %v must implement sdf.WatermarkEstimator.",
+ createWatermarkEstimatorName, 0, method.Ret[0].T, createWatermarkEstimatorName)
+ }
+ }
+ }
+ }
+
+ return nil
+}
+
// CombineFn represents a CombineFn.
type CombineFn Fn
diff --git a/sdks/go/pkg/beam/core/graph/fn_test.go b/sdks/go/pkg/beam/core/graph/fn_test.go
index 84d318e866f..e38c9a34af6 100644
--- a/sdks/go/pkg/beam/core/graph/fn_test.go
+++ b/sdks/go/pkg/beam/core/graph/fn_test.go
@@ -19,6 +19,7 @@ import (
"context"
"reflect"
"testing"
+ "time"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
)
@@ -227,6 +228,59 @@ func TestNewDoFnSdf(t *testing.T) {
})
}
+func TestNewDoFnWatermarkEstimating(t *testing.T) {
+ t.Run("valid", func(t *testing.T) {
+ tests := []struct {
+ dfn interface{}
+ main mainInputs
+ }{
+ {dfn: &GoodWatermarkEstimating{}, main: MainSingle},
+ }
+
+ for _, test := range tests {
+ t.Run(reflect.TypeOf(test.dfn).String(), func(t *testing.T) {
+ // Valid DoFns should pass validation with and without KV info.
+ if _, err := NewDoFn(test.dfn); err != nil {
+ t.Fatalf("NewDoFn with Watermark Estimation failed: %v", err)
+ }
+ if _, err := NewDoFn(test.dfn, NumMainInputs(test.main)); err != nil {
+ t.Fatalf("NewDoFn(NumMainInputs(%v)) with Watermark Estimation failed: %v", test.main, err)
+ }
+ })
+ }
+ })
+ t.Run("invalid", func(t *testing.T) {
+ tests := []struct {
+ dfn interface{}
+ }{
+ {dfn: &BadWatermarkEstimatingNonSdf{}},
+ {dfn: &BadWatermarkEstimatingCreateWatermarkEstimatorReturnType{}},
+ }
+ for _, test := range tests {
+ t.Run(reflect.TypeOf(test.dfn).String(), func(t *testing.T) {
+ if cfn, err := NewDoFn(test.dfn); err != nil {
+ t.Logf("NewDoFn with SDF failed as expected:\n%v", err)
+ } else {
+ t.Errorf("NewDoFn(%v) = %v, want failure", cfn.Name(), cfn)
+ }
+ // If validation fails with unknown main inputs, then it should
+ // always fail for any known number of main inputs, so test them
+ // all. Error messages won't necessarily match.
+ if cfn, err := NewDoFn(test.dfn, NumMainInputs(MainSingle)); err != nil {
+ t.Logf("NewDoFn(NumMainInputs(MainSingle)) with SDF failed as expected:\n%v", err)
+ } else {
+ t.Errorf("NewDoFn(%v, NumMainInputs(MainSingle)) = %v, want failure", cfn.Name(), cfn)
+ }
+ if cfn, err := NewDoFn(test.dfn, NumMainInputs(MainKv)); err != nil {
+ t.Logf("NewDoFn(NumMainInputs(MainKv)) with SDF failed as expected:\n%v", err)
+ } else {
+ t.Errorf("NewDoFn(%v, NumMainInputs(MainKv)) = %v, want failure", cfn.Name(), cfn)
+ }
+ })
+ }
+ })
+}
+
func TestNewCombineFn(t *testing.T) {
t.Run("valid", func(t *testing.T) {
tests := []struct {
@@ -687,6 +741,40 @@ func (fn *GoodSdfKv) ProcessElement(*RTrackerT, int, int) int {
return 0
}
+type WatermarkEstimatorT struct{}
+
+func (e WatermarkEstimatorT) CurrentWatermark() time.Time {
+ return time.Now()
+}
+
+type GoodWatermarkEstimating struct {
+ *GoodDoFn
+}
+
+func (fn *GoodWatermarkEstimating) CreateInitialRestriction(int) RestT {
+ return RestT{}
+}
+
+func (fn *GoodWatermarkEstimating) SplitRestriction(int, RestT) []RestT {
+ return []RestT{}
+}
+
+func (fn *GoodWatermarkEstimating) RestrictionSize(int, RestT) float64 {
+ return 0
+}
+
+func (fn *GoodWatermarkEstimating) CreateTracker(RestT) *RTrackerT {
+ return &RTrackerT{}
+}
+
+func (fn *GoodWatermarkEstimating) ProcessElement(*RTrackerT, int) int {
+ return 0
+}
+
+func (fn *GoodWatermarkEstimating) CreateWatermarkEstimator() WatermarkEstimatorT {
+ return WatermarkEstimatorT{}
+}
+
// Examples of incorrect SDF signatures.
// Examples with missing methods.
@@ -828,6 +916,14 @@ func (fn *BadSdfRestTCreateTracker) CreateTracker(BadRestT) *RTrackerT {
return &RTrackerT{}
}
+type BadWatermarkEstimatingNonSdf struct {
+ *GoodDoFn
+}
+
+func (fn *BadWatermarkEstimatingNonSdf) CreateWatermarkEstimator() WatermarkEstimatorT {
+ return WatermarkEstimatorT{}
+}
+
// Examples of other type validation that needs to be done.
type BadSdfRestSizeReturn struct {
@@ -868,6 +964,14 @@ func (fn *BadSdfMismatchedRTracker) ProcessElement(*OtherRTrackerT, int) int {
return 0
}
+type BadWatermarkEstimatingCreateWatermarkEstimatorReturnType struct {
+ *GoodWatermarkEstimating
+}
+
+func (fn *BadWatermarkEstimatingCreateWatermarkEstimatorReturnType) CreateWatermarkEstimator() int {
+ return 5
+}
+
// Examples of correct CombineFn signatures
type MyAccum struct{}
diff --git a/sdks/go/pkg/beam/core/runtime/exec/datasource.go b/sdks/go/pkg/beam/core/runtime/exec/datasource.go
index 4b7ada9995a..b7bf2674e8a 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/datasource.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/datasource.go
@@ -387,6 +387,8 @@ func (n *DataSource) Split(splits []int64, frac float64, bufSize int64) (SplitRe
n.splitIdx = s
return SplitResult{PI: s - 1, RI: s}, nil
}
+ // Get the output watermark before splitting to avoid accidentally overestimating
+ ow := su.GetOutputWatermark()
// Otherwise, perform a sub-element split.
ps, rs, err := su.Split(fr)
if err != nil {
@@ -431,6 +433,7 @@ func (n *DataSource) Split(splits []int64, frac float64, bufSize int64) (SplitRe
RS: rsEnc,
TId: su.GetTransformId(),
InId: su.GetInputId(),
+ OW: ow,
}
return res, nil
}
diff --git a/sdks/go/pkg/beam/core/runtime/exec/datasource_test.go b/sdks/go/pkg/beam/core/runtime/exec/datasource_test.go
index 85919334f39..bf367ae1cb4 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/datasource_test.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/datasource_test.go
@@ -21,12 +21,14 @@ import (
"io"
"math"
"testing"
+ "time"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
"github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
+ "google.golang.org/protobuf/types/known/timestamppb"
)
func TestDataSource_PerElement(t *testing.T) {
@@ -507,6 +509,9 @@ func TestDataSource_Split(t *testing.T) {
if got, want := splitRes.InId, testInputId; got != want {
t.Errorf("error in Split: got incorrect Input Id = %v, want %v", got, want)
}
+ if _, ok := splitRes.OW["output1"]; !ok {
+ t.Errorf("error in Split: no output watermark for output1")
+ }
}
// Check that split indices are correct, for both sub-element and channel splits.
@@ -621,6 +626,14 @@ func (n *TestSplittableUnit) GetInputId() string {
return testInputId
}
+// GetOutputWatermark gets the current output watermark of the splittable unit
+// if one is defined, or returns nil otherwise
+func (n *TestSplittableUnit) GetOutputWatermark() map[string]*timestamppb.Timestamp {
+ ow := make(map[string]*timestamppb.Timestamp)
+ ow["output1"] = timestamppb.New(time.Date(2022, time.January, 1, 1, 0, 0, 0, time.UTC))
+ return ow
+}
+
func floatEquals(a, b, epsilon float64) bool {
return math.Abs(a-b) < epsilon
}
diff --git a/sdks/go/pkg/beam/core/runtime/exec/pardo.go b/sdks/go/pkg/beam/core/runtime/exec/pardo.go
index 6a83005f0fa..03a81ab5a62 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/pardo.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/pardo.go
@@ -44,6 +44,7 @@ type ParDo struct {
ctx context.Context
inv *invoker
bf *bundleFinalizer
+ we sdf.WatermarkEstimator
reader StateReader
cache *cacheElm
diff --git a/sdks/go/pkg/beam/core/runtime/exec/plan.go b/sdks/go/pkg/beam/core/runtime/exec/plan.go
index abda426f639..542e4be7be8 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/plan.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/plan.go
@@ -24,6 +24,7 @@ import (
"time"
"github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
+ "google.golang.org/protobuf/types/known/timestamppb"
)
// Plan represents the bundle execution plan. It will generally be constructed
@@ -260,6 +261,8 @@ type SplitResult struct {
RS [][]byte // Residual splits. If an element is split, these are the encoded residuals.
TId string // Transform ID of the transform receiving the split elements.
InId string // Input ID of the input the split elements are received from.
+
+ OW map[string]*timestamppb.Timestamp // Map of outputs to output watermark for the plan being split
}
// Split takes a set of potential split indexes, and if successful returns
diff --git a/sdks/go/pkg/beam/core/runtime/exec/sdf.go b/sdks/go/pkg/beam/core/runtime/exec/sdf.go
index 7a860add9ff..cd1234c29e3 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/sdf.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/sdf.go
@@ -25,6 +25,7 @@ import (
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/sdf"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
"github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
+ "google.golang.org/protobuf/types/known/timestamppb"
)
// PairWithRestriction is an executor for the expanded SDF step of the same
@@ -221,6 +222,7 @@ type ProcessSizedElementsAndRestrictions struct {
TfId string // Transform ID. Needed for splitting.
ctInv *ctInvoker
sizeInv *rsInvoker
+ cweInv *cweInvoker
// SU is a buffered channel for indicating when this unit is splittable.
// When this unit is processing an element, it sends a SplittableUnit
@@ -249,6 +251,9 @@ type ProcessSizedElementsAndRestrictions struct {
// This can change during processing due to splits, but it should always be
// set greater than currW.
numW int
+
+ // List of PTransforms that this Sdf outputs into.
+ outputs []string
}
// ID calls the ParDo's ID method.
@@ -267,6 +272,12 @@ func (n *ProcessSizedElementsAndRestrictions) Up(ctx context.Context) error {
if n.sizeInv, err = newRestrictionSizeInvoker(fn); err != nil {
return errors.WithContextf(err, "%v", n)
}
+ if (*graph.SplittableDoFn)(n.PDo.Fn).IsWatermarkEstimating() {
+ fn = (*graph.SplittableDoFn)(n.PDo.Fn).CreateWatermarkEstimatorFn()
+ if n.cweInv, err = newCreateWatermarkEstimatorInvoker(fn); err != nil {
+ return errors.WithContextf(err, "%v", n)
+ }
+ }
n.SU = make(chan SplittableUnit, 1)
return n.PDo.Up(ctx)
}
@@ -331,6 +342,10 @@ func (n *ProcessSizedElementsAndRestrictions) ProcessElement(_ context.Context,
}
}
+ if n.cweInv != nil {
+ n.PDo.we = n.cweInv.Invoke()
+ }
+
// Begin processing elements, exploding windows if necessary.
n.currW = 0
if !mustExplodeWindows(n.PDo.inv.fn, elm, len(n.PDo.Side) > 0) {
@@ -384,6 +399,9 @@ func (n *ProcessSizedElementsAndRestrictions) ProcessElement(_ context.Context,
func (n *ProcessSizedElementsAndRestrictions) FinishBundle(ctx context.Context) error {
n.ctInv.Reset()
n.sizeInv.Reset()
+ if n.cweInv != nil {
+ n.cweInv.Reset()
+ }
return n.PDo.FinishBundle(ctx)
}
@@ -422,6 +440,10 @@ type SplittableUnit interface {
// GetInputId returns the local input ID of the input that the element being
// split was received from.
GetInputId() string
+
+ // GetOutputWatermark gets the current output watermark of the splittable unit
+ // if one is defined, or nil otherwise.
+ GetOutputWatermark() map[string]*timestamppb.Timestamp
}
// Split splits the currently processing element using its restriction tracker.
@@ -672,6 +694,21 @@ func (n *ProcessSizedElementsAndRestrictions) GetInputId() string {
return indexToInputId(0)
}
+// GetOutputWatermark gets the current output watermark of the splittable unit
+// if one is defined, or returns nil otherwise.
+func (n *ProcessSizedElementsAndRestrictions) GetOutputWatermark() map[string]*timestamppb.Timestamp {
+ if n.PDo.we != nil {
+ ow := timestamppb.New(n.PDo.we.CurrentWatermark())
+ owMap := make(map[string]*timestamppb.Timestamp)
+ for _, out := range n.outputs {
+ owMap[out] = ow
+ }
+ return owMap
+ }
+
+ return nil
+}
+
// SdfFallback is an executor used when an SDF isn't expanded into steps by the
// runner, indicating that the runner doesn't support splitting. It executes all
// the SDF steps together in one unit.
diff --git a/sdks/go/pkg/beam/core/runtime/exec/sdf_invokers.go b/sdks/go/pkg/beam/core/runtime/exec/sdf_invokers.go
index 90598400472..79eeda23afc 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/sdf_invokers.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/sdf_invokers.go
@@ -16,11 +16,12 @@
package exec
import (
+ "reflect"
+
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/funcx"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/sdf"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx"
"github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
- "reflect"
)
// This file contains invokers for SDF methods. These invokers are based off
@@ -299,3 +300,54 @@ func (n *ctInvoker) Reset() {
n.args[i] = nil
}
}
+
+// cweInvoker is an invoker for CreateWatermarkEstimator.
+type cweInvoker struct {
+ fn *funcx.Fn
+ args []interface{} // Cache to avoid allocating new slices per-element.
+ call func() sdf.WatermarkEstimator
+}
+
+func newCreateWatermarkEstimatorInvoker(fn *funcx.Fn) (*cweInvoker, error) {
+ n := &cweInvoker{
+ fn: fn,
+ args: make([]interface{}, len(fn.Param)),
+ }
+ if err := n.initCallFn(); err != nil {
+ return nil, errors.WithContext(err, "sdf CreateWatermarkEstimator invoker")
+ }
+ return n, nil
+}
+
+func (n *cweInvoker) initCallFn() error {
+ // Expects a signature of the form:
+ // () sdf.WatermarkEstimator
+ switch fnT := n.fn.Fn.(type) {
+ case reflectx.Func0x1:
+ n.call = func() sdf.WatermarkEstimator {
+ return fnT.Call0x1().(sdf.WatermarkEstimator)
+ }
+ default:
+ if len(n.fn.Param) != 0 {
+ return errors.Errorf("CreateWatermarkEstimator fn %v has unexpected number of parameters: %v",
+ n.fn.Fn.Name(), len(n.fn.Param))
+ }
+ n.call = func() sdf.WatermarkEstimator {
+ return n.fn.Fn.Call(n.args)[0].(sdf.WatermarkEstimator)
+ }
+ }
+ return nil
+}
+
+// Invoke calls CreateWatermarkEstimator given a restriction and returns an sdf.WatermarkEstimator.
+func (n *cweInvoker) Invoke() sdf.WatermarkEstimator {
+ return n.call()
+}
+
+// Reset zeroes argument entries in the cached slice to allow values to be
+// garbage collected after the bundle ends.
+func (n *cweInvoker) Reset() {
+ for i := range n.args {
+ n.args[i] = nil
+ }
+}
diff --git a/sdks/go/pkg/beam/core/runtime/exec/sdf_invokers_test.go b/sdks/go/pkg/beam/core/runtime/exec/sdf_invokers_test.go
index 31faabd61ba..10f5f899384 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/sdf_invokers_test.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/sdf_invokers_test.go
@@ -17,6 +17,7 @@ package exec
import (
"testing"
+ "time"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph"
"github.com/google/go-cmp/cmp"
@@ -228,6 +229,25 @@ func TestInvokes(t *testing.T) {
})
}
})
+
+ t.Run("CreateWatermarkEstimator Invoker (cweInvoker)", func(t *testing.T) {
+ fn := sdf.CreateWatermarkEstimatorFn()
+ invoker, err := newCreateWatermarkEstimatorInvoker(fn)
+ if err != nil {
+ t.Fatalf("newCreateWatermarkEstimatorInvoker failed: %v", err)
+ }
+ got := invoker.Invoke()
+ want := &VetWatermarkEstimator{}
+ if !cmp.Equal(got, want) {
+ t.Errorf("Invoke() has incorrect output: got: %v, want: %v", got, want)
+ }
+ invoker.Reset()
+ for i, arg := range invoker.args {
+ if arg != nil {
+ t.Errorf("Reset() failed to empty all args. args[%v] = %v", i, arg)
+ }
+ }
+ })
}
// VetRestriction is a restriction used for validating that SDF methods get
@@ -268,6 +288,12 @@ func (rt *VetRTracker) TrySplit(_ float64) (interface{}, interface{}, error) {
return nil, nil, nil
}
+type VetWatermarkEstimator struct{}
+
+func (e *VetWatermarkEstimator) CurrentWatermark() time.Time {
+ return time.Date(2022, time.January, 1, 1, 0, 0, 0, time.UTC)
+}
+
// VetSdf runs an SDF In order to test that these methods get called properly,
// each method will flip the corresponding flag in the passed in VetRestriction,
// overwrite the restriction's Key and Val with the last seen input elements,
@@ -312,6 +338,11 @@ func (fn *VetSdf) CreateTracker(rest *VetRestriction) *VetRTracker {
return &VetRTracker{rest}
}
+// CreateWatermarkEstimator creates a watermark estimator to be used by the Sdf
+func (fn *VetSdf) CreateWatermarkEstimator() *VetWatermarkEstimator {
+ return &VetWatermarkEstimator{}
+}
+
// ProcessElement emits the restriction from the restriction tracker it
// received, with the appropriate flags flipped to track that this was called.
// Note that emitting restrictions is discouraged in normal usage. It is only
diff --git a/sdks/go/pkg/beam/core/runtime/exec/sdf_test.go b/sdks/go/pkg/beam/core/runtime/exec/sdf_test.go
index 9a2ee4049eb..259fad1f7a1 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/sdf_test.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/sdf_test.go
@@ -20,6 +20,7 @@ import (
"fmt"
"strings"
"testing"
+ "time"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window"
@@ -897,6 +898,64 @@ func TestAsSplittableUnit(t *testing.T) {
})
}
})
+
+ t.Run("WatermarkEstimation", func(t *testing.T) {
+ tests := []struct {
+ name string
+ fn *graph.DoFn
+ in FullValue
+ want time.Time
+ }{
+ {
+ name: "SingleElem",
+ fn: dfn,
+ in: FullValue{
+ Elm: &FullValue{
+ Elm: 1,
+ Elm2: &VetRestriction{ID: "Sdf"},
+ },
+ Elm2: 1.0,
+ Timestamp: testTimestamp,
+ Windows: testWindows,
+ },
+ want: time.Date(2022, time.January, 1, 1, 0, 0, 0, time.UTC),
+ },
+ }
+ for _, test := range tests {
+ test := test
+ t.Run(test.name, func(t *testing.T) {
+ capt := &CaptureNode{UID: 2}
+ n := &ParDo{UID: 1, Fn: test.fn, Out: []Node{capt}}
+ node := &ProcessSizedElementsAndRestrictions{PDo: n, outputs: []string{"output1", "output2"}}
+ root := &FixedRoot{UID: 0, Elements: []MainInput{{Key: test.in}}, Out: node}
+ units := []Unit{root, node, capt}
+ constructAndExecutePlan(t, units)
+
+ ow := node.GetOutputWatermark()
+ if ow == nil {
+ t.Errorf("ProcessSizedElementsAndRestrictions(%v), got: nil, want: output watermarks", test.in)
+ } else if got, want := len(ow), 2; got != want {
+ t.Errorf("ProcessSizedElementsAndRestrictions(%v) has incorrect number of watermarks, got: %v, want: %v",
+ test.in, len(ow), 2)
+ } else {
+ if got, ok := ow["output1"]; !ok {
+ t.Errorf("ProcessSizedElementsAndRestrictions(%v) has no watermark for ouptput1, want: %v",
+ test.in, test.want)
+ } else if got.AsTime() != test.want {
+ t.Errorf("ProcessSizedElementsAndRestrictions(%v) has incorrect watermark for output1: got: %v, want: %v",
+ test.in, got.AsTime(), test.want)
+ }
+ if got, ok := ow["output2"]; !ok {
+ t.Errorf("ProcessSizedElementsAndRestrictions(%v) has no watermark for ouptput2, want: %v",
+ test.in, test.want)
+ } else if got.AsTime() != test.want {
+ t.Errorf("ProcessSizedElementsAndRestrictions(%v) has incorrect watermark for output2: got: %v, want: %v",
+ test.in, got.AsTime(), test.want)
+ }
+ }
+ })
+ }
+ })
}
// TestMultiWindowProcessing tests that ProcessSizedElementsAndRestrictions
diff --git a/sdks/go/pkg/beam/core/runtime/exec/translate.go b/sdks/go/pkg/beam/core/runtime/exec/translate.go
index 3b17a5e12cc..5eaac7c5644 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/translate.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/translate.go
@@ -486,7 +486,13 @@ func (b *builder) makeLink(from string, id linkID) (Node, error) {
}
u = n
if urn == urnProcessSizedElementsAndRestrictions {
- u = &ProcessSizedElementsAndRestrictions{PDo: n, TfId: id.to}
+ outputs := make([]string, len(transform.GetOutputs()))
+ i := 0
+ for out := range transform.GetOutputs() {
+ outputs[i] = out
+ i++
+ }
+ u = &ProcessSizedElementsAndRestrictions{PDo: n, TfId: id.to, outputs: outputs}
} else if dofn.IsSplittable() {
u = &SdfFallback{PDo: n}
}
diff --git a/sdks/go/pkg/beam/core/runtime/genx/genx.go b/sdks/go/pkg/beam/core/runtime/genx/genx.go
index 67209eaaf61..422ee2900c7 100644
--- a/sdks/go/pkg/beam/core/runtime/genx/genx.go
+++ b/sdks/go/pkg/beam/core/runtime/genx/genx.go
@@ -112,6 +112,11 @@ func handleDoFn(fn *graph.DoFn, c cache) {
c.pullMethod(sdf.RestrictionSizeFn())
c.pullMethod(sdf.SplitRestrictionFn())
c.regType(sdf.RestrictionT())
+ if !sdf.IsWatermarkEstimating() {
+ return
+ }
+ c.pullMethod(sdf.CreateWatermarkEstimatorFn())
+ c.regType(sdf.WatermarkEstimatorT())
}
func handleCombineFn(fn *graph.CombineFn, c cache) {
diff --git a/sdks/go/pkg/beam/core/runtime/genx/genx_test.go b/sdks/go/pkg/beam/core/runtime/genx/genx_test.go
index 21527ede997..77dbe5c9e2c 100644
--- a/sdks/go/pkg/beam/core/runtime/genx/genx_test.go
+++ b/sdks/go/pkg/beam/core/runtime/genx/genx_test.go
@@ -40,6 +40,7 @@ func TestRegisterDoFn(t *testing.T) {
tA := reflect.TypeOf((*A)(nil)).Elem()
tO := reflect.TypeOf((*O)(nil)).Elem()
tRt := reflect.TypeOf((*sdf.LockRTracker)(nil)).Elem()
+ tWe := reflect.TypeOf((*sdf.WallTimeWatermarkEstimator)(nil)).Elem()
tests := []struct {
name string
@@ -68,7 +69,7 @@ func TestRegisterDoFn(t *testing.T) {
{"DoFn01 pointer reflect", reflect.TypeOf(&DoFn01{}), true, false, []reflect.Type{tDoFn01, tR, tS}},
{"DoFn02 reflect - filtered types", tDoFn02, true, false, []reflect.Type{tDoFn02}},
{"CombineFn01 reflect - combine methods", tCmbFn01, true, false, []reflect.Type{tCmbFn01, tA, tI, tO}},
- {"DoFn03 reflect - sdf methods", tDoFn03, true, false, []reflect.Type{tDoFn03, tRt, tR}},
+ {"DoFn03 reflect - sdf methods", tDoFn03, true, false, []reflect.Type{tDoFn03, tRt, tWe, tR}},
{"DoFn04 reflect - containers", tDoFn04, true, false, []reflect.Type{tDoFn04, tR, tS, tT, tA, tI, tO}},
}
@@ -227,6 +228,10 @@ func (fn *DoFn03) CreateTracker(rest R) *sdf.LockRTracker {
return &sdf.LockRTracker{Rt: RT{}}
}
+func (fn *DoFn03) CreateWatermarkEstimator() *sdf.WallTimeWatermarkEstimator {
+ return &sdf.WallTimeWatermarkEstimator{}
+}
+
type DoFn04 struct{}
func (*DoFn04) ProcessElement([4]R, map[S]T, func(*O) bool, func() func(*I) bool, func([]A)) {
diff --git a/sdks/go/pkg/beam/core/runtime/harness/harness.go b/sdks/go/pkg/beam/core/runtime/harness/harness.go
index 90f8ad8f097..c0b38bfbb0a 100644
--- a/sdks/go/pkg/beam/core/runtime/harness/harness.go
+++ b/sdks/go/pkg/beam/core/runtime/harness/harness.go
@@ -527,9 +527,10 @@ func (c *control) handleInstruction(ctx context.Context, req *fnpb.InstructionRe
for i, r := range sr.RS {
rRoots[i] = &fnpb.DelayedBundleApplication{
Application: &fnpb.BundleApplication{
- TransformId: sr.TId,
- InputId: sr.InId,
- Element: r,
+ TransformId: sr.TId,
+ InputId: sr.InId,
+ Element: r,
+ OutputWatermarks: sr.OW,
},
}
}
diff --git a/sdks/go/pkg/beam/core/sdf/sdf.go b/sdks/go/pkg/beam/core/sdf/sdf.go
index 1de3f5cc085..f546f0e760c 100644
--- a/sdks/go/pkg/beam/core/sdf/sdf.go
+++ b/sdks/go/pkg/beam/core/sdf/sdf.go
@@ -19,6 +19,8 @@
// likely to have bugs.
package sdf
+import "time"
+
// RTracker is an interface used to interact with restrictions while processing elements in
// splittable DoFns (specifically, in the ProcessElement method). Each RTracker tracks the progress
// of a single restriction.
@@ -86,3 +88,11 @@ type RTracker interface {
// is unavailable for some reason.
GetRestriction() interface{}
}
+
+// WatermarkEstimator is an interface used to represent a user defined watermark estimator.
+// Watermark estimators allow users to advance the output watermark of the current sdf.
+type WatermarkEstimator interface {
+ // CurrentWatermark returns the estimator's current watermark. It is called any time a DoFn
+ // splits or checkpoints to advance the output watermark of the restriction's stage.
+ CurrentWatermark() time.Time
+}
diff --git a/sdks/go/pkg/beam/core/sdf/watermark_estimator.go b/sdks/go/pkg/beam/core/sdf/watermark_estimator.go
new file mode 100644
index 00000000000..1aae0c710db
--- /dev/null
+++ b/sdks/go/pkg/beam/core/sdf/watermark_estimator.go
@@ -0,0 +1,29 @@
+// 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 sdf
+
+import "time"
+
+// WallTimeWatermarkEstimator is a watermark estimator that advances the
+// current DoFn's output watermark to the current wallclock time on splits
+// or checkpoints.
+type WallTimeWatermarkEstimator struct{}
+
+// CurrentWatermark returns the current time. It is used by the Sdk harness
+// to set the current DoFn's output watermark on splits and checkpoints.
+func (e *WallTimeWatermarkEstimator) CurrentWatermark() time.Time {
+ return time.Now()
+}