Add direct IO to diskBBQ

thecoop · thecoop · commit e45df42eee35 · 2025-09-23T13:24:33.000+01:00
diff --git a/qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java b/qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java
@@ -106,7 +106,11 @@ private static String formatIndexPath(CmdLineArgs args) {
     static Codec createCodec(CmdLineArgs args) {
         final KnnVectorsFormat format;
         if (args.indexType() == IndexType.IVF) {
-            format = new ES920DiskBBQVectorsFormat(args.ivfClusterSize(), ES920DiskBBQVectorsFormat.DEFAULT_CENTROIDS_PER_PARENT_CLUSTER);
+            format = new ES920DiskBBQVectorsFormat(
+                args.ivfClusterSize(),
+                ES920DiskBBQVectorsFormat.DEFAULT_CENTROIDS_PER_PARENT_CLUSTER,
+                false
+            );
         } else {
             if (args.quantizeBits() == 1) {
                 if (args.indexType() == IndexType.FLAT) {
diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsFormat.java
@@ -20,6 +20,7 @@
 import org.apache.lucene.index.SegmentWriteState;
 import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
+import org.elasticsearch.index.codec.vectors.es818.DirectIOLucene99FlatVectorsFormat;
 
 import java.io.IOException;
 import java.util.Collections;
@@ -63,7 +64,15 @@ public class ES920DiskBBQVectorsFormat extends KnnVectorsFormat {
     private static final FlatVectorsFormat rawVectorFormat = new Lucene99FlatVectorsFormat(
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
     );
-    private static final Map<String, FlatVectorsFormat> supportedFormats = Map.of(rawVectorFormat.getName(), rawVectorFormat);
+    private static final FlatVectorsFormat directIORawVectorFormat = new DirectIOLucene99FlatVectorsFormat(
+        FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
+    );
+    private static final Map<String, FlatVectorsFormat> supportedFormats = Map.of(
+        rawVectorFormat.getName(),
+        rawVectorFormat,
+        directIORawVectorFormat.getName(),
+        directIORawVectorFormat
+    );
 
     // This dynamically sets the cluster probe based on the `k` requested and the number of clusters.
     // useful when searching with 'efSearch' type parameters instead of requiring a specific ratio.
@@ -77,8 +86,9 @@ public class ES920DiskBBQVectorsFormat extends KnnVectorsFormat {
 
     private final int vectorPerCluster;
     private final int centroidsPerParentCluster;
+    private final boolean directRawDiskReads;
 
-    public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster) {
+    public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster, boolean directRawDiskReads) {
         super(NAME);
         if (vectorPerCluster < MIN_VECTORS_PER_CLUSTER || vectorPerCluster > MAX_VECTORS_PER_CLUSTER) {
             throw new IllegalArgumentException(
@@ -102,19 +112,21 @@ public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentClu
         }
         this.vectorPerCluster = vectorPerCluster;
         this.centroidsPerParentCluster = centroidsPerParentCluster;
+        this.directRawDiskReads = directRawDiskReads;
     }
 
     /** Constructs a format using the given graph construction parameters and scalar quantization. */
     public ES920DiskBBQVectorsFormat() {
-        this(DEFAULT_VECTORS_PER_CLUSTER, DEFAULT_CENTROIDS_PER_PARENT_CLUSTER);
+        this(DEFAULT_VECTORS_PER_CLUSTER, DEFAULT_CENTROIDS_PER_PARENT_CLUSTER, false);
     }
 
     @Override
     public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
+        FlatVectorsFormat rawFormat = directRawDiskReads ? directIORawVectorFormat : rawVectorFormat;
         return new ES920DiskBBQVectorsWriter(
-            rawVectorFormat.getName(),
+            rawFormat.getName(),
             state,
-            rawVectorFormat.fieldsWriter(state),
+            rawFormat.fieldsWriter(state),
             vectorPerCluster,
             centroidsPerParentCluster
         );
diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIOLucene99FlatVectorsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIOLucene99FlatVectorsFormat.java
@@ -41,11 +41,10 @@
  * Copied from Lucene99FlatVectorsFormat in Lucene 10.1
  *
  * This is copied to change the implementation of {@link #fieldsReader} only.
- * The codec format itself is not changed, so we keep the original {@link #NAME}
  */
 public class DirectIOLucene99FlatVectorsFormat extends AbstractFlatVectorsFormat {
 
-    static final String NAME = "Lucene99FlatVectorsFormat";
+    static final String NAME = "DirectIOLucene99FlatVectorsFormat";
 
     public static final int VERSION_START = 0;
     public static final int VERSION_CURRENT = VERSION_START;
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java
@@ -1544,8 +1544,10 @@ public DenseVectorIndexOptions parseIndexOptions(String fieldName, Map<String, ?
                         );
                     }
                 }
+                Object directRawVectorReadsNode = indexOptionsMap.remove("direct_raw_vector_reads");
+                boolean directRawVectorReads = XContentMapValues.nodeBooleanValue(directRawVectorReadsNode, false);
                 MappingParser.checkNoRemainingFields(fieldName, indexOptionsMap);
-                return new BBQIVFIndexOptions(clusterSize, visitPercentage, rescoreVector);
+                return new BBQIVFIndexOptions(clusterSize, visitPercentage, rescoreVector, directRawVectorReads);
             }
 
             @Override
@@ -2149,17 +2151,23 @@ public boolean validateDimension(int dim, boolean throwOnError) {
     static class BBQIVFIndexOptions extends QuantizedIndexOptions {
         final int clusterSize;
         final double defaultVisitPercentage;
+        final boolean directRawDiskReads;
 
-        BBQIVFIndexOptions(int clusterSize, double defaultVisitPercentage, RescoreVector rescoreVector) {
+        BBQIVFIndexOptions(int clusterSize, double defaultVisitPercentage, RescoreVector rescoreVector, boolean directRawDiskReads) {
             super(VectorIndexType.BBQ_DISK, rescoreVector);
             this.clusterSize = clusterSize;
             this.defaultVisitPercentage = defaultVisitPercentage;
+            this.directRawDiskReads = directRawDiskReads;
         }
 
         @Override
         KnnVectorsFormat getVectorsFormat(ElementType elementType) {
             assert elementType == ElementType.FLOAT;
-            return new ES920DiskBBQVectorsFormat(clusterSize, ES920DiskBBQVectorsFormat.DEFAULT_CENTROIDS_PER_PARENT_CLUSTER);
+            return new ES920DiskBBQVectorsFormat(
+                clusterSize,
+                ES920DiskBBQVectorsFormat.DEFAULT_CENTROIDS_PER_PARENT_CLUSTER,
+                directRawDiskReads
+            );
         }
 
         @Override
diff --git a/server/src/test/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsFormatTests.java b/server/src/test/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsFormatTests.java
@@ -62,13 +62,15 @@ public void setUp() throws Exception {
         if (rarely()) {
             format = new ES920DiskBBQVectorsFormat(
                 random().nextInt(2 * MIN_VECTORS_PER_CLUSTER, ES920DiskBBQVectorsFormat.MAX_VECTORS_PER_CLUSTER),
-                random().nextInt(8, ES920DiskBBQVectorsFormat.MAX_CENTROIDS_PER_PARENT_CLUSTER)
+                random().nextInt(8, ES920DiskBBQVectorsFormat.MAX_CENTROIDS_PER_PARENT_CLUSTER),
+                random().nextBoolean()
             );
         } else {
             // run with low numbers to force many clusters with parents
             format = new ES920DiskBBQVectorsFormat(
                 random().nextInt(MIN_VECTORS_PER_CLUSTER, 2 * MIN_VECTORS_PER_CLUSTER),
-                random().nextInt(MIN_CENTROIDS_PER_PARENT_CLUSTER, 8)
+                random().nextInt(MIN_CENTROIDS_PER_PARENT_CLUSTER, 8),
+                random().nextBoolean()
             );
         }
         super.setUp();
@@ -129,7 +131,7 @@ public void testToString() {
         FilterCodec customCodec = new FilterCodec("foo", Codec.getDefault()) {
             @Override
             public KnnVectorsFormat knnVectorsFormat() {
-                return new ES920DiskBBQVectorsFormat(128, 4);
+                return new ES920DiskBBQVectorsFormat(128, 4, false);
             }
         };
         String expectedPattern = "ES920DiskBBQVectorsFormat(vectorPerCluster=128)";
@@ -140,10 +142,10 @@ public KnnVectorsFormat knnVectorsFormat() {
     }
 
     public void testLimits() {
-        expectThrows(IllegalArgumentException.class, () -> new ES920DiskBBQVectorsFormat(MIN_VECTORS_PER_CLUSTER - 1, 16));
-        expectThrows(IllegalArgumentException.class, () -> new ES920DiskBBQVectorsFormat(MAX_VECTORS_PER_CLUSTER + 1, 16));
-        expectThrows(IllegalArgumentException.class, () -> new ES920DiskBBQVectorsFormat(128, MIN_CENTROIDS_PER_PARENT_CLUSTER - 1));
-        expectThrows(IllegalArgumentException.class, () -> new ES920DiskBBQVectorsFormat(128, MAX_CENTROIDS_PER_PARENT_CLUSTER + 1));
+        expectThrows(IllegalArgumentException.class, () -> new ES920DiskBBQVectorsFormat(MIN_VECTORS_PER_CLUSTER - 1, 16, false));
+        expectThrows(IllegalArgumentException.class, () -> new ES920DiskBBQVectorsFormat(MAX_VECTORS_PER_CLUSTER + 1, 16, false));
+        expectThrows(IllegalArgumentException.class, () -> new ES920DiskBBQVectorsFormat(128, MIN_CENTROIDS_PER_PARENT_CLUSTER - 1, false));
+        expectThrows(IllegalArgumentException.class, () -> new ES920DiskBBQVectorsFormat(128, MAX_CENTROIDS_PER_PARENT_CLUSTER + 1, false));
     }
 
     public void testSimpleOffHeapSize() throws IOException {
diff --git a/server/src/test/java/org/elasticsearch/search/vectors/AbstractDiversifyingChildrenIVFKnnVectorQueryTestCase.java b/server/src/test/java/org/elasticsearch/search/vectors/AbstractDiversifyingChildrenIVFKnnVectorQueryTestCase.java
@@ -98,7 +98,8 @@ public void setUp() throws Exception {
             random().nextInt(
                 ES920DiskBBQVectorsFormat.MIN_CENTROIDS_PER_PARENT_CLUSTER,
                 ES920DiskBBQVectorsFormat.MAX_CENTROIDS_PER_PARENT_CLUSTER
-            )
+            ),
+            random().nextBoolean()
         );
     }
 
diff --git a/server/src/test/java/org/elasticsearch/search/vectors/AbstractIVFKnnVectorQueryTestCase.java b/server/src/test/java/org/elasticsearch/search/vectors/AbstractIVFKnnVectorQueryTestCase.java
@@ -98,7 +98,7 @@ abstract class AbstractIVFKnnVectorQueryTestCase extends LuceneTestCase {
     @Before
     public void setUp() throws Exception {
         super.setUp();
-        format = new ES920DiskBBQVectorsFormat(128, 4);
+        format = new ES920DiskBBQVectorsFormat(128, 4, false);
     }
 
     abstract AbstractIVFKnnVectorQuery getKnnVectorQuery(String field, float[] query, int k, Query queryFilter, float visitRatio);

Original file line number	Diff line number	Diff line change
`@@ -98,7 +98,8 @@ public void setUp() throws Exception {`
`98`	`98`	`random().nextInt(`
`99`	`99`	`ES920DiskBBQVectorsFormat.MIN_CENTROIDS_PER_PARENT_CLUSTER,`
`100`	`100`	`ES920DiskBBQVectorsFormat.MAX_CENTROIDS_PER_PARENT_CLUSTER`
`101`		`- )`
	`101`	`+ ),`
	`102`	`+ random().nextBoolean()`
`102`	`103`	`);`
`103`	`104`	`}`
`104`	`105`
Original file line number	Diff line number	Diff line change
`@@ -98,7 +98,7 @@ abstract class AbstractIVFKnnVectorQueryTestCase extends LuceneTestCase {`
`98`	`98`	`@Before`
`99`	`99`	`public void setUp() throws Exception {`
`100`	`100`	`super.setUp();`
`101`		`- format = new ES920DiskBBQVectorsFormat(128, 4);`
	`101`	`+ format = new ES920DiskBBQVectorsFormat(128, 4, false);`
`102`	`102`	`}`
`103`	`103`
`104`	`104`	`abstract AbstractIVFKnnVectorQuery getKnnVectorQuery(String field, float[] query, int k, Query queryFilter, float visitRatio);`