[#23106] Add periodic.Sequence and periodic.Impulse transforms to Go SDK

CHANGES.md

@@ -64,6 +64,7 @@
 ## New Features / Improvements
 
 * The Flink runner now supports Flink 1.16.x ([#25046](https://github.com/apache/beam/issues/25046)).
+* The Go SDK adds new transforms periodic.Impulse and periodic.Sequence that extends support for slowly updating side input patterns. ([#23106](https://github.com/apache/beam/issues/23106))
 
 ## Breaking Changes
 

sdks/go/examples/slowly_updating_side_input/slowly_updating_side_input.go

@@ -0,0 +1,146 @@
+// 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 main
+
+import (
+	"context"
+	"flag"
+	"strings"
+	"time"
+
+	"cloud.google.com/go/pubsub"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window/trigger"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/io/pubsubio"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/log"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/options/gcpopts"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/register"
+	_ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/dataflow"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/transforms/periodic"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/util/pubsubx"
+)
+
+func init() {
+	register.Function4x0(update)
+	register.Function4x0(process)
+	register.Emitter2[int, string]()
+	register.Iter1[string]()
+}
+
+// update simulates an external call to get data for the side input.
+func update(ctx context.Context, t beam.EventTime, i int64, emit func(int, string)) {
+	log.Infof(ctx, "Making external call %d at %s", i, t.ToTime().Format(time.RFC3339))
+
+	// zero is the key used in beam.AddFixedKey which will be applied on the main input.
+	id, externalData := 0, "some fake data that changed at "+time.Now().Format(time.RFC3339)
+
+	emit(id, externalData)
+}
+
+// process simulates processing of main input. It reads side input by key
+func process(ctx context.Context, k int, v []byte, side func(int) func(*string) bool) {
+	log.Infof(ctx, "Processing (key:%d,value:%q)", k, v)
+
+	iter := side(k)
+
+	var externalData []string
+	var externalDatum string
+	for iter(&externalDatum) {
+		externalData = append(externalData, externalDatum)
+	}
+
+	log.Infof(ctx, "Processing (key:%d,value:%q) with external data %q", k, v, strings.Join(externalData, ","))
+}
+
+func fatalf(err error, format string, args ...interface{}) {
+	if err != nil {
+		log.Fatalf(context.TODO(), format, args...)
+	}
+}
+
+func main() {
+	var inputTopic, periodicSequenceStart, periodicSequenceEnd string
+	var periodicSequenceInterval time.Duration
+
+	now := time.Now()
+
+	flag.StringVar(&periodicSequenceStart, "periodic_sequence_start", now.Add(-1*time.Hour).Format(time.RFC3339),
+		"The time at which to start the periodic sequence.")
+
+	flag.StringVar(&periodicSequenceEnd, "periodic_sequence_end", now.Add(100*time.Hour).Format(time.RFC3339),
+		"The time at which to end the periodic sequence.")
+
+	flag.DurationVar(&periodicSequenceInterval, "periodic_sequence_interval", 1*time.Minute,
+		"The interval between periodic sequence output.")
+
+	flag.StringVar(&inputTopic, "input_topic", "input",
+		"The PubSub topic from which to read the main input data.")
+
+	flag.Parse()
+	beam.Init()
+	ctx := context.Background()
+	p, s := beam.NewPipelineWithRoot()
+
+	project := gcpopts.GetProject(ctx)
+	client, err := pubsub.NewClient(ctx, project)
+	fatalf(err, "Failed to create client: %v", err)
+	_, err = pubsubx.EnsureTopic(ctx, client, inputTopic)
+	fatalf(err, "Failed to ensure topic: %v", err)
+
+	mainInput := beam.WindowInto(
+		s,
+		window.NewFixedWindows(periodicSequenceInterval),
+		beam.AddFixedKey( // simulate keyed data by adding a fixed key
+			s,
+			pubsubio.Read(
+				s,
+				project,
+				inputTopic,
+				nil,
+			),
+		),
+		beam.Trigger(trigger.Repeat(trigger.Always())),
+		beam.PanesDiscard(),
+	)
+
+	startTime, _ := time.Parse(time.RFC3339, periodicSequenceStart)
+	endTime, _ := time.Parse(time.RFC3339, periodicSequenceEnd)
+	sideInput := beam.WindowInto(s, window.NewFixedWindows(periodicSequenceInterval),
+		beam.ParDo(
+			s,
+			update,
+			periodic.Impulse(
+				s,
+				startTime,
+				endTime,
+				periodicSequenceInterval,
+				false,
+			),
+		),
+		beam.Trigger(trigger.Repeat(trigger.Always())),
+		beam.PanesDiscard(),
+	)
+
+	beam.ParDo0(s, process, mainInput,
+		beam.SideInput{
+			Input: sideInput,
+		},
+	)
+
+	if _, err := beam.Run(context.Background(), "dataflow", p); err != nil {
+		log.Exitf(ctx, "Failed to run job: %v", err)
+	}
+}

sdks/go/pkg/beam/transforms/periodic/periodic.go

@@ -0,0 +1,170 @@
+// 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 periodic contains transformations for generating periodic sequences.
+package periodic
+
+import (
+	"context"
+	"reflect"
+	"time"
+
+	"github.com/apache/beam/sdks/v2/go/pkg/beam"
+	"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/sdf"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/io/rtrackers/offsetrange"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/register"
+)
+
+func init() {
+	register.DoFn5x2[context.Context, *sdf.ManualWatermarkEstimator, *sdf.LockRTracker, SequenceDefinition,
+		func(beam.EventTime, []byte),
+		sdf.ProcessContinuation, error](&sequenceGenDoFn{})
+	register.Emitter2[beam.EventTime, []byte]()
+	beam.RegisterType(reflect.TypeOf(SequenceDefinition{}))
+}
+
+// SequenceDefinition holds the configuration for generating a sequence of
+// timestamped elements at an interval.
+type SequenceDefinition struct {
+	Interval time.Duration
+	Start    int64
+	End      int64
+}
+
+type sequenceGenDoFn struct{}
+
+func (fn *sequenceGenDoFn) Setup() {
+}
+
+func (fn *sequenceGenDoFn) CreateInitialRestriction(sd SequenceDefinition) offsetrange.Restriction {
+	totalOutputs := mtime.Time(sd.End).ToTime().Sub(mtime.Time(sd.Start).ToTime()) / sd.Interval
+	return offsetrange.Restriction{
+		Start: int64(0),
+		End:   int64(totalOutputs),
+	}
+}
+
+func (fn *sequenceGenDoFn) CreateTracker(rest offsetrange.Restriction) *sdf.LockRTracker {
+	return sdf.NewLockRTracker(offsetrange.NewTracker(rest))
+}
+
+func (fn *sequenceGenDoFn) RestrictionSize(_ SequenceDefinition, rest offsetrange.Restriction) float64 {
+	return rest.Size()
+}
+
+func (fn *sequenceGenDoFn) SplitRestriction(_ SequenceDefinition, rest offsetrange.Restriction) []offsetrange.Restriction {
+	return []offsetrange.Restriction{rest}
+}
+
+// TruncateRestriction immediately truncates the entire restrication.
+func (fn *sequenceGenDoFn) TruncateRestriction(_ *sdf.LockRTracker, _ SequenceDefinition) offsetrange.Restriction {
+	return offsetrange.Restriction{}
+}
+
+func (fn *sequenceGenDoFn) CreateWatermarkEstimator() *sdf.ManualWatermarkEstimator {
+	return &sdf.ManualWatermarkEstimator{}
+}
+
+func (fn *sequenceGenDoFn) ProcessElement(ctx context.Context, we *sdf.ManualWatermarkEstimator, rt *sdf.LockRTracker, sd SequenceDefinition, emit func(beam.EventTime, []byte)) (sdf.ProcessContinuation, error) {
+	currentOutputIndex := rt.GetRestriction().(offsetrange.Restriction).Start
+	currentOutputTimestamp := mtime.Time(sd.Start).ToTime().Add(sd.Interval * time.Duration(currentOutputIndex))
+	currentTime := time.Now()
+	we.UpdateWatermark(currentOutputTimestamp)
+	for currentOutputTimestamp.Before(currentTime) {
+		if rt.TryClaim(currentOutputIndex) {
+			emit(mtime.FromTime(currentOutputTimestamp), []byte{})
+			currentOutputIndex += 1
+			currentOutputTimestamp = mtime.Time(sd.Start).ToTime().Add(sd.Interval * time.Duration(currentOutputIndex))
+			currentTime = time.Now()
+			we.UpdateWatermark(currentOutputTimestamp)
+		} else if err := rt.GetError(); err != nil || rt.IsDone() {
+			// Stop processing on error or completion
+			return sdf.StopProcessing(), rt.GetError()
+		} else {
+			return sdf.ResumeProcessingIn(sd.Interval), nil
+		}
+	}
+
+	return sdf.ResumeProcessingIn(time.Until(currentOutputTimestamp)), nil
+}
+
+// Impulse is a PTransform which generates a sequence of timestamped
+// elements at fixed runtime intervals. If applyWindow is specified, each
+// element will be assigned to its own fixed window of interval size.
+//
+// The transform behaves the same as [Sequence] transform, but can be
+// used as the first transform in a pipeline.
+//
+// The following applies to the arguments.
+//   - if start is a zero value [time.Time], start is set to the current time
+//   - if start is after end, start is set to end
+//   - start and end are normalized with [mtime.Normalize]
+//   - if interval <= 0 or interval > end.Sub(start), interval is set to end.Sub(start)
+//
+// The PCollection<[]byte> generated by Impulse is unbounded.
+func Impulse(s beam.Scope, start, end time.Time, interval time.Duration, applyWindow bool) beam.PCollection {
+	if start.IsZero() {
+		start = time.Now()
+	}
+	if start.After(end) {
+		start = end
+	}
+	start = mtime.Normalize(mtime.FromTime(start)).ToTime()
+	end = mtime.Normalize(mtime.FromTime(end)).ToTime()
+	if interval <= 0 || interval > end.Sub(start) {
+		interval = end.Sub(start)
+	}
+
+	return genImpulse(s.Scope("periodic.Impulse"), start, end, interval, applyWindow, &sequenceGenDoFn{})
+}
+
+func genImpulse(s beam.Scope, start, end time.Time, interval time.Duration, applyWindow bool, fn *sequenceGenDoFn) beam.PCollection {
+	sd := SequenceDefinition{Interval: interval, Start: start.UnixMilli(), End: end.UnixMilli()}
+	imp := beam.Create(s.Scope("ImpulseElement"), sd)
+	col := genSequence(s, imp, fn)
+	if applyWindow {
+		return beam.WindowInto(s.Scope("ApplyWindowing"),
+			window.NewFixedWindows(interval), col)
+	}
+	return col
+}
+
+// Sequence is a PTransform which generates a sequence of timestamped
+// elements at fixed runtime intervals.
+//
+// The transform assigns each element a timestamp and will only output an
+// element once the worker clock reach the output timestamp. Sequence is not
+// able to guarantee that elements are output at the their exact timestamp, but
+// it guarantees that elements will not be output prior to runtime timestamp.
+//
+// The transform will not output elements prior to the start time.
+//
+// Sequence receives [SequenceDefinition] elements and for each input element
+// received, it will start generating output elements in the following pattern:
+//
+//   - if element timestamp is less than current runtime then output element.
+//   - if element timestamp is greater than current runtime, wait until next
+//     element timestamp.
+//
+// The PCollection<[]byte> generated by Sequence is unbounded.
+func Sequence(s beam.Scope, col beam.PCollection) beam.PCollection {
+	return genSequence(s.Scope("periodic.Sequence"), col, &sequenceGenDoFn{})
+}
+
+func genSequence(s beam.Scope, col beam.PCollection, fn *sequenceGenDoFn) beam.PCollection {
+	return beam.ParDo(s.Scope("GenSequence"), fn, col)
+}

sdks/go/pkg/beam/transforms/periodic/periodic_test.go

@@ -0,0 +1,57 @@
+// 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 periodic
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/apache/beam/sdks/v2/go/pkg/beam"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/passert"
+)
+
+func TestSequence(t *testing.T) {
+	p, s := beam.NewPipelineWithRoot()
+	sd := SequenceDefinition{
+		Interval: time.Second,
+		Start:    0,
+		End:      time.Minute.Milliseconds(),
+	}
+	in := beam.Create(s, sd)
+	out := Sequence(s, in)
+	passert.Count(s, out, "SecondsInMinute", 60)
+	beam.Init()
+	_, err := prism.Execute(context.Background(), p)
+	if err != nil {
+		t.Fatalf("Failed to execute job: %v", err)
+	}
+}
+
+func TestImpulse(t *testing.T) {
+	p, s := beam.NewPipelineWithRoot()
+	interval := time.Second
+	start := time.Unix(0, 0)
+	end := start.Add(time.Minute)
+	out := Impulse(s, start, end, interval, false)
+	passert.Count(s, out, "SecondsInMinute", 60)
+	beam.Init()
+	_, err := prism.Execute(context.Background(), p)
+	if err != nil {
+		t.Fatalf("Failed to execute job: %v", err)
+	}
+}