Avoid side input when gathering unwindowed WriteFiles results

CHANGES.md

@@ -82,6 +82,7 @@
 ## Bugfixes
 
 * Fixed unbounded checkpoint state growth for splittable DoFns that self-checkpoint on the portable Flink runner (Java) ([#27648](https://github.com/apache/beam/issues/27648)).
+* Fixed bounded, unwindowed Java `WriteFiles` finalization to gather temporary-file results through a main-input shuffle instead of a global side input ([#39370](https://github.com/apache/beam/issues/39370)).
 * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 
 ## Security Fixes
@@ -2507,4 +2508,4 @@ Schema Options, it will be removed in version `2.23.0`. ([BEAM-9704](https://iss
 
 ## Highlights
 
-- For versions 2.19.0 and older release notes are available on [Apache Beam Blog](https://beam.apache.org/blog/).
+- For versions 2.19.0 and older release notes are available on [Apache Beam Blog](https://beam.apache.org/blog/).

sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileBasedSink.java

@@ -34,6 +34,7 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -662,8 +663,14 @@ protected final List<KV<FileResult<DestinationT>, ResourceId>> finalizeDestinati
       if (numShards != null) {
         resultsWithShardNumbers = Lists.newArrayList(completeResults);
       } else {
+        List<FileResult<DestinationT>> orderedResults = Lists.newArrayList(completeResults);
+        if (!windowedWrites) {
+          // Runner-determined shard numbers are positional. Use a stable order so retries after a
+          // partial rename cannot map the same temporary file to a different final shard.
+          orderedResults.sort(Comparator.comparing(result -> result.getTempFilename().toString()));
+        }
         int i = 0;
-        for (FileResult<DestinationT> res : completeResults) {
+        for (FileResult<DestinationT> res : orderedResults) {
           resultsWithShardNumbers.add(res.withShard(i++));
         }
       }

sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java

@@ -25,6 +25,7 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
@@ -40,6 +41,7 @@
 import org.apache.beam.sdk.coders.ShardedKeyCoder;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.coders.VoidCoder;
 import org.apache.beam.sdk.io.FileBasedSink.DynamicDestinations;
 import org.apache.beam.sdk.io.FileBasedSink.FileResult;
 import org.apache.beam.sdk.io.FileBasedSink.FileResultCoder;
@@ -49,6 +51,7 @@
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.options.ValueProvider;
 import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
+import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.DoFn.MultiOutputReceiver;
 import org.apache.beam.sdk.transforms.Flatten;
@@ -57,11 +60,9 @@
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.transforms.Reify;
 import org.apache.beam.sdk.transforms.Reshuffle;
 import org.apache.beam.sdk.transforms.SimpleFunction;
 import org.apache.beam.sdk.transforms.Values;
-import org.apache.beam.sdk.transforms.View;
 import org.apache.beam.sdk.transforms.WithKeys;
 import org.apache.beam.sdk.transforms.display.DisplayData;
 import org.apache.beam.sdk.transforms.errorhandling.BadRecord;
@@ -73,6 +74,7 @@
 import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
 import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
 import org.apache.beam.sdk.transforms.windowing.Window;
 import org.apache.beam.sdk.util.CoderUtils;
 import org.apache.beam.sdk.util.MoreFutures;
@@ -88,6 +90,7 @@
 import org.apache.beam.sdk.values.ShardedKey;
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.WindowingStrategy;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Objects;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ArrayListMultimap;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
@@ -553,29 +556,53 @@ public PCollection<List<ResultT>> expand(PCollection<ResultT> input) {
             // Reshuffle one more time to stabilize the contents of the bundle lists to finalize.
             .apply(Reshuffle.viaRandomKey());
       } else {
-        // Pass results via a side input rather than reshuffle, because we need to get an empty
-        // iterable to finalize if there are no results.
-        return input
-            .getPipeline()
-            .apply(
-                "AsPossiblyEmptyList",
-                Reify.viewInGlobalWindow(
-                    // Insert a reshuffle before taking the view to consolidate the (typically)
-                    // one-output-per-bundle writes.
-                    // This avoids producing a huge number of tiny files in the case that side
-                    // inputs are materialized to disk bundle-by-bundle.
-                    input.apply("Consolidate", Reshuffle.viaRandomKey()).apply(View.asIterable()),
-                    IterableCoder.of(resultCoder)))
-            // View.asIterable() can be (significantly) cheaper than View.asList(), as it does not
-            // create a backing indexable view, but we must return a list to maintain update
-            // compatibility for consumers that are shared between this path and the streaming one.
+        Coder<List<ResultT>> resultListCoder = ListCoder.of(resultCoder);
+        PCollection<List<ResultT>> resultBundles =
+            input
+                .apply("Gather bundles", ParDo.of(new GatherBundlesPerWindowFn<>()))
+                .setCoder(resultListCoder);
+
+        // Add an empty list so empty, unwindowed writes still reach finalization. Grouping the
+        // bundle-sized lists on the main-input path also checkpoints the FileResults against
+        // retries without materializing and reading them back as a global side input.
+        PCollection<List<ResultT>> emptyResultList =
+            input
+                .getPipeline()
+                .apply(
+                    "CreateEmptyResultList",
+                    Create.<List<ResultT>>of(Collections.singletonList(Collections.emptyList()))
+                        .withCoder(resultListCoder));
+
+        return PCollectionList.of(resultBundles)
+            .and(emptyResultList)
+            .apply("EnsureNonEmpty", Flatten.pCollections())
+            .apply("Add void key", WithKeys.of((Void) null))
+            .setCoder(KvCoder.of(VoidCoder.of(), resultListCoder))
+            // The old side-input path emitted its list at the minimum timestamp. Keep that
+            // observable timestamp by combining with the minimum-timestamp empty marker.
             .apply(
-                "IterableToList",
-                MapElements.via(
-                    new SimpleFunction<Iterable<ResultT>, List<ResultT>>(
-                        x -> ImmutableList.copyOf(x)) {}))
-            .setCoder(ListCoder.of(resultCoder));
+                "Preserve minimum timestamp",
+                Window.<KV<Void, List<ResultT>>>configure()
+                    .withTimestampCombiner(TimestampCombiner.EARLIEST))
+            .apply("Gather all results", GroupByKey.create())
+            .apply("Extract results", ParDo.of(new ExtractGatheredResultsFn<>()))
+            .setCoder(resultListCoder)
+            // EARLIEST is internal to gathering. The old side-input path exposed the global
+            // default strategy, so restore it before finalization and filename output.
+            .setWindowingStrategyInternal(WindowingStrategy.globalDefault());
+      }
+    }
+  }
+
+  private static class ExtractGatheredResultsFn<T>
+      extends DoFn<KV<Void, Iterable<List<T>>>, List<T>> {
+    @ProcessElement
+    public void process(ProcessContext c) {
+      List<T> results = new ArrayList<>();
+      for (List<T> bundleResults : c.element().getValue()) {
+        results.addAll(bundleResults);
       }
+      c.output(results);
     }
   }
 

sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileBasedSinkTest.java

@@ -205,6 +205,52 @@ public void testFinalizeWithIntermediateState() throws Exception {
     runFinalize(writeOp, files);
   }
 
+  @Test
+  public void testRunnerDeterminedShardAssignmentIsStableAcrossResultOrder() throws Exception {
+    SimpleSink.SimpleWriteOperation<Void> writeOp = buildWriteOperation();
+    ResourceId earlierTempFilename =
+        getBaseTempDirectory().resolve("a", StandardResolveOptions.RESOLVE_FILE);
+    ResourceId laterTempFilename =
+        getBaseTempDirectory().resolve("z", StandardResolveOptions.RESOLVE_FILE);
+    FileResult<Void> earlierResult =
+        new FileResult<>(
+            earlierTempFilename,
+            UNKNOWN_SHARDNUM,
+            GlobalWindow.INSTANCE,
+            PaneInfo.ON_TIME_AND_ONLY_FIRING,
+            null);
+    FileResult<Void> laterResult =
+        new FileResult<>(
+            laterTempFilename,
+            UNKNOWN_SHARDNUM,
+            GlobalWindow.INSTANCE,
+            PaneInfo.ON_TIME_AND_ONLY_FIRING,
+            null);
+
+    List<KV<FileResult<Void>, ResourceId>> forward =
+        writeOp.finalizeDestination(
+            null, GlobalWindow.INSTANCE, null, Arrays.asList(earlierResult, laterResult));
+    List<KV<FileResult<Void>, ResourceId>> reversed =
+        writeOp.finalizeDestination(
+            null, GlobalWindow.INSTANCE, null, Arrays.asList(laterResult, earlierResult));
+
+    assertEquals(tempToFinalFilenameMappings(forward), tempToFinalFilenameMappings(reversed));
+    assertEquals(earlierTempFilename, forward.get(0).getKey().getTempFilename());
+    assertEquals(0, forward.get(0).getKey().getShard());
+    assertEquals(laterTempFilename, forward.get(1).getKey().getTempFilename());
+    assertEquals(1, forward.get(1).getKey().getShard());
+  }
+
+  private static List<String> tempToFinalFilenameMappings(
+      List<KV<FileResult<Void>, ResourceId>> resultsToFinalFilenames) {
+    List<String> mappings = new ArrayList<>();
+    for (KV<FileResult<Void>, ResourceId> entry : resultsToFinalFilenames) {
+      mappings.add(entry.getKey().getTempFilename() + " -> " + entry.getValue());
+    }
+    Collections.sort(mappings);
+    return mappings;
+  }
+
   /** Generate n temporary files using the temporary file pattern of Writer. */
   private List<File> generateTemporaryFilesForFinalize(int numFiles) throws Exception {
     List<File> temporaryFiles = new ArrayList<>();

sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java

@@ -61,6 +61,7 @@
 import org.apache.beam.sdk.options.Description;
 import org.apache.beam.sdk.options.PipelineOptionsFactoryTest.TestPipelineOptions;
 import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
+import org.apache.beam.sdk.runners.TransformHierarchy;
 import org.apache.beam.sdk.testing.NeedsRunner;
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.testing.TestPipeline;
@@ -94,6 +95,7 @@
 import org.apache.beam.sdk.values.PCollectionView;
 import org.apache.beam.sdk.values.PDone;
 import org.apache.beam.sdk.values.ShardedKey;
+import org.apache.beam.sdk.values.WindowingStrategy;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Optional;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
@@ -189,6 +191,13 @@ public PDone expand(PCollection<KV<DestinationT, String>> input) {
     }
   }
 
+  private static class ExtractTimestampFn<T> extends DoFn<T, Long> {
+    @ProcessElement
+    public void process(ProcessContext c) {
+      c.output(c.timestamp().getMillis());
+    }
+  }
+
   private String getBaseOutputFilename() {
     return getBaseOutputDirectory().resolve("file", StandardResolveOptions.RESOLVE_FILE).toString();
   }
@@ -223,6 +232,53 @@ public void testEmptyWrite() throws IOException {
         true /* expectRemovedTempDirectory */);
   }
 
+  @Test
+  public void testUnwindowedGatherDoesNotMaterializeResultsAsSideInput() {
+    Pipeline pipeline = Pipeline.create();
+    pipeline.apply(Create.of("foo")).apply(WriteFiles.to(makeSimpleSink()));
+
+    List<String> resultViews = new ArrayList<>();
+    pipeline.traverseTopologically(
+        new Pipeline.PipelineVisitor.Defaults() {
+          @Override
+          public void visitPrimitiveTransform(TransformHierarchy.Node node) {
+            if (node.getFullName().contains("GatherTempFileResults")
+                && node.getTransform() instanceof View.CreatePCollectionView) {
+              resultViews.add(node.getFullName());
+            }
+          }
+        });
+
+    assertThat(resultViews, Matchers.empty());
+  }
+
+  @Test
+  public void testUnwindowedOutputFilenamesKeepGlobalDefaultWindowingStrategy() {
+    Pipeline pipeline = Pipeline.create();
+    WriteFilesResult<Void> result =
+        pipeline.apply(Create.of("foo")).apply(WriteFiles.to(makeSimpleSink()));
+
+    assertThat(
+        result.getPerDestinationOutputFilenames().getWindowingStrategy(),
+        equalTo(WindowingStrategy.globalDefault()));
+  }
+
+  @Test
+  @Category(NeedsRunner.class)
+  public void testUnwindowedOutputFilenamesKeepMinimumTimestamp() {
+    WriteFilesResult<Void> result =
+        p.apply(Create.of("foo")).apply(WriteFiles.to(makeSimpleSink()).withNumShards(1));
+
+    PCollection<Long> outputTimestamps =
+        result
+            .getPerDestinationOutputFilenames()
+            .apply("Extract output timestamps", ParDo.of(new ExtractTimestampFn<>()));
+
+    PAssert.that(outputTimestamps)
+        .containsInAnyOrder(BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis());
+    p.run();
+  }
+
   /**
    * Test that WriteFiles with a configured number of shards produces the desired number of shards
    * even when there are many elements.