Skip to content

Commit 2ab1b02

Browse files
authored
GH-192: Make LZ4 frame compression and decompression at least 10x faster with lz4-java (#1420)
## What's Changed `Lz4CompressionCodec` uses commons-compress's pure-Java LZ4, which is slow enough that compressing an 8 MiB buffer of sequential `int64` values takes several minutes (#192). This changes the codec to use [lz4-java](https://github.com/yawkat/lz4-java), the JNI bindings to liblz4 that Spark, Kafka and others use, while keeping every frame it could read before readable. **Compression.** Each buffer is written as one LZ4 frame of independent blocks of up to 4 MiB, without checksums, compressed with lz4-java directly from the buffer's memory. Every LZ4 frame decoder reads frames with independent blocks: liblz4 (Arrow C++, pyarrow), `lz4_flex` (arrow-rs), lz4-java, and commons-compress. **Decompression.** A frame with independent blocks is decompressed with lz4-java directly into the output buffer. That covers frames from this codec, from arrow-rs, and from the previous version of this codec. The new decoder accepts any block size, block and content checksums (which it verifies), and a content size. It checks the frame's block lengths, and that they can hold the recorded decompressed length, before allocating for it. A frame whose blocks are linked, or that uses a dictionary, still goes through commons-compress. lz4-java cannot decode a block that refers to earlier blocks, and Arrow C++ writes linked blocks: liblz4's default `LZ4F_preferences_t` is `LZ4F_blockLinked`, so pyarrow 26 writes every buffer larger than 64 KiB as linked 64 KiB blocks. The LZ4 files in arrow-testing are small enough that liblz4 writes them as single, independent blocks, so the integration tests would not catch this. `TestLz4CompressionCodec` covers linked frames written by commons-compress instead. **Dependency.** This adds `at.yawk.lz4:lz4-java` 1.12.1, the maintained fork of `org.lz4:lz4-java`. The original artifact's last code release, 1.8.0, has several published advisories (for example CVE-2025-12183 and CVE-2025-66566), and its 1.8.1 is a relocation POM that points to the fork. Spark moved to the fork in 4.2. Both artifacts use the `net.jpountz` packages and the `org.lz4.java` module name, and the codec uses only APIs present in both, so an application that pins `org.lz4:lz4-java` instead still works. The codec asks for lz4-java's JNI implementation and falls back to its Java one only when the native library cannot load, because `LZ4Factory.fastestInstance()` settles for Java whenever lz4-java is not loaded by the system class loader. ### Performance `CompressionCodecBenchmarks` (new, in `performance`) on an Apple M3 Ultra, JDK 17, 8 MiB buffers: | Benchmark | LZ4_FRAME | ZSTD | | --- | ---: | ---: | | compress, sequential `int64` | 8.5 ms | 12.5 ms | | compress, random decimal text | 9.9 ms | 44.6 ms | | decompress, sequential `int64` | 2.1 ms | 4.8 ms | | decompress, random decimal text | 1.9 ms | 8.5 ms | The previous codec is too slow for the same JMH run. From a simple timing harness on the same machine: | Input | Previous compress | New compress | Previous decompress | New decompress | | --- | ---: | ---: | ---: | ---: | | 8 MiB sequential `int64` | 278 s | 7.8 ms | 93.5 ms | 2.2 ms | | 8 MiB of `i % 1000` as text | 17.9 ms | 0.3 ms | 3.3 ms | 0.3 ms | | 8 MiB random bytes | 243 ms | 0.5 ms | (stored) | (stored) | | 1 MiB random decimal text | 38.5 s | 1.4 ms | 33.3 ms | 0.3 ms | ### Compression ratio On the inputs above, the ratio is unchanged except for random decimal text, where commons-compress's more thorough match search reached 1.61 and lz4-java's fast compressor reaches 1.10. The new output matches what Arrow C++ writes by default: pyarrow's LZ4 codec also reaches 1.10 on the same kind of data, and 1.63 at compression level 9. `CommonsCompressionFactory.createCodec(LZ4_FRAME, level)` ignores the level today. Mapping it to lz4-java's high-compression mode would give users who want the better ratio a way to ask for it; I left that out to keep this change focused, and can add it here or in a follow-up. ### Testing - The existing `TestCompressionCodec` and `TestArrowReaderWriterWithCompression` pass unchanged, including the decompressed-length check from #1207. - `TestLz4CompressionCodec` (new): - round-trips buffers across block boundaries, including one whose first block is incompressible and stored, and checks that commons-compress and lz4-java's `LZ4FrameInputStream` both read the frames this codec writes; - decompresses frames written by commons-compress for every block size, with linked and independent blocks, with and without checksums, and a frame written by lz4-java with a content size; - rejects corrupt frames (descriptor, block and content checksums, a missing end mark, blocks that run past the frame or exceed the maximum block size, undecodable blocks, a wrong content size, and recorded lengths that are too long or too short), and releases everything it allocated. - I also decoded pyarrow 26 output (buffers over 64 KiB, linked blocks) by hand to confirm that lz4-java rejects linked frames and commons-compress reads them, which is why the fallback stays. ### AI usage The code, tests and benchmark in this PR were written with an AI coding assistant. The frame behavior described above (pyarrow's linked blocks, lz4-java's rejection of them) and the benchmark numbers come from running them. Closes #192.
1 parent 3c9e49b commit 2ab1b02

7 files changed

Lines changed: 775 additions & 32 deletions

File tree

‎compression/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,5 +57,10 @@ under the License.
5757
<artifactId>zstd-jni</artifactId>
5858
<version>1.5.7-25</version>
5959
</dependency>
60+
<dependency>
61+
<groupId>at.yawk.lz4</groupId>
62+
<artifactId>lz4-java</artifactId>
63+
<version>1.12.1</version>
64+
</dependency>
6065
</dependencies>
6166
</project>

‎compression/src/main/java/module-info.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
requires org.apache.arrow.memory.core;
2525
requires org.apache.arrow.vector;
2626
requires org.apache.commons.compress;
27+
requires org.lz4.java;
2728

2829
// Also defined under META-INF/services to support non-modular applications
2930
provides CompressionCodec.Factory with

‎compression/src/main/java/org/apache/arrow/compression/Lz4CompressionCodec.java‎

Lines changed: 291 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -20,58 +20,268 @@
2020
import java.io.ByteArrayOutputStream;
2121
import java.io.IOException;
2222
import java.io.InputStream;
23-
import java.io.OutputStream;
23+
import net.jpountz.lz4.LZ4Compressor;
24+
import net.jpountz.lz4.LZ4Exception;
25+
import net.jpountz.lz4.LZ4Factory;
26+
import net.jpountz.lz4.LZ4SafeDecompressor;
27+
import net.jpountz.xxhash.XXHash32;
28+
import net.jpountz.xxhash.XXHashFactory;
2429
import org.apache.arrow.memory.ArrowBuf;
2530
import org.apache.arrow.memory.BufferAllocator;
31+
import org.apache.arrow.memory.util.MemoryUtil;
2632
import org.apache.arrow.util.Preconditions;
2733
import org.apache.arrow.vector.compression.AbstractCompressionCodec;
2834
import org.apache.arrow.vector.compression.CompressionUtil;
2935
import org.apache.commons.compress.compressors.lz4.FramedLZ4CompressorInputStream;
30-
import org.apache.commons.compress.compressors.lz4.FramedLZ4CompressorOutputStream;
3136
import org.apache.commons.compress.utils.IOUtils;
3237

33-
/** Compression codec for the LZ4 algorithm. */
38+
/**
39+
* Compression codec for the LZ4 frame format.
40+
*
41+
* <p>Each buffer is written as one frame of independent blocks of up to 4 MiB, without checksums,
42+
* which every LZ4 frame decoder reads. Compression, and the decompression of frames whose blocks
43+
* are independent, use lz4-java on the buffers' memory directly. A frame whose blocks are linked,
44+
* as Arrow C++ writes for a buffer larger than 64 KiB, or that uses a dictionary, is decompressed
45+
* with commons-compress, since lz4-java cannot decode a block that refers to earlier ones.
46+
*/
3447
public class Lz4CompressionCodec extends AbstractCompressionCodec {
3548

49+
private static final long PREFIX = CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH;
50+
51+
private static final int MAGIC = 0x184D2204;
52+
53+
// The frame descriptor this codec writes: FLG (version 01, independent blocks), BD (4 MiB
54+
// maximum block size) and HC, the second byte of the xxHash32 of FLG and BD with a zero seed.
55+
private static final byte WRITE_FLG = 0x60;
56+
private static final byte WRITE_BD = 0x70;
57+
private static final byte WRITE_HC = 0x73;
58+
private static final int WRITE_BLOCK_SIZE = 4 * 1024 * 1024;
59+
60+
private static final int FLG_VERSION_MASK = 0xC0;
61+
private static final int FLG_VERSION = 0x40;
62+
private static final int FLG_INDEPENDENT_BLOCKS = 0x20;
63+
private static final int FLG_BLOCK_CHECKSUM = 0x10;
64+
private static final int FLG_CONTENT_SIZE = 0x08;
65+
private static final int FLG_CONTENT_CHECKSUM = 0x04;
66+
private static final int FLG_RESERVED = 0x02;
67+
private static final int FLG_DICTIONARY_ID = 0x01;
68+
private static final int BD_RESERVED = 0x8F;
69+
70+
// A block length's high bit marks a block stored uncompressed. A zero length ends the frame.
71+
private static final int UNCOMPRESSED_BLOCK = 0x80000000;
72+
73+
// An LZ4 block cannot decompress to more than 255 times its compressed size.
74+
private static final long MAX_COMPRESSION_RATIO = 255;
75+
3676
@Override
3777
protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf uncompressedBuffer) {
38-
Preconditions.checkArgument(
39-
uncompressedBuffer.writerIndex() <= Integer.MAX_VALUE,
40-
"The uncompressed buffer size exceeds the integer limit %s.",
41-
Integer.MAX_VALUE);
78+
LZ4Compressor compressor = Lz4.COMPRESSOR;
79+
long length = uncompressedBuffer.writerIndex();
80+
long fullBlocks = length / WRITE_BLOCK_SIZE;
81+
int lastBlock = (int) (length % WRITE_BLOCK_SIZE);
82+
long bound =
83+
PREFIX
84+
+ 7
85+
+ fullBlocks * (Integer.BYTES + compressor.maxCompressedLength(WRITE_BLOCK_SIZE))
86+
+ (lastBlock > 0 ? Integer.BYTES + compressor.maxCompressedLength(lastBlock) : 0)
87+
+ Integer.BYTES;
88+
ArrowBuf frame = allocator.buffer(bound);
89+
try {
90+
setIntLE(frame, PREFIX, MAGIC);
91+
frame.setByte(PREFIX + 4, WRITE_FLG);
92+
frame.setByte(PREFIX + 5, WRITE_BD);
93+
frame.setByte(PREFIX + 6, WRITE_HC);
94+
long at = PREFIX + 7;
95+
for (long from = 0; from < length; ) {
96+
int n = (int) Math.min(WRITE_BLOCK_SIZE, length - from);
97+
int room = compressor.maxCompressedLength(n);
98+
int packed =
99+
compressor.compress(
100+
uncompressedBuffer.nioBuffer(from, n),
101+
0,
102+
n,
103+
frame.nioBuffer(at + Integer.BYTES, room),
104+
0,
105+
room);
106+
if (packed < n) {
107+
setIntLE(frame, at, packed);
108+
at += Integer.BYTES + packed;
109+
} else {
110+
setIntLE(frame, at, n | UNCOMPRESSED_BLOCK);
111+
frame.setBytes(at + Integer.BYTES, uncompressedBuffer, from, n);
112+
at += Integer.BYTES + n;
113+
}
114+
from += n;
115+
}
116+
setIntLE(frame, at, 0);
117+
frame.writerIndex(at + Integer.BYTES);
118+
return frame;
119+
} catch (RuntimeException | Error e) {
120+
frame.close();
121+
throw e;
122+
}
123+
}
42124

43-
byte[] inBytes = new byte[(int) uncompressedBuffer.writerIndex()];
44-
uncompressedBuffer.getBytes(/* index= */ 0, inBytes);
45-
ByteArrayOutputStream baos = new ByteArrayOutputStream();
46-
try (InputStream in = new ByteArrayInputStream(inBytes);
47-
OutputStream out = new FramedLZ4CompressorOutputStream(baos)) {
48-
IOUtils.copy(in, out);
49-
} catch (IOException e) {
50-
throw new RuntimeException(e);
125+
@Override
126+
protected ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf compressedBuffer) {
127+
long length = readUncompressedLength(compressedBuffer);
128+
long end = compressedBuffer.writerIndex();
129+
if (end - PREFIX < 7 || getIntLE(compressedBuffer, PREFIX) != MAGIC) {
130+
return decompressWithCommonsCompress(allocator, compressedBuffer);
131+
}
132+
int flg = compressedBuffer.getByte(PREFIX + 4) & 0xFF;
133+
int bd = compressedBuffer.getByte(PREFIX + 5) & 0xFF;
134+
int blockSizeId = (bd >>> 4) & 0x07;
135+
boolean contentChecksum = (flg & FLG_CONTENT_CHECKSUM) != 0;
136+
if ((flg & FLG_VERSION_MASK) != FLG_VERSION
137+
|| (flg & (FLG_RESERVED | FLG_DICTIONARY_ID)) != 0
138+
|| (flg & FLG_INDEPENDENT_BLOCKS) == 0
139+
|| (bd & BD_RESERVED) != 0
140+
|| blockSizeId < 4
141+
|| (contentChecksum && length > Integer.MAX_VALUE)) {
142+
return decompressWithCommonsCompress(allocator, compressedBuffer);
51143
}
144+
boolean blockChecksum = (flg & FLG_BLOCK_CHECKSUM) != 0;
145+
int maxBlockSize = 1 << (8 + 2 * blockSizeId);
52146

53-
byte[] outBytes = baos.toByteArray();
147+
long at = PREFIX + 6;
148+
if ((flg & FLG_CONTENT_SIZE) != 0) {
149+
if (end - at < Long.BYTES + 1) {
150+
throw corrupt("the frame descriptor is cut off");
151+
}
152+
long contentSize = getLongLE(compressedBuffer, at);
153+
if (contentSize != length) {
154+
throw corrupt(
155+
"the frame records a content size of "
156+
+ contentSize
157+
+ " bytes, but the buffer's decompressed length is "
158+
+ length);
159+
}
160+
at += Long.BYTES;
161+
}
162+
int descriptorLength = (int) (at - (PREFIX + 4));
163+
int headerChecksum =
164+
(Lz4.XXHASH.hash(
165+
compressedBuffer.nioBuffer(PREFIX + 4, descriptorLength),
166+
0,
167+
descriptorLength,
168+
0)
169+
>>> 8)
170+
& 0xFF;
171+
if ((compressedBuffer.getByte(at) & 0xFF) != headerChecksum) {
172+
throw corrupt("the frame descriptor's checksum does not match");
173+
}
174+
at += 1;
54175

55-
ArrowBuf compressedBuffer =
56-
allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + outBytes.length);
57-
compressedBuffer.setBytes(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH, outBytes);
58-
compressedBuffer.writerIndex(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + outBytes.length);
59-
return compressedBuffer;
176+
// Check the blocks' lengths, and that they can hold the decompressed length, before
177+
// allocating for it.
178+
long blocksStart = at;
179+
long maxDecompressedLength = 0;
180+
while (true) {
181+
if (end - at < Integer.BYTES) {
182+
throw corrupt("the frame ends before its end mark");
183+
}
184+
int word = getIntLE(compressedBuffer, at);
185+
at += Integer.BYTES;
186+
if (word == 0) {
187+
break;
188+
}
189+
int n = word & ~UNCOMPRESSED_BLOCK;
190+
if (n > maxBlockSize) {
191+
throw corrupt("a block of " + n + " bytes is larger than the frame's maximum block size");
192+
}
193+
if (n > end - at) {
194+
throw corrupt("a block of " + n + " bytes runs past the end of the frame");
195+
}
196+
maxDecompressedLength +=
197+
(word & UNCOMPRESSED_BLOCK) != 0 ? n : Math.min(maxBlockSize, n * MAX_COMPRESSION_RATIO);
198+
at += n;
199+
if (blockChecksum) {
200+
if (end - at < Integer.BYTES) {
201+
throw corrupt("the frame ends before a block's checksum");
202+
}
203+
at += Integer.BYTES;
204+
}
205+
}
206+
if (contentChecksum && end - at < Integer.BYTES) {
207+
throw corrupt("the frame ends before its content checksum");
208+
}
209+
if (length < 0 || length > maxDecompressedLength) {
210+
throw corrupt(
211+
"the frame decompresses to at most "
212+
+ maxDecompressedLength
213+
+ " bytes, not its recorded decompressed length of "
214+
+ length);
215+
}
216+
217+
ArrowBuf decompressedBuffer = allocator.buffer(length);
218+
try {
219+
LZ4SafeDecompressor decompressor = Lz4.DECOMPRESSOR;
220+
at = blocksStart;
221+
long to = 0;
222+
while (true) {
223+
int word = getIntLE(compressedBuffer, at);
224+
at += Integer.BYTES;
225+
if (word == 0) {
226+
break;
227+
}
228+
int n = word & ~UNCOMPRESSED_BLOCK;
229+
if (blockChecksum
230+
&& Lz4.XXHASH.hash(compressedBuffer.nioBuffer(at, n), 0, n, 0)
231+
!= getIntLE(compressedBuffer, at + n)) {
232+
throw corrupt("a block's checksum does not match");
233+
}
234+
if ((word & UNCOMPRESSED_BLOCK) != 0) {
235+
if (n > length - to) {
236+
throw corrupt(
237+
"the frame holds more than its recorded decompressed length of " + length);
238+
}
239+
decompressedBuffer.setBytes(to, compressedBuffer, at, n);
240+
to += n;
241+
} else {
242+
int room = (int) Math.min(maxBlockSize, length - to);
243+
try {
244+
to +=
245+
decompressor.decompress(
246+
compressedBuffer.nioBuffer(at, n),
247+
0,
248+
n,
249+
decompressedBuffer.nioBuffer(to, room),
250+
0,
251+
room);
252+
} catch (LZ4Exception e) {
253+
throw corrupt("a block does not decompress within " + room + " bytes", e);
254+
}
255+
}
256+
at += n + (blockChecksum ? Integer.BYTES : 0);
257+
}
258+
if (to != length) {
259+
throw wrongLength(length, to);
260+
}
261+
if (contentChecksum
262+
&& Lz4.XXHASH.hash(decompressedBuffer.nioBuffer(0, (int) length), 0, (int) length, 0)
263+
!= getIntLE(compressedBuffer, at)) {
264+
throw corrupt("the frame's content checksum does not match");
265+
}
266+
decompressedBuffer.writerIndex(length);
267+
return decompressedBuffer;
268+
} catch (RuntimeException | Error e) {
269+
decompressedBuffer.close();
270+
throw e;
271+
}
60272
}
61273

62-
@Override
63-
protected ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf compressedBuffer) {
274+
private ArrowBuf decompressWithCommonsCompress(
275+
BufferAllocator allocator, ArrowBuf compressedBuffer) {
64276
Preconditions.checkArgument(
65277
compressedBuffer.writerIndex() <= Integer.MAX_VALUE,
66278
"The compressed buffer size exceeds the integer limit %s",
67279
Integer.MAX_VALUE);
68280

69281
long decompressedLength = readUncompressedLength(compressedBuffer);
70282

71-
byte[] inBytes =
72-
new byte
73-
[(int) (compressedBuffer.writerIndex() - CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH)];
74-
compressedBuffer.getBytes(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH, inBytes);
283+
byte[] inBytes = new byte[(int) (compressedBuffer.writerIndex() - PREFIX)];
284+
compressedBuffer.getBytes(PREFIX, inBytes);
75285
ByteArrayOutputStream out = new ByteArrayOutputStream((int) decompressedLength);
76286
try (InputStream in = new FramedLZ4CompressorInputStream(new ByteArrayInputStream(inBytes))) {
77287
IOUtils.copy(in, out);
@@ -81,11 +291,7 @@ protected ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf compressedBu
81291

82292
byte[] outBytes = out.toByteArray();
83293
if (outBytes.length != decompressedLength) {
84-
throw new RuntimeException(
85-
"Expected != actual decompressed length: "
86-
+ decompressedLength
87-
+ " != "
88-
+ outBytes.length);
294+
throw wrongLength(decompressedLength, outBytes.length);
89295
}
90296
ArrowBuf decompressedBuffer = allocator.buffer(outBytes.length);
91297
decompressedBuffer.setBytes(/* index= */ 0, outBytes);
@@ -97,4 +303,57 @@ protected ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf compressedBu
97303
public CompressionUtil.CodecType getCodecType() {
98304
return CompressionUtil.CodecType.LZ4_FRAME;
99305
}
306+
307+
private static RuntimeException corrupt(String reason) {
308+
return new IllegalArgumentException("Corrupt LZ4 frame: " + reason);
309+
}
310+
311+
private static RuntimeException corrupt(String reason, Throwable cause) {
312+
return new IllegalArgumentException("Corrupt LZ4 frame: " + reason, cause);
313+
}
314+
315+
private static RuntimeException wrongLength(long expected, long actual) {
316+
return new RuntimeException(
317+
"Expected != actual decompressed length: " + expected + " != " + actual);
318+
}
319+
320+
private static int getIntLE(ArrowBuf buf, long index) {
321+
int value = buf.getInt(index);
322+
return MemoryUtil.LITTLE_ENDIAN ? value : Integer.reverseBytes(value);
323+
}
324+
325+
private static long getLongLE(ArrowBuf buf, long index) {
326+
long value = buf.getLong(index);
327+
return MemoryUtil.LITTLE_ENDIAN ? value : Long.reverseBytes(value);
328+
}
329+
330+
private static void setIntLE(ArrowBuf buf, long index, int value) {
331+
buf.setInt(index, MemoryUtil.LITTLE_ENDIAN ? value : Integer.reverseBytes(value));
332+
}
333+
334+
/** lz4-java's implementations, loaded on first use. */
335+
private static final class Lz4 {
336+
// fastestInstance() picks the JNI implementations only when lz4-java was loaded by the system
337+
// class loader, and otherwise settles for the Java ones, so ask for the JNI ones directly.
338+
static final LZ4Factory FACTORY = lz4Factory();
339+
static final LZ4Compressor COMPRESSOR = FACTORY.fastCompressor();
340+
static final LZ4SafeDecompressor DECOMPRESSOR = FACTORY.safeDecompressor();
341+
static final XXHash32 XXHASH = xxHashFactory().hash32();
342+
343+
private static LZ4Factory lz4Factory() {
344+
try {
345+
return LZ4Factory.nativeInstance();
346+
} catch (LinkageError | RuntimeException e) {
347+
return LZ4Factory.fastestJavaInstance();
348+
}
349+
}
350+
351+
private static XXHashFactory xxHashFactory() {
352+
try {
353+
return XXHashFactory.nativeInstance();
354+
} catch (LinkageError | RuntimeException e) {
355+
return XXHashFactory.fastestJavaInstance();
356+
}
357+
}
358+
}
100359
}

0 commit comments

Comments
 (0)