Quantization Framework Implementation with 1bit and MultiBit Binary Q…

…uantizer

Signed-off-by: VIKASH TIWARI <viktari@amazon.com>

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,
+
+    /**
+     * Represents value quantization, typically involving the conversion of continuous
+     * values into discrete ones.
+     */
+    VALUE,
+}

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

@@ -0,0 +1,34 @@
+/*
+ * 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 ScalarQuantizationType {
+    /**
+     * 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,18 @@
+/*
+ * 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.
+     */
+    SCALAR
+}

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,48 @@
+/*
+ * 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.ScalarQuantizationType;
+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, ScalarQuantizationType.ONE_BIT, OneBitScalarQuantizer::new);
+        // Register MultiBitScalarQuantizer for SQParams with VALUE_QUANTIZATION with bit per co-ordinate = 2
+        QuantizerRegistry.register(
+            SQParams.class,
+            QuantizationType.VALUE,
+            ScalarQuantizationType.TWO_BIT,
+            () -> new MultiBitScalarQuantizer(2)
+        );
+        // Register MultiBitScalarQuantizer for SQParams with VALUE_QUANTIZATION with bit per co-ordinate = 4
+        QuantizerRegistry.register(
+            SQParams.class,
+            QuantizationType.VALUE,
+            ScalarQuantizationType.FOUR_BIT,
+            () -> new MultiBitScalarQuantizer(4)
+        );
+    }
+}

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

@@ -0,0 +1,82 @@
+/*
+ * 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.ScalarQuantizationType;
+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 ScalarQuantizationType 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 ScalarQuantizationType 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,20 @@
+/*
+ * 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();
+}

src/main/java/org/opensearch/knn/quantization/models/quantizationParams/QuantizationParams.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.quantization.models.quantizationParams;
+
+import org.opensearch.knn.quantization.enums.QuantizationType;
+
+import java.io.Serializable;
+
+/**
+ * Interface for quantization parameters.
+ * This interface defines the basic contract for all quantization parameter types.
+ * It provides methods to retrieve the quantization type and a unique type identifier.
+ * Implementations of this interface are expected to provide specific configurations
+ * for various quantization strategies.
+ */
+public interface QuantizationParams extends Serializable {
+
+    /**
+     * Gets the quantization type associated with the parameters.
+     * The quantization type defines the overall strategy or method used
+     * for quantization, such as VALUE_QUANTIZATION or SPACE_QUANTIZATION.
+     *
+     * @return the {@link QuantizationType} indicating the quantization method.
+     */
+    QuantizationType getQuantizationType();
+
+    /**
+     * Provides a unique identifier for the quantization parameters.
+     * This identifier is typically a combination of the quantization type
+     * and additional specifics, and it serves to distinguish between different
+     * configurations or modes of quantization.
+     *
+     * @return a string representing the unique type identifier.
+     */
+    String getTypeIdentifier();
+}