Skip to content

Commit 601b6ce

Browse files
dfa1claude
dfa1
and
claude
committed
feat: dictionary-aware streaming
ZstdOutputStream / ZstdInputStream gain dictionary constructors, so large or incremental streams of small, similar records (logs, JSON over a socket) compress against a trained dictionary. - bind ZSTD_CCtx_loadDictionary / ZSTD_DCtx_loadDictionary; load into the context before streaming starts (dict bytes copied, source may be released) - new ctors: ZstdOutputStream(out[, level], dict), ZstdInputStream(in, dict) - compatibility.md now 39/186 Tests: dictionary stream round-trip; dictionary shrinks a streamed record below the plain stream frame; a dictionary stream frame also decodes with the one-shot ZstdDecompressCtx + same dictionary. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent c2a1138 commit 601b6ce

5 files changed

Lines changed: 132 additions & 6 deletions

File tree

‎docs/compatibility.md‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ and semantics follow the [official manual](https://facebook.github.io/zstd/doc/a
55

66
- zstd version: **1.6.0** (vendored `third_party/zstd`)
77
- Public symbols exported by `libzstd`: **186**
8-
- Bound so far: **37** (~20%)
8+
- Bound so far: **39** (~21%)
99

1010
"Bound" means the symbol has a `MethodHandle` in `Bindings` and is reachable
1111
through the public Java API. The rest are reachable from native code but not yet
@@ -27,7 +27,7 @@ low-level variants that an idiomatic Java API does not need.
2727
| Dictionary training (ZDICT) | 4 / 12 | `trainFromBuffer`; cover/fastCover optimizers not bound |
2828
| Streaming — compress | 3 / 22 | `ZstdOutputStream` (compressStream2 + buffer sizes) |
2929
| Streaming — decompress | 3 / 15 | `ZstdInputStream` (decompressStream + buffer sizes) |
30-
| Advanced parameters | 2 / 38 | `CCtx_setParameter` + `compress2` via `ZstdCompressCtx` (level, checksum, LDM, windowLog…); MT/getBounds not bound |
30+
| Advanced parameters | 4 / 38 | `CCtx_setParameter` + `compress2` (level, checksum, LDM, windowLog) and `C/DCtx_loadDictionary` (dictionary streams); MT/getBounds not bound |
3131
| Frame inspection | 1 / 13 | only `getFrameContentSize`; header/skippable/dictID-from-frame not bound |
3232
| Memory sizing | 0 / 14 | `sizeof_*` / `estimate*` accounting not bound |
3333
| Low-level block | 0 / 12 | expert block/continue API not bound |
@@ -55,10 +55,11 @@ low-level variants that an idiomatic Java API does not need.
5555
| `ZSTD_compressStream2`, `ZSTD_CStreamInSize`, `ZSTD_CStreamOutSize`, `ZSTD_CCtx_setParameter` | `ZstdOutputStream` |
5656
| `ZSTD_decompressStream`, `ZSTD_DStreamInSize`, `ZSTD_DStreamOutSize` | `ZstdInputStream` |
5757
| `ZSTD_compress2`, `ZSTD_CCtx_setParameter` | `ZstdCompressCtx.parameter` / `checksum` / `longDistanceMatching` / `windowLog` (+ `ZstdCompressParameter`) |
58+
| `ZSTD_CCtx_loadDictionary`, `ZSTD_DCtx_loadDictionary` | `ZstdOutputStream` / `ZstdInputStream` dictionary constructors |
5859

5960
## Roadmap (priority order)
6061

61-
1. ~~**Streaming**~~ — done: `ZstdOutputStream` / `ZstdInputStream` (`compressStream2`, `decompressStream`, bounded buffers). Remaining: dictionary-on-stream, `MemorySegment`-buffer driver, `pledgedSrcSize`.
62+
1. ~~**Streaming**~~ — done: `ZstdOutputStream` / `ZstdInputStream` (`compressStream2`, `decompressStream`, bounded buffers, dictionary constructors). Remaining: `MemorySegment`-buffer driver, `pledgedSrcSize`.
6263
2. **Advanced parameters** — done for compression: `CCtx_setParameter` + `compress2` via `ZstdCompressCtx` (`checksum`, `longDistanceMatching`, `windowLog`, generic `parameter`). Remaining: `cParam_getBounds`, `pledgedSrcSize`, and `nbWorkers` (needs a multithreaded native build).
6364
3. **Frame inspection** — `getFrameHeader`, `isFrame`, `findFrameCompressedSize`, `decompressBound`, `getDictID_from*`, skippable frames.
6465
4. **Better dictionaries** — `ZDICT_optimizeTrainFromBuffer_cover` / `_fastCover`, `finalizeDictionary`.
@@ -201,7 +202,7 @@ low-level variants that an idiomatic Java API does not need.
201202
| `ZSTD_resetDStream` | — |
202203
| `ZSTD_sizeof_DStream` | — |
203204

204-
### Advanced parameters (2/38)
205+
### Advanced parameters (4/38)
205206

206207
| Symbol | Bound |
207208
|---|:---:|
@@ -212,7 +213,7 @@ low-level variants that an idiomatic Java API does not need.
212213
| `ZSTD_CCtxParams_reset` | — |
213214
| `ZSTD_CCtxParams_setParameter` | — |
214215
| `ZSTD_CCtx_getParameter` | — |
215-
| `ZSTD_CCtx_loadDictionary` | — |
216+
| `ZSTD_CCtx_loadDictionary` | ✅ |
216217
| `ZSTD_CCtx_loadDictionary_advanced` | — |
217218
| `ZSTD_CCtx_loadDictionary_byReference` | — |
218219
| `ZSTD_CCtx_refCDict` | — |
@@ -227,7 +228,7 @@ low-level variants that an idiomatic Java API does not need.
227228
| `ZSTD_CCtx_setParams` | — |
228229
| `ZSTD_CCtx_setPledgedSrcSize` | — |
229230
| `ZSTD_DCtx_getParameter` | — |
230-
| `ZSTD_DCtx_loadDictionary` | — |
231+
| `ZSTD_DCtx_loadDictionary` | ✅ |
231232
| `ZSTD_DCtx_loadDictionary_advanced` | — |
232233
| `ZSTD_DCtx_loadDictionary_byReference` | — |
233234
| `ZSTD_DCtx_refDDict` | — |

‎zstd/src/main/java/io/github/dfa1/zstdffm/Bindings.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,15 @@ final class Bindings {
107107
NativeLibrary.lookup("ZSTD_decompressStream",
108108
FunctionDescriptor.of(JAVA_LONG, ADDRESS, ADDRESS, ADDRESS));
109109

110+
// size_t ZSTD_CCtx_loadDictionary(ZSTD_CCtx*, const void* dict, size_t dictSize)
111+
static final MethodHandle CCTX_LOAD_DICTIONARY =
112+
NativeLibrary.lookup("ZSTD_CCtx_loadDictionary",
113+
FunctionDescriptor.of(JAVA_LONG, ADDRESS, ADDRESS, JAVA_LONG));
114+
// size_t ZSTD_DCtx_loadDictionary(ZSTD_DCtx*, const void* dict, size_t dictSize)
115+
static final MethodHandle DCTX_LOAD_DICTIONARY =
116+
NativeLibrary.lookup("ZSTD_DCtx_loadDictionary",
117+
FunctionDescriptor.of(JAVA_LONG, ADDRESS, ADDRESS, JAVA_LONG));
118+
110119
// size_t ZSTD_CStreamInSize(void) / ZSTD_CStreamOutSize(void)
111120
static final MethodHandle CSTREAM_IN_SIZE =
112121
NativeLibrary.lookup("ZSTD_CStreamInSize", FunctionDescriptor.of(JAVA_LONG));

‎zstd/src/main/java/io/github/dfa1/zstdffm/ZstdInputStream.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,12 +44,23 @@ public final class ZstdInputStream extends InputStream {
4444
///
4545
/// @param in the stream to read the compressed frame from
4646
public ZstdInputStream(InputStream in) {
47+
this(in, null);
48+
}
49+
50+
/// Wraps `in`, decompressing a frame compressed against `dictionary`.
51+
///
52+
/// @param in the stream to read the compressed frame from
53+
/// @param dictionary the dictionary the frame was compressed against, or `null` for none
54+
public ZstdInputStream(InputStream in, ZstdDictionary dictionary) {
4755
this.in = in;
4856
try {
4957
this.dctx = (MemorySegment) Bindings.CREATE_DCTX.invokeExact();
5058
if (MemorySegment.NULL.equals(dctx)) {
5159
throw new ZstdException("ZSTD_createDCtx returned NULL");
5260
}
61+
if (dictionary != null) {
62+
loadDictionary(dictionary);
63+
}
5364
this.inCap = (long) Bindings.DSTREAM_IN_SIZE.invokeExact();
5465
this.outCap = (long) Bindings.DSTREAM_OUT_SIZE.invokeExact();
5566
} catch (Throwable t) {
@@ -62,6 +73,15 @@ public ZstdInputStream(InputStream in) {
6273
this.hold = new byte[Math.toIntExact(outCap)];
6374
}
6475

76+
private void loadDictionary(ZstdDictionary dictionary) {
77+
try (Arena staging = Arena.ofConfined()) {
78+
byte[] raw = dictionary.raw();
79+
MemorySegment dictSeg = Zstd.copyIn(staging, raw);
80+
Zstd.call(() -> (long) Bindings.DCTX_LOAD_DICTIONARY.invokeExact(
81+
dctx, dictSeg, (long) raw.length));
82+
}
83+
}
84+
6585
@Override
6686
public int read() throws IOException {
6787
byte[] one = new byte[1];

‎zstd/src/main/java/io/github/dfa1/zstdffm/ZstdOutputStream.java‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,23 @@ public ZstdOutputStream(OutputStream out) {
5555
/// @param out the stream to write the compressed frame to
5656
/// @param level the compression level
5757
public ZstdOutputStream(OutputStream out, int level) {
58+
this(out, level, null);
59+
}
60+
61+
/// Wraps `out`, compressing against `dictionary` at the default level.
62+
///
63+
/// @param out the stream to write the compressed frame to
64+
/// @param dictionary the dictionary to compress against
65+
public ZstdOutputStream(OutputStream out, ZstdDictionary dictionary) {
66+
this(out, Zstd.defaultCompressionLevel(), dictionary);
67+
}
68+
69+
/// Wraps `out`, compressing against `dictionary` at `level`.
70+
///
71+
/// @param out the stream to write the compressed frame to
72+
/// @param level the compression level
73+
/// @param dictionary the dictionary to compress against, or `null` for none
74+
public ZstdOutputStream(OutputStream out, int level, ZstdDictionary dictionary) {
5875
this.out = out;
5976
try {
6077
this.cctx = (MemorySegment) Bindings.CREATE_CCTX.invokeExact();
@@ -63,6 +80,9 @@ public ZstdOutputStream(OutputStream out, int level) {
6380
}
6481
Zstd.call(() -> (long) Bindings.CCTX_SET_PARAMETER.invokeExact(
6582
cctx, ZSTD_C_COMPRESSION_LEVEL, level));
83+
if (dictionary != null) {
84+
loadDictionary(dictionary);
85+
}
6686
this.inCap = (long) Bindings.CSTREAM_IN_SIZE.invokeExact();
6787
this.outCap = (long) Bindings.CSTREAM_OUT_SIZE.invokeExact();
6888
} catch (Throwable t) {
@@ -74,6 +94,15 @@ public ZstdOutputStream(OutputStream out, int level) {
7494
this.drain = new byte[Math.toIntExact(outCap)];
7595
}
7696

97+
private void loadDictionary(ZstdDictionary dictionary) {
98+
try (Arena staging = Arena.ofConfined()) {
99+
byte[] raw = dictionary.raw();
100+
MemorySegment dictSeg = Zstd.copyIn(staging, raw);
101+
Zstd.call(() -> (long) Bindings.CCTX_LOAD_DICTIONARY.invokeExact(
102+
cctx, dictSeg, (long) raw.length));
103+
}
104+
}
105+
77106
@Override
78107
public void write(int b) throws IOException {
79108
write(new byte[]{(byte) b}, 0, 1);

‎zstd/src/test/java/io/github/dfa1/zstdffm/ZstdStreamTest.java‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,73 @@ void oneShotOutputDecodesWithStream() throws IOException {
8181
}
8282
}
8383

84+
@Nested
85+
class Dictionary {
86+
87+
@Test
88+
void roundTripsWithDictionary() throws IOException {
89+
// Given a dictionary and a small record
90+
ZstdDictionary dict = trainDict();
91+
byte[] record = record(42);
92+
93+
// When streamed through compress and decompress with the same dictionary
94+
ByteArrayOutputStream sink = new ByteArrayOutputStream();
95+
try (ZstdOutputStream zout = new ZstdOutputStream(sink, dict)) {
96+
zout.write(record);
97+
}
98+
byte[] restored;
99+
try (ZstdInputStream zin = new ZstdInputStream(new ByteArrayInputStream(sink.toByteArray()), dict)) {
100+
restored = zin.readAllBytes();
101+
}
102+
103+
// Then the record is recovered
104+
assertThat(restored).isEqualTo(record);
105+
}
106+
107+
@Test
108+
void dictionaryShrinksStreamedRecord() throws IOException {
109+
ZstdDictionary dict = trainDict();
110+
byte[] record = record(123);
111+
112+
ByteArrayOutputStream withDict = new ByteArrayOutputStream();
113+
try (ZstdOutputStream zout = new ZstdOutputStream(withDict, dict)) {
114+
zout.write(record);
115+
}
116+
117+
// a dictionary frame of a tiny record is smaller than a plain stream frame
118+
assertThat(withDict.size()).isLessThan(streamCompress(record, Zstd.defaultCompressionLevel()).length);
119+
}
120+
121+
@Test
122+
void streamWithDictionaryDecodesWithOneShot() throws IOException {
123+
// Given a dictionary frame produced by the streaming compressor
124+
ZstdDictionary dict = trainDict();
125+
byte[] record = record(7);
126+
ByteArrayOutputStream sink = new ByteArrayOutputStream();
127+
try (ZstdOutputStream zout = new ZstdOutputStream(sink, dict)) {
128+
zout.write(record);
129+
}
130+
131+
// Then the one-shot context decodes it with the same dictionary
132+
try (ZstdDecompressCtx ctx = new ZstdDecompressCtx()) {
133+
assertThat(ctx.decompress(sink.toByteArray(), record.length, dict)).isEqualTo(record);
134+
}
135+
}
136+
137+
private ZstdDictionary trainDict() {
138+
java.util.List<byte[]> samples = new java.util.ArrayList<>();
139+
for (int i = 0; i < 3000; i++) {
140+
samples.add(record(i));
141+
}
142+
return ZstdDictionary.train(samples, 8 * 1024);
143+
}
144+
145+
private byte[] record(int i) {
146+
return ("{\"id\":" + i + ",\"user\":\"user_" + (i % 40) + "\",\"event\":\"click\"}")
147+
.getBytes(StandardCharsets.UTF_8);
148+
}
149+
}
150+
84151
private static byte[] streamCompress(byte[] data, int level) throws IOException {
85152
ByteArrayOutputStream sink = new ByteArrayOutputStream();
86153
try (ZstdOutputStream zout = new ZstdOutputStream(sink, level)) {

0 commit comments

Comments
 (0)