GGraziadei opened a new issue, #17758:
URL: https://github.com/apache/iceberg/issues/17758
### 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
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]