@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 testSimpleTagLocationWithInvalidCommit() throws Exception {
    // Load to memory
    HoodieWriteConfig config = getConfigBuilder(100, false, false).withRollbackUsingMarkers(false).build();
    SparkHoodieHBaseIndex index = new SparkHoodieHBaseIndex(config);
    try (SparkRDDWriteClient writeClient = getHoodieWriteClient(config)) {

      String newCommitTime = writeClient.startCommit();
      // make a commit with 199 records
      JavaRDD<HoodieRecord> writeRecords = generateAndCommitRecords(writeClient, 199, newCommitTime);

      // make a second commit with a single record
      String invalidCommit = writeClient.startCommit();
      JavaRDD<HoodieRecord> invalidWriteRecords = generateAndCommitRecords(writeClient, 1, invalidCommit);

      // verify location is tagged.
      HoodieTable hoodieTable = HoodieSparkTable.create(config, context, metaClient);
      JavaRDD<HoodieRecord> javaRDD0 = tagLocation(index, invalidWriteRecords, hoodieTable);
      assert (javaRDD0.collect().size() == 1);   // one record present
      assert (javaRDD0.filter(HoodieRecord::isCurrentLocationKnown).collect().size() == 1); // it is tagged
      assert (javaRDD0.collect().get(0).getCurrentLocation().getInstantTime().equals(invalidCommit));

      // rollback the invalid commit, so that hbase will be left with a stale entry.
      writeClient.rollback(invalidCommit);

      // Now tagLocation for the valid records, hbaseIndex should tag them
      metaClient = HoodieTableMetaClient.reload(metaClient);
      hoodieTable = HoodieSparkTable.create(config, context, metaClient);
      JavaRDD<HoodieRecord> javaRDD1 = tagLocation(index, writeRecords, hoodieTable);
      assert (javaRDD1.filter(HoodieRecord::isCurrentLocationKnown).collect().size() == 199);

      // tagLocation for the invalid record - commit is not present in timeline due to rollback.
      JavaRDD<HoodieRecord> javaRDD2 = tagLocation(index, invalidWriteRecords, hoodieTable);
      assert (javaRDD2.collect().size() == 1);   // one record present
      assert (javaRDD2.filter(HoodieRecord::isCurrentLocationKnown).collect().size() == 0); // it is not tagged
    }
  }