diff --git a/core/src/main/java/io/github/dfa1/vortex/core/model/PType.java b/core/src/main/java/io/github/dfa1/vortex/core/model/PType.java index a682f74b..342b767c 100644 --- a/core/src/main/java/io/github/dfa1/vortex/core/model/PType.java +++ b/core/src/main/java/io/github/dfa1/vortex/core/model/PType.java @@ -94,4 +94,23 @@ public static PType fromOrdinal(int ordinal) { } return all[ordinal]; } + + /// Returns the narrowest unsigned integer ptype that can hold every value in + /// `0..maxValue` — `U8` up to `0xFF`, `U16` up to `0xFFFF`, else `U32`. + /// + /// Used writer-side to size index/length/offset buffers (patch indices, FSST row + /// lengths and code offsets) to the smallest width that still fits, instead of + /// always paying for a 4-byte `U32`. + /// + /// @param maxValue the largest value the buffer must represent, `>= 0` + /// @return `U8`, `U16`, or `U32`, whichever is smallest and still fits + public static PType narrowestUnsigned(long maxValue) { + if (maxValue <= 0xFFL) { + return U8; + } + if (maxValue <= 0xFFFFL) { + return U16; + } + return U32; + } } diff --git a/core/src/test/java/io/github/dfa1/vortex/core/model/PTypeTest.java b/core/src/test/java/io/github/dfa1/vortex/core/model/PTypeTest.java index 48025834..676d2d26 100644 --- a/core/src/test/java/io/github/dfa1/vortex/core/model/PTypeTest.java +++ b/core/src/test/java/io/github/dfa1/vortex/core/model/PTypeTest.java @@ -95,4 +95,21 @@ void fromOrdinal_zeroIsU8() { // Given / When / Then — wire-format anchor: U8 must remain ordinal 0 assertThat(PType.fromOrdinal(0)).isSameAs(PType.U8); } + + @ParameterizedTest + @CsvSource({ + // boundary values: each tier's inclusive ceiling and the first value that spills to the next + "0, U8", + "255, U8", // 0xFF — last U8 + "256, U16", // 0x100 — first U16 + "65535, U16", // 0xFFFF — last U16 + "65536, U32", // 0x10000 — first U32 + "4294967295, U32", // 0xFFFFFFFF — beyond U16, still U32 + }) + void narrowestUnsigned_picksSmallestFittingTier(long maxValue, PType expected) { + // Given / When + PType result = PType.narrowestUnsigned(maxValue); + // Then + assertThat(result).isSameAs(expected); + } } diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/BitpackedEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/BitpackedEncodingEncoder.java index a04bdb45..60e02a11 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/BitpackedEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/BitpackedEncodingEncoder.java @@ -118,7 +118,7 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { } int numPatches = patchIdx.size(); - PType idxPtype = chooseIdxPtype(n); + PType idxPtype = PType.narrowestUnsigned(n); MemorySegment idxBuf = buildPatchIdxBuf(patchIdx, idxPtype, ctx.arena()); MemorySegment valBuf = buildPatchValBuf(patchVal, ptype, ctx.arena()); @@ -157,16 +157,6 @@ private static int bestBitWidth(int[] bitWidthFreq, int bytesPerException, int n return bestWidth; } - private static PType chooseIdxPtype(int n) { - if (n <= 0xFF) { - return PType.U8; - } - if (n <= 0xFFFF) { - return PType.U16; - } - return PType.U32; - } - private static MemorySegment buildPatchIdxBuf(List idx, PType idxPtype, Arena arena) { int numPatches = idx.size(); int elemBytes = idxPtype.byteSize(); diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FsstEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FsstEncodingEncoder.java index f0a2ddfd..687d84c7 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FsstEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FsstEncodingEncoder.java @@ -118,8 +118,8 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { // typically far below the 4-byte ceiling this always used to pay (e.g. a 6-byte string // column needs only U8 lengths, not I32), and the wire format carries the chosen ptype // per FSSTMetadata specifically so a reader never has to guess. - PType uncompLenPType = narrowestUnsignedPType(maxUncompLen); - PType codesOffPType = narrowestUnsignedPType(totalCompressed); + PType uncompLenPType = PType.narrowestUnsigned(maxUncompLen); + PType codesOffPType = PType.narrowestUnsigned(totalCompressed); MemorySegment uncompLenBuf = arena.allocate(Math.max((long) n * uncompLenPType.byteSize(), 1)); for (int i = 0; i < n; i++) { @@ -152,24 +152,10 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { null, null); } - /// Narrowest unsigned `PType` that can hold every value up to `maxValue`. - /// - /// @param maxValue the largest value the buffer must represent, `>= 0` - /// @return `U8`, `U16`, or `U32`, whichever is smallest and still fits - private static PType narrowestUnsignedPType(int maxValue) { - if (maxValue <= 0xFF) { - return PType.U8; - } - if (maxValue <= 0xFFFF) { - return PType.U16; - } - return PType.U32; - } - /// Writes `value` into `seg` at row `idx`, using `ptype`'s byte width. /// /// @param seg destination segment - /// @param ptype `U8`, `U16`, or `U32` (whatever [#narrowestUnsignedPType(int)] returned) + /// @param ptype `U8`, `U16`, or `U32` (whatever [PType#narrowestUnsigned(long)] returned) /// @param idx row index (not a byte offset) /// @param value the value to write private static void writeUnsigned(MemorySegment seg, PType ptype, long idx, long value) { diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/SparseEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/SparseEncodingEncoder.java index 524301d4..7cf3d14b 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/SparseEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/SparseEncodingEncoder.java @@ -75,7 +75,7 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { return CascadeStep.notApplicable(); } int elemBytes = ptype.byteSize(); - PType idxPtype = chooseIdxPtype(n); + PType idxPtype = PType.narrowestUnsigned(n); int idxBytes = idxPtype.byteSize(); int patchCost = idxBytes + elemBytes; int maxPatches = (int) Math.min(Integer.MAX_VALUE, ((long) n * elemBytes / 2L) / patchCost); @@ -227,7 +227,7 @@ static EncodeResult encodeBool(boolean[] validity, EncodeContext ctx) { boolean[] patchVals = new boolean[numPatches]; java.util.Arrays.fill(patchVals, !fillValue); - PType idxPtype = chooseIdxPtype(n); + PType idxPtype = PType.narrowestUnsigned(n); Object idxArr = idxArr(patchIdx, idxPtype); ProtoScalarValue fillScalar = ProtoScalarValue.ofBoolValue(fillValue); @@ -299,7 +299,7 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { } int numPatches = patchIdx.size(); - PType idxPtype = chooseIdxPtype(n); + PType idxPtype = PType.narrowestUnsigned(n); ProtoScalarValue fillScalar = zeroScalar(ptype); byte[] fillBytes = fillScalar.encode(); @@ -353,16 +353,6 @@ private static long readBits(Object data, PType ptype, int i) { }; } - private static PType chooseIdxPtype(int n) { - if (n <= 0xFF) { - return PType.U8; - } else if (n <= 0xFFFF) { - return PType.U16; - } else { - return PType.U32; - } - } - private static ProtoScalarValue zeroScalar(PType ptype) { return switch (ptype) { case I8, I16, I32, I64 -> ProtoScalarValue.ofInt64Value(0L);