HBASE-26353 Support loadable dictionaries in hbase-compression-zstd

ZStandard supports initialization of compressors and decompressors with a
precomputed dictionary, which can dramatically improve and speed up compression
of tables with small values. For more details, please see

  The Case For Small Data Compression
  https://github.com/facebook/zstd#the-case-for-small-data-compression

Author: apurtell

hbase-compression/hbase-compression-zstd/src/main/java/org/apache/hadoop/hbase/io/compress/zstd/ZstdCodec.java

@@ -16,20 +16,32 @@
  */
 package org.apache.hadoop.hbase.io.compress.zstd;
 
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
 import org.apache.hadoop.conf.Configurable;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.CommonConfigurationKeys;
+import org.apache.hadoop.fs.FSDataInputStream;
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.io.compress.BlockCompressorStream;
 import org.apache.hadoop.io.compress.BlockDecompressorStream;
 import org.apache.hadoop.io.compress.CompressionCodec;
 import org.apache.hadoop.io.compress.CompressionInputStream;
 import org.apache.hadoop.io.compress.CompressionOutputStream;
 import org.apache.hadoop.io.compress.Compressor;
 import org.apache.hadoop.io.compress.Decompressor;
+import org.apache.hbase.thirdparty.com.google.common.cache.CacheBuilder;
+import org.apache.hbase.thirdparty.com.google.common.cache.CacheLoader;
+import org.apache.hbase.thirdparty.com.google.common.cache.LoadingCache;
 import org.apache.yetus.audience.InterfaceAudience;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * Hadoop ZStandard codec implemented with zstd-jni.
@@ -41,7 +53,13 @@ public class ZstdCodec implements Configurable, CompressionCodec {
 
   public static final String ZSTD_LEVEL_KEY = "hbase.io.compress.zstd.level";
   public static final String ZSTD_BUFFER_SIZE_KEY = "hbase.io.compress.zstd.buffersize";
+  public static final String ZSTD_DICTIONARY_KEY = "hbase.io.compress.zstd.dictionary";
+  public static final String ZSTD_DICTIONARY_MAX_SIZE_KEY =
+    "hbase.io.compress.zstd.dictionary.max.size";
+  public static final int DEFAULT_ZSTD_DICTIONARY_MAX_SIZE = 10 * 1024 * 1024;
 
+  private static final Logger LOG = LoggerFactory.getLogger(ZstdCodec.class);
+  private static volatile LoadingCache<String, byte[]> DICTIONARY_CACHE;
   private Configuration conf;
 
   public ZstdCodec() {
@@ -60,12 +78,12 @@ public void setConf(Configuration conf) {
 
   @Override
   public Compressor createCompressor() {
-    return new ZstdCompressor(getLevel(conf), getBufferSize(conf));
+    return new ZstdCompressor(getLevel(conf), getBufferSize(conf), getDictionary(conf));
   }
 
   @Override
   public Decompressor createDecompressor() {
-    return new ZstdDecompressor(getBufferSize(conf));
+    return new ZstdDecompressor(getBufferSize(conf), getDictionary(conf));
   }
 
   @Override
@@ -123,4 +141,59 @@ static int getBufferSize(Configuration conf) {
     return size > 0 ? size : 256 * 1024; // Don't change this default
   }
 
+  static byte[] getDictionary(final Configuration conf) {
+    // Get the dictionary path, if set
+    final String s = conf.get(ZSTD_DICTIONARY_KEY);
+    if (s == null) {
+      return null;
+    }
+
+    // Create the dictionary loading cache if we haven't already
+    if (DICTIONARY_CACHE == null) {
+      synchronized (ZstdCodec.class) {
+        if (DICTIONARY_CACHE == null) {
+          DICTIONARY_CACHE = CacheBuilder.newBuilder()
+            .expireAfterAccess(1, TimeUnit.HOURS)
+            .build(
+              new CacheLoader<String, byte[]>() {
+                public byte[] load(String s) throws Exception {
+                  final Path path = new Path(s);
+                  final FileSystem fs = FileSystem.get(path.toUri(), conf);
+                  final FileStatus stat = fs.getFileStatus(path);
+                  if (!stat.isFile()) {
+                    throw new IllegalArgumentException(s + " is not a file");
+                  }
+                  final int limit = conf.getInt(ZSTD_DICTIONARY_MAX_SIZE_KEY,
+                    DEFAULT_ZSTD_DICTIONARY_MAX_SIZE);
+                  if (stat.getLen() > limit) {
+                    throw new IllegalArgumentException("Dictionary " + s + " is too large" +
+                      ", size=" + stat.getLen() + ", limit=" + limit);
+                  }
+                  final ByteArrayOutputStream baos = new ByteArrayOutputStream();
+                  final byte[] buffer = new byte[8192];
+                  try (final FSDataInputStream in = fs.open(path)) {
+                    int n;
+                    do {
+                      n = in.read(buffer);
+                      if (n > 0) {
+                        baos.write(buffer, 0, n);
+                         }
+                    } while (n > 0);
+                  }
+                  LOG.info("Loaded {} from {} (size {})", ZSTD_DICTIONARY_KEY, s, stat.getLen());
+                  return baos.toByteArray();
+                }
+            });
+        }
+      }
+    }
+
+    // Get or load the dictionary for the given path
+    try {
+      return DICTIONARY_CACHE.get(s);
+    } catch (ExecutionException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
 }

hbase-compression/hbase-compression-zstd/src/main/java/org/apache/hadoop/hbase/io/compress/zstd/ZstdCompressor.java

@@ -25,6 +25,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import com.github.luben.zstd.Zstd;
+import com.github.luben.zstd.ZstdDictCompress;
 
 /**
  * Hadoop compressor glue for zstd-jni.
@@ -37,17 +38,25 @@ public class ZstdCompressor implements Compressor {
   protected ByteBuffer inBuf, outBuf;
   protected boolean finish, finished;
   protected long bytesRead, bytesWritten;
+  protected ZstdDictCompress dict;
 
-  ZstdCompressor(int level, int bufferSize) {
+  ZstdCompressor(final int level, final int bufferSize, final byte[] dict) {
     this.level = level;
     this.bufferSize = bufferSize;
     this.inBuf = ByteBuffer.allocateDirect(bufferSize);
     this.outBuf = ByteBuffer.allocateDirect(bufferSize);
     this.outBuf.position(bufferSize);
+    if (dict != null) {
+      this.dict = new ZstdDictCompress(dict, level);
+    }
+  }
+
+  ZstdCompressor(final int level, final int bufferSize) {
+    this(level, bufferSize, null);
   }
 
   @Override
-  public int compress(byte[] b, int off, int len) throws IOException {
+  public int compress(final byte[] b, final int off, final int len) throws IOException {
     // If we have previously compressed our input and still have some buffered bytes
     // remaining, provide them to the caller.
     if (outBuf.hasRemaining()) {
@@ -71,7 +80,12 @@ public int compress(byte[] b, int off, int len) throws IOException {
         } else {
           outBuf.clear();
         }
-        int written = Zstd.compress(outBuf, inBuf, level, true);
+        int written;
+        if (dict != null) {
+          written = Zstd.compress(outBuf, inBuf, dict);
+        } else {
+          written = Zstd.compress(outBuf, inBuf, level);
+        }
         bytesWritten += written;
         inBuf.clear();
         LOG.trace("compress: compressed {} -> {} (level {})", uncompressed, written, level);
@@ -125,11 +139,16 @@ public boolean needsInput() {
   }
 
   @Override
-  public void reinit(Configuration conf) {
+  public void reinit(final Configuration conf) {
     LOG.trace("reinit");
     if (conf != null) {
       // Level might have changed
       level = ZstdCodec.getLevel(conf);
+      // Dictionary may have changed
+      byte[] b = ZstdCodec.getDictionary(conf);
+      if (b != null) {
+        dict = new ZstdDictCompress(b, level);
+      }
       // Buffer size might have changed
       int newBufferSize = ZstdCodec.getBufferSize(conf);
       if (bufferSize != newBufferSize) {
@@ -154,12 +173,12 @@ public void reset() {
   }
 
   @Override
-  public void setDictionary(byte[] b, int off, int len) {
+  public void setDictionary(final byte[] b, final int off, final int len) {
     throw new UnsupportedOperationException("setDictionary is not supported");
   }
 
   @Override
-  public void setInput(byte[] b, int off, int len) {
+  public void setInput(final byte[] b, final int off, final int len) {
     LOG.trace("setInput: off={} len={}", off, len);
     if (inBuf.remaining() < len) {
       // Get a new buffer that can accomodate the accumulated input plus the additional
@@ -179,7 +198,7 @@ public void setInput(byte[] b, int off, int len) {
 
   // Package private
 
-  int maxCompressedLength(int len) {
+  int maxCompressedLength(final int len) {
     return (int) Zstd.compressBound(len);
   }
 

hbase-compression/hbase-compression-zstd/src/main/java/org/apache/hadoop/hbase/io/compress/zstd/ZstdDecompressor.java

@@ -24,6 +24,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import com.github.luben.zstd.Zstd;
+import com.github.luben.zstd.ZstdDictDecompress;
 
 /**
  * Hadoop decompressor glue for zstd-java.
@@ -35,15 +36,23 @@ public class ZstdDecompressor implements Decompressor {
   protected ByteBuffer inBuf, outBuf;
   protected int inLen;
   protected boolean finished;
+  protected ZstdDictDecompress dict;
 
-  ZstdDecompressor(int bufferSize) {
+  ZstdDecompressor(final int bufferSize, final byte[] dictionary) {
     this.inBuf = ByteBuffer.allocateDirect(bufferSize);
     this.outBuf = ByteBuffer.allocateDirect(bufferSize);
     this.outBuf.position(bufferSize);
+    if (dictionary != null) {
+      this.dict = new ZstdDictDecompress(dictionary);
+    }
+  }
+
+  ZstdDecompressor(final int bufferSize) {
+    this(bufferSize, null);
   }
 
   @Override
-  public int decompress(byte[] b, int off, int len) throws IOException {
+  public int decompress(final byte[] b, final int off, final int len) throws IOException {
     if (outBuf.hasRemaining()) {
       int remaining = outBuf.remaining(), n = Math.min(remaining, len);
       outBuf.get(b, off, n);
@@ -55,7 +64,12 @@ public int decompress(byte[] b, int off, int len) throws IOException {
       int remaining = inBuf.remaining();
       inLen -= remaining;
       outBuf.clear();
-      int written = Zstd.decompress(outBuf, inBuf);
+      int written;
+      if (dict != null) {
+        written = Zstd.decompress(outBuf, inBuf, dict);
+      } else {
+        written = Zstd.decompress(outBuf, inBuf);
+      }
       inBuf.clear();
       LOG.trace("decompress: decompressed {} -> {}", remaining, written);
       outBuf.flip();
@@ -104,24 +118,24 @@ public void reset() {
 
   @Override
   public boolean needsInput() {
-    boolean b = (inBuf.position() == 0);
+    final boolean b = (inBuf.position() == 0);
     LOG.trace("needsInput: {}", b);
     return b;
   }
 
   @Override
-  public void setDictionary(byte[] b, int off, int len) {
+  public void setDictionary(final byte[] b, final int off, final int len) {
     throw new UnsupportedOperationException("setDictionary is not supported");
   }
 
   @Override
-  public void setInput(byte[] b, int off, int len) {
+  public void setInput(final byte[] b, final int off, final int len) {
     LOG.trace("setInput: off={} len={}", off, len);
     if (inBuf.remaining() < len) {
       // Get a new buffer that can accomodate the accumulated input plus the additional
       // input that would cause a buffer overflow without reallocation.
       // This condition should be fortunately rare, because it is expensive.
-      int needed = CompressionUtil.roundInt2(inBuf.capacity() + len);
+      final int needed = CompressionUtil.roundInt2(inBuf.capacity() + len);
       LOG.trace("setInput: resize inBuf {}", needed);
       ByteBuffer newBuf = ByteBuffer.allocateDirect(needed);
       inBuf.flip();