Fix dibit striping for quantized vector formats (#15578)

* Fix dibit striping for quantized vector formats

* tidy

Author: benwtrent

lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104ScalarQuantizedVectorsFormat.java

@@ -145,7 +145,28 @@ public enum ScalarEncoding {
      * between the compression of {@link #SINGLE_BIT_QUERY_NIBBLE} and the accuracy of {@link
      * #PACKED_NIBBLE}.
      */
-    DIBIT_QUERY_NIBBLE(4, (byte) 2, 2, (byte) 4, 4);
+    DIBIT_QUERY_NIBBLE(4, (byte) 2, 2, (byte) 4, 4) {
+      @Override
+      public int getDiscreteDimensions(int dimensions) {
+        int queryDiscretized = (dimensions * 4 + 7) / 8 * 8 / 4;
+        // we want to force dibit packing to byte boundaries assuming single bit striping
+        // so we discretize to the same as single bit encoding
+        int docDiscretized = (dimensions + 7) / 8 * 8;
+        int maxDiscretized = Math.max(queryDiscretized, docDiscretized);
+        assert maxDiscretized % (8.0 / 4) == 0
+            : "bad discretized=" + maxDiscretized + " for dim=" + dimensions;
+        assert maxDiscretized % (8.0 / 2) == 0
+            : "bad discretized=" + maxDiscretized + " for dim=" + dimensions;
+        return maxDiscretized;
+      }
+
+      @Override
+      public int getDocPackedLength(int dimensions) {
+        int discretized = getDiscreteDimensions(dimensions);
+        // DIBIT should be stored as two single bits striped
+        return 2 * ((discretized + 7) / 8);
+      }
+    };
 
     public static ScalarEncoding fromNumBits(int bits) {
       for (ScalarEncoding encoding : values()) {

lucene/core/src/test/org/apache/lucene/util/quantization/TestOptimizedScalarQuantizer.java

@@ -194,6 +194,20 @@ public void testUnpackBinary() {
     assertArrayEquals(scratch, unpacked);
   }
 
+  public void testPackTransposeDibit() {
+    int dim = randomIntBetween(1, 4096);
+    ScalarEncoding encoding = ScalarEncoding.DIBIT_QUERY_NIBBLE;
+    byte[] scratch = new byte[encoding.getDiscreteDimensions(dim)];
+    for (int i = 0; i < scratch.length; i++) {
+      scratch[i] = (byte) randomIntBetween(0, 3);
+    }
+    byte[] packed = new byte[encoding.getDocPackedLength(scratch.length)];
+    byte[] unpacked = new byte[scratch.length];
+    OptimizedScalarQuantizer.transposeDibit(scratch, packed);
+    OptimizedScalarQuantizer.untransposeDibit(packed, unpacked);
+    assertArrayEquals(scratch, unpacked);
+  }
+
   static void assertValidQuantizedRange(byte[] quantized, byte bits) {
     for (byte b : quantized) {
       if (bits < 8) {