HBASE-30016 Cached table regions locations are not using in TableSnapshotInputFormatImpl#getSplits

Author: srinireddy2020

hbase-mapreduce/src/test/java/org/apache/hadoop/hbase/mapreduce/TestTableSnapshotInputFormat.java

@@ -39,6 +39,7 @@
 import org.apache.hadoop.hbase.HDFSBlocksDistribution;
 import org.apache.hadoop.hbase.TableName;
 import org.apache.hadoop.hbase.client.Admin;
+import org.apache.hadoop.hbase.client.RegionInfo;
 import org.apache.hadoop.hbase.client.Result;
 import org.apache.hadoop.hbase.client.Scan;
 import org.apache.hadoop.hbase.client.Scan.ReadType;
@@ -47,6 +48,7 @@
 import org.apache.hadoop.hbase.io.ImmutableBytesWritable;
 import org.apache.hadoop.hbase.mapreduce.TableSnapshotInputFormat.TableSnapshotRegionSplit;
 import org.apache.hadoop.hbase.snapshot.RestoreSnapshotHelper;
+import org.apache.hadoop.hbase.snapshot.SnapshotManifest;
 import org.apache.hadoop.hbase.snapshot.SnapshotTestingUtils;
 import org.apache.hadoop.hbase.testclassification.LargeTests;
 import org.apache.hadoop.hbase.testclassification.VerySlowMapReduceTests;
@@ -683,4 +685,49 @@ public void testReadFromRestoredSnapshotViaMR() throws Exception {
       }
     }
   }
+
+  /**
+   * Tests that TableSnapshotInputFormatImpl correctly caches and uses region locations when
+   * locality is enabled
+   */
+  @Test
+  public void testRegionLocatorUsesCache() throws Exception {
+    final TableName tableName = TableName.valueOf("testRegionLocatorCache");
+    Configuration conf = UTIL.getConfiguration();
+    try {
+      // Create table with multiple regions
+      createTableAndSnapshot(UTIL, tableName, "snapshot", getStartRow(), getEndRow(), 3);
+
+      conf.setBoolean(SNAPSHOT_INPUTFORMAT_LOCALITY_BY_REGION_LOCATION, true);
+
+      Path tmpTableDir = UTIL.getDataTestDirOnTestFS("snapshot");
+
+      // Get snapshot manifest and region info
+      SnapshotManifest manifest = TableSnapshotInputFormatImpl.getSnapshotManifest(conf, "snapshot",
+        CommonFSUtils.getRootDir(conf), UTIL.getTestFileSystem());
+      List<RegionInfo> regionInfos =
+        TableSnapshotInputFormatImpl.getRegionInfosFromManifest(manifest);
+
+      Scan scan = new Scan();
+
+      List<TableSnapshotInputFormatImpl.InputSplit> splits = TableSnapshotInputFormatImpl.getSplits(
+        scan, manifest, regionInfos, tmpTableDir, conf, new RegionSplitter.UniformSplit(), 1);
+
+      // Verify that splits contain proper locality information
+      Assert.assertNotNull(splits);
+      Assert.assertTrue(splits.size() > 0);
+
+      // Verify locations are populated from cache
+      for (TableSnapshotInputFormatImpl.InputSplit split : splits) {
+        String[] locations = split.getLocations();
+        // Locations should be populated from cache
+        Assert.assertNotNull(locations);
+      }
+
+    } finally {
+      UTIL.getAdmin().deleteSnapshot("snapshot");
+      UTIL.deleteTable(tableName);
+      conf.unset(SNAPSHOT_INPUTFORMAT_LOCALITY_BY_REGION_LOCATION);
+    }
+  }
 }