Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

WIP: [Go SDK] Handle PaneInfo correctly #31174

Draft
wants to merge 6 commits into
base: master
Choose a base branch
from
Draft

WIP: [Go SDK] Handle PaneInfo correctly #31174

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
530a89e
pass pane through emitter. ensure prism adds window coders
camphillips22 May 2, 2024
220dac4
try generate
camphillips22 May 2, 2024
bb5661b
more panes
camphillips22 May 2, 2024
e349dbd
cogbk restream pane
camphillips22 May 2, 2024
0954f09
more pane
camphillips22 May 3, 2024
a7b8efc
emit plan from logs
camphillips22 May 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
38 changes: 9 additions & 29 deletions sdks/go/pkg/beam/beam.shims.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion sdks/go/pkg/beam/core/runtime/exec/coder.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -313,7 +313,7 @@ func (*bytesDecoder) DecodeTo(r io.Reader, fv *FullValue) error {
if err != nil {
return err
}
*fv = FullValue{Elm: data}
fv.Elm = data
return nil
}

Expand Down
2 changes: 2 additions & 0 deletions sdks/go/pkg/beam/core/runtime/exec/cogbk.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@ func (n *Inject) ProcessElement(ctx context.Context, elm *FullValue, values ...R
},
Timestamp: elm.Timestamp,
Windows: elm.Windows,
Pane: elm.Pane,
}
return n.Out.ProcessElement(ctx, v, values...)
}
Expand Down Expand Up @@ -173,6 +174,7 @@ func (f *filterStream) Read() (*FullValue, error) {
}
v.Timestamp = elm.Timestamp
v.Windows = elm.Windows
v.Pane = elm.Pane
return v, nil
}
}
8 changes: 4 additions & 4 deletions sdks/go/pkg/beam/core/runtime/exec/combine.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,7 +175,7 @@ func (n *Combine) ProcessElement(ctx context.Context, value *FullValue, values .
if err != nil {
return n.fail(err)
}
return n.Out.ProcessElement(n.ctx, &FullValue{Windows: value.Windows, Elm: value.Elm, Elm2: out, Timestamp: value.Timestamp})
return n.Out.ProcessElement(n.ctx, &FullValue{Windows: value.Windows, Elm: value.Elm, Elm2: out, Timestamp: value.Timestamp, Pane: value.Pane})
}

// FinishBundle completes this node's processing of a bundle.
Expand Down Expand Up @@ -606,7 +606,7 @@ func (n *MergeAccumulators) ProcessElement(ctx context.Context, value *FullValue
return err
}
}
return n.Out.ProcessElement(n.Combine.ctx, &FullValue{Windows: value.Windows, Elm: value.Elm, Elm2: a, Timestamp: value.Timestamp})
return n.Out.ProcessElement(n.Combine.ctx, &FullValue{Windows: value.Windows, Elm: value.Elm, Elm2: a, Timestamp: value.Timestamp, Pane: value.Pane})
}

// Up eagerly gets the optimized binary merge function.
Expand Down Expand Up @@ -637,7 +637,7 @@ func (n *ExtractOutput) ProcessElement(ctx context.Context, value *FullValue, va
if err != nil {
return n.fail(err)
}
return n.Out.ProcessElement(n.Combine.ctx, &FullValue{Windows: value.Windows, Elm: value.Elm, Elm2: out, Timestamp: value.Timestamp})
return n.Out.ProcessElement(n.Combine.ctx, &FullValue{Windows: value.Windows, Elm: value.Elm, Elm2: out, Timestamp: value.Timestamp, Pane: value.Pane})
}

// ConvertToAccumulators is an executor for converting an input value to an accumulator value.
Expand Down Expand Up @@ -665,5 +665,5 @@ func (n *ConvertToAccumulators) ProcessElement(ctx context.Context, value *FullV
if err != nil {
return n.fail(err)
}
return n.Out.ProcessElement(n.Combine.ctx, &FullValue{Windows: value.Windows, Elm: value.Elm, Elm2: a, Timestamp: value.Timestamp})
return n.Out.ProcessElement(n.Combine.ctx, &FullValue{Windows: value.Windows, Elm: value.Elm, Elm2: a, Timestamp: value.Timestamp, Pane: value.Pane})
}
8 changes: 5 additions & 3 deletions sdks/go/pkg/beam/core/runtime/exec/emit.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ import (
// emit event time.
type ReusableEmitter interface {
// Init resets the value. Can be called multiple times.
Init(ctx context.Context, ws []typex.Window, t typex.EventTime) error
Init(ctx context.Context, pn typex.PaneInfo, ws []typex.Window, t typex.EventTime) error
// Value returns the side input value. Constant value.
Value() any
}
Expand Down Expand Up @@ -98,12 +98,14 @@ type emitValue struct {
ctx context.Context
ws []typex.Window
et typex.EventTime
pn typex.PaneInfo
}

func (e *emitValue) Init(ctx context.Context, ws []typex.Window, et typex.EventTime) error {
func (e *emitValue) Init(ctx context.Context, pn typex.PaneInfo, ws []typex.Window, et typex.EventTime) error {
e.ctx = ctx
e.ws = ws
e.et = et
e.pn = pn
return nil
}

Expand All @@ -116,7 +118,7 @@ func (e *emitValue) AttachEstimator(est *sdf.WatermarkEstimator) {
}

func (e *emitValue) invoke(args []reflect.Value) []reflect.Value {
value := &FullValue{Windows: e.ws, Timestamp: e.et}
value := &FullValue{Windows: e.ws, Timestamp: e.et, Pane: e.pn}
isKey := true
for i, t := range e.types {
switch {
Expand Down