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-common/src/main/java/org/apache/hadoop/hbase/io/compress/CanReinit.java

@@ -0,0 +1,31 @@
+/*
+ * 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.io.compress;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.yetus.audience.InterfaceAudience;
+
+/**
+ * This is a marker interface that indicates if a compressor or decompressor
+ * type can support reinitialization via reinit(Configuration conf).
+ */
+@InterfaceAudience.Private
+public interface CanReinit {
+
+  void reinit(Configuration conf);
+
+}

hbase-common/src/main/java/org/apache/hadoop/hbase/io/encoding/AbstractDataBlockEncoder.java

@@ -18,6 +18,8 @@
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
+
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.ByteBufferKeyOnlyKeyValue;
 import org.apache.hadoop.hbase.Cell;
 import org.apache.hadoop.hbase.KeyValue;
@@ -35,8 +37,8 @@ public HFileBlockEncodingContext newDataBlockEncodingContext(
   }
 
   @Override
-  public HFileBlockDecodingContext newDataBlockDecodingContext(HFileContext meta) {
-    return new HFileBlockDefaultDecodingContext(meta);
+  public HFileBlockDecodingContext newDataBlockDecodingContext(Configuration conf, HFileContext meta) {
+    return new HFileBlockDefaultDecodingContext(conf, meta);
   }
 
   protected void postEncoding(HFileBlockEncodingContext encodingCtx)

hbase-common/src/main/java/org/apache/hadoop/hbase/io/encoding/DataBlockEncoder.java

@@ -20,6 +20,8 @@
 import java.io.DataOutputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
+
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.Cell;
 import org.apache.hadoop.hbase.CellComparator;
 import org.apache.hadoop.hbase.io.hfile.HFileContext;
@@ -107,11 +109,13 @@ HFileBlockEncodingContext newDataBlockEncodingContext(
    * Creates an encoder specific decoding context, which will prepare the data
    * before actual decoding
    *
+   * @param conf
+   *          store configuration
    * @param meta
    *          HFile meta data        
    * @return a newly created decoding context
    */
-  HFileBlockDecodingContext newDataBlockDecodingContext(HFileContext meta);
+  HFileBlockDecodingContext newDataBlockDecodingContext(Configuration conf, HFileContext meta);
 
   /**
    * An interface which enable to seek while underlying data is encoded.

hbase-common/src/main/java/org/apache/hadoop/hbase/io/encoding/EncodedDataBlock.java

@@ -27,6 +27,7 @@
 import java.util.Iterator;
 import java.util.List;
 import org.apache.commons.lang3.NotImplementedException;
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.Cell;
 import org.apache.hadoop.hbase.HConstants;
 import org.apache.hadoop.hbase.KeyValue;
@@ -57,6 +58,7 @@ public class EncodedDataBlock {
   private HFileContext meta;
 
   private final DataBlockEncoding encoding;
+  private final Configuration conf;
 
   // The is for one situation that there are some cells includes tags and others are not.
   // isTagsLenZero stores if cell tags length is zero before doing encoding since we need
@@ -70,11 +72,12 @@ public class EncodedDataBlock {
    * Create a buffer which will be encoded using dataBlockEncoder.
    * @param dataBlockEncoder Algorithm used for compression.
    * @param encoding encoding type used
-   * @param rawKVs
-   * @param meta
+   * @param rawKVs raw KVs
+   * @param meta hfile context
+   * @param conf store configuration
    */
   public EncodedDataBlock(DataBlockEncoder dataBlockEncoder, DataBlockEncoding encoding,
-      byte[] rawKVs, HFileContext meta) {
+      byte[] rawKVs, HFileContext meta, Configuration conf) {
     Preconditions.checkNotNull(encoding,
         "Cannot create encoded data block with null encoder");
     this.dataBlockEncoder = dataBlockEncoder;
@@ -83,6 +86,7 @@ public EncodedDataBlock(DataBlockEncoder dataBlockEncoder, DataBlockEncoding enc
         HConstants.HFILEBLOCK_DUMMY_HEADER, meta);
     this.rawKVs = rawKVs;
     this.meta = meta;
+    this.conf = conf;
   }
 
   /**
@@ -115,7 +119,7 @@ public Cell next() {
         if (decompressedData == null) {
           try {
             decompressedData = dataBlockEncoder.decodeKeyValues(dis, dataBlockEncoder
-                .newDataBlockDecodingContext(meta));
+                .newDataBlockDecodingContext(conf, meta));
           } catch (IOException e) {
             throw new RuntimeException("Problem with data block encoder, " +
                 "most likely it requested more bytes than are available.", e);

hbase-common/src/main/java/org/apache/hadoop/hbase/io/encoding/HFileBlockDefaultDecodingContext.java

@@ -20,8 +20,10 @@
 import java.io.IOException;
 import java.io.InputStream;
 import org.apache.commons.io.IOUtils;
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.io.ByteBuffInputStream;
 import org.apache.hadoop.hbase.io.TagCompressionContext;
+import org.apache.hadoop.hbase.io.compress.CanReinit;
 import org.apache.hadoop.hbase.io.compress.Compression;
 import org.apache.hadoop.hbase.io.crypto.Cipher;
 import org.apache.hadoop.hbase.io.crypto.Decryptor;
@@ -30,6 +32,7 @@
 import org.apache.hadoop.hbase.io.util.BlockIOUtils;
 import org.apache.hadoop.hbase.nio.ByteBuff;
 import org.apache.hadoop.hbase.util.Bytes;
+import org.apache.hadoop.io.compress.Decompressor;
 import org.apache.yetus.audience.InterfaceAudience;
 
 /**
@@ -41,10 +44,12 @@
  */
 @InterfaceAudience.Private
 public class HFileBlockDefaultDecodingContext implements HFileBlockDecodingContext {
+  private final Configuration conf;
   private final HFileContext fileContext;
   private TagCompressionContext tagCompressionContext;
 
-  public HFileBlockDefaultDecodingContext(HFileContext fileContext) {
+  public HFileBlockDefaultDecodingContext(Configuration conf, HFileContext fileContext) {
+    this.conf = conf;
     this.fileContext = fileContext;
   }
 
@@ -87,8 +92,24 @@ public void prepareDecoding(int onDiskSizeWithoutHeader, int uncompressedSizeWit
 
       Compression.Algorithm compression = fileContext.getCompression();
       if (compression != Compression.Algorithm.NONE) {
-        Compression.decompress(blockBufferWithoutHeader, dataInputStream,
-          uncompressedSizeWithoutHeader, compression);
+        Decompressor decompressor = null;
+        try {
+          decompressor = compression.getDecompressor();
+          // Some algorithms don't return decompressors and accept null as a valid parameter for
+          // same when creating decompression streams. We can ignore these cases wrt reinit.
+          if (decompressor != null && decompressor instanceof CanReinit) {
+            ((CanReinit)decompressor).reinit(conf);
+          }
+          try (InputStream is =
+              compression.createDecompressionStream(dataInputStream, decompressor, 0)) {
+            BlockIOUtils.readFullyWithHeapBuffer(is, blockBufferWithoutHeader,
+              uncompressedSizeWithoutHeader);
+          }
+        } finally {
+          if (decompressor != null) {
+            compression.returnDecompressor(decompressor);
+          }
+        }
       } else {
         BlockIOUtils.readFullyWithHeapBuffer(dataInputStream, blockBufferWithoutHeader,
           onDiskSizeWithoutHeader);

hbase-compression/hbase-compression-aircompressor/src/main/java/org/apache/hadoop/hbase/io/compress/aircompressor/HadoopCompressor.java

@@ -19,6 +19,7 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.hbase.io.compress.CanReinit;
 import org.apache.hadoop.hbase.io.compress.CompressionUtil;
 import org.apache.yetus.audience.InterfaceAudience;
 import org.slf4j.Logger;
@@ -30,7 +31,7 @@
  */
 @InterfaceAudience.Private
 public abstract class HadoopCompressor<T extends Compressor>
-    implements org.apache.hadoop.io.compress.Compressor {
+    implements CanReinit, org.apache.hadoop.io.compress.Compressor {
 
   protected static final Logger LOG = LoggerFactory.getLogger(HadoopCompressor.class);
   protected T compressor;

hbase-compression/hbase-compression-lz4/src/main/java/org/apache/hadoop/hbase/io/compress/lz4/Lz4Compressor.java

@@ -19,6 +19,7 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.hbase.io.compress.CanReinit;
 import org.apache.hadoop.hbase.io.compress.CompressionUtil;
 import org.apache.hadoop.io.compress.Compressor;
 import org.apache.yetus.audience.InterfaceAudience;
@@ -31,7 +32,7 @@
  * Hadoop compressor glue for lz4-java.
  */
 @InterfaceAudience.Private
-public class Lz4Compressor implements Compressor {
+public class Lz4Compressor implements CanReinit, Compressor {
 
   protected static final Logger LOG = LoggerFactory.getLogger(Lz4Compressor.class);
   protected LZ4Compressor compressor;

hbase-compression/hbase-compression-snappy/src/main/java/org/apache/hadoop/hbase/io/compress/xerial/SnappyCompressor.java

@@ -19,6 +19,7 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.hbase.io.compress.CanReinit;
 import org.apache.hadoop.hbase.io.compress.CompressionUtil;
 import org.apache.hadoop.io.compress.Compressor;
 import org.apache.yetus.audience.InterfaceAudience;
@@ -30,7 +31,7 @@
  * Hadoop compressor glue for Xerial Snappy.
  */
 @InterfaceAudience.Private
-public class SnappyCompressor implements Compressor {
+public class SnappyCompressor implements CanReinit, Compressor {
 
   protected static final Logger LOG = LoggerFactory.getLogger(SnappyCompressor.class);
   protected ByteBuffer inBuf, outBuf;

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 || s.isEmpty()) {
+      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);
+    }
+  }
+
 }