Add support for determining off-heap memory requirements for KnnVectorsReader

lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene90/Lucene90HnswVectorsReader.java

@@ -306,6 +306,16 @@ private HnswGraph getGraphValues(FieldEntry entry) throws IOException {
     return new OffHeapHnswGraph(entry, bytesSlice);
   }
 
+  @Override
+  public long offHeapBytes() {
+    long bytes = 0L;
+    for (var field : fields.values()) {
+      var f = field.value;
+      bytes += f.vectorDataLength() + f.indexDataLength();
+    }
+    return bytes;
+  }
+
   @Override
   public void close() throws IOException {
     IOUtils.close(vectorData, vectorIndex);

lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene91/Lucene91HnswVectorsReader.java

@@ -297,6 +297,16 @@ private HnswGraph getGraph(FieldEntry entry) throws IOException {
     return new OffHeapHnswGraph(entry, bytesSlice);
   }
 
+  @Override
+  public long offHeapBytes() {
+    long bytes = 0L;
+    for (var field : fields.values()) {
+      var f = field.value;
+      bytes += f.vectorDataLength() + f.vectorIndexLength();
+    }
+    return bytes;
+  }
+
   @Override
   public void close() throws IOException {
     IOUtils.close(vectorData, vectorIndex);

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

@@ -263,6 +263,16 @@ private HnswGraph getGraph(FieldEntry entry) throws IOException {
     return new OffHeapHnswGraph(entry, bytesSlice);
   }
 
+  @Override
+  public long offHeapBytes() {
+    long bytes = 0L;
+    for (var field : fields.values()) {
+      var f = field.value;
+      bytes += f.vectorDataLength() + f.vectorIndexLength();
+    }
+    return bytes;
+  }
+
   @Override
   public void close() throws IOException {
     IOUtils.close(vectorData, vectorIndex);

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

@@ -305,6 +305,16 @@ private HnswGraph getGraph(FieldEntry entry) throws IOException {
     return new OffHeapHnswGraph(entry, bytesSlice);
   }
 
+  @Override
+  public long offHeapBytes() {
+    long bytes = 0L;
+    for (var field : fields.values()) {
+      var f = field.value;
+      bytes += f.vectorDataLength() + f.vectorIndexLength();
+    }
+    return bytes;
+  }
+
   @Override
   public void close() throws IOException {
     IOUtils.close(vectorData, vectorIndex);

lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene95/Lucene95HnswVectorsReader.java

@@ -358,6 +358,16 @@ private HnswGraph getGraph(FieldEntry entry) throws IOException {
     return new OffHeapHnswGraph(entry, vectorIndex);
   }
 
+  @Override
+  public long offHeapBytes() {
+    long bytes = 0L;
+    for (var field : fields.values()) {
+      var f = field.value;
+      bytes += f.vectorDataLength() + f.vectorIndexLength();
+    }
+    return bytes;
+  }
+
   @Override
   public void close() throws IOException {
     IOUtils.close(vectorData, vectorIndex);

lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextKnnVectorsReader.java

@@ -275,6 +275,11 @@ public void checkIntegrity() throws IOException {
     }
   }
 
+  @Override
+  public long offHeapBytes() {
+    return 0L;
+  }
+
   @Override
   public void close() throws IOException {
     dataIn.close();

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

@@ -150,6 +150,11 @@ public void search(
 
             @Override
             public void close() {}
+
+            @Override
+            public long offHeapBytes() {
+              return 0L;
+            }
           };
         }
 

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

@@ -27,9 +27,10 @@
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.TotalHits;
 import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.OffHeapAccountable;
 
 /** Reads vectors from an index. */
-public abstract class KnnVectorsReader implements Closeable {
+public abstract class KnnVectorsReader implements Closeable, OffHeapAccountable {
 
   /** Sole constructor */
   protected KnnVectorsReader() {}

lucene/core/src/java/org/apache/lucene/codecs/lucene102/Lucene102BinaryQuantizedVectorsReader.java

@@ -257,6 +257,15 @@ public long ramBytesUsed() {
     return size;
   }
 
+  @Override
+  public long offHeapBytes() {
+    long bytes = 0L;
+    for (var field : fields.values()) {
+      bytes += field.vectorDataLength();
+    }
+    return bytes;
+  }
+
   public float[] getCentroid(String field) {
     FieldEntry fieldEntry = fields.get(field);
     if (fieldEntry != null) {

lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99FlatVectorsReader.java

@@ -166,6 +166,15 @@ public long ramBytesUsed() {
     return Lucene99FlatVectorsReader.SHALLOW_SIZE + fields.ramBytesUsed();
   }
 
+  @Override
+  public long offHeapBytes() {
+    long bytes = 0L;
+    for (var field : fields.values()) {
+      bytes += field.value.vectorDataLength();
+    }
+    return bytes;
+  }
+
   @Override
   public void checkIntegrity() throws IOException {
     CodecUtil.checksumEntireFile(vectorData);

lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java

@@ -21,6 +21,7 @@
 
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.List;
 import org.apache.lucene.codecs.CodecUtil;
 import org.apache.lucene.codecs.KnnVectorsReader;
@@ -47,6 +48,7 @@
 import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.IOSupplier;
 import org.apache.lucene.util.IOUtils;
+import org.apache.lucene.util.OffHeapAccountable;
 import org.apache.lucene.util.RamUsageEstimator;
 import org.apache.lucene.util.hnsw.HnswGraph;
 import org.apache.lucene.util.hnsw.HnswGraphSearcher;
@@ -248,6 +250,20 @@ public long ramBytesUsed() {
         + flatVectorsReader.ramBytesUsed();
   }
 
+  @Override
+  public long offHeapBytes() {
+    long bytes = 0L;
+    for (var field : fields.values()) {
+      bytes += field.value.vectorIndexLength();
+    }
+    return bytes;
+  }
+
+  @Override
+  public Collection<OffHeapAccountable> getChildOffHeapResources() {
+    return List.of(flatVectorsReader);
+  }
+
   @Override
   public void checkIntegrity() throws IOException {
     flatVectorsReader.checkIntegrity();
@@ -464,7 +480,7 @@ static FieldEntry create(
   }
 
   /** Read the nearest-neighbors graph from the index input */
-  private static final class OffHeapHnswGraph extends HnswGraph {
+  private static final class OffHeapHnswGraph extends HnswGraph implements OffHeapAccountable {
 
     final IndexInput dataIn;
     final int[][] nodesByLevel;
@@ -567,5 +583,10 @@ public NodesIterator getNodesOnLevel(int level) {
         return new ArrayNodesIterator(nodesByLevel[level], nodesByLevel[level].length);
       }
     }
+
+    @Override
+    public long offHeapBytes() {
+      return dataIn.length();
+    }
   }
 }

lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99ScalarQuantizedVectorsReader.java

@@ -21,6 +21,8 @@
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readVectorEncoding;
 
 import java.io.IOException;
+import java.util.Collection;
+import java.util.List;
 import org.apache.lucene.codecs.CodecUtil;
 import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
@@ -41,6 +43,7 @@
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.store.ReadAdvice;
 import org.apache.lucene.util.IOUtils;
+import org.apache.lucene.util.OffHeapAccountable;
 import org.apache.lucene.util.RamUsageEstimator;
 import org.apache.lucene.util.hnsw.RandomVectorScorer;
 import org.apache.lucene.util.quantization.QuantizedByteVectorValues;
@@ -282,6 +285,20 @@ public long ramBytesUsed() {
     return SHALLOW_SIZE + fields.ramBytesUsed() + rawVectorsReader.ramBytesUsed();
   }
 
+  @Override
+  public long offHeapBytes() {
+    long bytes = 0L;
+    for (var field : fields.values()) {
+      bytes += field.value.vectorDataLength();
+    }
+    return bytes;
+  }
+
+  @Override
+  public Collection<OffHeapAccountable> getChildOffHeapResources() {
+    return List.of(rawVectorsReader);
+  }
+
   private FieldEntry readField(IndexInput input, int versionMeta, FieldInfo info)
       throws IOException {
     VectorEncoding vectorEncoding = readVectorEncoding(input);

lucene/core/src/java/org/apache/lucene/codecs/perfield/PerFieldKnnVectorsFormat.java

@@ -343,6 +343,15 @@ public void close() throws IOException {
       }
       IOUtils.close(readers);
     }
+
+    @Override
+    public long offHeapBytes() {
+      long bytes = 0L;
+      for (var field : fields.values()) {
+        bytes += field.value.offHeapBytes();
+      }
+      return bytes;
+    }
   }
 
   static String getSuffix(String formatName, String suffix) {

lucene/core/src/java/org/apache/lucene/index/SlowCodecReaderWrapper.java

@@ -189,6 +189,11 @@ public void checkIntegrity() {
         // We already checkIntegrity the entire reader up front
       }
 
+      @Override
+      public long offHeapBytes() {
+        return 0;
+      }
+
       @Override
       public void close() {}
     };

lucene/core/src/java/org/apache/lucene/index/SlowCompositeCodecReaderWrapper.java

@@ -852,6 +852,11 @@ public FloatVectorValues getFloatVectorValues(String field) throws IOException {
       return new MergedFloatVectorValues(dimension, size, subs);
     }
 
+    @Override
+    public long offHeapBytes() {
+      return Arrays.stream(readers).mapToLong(KnnVectorsReader::offHeapBytes).sum();
+    }
+
     class MergedFloatVectorValues extends FloatVectorValues {
       final int dimension;
       final int size;

lucene/core/src/java/org/apache/lucene/index/SortingCodecReader.java

@@ -590,6 +590,11 @@ public void search(String field, byte[] target, KnnCollector knnCollector, Bits
         throw new UnsupportedOperationException();
       }
 
+      @Override
+      public long offHeapBytes() {
+        return delegate.offHeapBytes();
+      }
+
       @Override
       public void close() throws IOException {
         delegate.close();

lucene/core/src/java/org/apache/lucene/util/OffHeapAccountable.java

@@ -0,0 +1,44 @@
+/*
+ * 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.util;
+
+import java.util.Collection;
+import java.util.Collections;
+
+/**
+ * An object whose off-heap memory requirements can be computed.
+ *
+ * @lucene.internal
+ */
+public interface OffHeapAccountable {
+
+  /**
+   * Returns the size of the off-heap memory requirements in bytes.
+   *
+   * @return a non-negative value indicate the size of the off-heap memory requirements.
+   */
+  long offHeapBytes();
+
+  /**
+   * Returns nested resources of this class. The result should be a point-in-time snapshot (to avoid
+   * race conditions).
+   */
+  default Collection<OffHeapAccountable> getChildOffHeapResources() {
+    return Collections.emptyList();
+  }
+}

lucene/test-framework/src/java/org/apache/lucene/tests/codecs/asserting/AssertingKnnVectorsFormat.java

@@ -218,6 +218,11 @@ public void finishMerge() throws IOException {
       finishMergeCount.incrementAndGet();
     }
 
+    @Override
+    public long offHeapBytes() {
+      return delegate.offHeapBytes();
+    }
+
     @Override
     public void close() throws IOException {
       assert !mergeInstance;