fix(metadata): use committed Spark write statuses after retries

hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/client/SparkStreamingMetadataWriteHandler.java

@@ -42,26 +42,32 @@
     Option<HoodieTableMetadataWriter> metadataWriterOpt = getMetadataWriter(instantTime, table);
     ValidationUtils.checkState(metadataWriterOpt.isPresent(),
         "Cannot instantiate metadata writer for the table of interest " + table.getMetaClient().getBasePath());
-    return streamWriteToMetadataTable(dataTableWriteStatuses, metadataWriterOpt.get(), table, instantTime, coalesceDivisorForDataTableWrites);
+    return streamWriteToMetadataTable(dataTableWriteStatuses, metadataWriterOpt.get(), table, instantTime,
+        coalesceDivisorForDataTableWrites);
   }
 
   private HoodieData<WriteStatus> streamWriteToMetadataTable(HoodieData<WriteStatus> dataTableWriteStatuses,
                                                              HoodieTableMetadataWriter metadataWriter,
                                                              HoodieTable table,
                                                              String instantTime,
                                                              int coalesceDivisorForDataTableWrites) {
-    HoodieData<WriteStatus> mdtWriteStatuses = metadataWriter.streamWriteToMetadataPartitions(dataTableWriteStatuses, instantTime);
-    mdtWriteStatuses.persist("MEMORY_AND_DISK_SER", table.getContext(), HoodieData.HoodieDataCacheKey.of(table.getMetaClient().getBasePath().toString(), instantTime));
-    HoodieData<WriteStatus> coalescedDataWriteStatuses;
-    int coalesceParallelism = Math.max(1, dataTableWriteStatuses.getNumPartitions() / coalesceDivisorForDataTableWrites);
-    // lets coalesce to lesser number of spark tasks so that, when unioned along with metadata table write status,
-    // we only allocate very less number of tasks for data table write statuses.
-    // In fact, data table writes should have triggered in previous stage before coalesce (partition by below forces the writes
-    // to data table is triggered in previous stage and with the coalesced stage)
-    coalescedDataWriteStatuses = HoodieJavaRDD.of(HoodieJavaRDD.getJavaRDD(dataTableWriteStatuses)
-            .mapToPair((PairFunction<WriteStatus, String, WriteStatus>) writeStatus -> new Tuple2(writeStatus.getStat().getPath(), writeStatus))
+    int coalesceParallelism = Math.max(1,
+        dataTableWriteStatuses.getNumPartitions() / coalesceDivisorForDataTableWrites);
+
+    // Materialize data-table statuses behind a shuffle before deriving metadata records. Spark's
+    // scheduler exposes only the successful attempt for each shuffle partition, so failed or losing
+    // speculative attempts cannot contribute stale RLI/SI locations to the metadata write.
+    HoodieData<WriteStatus> successfulDataWriteStatuses = HoodieJavaRDD.of(
+        HoodieJavaRDD.getJavaRDD(dataTableWriteStatuses)
+            .mapToPair((PairFunction<WriteStatus, String, WriteStatus>) writeStatus ->
+                new Tuple2<>(writeStatus.getStat().getPath(), writeStatus))
             .partitionBy(new CoalescingPartitioner(coalesceParallelism))
             .map((Function<Tuple2<String, WriteStatus>, WriteStatus>) entry -> entry._2));
-    return coalescedDataWriteStatuses.union(mdtWriteStatuses);
+
+    HoodieData<WriteStatus> metadataWriteStatuses =
+        metadataWriter.streamWriteToMetadataPartitions(committedDataWriteStatuses, instantTime);
+    metadataWriteStatuses.persist("MEMORY_AND_DISK_SER", table.getContext(),
+        HoodieData.HoodieDataCacheKey.of(table.getMetaClient().getBasePath().toString(), instantTime));
+    return committedDataWriteStatuses.union(metadataWriteStatuses);
   }
 }

hudi-client/hudi-spark-client/src/test/java/org/apache/hudi/client/TestSparkStreamingMetadataWriteHandler.java

@@ -21,6 +21,7 @@
 
 import org.apache.hudi.common.data.HoodieData;
 import org.apache.hudi.common.engine.HoodieEngineContext;
+import org.apache.hudi.common.model.HoodieWriteStat;
 import org.apache.hudi.common.table.HoodieTableMetaClient;
 import org.apache.hudi.common.util.Option;
 import org.apache.hudi.data.HoodieJavaRDD;
@@ -29,23 +30,36 @@
 import org.apache.hudi.table.HoodieTable;
 import org.apache.hudi.testutils.SparkClientFunctionalTestHarness;
 
+import org.apache.spark.SparkConf;
+import org.apache.spark.TaskContext;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
 
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Stream;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 public class TestSparkStreamingMetadataWriteHandler extends SparkClientFunctionalTestHarness {
 
+  @Override
+  public SparkConf conf() {
+    // Permit one retry so the regression can distinguish a failed task attempt from the
+    // scheduler-selected successful output consumed by the metadata completion path.
+    return conf(java.util.Collections.singletonMap("spark.master", "local[8,2]"));
+  }
+
   private final HoodieTable<?, ?, ?, ?> mockHoodieTable = mock(HoodieTable.class);
   private HoodieTableMetaClient metaClient;
 
@@ -60,27 +74,93 @@ void setUp() {
 
   private static Stream<Arguments> coalesceDivisorTestArgs() {
     return Arrays.stream(new Object[][] {
-        {100, 20, 1000, true},
-        {1, 1, 1000, true},
-        {10000, 100, 5000, true},
-        {10000, 100, 5000, true},
-        {10000, 100, 20000, true},
-        {10000, 100, 20000, true}
+        {100, 1000},
+        {1, 1},
+        {10000, 1},
+        {10000, 5000},
+        {10001, 5000},
+        {10000, 20000}
     }).map(Arguments::of);
   }
 
   @ParameterizedTest
   @MethodSource("coalesceDivisorTestArgs")
-  public void testCoalesceDividentConfig(int numDataTableWriteStatuses, int numMdtWriteStatus, int coalesceDividentForDataTableWrites) {
+  public void testCoalesceDividentConfig(int numDataTableWriteStatuses, int coalesceDividentForDataTableWrites) {
     HoodieData<WriteStatus> dataTableWriteStatus = mockWriteStatuses(numDataTableWriteStatuses);
-    HoodieData<WriteStatus> mdtWriteStatus = mockWriteStatuses(numMdtWriteStatus);
-    HoodieTableMetadataWriter mdtWriter = mock(HoodieTableMetadataWriter.class);
-    when(mdtWriter.streamWriteToMetadataPartitions(any(), any())).thenReturn(mdtWriteStatus);
+    HoodieTableMetadataWriter mdtWriter = metadataWriterReturningEmptyStatuses();
     SparkStreamingMetadataWriteHandler metadataWriteHandler = new MockSparkStreamingMetadataWriteHandler(mdtWriter);
 
     HoodieData<WriteStatus> allWriteStatuses = metadataWriteHandler.streamWriteToMetadataTable(mockHoodieTable, dataTableWriteStatus, "00001",
         coalesceDividentForDataTableWrites);
-    assertEquals(Math.max(1, numDataTableWriteStatuses / coalesceDividentForDataTableWrites) + numMdtWriteStatus, allWriteStatuses.getNumPartitions());
+    assertEquals(Math.max(1, numDataTableWriteStatuses / coalesceDividentForDataTableWrites),
+        allWriteStatuses.getNumPartitions());
+    verify(mdtWriter).streamWriteToMetadataPartitions(any(), any());
+  }
+
+  @Test
+  void testSparkCollectsOnlySuccessfulRetriedTaskAttempt() {
+    HoodieData<WriteStatus> retryingDataWriteStatus = HoodieJavaRDD.of(jsc().parallelize(Arrays.asList(0, 1), 2)
+        .map(partition -> {
+          int attempt = TaskContext.get().attemptNumber();
+          if (partition == 0 && attempt == 0) {
+            throw new IllegalStateException("intentional first-attempt failure");
+          }
+          return writeStatus("file-" + partition + "-attempt-" + attempt, false);
+        }));
+    HoodieTableMetadataWriter mdtWriter = mock(HoodieTableMetadataWriter.class);
+    AtomicReference<HoodieData<WriteStatus>> metadataInput = new AtomicReference<>();
+    when(mdtWriter.streamWriteToMetadataPartitions(any(), any())).thenAnswer(invocation -> {
+      metadataInput.set(invocation.getArgument(0));
+      return HoodieJavaRDD.of(jsc().emptyRDD());
+    });
+
+    List<WriteStatus> statuses = new MockSparkStreamingMetadataWriteHandler(mdtWriter)
+        .streamWriteToMetadataTable(mockHoodieTable, retryingDataWriteStatus, "00001", 1000)
+        .collectAsList();
+    List<String> paths = statuses.stream().map(status -> status.getStat().getPath()).collect(java.util.stream.Collectors.toList());
+
+    List<String> metadataPaths = metadataInput.get().collectAsList().stream()
+        .map(status -> status.getStat().getPath())
+        .collect(java.util.stream.Collectors.toList());
+    assertTrue(paths.contains("file-0-attempt-1"), "Collected output must come from the successful attempt");
+    assertTrue(paths.stream().noneMatch(path -> path.equals("file-0-attempt-0")));
+    assertTrue(metadataPaths.contains("file-0-attempt-1"),
+        "Metadata generation must consume the successful attempt");
+    assertTrue(metadataPaths.stream().noneMatch(path -> path.equals("file-0-attempt-0")));
+    verify(mdtWriter).streamWriteToMetadataPartitions(any(), any());
+  }
+
+  @Test
+  void testSparkUnionsMetadataWriteStatusesAfterCommittedDataStatuses() {
+    HoodieData<WriteStatus> lazyDataWriteStatus = HoodieJavaRDD.of(jsc().parallelize(Arrays.asList(0, 1), 2)
+        .map(partition -> writeStatus("file-" + partition, false)));
+    HoodieTableMetadataWriter mdtWriter = mock(HoodieTableMetadataWriter.class);
+    when(mdtWriter.streamWriteToMetadataPartitions(any(), any()))
+        .thenReturn(HoodieJavaRDD.of(jsc().parallelize(
+            java.util.Collections.singletonList(writeStatus("metadata-file", true)), 1)));
+
+    List<WriteStatus> statuses = new MockSparkStreamingMetadataWriteHandler(mdtWriter)
+        .streamWriteToMetadataTable(mockHoodieTable, lazyDataWriteStatus, "00001", 1000)
+        .collectAsList();
+
+    assertEquals(3, statuses.size());
+    assertEquals(1, statuses.stream().filter(WriteStatus::isMetadataTable).count());
+    verify(mdtWriter).streamWriteToMetadataPartitions(any(), any());
+  }
+
+  private HoodieTableMetadataWriter metadataWriterReturningEmptyStatuses() {
+    HoodieTableMetadataWriter mdtWriter = mock(HoodieTableMetadataWriter.class);
+    when(mdtWriter.streamWriteToMetadataPartitions(any(), any()))
+        .thenReturn(HoodieJavaRDD.of(jsc().emptyRDD()));
+    return mdtWriter;
+  }
+
+  private static WriteStatus writeStatus(String path, boolean isMetadataTable) {
+    HoodieWriteStat stat = new HoodieWriteStat();
+    stat.setPath(path);
+    WriteStatus status = new WriteStatus(true, 0.0, isMetadataTable);
+    status.setStat(stat);
+    return status;
   }
 
   private HoodieData<WriteStatus> mockWriteStatuses(int size) {