jonathansudhakar1 opened a new pull request, #40991:
URL: https://github.com/apache/beam/pull/40991
Fixes #23550.
`register.Combiner1/2/3` panicked with `Failed to optimize MergeAccumulators
for combiner &{}. Failed to infer types` (and likewise for the other methods)
when a CombineFn method takes a leading `context.Context`. CombineFn validation
in `graph` already accepts that parameter, and it's how combiners use metrics
and logging. Only `register` rejected it.
This follows the plan @lostluck agreed to on the issue (thank you!):
- **`register.tmpl`**: adds ctx-first variants of the combiner extraction
interfaces (`createAccumulator1x1/1x2`, `addInput3x1/3x2`,
`mergeAccumulators3x1/3x2`, `extractOutput2x1/2x2`). In each `CombinerN`, adds
a branch per variant, next to its non-ctx counterpart. Each branch registers
the `caller` for the `func(context.Context, ...)` type and builds the method
wrapper, the same way the existing branches do.
- **`register.go`**: regenerated from the template. Only the combiner
section changes.
- **Compatibility**: combiners without a context take the same branches as
before. Signatures that are still unsupported still panic with "Failed to
optimize".
## Testing
All commands run from `sdks/`, Go 1.26.
- **Regeneration.** I built `cmd/specialize` and ran it the way `doc.go`'s
`go:generate` lines do (`specialize --package=register --input=register.tmpl
--x=data,universals --imports=typex`, then gofmt). As a check, I first
regenerated from the unchanged template: the output was byte-for-byte identical
to the checked-in `register.go`. Then I regenerated from the edited template.
- **New `TestCombiner_ContextCombiners`** in `register_test.go`,
table-driven, with 7 subtests:
- **Combiners covered:** ctx in every method, with and without `error`
returns, for `Combiner1`, `Combiner2` and `Combiner3`. Together these cover all
eight new interfaces, for each type pairing `[T0, T0/T1/T2]` that has an
error-return variant. A seventh combiner, with only `MergeAccumulators(ctx, a,
b)`, is the exact repro from the issue.
- **Checks on the wrappers from `reflectx.WrapMethods`:**
- each wrapper has a leading `context.Context` parameter;
- each uses the registered caller, not the reflection fallback;
- each forwards the caller's ctx (the methods add a value read from ctx);
- each returns the right results.
- **Checks on `graph.NewCombineFn`:** it accepts the wrapped methods, and
`funcx` detects their context parameter.
- `go test ./go/pkg/beam/register/ -run TestCombiner_ContextCombiners -v`:
all 7 subtests PASS.
With the pre-change `register.go` from `upstream/master`, the same test
fails with the issue's panic: `panic: Failed to optimize MergeAccumulators for
combiner &{}. Failed to infer types`.
- `go test ./go/pkg/beam/register/...`: `ok`. `go vet
./go/pkg/beam/register/`: clean. `gofmt -l go/pkg/beam/register/`: clean.
- `go test ./go/pkg/beam/transforms/stats/...
./go/pkg/beam/transforms/top/... ./go/examples/snippets/...` (other users of
`register.CombinerN`): all `ok`.
- **End to end on Prism (local only, not committed):**
- I ran `beam.Combine` with a registered combiner that takes a ctx in
every method. Its `AddInput` increments a counter through that ctx.
- It PASSES: the output is `"55"`, and the counter is committed as 10.
- Two mutants fail as expected: expecting `"56"` fails the passert, and
expecting 11 for the counter fails the check.
- On the old `register.go` it panics at `init`.
<details><summary>Pipeline test source</summary>
```go
package register_test
import (
"context"
"strconv"
"testing"
"github.com/apache/beam/sdks/v2/go/pkg/beam"
"github.com/apache/beam/sdks/v2/go/pkg/beam/core/metrics"
"github.com/apache/beam/sdks/v2/go/pkg/beam/register"
"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/passert"
"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/ptest"
)
func TestMain(m *testing.M) {
ptest.Main(m)
}
func init() {
register.Combiner2[int, string](&sumWithContextFn{})
}
var sumWithContextInputs = beam.NewCounter("register_test",
"sumWithContextInputs")
// sumWithContextFn sums ints and formats the result. Every method takes a
// context.Context, and AddInput uses it to update a counter.
type sumWithContextFn struct{}
func (fn *sumWithContextFn) CreateAccumulator(_ context.Context) int {
return 0
}
func (fn *sumWithContextFn) AddInput(ctx context.Context, a, i int) (int,
error) {
sumWithContextInputs.Inc(ctx, 1)
return a + i, nil
}
func (fn *sumWithContextFn) MergeAccumulators(_ context.Context, a, b int)
int {
return a + b
}
func (fn *sumWithContextFn) ExtractOutput(_ context.Context, a int) (string,
error) {
return strconv.Itoa(a), nil
}
func TestCombiner_ContextMethodsInPipeline(t *testing.T) {
p, s, col := ptest.CreateList([]int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})
sum := beam.Combine(s, &sumWithContextFn{}, col)
passert.Equals(s, sum, "55")
pr := ptest.RunAndValidate(t, p)
qr := pr.Metrics().Query(func(sr metrics.SingleResult) bool {
return sr.Name() == "sumWithContextInputs"
})
counters := qr.Counters()
if len(counters) != 1 {
t.Fatalf("pr.Metrics().Query(Name = \"sumWithContextInputs\")
returned %d counters, want 1", len(counters))
}
if got, want := counters[0].Committed, int64(10); got != want {
t.Errorf("sumWithContextInputs counter = %v, want %v", got,
want)
}
}
```
</details>
## Open questions
- I left the Prism pipeline test out of the PR. In the `register` package's
test binary, running any pipeline calls `beam.Init`. After that, the existing
unit tests that call `runtime.RegisterType` inside test functions panic with
"Init hooks have already run" (I tried it: `TestEmitter1` fails). If you'd like
an end-to-end test in the repo, I'm happy to add it somewhere better suited,
for example `sdks/go/test/integration/primitives`. Let me know where you'd
prefer it.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
