-
Notifications
You must be signed in to change notification settings - Fork 28.4k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Davies Liu
committed
Dec 17, 2015
1 parent
f49f1ef
commit 13390e8
Showing
4 changed files
with
266 additions
and
8 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
259 changes: 259 additions & 0 deletions
259
core/src/main/scala/org/apache/spark/io/LZ4BlockInputStream.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,259 @@ | ||
| package org.apache.spark.io; | ||
|
|
||
| /* | ||
| * Licensed 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. | ||
| */ | ||
|
|
||
| import java.io.EOFException; | ||
| import java.io.FilterInputStream; | ||
| import java.io.IOException; | ||
| import java.io.InputStream; | ||
| import java.util.zip.Checksum; | ||
|
|
||
| import net.jpountz.lz4.LZ4BlockOutputStream; | ||
| import net.jpountz.lz4.LZ4Exception; | ||
| import net.jpountz.lz4.LZ4Factory; | ||
| import net.jpountz.lz4.LZ4FastDecompressor; | ||
| import net.jpountz.util.SafeUtils; | ||
| import net.jpountz.xxhash.StreamingXXHash32; | ||
| import net.jpountz.xxhash.XXHash32; | ||
| import net.jpountz.xxhash.XXHashFactory; | ||
|
|
||
| /** | ||
| * {@link InputStream} implementation to decode data written with | ||
| * {@link LZ4BlockOutputStream}. This class is not thread-safe and does not | ||
| * support {@link #mark(int)}/{@link #reset()}. | ||
| * @see LZ4BlockOutputStream | ||
| * | ||
| * This is based on net.jpountz.lz4.LZ4BlockInputStream | ||
| */ | ||
| public final class LZ4BlockInputStream extends FilterInputStream { | ||
|
|
||
| // Copied from net.jpountz.lz4.LZ4BlockOutputStream | ||
| static final byte[] MAGIC = new byte[] { 'L', 'Z', '4', 'B', 'l', 'o', 'c', 'k' }; | ||
| static final int MAGIC_LENGTH = MAGIC.length; | ||
|
|
||
| static final int HEADER_LENGTH = | ||
| MAGIC_LENGTH // magic bytes | ||
| + 1 // token | ||
| + 4 // compressed length | ||
| + 4 // decompressed length | ||
| + 4; // checksum | ||
|
|
||
| static final int COMPRESSION_LEVEL_BASE = 10; | ||
|
|
||
| static final int COMPRESSION_METHOD_RAW = 0x10; | ||
| static final int COMPRESSION_METHOD_LZ4 = 0x20; | ||
|
|
||
| static final int DEFAULT_SEED = 0x9747b28c; | ||
|
|
||
| private final LZ4FastDecompressor decompressor; | ||
| private final Checksum checksum; | ||
| private byte[] buffer; | ||
| private byte[] compressedBuffer; | ||
| private int originalLen; | ||
| private int o; | ||
| private boolean finished; | ||
|
|
||
| /** | ||
| * Create a new {@link InputStream}. | ||
| * | ||
| * @param in the {@link InputStream} to poll | ||
| * @param decompressor the {@link LZ4FastDecompressor decompressor} instance to | ||
| * use | ||
| * @param checksum the {@link Checksum} instance to use, must be | ||
| * equivalent to the instance which has been used to | ||
| * write the stream | ||
| */ | ||
| public LZ4BlockInputStream(InputStream in, LZ4FastDecompressor decompressor, Checksum checksum) { | ||
| super(in); | ||
| this.decompressor = decompressor; | ||
| this.checksum = checksum; | ||
| this.buffer = new byte[0]; | ||
| this.compressedBuffer = new byte[HEADER_LENGTH]; | ||
| o = originalLen = 0; | ||
| finished = false; | ||
| } | ||
|
|
||
| /** | ||
| * Create a new instance using {@link XXHash32} for checksuming. | ||
| * @see #LZ4BlockInputStream(InputStream, LZ4FastDecompressor, Checksum) | ||
| * @see StreamingXXHash32#asChecksum() | ||
| */ | ||
| public LZ4BlockInputStream(InputStream in, LZ4FastDecompressor decompressor) { | ||
| this(in, decompressor, XXHashFactory.fastestInstance().newStreamingHash32(DEFAULT_SEED).asChecksum()); | ||
| } | ||
|
|
||
| /** | ||
| * Create a new instance which uses the fastest {@link LZ4FastDecompressor} available. | ||
| * @see LZ4Factory#fastestInstance() | ||
| * @see #LZ4BlockInputStream(InputStream, LZ4FastDecompressor) | ||
| */ | ||
| public LZ4BlockInputStream(InputStream in) { | ||
| this(in, LZ4Factory.fastestInstance().fastDecompressor()); | ||
| } | ||
|
|
||
| @Override | ||
| public int available() throws IOException { | ||
| refill(); | ||
| return originalLen - o; | ||
| } | ||
|
|
||
| @Override | ||
| public int read() throws IOException { | ||
| refill(); | ||
| if (finished) { | ||
| return -1; | ||
| } | ||
| return buffer[o++] & 0xFF; | ||
| } | ||
|
|
||
| @Override | ||
| public int read(byte[] b, int off, int len) throws IOException { | ||
| SafeUtils.checkRange(b, off, len); | ||
| refill(); | ||
| if (finished) { | ||
| return -1; | ||
| } | ||
| len = Math.min(len, originalLen - o); | ||
| System.arraycopy(buffer, o, b, off, len); | ||
| o += len; | ||
| return len; | ||
| } | ||
|
|
||
| @Override | ||
| public int read(byte[] b) throws IOException { | ||
| return read(b, 0, b.length); | ||
| } | ||
|
|
||
| @Override | ||
| public long skip(long n) throws IOException { | ||
| refill(); | ||
| if (finished) { | ||
| return -1; | ||
| } | ||
| final int skipped = (int) Math.min(n, originalLen - o); | ||
| o += skipped; | ||
| return skipped; | ||
| } | ||
|
|
||
| private void refill() throws IOException { | ||
| if (finished || o < originalLen) { | ||
| return; | ||
| } | ||
| try { | ||
| readFully(compressedBuffer, HEADER_LENGTH); | ||
| } catch (EOFException e) { | ||
| finished = true; | ||
| return; | ||
| } | ||
| for (int i = 0; i < MAGIC_LENGTH; ++i) { | ||
| if (compressedBuffer[i] != MAGIC[i]) { | ||
| throw new IOException("Stream is corrupted"); | ||
| } | ||
| } | ||
| final int token = compressedBuffer[MAGIC_LENGTH] & 0xFF; | ||
| final int compressionMethod = token & 0xF0; | ||
| final int compressionLevel = COMPRESSION_LEVEL_BASE + (token & 0x0F); | ||
| if (compressionMethod != COMPRESSION_METHOD_RAW && compressionMethod != COMPRESSION_METHOD_LZ4) | ||
| { | ||
| throw new IOException("Stream is corrupted"); | ||
| } | ||
| final int compressedLen = SafeUtils.readIntLE(compressedBuffer, MAGIC_LENGTH + 1); | ||
| originalLen = SafeUtils.readIntLE(compressedBuffer, MAGIC_LENGTH + 5); | ||
| final int check = SafeUtils.readIntLE(compressedBuffer, MAGIC_LENGTH + 9); | ||
| assert HEADER_LENGTH == MAGIC_LENGTH + 13; | ||
| if (originalLen > 1 << compressionLevel | ||
| || originalLen < 0 | ||
| || compressedLen < 0 | ||
| || (originalLen == 0 && compressedLen != 0) | ||
| || (originalLen != 0 && compressedLen == 0) | ||
| || (compressionMethod == COMPRESSION_METHOD_RAW && originalLen != compressedLen)) { | ||
| throw new IOException("Stream is corrupted"); | ||
| } | ||
| if (originalLen == 0 && compressedLen == 0) { | ||
| if (check != 0) { | ||
| throw new IOException("Stream is corrupted"); | ||
| } | ||
| refill(); | ||
| return; | ||
| } | ||
| if (buffer.length < originalLen) { | ||
| buffer = new byte[Math.max(originalLen, buffer.length * 3 / 2)]; | ||
| } | ||
| switch (compressionMethod) { | ||
| case COMPRESSION_METHOD_RAW: | ||
| readFully(buffer, originalLen); | ||
| break; | ||
| case COMPRESSION_METHOD_LZ4: | ||
| if (compressedBuffer.length < originalLen) { | ||
| compressedBuffer = new byte[Math.max(compressedLen, compressedBuffer.length * 3 / 2)]; | ||
| } | ||
| readFully(compressedBuffer, compressedLen); | ||
| try { | ||
| final int compressedLen2 = | ||
| decompressor.decompress(compressedBuffer, 0, buffer, 0, originalLen); | ||
| if (compressedLen != compressedLen2) { | ||
| throw new IOException("Stream is corrupted"); | ||
| } | ||
| } catch (LZ4Exception e) { | ||
| throw new IOException("Stream is corrupted", e); | ||
| } | ||
| break; | ||
| default: | ||
| throw new AssertionError(); | ||
| } | ||
| checksum.reset(); | ||
| checksum.update(buffer, 0, originalLen); | ||
| if ((int) checksum.getValue() != check) { | ||
| throw new IOException("Stream is corrupted"); | ||
| } | ||
| o = 0; | ||
| } | ||
|
|
||
| private void readFully(byte[] b, int len) throws IOException { | ||
| int read = 0; | ||
| while (read < len) { | ||
| final int r = in.read(b, read, len - read); | ||
| if (r < 0) { | ||
| throw new EOFException("Stream ended prematurely"); | ||
| } | ||
| read += r; | ||
| } | ||
| assert len == read; | ||
| } | ||
|
|
||
| @Override | ||
| public boolean markSupported() { | ||
| return false; | ||
| } | ||
|
|
||
| @SuppressWarnings("sync-override") | ||
| @Override | ||
| public void mark(int readlimit) { | ||
| // unsupported | ||
| } | ||
|
|
||
| @SuppressWarnings("sync-override") | ||
| @Override | ||
| public void reset() throws IOException { | ||
| throw new IOException("mark/reset not supported"); | ||
| } | ||
|
|
||
| @Override | ||
| public String toString() { | ||
| return getClass().getSimpleName() + "(in=" + in | ||
| + ", decompressor=" + decompressor + ", checksum=" + checksum + ")"; | ||
| } | ||
|
|
||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters