From 0bba43969f6c79a5ee45ecfbee79bc46011e579f Mon Sep 17 00:00:00 2001
From: Robert Burke <lostluck@users.noreply.github.com>
Date: Thu, 29 Dec 2022 15:25:05 -0800
Subject: [PATCH] [#24789] Spot fix fullvalue wrapping for SDF. (#24826)

* [#24789] Spot fix fullvalue wrapping for SDF.

* Fix indenting on diagrams

* Reinstate comment
---
 sdks/go/pkg/beam/core/runtime/exec/sdf.go | 43 ++++++++++++-----------
 1 file changed, 23 insertions(+), 20 deletions(-)

diff --git a/sdks/go/pkg/beam/core/runtime/exec/sdf.go b/sdks/go/pkg/beam/core/runtime/exec/sdf.go
index be4f491b04cf..d4b1b32d257d 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/sdf.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/sdf.go
@@ -159,13 +159,13 @@ func (n *SplitAndSizeRestrictions) StartBundle(ctx context.Context, id string, d
 // Input Diagram:
 //
 //	*FullValue {
-//	  Elm: *FullValue (original input)
-//	  Elm2: *FullValue {
-//	    Elm: Restriction
-//	    Elm2: Watermark estimator state
-//	  }
-//	  Windows
-//	  Timestamps
+//		Elm: *FullValue (original input)
+//		Elm2: *FullValue {
+//			Elm: Restriction
+//			Elm2: Watermark estimator state
+//		}
+//		Windows
+//		Timestamps
 //	}
 //
 // ProcessElement splits the given restriction into one or more restrictions and
@@ -175,22 +175,25 @@ func (n *SplitAndSizeRestrictions) StartBundle(ctx context.Context, id string, d
 //
 // Output Diagram:
 //
-//	  *FullValue {
-//	    Elm: *FullValue {
-//	      Elm:  *FullValue (original input)
-//	      Elm2: *FullValue {
-//			   Elm: Restriction
-//	        Elm2: Watermark estimator state
-//			 }
-//	    }
-//	    Elm2: float64 (size)
-//	    Windows
-//	    Timestamps
-//	  }
+//	*FullValue {
+//		Elm: *FullValue {
+//			Elm:  *FullValue (original input)
+//			Elm2: *FullValue {
+//				Elm: Restriction
+//				Elm2: Watermark estimator state
+//			}
+//		}
+//		Elm2: float64 (size)
+//		Windows
+//		Timestamps
+//	}
 func (n *SplitAndSizeRestrictions) ProcessElement(ctx context.Context, elm *FullValue, values ...ReStream) error {
 	rest := elm.Elm2.(*FullValue).Elm
 	ws := elm.Elm2.(*FullValue).Elm2
-	mainElm := elm.Elm.(*FullValue)
+
+	// If receiving directly from a datasource,
+	// the element may not be wrapped in a *FullValue
+	mainElm := convertIfNeeded(elm.Elm, &FullValue{})
 
 	splitRests := n.splitInv.Invoke(mainElm, rest)