Core: ZOrderByteUtils.interleaveBits interleaves one bit at a time; a bit-spreading table is faster
- Dominant language
- Java
- Stars
- 9.2k
- Forks
- 3.5k
- Avg merge
- 2d 11h
- Merged PRs (30d)
- 132
Description
### Feature Request / Improvement
## Summary
`ZOrderByteUtils.interleaveBits` is implemented as a per-bit loop. It is the innermost kernel of Z-order clustering: it runs once per row for every `rewrite_data_files` or `sort` invocation that uses `zorder(...)`, so its cost scales with the full row count of every table being clustered. Replacing that loop with a precomputed bit-spreading table makes the common case — every column contributing the same number of bytes — **7-9x faster**, while producing bit-identical output and requiring no format or API change.
## Current implementation and why it is slow
`core/src/main/java/org/apache/iceberg/util/ZOrderByteUtils.java` (`interleaveBits`, current `main`):
```java
while (interleaveByte < interleavedSize) {
interleavedBytes[interleaveByte] |=
(columnsBinary[sourceColumn][sourceByte] & 1 << sourceBit) >>> sourceBit << interleaveBit;
--interleaveBit;
if (interleaveBit == -1) { interleaveByte++; interleaveBit = 7; }
if (interleaveByte == interleavedSize) { break; }
do {
++sourceColumn;
if (sourceColumn == columnsBinary.length) {
sourceColumn = 0;
--sourceBit;
if (sourceBit == -1) { sourceByte++; sourceBit = 7; }
}
} while (columnsBinary[sourceColumn].length <= sourceByte);
}
```
The loop body executes **once per output bit**, i.e. `8 * interleavedSize` times. For a typical Spark configuration — four Z-ordered columns and `PRIMITIVE_BUFFER_SIZE = 8`, giving a 32-byte output — that is **256 iterations per row**, each of which performs:
- two dependent loads (`columnsBinary[sourceColumn]`, then `[sourceByte]`), plus a read-modify-write of `interleavedBytes[interleaveByte]`;
- three shifts and a mask to move a single bit;
- an inner `do { ... } while (columnsBinary[sourceColumn].length <= sourceByte)` that re-reads an array length and whose trip count depends on the data. Combined with the loop-carried dependency chain — `interleaveBit`, `sourceBit`, and `sourceColumn` all update conditionally — this leaves the hardware almost no instruction-level parallelism to extract.
The generality that this per-bit dispatch buys — support for columns of differing lengths, where a column that runs out of bytes is skipped — is not exercised by any caller in the repository. `SparkZOrderUDF` (v3.5, v4.0, v4.1) always produces fixed-width contributions: `PRIMITIVE_BUFFER_SIZE` for primitives and a single configured `varLengthContribution` for strings/binary. In other words, the hot path pays for a branch structure that only the case it never takes requires.
## Proposed implementation: precomputed bit-spreading table
For a fixed column count `n`, the mapping from a source byte to its output bits is a fixed permutation: source bit `i` (counted from the MSB) of byte `j` of column `c` lands at output bit `i * n + c` of the `n`-byte output group that starts at byte `j * n`. The permutation depends only on the pair `(n, byteValue)`, so it can be tabulated.
Define `SPREAD[n][b]` as the byte `b` with its 8 bits spread `n` positions apart. The highest bit set is `8n - 1`, so the result fits in a `long` for all `n <= 8`:
```java
private static final int MAX_LUT_COLUMNS = 8;
private static final long[][] SPREAD = buildSpread();
private static long[][] buildSpread() {
long[][] tables = new long[MAX_LUT_COLUMNS + 1][];
for (int n = 1; n <= MAX_LUT_COLUMNS; n++) {
long[] table = new long[256];
for (int b = 0; b < 256; b++) {
long spread = 0L;
for (int i = 0; i < 8; i++) { // i counts source bits from the MSB
if ((b & (1 << (7 - i))) != 0) {
spread |= 1L << (8 * n - 1 - i * n);
}
}
table[b] = spread;
}
tables[n] = table;
}
return tables;
}
```
Interleaving one group of `n` output bytes then costs `n` table lookups and `n` shift-or operations. The shift by `c` encodes exactly the offset by which column `c` is displaced within each group:
```java
long chunk = 0L;
for (int c = 0; c < numColumns; c++) {
chunk |= spread[columnsBinary[c][j] & 0xFF] >>> c;
}
for (int k = numColumns - 1; k >= 0; k--) {
interleavedBytes[out++] = (byte) (chunk >>> (8 * k));
}
```
This turns `8n` loop iterations per group into `n` lookups. The loads are independent, the shift counts are known per column rather than recomputed per bit, and no branch depends on the data.
### Preconditions and fallback
The table applies when the layout is uniform:
- `1 <= columnsBinary.length <= 8` (bounded by the 64-bit accumulator);
- every column has the same, non-zero length;
- `interleavedSize <= columnLength * numColumns`.
Verifying uniformity costs `n` length comparisons. Every other input — in particular the ragged case, where a short column drops out of the interleaving partway through — falls through to the existing loop, which remains unchanged. Truncated output, where `interleavedSize` is smaller than the full interleaving, is handled by emitting only the leading bytes of the final group: with uniform columns no column is exhausted early, so the bit order matches exactly what the current loop produces before it stops.
**This is not a format change.** The output is byte-for-byte identical to the current implementation's for every input, so existing Z-ordered data and the sort order it induces are unaffected.
## Measurements
Measured with a standalone harness in which both implementations were extracted verbatim, using the same input shape as the in-repo `ZOrderByteUtilsBenchmark`: random 8-byte columns, a single reused output buffer, 1M rows, best of 5 runs after 3 warmup rounds. JDK 17.0.19, Intel i7-13700H.
| columns | output | current | table | speedup |
|---|---|---|---|---|
| 2 | 16 B | 348.8 ns/row | 49.6 ns/row | 7.03x |
| 2 | 8 B | 166.3 ns/row | 25.3 ns/row | 6.56x |
| 3 | 24 B | 518.2 ns/row | 67.6 ns/row | 7.67x |
| 3 | 8 B | 174.2 ns/row | 25.3 ns/row | 6.90x |
| 4 | 32 B | 665.1 ns/row | 73.1 ns/row | 9.11x |
| 4 | 8 B | 168.4 ns/row | 21.8 ns/row | 7.71x |
The speedup grows with output width, as expected, because the fixed per-row costs — the uniformity check and the `Arrays.fill` — amortize over more groups.
Correctness was checked by differential testing against the current implementation. Across 250,000 randomized cases — 1 to 8 columns, column lengths from 1 to 10, every output size from 1 byte up to the full interleaving, plus ragged-length inputs that exercise the fallback — both implementations produced identical byte arrays.
I will also run the repository's own `ZOrderByteUtilsBenchmark` (`core/src/jmh/java/org/apache/iceberg/util/ZOrderByteUtilsBenchmark.java`) before and after the change, and post the JMH output alongside the pull request.
## Cost
- **Memory**: `long[9][256]` is roughly 18 KB of static data. Only the 2 KB row for the active column count is ever touched, so the working set stays resident in L1. If 18 KB is considered too much for a utility class, a nibble-indexed variant (`long[9][16]`, roughly 1 KB) produces the same result with two lookups per byte, at a small throughput cost.
- **Code**: roughly 40 additional lines, with the existing loop retained as the fallback.
- **Alternatives**: `Long.expand`, which compiles down to `PDEP` on x86, would express the same operation without a table, but it requires Java 19 and Iceberg targets Java 17 (`build.gradle`, `sourceCompatibility = "17"`). It is therefore not available today. The table also has the advantage of being portable across architectures: on hardware where `PDEP` is microcoded and slow, such as AMD before Zen 3, the intrinsic would be the worse choice.
## Scope
The change is limited to `ZOrderByteUtils.interleaveBits`. The `*ToOrderedBytes` conversions and the public method signature are untouched, so the diff is confined to a single method, its lookup table, and the accompanying tests.
### Query engine
Spark
### Willingness to contribute
- [x] I can contribute this improvement/feature independently
Contributor guide
Research direction
Start in core/src/main/java/org/apache/iceberg/util/ZOrderByteUtils.java at interleaveBits, then inspect core/src/jmh/java/org/apache/iceberg/util/ZOrderByteUtilsBenchmark.java and the existing callers described in the issue. Run the benchmark and verify that the optimized uniform-layout path preserves byte-for-byte output, while other inputs retain the existing fallback behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- java, spark
- Domain
- data-engineering, performance
- Issue type
- Refactor
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 76/100