Enabling use of graph only for binary quantization for gpus. Search is

rchitale7 · rchitale7 · commit 338ec1768729 · 2026-03-02T09:19:21.000Z
expected to fail
diff --git a/jni/external/faiss b/jni/external/faiss
@@ -1 +1 @@
-Subproject commit 5616caad4cea4e8326cdef90aa177bd618af8531
+Subproject commit 9107354294501f067ffd511e999d703c2802714f
diff --git a/jni/external/nmslib b/jni/external/nmslib
@@ -1 +1 @@
-Subproject commit a2d6624e1315402662025debfdd614b505d9c3ef
+Subproject commit 5cf882911d937221e7d7bc30d2b6db534a7ff6d3
diff --git a/jni/src/faiss_wrapper.cpp b/jni/src/faiss_wrapper.cpp
@@ -460,6 +460,16 @@ jlong knn_jni::faiss_wrapper::LoadIndexWithStream(faiss::IOReader* ioReader) {
                         | faiss::IO_FLAG_PQ_SKIP_SDC_TABLE
                         | faiss::IO_FLAG_SKIP_PRECOMPUTE_TABLE);
 
+    // Graph-only indexes have null flat storage and cannot be searched via JNI.
+    auto* idMap = dynamic_cast<faiss::IndexIDMap*>(indexReader);
+    if (idMap != nullptr) {
+        auto* hnsw = dynamic_cast<faiss::IndexHNSW*>(idMap->index);
+        if (hnsw != nullptr && hnsw->storage == nullptr) {
+            delete indexReader;
+            throw std::runtime_error("Cannot load a graph-only index: flat vector storage is null");
+        }
+    }
+
     return (jlong) indexReader;
 }
 jlong knn_jni::faiss_wrapper::LoadIndexWithStreamADCParams(faiss::IOReader* ioReader, knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jobject methodParamsJ) {
diff --git a/remote-index-build-client/src/main/java/org/opensearch/remoteindexbuild/constants/KNNRemoteConstants.java b/remote-index-build-client/src/main/java/org/opensearch/remoteindexbuild/constants/KNNRemoteConstants.java
@@ -50,4 +50,5 @@ public class KNNRemoteConstants {
     public static final String DIMENSION = "dimension";
     public static final String VECTOR_DATA_TYPE_FIELD = "data_type";
     public static final String KNN_ENGINE = "engine";
+    public static final String GRAPH_ONLY = "graph_only";
 }
diff --git a/remote-index-build-client/src/main/java/org/opensearch/remoteindexbuild/model/RemoteBuildRequest.java b/remote-index-build-client/src/main/java/org/opensearch/remoteindexbuild/model/RemoteBuildRequest.java
@@ -22,6 +22,7 @@
 import static org.opensearch.remoteindexbuild.constants.KNNRemoteConstants.TENANT_ID;
 import static org.opensearch.remoteindexbuild.constants.KNNRemoteConstants.VECTOR_DATA_TYPE_FIELD;
 import static org.opensearch.remoteindexbuild.constants.KNNRemoteConstants.VECTOR_PATH;
+import static org.opensearch.remoteindexbuild.constants.KNNRemoteConstants.GRAPH_ONLY;
 
 /**
  * Request object for sending build requests to the remote build service, encapsulating all the required parameters
@@ -40,6 +41,7 @@ public class RemoteBuildRequest implements ToXContentObject {
     protected String vectorDataType;
     protected String engine;
     protected RemoteIndexParameters indexParameters;
+    protected boolean graphOnly;
 
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
@@ -54,6 +56,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         builder.field(VECTOR_DATA_TYPE_FIELD, vectorDataType);
         builder.field(KNN_ENGINE, engine);
         builder.field(INDEX_PARAMETERS, indexParameters);
+        builder.field(GRAPH_ONLY, graphOnly);
         builder.endObject();
         return builder;
     }
diff --git a/remote-index-build-client/src/test/java/org/opensearch/remoteindexbuild/model/RemoteBuildRequestTests.java b/remote-index-build-client/src/test/java/org/opensearch/remoteindexbuild/model/RemoteBuildRequestTests.java
@@ -42,6 +42,7 @@
 import static org.opensearch.remoteindexbuild.constants.KNNRemoteConstants.VECTOR_BLOB_FILE_EXTENSION;
 import static org.opensearch.remoteindexbuild.constants.KNNRemoteConstants.VECTOR_DATA_TYPE_FIELD;
 import static org.opensearch.remoteindexbuild.constants.KNNRemoteConstants.VECTOR_PATH;
+import static org.opensearch.remoteindexbuild.constants.KNNRemoteConstants.GRAPH_ONLY;
 
 public class RemoteBuildRequestTests extends OpenSearchSingleNodeTestCase {
     public static final String MOCK_FULL_PATH = "vectors/1_1_25/SIRKos4rOWlMA62PX2p75m_vectors/SIRKos4rOWlMA62PX2p75m_target_field__3l";
@@ -190,7 +191,10 @@ public String getMockExpectedJson() {
             + METHOD_PARAMETER_M
             + "\":14"
             + "}"
-            + "}"
+            + "},"
+            + "\""
+            + GRAPH_ONLY
+            + "\":false"
             + "}";
     }
 }
diff --git a/src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteIndexBuildStrategy.java b/src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteIndexBuildStrategy.java
@@ -24,7 +24,6 @@
 import org.opensearch.knn.index.remote.RemoteIndexWaiter;
 import org.opensearch.knn.index.remote.RemoteIndexWaiterFactory;
 import org.opensearch.knn.index.vectorvalues.KNNVectorValues;
-import org.opensearch.knn.index.vectorvalues.QuantizedKNNBinaryVectorValues;
 import org.opensearch.remoteindexbuild.client.RemoteIndexClient;
 import org.opensearch.remoteindexbuild.client.RemoteIndexClientFactory;
 import org.opensearch.remoteindexbuild.model.RemoteBuildRequest;
@@ -186,12 +185,14 @@ public void buildAndWriteIndex(BuildIndexParams indexInfo) throws IOException {
     private void writeToRepository(RepositoryContext repositoryContext, BuildIndexParams indexInfo) {
         VectorRepositoryAccessor vectorRepositoryAccessor = repositoryContext.vectorRepositoryAccessor;
         boolean success = false;
+        // For the BQ quantization path, send raw fp32 vectors instead of quantized binary
+        final VectorDataType writeDataType = isGraphOnlyBuild(indexInfo) ? VectorDataType.FLOAT : indexInfo.getVectorDataType();
         metrics.startRepositoryWriteMetrics();
         try {
             vectorRepositoryAccessor.writeToRepository(
                 repositoryContext.blobName,
                 indexInfo.getTotalLiveDocs(),
-                indexInfo.getVectorDataType(),
+                writeDataType,
                 decorateVectorValuesSupplier(indexInfo)
             );
             success = true;
@@ -203,13 +204,19 @@ private void writeToRepository(RepositoryContext repositoryContext, BuildIndexPa
     }
 
     private static Supplier<KNNVectorValues<?>> decorateVectorValuesSupplier(final BuildIndexParams indexInfo) {
-        if (indexInfo.getVectorDataType() == VectorDataType.BINARY && indexInfo.getQuantizationState() != null) {
-            return () -> new QuantizedKNNBinaryVectorValues(indexInfo.getKnnVectorValuesSupplier().get(), indexInfo);
-        }
-
+        // When quantization is present, skip wrapping — send raw fp32 vectors for graph-only build
         return indexInfo.getKnnVectorValuesSupplier();
     }
 
+    /**
+     * Returns true when the build should produce a graph-only index (no flat vector storage).
+     * Currently applies to the binary quantization path where fp32 vectors are quantized to binary.
+     * TODO: Add integration tests for graph-only build once search-side support is implemented.
+     */
+    private static boolean isGraphOnlyBuild(final BuildIndexParams indexInfo) {
+        return indexInfo.getVectorDataType() == VectorDataType.BINARY && indexInfo.getQuantizationState() != null;
+    }
+
     /**
      * Submits a remote build request to the remote index build service
      * @return RemoteBuildResponse containing the response from the remote service
@@ -364,6 +371,15 @@ static RemoteBuildRequest buildRemoteBuildRequest(
 
         final String vectorDataType = determineVectorDataType(indexInfo.getVectorDataType(), parameters);
 
+        // For the binary quantization path (fp32 vectors quantized to binary), skip quantization
+        // and send raw fp32 vectors with graph_only=true. The GPU builds the CAGRA graph from fp32
+        // and returns only the graph structure without flat vector storage.
+        final boolean graphOnly = isGraphOnlyBuild(indexInfo);
+        final String resolvedVectorDataType = graphOnly ? VectorDataType.FLOAT.getValue() : vectorDataType;
+        if (graphOnly) {
+            log.info("Graph-only build: sending fp32 vectors with graph_only=true for field [{}]", indexInfo.getFieldName());
+        }
+
         KNNVectorValues<?> vectorValues = decorateVectorValuesSupplier(indexInfo).get();
         initializeVectorValues(vectorValues);
         assert (vectorValues.dimension() > 0);
@@ -376,9 +392,10 @@ static RemoteBuildRequest buildRemoteBuildRequest(
             .tenantId(indexSettings.getSettings().get(ClusterName.CLUSTER_NAME_SETTING.getKey()))
             .dimension(vectorValues.dimension())
             .docCount(indexInfo.getTotalLiveDocs())
-            .vectorDataType(vectorDataType)
+            .vectorDataType(resolvedVectorDataType)
             .engine(indexInfo.getKnnEngine().getName())
             .indexParameters(indexInfo.getKnnEngine().createRemoteIndexingParameters(parameters))
+            .graphOnly(graphOnly)
             .build();
     }
 }
diff --git a/src/main/java/org/opensearch/knn/index/store/IndexInputWithBuffer.java b/src/main/java/org/opensearch/knn/index/store/IndexInputWithBuffer.java
@@ -45,6 +45,19 @@ private long remainingBytes() {
         return contentLength - indexInput.getFilePointer();
     }
 
+    /**
+     * Reads the first 4 bytes (fourcc) of the index file and resets the read position.
+     * Used to determine the index file type (binary vs non-binary) at load time,
+     * rather than relying on index metadata which may not reflect the actual file contents
+     * (e.g. graph-only builds produce non-binary files for binary-quantized indexes).
+     */
+    public String peekFourcc() throws IOException {
+        byte[] fourcc = new byte[4];
+        indexInput.readBytes(fourcc, 0, 4);
+        indexInput.seek(0);
+        return new String(fourcc);
+    }
+
     @Override
     public String toString() {
         return "{indexInput=" + indexInput + ", len(buffer)=" + buffer.length + "}";
diff --git a/src/main/java/org/opensearch/knn/jni/JNIService.java b/src/main/java/org/opensearch/knn/jni/JNIService.java
@@ -20,6 +20,7 @@
 import org.opensearch.knn.index.store.IndexOutputWithBuffer;
 import org.opensearch.knn.index.util.IndexUtil;
 
+import java.io.IOException;
 import java.util.Locale;
 import java.util.Map;
 
@@ -195,8 +196,16 @@ public static void createIndexFromTemplate(
      */
     public static long loadIndex(IndexInputWithBuffer readStream, Map<String, Object> parameters, KNNEngine knnEngine) {
         if (KNNEngine.FAISS == knnEngine) {
-            if (IndexUtil.isBinaryIndex(knnEngine, parameters)) {
-                return FaissService.loadBinaryIndexWithStream(readStream);
+            // Use the file's fourcc to determine binary vs non-binary loader.
+            // Graph-only builds for binary quantization produce fp32 index files,
+            // so we can't rely on index metadata to make this decision.
+            try {
+                String fourcc = readStream.peekFourcc();
+                if (isBinaryFourcc(fourcc)) {
+                    return FaissService.loadBinaryIndexWithStream(readStream);
+                }
+            } catch (IOException e) {
+                throw new RuntimeException("Failed to peek at index file fourcc", e);
             }
 
             if (IndexUtil.isADCEnabled(knnEngine, parameters)) {
@@ -213,6 +222,10 @@ public static long loadIndex(IndexInputWithBuffer readStream, Map<String, Object
         );
     }
 
+    private static boolean isBinaryFourcc(String fourcc) {
+        return fourcc != null && fourcc.startsWith("IB");
+    }
+
     /**
      * Determine if index contains shared state. Currently, we cannot do this in the plugin because we do not store the
      * model definition anywhere. Only faiss supports indices that have shared state. So for all other engines it will
diff --git a/src/test/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteIndexBuildStrategyTests.java b/src/test/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteIndexBuildStrategyTests.java
@@ -17,7 +17,10 @@
 import org.opensearch.knn.common.exception.TerminalIOException;
 import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.VectorDataType;
+import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams;
+import org.opensearch.knn.index.engine.KNNEngine;
 import org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue;
+import org.opensearch.knn.quantization.models.quantizationState.QuantizationState;
 import org.opensearch.remoteindexbuild.model.RemoteBuildRequest;
 import org.opensearch.repositories.RepositoriesService;
 import org.opensearch.repositories.RepositoryMissingException;
@@ -196,6 +199,32 @@ public void testBuildRequest() throws IOException {
         assertEquals(TEST_CLUSTER, request.getTenantId());
         assertEquals(3, request.getDocCount());
         assertEquals(2, request.getDimension());
+        assertFalse(request.isGraphOnly());
+    }
+
+    public void testBuildRequestGraphOnly() throws IOException {
+        // Build params with BINARY vectorDataType + quantizationState to trigger graph-only
+        BuildIndexParams bqBuildParams = BuildIndexParams.builder()
+            .indexOutputWithBuffer(indexOutputWithBuffer)
+            .knnEngine(KNNEngine.FAISS)
+            .vectorDataType(VectorDataType.BINARY)
+            .parameters(Map.of("index", "param"))
+            .knnVectorValuesSupplier(knnVectorValuesSupplier)
+            .totalLiveDocs((int) knnVectorValues.totalLiveDocs())
+            .segmentWriteState(segmentWriteState)
+            .quantizationState(mock(QuantizationState.class))
+            .isFlush(randomBoolean())
+            .build();
+
+        RemoteBuildRequest request = RemoteIndexBuildStrategy.buildRemoteBuildRequest(
+            createTestIndexSettings(),
+            bqBuildParams,
+            createTestRepositoryMetadata(),
+            MOCK_FULL_PATH,
+            getMockParameterMap()
+        );
+        assertTrue(request.isGraphOnly());
+        assertEquals(VectorDataType.FLOAT.getValue(), request.getVectorDataType());
     }
 
     public Map<String, Object> getMockParameterMap() {
diff --git a/src/test/java/org/opensearch/knn/index/store/IndexInputWithBufferTests.java b/src/test/java/org/opensearch/knn/index/store/IndexInputWithBufferTests.java
@@ -0,0 +1,45 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.store;
+
+import org.apache.lucene.store.ByteBuffersDirectory;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.store.IndexOutput;
+import org.opensearch.knn.KNNTestCase;
+
+import java.io.IOException;
+
+public class IndexInputWithBufferTests extends KNNTestCase {
+
+    public void testPeekFourcc() throws IOException {
+        ByteBuffersDirectory dir = new ByteBuffersDirectory();
+        try (IndexOutput out = dir.createOutput("test.bin", IOContext.DEFAULT)) {
+            out.writeBytes(new byte[] { 'I', 'x', 'M', 'p' }, 4);
+            out.writeBytes(new byte[] { 0, 0, 0, 0 }, 4); // extra bytes
+        }
+        try (IndexInput in = dir.openInput("test.bin", IOContext.DEFAULT)) {
+            IndexInputWithBuffer buffer = new IndexInputWithBuffer(in);
+            assertEquals("IxMp", buffer.peekFourcc());
+            // Verify position is reset to 0
+            assertEquals("IxMp", buffer.peekFourcc());
+        }
+    }
+
+    public void testPeekFourccBinaryIndex() throws IOException {
+        ByteBuffersDirectory dir = new ByteBuffersDirectory();
+        try (IndexOutput out = dir.createOutput("test.bin", IOContext.DEFAULT)) {
+            out.writeBytes(new byte[] { 'I', 'B', 'M', 'p' }, 4);
+            out.writeBytes(new byte[] { 0, 0, 0, 0 }, 4);
+        }
+        try (IndexInput in = dir.openInput("test.bin", IOContext.DEFAULT)) {
+            IndexInputWithBuffer buffer = new IndexInputWithBuffer(in);
+            String fourcc = buffer.peekFourcc();
+            assertEquals("IBMp", fourcc);
+            assertTrue(fourcc.startsWith("IB"));
+        }
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -50,4 +50,5 @@ public class KNNRemoteConstants {`
`50`	`50`	`public static final String DIMENSION = "dimension";`
`51`	`51`	`public static final String VECTOR_DATA_TYPE_FIELD = "data_type";`
`52`	`52`	`public static final String KNN_ENGINE = "engine";`
	`53`	`+ public static final String GRAPH_ONLY = "graph_only";`
`53`	`54`	`}`