diff --git a/v1/src/main/java/com/google/cloud/teleport/spanner/spannerio/SpannerIO.java b/v1/src/main/java/com/google/cloud/teleport/spanner/spannerio/SpannerIO.java index d9d12d69aa..ca52dd80d4 100644 --- a/v1/src/main/java/com/google/cloud/teleport/spanner/spannerio/SpannerIO.java +++ b/v1/src/main/java/com/google/cloud/teleport/spanner/spannerio/SpannerIO.java @@ -2128,20 +2128,34 @@ public int compareTo(MutationGroupContainer o) { private static class OutputReceiverForFinishBundle implements OutputReceiver> { - private final FinishBundleContext c; - - OutputReceiverForFinishBundle(FinishBundleContext c) { - this.c = c; - } - - @Override - public void output(Iterable output) { - outputWithTimestamp(output, Instant.now()); + private final OutputBuilderSupplier outputBuilderSupplier; + private final DoFn>.FinishBundleContext context; + + OutputReceiverForFinishBundle(FinishBundleContext context) { + this.context = context; + this.outputBuilderSupplier = + new OutputBuilderSupplier() { + @Override + public WindowedValues.Builder builder(OutputT value) { + return WindowedValues.builder() + .setValue(value) + .setTimestamp(Instant.now()) + .setPaneInfo(PaneInfo.NO_FIRING) + .setWindow(GlobalWindow.INSTANCE); + } + }; } @Override - public void outputWithTimestamp(Iterable output, Instant timestamp) { - c.output(output, timestamp, GlobalWindow.INSTANCE); + public OutputBuilder> builder(Iterable value) { + return outputBuilderSupplier + .builder(value) + .setReceiver( + wv -> { + for (BoundedWindow window : wv.getWindows()) { + context.output(wv.getValue(), wv.getTimestamp(), window); + } + }); } } }