[Prism] Drop elements for windows whose trigger has finished

sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go

@@ -1515,6 +1515,20 @@ func (ss *stageState) injectTriggeredBundlesIfReady(em *ElementManager, window t
 	return count
 }
 
+// isTriggerFinished returns true if the root trigger for the given window and key
+// has finished, meaning the window is closed and no further panes will be produced.
+// Callers must hold ss.mu.
+func (ss *stageState) isTriggerFinished(window typex.Window, key string) bool {
+	if ss.strat.Trigger == nil {
+		return false
+	}
+	state, ok := ss.state[LinkID{}][window][key]
+	if !ok {
+		return false
+	}
+	return state.getTriggerState(ss.strat.Trigger).finished
+}
+
 // addPending for aggregate stages behaves likes stateful stages, but don't need to handle timers or a separate window
 // expiration condition.
 func (*aggregateStageKind) addPending(ss *stageState, em *ElementManager, newPending []element) int {
@@ -1543,6 +1557,13 @@ func (*aggregateStageKind) addPending(ss *stageState, em *ElementManager, newPen
 
 	count := 0
 	for _, e := range newPending {
+		if len(e.keyBytes) > 0 && ss.isTriggerFinished(e.window, string(e.keyBytes)) {
+			// The trigger for this window and key has already finished, so the
+			// window is closed. Drop the element, matching the Java ReduceFnRunner
+			// (droppedDueToClosedWindow). Otherwise the element would be emitted
+			// in an extra pane when the watermark passes the end of the window.
+			continue
+		}
 		count++
 		if len(e.keyBytes) == 0 {
 			panic(fmt.Sprintf("zero length key: %v %v", ss.ID, ss.inputID))

sdks/go/pkg/beam/runners/prism/internal/unimplemented_test.go

@@ -42,13 +42,6 @@ func TestUnimplemented(t *testing.T) {
 		pipeline func(s beam.Scope)
 	}{
 		// {pipeline: primitives.Drain}, // Can't test drain automatically yet.
-
-		// Implemented but the Go SDK doesn't fully handle panes and
-		// their associated valid behaviors for these triggers, leading
-		// to variable results.
-		// See https://github.com/apache/beam/issues/31153.
-		{pipeline: primitives.TriggerElementCount},
-		{pipeline: primitives.TriggerOrFinally},
 	}
 
 	for _, test := range tests {
@@ -91,6 +84,10 @@ func TestImplemented(t *testing.T) {
 		{pipeline: primitives.TriggerAfterProcessingTime},
 		{pipeline: primitives.TriggerAfterProcessingTimeNotTriggered},
 		{pipeline: primitives.TriggerAfterSynchronizedProcessingTime},
+		// These rely on elements being dropped once the window's trigger has
+		// finished (closed window), rather than emitted in an extra pane.
+		{pipeline: primitives.TriggerElementCount},
+		{pipeline: primitives.TriggerOrFinally},
 	}
 
 	for _, test := range tests {

sdks/go/test/integration/integration.go

@@ -152,9 +152,6 @@ var portableFilters = []string{
 var prismFilters = []string{
 	// The prism runner does not yet support Java's CoGBK.
 	"TestXLang_CoGroupBy",
-	// The trigger and pane tests uses TestStream
-	"TestTrigger.*",
-	"TestPanes",
 
 	// TODO(https://github.com/apache/beam/issues/21058): Xlang ios don't yet work on prism.
 	"TestKafkaIO.*",

sdks/go/test/integration/primitives/windowinto.go

@@ -235,7 +235,9 @@ func TriggerAfterProcessingTimeNotTriggered(s beam.Scope) {
 		}, 10.0)
 }
 
-// TriggerAfterProcessingTime tests the AfterProcessingTime Trigger. It fires output panes once 't' processing time has passed
+// TriggerAfterProcessingTime tests the AfterProcessingTime Trigger. It fires output panes once 't' processing time has passed.
+// Since the trigger is not repeated, it finishes after the first firing and closes the window, so the
+// element arriving after the firing (4.0) is dropped.
 // Not yet supported by the flink runner:
 // java.lang.UnsupportedOperationException: Advancing Processing time is not supported by the Flink Runner.
 func TriggerAfterProcessingTime(s beam.Scope) {
@@ -250,7 +252,7 @@ func TriggerAfterProcessingTime(s beam.Scope) {
 	validateEquals(s.Scope("Global"), window.NewGlobalWindows(), col,
 		[]beam.WindowIntoOption{
 			beam.Trigger(trigger.AfterProcessingTime().PlusDelay(5 * time.Second)),
-		}, 6.0, 4.0)
+		}, 6.0)
 }
 
 // TriggerRepeat tests the repeat trigger. As of now is it is configure to take only one trigger as a subtrigger.

sdks/python/apache_beam/runners/portability/prism_runner_test.py

@@ -288,15 +288,17 @@ def test_after_count_trigger_streaming(self):
           )
           | beam.GroupByKey()
           | beam.MapTuple(PrismRunnerTest.format_result))
+      # AfterCount(3) is not repeated, so it finishes after its first firing
+      # and closes the window. Later elements for that window and key
+      # (A4, A5 and B9) are dropped. Windows whose trigger never fires
+      # (A in [10, 20)) still emit a final pane when they expire.
       assert_that(
           result,
           equal_to(
               list([
                   ('A-3', {1, 2, 3}),
-                  ('A-2', {4, 5}),
                   ('A-2', {10, 11}),
                   ('B-3', {6, 7, 8}),
-                  ('B-1', {9}),
                   ('B-3', {10, 15, 16}),
               ])))