HBASE-29272 Fix TableSnapshotInputFormatImpl returning 0 for InputSplit length (branch-2.6)

hbase-mapreduce/src/main/java/org/apache/hadoop/hbase/mapred/TableSnapshotInputFormat.java

@@ -56,10 +56,33 @@ public TableSnapshotRegionSplit(TableSnapshotInputFormatImpl.InputSplit delegate
       this.delegate = delegate;
     }
 
+    /**
+     * @deprecated since 2.6.7 and will be removed in 4.0.0.
+     * Use {@link #TableSnapshotRegionSplit(TableSnapshotInputFormatImpl.InputSplit)} instead.
+     * @see <a href="https://issues.apache.org/jira/browse/HBASE-29272">HBASE-29272</a>
+     */
+    @Deprecated
     public TableSnapshotRegionSplit(HTableDescriptor htd, HRegionInfo regionInfo,
       List<String> locations, Scan scan, Path restoreDir) {
+      this(htd, regionInfo, locations, scan, restoreDir, 1);
+    }
+
+    /**
+     * Creates a TableSnapshotRegionSplit with the given region information and the estimated length
+     * of the region in bytes. The length is used by the MapReduce framework to balance the
+     * distribution of splits; a value of 0 may cause the framework to skip the split.
+     * @param htd the table descriptor
+     * @param regionInfo the region info
+     * @param locations the locations of the region
+     * @param scan the scan to use
+     * @param restoreDir the directory to restore the snapshot to
+     * @param length the estimated size of the region in bytes
+     * @see <a href="https://issues.apache.org/jira/browse/HBASE-29272">HBASE-29272</a>
+     */
+    public TableSnapshotRegionSplit(HTableDescriptor htd, HRegionInfo regionInfo,
+      List<String> locations, Scan scan, Path restoreDir, long length) {
       this.delegate =
-        new TableSnapshotInputFormatImpl.InputSplit(htd, regionInfo, locations, scan, restoreDir);
+        new TableSnapshotInputFormatImpl.InputSplit(htd, regionInfo, locations, scan, restoreDir, length);
     }
 
     @Override

hbase-mapreduce/src/main/java/org/apache/hadoop/hbase/mapreduce/TableSnapshotInputFormat.java

@@ -95,10 +95,33 @@ public TableSnapshotRegionSplit(TableSnapshotInputFormatImpl.InputSplit delegate
       this.delegate = delegate;
     }
 
+    /**
+     * @deprecated since 2.6.7 and will be removed in 4.0.0.
+     * Use {@link #TableSnapshotRegionSplit(TableSnapshotInputFormatImpl.InputSplit)} instead.
+     * @see <a href="https://issues.apache.org/jira/browse/HBASE-29272">HBASE-29272</a>
+     */
+    @Deprecated
     public TableSnapshotRegionSplit(HTableDescriptor htd, HRegionInfo regionInfo,
       List<String> locations, Scan scan, Path restoreDir) {
+      this(htd, regionInfo, locations, scan, restoreDir, 1);
+    }
+
+    /**
+     * Creates a TableSnapshotRegionSplit with the given region information and the estimated length
+     * of the region in bytes. The length is used by the MapReduce framework to balance the
+     * distribution of splits; a value of 0 may cause the framework to skip the split.
+     * @param htd the table descriptor
+     * @param regionInfo the region info
+     * @param locations the locations of the region
+     * @param scan the scan to use
+     * @param restoreDir the directory to restore the snapshot to
+     * @param length the estimated size of the region in bytes
+     * @see <a href="https://issues.apache.org/jira/browse/HBASE-29272">HBASE-29272</a>
+     */
+    public TableSnapshotRegionSplit(HTableDescriptor htd, HRegionInfo regionInfo,
+      List<String> locations, Scan scan, Path restoreDir, long length) {
       this.delegate =
-        new TableSnapshotInputFormatImpl.InputSplit(htd, regionInfo, locations, scan, restoreDir);
+        new TableSnapshotInputFormatImpl.InputSplit(htd, regionInfo, locations, scan, restoreDir, length);
     }
 
     @Override

hbase-mapreduce/src/main/java/org/apache/hadoop/hbase/mapreduce/TableSnapshotInputFormatImpl.java

@@ -20,6 +20,7 @@
 import java.io.ByteArrayOutputStream;
 import java.io.DataInput;
 import java.io.DataOutput;
+import java.io.EOFException;
 import java.io.IOException;
 import java.lang.reflect.InvocationTargetException;
 import java.util.ArrayList;
@@ -45,8 +46,10 @@
 import org.apache.hadoop.hbase.io.ImmutableBytesWritable;
 import org.apache.hadoop.hbase.regionserver.HRegion;
 import org.apache.hadoop.hbase.snapshot.RestoreSnapshotHelper;
+import org.apache.hadoop.hbase.snapshot.RegionSizes;
 import org.apache.hadoop.hbase.snapshot.SnapshotDescriptionUtils;
 import org.apache.hadoop.hbase.snapshot.SnapshotManifest;
+import org.apache.hadoop.hbase.snapshot.SnapshotRegionSizeCalculator;
 import org.apache.hadoop.hbase.util.Bytes;
 import org.apache.hadoop.hbase.util.CommonFSUtils;
 import org.apache.hadoop.hbase.util.RegionSplitter;
@@ -145,13 +148,14 @@ public static class InputSplit implements Writable {
     private String[] locations;
     private String scan;
     private String restoreDir;
+    private long length;
 
     // constructor for mapreduce framework / Writable
     public InputSplit() {
     }
 
     public InputSplit(TableDescriptor htd, HRegionInfo regionInfo, List<String> locations,
-      Scan scan, Path restoreDir) {
+      Scan scan, Path restoreDir, long length) {
       this.htd = htd;
       this.regionInfo = regionInfo;
       if (locations == null || locations.isEmpty()) {
@@ -166,6 +170,7 @@ public InputSplit(TableDescriptor htd, HRegionInfo regionInfo, List<String> loca
       }
 
       this.restoreDir = restoreDir.toString();
+      this.length = length;
     }
 
     public TableDescriptor getHtd() {
@@ -181,8 +186,7 @@ public String getRestoreDir() {
     }
 
     public long getLength() {
-      // TODO: We can obtain the file sizes of the snapshot here.
-      return 0;
+      return length;
     }
 
     public String[] getLocations() {
@@ -219,6 +223,7 @@ public void write(DataOutput out) throws IOException {
 
       Bytes.writeByteArray(out, Bytes.toBytes(scan));
       Bytes.writeByteArray(out, Bytes.toBytes(restoreDir));
+      out.writeLong(length);
 
     }
 
@@ -235,6 +240,12 @@ public void readFields(DataInput in) throws IOException {
 
       this.scan = Bytes.toString(Bytes.readByteArray(in));
       this.restoreDir = Bytes.toString(Bytes.readByteArray(in));
+      try {
+        this.length = in.readLong();
+      } catch (EOFException e) {
+        // Older serialized splits do not carry length and keep the previous behavior
+        this.length = 0L;
+      }
     }
   }
 
@@ -441,6 +452,8 @@ public static List<InputSplit> getSplits(Scan scan, SnapshotManifest manifest,
     }
 
     List<InputSplit> splits = new ArrayList<>();
+    RegionSizes regionSizes =
+      SnapshotRegionSizeCalculator.calculateRegionSizes(conf, manifest);
     for (HRegionInfo hri : regionManifests) {
       // load region descriptor
       List<String> hosts = null;
@@ -456,6 +469,8 @@ public static List<InputSplit> getSplits(Scan scan, SnapshotManifest manifest,
         }
       }
 
+      long snapshotRegionSize = regionSizes.getRegionSize(hri.getEncodedName());
+
       if (numSplits > 1) {
         byte[][] sp = sa.split(hri.getStartKey(), hri.getEndKey(), numSplits, true);
         for (int i = 0; i < sp.length - 1; i++) {
@@ -478,7 +493,8 @@ public static List<InputSplit> getSplits(Scan scan, SnapshotManifest manifest,
                 Bytes.compareTo(scan.getStopRow(), sp[i + 1]) < 0 ? scan.getStopRow() : sp[i + 1]);
             }
 
-            splits.add(new InputSplit(htd, hri, hosts, boundedScan, restoreDir));
+            splits.add(new InputSplit(htd, hri, hosts, boundedScan, restoreDir,
+              snapshotRegionSize / numSplits));
           }
         }
       } else {
@@ -487,7 +503,7 @@ public static List<InputSplit> getSplits(Scan scan, SnapshotManifest manifest,
             hri.getEndKey())
         ) {
 
-          splits.add(new InputSplit(htd, hri, hosts, scan, restoreDir));
+          splits.add(new InputSplit(htd, hri, hosts, scan, restoreDir, snapshotRegionSize));
         }
       }
     }

hbase-server/src/main/java/org/apache/hadoop/hbase/snapshot/RegionSizes.java

@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.hadoop.hbase.snapshot;
+
+import java.util.Collections;
+import java.util.Map;
+import org.apache.yetus.audience.InterfaceAudience;
+
+/**
+ * Holds the size of each region in a snapshot.
+ */
+@InterfaceAudience.Private
+public class RegionSizes {
+
+  private final Map<String, Long> regionSizeMap;
+
+  RegionSizes(Map<String, Long> regionSizeMap) {
+    this.regionSizeMap = regionSizeMap;
+  }
+
+  /**
+   * Returns the total size of the given region in bytes.
+   * @param encodedRegionName the encoded region name
+   * @return the size of the region, or 0 if the region is not present in the snapshot
+   */
+  public long getRegionSize(String encodedRegionName) {
+    return regionSizeMap.getOrDefault(encodedRegionName, 0L);
+  }
+
+  /**
+   * Returns an unmodifiable view of the region size map.
+   * @return a map of region encoded names to their total size in bytes
+   */
+  public Map<String, Long> getRegionSizeMap() {
+    return Collections.unmodifiableMap(regionSizeMap);
+  }
+}

hbase-server/src/main/java/org/apache/hadoop/hbase/snapshot/SnapshotInfo.java

@@ -160,6 +160,7 @@ String getStateToString() {
     private AtomicLong hfilesMobSize = new AtomicLong();
     private AtomicLong nonSharedHfilesArchiveSize = new AtomicLong();
     private AtomicLong logSize = new AtomicLong();
+    private Map<String, Long> regionSizeMap = new ConcurrentHashMap<>();
 
     private final SnapshotProtos.SnapshotDescription snapshot;
     private final TableName snapshotTable;
@@ -182,6 +183,14 @@ String getStateToString() {
       this.fs = fs;
     }
 
+    /**
+     * Returns the map containing region sizes.
+     * @return A map where keys are region names and values are their corresponding sizes.
+     */
+    public Map<String, Long> getRegionSizeMap() {
+      return regionSizeMap;
+    }
+
     /** Returns the snapshot descriptor */
     public SnapshotDescription getSnapshotDescription() {
       return ProtobufUtil.createSnapshotDesc(this.snapshot);
@@ -353,6 +362,12 @@ FileInfo addStoreFile(final RegionInfo region, final String family,
       return new FileInfo(inArchive, size, isCorrupted);
     }
 
+    void updateRegionSizeMap(final RegionInfo region,
+      final SnapshotRegionManifest.StoreFile storeFile) {
+      long currentSize = regionSizeMap.getOrDefault(region.getEncodedName(), 0L);
+      regionSizeMap.put(region.getEncodedName(), currentSize + storeFile.getFileSize());
+    }
+
     /**
      * Add the specified log file to the stats
      * @param server  server name
@@ -633,6 +648,7 @@ public void storeFile(final RegionInfo regionInfo, final String family,
           final SnapshotRegionManifest.StoreFile storeFile) throws IOException {
           if (!storeFile.hasReference()) {
             stats.addStoreFile(regionInfo, family, storeFile, filesMap);
+            stats.updateRegionSizeMap(regionInfo, storeFile);
           }
         }
       });

hbase-server/src/main/java/org/apache/hadoop/hbase/snapshot/SnapshotRegionSizeCalculator.java

@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.hadoop.hbase.snapshot;
+
+import java.io.IOException;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.hbase.snapshot.SnapshotInfo.SnapshotStats;
+import org.apache.hadoop.hbase.util.CommonFSUtils;
+import org.apache.yetus.audience.InterfaceAudience;
+
+import org.apache.hadoop.hbase.shaded.protobuf.generated.SnapshotProtos;
+
+/**
+ * Utility class to calculate the size of each region in a snapshot.
+ */
+@InterfaceAudience.Private
+public class SnapshotRegionSizeCalculator {
+
+  /**
+   * Calculate the size of each region in the snapshot.
+   * @return A map of region encoded names to their total size in bytes.
+   * @throws IOException If an error occurs during calculation.
+   */
+  public static RegionSizes calculateRegionSizes(Configuration conf,
+    SnapshotManifest manifest) throws IOException {
+    FileSystem fs = FileSystem.get(CommonFSUtils.getRootDir(conf).toUri(), conf);
+    SnapshotProtos.SnapshotDescription snapshot =
+      SnapshotDescriptionUtils.readSnapshotInfo(fs, manifest.getSnapshotDir());
+    SnapshotStats stats = SnapshotInfo.getSnapshotStats(conf, snapshot, null);
+    return new RegionSizes(stats.getRegionSizeMap());
+  }
+
+}