From 758afcd5f69aba8f0ff437107e980ea31a32726d Mon Sep 17 00:00:00 2001
From: Yi Hu <yathu@google.com>
Date: Tue, 4 Aug 2026 22:45:40 -0400
Subject: [PATCH] Support sharded coder for Prism runner cross-lang

---
 .../go/pkg/beam/runners/prism/internal/coders.go | 11 +++++++++++
 .../beam/runners/prism/internal/coders_test.go   | 16 ++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/sdks/go/pkg/beam/runners/prism/internal/coders.go b/sdks/go/pkg/beam/runners/prism/internal/coders.go
index d326a332b8d3..0f770849a984 100644
--- a/sdks/go/pkg/beam/runners/prism/internal/coders.go
+++ b/sdks/go/pkg/beam/runners/prism/internal/coders.go
@@ -367,6 +367,17 @@ func pullDecoderNoAlloc(c *pipepb.Coder, coders map[string]*pipepb.Coder) func(i
 			ed(r)
 			wd(r)
 		}
+	case urns.CoderShardedKey:
+		ccids := c.GetComponentCoderIds()
+		if len(ccids) != 1 {
+			panic(fmt.Sprintf("ShardedKey coder must have only 1 component: %s", prototext.Format(c)))
+		}
+		kd := pullDecoderNoAlloc(coders[ccids[0]], coders)
+		return func(r io.Reader) {
+			l, _ := coder.DecodeVarInt(r)
+			ioutilx.ReadN(r, int(l))
+			kd(r)
+		}
 	case urns.CoderRow:
 		panic(fmt.Sprintf("Runner forgot to LP this Row Coder. %v", prototext.Format(c)))
 	default:
diff --git a/sdks/go/pkg/beam/runners/prism/internal/coders_test.go b/sdks/go/pkg/beam/runners/prism/internal/coders_test.go
index 4656a94e03ec..1d1a8b6d4596 100644
--- a/sdks/go/pkg/beam/runners/prism/internal/coders_test.go
+++ b/sdks/go/pkg/beam/runners/prism/internal/coders_test.go
@@ -370,6 +370,22 @@ func Test_pullDecoder(t *testing.T) {
 				},
 			},
 			[]byte{3, 0},
+		}, {
+			"sharded_key",
+			&pipepb.Coder{
+				Spec: &pipepb.FunctionSpec{
+					Urn: urns.CoderShardedKey,
+				},
+				ComponentCoderIds: []string{"key"},
+			},
+			map[string]*pipepb.Coder{
+				"key": {
+					Spec: &pipepb.FunctionSpec{
+						Urn: urns.CoderVarInt,
+					},
+				},
+			},
+			[]byte{3, 1, 2, 3, 255, 3},
 		},
 	}
 	for _, test := range tests {