7979 * <ul>
8080 * <li><b>int</b> the field number
8181 * <li><b>int</b> the vector encoding ordinal
82- * <li><b>int</b> the query encoding ordinal
8382 * <li><b>int</b> the vector similarity ordinal
8483 * <li><b>vint</b> the vector dimensions
8584 * <li><b>vlong</b> the offset to the vector data in the .veq file
@@ -110,7 +109,6 @@ public class Lucene104ScalarQuantizedVectorsFormat extends FlatVectorsFormat {
110109 new Lucene104ScalarQuantizedVectorScorer (FlatVectorScorerUtil .getLucene99FlatVectorsScorer ());
111110
112111 private final ScalarEncoding encoding ;
113- private final ScalarEncoding queryEncoding ;
114112
115113 /**
116114 * Allowed encodings for scalar quantization.
@@ -132,14 +130,13 @@ public enum ScalarEncoding {
132130 */
133131 SEVEN_BIT (2 , (byte ) 7 , 8 ),
134132 /**
135- * Each dimension is quantized to a single bit and packed into bytes.
133+ * Each dimension is quantized to a single bit and packed into bytes. During query time, the
134+ * query vector is quantized to 4 bits per dimension.
136135 *
137136 * <p>This is the most space efficient encoding, and will produce an index 8x smaller than
138- * {@link #UNSIGNED_BYTE}. However, this comes at the cost of accuracy. This encoding is
139- * recommended for use when the number of dimensions is high (e.g. > 128) and with an
140- * asymmetric quantization scheme where query vectors are quantized to 4 bits.
137+ * {@link #UNSIGNED_BYTE}. However, this comes at the cost of accuracy.
141138 */
142- SINGLE_BIT (3 , (byte ) 1 , 1 );
139+ SINGLE_BIT_QUERY_NIBBLE (3 , (byte ) 1 , 1 , ( byte ) 4 , 4 );
143140
144141 public static ScalarEncoding fromNumBits (int bits ) {
145142 for (ScalarEncoding encoding : values ()) {
@@ -153,13 +150,27 @@ public static ScalarEncoding fromNumBits(int bits) {
153150 /** The number used to identify this encoding on the wire, rather than relying on ordinal. */
154151 private final int wireNumber ;
155152
156- private final byte bits ;
157- private final int bitsPerDim ;
153+ private final byte bits , queryBits ;
154+ private final int bitsPerDim , queryBitsPerDim ;
158155
159156 ScalarEncoding (int wireNumber , byte bits , int bitsPerDim ) {
160157 this .wireNumber = wireNumber ;
161158 this .bits = bits ;
159+ this .queryBits = bits ;
162160 this .bitsPerDim = bitsPerDim ;
161+ this .queryBitsPerDim = bitsPerDim ;
162+ }
163+
164+ ScalarEncoding (int wireNumber , byte bits , int bitsPerDim , byte queryBits , int queryBitsPerDim ) {
165+ this .wireNumber = wireNumber ;
166+ this .bits = bits ;
167+ this .queryBits = queryBits ;
168+ this .bitsPerDim = bitsPerDim ;
169+ this .queryBitsPerDim = queryBitsPerDim ;
170+ }
171+
172+ boolean isAsymmetric () {
173+ return bits != queryBits ;
163174 }
164175
165176 int getWireNumber () {
@@ -171,20 +182,48 @@ public byte getBits() {
171182 return bits ;
172183 }
173184
185+ public byte getQueryBits () {
186+ return queryBits ;
187+ }
188+
174189 /** Return the number of dimensions rounded up to fit into whole bytes. */
175190 public int getDiscreteDimensions (int dimensions ) {
176- int totalBits = dimensions * bitsPerDim ;
177- return (totalBits + 7 ) / 8 * 8 / bitsPerDim ;
191+ if (queryBits == bits ) {
192+ int totalBits = dimensions * bitsPerDim ;
193+ return (totalBits + 7 ) / 8 * 8 / bitsPerDim ;
194+ }
195+ int queryDiscretized = (dimensions * queryBitsPerDim + 7 ) / 8 * 8 / queryBitsPerDim ;
196+ int docDiscretized = (dimensions * bitsPerDim + 7 ) / 8 * 8 / bitsPerDim ;
197+ int maxDiscretized = Math .max (queryDiscretized , docDiscretized );
198+ assert maxDiscretized % (8.0 / queryBitsPerDim ) == 0
199+ : "bad discretized=" + maxDiscretized + " for dim=" + dimensions ;
200+ assert maxDiscretized % (8.0 / bitsPerDim ) == 0
201+ : "bad discretized=" + maxDiscretized + " for dim=" + dimensions ;
202+ return maxDiscretized ;
178203 }
179204
180205 /** Return the number of dimensions that can be packed into a single byte. */
181- public int getBitsPerDim () {
206+ public int getDocBitsPerDim () {
182207 return this .bitsPerDim ;
183208 }
184209
210+ public int getQueryBitsPerDim () {
211+ return this .queryBitsPerDim ;
212+ }
213+
185214 /** Return the number of bytes required to store a packed vector of the given dimensions. */
186- public int getPackedLength (int dimensions ) {
187- return (dimensions * bitsPerDim + 7 ) / 8 ;
215+ public int getDocPackedLength (int dimensions ) {
216+ int discretized = getDiscreteDimensions (dimensions );
217+ // how many bytes do we need to store the quantized vector?
218+ int totalBits = discretized * bitsPerDim ;
219+ return (totalBits + 7 ) / 8 ;
220+ }
221+
222+ public int getQueryPackedLength (int dimensions ) {
223+ int discretized = getDiscreteDimensions (dimensions );
224+ // how many bytes do we need to store the quantized vector?
225+ int totalBits = discretized * queryBitsPerDim ;
226+ return (totalBits + 7 ) / 8 ;
188227 }
189228
190229 /** Returns the encoding for the given wire number, or empty if unknown. */
@@ -203,35 +242,16 @@ public Lucene104ScalarQuantizedVectorsFormat() {
203242 this (ScalarEncoding .UNSIGNED_BYTE );
204243 }
205244
206- /** Creates a new instance with the chosen symmetric quantization encoding. */
245+ /** Creates a new instance with the chosen quantization encoding. */
207246 public Lucene104ScalarQuantizedVectorsFormat (ScalarEncoding encoding ) {
208- this (encoding , encoding );
209- }
210-
211- /** Creates a new instance with the chosen asymmetric quantization encoding. */
212- public Lucene104ScalarQuantizedVectorsFormat (
213- ScalarEncoding encoding , ScalarEncoding queryEncoding ) {
214247 super (NAME );
215248 this .encoding = encoding ;
216- this .queryEncoding = queryEncoding ;
217- // until we have optimized scorers for various other asymmetric encodings, maybe we only allow 1
218- // bit -> 4 bit
219- // Technically, we should be able to do 2 bit -> 4 bit, and 1, 2 -> 8, and 4 -> 8. But these
220- // will take time to
221- // have optimized scorers, and we don't want users to accidentally use poorly optimized
222- // combinations.
223- if (encoding != queryEncoding ) {
224- if (encoding != ScalarEncoding .SINGLE_BIT || queryEncoding != ScalarEncoding .PACKED_NIBBLE ) {
225- throw new IllegalArgumentException (
226- "Only SINGLE_BIT -> PACKED_NIBBLE asymmetric encoding is supported" );
227- }
228- }
229249 }
230250
231251 @ Override
232252 public FlatVectorsWriter fieldsWriter (SegmentWriteState state ) throws IOException {
233253 return new Lucene104ScalarQuantizedVectorsWriter (
234- state , encoding , queryEncoding , rawVectorFormat .fieldsWriter (state ), scorer );
254+ state , encoding , rawVectorFormat .fieldsWriter (state ), scorer );
235255 }
236256
237257 @ Override
@@ -251,8 +271,6 @@ public String toString() {
251271 + NAME
252272 + ", encoding="
253273 + encoding
254- + ", queryEncoding="
255- + queryEncoding
256274 + ", flatVectorScorer="
257275 + scorer
258276 + ", rawVectorFormat="
0 commit comments