Quantization Framework Implementation with 1bit and MultiBit Binary Q…

…uantizer

CHANGELOG.md

@@ -27,4 +27,5 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 * Clean up parsing for query [#1824](https://github.com/opensearch-project/k-NN/pull/1824)
 * Refactor engine package structure [#1913](https://github.com/opensearch-project/k-NN/pull/1913)
 * Refactor method structure and definitions [#1920](https://github.com/opensearch-project/k-NN/pull/1920)
-* Generalize lib interface to return context objects [#1925](https://github.com/opensearch-project/k-NN/pull/1925)
+* Generalize lib interface to return context objects [#1925](https://github.com/opensearch-project/k-NN/pull/1925)
+* Quantization Framework For Disk Optimized Vector Search and Implementation of Binary 1Bit and multibit quantizer[#1889](https://github.com/opensearch-project/k-NN/issues/1889)

src/main/java/org/opensearch/knn/quantization/enums/QuantizationType.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.quantization.enums;
+
+/**
+ * The QuantizationType enum represents the different types of quantization
+ * that can be applied in the KNN.
+ *
+ * <ul>
+ *   <li><b>SPACE_QUANTIZATION:</b> This type of quantization focuses on the space
+ *       or the representation of the data vectors. It is commonly used for techniques
+ *       that reduce the dimensionality or discretize the data space.</li>
+ *   <li><b>VALUE_QUANTIZATION:</b> This type of quantization focuses on the values
+ *       within the data vectors. It involves mapping continuous values into discrete
+ *       values, which can be useful for compressing data or reducing the precision
+ *       of the representation.</li>
+ * </ul>
+ */
+public enum QuantizationType {
+    /**
+     * Represents space quantization, typically involving dimensionality reduction
+     * or space partitioning techniques.
+     */
+    SPACE_QUANTIZATION,
+
+    /**
+     * Represents value quantization, typically involving the conversion of continuous
+     * values into discrete ones.
+     */
+    VALUE_QUANTIZATION,
+}

src/main/java/org/opensearch/knn/quantization/enums/SQTypes.java

@@ -0,0 +1,60 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.quantization.enums;
+
+/**
+ * The SQTypes enum defines the various scalar quantization types that can be used
+ * in the KNN  for vector quantization.
+ * Each type corresponds to a different bit-width representation of the quantized values.
+ */
+public enum SQTypes {
+    /**
+     * FP16 quantization uses 16-bit floating-point representation.
+     */
+    FP16,
+
+    /**
+     * FP8 quantization uses 8-bit floating-point representation.
+     */
+    FP8,
+
+    /**
+     * INT8 quantization uses 8-bit integer representation.
+     */
+    INT8,
+
+    /**
+     * INT6 quantization uses 6-bit integer representation.
+     */
+    INT6,
+
+    /**
+     * INT4 quantization uses 4-bit integer representation.
+     */
+    INT4,
+
+    /**
+     * ONE_BIT quantization uses a single bit per coordinate.
+     */
+    ONE_BIT,
+
+    /**
+     * TWO_BIT quantization uses two bits per coordinate.
+     */
+    TWO_BIT,
+
+    /**
+     * FOUR_BIT quantization uses four bits per coordinate.
+     */
+    FOUR_BIT,
+
+    /**
+     * UNSUPPORTED_TYPE is used to denote quantization types that are not supported.
+     * This can be used as a placeholder or default value.
+     */
+    UNSUPPORTED_TYPE
+}
+

src/main/java/org/opensearch/knn/quantization/enums/ValueQuantizationType.java

@@ -0,0 +1,20 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.quantization.enums;
+
+/**
+ * The ValueQuantizationType enum defines the types of value quantization techniques
+ * that can be applied in the KNN.
+ */
+public enum ValueQuantizationType {
+    /**
+     * SQ (Scalar Quantization) represents a method where each coordinate of the vector is quantized
+     * independently.
+     */
+    SQ
+}
+
+

src/main/java/org/opensearch/knn/quantization/factory/QuantizerFactory.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.quantization.factory;
+
+import org.opensearch.knn.quantization.models.quantizationParams.QuantizationParams;
+import org.opensearch.knn.quantization.quantizer.Quantizer;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * The QuantizerFactory class is responsible for creating instances of {@link Quantizer}
+ * based on the provided {@link QuantizationParams}. It uses a registry to look up the
+ * appropriate quantizer implementation for the given quantization parameters.
+ */
+public final class QuantizerFactory {
+    private static final AtomicBoolean isRegistered = new AtomicBoolean(false);
+
+    // Private constructor to prevent instantiation
+    private QuantizerFactory() {}
+
+    /**
+     * Ensures that default quantizers are registered.
+     */
+    private static void ensureRegistered() {
+        if (!isRegistered.get()) {
+            synchronized (QuantizerFactory.class) {
+                if (!isRegistered.get()) {
+                    QuantizerRegistrar.registerDefaultQuantizers();
+                    isRegistered.set(true);
+                }
+            }
+        }
+    }
+
+    /**
+     * Retrieves a quantizer instance based on the provided quantization parameters.
+     *
+     * @param params the quantization parameters used to determine the appropriate quantizer
+     * @param <P>    the type of quantization parameters, extending {@link QuantizationParams}
+     * @param <Q>    the type of the quantized output
+     * @return an instance of {@link Quantizer} corresponding to the provided parameters
+     */
+    public static <P extends QuantizationParams, Q> Quantizer<P, Q> getQuantizer(final P params) {
+        if (params == null) {
+            throw new IllegalArgumentException("Quantization parameters must not be null.");
+        }
+        // Lazy Registration instead of static block as class level;
+        ensureRegistered();
+        return QuantizerRegistry.getQuantizer(params);
+    }
+}

src/main/java/org/opensearch/knn/quantization/factory/QuantizerRegistrar.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.quantization.factory;
+
+import org.opensearch.knn.quantization.enums.QuantizationType;
+import org.opensearch.knn.quantization.enums.SQTypes;
+import org.opensearch.knn.quantization.models.quantizationParams.SQParams;
+import org.opensearch.knn.quantization.quantizer.MultiBitScalarQuantizer;
+import org.opensearch.knn.quantization.quantizer.OneBitScalarQuantizer;
+
+/**
+ * The QuantizerRegistrar class is responsible for registering default quantizers.
+ * This class ensures that the registration happens only once in a thread-safe manner.
+ */
+final class QuantizerRegistrar {
+
+    // Private constructor to prevent instantiation
+    private QuantizerRegistrar() {
+    }
+
+    /**
+     * Registers default quantizers if not already registered.
+     * <p>
+     * This method is synchronized to ensure that registration occurs only once,
+     * even in a multi-threaded environment.
+     * </p>
+     */
+    public static synchronized void registerDefaultQuantizers() {
+        // Register OneBitScalarQuantizer for SQParams with VALUE_QUANTIZATION and SQTypes.ONE_BIT
+        QuantizerRegistry.register(
+                SQParams.class,
+                QuantizationType.VALUE_QUANTIZATION,
+                SQTypes.ONE_BIT,
+                OneBitScalarQuantizer::new
+        );
+        // Register MultiBitScalarQuantizer for SQParams with VALUE_QUANTIZATION with bit per co-ordinate = 2
+        QuantizerRegistry.register(
+                SQParams.class,
+                QuantizationType.VALUE_QUANTIZATION,
+                SQTypes.TWO_BIT,
+                () -> new MultiBitScalarQuantizer(2)
+        );
+        // Register MultiBitScalarQuantizer for SQParams with VALUE_QUANTIZATION with bit per co-ordinate = 4
+        QuantizerRegistry.register(
+                SQParams.class,
+                QuantizationType.VALUE_QUANTIZATION,
+                SQTypes.FOUR_BIT,
+                () -> new MultiBitScalarQuantizer(4)
+        );
+    }
+}

src/main/java/org/opensearch/knn/quantization/factory/QuantizerRegistry.java

@@ -0,0 +1,79 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.quantization.factory;
+
+import org.opensearch.knn.quantization.enums.QuantizationType;
+import org.opensearch.knn.quantization.enums.SQTypes;
+import org.opensearch.knn.quantization.models.quantizationParams.QuantizationParams;
+import org.opensearch.knn.quantization.quantizer.Quantizer;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Supplier;
+
+/**
+ * The QuantizerRegistry class is responsible for managing the registration and retrieval
+ * of quantizer instances. Quantizers are registered with specific quantization parameters
+ * and type identifiers, allowing for efficient lookup and instantiation.
+ */
+final class QuantizerRegistry {
+
+    // Private constructor to prevent instantiation
+    private QuantizerRegistry() {}
+
+    //ConcurrentHashMap for thread-safe access
+    private static final Map<String, Supplier<? extends Quantizer<?, ?>>> registry = new ConcurrentHashMap<>();
+
+    /**
+     * Registers a quantizer with the registry.
+     *
+     * @param paramClass        the class of the quantization parameters
+     * @param quantizationType  the quantization type (e.g., VALUE_QUANTIZATION)
+     * @param sqType            the specific quantization subtype (e.g., ONE_BIT, TWO_BIT)
+     * @param quantizerSupplier a supplier that provides instances of the quantizer
+     * @param <P>               the type of quantization parameters
+     */
+    public static <P extends QuantizationParams> void register( final Class<P> paramClass,
+                                                                final QuantizationType quantizationType,
+                                                                final SQTypes sqType,
+                                                                final Supplier<? extends Quantizer<?, ?>> quantizerSupplier) {
+        String identifier = createIdentifier(quantizationType, sqType);
+        // Ensure that the quantizer for this identifier is registered only once
+        registry.computeIfAbsent(identifier, key -> quantizerSupplier);
+    }
+
+    /**
+     * Retrieves a quantizer instance based on the provided quantization parameters.
+     *
+     * @param params the quantization parameters used to determine the appropriate quantizer
+     * @param <P>    the type of quantization parameters
+     * @param <Q>    the type of the quantized output
+     * @return an instance of {@link Quantizer} corresponding to the provided parameters
+     * @throws IllegalArgumentException if no quantizer is registered for the given parameters
+     */
+    public static <P extends QuantizationParams, Q> Quantizer<P, Q> getQuantizer(final P params) {
+        String identifier = params.getTypeIdentifier();
+        Supplier<? extends Quantizer<?, ?>> supplier = registry.get(identifier);
+        if (supplier == null) {
+            throw new IllegalArgumentException("No quantizer registered for type identifier: " + identifier +
+                    ". Available quantizers: " + registry.keySet());
+        }
+        @SuppressWarnings("unchecked")
+        Quantizer<P, Q> quantizer = (Quantizer<P, Q>) supplier.get();
+        return quantizer;
+    }
+
+    /**
+     * Creates a unique identifier for the quantizer based on the quantization type and specific quantization subtype.
+     *
+     * @param quantizationType the quantization type
+     * @param sqType           the specific quantization subtype
+     * @return a string identifier
+     */
+    private static String createIdentifier(final QuantizationType quantizationType, final SQTypes sqType) {
+        return quantizationType.name() + "_" + sqType.name();
+    }
+}

src/main/java/org/opensearch/knn/quantization/models/quantizationOutput/BinaryQuantizationOutput.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.quantization.models.quantizationOutput;
+
+/**
+ * The BinaryQuantizationOutput class represents the output of a quantization process in binary format.
+ * It implements the QuantizationOutput interface to handle byte arrays specifically.
+ */
+public class BinaryQuantizationOutput implements QuantizationOutput<byte[]> {
+    private final byte[] quantizedVector;
+
+    /**
+     * Constructs a BinaryQuantizationOutput instance with the specified quantized vector.
+     *
+     * @param quantizedVector the quantized vector represented as a byte array.
+     */
+    public BinaryQuantizationOutput(final byte[] quantizedVector) {
+        if (quantizedVector == null) {
+            throw new IllegalArgumentException("Quantized vector cannot be null");
+        }
+        this.quantizedVector = quantizedVector;
+    }
+
+    @Override
+    public byte[] getQuantizedVector() {
+        return quantizedVector;
+    }
+}

src/main/java/org/opensearch/knn/quantization/models/quantizationOutput/QuantizationOutput.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.quantization.models.quantizationOutput;
+
+/**
+ * The QuantizationOutput interface defines the contract for quantization output data.
+ *
+ * @param <T> The type of the quantized data.
+ */
+public interface QuantizationOutput<T> {
+    /**
+     * Returns the quantized vector.
+     *
+     * @return the quantized data.
+     */
+    T getQuantizedVector();
+}
+
+