LUCENE-8739: custom codec providing Zstandard compression/decompression

gradle/testing/randomization/policies/tests.policy

@@ -58,6 +58,10 @@ grant {
   // needed by cyberneko usage by benchmarks on J9
   permission java.lang.RuntimePermission "accessClassInPackage.org.apache.xerces.util";
   permission java.lang.RuntimePermission "getClassLoader";
+  permission java.lang.RuntimePermission "setContextClassLoader";
+  // permission needed to load Zstd JNI library
+  permission java.lang.RuntimePermission "loadLibrary.*";
+
 
   // Needed for loading native library (lucene:misc:native) in lucene:misc
   permission java.lang.RuntimePermission "getFileStoreAttributes";

lucene/codecs/build.gradle

@@ -21,5 +21,7 @@ description = 'Lucene codecs and postings formats'
 
 dependencies {
     moduleImplementation project(':lucene:core')
+    moduleImplementation 'com.github.luben:zstd-jni:1.5.0-4'
+    testImplementation project(':lucene:test-framework')
     moduleTestImplementation project(':lucene:test-framework')
 }

lucene/codecs/src/java/module-info.java

@@ -16,8 +16,10 @@
  */
 
 /** Lucene codecs and postings formats */
+@SuppressWarnings({"requires-automatic", "requires-transitive-automatic"})
 module org.apache.lucene.codecs {
   requires org.apache.lucene.core;
+  requires com.github.luben.zstd_jni;
 
   exports org.apache.lucene.codecs.blockterms;
   exports org.apache.lucene.codecs.blocktreeords;
@@ -26,6 +28,7 @@
   exports org.apache.lucene.codecs.simpletext;
   exports org.apache.lucene.codecs.uniformsplit;
   exports org.apache.lucene.codecs.uniformsplit.sharedterms;
+  exports org.apache.lucene.codecs.customcompression;
 
   provides org.apache.lucene.codecs.PostingsFormat with
       org.apache.lucene.codecs.blocktreeords.BlockTreeOrdsPostingsFormat,
@@ -35,5 +38,6 @@
       org.apache.lucene.codecs.uniformsplit.UniformSplitPostingsFormat,
       org.apache.lucene.codecs.uniformsplit.sharedterms.STUniformSplitPostingsFormat;
   provides org.apache.lucene.codecs.Codec with
-      org.apache.lucene.codecs.simpletext.SimpleTextCodec;
+      org.apache.lucene.codecs.simpletext.SimpleTextCodec,
+      org.apache.lucene.codecs.customcompression.Lucene90CustomCompressionCodec;
 }

lucene/codecs/src/java/org/apache/lucene/codecs/customcompression/Lucene90CustomCompressionCodec.java

@@ -0,0 +1,83 @@
+/*
+ * 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.lucene.codecs.customcompression;
+
+import org.apache.lucene.codecs.FilterCodec;
+import org.apache.lucene.codecs.StoredFieldsFormat;
+import org.apache.lucene.codecs.lucene90.Lucene90Codec;
+
+/** Custom codec for different compression algorithm */
+public final class Lucene90CustomCompressionCodec extends FilterCodec {
+
+  private final StoredFieldsFormat storedFieldsFormat;
+  private int compressionLevel;
+  public static final int defaultCompressionLevel = 6;
+
+  /** Compression modes */
+  public static enum Mode {
+
+    // Currently Zstandard is supported, other compression algorithms can be implemented and
+    // respective modes can be added for e.g. ZSTD, ZSTD_DICT, ZSTD_FAST or any other compression
+    // algos
+
+    // Zstandard without dictionary
+    ZSTD,
+    // Zstandard with dictionary
+    ZSTD_DICT
+  }
+  /** Default codec */
+  public Lucene90CustomCompressionCodec() {
+    this(Mode.ZSTD_DICT, defaultCompressionLevel);
+  }
+
+  /** new codec for a given compression algorithm and compression level */
+  public Lucene90CustomCompressionCodec(Mode compressionMode, int compressionLevel) {
+    super("Lucene90CustomCompression", new Lucene90Codec());
+    this.compressionLevel = compressionLevel;
+
+    switch (compressionMode) {
+      case ZSTD:
+        if (this.compressionLevel < 1 || this.compressionLevel > 22)
+          throw new IllegalArgumentException("Invalid compression level");
+
+        this.storedFieldsFormat =
+            new Lucene90CustomCompressionStoredFieldsFormat(Mode.ZSTD, compressionLevel);
+        break;
+
+      case ZSTD_DICT:
+        if (this.compressionLevel < 1 || this.compressionLevel > 22)
+          throw new IllegalArgumentException("Invalid compression level");
+
+        this.storedFieldsFormat =
+            new Lucene90CustomCompressionStoredFieldsFormat(Mode.ZSTD_DICT, compressionLevel);
+        break;
+
+      default:
+        throw new IllegalArgumentException("Chosen compression mode does not exist");
+    }
+  }
+
+  @Override
+  public StoredFieldsFormat storedFieldsFormat() {
+    return storedFieldsFormat;
+  }
+
+  @Override
+  public String toString() {
+    return getClass().getSimpleName();
+  }
+}

lucene/codecs/src/java/org/apache/lucene/codecs/customcompression/Lucene90CustomCompressionStoredFieldsFormat.java

@@ -0,0 +1,102 @@
+/*
+ * 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.lucene.codecs.customcompression;
+
+import java.io.IOException;
+import java.util.Objects;
+import org.apache.lucene.codecs.StoredFieldsFormat;
+import org.apache.lucene.codecs.StoredFieldsReader;
+import org.apache.lucene.codecs.StoredFieldsWriter;
+import org.apache.lucene.codecs.compressing.CompressionMode;
+import org.apache.lucene.codecs.lucene90.compressing.Lucene90CompressingStoredFieldsFormat;
+import org.apache.lucene.index.FieldInfos;
+import org.apache.lucene.index.SegmentInfo;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.IOContext;
+
+/** Stored field format used by plugaable codec */
+public class Lucene90CustomCompressionStoredFieldsFormat extends StoredFieldsFormat {
+  // chunk size for zstandard
+  private static final int ZSTD_BLOCK_LENGTH = 10 * 48 * 1024;
+
+  public static final String MODE_KEY =
+      Lucene90CustomCompressionStoredFieldsFormat.class.getSimpleName() + ".mode";
+
+  final Lucene90CustomCompressionCodec.Mode mode;
+
+  private int compressionLevel;
+
+  /** default constructor */
+  public Lucene90CustomCompressionStoredFieldsFormat() {
+    this(
+        Lucene90CustomCompressionCodec.Mode.ZSTD_DICT,
+        Lucene90CustomCompressionCodec.defaultCompressionLevel);
+  }
+
+  /** Stored fields format with specified compression algo. */
+  public Lucene90CustomCompressionStoredFieldsFormat(
+      Lucene90CustomCompressionCodec.Mode mode, int compressionLevel) {
+    this.mode = Objects.requireNonNull(mode);
+    this.compressionLevel = compressionLevel;
+  }
+
+  @Override
+  public StoredFieldsReader fieldsReader(
+      Directory directory, SegmentInfo si, FieldInfos fn, IOContext context) throws IOException {
+    String value = si.getAttribute(MODE_KEY);
+    if (value == null) {
+      throw new IllegalStateException("missing value for " + MODE_KEY + " for segment: " + si.name);
+    }
+    Lucene90CustomCompressionCodec.Mode mode = Lucene90CustomCompressionCodec.Mode.valueOf(value);
+    return impl(mode).fieldsReader(directory, si, fn, context);
+  }
+
+  @Override
+  public StoredFieldsWriter fieldsWriter(Directory directory, SegmentInfo si, IOContext context)
+      throws IOException {
+    String previous = si.putAttribute(MODE_KEY, mode.name());
+    if (previous != null && previous.equals(mode.name()) == false) {
+      throw new IllegalStateException(
+          "found existing value for "
+              + MODE_KEY
+              + " for segment: "
+              + si.name
+              + "old="
+              + previous
+              + ", new="
+              + mode.name());
+    }
+    return impl(mode).fieldsWriter(directory, si, context);
+  }
+
+  StoredFieldsFormat impl(Lucene90CustomCompressionCodec.Mode mode) {
+    switch (mode) {
+      case ZSTD:
+        return new Lucene90CompressingStoredFieldsFormat(
+            "CustomCompressionStoredFieldsZSTD", ZSTD_MODE_NO_DICT, ZSTD_BLOCK_LENGTH, 4096, 10);
+
+      case ZSTD_DICT:
+        return new Lucene90CompressingStoredFieldsFormat(
+            "CustomCompressionStoredFieldsZSTD", ZSTD_MODE_DICT, ZSTD_BLOCK_LENGTH, 4096, 10);
+      default:
+        throw new AssertionError();
+    }
+  }
+
+  public final CompressionMode ZSTD_MODE_NO_DICT = new ZstdCompressionMode(compressionLevel);
+  public final CompressionMode ZSTD_MODE_DICT = new ZstdDictCompressionMode(compressionLevel);
+}