Skip to content

Fix coder corruption in prism while handling flatten #34582

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

Merged
jrmccluskey merged 3 commits into from
Apr 9, 2025
Merged

Fix coder corruption in prism while handling flatten #34582

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a345148
Replace coder substitution with pcoll substitution in handling flatten.
shunping Apr 8, 2025
581bb0e
Re-enable some tests for flatten.
shunping Apr 8, 2025
0ce6930
Revert "Re-enable some tests for flatten."
shunping Apr 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions runners/prism/java/build.gradle
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,10 +111,10 @@ def sickbayTests = [

// Java side dying during execution.
// https://github.com/apache/beam/issues/32930
'org.apache.beam.sdk.transforms.FlattenTest.testFlattenMultipleCoders',
// 'org.apache.beam.sdk.transforms.FlattenTest.testFlattenMultipleCoders',
// Stream corruption error java side: failed:java.io.StreamCorruptedException: invalid stream header: 206E6F74
// Likely due to prism't coder changes.
'org.apache.beam.sdk.transforms.FlattenTest.testFlattenWithDifferentInputAndOutputCoders2',
// 'org.apache.beam.sdk.transforms.FlattenTest.testFlattenWithDifferentInputAndOutputCoders2',

// java.lang.IllegalStateException: Output with tag Tag<output> must have a schema in order to call getRowReceiver
// Ultimately because getRoeReceiver code path SDK side isn't friendly to LengthPrefix wrapping of row coders.
Expand Down
9 changes: 5 additions & 4 deletions sdks/go/pkg/beam/runners/prism/internal/handlerunner.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ import (
"github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal/urns"
"github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal/worker"
"google.golang.org/protobuf/encoding/prototext"
"google.golang.org/protobuf/proto"
)

// This file retains the logic for the pardo handler
Expand Down Expand Up @@ -107,12 +108,12 @@ func (h *runner) handleFlatten(tid string, t *pipepb.PTransform, comps *pipepb.C
// they're written out to the runner in the same fashion.
// This may stop being necessary once Flatten Unzipping happens in the optimizer.
outPCol := comps.GetPcollections()[outColID]
outCoder := comps.GetCoders()[outPCol.GetCoderId()]
coderSubs := map[string]*pipepb.Coder{}
pcollSubs := map[string]*pipepb.PCollection{}
for _, p := range t.GetInputs() {
inPCol := comps.GetPcollections()[p]
if inPCol.CoderId != outPCol.CoderId {
coderSubs[inPCol.CoderId] = outCoder
pcollSubs[p] = proto.Clone(inPCol).(*pipepb.PCollection)
pcollSubs[p].CoderId = outPCol.CoderId
}
}

Expand All @@ -123,7 +124,7 @@ func (h *runner) handleFlatten(tid string, t *pipepb.PTransform, comps *pipepb.C
Transforms: map[string]*pipepb.PTransform{
tid: t,
},
Coders: coderSubs,
Pcollections: pcollSubs,
},
RemovedLeaves: nil,
ForcedRoots: forcedRoots,
Expand Down
Loading