amogh-jahagirdar commented on code in PR #16747: URL: https://github.com/apache/iceberg/pull/16747#discussion_r4006181988
########## core/src/test/java/org/apache/iceberg/mumbling/TestMumblingBitmap.java: ########## @@ -0,0 +1,352 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg.mumbling; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import java.nio.ByteBuffer; +import java.util.Arrays; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.junit.jupiter.api.Test; + +class TestMumblingBitmap { + + @Test + void testEmptyBitmap() { + MumblingBitmap bitmap = bitmap(); + assertThat(bitmap.cardinality()).isEqualTo(0); + + // all positions beyond the bitmap range are false + assertThat(bitmap.isSet(0)).isFalse(); + assertThat(bitmap.isSet(255)).isFalse(); + assertThat(bitmap.isSet(256)).isFalse(); + } + + @Test + void testInvalidPosition() { + MumblingBitmap bitmap = bitmap(); + assertThat(bitmap.cardinality()).isEqualTo(0); + assertThat(bitmap.isSet(0)).isFalse(); + assertThatThrownBy(() -> bitmap.isSet(-1)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Invalid bit position: -1 < 0"); + } + + @Test + void testEmptySparseContainer() { + MumblingBitmap bitmap = bitmap(sparse()); + assertThat(bitmap.cardinality()).isEqualTo(0); + assertThat(bitmap.isSet(0)).isFalse(); + assertThat(bitmap.isSet(100)).isFalse(); + assertThat(bitmap.isSet(255)).isFalse(); + } + + @Test + void testSparseContainerSetPositions() { + MumblingBitmap bitmap = bitmap(sparse(0, 5, 100, 255)); + assertThat(bitmap.cardinality()).isEqualTo(4); + + assertThat(bitmap.isSet(0)).isTrue(); + assertThat(bitmap.isSet(5)).isTrue(); + assertThat(bitmap.isSet(100)).isTrue(); + assertThat(bitmap.isSet(255)).isTrue(); + + assertThat(bitmap.isSet(1)).isFalse(); + assertThat(bitmap.isSet(4)).isFalse(); + assertThat(bitmap.isSet(6)).isFalse(); + assertThat(bitmap.isSet(99)).isFalse(); + assertThat(bitmap.isSet(101)).isFalse(); + assertThat(bitmap.isSet(254)).isFalse(); + assertThat(bitmap.isSet(256)).isFalse(); + } + + @Test + void testFullSparseContainer() { + int[] positions = new int[31]; + for (int i = 0; i < 31; i += 1) { + positions[i] = i * 8; // 0, 8, 16, ..., 240 + } + + MumblingBitmap bitmap = bitmap(sparse(positions)); + assertThat(bitmap.cardinality()).isEqualTo(31); + + for (int p : positions) { + assertThat(bitmap.isSet(p)).isTrue(); + } + + assertThat(bitmap.isSet(1)).isFalse(); + assertThat(bitmap.isSet(7)).isFalse(); + assertThat(bitmap.isSet(255)).isFalse(); + } + + @Test + void testFullDenseContainer() { + byte[] container = new byte[32]; + Arrays.fill(container, (byte) 0xFF); + + MumblingBitmap bitmap = bitmap(dense(container)); + assertThat(bitmap.cardinality()).isEqualTo(256); + + for (int i = 0; i < 256; i += 1) { + assertThat(bitmap.isSet(i)).isTrue(); + } + + assertThat(bitmap.isSet(256)).isFalse(); + } + + // Example 1: positions 0-31: `FF FF FF FF 00 ... 00` + @Test + void testDenseSpecExample1() { + byte[] container = new byte[32]; + container[0] = (byte) 0xFF; + container[1] = (byte) 0xFF; + container[2] = (byte) 0xFF; + container[3] = (byte) 0xFF; + MumblingBitmap bitmap = bitmap(dense(container)); + assertThat(bitmap.cardinality()).isEqualTo(32); + + for (int i = 0; i <= 31; i += 1) { + assertThat(bitmap.isSet(i)).isTrue(); + } + + assertThat(bitmap.isSet(32)).isFalse(); + assertThat(bitmap.isSet(255)).isFalse(); + } + + // Example 2: positions 0-32: `FF FF FF FF 80 00 ... 00` + @Test + void testDenseSpecExample2() { + byte[] container = new byte[32]; + container[0] = (byte) 0xFF; + container[1] = (byte) 0xFF; + container[2] = (byte) 0xFF; + container[3] = (byte) 0xFF; + container[4] = (byte) 0x80; + + MumblingBitmap bitmap = bitmap(dense(container)); + assertThat(bitmap.cardinality()).isEqualTo(33); + + for (int i = 0; i <= 32; i += 1) { + assertThat(bitmap.isSet(i)).isTrue(); + } + + assertThat(bitmap.isSet(33)).isFalse(); + assertThat(bitmap.isSet(255)).isFalse(); + } + + // Example 3: positions 0-15 and 240-255: `FF FF 00 ... 00 FF FF` + @Test + void testDenseSpecExample3() { + byte[] container = new byte[32]; + container[0] = (byte) 0xFF; + container[1] = (byte) 0xFF; + container[30] = (byte) 0xFF; + container[31] = (byte) 0xFF; + + MumblingBitmap bitmap = bitmap(dense(container)); + assertThat(bitmap.cardinality()).isEqualTo(32); + + for (int i = 0; i <= 15; i += 1) { + assertThat(bitmap.isSet(i)).isTrue(); + } + for (int i = 240; i <= 255; i += 1) { + assertThat(bitmap.isSet(i)).isTrue(); + } + assertThat(bitmap.isSet(16)).isFalse(); + assertThat(bitmap.isSet(239)).isFalse(); + assertThat(bitmap.isSet(256)).isFalse(); + } + + // Example 4: even positions 0, 2, 4, ...: `AA AA ... AA AA` + @Test + void testDenseSpecExample4() { + byte[] container = new byte[32]; + Arrays.fill(container, (byte) 0xAA); + + MumblingBitmap bitmap = bitmap(dense(container)); + assertThat(bitmap.cardinality()).isEqualTo(128); + + for (int i = 0; i < 256; i += 1) { + assertThat(bitmap.isSet(i)).isEqualTo(i % 2 == 0); + } + + assertThat(bitmap.isSet(256)).isFalse(); + } + + @Test + void testMultipleContainers() { + MumblingBitmap bitmap = bitmap(sparse(5), sparse(), sparse(10)); + assertThat(bitmap.cardinality()).isEqualTo(2); + + assertThat(bitmap.isSet(5)).isTrue(); // container 0, pos 5 + assertThat(bitmap.isSet(256)).isFalse(); // container 1 + assertThat(bitmap.isSet(522)).isTrue(); // container 2, pos 10 + + assertThat(bitmap.isSet(512)).isFalse(); + assertThat(bitmap.isSet(4)).isFalse(); + assertThat(bitmap.isSet(265)).isFalse(); + assertThat(bitmap.isSet(267)).isFalse(); + } + + @Test + void testMixedSparseAndDense() { Review Comment: `testMixedSparseAndDense` doesn't seem to hit any PFOR exception cases. Could we add one, e.g. descriptors like `[1, 1, ..., 1, 32]` where the `32` becomes an exception? Or if it's not apppropriate for this unit test, a separate one. I see the PFOR encoding tests already cover exception cases, but adding one here would catch any issues around decoding the descriptor array at a non-zero offset. ########## core/src/test/java/org/apache/iceberg/mumbling/TestMumblingBitmap.java: ########## @@ -0,0 +1,352 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg.mumbling; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import java.nio.ByteBuffer; +import java.util.Arrays; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.junit.jupiter.api.Test; + +class TestMumblingBitmap { + + @Test + void testEmptyBitmap() { + MumblingBitmap bitmap = bitmap(); + assertThat(bitmap.cardinality()).isEqualTo(0); + + // all positions beyond the bitmap range are false + assertThat(bitmap.isSet(0)).isFalse(); + assertThat(bitmap.isSet(255)).isFalse(); + assertThat(bitmap.isSet(256)).isFalse(); + } + + @Test + void testInvalidPosition() { + MumblingBitmap bitmap = bitmap(); + assertThat(bitmap.cardinality()).isEqualTo(0); + assertThat(bitmap.isSet(0)).isFalse(); + assertThatThrownBy(() -> bitmap.isSet(-1)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Invalid bit position: -1 < 0"); + } + + @Test + void testEmptySparseContainer() { + MumblingBitmap bitmap = bitmap(sparse()); + assertThat(bitmap.cardinality()).isEqualTo(0); + assertThat(bitmap.isSet(0)).isFalse(); + assertThat(bitmap.isSet(100)).isFalse(); + assertThat(bitmap.isSet(255)).isFalse(); + } + + @Test + void testSparseContainerSetPositions() { + MumblingBitmap bitmap = bitmap(sparse(0, 5, 100, 255)); + assertThat(bitmap.cardinality()).isEqualTo(4); + + assertThat(bitmap.isSet(0)).isTrue(); + assertThat(bitmap.isSet(5)).isTrue(); + assertThat(bitmap.isSet(100)).isTrue(); + assertThat(bitmap.isSet(255)).isTrue(); + + assertThat(bitmap.isSet(1)).isFalse(); + assertThat(bitmap.isSet(4)).isFalse(); + assertThat(bitmap.isSet(6)).isFalse(); + assertThat(bitmap.isSet(99)).isFalse(); + assertThat(bitmap.isSet(101)).isFalse(); + assertThat(bitmap.isSet(254)).isFalse(); + assertThat(bitmap.isSet(256)).isFalse(); + } + + @Test + void testFullSparseContainer() { + int[] positions = new int[31]; + for (int i = 0; i < 31; i += 1) { + positions[i] = i * 8; // 0, 8, 16, ..., 240 + } + + MumblingBitmap bitmap = bitmap(sparse(positions)); + assertThat(bitmap.cardinality()).isEqualTo(31); + + for (int p : positions) { + assertThat(bitmap.isSet(p)).isTrue(); + } + + assertThat(bitmap.isSet(1)).isFalse(); + assertThat(bitmap.isSet(7)).isFalse(); + assertThat(bitmap.isSet(255)).isFalse(); + } + + @Test + void testFullDenseContainer() { + byte[] container = new byte[32]; + Arrays.fill(container, (byte) 0xFF); + + MumblingBitmap bitmap = bitmap(dense(container)); + assertThat(bitmap.cardinality()).isEqualTo(256); + + for (int i = 0; i < 256; i += 1) { + assertThat(bitmap.isSet(i)).isTrue(); + } + + assertThat(bitmap.isSet(256)).isFalse(); + } + + // Example 1: positions 0-31: `FF FF FF FF 00 ... 00` + @Test + void testDenseSpecExample1() { + byte[] container = new byte[32]; + container[0] = (byte) 0xFF; + container[1] = (byte) 0xFF; + container[2] = (byte) 0xFF; + container[3] = (byte) 0xFF; + MumblingBitmap bitmap = bitmap(dense(container)); + assertThat(bitmap.cardinality()).isEqualTo(32); + + for (int i = 0; i <= 31; i += 1) { + assertThat(bitmap.isSet(i)).isTrue(); + } + + assertThat(bitmap.isSet(32)).isFalse(); + assertThat(bitmap.isSet(255)).isFalse(); + } + + // Example 2: positions 0-32: `FF FF FF FF 80 00 ... 00` + @Test + void testDenseSpecExample2() { + byte[] container = new byte[32]; + container[0] = (byte) 0xFF; + container[1] = (byte) 0xFF; + container[2] = (byte) 0xFF; + container[3] = (byte) 0xFF; + container[4] = (byte) 0x80; + + MumblingBitmap bitmap = bitmap(dense(container)); + assertThat(bitmap.cardinality()).isEqualTo(33); + + for (int i = 0; i <= 32; i += 1) { + assertThat(bitmap.isSet(i)).isTrue(); + } + + assertThat(bitmap.isSet(33)).isFalse(); + assertThat(bitmap.isSet(255)).isFalse(); + } + + // Example 3: positions 0-15 and 240-255: `FF FF 00 ... 00 FF FF` + @Test + void testDenseSpecExample3() { + byte[] container = new byte[32]; + container[0] = (byte) 0xFF; + container[1] = (byte) 0xFF; + container[30] = (byte) 0xFF; + container[31] = (byte) 0xFF; + + MumblingBitmap bitmap = bitmap(dense(container)); + assertThat(bitmap.cardinality()).isEqualTo(32); + + for (int i = 0; i <= 15; i += 1) { + assertThat(bitmap.isSet(i)).isTrue(); + } + for (int i = 240; i <= 255; i += 1) { + assertThat(bitmap.isSet(i)).isTrue(); + } + assertThat(bitmap.isSet(16)).isFalse(); + assertThat(bitmap.isSet(239)).isFalse(); + assertThat(bitmap.isSet(256)).isFalse(); + } + + // Example 4: even positions 0, 2, 4, ...: `AA AA ... AA AA` + @Test + void testDenseSpecExample4() { + byte[] container = new byte[32]; + Arrays.fill(container, (byte) 0xAA); + + MumblingBitmap bitmap = bitmap(dense(container)); + assertThat(bitmap.cardinality()).isEqualTo(128); + + for (int i = 0; i < 256; i += 1) { + assertThat(bitmap.isSet(i)).isEqualTo(i % 2 == 0); + } + + assertThat(bitmap.isSet(256)).isFalse(); + } + + @Test + void testMultipleContainers() { + MumblingBitmap bitmap = bitmap(sparse(5), sparse(), sparse(10)); + assertThat(bitmap.cardinality()).isEqualTo(2); + + assertThat(bitmap.isSet(5)).isTrue(); // container 0, pos 5 + assertThat(bitmap.isSet(256)).isFalse(); // container 1 + assertThat(bitmap.isSet(522)).isTrue(); // container 2, pos 10 + + assertThat(bitmap.isSet(512)).isFalse(); + assertThat(bitmap.isSet(4)).isFalse(); + assertThat(bitmap.isSet(265)).isFalse(); + assertThat(bitmap.isSet(267)).isFalse(); + } + + @Test + void testMixedSparseAndDense() { + byte[] denseContainer = new byte[32]; + denseContainer[0] = (byte) 0xFF; + denseContainer[1] = (byte) 0xFF; + denseContainer[2] = (byte) 0xFF; + denseContainer[3] = (byte) 0xFF; + + MumblingBitmap bitmap = bitmap(dense(denseContainer), sparse(1)); + assertThat(bitmap.cardinality()).isEqualTo(33); + + for (int i = 0; i < 32; i += 1) { + assertThat(bitmap.isSet(i)).isTrue(); + } + assertThat(bitmap.isSet(32)).isFalse(); + + assertThat(bitmap.isSet(256)).isFalse(); + assertThat(bitmap.isSet(257)).isTrue(); // container 1, pos 1 + assertThat(bitmap.isSet(258)).isFalse(); + } + + @Test + void testBufferWithOffset() { + // Prepend 4 bytes of garbage before the actual bitmap data + ByteBuffer buffer = build(sparse(42)); + byte[] rawBytes = new byte[buffer.remaining()]; + buffer.get(rawBytes); + + ByteBuffer padded = ByteBuffer.allocate(4 + rawBytes.length); + padded.position(4); + padded.put(rawBytes); + padded.position(4); // position the buffer at the start of bitmap data + + MumblingBitmap bitmap = new MumblingBitmap(padded); + assertThat(bitmap.cardinality()).isEqualTo(1); + + assertThat(bitmap.isSet(41)).isFalse(); + assertThat(bitmap.isSet(42)).isTrue(); + assertThat(bitmap.isSet(43)).isFalse(); + } + + private static Container sparse(int... positions) { + byte[] bytes = new byte[positions.length]; + for (int i = 0; i < positions.length; i += 1) { + if (i > 0) { + Preconditions.checkArgument( + positions[i] < 256, "Invalid position in container: %s", positions[i]); + Preconditions.checkArgument( + positions[i] > positions[i - 1], + "Invalid sparse container: pos %s=%s >= pos %s=%s", + i - 1, + positions[i - 1], + i, + positions[i]); + } + + bytes[i] = (byte) positions[i]; + } + + return new Container(bytes); + } + + /** Descriptor + bytes for a dense container. */ + private static Container dense(byte[] container) { + Preconditions.checkArgument(container.length == 32, "Dense container must be 32 bytes"); + return new Container(container); + } + + private static class Container { + private final byte[] bytes; + private final int descriptor; + private final int cardinality; + + Container(byte[] bytes) { + this.bytes = bytes; + this.descriptor = bytes.length; Review Comment: If I follow right, this is always hardcoded to 0x20. If that's right, I don't think the tests cover a case where we have dense descriptor and with the low bits set. The spec says "The two most significant bits are reserved for future use and implementations must ignore the least-significant bits for a dense container if they are set." . Should we add a case, something like 0x2F and make sure the container is just read like a normal dense bitmap? ########## core/src/main/java/org/apache/iceberg/mumbling/MumblingBitmap.java: ########## @@ -0,0 +1,168 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg.mumbling; + +import java.nio.ByteBuffer; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; + +/** + * Read-only view of a Mumbling compressed bitmap stored in a {@link ByteBuffer}. + * + * <p>The bitmap is lazy: no decoding is done at construction time. On the first call to {@link + * #isSet}, the PFOR-encoded descriptor array is decoded and used to build an offsets array that + * maps each container index to its absolute byte position in the buffer. This offsets array is the + * only derived state kept by this class. + * + * <p>Format (all integers unsigned, little-endian): + * + * <ul> + * <li>Header (6 bytes): version (1), cardinality (3), container count (2) + * <li>Descriptor array: PFOR-encoded, one byte per container + * <li>Containers: concatenated sparse (0–31 bytes) or dense (32 bytes) containers + * </ul> + */ +class MumblingBitmap { + private static final int VERSION = 1; + private static final int HEADER_SIZE = 6; + private static final int DENSE_CONTAINER_BIT = 0b0010_0000; + + private final ByteBuffer data; + private final int cardinality; + private final int containerCount; + private int[] descriptors = null; + private int[] offsets = null; Review Comment: Just to harden this path a bit, should these be made as `volatile` and these fields are only set after they are actually populated. I think that avoids any subtle concurrency bugs in case for whatever reason a single bitmap reference is decoded and acceessed across different threads. While I wouldn't expect there to be more 1 thread reading a given bitmap, seems like we can avoid those cases in a simple manner and allow for flexible consumption patterns to callers. ########## core/src/main/java/org/apache/iceberg/mumbling/PFOREncoding.java: ########## @@ -0,0 +1,402 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg.mumbling; + +import java.nio.ByteBuffer; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.util.ByteBuffers; +import org.apache.iceberg.util.Pair; + +/** + * Patched Frame of Reference (PFOR) encoding for arrays of unsigned byte values. + * + * <p>Implements the encoding described in Appendix A of the Mumbling bitmap specification. The + * input array is split into 256-value chunks (the last chunk may be shorter). Each chunk is + * independently encoded using 4 configuration values: + * + * <ul> + * <li>{@code b1}: number of bits stored in the primary array for every normalized value + * <li>{@code b2}: number of bits stored per exception value (normalized value of > b1 bits) + * <li>{@code e}: number of exceptions with more than b1 bits + * <li>{@code m}: chunk-local minimum value, subtracted from all values to normalize + * </ul> + * + * <p>Each chunk is stored as: + * + * <ul> + * <li>3-byte header: {@code b1|b2} primary and exception bit widths (byte 0), {@code e} exception + * count (byte 1), {@code m} normalization base value (byte 2) + * <li>Primary array: the low {@code b1} bits of every normalized value, packed MSB-first ({@code + * b1 * n} bits, padded to a byte) + * <li>Exception offsets: chunk-relative positions of exception values ({@code e} bytes) + * <li>Exception values: the high {@code b2} bits of every exception value, packed MSB-first + * ({@code e * b2} bits, padded to a byte. + * </ul> + */ +class PFOREncoding { + private static final int CHUNK_SIZE = 256; + + private PFOREncoding() {} + + /** + * Encodes {@code count} values from an array of unsigned byte values. + * + * @param values unsigned byte values to encode + * @param count number of values to encode + * @return a {@link ByteBuffer} of the encoded values with position and limit set for reading + */ + static ByteBuffer encode(int[] values, int count) { + ByteBuffer out = ByteBuffer.allocate(estimateEncodedSize(count)); + int bytesWritten = encode(values, 0, out, 0, count); + return out.slice(0, bytesWritten); + } + + /** + * Encode {@code count} unsigned byte values from {@code values} into a buffer. + * + * <p>The buffer's position and limit are not modified. + * + * @param values unsigned byte values to encode + * @param valueOffset starting offset of values to encode + * @param out buffer to write encoded values to + * @param outOffset starting offset in the output buffer + * @param count number of values to encode + * @return the number of bytes written to the buffer + */ + static int encode(int[] values, int valueOffset, ByteBuffer out, int outOffset, int count) { + // outOffset is relative to the buffer's position; check the encoded data fits + Preconditions.checkArgument(outOffset >= 0, "Cannot encode at negative offset: %s", outOffset); + Preconditions.checkArgument( + estimateEncodedSize(count) <= out.remaining() - outOffset, + "Cannot encode %s values to buffer with %s remaining space", + count, + out.remaining() - outOffset); + + int bytesWritten = 0; + int valuesEncoded = 0; + + while (valuesEncoded < count) { + int chunkLength = Math.min(CHUNK_SIZE, count - valuesEncoded); + bytesWritten += + encodeChunk( + values, valueOffset + valuesEncoded, out, outOffset + bytesWritten, chunkLength); + valuesEncoded += chunkLength; + } + + return bytesWritten; + } + + /** + * Decode to produce unsigned byte values. + * + * <p>Decodes starting at {@code encoded.position()} and does not modify the input buffer. + * + * @param encoded PFOR-encoded ByteBuffer produced by {@link #encode} + * @param count total number of values to decode + * @return decoded unsigned byte values + */ + static int[] decode(ByteBuffer encoded, int count) { + int[] out = new int[count]; + decode(encoded, 0, out, 0, count); + return out; + } + + /** + * Decode {@code count} unsigned bytes from a buffer into {@code out}. + * + * <p>This does not modify the input buffer. + * + * @param encoded a buffer containing encoded data + * @param offset starting offset of encoded values + * @param out an output value array + * @param outOffset starting offset in the output array + * @param count number of values to decode + * @return the number of bytes read from the encoded buffer + */ + static int decode(ByteBuffer encoded, int offset, int[] out, int outOffset, int count) { + Preconditions.checkArgument(offset >= 0, "Cannot decode at negative offset: %s", offset); + + int bytesRead = 0; + int valuesRead = 0; + + while (valuesRead < count) { + int chunkSize = Math.min(CHUNK_SIZE, count - valuesRead); + bytesRead += decodeChunk(encoded, offset + bytesRead, out, outOffset + valuesRead, chunkSize); + valuesRead += chunkSize; + } + + return bytesRead; + } + + /** + * Encode one chunk into {@code out} starting at absolute position {@code outPos}. + * + * @param values array containing source values to encode + * @param valueOffset starting index of values to encode + * @param out an output {@link ByteBuffer} + * @param outOffset starting index for output in the out buffer + * @param count number of values to encode + * @return the number of bytes written to the output buffer + */ + private static int encodeChunk( + int[] values, int valueOffset, ByteBuffer out, int outOffset, int count) { + Preconditions.checkArgument(count >= 0, "Invalid value count to encode: %s", count); + Preconditions.checkArgument( + valueOffset + count <= values.length, + "Cannot encode %s values starting at %s from int[%s]: not enough values", + count, + valueOffset, + values.length); + + // find base=min(values) for normalization + int base = min(values, valueOffset, count); + + // normalize by subtracting base + int[] normalized = new int[count]; + int setBits = 0; + int normalizedSetBits = 0; + for (int i = 0; i < count; i += 1) { + setBits |= values[valueOffset + i]; + normalized[i] = values[valueOffset + i] - base; + normalizedSetBits |= normalized[i]; + } + + Preconditions.checkArgument( + width(setBits) <= 8, + "Cannot encode values wider than 8 bits: %s bits needed", + width(setBits)); + + // Choose b1 to minimize total encoded data size (excluding 3-byte header) + int maxWidth = width(normalizedSetBits); + Pair<Integer, Integer> widthAndExcCount = chooseBitWidth(normalized, count, maxWidth); + int b1 = widthAndExcCount.first(); + int b2 = maxWidth - b1; + int excCount = widthAndExcCount.second(); + + // check that there is enough space in the buffer for the encoded data + int requiredSize = encodedSize(count, b1, b2, excCount); + Preconditions.checkArgument( + outOffset + requiredSize <= out.remaining(), + "Cannot encode %s values (%s bytes) into buffer with %s remaining bytes", + requiredSize, + out.remaining() - outOffset); + + // Special case: b1=8 means store original values as raw bytes with b2, e, and m set to 0. + if (b1 == 8) { + writeHeader(out, outOffset, b1, 0 /* b2 */, 0 /* excCount */, 0 /* m */); + return 3 + BitPacking.packBits(8, values, valueOffset, out, outOffset + 3, count); + } + + int bytesWritten = writeHeader(out, outOffset, b1, b2, excCount, base); + + // Primary array: low b1 bits of every value + bytesWritten += BitPacking.packBits(b1, normalized, 0, out, outOffset + bytesWritten, count); + + // b2 is the bit width of exception values: (maxWidth - b1) bits of each exception + if (excCount > 0) { + int[] excOffsets = new int[excCount]; + int[] excValues = new int[excCount]; + + // Collect exceptions (values that do not fit in b1 bits) + int excIndex = 0; + int threshold = 1 << b1; + for (int i = 0; i < count; i += 1) { + if (normalized[i] >= threshold) { + excOffsets[excIndex] = i; + excValues[excIndex] = normalized[i] >>> b1; + excIndex += 1; + } + } + + // Exception offsets (one byte per exception) + bytesWritten += + BitPacking.packBits(8, excOffsets, 0, out, outOffset + bytesWritten, excCount); + + // Exception values: remaining high b2 bits of each exception + bytesWritten += + BitPacking.packBits(b2, excValues, 0, out, outOffset + bytesWritten, excCount); + } + + return bytesWritten; + } + + /** + * Decode one chunk of encoded data, writing decoded values into an output array. + * + * @param data buffer containing source data to decode + * @param dataOffset starting index in the buffer to decode + * @param out an output {@link ByteBuffer} + * @param outOffset starting index for output in the out buffer + * @param count number of values to decode + * @return the number of bytes read from {@code data} + */ + private static int decodeChunk( + ByteBuffer data, int dataOffset, int[] out, int outOffset, int count) { + Preconditions.checkArgument(count >= 0, "Invalid value count to decode: %s", count); + Preconditions.checkArgument( + outOffset + count <= out.length, + "Cannot decode %s values starting at %s into int[%s]: not enough space", + count, + outOffset, + out.length); + + int b1 = ByteBuffers.readByte(data, dataOffset) & 0x0F; + int b2 = (ByteBuffers.readByte(data, dataOffset) >>> 4) & 0x0F; + int excCount = ByteBuffers.readByte(data, dataOffset + 1); + int base = ByteBuffers.readByte(data, dataOffset + 2); + int bytesRead = 3; + + // after reading the header, check that the full chunk is present + int expectedSize = encodedSize(count, b1, b2, excCount); + Preconditions.checkArgument( + dataOffset + expectedSize <= data.remaining(), + "Cannot decode %s values from buffer with %s remaining bytes", + expectedSize, + data.remaining() - dataOffset); + + // Read primary array: low b1 bits of each value + bytesRead += BitPacking.unpackBits(b1, data, dataOffset + bytesRead, out, outOffset, count); + + // Read exceptions and update output values + if (excCount > 0) { + int[] excOffsets = new int[excCount]; + int[] excValues = new int[excCount]; + int excListOffset = dataOffset + bytesRead; + int excDataOffset = dataOffset + bytesRead + excCount; + + // Read exception indexes + bytesRead += BitPacking.unpackBits(8, data, excListOffset, excOffsets, 0, excCount); + + // Read exception values and patch the primary values + bytesRead += BitPacking.unpackBits(b2, data, excDataOffset, excValues, 0, excCount); + + // Update output values + for (int i = 0; i < excCount; i += 1) { + out[outOffset + excOffsets[i]] |= excValues[i] << b1; + } + } + + // Add back the chunk minimum + for (int i = 0; i < count; i += 1) { + out[outOffset + i] += base; + } + + return bytesRead; + } + + private static int writeHeader( + ByteBuffer out, int outOffset, int b1, int b2, int excCount, int base) { + // Header: b1 in low nibble, b2 in high nibble, then e, then m + ByteBuffers.writeByte(out, (b2 << 4) | (b1 & 0b1111), outOffset); + ByteBuffers.writeByte(out, excCount, outOffset + 1); + ByteBuffers.writeByte(out, base, outOffset + 2); + + return 3; + } + + /** + * Choose the primary bit width {@code b1} that minimizes total encoded chunk size. + * + * <p>This produces the width that results in the smallest total size and the number of exceptions Review Comment: The spec example encodes `[6, 34, 8, 7]` as `b1=2` with one exception, but this function picks `b1=5` with no exceptions. Both come out to the same size and both decode correctly, so this is a valid choice, just different from the example bytes in the spec. -- 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]
