Refactoring HNSW to use a new internal FlatVectorFormat (#12729)

Currently the HNSW codec does too many things, it not only indexes vectors, but stores them and determines how to store them given the vector type.

This PR extracts out the vector storage into a new format `Lucene99FlatVectorsFormat` and adds new base class called `FlatVectorsFormat`. This allows for some additional helper functions that allow an indexing codec (like HNSW) take advantage of the flat formats.

Additionally, this PR refactors the new `Lucene99ScalarQuantizedVectorsFormat` to be a `FlatVectorsFormat`.

Now, `Lucene99HnswVectorsFormat` is constructed with a `Lucene99FlatVectorsFormat` and a new `Lucene99HnswScalarQuantizedVectorsFormat` that uses `Lucene99ScalarQuantizedVectorsFormat`

Author: benwtrent

lucene/CHANGES.txt

@@ -184,6 +184,10 @@ New Features
 * GITHUB#12660: HNSW graph now can be merged with multiple thread. Configurable in Lucene99HnswVectorsFormat.
   (Patrick Zhai)
 
+* GITHUB#12729: Add new Lucene99FlatVectorsFormat for writing vectors in a flat format and refactor
+  Lucene99ScalarQuantizedVectorsFormat & Lucene99HnswVectorsFormat to reuse the flat formats.
+  Additionally, this allows flat formats to be pluggable independent of HNSW. (Ben Trent)
+
 Improvements
 ---------------------
 * GITHUB#12523: TaskExecutor waits for all tasks to complete before returning when Exceptions

lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene92/OffHeapFloatVectorValues.java

@@ -80,8 +80,6 @@ static OffHeapFloatVectorValues load(
     }
   }
 
-  abstract Bits getAcceptOrds(Bits acceptDocs);
-
   static class DenseOffHeapVectorValues extends OffHeapFloatVectorValues {
 
     private int doc = -1;
@@ -120,7 +118,7 @@ public RandomAccessVectorValues<float[]> copy() throws IOException {
     }
 
     @Override
-    Bits getAcceptOrds(Bits acceptDocs) {
+    public Bits getAcceptOrds(Bits acceptDocs) {
       return acceptDocs;
     }
   }
@@ -184,7 +182,7 @@ public int ordToDoc(int ord) {
     }
 
     @Override
-    Bits getAcceptOrds(Bits acceptDocs) {
+    public Bits getAcceptOrds(Bits acceptDocs) {
       if (acceptDocs == null) {
         return null;
       }
@@ -256,7 +254,7 @@ public int ordToDoc(int ord) {
     }
 
     @Override
-    Bits getAcceptOrds(Bits acceptDocs) {
+    public Bits getAcceptOrds(Bits acceptDocs) {
       return null;
     }
   }

lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene94/OffHeapByteVectorValues.java

@@ -89,8 +89,6 @@ static OffHeapByteVectorValues load(
     }
   }
 
-  abstract Bits getAcceptOrds(Bits acceptDocs);
-
   static class DenseOffHeapVectorValues extends OffHeapByteVectorValues {
 
     private int doc = -1;
@@ -129,7 +127,7 @@ public RandomAccessVectorValues<byte[]> copy() throws IOException {
     }
 
     @Override
-    Bits getAcceptOrds(Bits acceptDocs) {
+    public Bits getAcceptOrds(Bits acceptDocs) {
       return acceptDocs;
     }
   }
@@ -196,7 +194,7 @@ public int ordToDoc(int ord) {
     }
 
     @Override
-    Bits getAcceptOrds(Bits acceptDocs) {
+    public Bits getAcceptOrds(Bits acceptDocs) {
       if (acceptDocs == null) {
         return null;
       }
@@ -268,7 +266,7 @@ public int ordToDoc(int ord) {
     }
 
     @Override
-    Bits getAcceptOrds(Bits acceptDocs) {
+    public Bits getAcceptOrds(Bits acceptDocs) {
       return null;
     }
   }

lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene94/OffHeapFloatVectorValues.java

@@ -86,8 +86,6 @@ static OffHeapFloatVectorValues load(
     }
   }
 
-  abstract Bits getAcceptOrds(Bits acceptDocs);
-
   static class DenseOffHeapVectorValues extends OffHeapFloatVectorValues {
 
     private int doc = -1;
@@ -126,7 +124,7 @@ public RandomAccessVectorValues<float[]> copy() throws IOException {
     }
 
     @Override
-    Bits getAcceptOrds(Bits acceptDocs) {
+    public Bits getAcceptOrds(Bits acceptDocs) {
       return acceptDocs;
     }
   }
@@ -193,7 +191,7 @@ public int ordToDoc(int ord) {
     }
 
     @Override
-    Bits getAcceptOrds(Bits acceptDocs) {
+    public Bits getAcceptOrds(Bits acceptDocs) {
       if (acceptDocs == null) {
         return null;
       }
@@ -265,7 +263,7 @@ public int ordToDoc(int ord) {
     }
 
     @Override
-    Bits getAcceptOrds(Bits acceptDocs) {
+    public Bits getAcceptOrds(Bits acceptDocs) {
       return null;
     }
   }

lucene/core/src/java/module-info.java

@@ -16,7 +16,6 @@
  */
 
 import org.apache.lucene.codecs.lucene99.Lucene99Codec;
-import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat;
 
 /** Lucene Core. */
 @SuppressWarnings("module") // the test framework is compiled after the core...
@@ -70,7 +69,8 @@
   provides org.apache.lucene.codecs.DocValuesFormat with
       org.apache.lucene.codecs.lucene90.Lucene90DocValuesFormat;
   provides org.apache.lucene.codecs.KnnVectorsFormat with
-      Lucene99HnswVectorsFormat;
+      org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat,
+      org.apache.lucene.codecs.lucene99.Lucene99HnswScalarQuantizedVectorsFormat;
   provides org.apache.lucene.codecs.PostingsFormat with
       org.apache.lucene.codecs.lucene99.Lucene99PostingsFormat;
   provides org.apache.lucene.index.SortFieldProvider with

lucene/core/src/java/org/apache/lucene/codecs/FlatFieldVectorsWriter.java

@@ -0,0 +1,43 @@
+/*
+ * 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.lucene.codecs;
+
+/**
+ * Vectors' writer for a field
+ *
+ * @param <T> an array type; the type of vectors to be written
+ * @lucene.experimental
+ */
+public abstract class FlatFieldVectorsWriter<T> extends KnnFieldVectorsWriter<T> {
+
+  /**
+   * The delegate to write to, can be null When non-null, all vectors seen should be written to the
+   * delegate along with being written to the flat vectors.
+   */
+  protected final KnnFieldVectorsWriter<T> indexingDelegate;
+
+  /**
+   * Sole constructor that expects some indexingDelegate. All vectors seen should be written to the
+   * delegate along with being written to the flat vectors.
+   *
+   * @param indexingDelegate the delegate to write to, can be null
+   */
+  protected FlatFieldVectorsWriter(KnnFieldVectorsWriter<T> indexingDelegate) {
+    this.indexingDelegate = indexingDelegate;
+  }
+}

lucene/core/src/java/org/apache/lucene/codecs/FlatVectorsFormat.java

@@ -0,0 +1,39 @@
+/*
+ * 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.lucene.codecs;
+
+import java.io.IOException;
+import org.apache.lucene.index.SegmentReadState;
+import org.apache.lucene.index.SegmentWriteState;
+
+/**
+ * Encodes/decodes per-document vectors
+ *
+ * @lucene.experimental
+ */
+public abstract class FlatVectorsFormat {
+
+  /** Sole constructor */
+  protected FlatVectorsFormat() {}
+
+  /** Returns a {@link FlatVectorsWriter} to write the vectors to the index. */
+  public abstract FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException;
+
+  /** Returns a {@link KnnVectorsReader} to read the vectors from the index. */
+  public abstract FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException;
+}

lucene/core/src/java/org/apache/lucene/codecs/FlatVectorsReader.java

@@ -0,0 +1,92 @@
+/*
+ * 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.lucene.codecs;
+
+import java.io.Closeable;
+import java.io.IOException;
+import org.apache.lucene.index.ByteVectorValues;
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.hnsw.RandomVectorScorer;
+
+/**
+ * Reads vectors from an index. When searching this reader, it iterates every vector in the index
+ * and scores them
+ *
+ * <p>This class is useful when:
+ *
+ * <ul>
+ *   <li>the number of vectors is small
+ *   <li>when used along side some additional indexing structure that can be used to better search
+ *       the vectors (like HNSW).
+ * </ul>
+ *
+ * @lucene.experimental
+ */
+public abstract class FlatVectorsReader implements Closeable, Accountable {
+
+  /** Sole constructor */
+  protected FlatVectorsReader() {}
+
+  /**
+   * Returns a {@link RandomVectorScorer} for the given field and target vector.
+   *
+   * @param field the field to search
+   * @param target the target vector
+   * @return a {@link RandomVectorScorer} for the given field and target vector.
+   * @throws IOException if an I/O error occurs when reading from the index.
+   */
+  public abstract RandomVectorScorer getRandomVectorScorer(String field, float[] target)
+      throws IOException;
+
+  /**
+   * Returns a {@link RandomVectorScorer} for the given field and target vector.
+   *
+   * @param field the field to search
+   * @param target the target vector
+   * @return a {@link RandomVectorScorer} for the given field and target vector.
+   * @throws IOException if an I/O error occurs when reading from the index.
+   */
+  public abstract RandomVectorScorer getRandomVectorScorer(String field, byte[] target)
+      throws IOException;
+
+  /**
+   * Checks consistency of this reader.
+   *
+   * <p>Note that this may be costly in terms of I/O, e.g. may involve computing a checksum value
+   * against large data files.
+   *
+   * @lucene.internal
+   */
+  public abstract void checkIntegrity() throws IOException;
+
+  /**
+   * Returns the {@link FloatVectorValues} for the given {@code field}. The behavior is undefined if
+   * the given field doesn't have KNN vectors enabled on its {@link FieldInfo}. The return value is
+   * never {@code null}.
+   */
+  public abstract FloatVectorValues getFloatVectorValues(String field) throws IOException;
+
+  /**
+   * Returns the {@link ByteVectorValues} for the given {@code field}. The behavior is undefined if
+   * the given field doesn't have KNN vectors enabled on its {@link FieldInfo}. The return value is
+   * never {@code null}.
+   */
+  public abstract ByteVectorValues getByteVectorValues(String field) throws IOException;
+}

lucene/core/src/java/org/apache/lucene/codecs/FlatVectorsWriter.java

@@ -0,0 +1,74 @@
+/*
+ * 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.lucene.codecs;
+
+import java.io.Closeable;
+import java.io.IOException;
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.MergeState;
+import org.apache.lucene.index.Sorter;
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.IOUtils;
+import org.apache.lucene.util.hnsw.CloseableRandomVectorScorerSupplier;
+
+/**
+ * Vectors' writer for a field that allows additional indexing logic to be implemented by the caller
+ *
+ * @lucene.experimental
+ */
+public abstract class FlatVectorsWriter implements Accountable, Closeable {
+
+  /** Sole constructor */
+  protected FlatVectorsWriter() {}
+
+  /**
+   * Add a new field for indexing, allowing the user to provide a writer that the flat vectors
+   * writer can delegate to if additional indexing logic is required.
+   *
+   * @param fieldInfo fieldInfo of the field to add
+   * @param indexWriter the writer to delegate to, can be null
+   * @return a writer for the field
+   * @throws IOException if an I/O error occurs when adding the field
+   */
+  public abstract FlatFieldVectorsWriter<?> addField(
+      FieldInfo fieldInfo, KnnFieldVectorsWriter<?> indexWriter) throws IOException;
+
+  /**
+   * Write the field for merging, providing a scorer over the newly merged flat vectors. This way
+   * any additional merging logic can be implemented by the user of this class.
+   *
+   * @param fieldInfo fieldInfo of the field to merge
+   * @param mergeState mergeState of the segments to merge
+   * @return a scorer over the newly merged flat vectors, which should be closed as it holds a
+   *     temporary file handle to read over the newly merged vectors
+   * @throws IOException if an I/O error occurs when merging
+   */
+  public abstract CloseableRandomVectorScorerSupplier mergeOneFieldToIndex(
+      FieldInfo fieldInfo, MergeState mergeState) throws IOException;
+
+  /** Write field for merging */
+  public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException {
+    IOUtils.close(mergeOneFieldToIndex(fieldInfo, mergeState));
+  }
+
+  /** Called once at the end before close */
+  public abstract void finish() throws IOException;
+
+  /** Flush all buffered data on disk * */
+  public abstract void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException;
+}

lucene/core/src/java/org/apache/lucene/codecs/lucene95/OffHeapByteVectorValues.java

@@ -94,8 +94,6 @@ public static OffHeapByteVectorValues load(
     }
   }
 
-  public abstract Bits getAcceptOrds(Bits acceptDocs);
-
   /**
    * Dense vector values that are stored off-heap. This is the most common case when every doc has a
    * vector.