luxugang / lucene-7.x-9.x Goto Github PK

public class ZstdWithPresetDictCompressionMode extends CompressionMode {
    private static final int NUM_SUB_BLOCKS = 10;
    private final int level;
    public static final int defaultLevel = 3;
    private static final int DICT_SIZE_FACTOR = 6;

    /**
     * default constructor
     */
    ZstdWithPresetDictCompressionMode() {
        this.level = defaultLevel;
    }

    /**
     * compression mode for a given compression level
     */
    ZstdWithPresetDictCompressionMode(int level) {
        this.level = level;
    }

    @Override
    public Compressor newCompressor() {
        return new ZSTDCompressor(level);
    }

    @Override
    public Decompressor newDecompressor() {
        return new ZSTDDecompressor();
    }

    /**
     * zstandard compressor
     */
    private static final class ZSTDCompressor extends Compressor {

        int compressionLevel;
        byte[] compressedBuffer;

        /**
         * compressor with a given compresion level
         */
        public ZSTDCompressor(int compressionLevel) {
            this.compressionLevel = compressionLevel;
            compressedBuffer = BytesRef.EMPTY_BYTES;
        }

        @Override
        public void close() throws IOException {
        }

        /*resuable compress function*/
        private void doCompress(byte[] bytes, int off, int len, ZstdCompressCtx cctx, DataOutput out)
            throws IOException {
            if (len == 0) {
                out.writeVInt(0);
                return;
            }
            final int maxCompressedLength = (int) Zstd.compressBound(len);
            compressedBuffer = ArrayUtil.grow(compressedBuffer, maxCompressedLength);

            int compressedSize =
                cctx.compressByteArray(compressedBuffer, 0, compressedBuffer.length, bytes, off, len);

            out.writeVInt(compressedSize);
            out.writeBytes(compressedBuffer, compressedSize);
        }

        @Override
        public void compress(byte[] bytes, int off, int len, DataOutput out) throws IOException {
            final int dictLength = len / (NUM_SUB_BLOCKS * DICT_SIZE_FACTOR);
            final int blockLength = (len - dictLength + NUM_SUB_BLOCKS - 1) / NUM_SUB_BLOCKS;

            // ZSTD写入魔数-666666
            out.writeVInt(Comporessions.ZSTD_MAGIC_NUMBER);

            out.writeVInt(dictLength);
            out.writeVInt(blockLength);

            final int end = off + len;

            try (ZstdCompressCtx cctx = new ZstdCompressCtx()) {
                cctx.setLevel(this.compressionLevel);

                // dictionary compression first
                doCompress(bytes, off, dictLength, cctx, out);
                cctx.loadDict(new ZstdDictCompress(bytes, off, dictLength, this.compressionLevel));

                for (int start = off + dictLength; start < end; start += blockLength) {
                    int l = Math.min(blockLength, off + len - start);
                    doCompress(bytes, start, l, cctx, out);
                }
            }
        }
    }

    /**
     * zstandard decompressor
     */
    private static final class ZSTDDecompressor extends Decompressor {

        private byte[] compressed;

        /**
         * default decompressor
         */
        public ZSTDDecompressor() {
            compressed = BytesRef.EMPTY_BYTES;
        }

        /*resuable decompress function*/
        private void doDecompress(
            DataInput in, ZstdDecompressCtx dctx, BytesRef bytes, int decompressedLen)
            throws IOException {
            final int compressedLength = in.readVInt();
            if (compressedLength == 0) {
                return;
            }

            compressed = ArrayUtil.grow(compressed, compressedLength);
            in.readBytes(compressed, 0, compressedLength);

            bytes.bytes = ArrayUtil.grow(bytes.bytes, bytes.length + decompressedLen);
            int uncompressed =
                dctx.decompressByteArray(
                    bytes.bytes, bytes.length, decompressedLen, compressed, 0, compressedLength);

            if (decompressedLen != uncompressed) {
                throw new IllegalStateException(decompressedLen + " " + uncompressed);
            }
            bytes.length += uncompressed;
        }

        @Override
        public void decompress(DataInput in, int originalLength, int offset, int length, BytesRef bytes)
            throws IOException {
            assert offset + length <= originalLength;

            if (length == 0) {
                bytes.length = 0;
                return;
            }

            final int magicNumber = in.readVInt();
            // 这里如果非-666666魔数，说明逻辑有问题
            if (magicNumber != Comporessions.ZSTD_MAGIC_NUMBER) {
                throw new IllegalArgumentException("Unknown magicNumber:[" + magicNumber + "]");
            }

            final int dictLength = in.readVInt();
            final int blockLength = in.readVInt();
            bytes.bytes = ArrayUtil.grow(bytes.bytes, dictLength);
            bytes.offset = bytes.length = 0;

            try (ZstdDecompressCtx dctx = new ZstdDecompressCtx()) {

                // decompress dictionary first
                doDecompress(in, dctx, bytes, dictLength);

                dctx.loadDict(new ZstdDictDecompress(bytes.bytes, 0, dictLength));

                int offsetInBlock = dictLength;
                int offsetInBytesRef = offset;

                // Skip unneeded blocks
                while (offsetInBlock + blockLength < offset) {
                    final int compressedLength = in.readVInt();
                    in.skipBytes(compressedLength);
                    offsetInBlock += blockLength;
                    offsetInBytesRef -= blockLength;
                }

                // Read blocks that intersect with the interval we need
                while (offsetInBlock < offset + length) {
                    bytes.bytes = ArrayUtil.grow(bytes.bytes, bytes.length + blockLength);
                    int l = Math.min(blockLength, originalLength - offsetInBlock);
                    doDecompress(in, dctx, bytes, l);
                    offsetInBlock += blockLength;
                }

                bytes.offset = offsetInBytesRef;
                bytes.length = length;
                assert bytes.isValid();
            }
        }

        @Override
        public Decompressor clone() {
            return new ZSTDDecompressor();
        }
    }
}

public class ZstdProxyDecompressor extends Decompressor {

    private final ZstdAdapterDecompressor zstdAdapterDecompressor;

    private final AdapterDecompressor adapterDecompressor;


    public ZstdProxyDecompressor(ZstdAdapterDecompressor zstdAdapterDecompressor, AdapterDecompressor adapterDecompressor) {
        this.zstdAdapterDecompressor = zstdAdapterDecompressor;
        this.adapterDecompressor = adapterDecompressor;
    }

    @Override
    public void decompress(DataInput in, int originalLength, int offset, int length, BytesRef bytes)
        throws IOException {

        assert offset + length <= originalLength;

        if (length == 0) {
            bytes.length = 0;
            return;
        }

        final int magicNumber = in.readVInt();
        // 这里如果非zstd压缩则使用原来的压缩算法进行解压
        if (magicNumber == Comporessions.ZSTD_MAGIC_NUMBER) {
            zstdAdapterDecompressor.decompress(in, originalLength, offset, length, bytes, magicNumber);
        } else {
            adapterDecompressor.decompress(in, originalLength, offset, length, bytes, magicNumber);
        }

    }

    @Override
    public Decompressor clone() {
        return new ZstdProxyDecompressor(this.zstdAdapterDecompressor.clone(), this.adapterDecompressor.clone());
    }

}

public final class LZ4AdapterDecompressor extends AdapterDecompressor {

    private int[] compressedLengths;
    private byte[] buffer;

    public LZ4AdapterDecompressor() {
        compressedLengths = new int[0];
        buffer = new byte[0];
    }

    private int readCompressedLengths(DataInput in, int originalLength, int dictLength, int blockLength) throws IOException {
        in.readVInt(); // compressed length of the dictionary, unused
        int totalLength = dictLength;
        int i = 0;
        while (totalLength < originalLength) {
            compressedLengths = ArrayUtil.grow(compressedLengths, i + 1);
            compressedLengths[i++] = in.readVInt();
            totalLength += blockLength;
        }
        return i;
    }

    @Override
    public void decompress(DataInput in, int originalLength, int offset, int length, BytesRef bytes, int firstVInt)
        throws IOException {
        assert offset + length <= originalLength;

        if (length == 0) {
            bytes.length = 0;
            return;
        }

        final int dictLength = firstVInt;
        final int blockLength = in.readVInt();

        final int numBlocks = readCompressedLengths(in, originalLength, dictLength, blockLength);

        buffer = ArrayUtil.grow(buffer, dictLength + blockLength);
        bytes.length = 0;
        // Read the dictionary
        if (LZ4.decompress(in, dictLength, buffer, 0) != dictLength) {
            throw new CorruptIndexException("Illegal dict length", in);
        }

        int offsetInBlock = dictLength;
        int offsetInBytesRef = offset;
        if (offset >= dictLength) {
            offsetInBytesRef -= dictLength;

            // Skip unneeded blocks
            int numBytesToSkip = 0;
            for (int i = 0; i < numBlocks && offsetInBlock + blockLength < offset; ++i) {
                int compressedBlockLength = compressedLengths[i];
                numBytesToSkip += compressedBlockLength;
                offsetInBlock += blockLength;
                offsetInBytesRef -= blockLength;
            }
            in.skipBytes(numBytesToSkip);
        } else {
            // The dictionary contains some bytes we need, copy its content to the BytesRef
            bytes.bytes = ArrayUtil.grow(bytes.bytes, dictLength);
            System.arraycopy(buffer, 0, bytes.bytes, 0, dictLength);
            bytes.length = dictLength;
        }

        // Read blocks that intersect with the interval we need
        while (offsetInBlock < offset + length) {
            final int bytesToDecompress = Math.min(blockLength, offset + length - offsetInBlock);
            LZ4.decompress(in, bytesToDecompress, buffer, dictLength);
            bytes.bytes = ArrayUtil.grow(bytes.bytes, bytes.length + bytesToDecompress);
            System.arraycopy(buffer, dictLength, bytes.bytes, bytes.length, bytesToDecompress);
            bytes.length += bytesToDecompress;
            offsetInBlock += blockLength;
        }

        bytes.offset = offsetInBytesRef;
        bytes.length = length;
        assert bytes.isValid();

    }

    @Override
    public AdapterDecompressor clone() {
        return new LZ4AdapterDecompressor();
    }

}

public final class ZstdAdapterDecompressor extends AdapterDecompressor {

    byte[] compressed;

    /**
     * default decompressor
     */
    public ZstdAdapterDecompressor() {
        compressed = BytesRef.EMPTY_BYTES;
    }

    /*resuable decompress function*/
    private void doDecompress(
        DataInput in, ZstdDecompressCtx dctx, BytesRef bytes, int decompressedLen)
        throws IOException {
        final int compressedLength = in.readVInt();
        if (compressedLength == 0) {
            return;
        }

        compressed = ArrayUtil.grow(compressed, compressedLength);
        in.readBytes(compressed, 0, compressedLength);

        bytes.bytes = ArrayUtil.grow(bytes.bytes, bytes.length + decompressedLen);
        int uncompressed =
            dctx.decompressByteArray(
                bytes.bytes, bytes.length, decompressedLen, compressed, 0, compressedLength);

        if (decompressedLen != uncompressed) {
            throw new IllegalStateException(decompressedLen + " " + uncompressed);
        }
        bytes.length += uncompressed;
    }

    @Override
    public void decompress(DataInput in, int originalLength, int offset, int length, BytesRef bytes, int firstVInt)
        throws IOException {
        assert offset + length <= originalLength;

        if (length == 0) {
            bytes.length = 0;
            return;
        }

        final int magicNumber = firstVInt;
        // 这里如果非-666666魔数说明逻辑有问题
        if (magicNumber != Comporessions.ZSTD_MAGIC_NUMBER) {
            throw new IllegalArgumentException("Unknown magicNumber:[" + magicNumber + "]");
        }

        final int dictLength = in.readVInt();
        final int blockLength = in.readVInt();
        bytes.bytes = ArrayUtil.grow(bytes.bytes, dictLength);
        bytes.offset = bytes.length = 0;

        try (ZstdDecompressCtx dctx = new ZstdDecompressCtx()) {

            // decompress dictionary first
            doDecompress(in, dctx, bytes, dictLength);

            dctx.loadDict(new ZstdDictDecompress(bytes.bytes, 0, dictLength));

            int offsetInBlock = dictLength;
            int offsetInBytesRef = offset;

            // Skip unneeded blocks
            while (offsetInBlock + blockLength < offset) {
                final int compressedLength = in.readVInt();
                in.skipBytes(compressedLength);
                offsetInBlock += blockLength;
                offsetInBytesRef -= blockLength;
            }

            // Read blocks that intersect with the interval we need
            while (offsetInBlock < offset + length) {
                bytes.bytes = ArrayUtil.grow(bytes.bytes, bytes.length + blockLength);
                int l = Math.min(blockLength, originalLength - offsetInBlock);
                doDecompress(in, dctx, bytes, l);
                offsetInBlock += blockLength;
            }

            bytes.offset = offsetInBytesRef;
            bytes.length = length;
            assert bytes.isValid();
        }
    }

    @Override
    public ZstdAdapterDecompressor clone() {
        return new ZstdAdapterDecompressor();
    }

}