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]

Reply via email to