[FP16 PR-5] Enable half_float as datatype

src/main/java/org/opensearch/knn/common/KNNValidationUtil.java

@@ -17,6 +17,8 @@
 import org.opensearch.knn.index.VectorDataType;
 
 import static org.opensearch.knn.common.KNNConstants.VECTOR_DATA_TYPE_FIELD;
+import static org.opensearch.knn.common.KNNConstants.FP16_MIN_VALUE;
+import static org.opensearch.knn.common.KNNConstants.FP16_MAX_VALUE;
 
 @NoArgsConstructor(access = AccessLevel.PRIVATE)
 public class KNNValidationUtil {
@@ -68,6 +70,28 @@ public static void validateByteVectorValue(float value, final VectorDataType dat
         }
     }
 
+    /**
+     * Validate the half float vector value and throw exception if it is not a number, not in the finite range,
+     * or outside the representable range of half-precision floating point.
+     *
+     * @param value  half float vector value
+     */
+    public static void validateHalfFloatVectorValue(float value) {
+        validateFloatVectorValue(value); // Check for NaN and Infinity
+
+        if (value < FP16_MIN_VALUE || value > FP16_MAX_VALUE) {
+            throw new IllegalArgumentException(
+                String.format(
+                    Locale.ROOT,
+                    "[%s] field was set as HALF_FLOAT in index mapping. But, KNN vector values are not within in the half_float range [%f, %f]",
+                    VECTOR_DATA_TYPE_FIELD,
+                    FP16_MIN_VALUE,
+                    FP16_MAX_VALUE
+                )
+            );
+        }
+    }
+
     /**
      * Validate if the given vector size matches with the dimension provided in mapping.
      *

src/main/java/org/opensearch/knn/index/KNNVectorDVLeafFieldData.java

@@ -160,7 +160,7 @@ public int docValueCount() {
 
             @Override
             public Object nextValue() throws IOException {
-                if (vectorDataType == VectorDataType.FLOAT) {
+                if (vectorDataType == VectorDataType.FLOAT || vectorDataType == VectorDataType.HALF_FLOAT) {
                     if (isBinary) {
                         // Convert float[] to little-endian byte[]; XContentBuilder will base64-encode it
                         return KNNVectorDocValueFormat.floatToLittleEndianBytes((float[]) vectorValues.getVector());

src/main/java/org/opensearch/knn/index/KNNVectorScriptDocValues.java

@@ -176,7 +176,7 @@ protected T doGetValue() throws IOException {
      * @return An empty KNNVectorScriptDocValues object.
      */
     public static KNNVectorScriptDocValues<?> emptyValues(String fieldName, VectorDataType type) {
-        if (type == VectorDataType.FLOAT) {
+        if (type == VectorDataType.FLOAT || type == VectorDataType.HALF_FLOAT) {
             return new KNNVectorScriptDocValues<float[]>(DocIdSetIterator.empty(), fieldName, type) {
                 @Override
                 protected float[] doGetValue() throws IOException {

src/main/java/org/opensearch/knn/index/SpaceType.java

@@ -248,7 +248,9 @@ public void validateVector(float[] vector) {
      * @param vectorDataType the given vector data type
      */
     public void validateVectorDataType(VectorDataType vectorDataType) {
-        if (VectorDataType.FLOAT != vectorDataType && VectorDataType.BYTE != vectorDataType) {
+        if (VectorDataType.FLOAT != vectorDataType
+            && VectorDataType.BYTE != vectorDataType
+            && VectorDataType.HALF_FLOAT != vectorDataType) {
             throw new IllegalArgumentException(
                 String.format(Locale.ROOT, "Space type [%s] is not supported with [%s] data type", getValue(), vectorDataType.getValue())
             );

src/main/java/org/opensearch/knn/index/VectorDataType.java

@@ -13,6 +13,7 @@
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.util.BytesRef;
 import org.opensearch.knn.index.codec.util.KNNVectorAsCollectionOfFloatsSerializer;
+import org.opensearch.knn.index.codec.util.KNNVectorAsCollectionOfHalfFloatsSerializer;
 import org.opensearch.knn.index.codec.util.KNNVectorSerializer;
 import org.opensearch.knn.index.memory.NativeMemoryAllocation;
 import org.opensearch.knn.jni.JNICommons;
@@ -105,6 +106,30 @@ public void freeNativeMemory(long memoryAddress) {
             JNICommons.freeVectorData(memoryAddress);
         }
 
+    },
+    HALF_FLOAT("half_float") {
+
+        @Override
+        public FieldType createKnnVectorFieldType(int dimension, KNNVectorSimilarityFunction knnVectorSimilarityFunction) {
+            return KnnFloatVectorField.createFieldType(dimension, knnVectorSimilarityFunction.getVectorSimilarityFunction());
+        }
+
+        @Override
+        public float[] getVectorFromBytesRef(BytesRef binaryValue) {
+            final KNNVectorAsCollectionOfHalfFloatsSerializer vectorSerializer = KNNVectorAsCollectionOfHalfFloatsSerializer.INSTANCE;
+            return vectorSerializer.byteToFloatArray(binaryValue);
+        }
+
+        @Override
+        public TrainingDataConsumer getTrainingDataConsumer(NativeMemoryAllocation.TrainingDataAllocation trainingDataAllocation) {
+            return new FloatTrainingDataConsumer(trainingDataAllocation);
+        }
+
+        @Override
+        public void freeNativeMemory(long memoryAddress) {
+            JNICommons.freeVectorData(memoryAddress);
+        }
+
     };
 
     public static final String SUPPORTED_VECTOR_DATA_TYPES = Arrays.stream(VectorDataType.values())

src/main/java/org/opensearch/knn/index/codec/KNN1040BasePerFieldKnnVectorsFormat.java

@@ -11,6 +11,7 @@
 import org.apache.lucene.codecs.KnnVectorsFormat;
 import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat;
 import org.opensearch.index.mapper.MapperService;
+import org.opensearch.knn.index.VectorDataType;
 import org.opensearch.knn.index.codec.nativeindex.NativeIndexBuildStrategyFactory;
 import org.opensearch.knn.index.engine.CodecFormatResolver;
 import org.opensearch.knn.index.engine.KNNEngine;
@@ -84,6 +85,7 @@ public KnnVectorsFormat getKnnVectorsFormatForField(final String field) {
         ).fieldType(field);
 
         final KNNMappingConfig knnMappingConfig = mappedFieldType.getKnnMappingConfig();
+        final VectorDataType vectorDataType = mappedFieldType.getVectorDataType();
         if (knnMappingConfig.getModelId().isPresent()) {
             return nativeFormatResolver.resolve();
         }
@@ -96,10 +98,27 @@ public KnnVectorsFormat getKnnVectorsFormatForField(final String field) {
         final ResolvedIndexSpec resolvedSpec = mappedFieldType.getResolvedSpec();
 
         if (engine == KNNEngine.LUCENE) {
-            return luceneFormatResolver.resolve(field, knnMethodContext, params, defaultMaxConnections, defaultBeamWidth, resolvedSpec);
+            return luceneFormatResolver.resolve(
+                field,
+                knnMethodContext,
+                params,
+                defaultMaxConnections,
+                defaultBeamWidth,
+                resolvedSpec,
+                vectorDataType
+            );
         }
 
-        return nativeFormatResolver.resolve(field, knnMethodContext, params, defaultMaxConnections, defaultBeamWidth, resolvedSpec);
+        // Native engines — pass params so the resolver can detect SQ encoder
+        return nativeFormatResolver.resolve(
+            field,
+            knnMethodContext,
+            params,
+            defaultMaxConnections,
+            defaultBeamWidth,
+            resolvedSpec,
+            vectorDataType
+        );
     }
 
     @Override

src/main/java/org/opensearch/knn/index/codec/KNN1040Codec/KNN1040PerFieldKnnVectorsFormat.java

@@ -15,6 +15,7 @@
 import org.opensearch.index.mapper.MapperService;
 import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.SpaceType;
+import org.opensearch.knn.index.VectorDataType;
 import org.opensearch.knn.index.codec.KNN1040BasePerFieldKnnVectorsFormat;
 import org.opensearch.knn.index.codec.KnnVectorsFormatContext;
 import org.opensearch.knn.index.codec.LuceneVectorsFormatType;
@@ -91,6 +92,11 @@ private static Map<LuceneVectorsFormatType, Function<KnnVectorsFormatContext, Kn
             if (p.getSpaceType() == SpaceType.HAMMING) {
                 return new KNN9120HnswBinaryVectorsFormat(p.getMaxConnections(), p.getBeamWidth(), merge.v1(), merge.v2(), threshold);
             }
+            // TODO: This branches on data type alone. Once x16 (SQ over FP16) lands, half_float will
+            // also need to select a quantized format, so this must additionally gate on compression level.
+            if (ctx.getVectorDataType() == VectorDataType.HALF_FLOAT) {
+                return new KNN1040HnswHalfFloatVectorsFormat(p.getMaxConnections(), p.getBeamWidth(), merge.v1(), merge.v2(), threshold);
+            }
             return new Lucene99HnswVectorsFormat(p.getMaxConnections(), p.getBeamWidth(), merge.v1(), merge.v2(), threshold);
         }, LuceneVectorsFormatType.SCALAR_QUANTIZED, ctx -> {
             final KNNScalarQuantizedVectorsFormatParams p = new KNNScalarQuantizedVectorsFormatParams(
@@ -120,7 +126,14 @@ private static Map<LuceneVectorsFormatType, Function<KnnVectorsFormatContext, Kn
                 merge.v2(),
                 threshold
             );
-        }, LuceneVectorsFormatType.FLAT, ctx -> new KNN1040ScalarQuantizedVectorsFormat(ScalarEncoding.SINGLE_BIT_QUERY_NIBBLE));
+        }, LuceneVectorsFormatType.FLAT, ctx -> {
+            // TODO: This branches on data type alone. Once x16 (SQ over FP16) lands, half_float will
+            // also need to select a quantized format, so this must additionally gate on compression level.
+            if (ctx.getVectorDataType() == VectorDataType.HALF_FLOAT) {
+                return new KNN1040HalfFloatFlatVectorsFormat();
+            }
+            return new KNN1040ScalarQuantizedVectorsFormat(ScalarEncoding.SINGLE_BIT_QUERY_NIBBLE);
+        });
     }
 
     @Override

src/main/java/org/opensearch/knn/index/codec/KnnVectorsFormatContext.java

@@ -6,6 +6,7 @@
 package org.opensearch.knn.index.codec;
 
 import lombok.Value;
+import org.opensearch.knn.index.VectorDataType;
 import org.opensearch.knn.index.engine.KNNMethodContext;
 
 import java.util.Map;
@@ -55,4 +56,9 @@ public class KnnVectorsFormatContext {
      * </ul>
      */
     int approximateThreshold;
+
+    /**
+     * The vector data type for the field (FLOAT, BYTE, BINARY, HALF_FLOAT).
+     */
+    VectorDataType vectorDataType;
 }

src/main/java/org/opensearch/knn/index/codec/LuceneVectorsFormatType.java

@@ -32,7 +32,7 @@ public enum LuceneVectorsFormatType {
     SCALAR_QUANTIZED,
 
     /**
-     * Flat vector format (e.g., SQ flat via Lucene).
+     * Flat vector format (e.g., SQ flat via Lucene, or half-float flat).
      */
     FLAT
 }

src/main/java/org/opensearch/knn/index/codec/transfer/OffHeapVectorTransferFactory.java

@@ -30,6 +30,7 @@ public static <T> OffHeapVectorTransfer<T> getVectorTransfer(
     ) {
         switch (vectorDataType) {
             case FLOAT:
+            case HALF_FLOAT:
                 return (OffHeapVectorTransfer<T>) new OffHeapFloatVectorTransfer(bytesPerVector, totalVectorsToTransfer);
             case BINARY:
                 return (OffHeapVectorTransfer<T>) new OffHeapBinaryVectorTransfer(bytesPerVector, totalVectorsToTransfer);

src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java

@@ -29,6 +29,8 @@
 public class KNNCodecUtil {
     // Floats are 4 bytes in size
     public static final int FLOAT_BYTE_SIZE = 4;
+    // Half floats are 2 bytes in size
+    public static final int HALF_FLOAT_BYTE_SIZE = 2;
 
     /**
      * This method provides a rough estimate of the number of bytes used for storing an array with the given parameters.
@@ -39,12 +41,14 @@ public class KNNCodecUtil {
      */
     public static long calculateArraySize(int numVectors, int vectorLength, VectorDataType vectorDataType) {
         if (vectorDataType == VectorDataType.FLOAT) {
-            return numVectors * vectorLength * FLOAT_BYTE_SIZE;
+            return (long) numVectors * vectorLength * FLOAT_BYTE_SIZE;
+        } else if (vectorDataType == VectorDataType.HALF_FLOAT) {
+            return (long) numVectors * vectorLength * HALF_FLOAT_BYTE_SIZE;
         } else if (vectorDataType == VectorDataType.BINARY || vectorDataType == VectorDataType.BYTE) {
-            return numVectors * vectorLength;
+            return (long) numVectors * vectorLength;
         } else {
             throw new IllegalArgumentException(
-                "Float, binary, and byte are the only supported vector data types for array size calculation."
+                "Float, half_float, binary, and byte are the only supported vector data types for array size calculation."
             );
         }
     }

src/main/java/org/opensearch/knn/index/engine/AbstractKNNMethod.java

@@ -98,6 +98,10 @@ protected PerDimensionValidator doGetPerDimensionValidator(
         if (VectorDataType.BYTE == vectorDataType) {
             return PerDimensionValidator.DEFAULT_BYTE_VALIDATOR;
         }
+
+        if (VectorDataType.HALF_FLOAT == vectorDataType) {
+            return PerDimensionValidator.DEFAULT_HALF_FLOAT_VALIDATOR;
+        }
         return PerDimensionValidator.DEFAULT_FLOAT_VALIDATOR;
     }
 

src/main/java/org/opensearch/knn/index/engine/AbstractMethodResolver.java

@@ -146,16 +146,38 @@ protected ValidationException validateNotTrainingContext(
         return validationException;
     }
 
+    /**
+     * Rejects a compression level that is not in the supported set for the given engine.
+     *
+     * <p>{@code vectorDataType} is named in the error when known, because an engine can support a level for
+     * one data type and not another -- lucene supports x4 for float but not for half_float -- and an
+     * engine-only message would be misleading there. It may be null (see {@link KNNMethodConfigContext#EMPTY}),
+     * in which case the message names only the engine.
+     */
     protected ValidationException validateCompressionSupported(
         CompressionLevel compressionLevel,
         Set<CompressionLevel> supportedCompressionLevels,
         KNNEngine knnEngine,
+        VectorDataType vectorDataType,
         ValidationException validationException
     ) {
         if (CompressionLevel.isConfigured(compressionLevel) && supportedCompressionLevels.contains(compressionLevel) == false) {
             validationException = validationException == null ? new ValidationException() : validationException;
             validationException.addValidationError(
-                String.format(Locale.ROOT, "\"%s\" does not support \"%s\" compression", knnEngine.getName(), compressionLevel.getName())
+                vectorDataType == null
+                    ? String.format(
+                        Locale.ROOT,
+                        "\"%s\" does not support \"%s\" compression",
+                        knnEngine.getName(),
+                        compressionLevel.getName()
+                    )
+                    : String.format(
+                        Locale.ROOT,
+                        "\"%s\" with \"%s\" data type does not support \"%s\" compression",
+                        knnEngine.getName(),
+                        vectorDataType.getValue(),
+                        compressionLevel.getName()
+                    )
             );
         }
         return validationException;

src/main/java/org/opensearch/knn/index/engine/CodecFormatResolver.java

@@ -6,6 +6,7 @@
 package org.opensearch.knn.index.engine;
 
 import org.apache.lucene.codecs.KnnVectorsFormat;
+import org.opensearch.knn.index.VectorDataType;
 
 import java.util.Map;
 
@@ -24,6 +25,7 @@ public interface CodecFormatResolver {
      * @param defaultMaxConnections default max connections for HNSW
      * @param defaultBeamWidth      default beam width for HNSW
      * @param resolvedSpec          the resolved index spec
+     * @param vectorDataType        the vector data type for the field
      * @return the resolved {@link KnnVectorsFormat}
      */
     KnnVectorsFormat resolve(
@@ -32,7 +34,8 @@ KnnVectorsFormat resolve(
         Map<String, Object> params,
         int defaultMaxConnections,
         int defaultBeamWidth,
-        ResolvedIndexSpec resolvedSpec
+        ResolvedIndexSpec resolvedSpec,
+        VectorDataType vectorDataType
     );
 
     KnnVectorsFormat resolve();

src/main/java/org/opensearch/knn/index/engine/faiss/AbstractFaissMethod.java

@@ -57,6 +57,10 @@ protected PerDimensionValidator doGetPerDimensionValidator(
             return PerDimensionValidator.DEFAULT_FLOAT_VALIDATOR;
         }
 
+        if (VectorDataType.HALF_FLOAT == vectorDataType) {
+            return PerDimensionValidator.DEFAULT_HALF_FLOAT_VALIDATOR;
+        }
+
         throw new IllegalStateException("Unsupported vector data type " + vectorDataType);
     }
 
@@ -82,6 +86,10 @@ protected PerDimensionProcessor doGetPerDimensionProcessor(
             return PerDimensionProcessor.NOOP_PROCESSOR;
         }
 
+        if (VectorDataType.HALF_FLOAT == vectorDataType) {
+            return PerDimensionProcessor.NOOP_PROCESSOR;
+        }
+
         throw new IllegalStateException("Unsupported vector data type " + vectorDataType);
     }
 

src/main/java/org/opensearch/knn/index/engine/faiss/FaissCodecFormatResolver.java

@@ -9,6 +9,7 @@
 import org.opensearch.common.Nullable;
 import org.opensearch.index.mapper.MapperService;
 import org.opensearch.knn.index.KNNSettings;
+import org.opensearch.knn.index.VectorDataType;
 import org.opensearch.knn.index.codec.KNN1040Codec.Faiss1040ScalarQuantizedKnnVectorsFormat;
 import org.opensearch.knn.index.codec.KNN990Codec.NativeEngines990KnnVectorsFormat;
 import org.opensearch.knn.index.codec.nativeindex.NativeIndexBuildStrategyFactory;
@@ -51,7 +52,8 @@ public KnnVectorsFormat resolve(
         Map<String, Object> params,
         int defaultMaxConnections,
         int defaultBeamWidth,
-        ResolvedIndexSpec resolvedSpec
+        ResolvedIndexSpec resolvedSpec,
+        VectorDataType vectorDataType
     ) {
         if (resolvedSpec.isFaissSQOneBit()) {
             return new Faiss1040ScalarQuantizedKnnVectorsFormat(

src/main/java/org/opensearch/knn/index/engine/faiss/FaissMethodResolver.java

@@ -169,6 +169,7 @@ private void validateConfig(KNNMethodConfigContext knnMethodConfigContext) {
             compressionLevel,
             SUPPORTED_COMPRESSION_LEVELS,
             KNNEngine.FAISS,
+            knnMethodConfigContext.getVectorDataType(),
             null
         );
         if (validationException != null) {