@Override
  public <R> HoodieData<HoodieRecord<R>> tagLocation(
      HoodieData<HoodieRecord<R>> records, HoodieEngineContext context,
      HoodieTable hoodieTable) {
    return HoodieJavaRDD.of(HoodieJavaRDD.getJavaRDD(records)
        .mapPartitionsWithIndex(locationTagFunction(hoodieTable.getMetaClient()), true));
  }
@Test
  public void testSimpleTagLocationAndUpdateWithRollback() throws Exception {
    // Load to memory
    HoodieWriteConfig config = getConfigBuilder(100, false, false)
        .withRollbackUsingMarkers(false).build();
    SparkHoodieHBaseIndex index = new SparkHoodieHBaseIndex(config);
    try (SparkRDDWriteClient writeClient = getHoodieWriteClient(config)) {

      final String newCommitTime = writeClient.startCommit();
      final int numRecords = 10;
      List<HoodieRecord> records = dataGen.generateInserts(newCommitTime, numRecords);
      JavaRDD<HoodieRecord> writeRecords = jsc().parallelize(records, 1);
      metaClient = HoodieTableMetaClient.reload(metaClient);

      // Insert 200 records
      JavaRDD<WriteStatus> writeStatues = writeClient.upsert(writeRecords, newCommitTime);
      assertNoWriteErrors(writeStatues.collect());

      // commit this upsert
      writeClient.commit(newCommitTime, writeStatues);
      HoodieTable hoodieTable = HoodieSparkTable.create(config, context, metaClient);
      // Now tagLocation for these records, hbaseIndex should tag them
      List<HoodieRecord> records2 = tagLocation(index, writeRecords, hoodieTable).collect();
      assertEquals(numRecords, records2.stream().filter(HoodieRecord::isCurrentLocationKnown).count());

      // check tagged records are tagged with correct fileIds
      List<String> fileIds = writeStatues.map(WriteStatus::getFileId).collect();
      assertEquals(0, records2.stream().filter(record -> record.getCurrentLocation().getFileId() == null).count());
      List<String> taggedFileIds = records2.stream().map(record -> record.getCurrentLocation().getFileId()).distinct().collect(Collectors.toList());

      // both lists should match
      assertTrue(taggedFileIds.containsAll(fileIds) && fileIds.containsAll(taggedFileIds));
      // Rollback the last commit
      writeClient.rollback(newCommitTime);

      hoodieTable = HoodieSparkTable.create(config, context, metaClient);
      // Now tagLocation for these records, hbaseIndex should not tag them since it was a rolled
      // back commit
      List<HoodieRecord> records3 = tagLocation(index, writeRecords, hoodieTable).collect();
      assertEquals(0, records3.stream().filter(HoodieRecord::isCurrentLocationKnown).count());
      assertEquals(0, records3.stream().filter(record -> record.getCurrentLocation() != null).count());
    }
  }