22
33import io .github .dfa1 .vortex .core .DType ;
44import io .github .dfa1 .vortex .core .PType ;
5- import io .github .dfa1 .vortex .proto .ALPRDMetadata ;
6- import io .github .dfa1 .vortex .proto .PatchesMetadata ;
75
8- import java .lang .foreign .MemorySegment ;
9- import java .nio .ByteBuffer ;
10- import java .util .ArrayList ;
11- import java .util .HashMap ;
12- import java .util .List ;
13- import java .util .Map ;
14-
15- /// Encoder/decoder for {@code vortex.alprd} — ALP Real Doubles.
16- ///
17- /// <p>Splits each float's bit pattern at a cut point {@code p} (1..16 MSBs):
18- /// left {@code p} bits are dictionary-coded (≤8 entries, codes bitpacked to 1–3 bits);
19- /// right {@code BITS-p} bits are bitpacked directly.
20- /// Values whose left bits fall outside the dictionary are stored as exceptions (patches).
6+ /// Metadata descriptor for {@code vortex.alprd} — ALP Real Doubles.
217///
22- /// <p>Metadata: protobuf {@code ALPRDMetadata} — {@code right_bit_width u32} (tag 1),
23- /// {@code dict_len u32} (tag 2), {@code dict repeated u32} (tag 3),
24- /// {@code left_parts_ptype PType} (tag 4), optional {@code patches PatchesMetadata} (tag 5).
8+ /// <p>Encoder: AlpRdEncodingEncoder.
9+ /// Decoder: AlpRdEncodingDecoder.
2510///
26- /// <p>Children:
27- /// <ul>
28- /// <li>0: left_parts — bitpacked U16 dictionary codes</li>
29- /// <li>1: right_parts — bitpacked U32 (F32) or U64 (F64) right bit-patterns</li>
30- /// <li>2: patch_indices (optional) — bitpacked U64 exception positions</li>
31- /// <li>3: patch_values (optional) — U16 raw left bit-patterns for exceptions</li>
32- /// </ul>
11+ /// <p>Shared dtype constants used by both encoder and decoder are kept here.
3312public final class AlpRdEncoding implements Encoding {
3413
14+ /// Dtype for U16 dictionary-code children.
3515 public static final DType U16_DTYPE = new DType .Primitive (PType .U16 , false );
16+ /// Dtype for U32 right-part children.
3617 public static final DType U32_DTYPE = new DType .Primitive (PType .U32 , false );
18+ /// Dtype for U64 right-part children.
3719 public static final DType U64_DTYPE = new DType .Primitive (PType .U64 , false );
3820
3921 /// Creates a new {@code AlpRdEncoding} instance; use via {@link Registry}.
@@ -52,312 +34,4 @@ public boolean accepts(DType dtype) {
5234 }
5335 return p .ptype () == PType .F32 || p .ptype () == PType .F64 ;
5436 }
55-
56- @ Override
57- public EncodeResult encode (DType dtype , Object data , EncodeContext ctx ) {
58- return Encoder .encode (dtype , data , ctx );
59- }
60-
61- // -------------------------------------------------------------------------
62- // Encoder
63- // -------------------------------------------------------------------------
64-
65- private static final class Encoder {
66-
67- private static final int SAMPLE_SIZE = 512 ;
68- private static final int MAX_CUT = 16 ;
69- private static final int MAX_DICT_SIZE = 8 ;
70-
71- static EncodeResult encode (DType dtype , Object data , EncodeContext ctx ) {
72- PType ptype = ((DType .Primitive ) dtype ).ptype ();
73- return switch (ptype ) {
74- case F64 -> encodeF64 ((double []) data , ctx );
75- case F32 -> encodeF32 ((float []) data , ctx );
76- default -> throw new UnsupportedOperationException ("ALP-RD encode not supported for " + ptype );
77- };
78- }
79-
80- // --- F64 ---
81-
82- private static EncodeResult encodeF64 (double [] values , EncodeContext ctx ) {
83- int n = values .length ;
84- if (n == 0 ) {
85- return emptyResult (U64_DTYPE , ctx );
86- }
87-
88- int sampleLen = Math .min (SAMPLE_SIZE , n );
89- Dictionary64 best = findBestDictionaryF64 (values , sampleLen );
90-
91- Map <Short , Short > lookup = buildLookup (best .dict );
92- long rightMask = -1L >>> (64 - best .rightBitWidth );
93-
94- short [] leftCodes = new short [n ];
95- long [] rightParts = new long [n ];
96- List <Long > excPos = new ArrayList <>();
97- List <Short > excVals = new ArrayList <>();
98-
99- for (int i = 0 ; i < n ; i ++) {
100- long bits = Double .doubleToRawLongBits (values [i ]);
101- short leftU16 = (short ) (bits >>> best .rightBitWidth );
102- rightParts [i ] = bits & rightMask ;
103- Short code = lookup .get (leftU16 );
104- if (code != null ) {
105- leftCodes [i ] = code ;
106- } else {
107- leftCodes [i ] = 0 ;
108- excPos .add ((long ) i );
109- excVals .add (leftU16 );
110- }
111- }
112-
113- return buildEncodeResult (
114- best .dict , best .rightBitWidth , leftCodes , rightParts ,
115- U64_DTYPE , excPos , excVals , ctx );
116- }
117-
118- private static Dictionary64 findBestDictionaryF64 (double [] values , int sampleLen ) {
119- double bestEstSize = Double .MAX_VALUE ;
120- int bestRightBw = 48 ;
121- short [] bestDict = new short []{0 };
122-
123- for (int p = 1 ; p <= MAX_CUT ; p ++) {
124- int rightBw = 64 - p ;
125- Map <Short , Integer > counts = new HashMap <>();
126- for (int i = 0 ; i < sampleLen ; i ++) {
127- long bits = Double .doubleToRawLongBits (values [i ]);
128- short leftU16 = (short ) (bits >>> rightBw );
129- counts .merge (leftU16 , 1 , Integer ::sum );
130- }
131- short [] dict = topKByCount (counts );
132- int excCount = countExceptions (values , sampleLen , dict , rightBw );
133- int maxCode = dict .length - 1 ;
134- int leftBw = maxCode == 0 ? 1 : (Integer .SIZE - Integer .numberOfLeadingZeros (maxCode ));
135- double estSize = rightBw + leftBw + (double ) (excCount * 32 ) / sampleLen ;
136- if (estSize < bestEstSize ) {
137- bestEstSize = estSize ;
138- bestRightBw = rightBw ;
139- bestDict = dict ;
140- }
141- }
142- return new Dictionary64 (bestDict , bestRightBw );
143- }
144-
145- private static int countExceptions (double [] values , int sampleLen , short [] dict , int rightBw ) {
146- Map <Short , Boolean > dictSet = new HashMap <>();
147- for (short d : dict ) {
148- dictSet .put (d , Boolean .TRUE );
149- }
150- int count = 0 ;
151- for (int i = 0 ; i < sampleLen ; i ++) {
152- long bits = Double .doubleToRawLongBits (values [i ]);
153- short leftU16 = (short ) (bits >>> rightBw );
154- if (!dictSet .containsKey (leftU16 )) {
155- count ++;
156- }
157- }
158- return count ;
159- }
160-
161- private static EncodeResult encodeF32 (float [] values , EncodeContext ctx ) {
162- int n = values .length ;
163- if (n == 0 ) {
164- return emptyResult (U32_DTYPE , ctx );
165- }
166-
167- int sampleLen = Math .min (SAMPLE_SIZE , n );
168- Dictionary32 best = findBestDictionaryF32 (values , sampleLen );
169-
170- Map <Short , Short > lookup = buildLookup (best .dict );
171- int rightMask = -1 >>> (32 - best .rightBitWidth );
172-
173- short [] leftCodes = new short [n ];
174- int [] rightParts = new int [n ];
175- List <Long > excPos = new ArrayList <>();
176- List <Short > excVals = new ArrayList <>();
177-
178- for (int i = 0 ; i < n ; i ++) {
179- int bits = Float .floatToRawIntBits (values [i ]);
180- short leftU16 = (short ) (bits >>> best .rightBitWidth );
181- rightParts [i ] = bits & rightMask ;
182- Short code = lookup .get (leftU16 );
183- if (code != null ) {
184- leftCodes [i ] = code ;
185- } else {
186- leftCodes [i ] = 0 ;
187- excPos .add ((long ) i );
188- excVals .add (leftU16 );
189- }
190- }
191-
192- return buildEncodeResult (
193- best .dict , best .rightBitWidth , leftCodes , rightParts ,
194- U32_DTYPE , excPos , excVals , ctx );
195- }
196-
197- // --- F32 ---
198-
199- private static Dictionary32 findBestDictionaryF32 (float [] values , int sampleLen ) {
200- double bestEstSize = Double .MAX_VALUE ;
201- int bestRightBw = 16 ;
202- short [] bestDict = new short []{0 };
203-
204- for (int p = 1 ; p <= MAX_CUT ; p ++) {
205- int rightBw = 32 - p ;
206- Map <Short , Integer > counts = new HashMap <>();
207- for (int i = 0 ; i < sampleLen ; i ++) {
208- int bits = Float .floatToRawIntBits (values [i ]);
209- short leftU16 = (short ) (bits >>> rightBw );
210- counts .merge (leftU16 , 1 , Integer ::sum );
211- }
212- short [] dict = topKByCount (counts );
213- int excCount = countExceptionsF32 (values , sampleLen , dict , rightBw );
214- int maxCode = dict .length - 1 ;
215- int leftBw = maxCode == 0 ? 1 : (Integer .SIZE - Integer .numberOfLeadingZeros (maxCode ));
216- double estSize = rightBw + leftBw + (double ) (excCount * 32 ) / sampleLen ;
217- if (estSize < bestEstSize ) {
218- bestEstSize = estSize ;
219- bestRightBw = rightBw ;
220- bestDict = dict ;
221- }
222- }
223- return new Dictionary32 (bestDict , bestRightBw );
224- }
225-
226- private static int countExceptionsF32 (float [] values , int sampleLen , short [] dict , int rightBw ) {
227- Map <Short , Boolean > dictSet = new HashMap <>();
228- for (short d : dict ) {
229- dictSet .put (d , Boolean .TRUE );
230- }
231- int count = 0 ;
232- for (int i = 0 ; i < sampleLen ; i ++) {
233- int bits = Float .floatToRawIntBits (values [i ]);
234- short leftU16 = (short ) (bits >>> rightBw );
235- if (!dictSet .containsKey (leftU16 )) {
236- count ++;
237- }
238- }
239- return count ;
240- }
241-
242- private static short [] topKByCount (Map <Short , Integer > counts ) {
243- List <Map .Entry <Short , Integer >> sorted = new ArrayList <>(counts .entrySet ());
244- sorted .sort ((a , b ) -> b .getValue () - a .getValue ());
245- int dictSize = Math .min (sorted .size (), Encoder .MAX_DICT_SIZE );
246- short [] dict = new short [dictSize ];
247- for (int i = 0 ; i < dictSize ; i ++) {
248- dict [i ] = sorted .get (i ).getKey ();
249- }
250- return dict ;
251- }
252-
253- private static Map <Short , Short > buildLookup (short [] dict ) {
254- Map <Short , Short > lookup = new HashMap <>();
255- for (short i = 0 ; i < dict .length ; i ++) {
256- lookup .put (dict [i ], i );
257- }
258- return lookup ;
259- }
260-
261- // --- shared helpers ---
262-
263- private static EncodeResult buildEncodeResult (
264- short [] dict , int rightBitWidth ,
265- short [] leftCodes , Object rightPartsData , DType rightDtype ,
266- List <Long > excPos , List <Short > excVals , EncodeContext ctx ) {
267-
268- Encoding bp = ctx .lookupEncoding (EncodingId .FASTLANES_BITPACKED );
269- EncodeResult leftResult = bp .encode (U16_DTYPE , leftCodes , ctx );
270- EncodeResult rightResult = bp .encode (rightDtype , rightPartsData , ctx );
271-
272- List <MemorySegment > allBuffers = new ArrayList <>(leftResult .buffers ());
273- int leftBufCount = allBuffers .size ();
274- allBuffers .addAll (rightResult .buffers ());
275-
276- EncodeNode leftNode = EncodeNode .remapBufferIndices (leftResult .rootNode (), 0 );
277- EncodeNode rightNode = EncodeNode .remapBufferIndices (rightResult .rootNode (), leftBufCount );
278-
279- java .util .List <Integer > dictList = new ArrayList <>(dict .length );
280- for (short d : dict ) {
281- dictList .add (d & 0xFFFF );
282- }
283-
284- EncodeNode [] children ;
285- PatchesMetadata patchesMeta = null ;
286- if (excPos .isEmpty ()) {
287- children = new EncodeNode []{leftNode , rightNode };
288- } else {
289- long [] excPosArr = excPos .stream ().mapToLong (Long ::longValue ).toArray ();
290- short [] excValsArr = new short [excVals .size ()];
291- for (int i = 0 ; i < excVals .size (); i ++) {
292- excValsArr [i ] = excVals .get (i );
293- }
294-
295- EncodeResult idxResult = bp .encode (U64_DTYPE , excPosArr , ctx );
296- EncodeResult valResult = bp .encode (U16_DTYPE , excValsArr , ctx );
297-
298- int idxOffset = allBuffers .size ();
299- allBuffers .addAll (idxResult .buffers ());
300- int idxBufCount = idxResult .buffers ().size ();
301- allBuffers .addAll (valResult .buffers ());
302-
303- EncodeNode idxNode = EncodeNode .remapBufferIndices (idxResult .rootNode (), idxOffset );
304- EncodeNode valNode = EncodeNode .remapBufferIndices (valResult .rootNode (), idxOffset + idxBufCount );
305-
306- patchesMeta = new PatchesMetadata (
307- (long ) excPos .size (),
308- 0L ,
309- io .github .dfa1 .vortex .proto .PType .fromValue (PType .U64 .ordinal ()),
310- null , null , null );
311- children = new EncodeNode []{leftNode , rightNode , idxNode , valNode };
312- }
313-
314- byte [] metaBytes = new ALPRDMetadata (
315- rightBitWidth ,
316- dict .length ,
317- dictList ,
318- io .github .dfa1 .vortex .proto .PType .fromValue (PType .U16 .ordinal ()),
319- patchesMeta
320- ).encode ();
321- EncodeNode root = new EncodeNode (
322- EncodingId .VORTEX_ALPRD , ByteBuffer .wrap (metaBytes ), children , new int []{});
323- return new EncodeResult (root , List .copyOf (allBuffers ), null , null );
324- }
325-
326- private static EncodeResult emptyResult (DType rightDtype , EncodeContext ctx ) {
327- Encoding bp = ctx .lookupEncoding (EncodingId .FASTLANES_BITPACKED );
328- EncodeResult leftResult = bp .encode (AlpRdEncoding .U16_DTYPE , new short [0 ], ctx );
329- EncodeResult rightResult = bp .encode (rightDtype ,
330- rightDtype .equals (U32_DTYPE ) ? new int [0 ] : new long [0 ], ctx );
331-
332- List <MemorySegment > allBuffers = new ArrayList <>(leftResult .buffers ());
333- int leftBufCount = allBuffers .size ();
334- allBuffers .addAll (rightResult .buffers ());
335-
336- EncodeNode leftNode = EncodeNode .remapBufferIndices (leftResult .rootNode (), 0 );
337- EncodeNode rightNode = EncodeNode .remapBufferIndices (rightResult .rootNode (), leftBufCount );
338-
339- byte [] metaBytes = new ALPRDMetadata (
340- 48 ,
341- 0 ,
342- java .util .List .of (),
343- io .github .dfa1 .vortex .proto .PType .fromValue (PType .U16 .ordinal ()),
344- null ).encode ();
345-
346- EncodeNode root = new EncodeNode (
347- EncodingId .VORTEX_ALPRD , ByteBuffer .wrap (metaBytes ),
348- new EncodeNode []{leftNode , rightNode }, new int []{});
349- return new EncodeResult (root , List .copyOf (allBuffers ), null , null );
350- }
351-
352- private record Dictionary64 (short [] dict , int rightBitWidth ) {
353- }
354-
355- private record Dictionary32 (short [] dict , int rightBitWidth ) {
356- }
357- }
358-
359- // -------------------------------------------------------------------------
360- // Decoder
361- // -------------------------------------------------------------------------
362-
36337}
0 commit comments