[ES|QL]: Update Vector Similarity To Support BFLOAT16

docs/changelog/139113.yaml

@@ -0,0 +1,5 @@
+pr: 139113
+summary: "[ES|QL]: Update Vector Similarity To Support BFLOAT16"
+area: "ES|QL"
+type: bug
+issues: []

x-pack/plugin/esql/qa/server/multi-clusters/src/javaRestTest/java/org/elasticsearch/xpack/esql/ccq/MultiClusterSpecIT.java

@@ -414,6 +414,19 @@ protected boolean supportsTDigestField() {
         }
     }
 
+    @Override
+    protected boolean supportsBFloat16ElementType() {
+        try {
+            return RestEsqlTestCase.hasCapabilities(client(), List.of(EsqlCapabilities.Cap.GENERIC_VECTOR_FORMAT.capabilityName()))
+                && RestEsqlTestCase.hasCapabilities(
+                    remoteClusterClient(),
+                    List.of(EsqlCapabilities.Cap.GENERIC_VECTOR_FORMAT.capabilityName())
+                );
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     /**
      * Convert index patterns and subqueries in FROM commands to use remote indices for a given test case.
      */

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/EsqlSpecTestCase.java

@@ -186,7 +186,8 @@ public void setup() {
                 supportsSemanticTextInference(),
                 false,
                 supportsExponentialHistograms(),
-                supportsTDigestField()
+                supportsTDigestField(),
+                supportsBFloat16ElementType()
             );
             return null;
         });
@@ -318,6 +319,10 @@ protected boolean supportsTDigestField() {
         return RestEsqlTestCase.hasCapabilities(client(), List.of(EsqlCapabilities.Cap.TDIGEST_FIELD_TYPE_SUPPORT_V3.capabilityName()));
     }
 
+    protected boolean supportsBFloat16ElementType() {
+        return RestEsqlTestCase.hasCapabilities(client(), List.of(EsqlCapabilities.Cap.GENERIC_VECTOR_FORMAT.capabilityName()));
+    }
+
     protected void doTest() throws Throwable {
         doTest(testCase.query);
     }

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java

@@ -265,7 +265,7 @@ private static List<String> originalTypes(Map<String, ?> x) {
     }
 
     private List<String> availableIndices() throws IOException {
-        return availableDatasetsForEs(true, supportsSourceFieldMapping(), false, requiresTimeSeries(), false, false).stream()
+        return availableDatasetsForEs(true, supportsSourceFieldMapping(), false, requiresTimeSeries(), false, false, false).stream()
             .filter(x -> x.requiresInferenceEndpoint() == false)
             .map(x -> x.indexName())
             .toList();

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvTestsDataLoader.java

@@ -173,6 +173,7 @@ public class CsvTestsDataLoader {
     private static final TestDataset DENSE_VECTOR_TEXT = new TestDataset("dense_vector_text");
     private static final TestDataset MV_TEXT = new TestDataset("mv_text");
     private static final TestDataset DENSE_VECTOR = new TestDataset("dense_vector");
+    private static final TestDataset DENSE_VECTOR_BFLOAT16 = new TestDataset("dense_vector_bfloat16");
     private static final TestDataset COLORS = new TestDataset("colors");
     private static final TestDataset COLORS_CMYK_LOOKUP = new TestDataset("colors_cmyk").withSetting("lookup-settings.json");
     private static final TestDataset BASE_CONVERSION = new TestDataset("base_conversion");
@@ -248,6 +249,7 @@ public class CsvTestsDataLoader {
         Map.entry(DENSE_VECTOR_TEXT.indexName, DENSE_VECTOR_TEXT),
         Map.entry(MV_TEXT.indexName, MV_TEXT),
         Map.entry(DENSE_VECTOR.indexName, DENSE_VECTOR),
+        Map.entry(DENSE_VECTOR_BFLOAT16.indexName, DENSE_VECTOR_BFLOAT16),
         Map.entry(COLORS.indexName, COLORS),
         Map.entry(COLORS_CMYK_LOOKUP.indexName, COLORS_CMYK_LOOKUP),
         Map.entry(BASE_CONVERSION.indexName, BASE_CONVERSION),
@@ -348,7 +350,7 @@ public static void main(String[] args) throws IOException {
         }
 
         try (RestClient client = builder.build()) {
-            loadDataSetIntoEs(client, true, true, false, false, true, true, (restClient, indexName, indexMapping, indexSettings) -> {
+            loadDataSetIntoEs(client, true, true, false, false, true, true, true, (restClient, indexName, indexMapping, indexSettings) -> {
                 // don't use ESRestTestCase methods here or, if you do, test running the main method before making the change
                 StringBuilder jsonBody = new StringBuilder("{");
                 if (indexSettings != null && indexSettings.isEmpty() == false) {
@@ -373,7 +375,8 @@ public static Set<TestDataset> availableDatasetsForEs(
         boolean inferenceEnabled,
         boolean requiresTimeSeries,
         boolean exponentialHistogramFieldSupported,
-        boolean tDigestFieldSupported
+        boolean tDigestFieldSupported,
+        boolean bFloat16ElementTypeSupported
     ) throws IOException {
         Set<TestDataset> testDataSets = new HashSet<>();
 
@@ -383,7 +386,8 @@ public static Set<TestDataset> availableDatasetsForEs(
                 && (supportsSourceFieldMapping || isSourceMappingDataset(dataset) == false)
                 && (requiresTimeSeries == false || isTimeSeries(dataset))
                 && (exponentialHistogramFieldSupported || containsExponentialHistogramFields(dataset) == false)
-                && (tDigestFieldSupported || containsTDigestFields(dataset) == false)) {
+                && (tDigestFieldSupported || containsTDigestFields(dataset) == false)
+                && (bFloat16ElementTypeSupported || containsBFloat16ElementType(dataset) == false)) {
                 testDataSets.add(dataset);
             }
         }
@@ -408,44 +412,33 @@ private static boolean isSourceMappingDataset(TestDataset dataset) throws IOExce
     }
 
     private static boolean containsExponentialHistogramFields(TestDataset dataset) throws IOException {
-        if (dataset.mappingFileName() == null) {
-            return false;
-        }
-        String mappingJsonText = readTextFile(getResource("/" + dataset.mappingFileName()));
-        JsonNode mappingNode = new ObjectMapper().readTree(mappingJsonText);
-        JsonNode properties = mappingNode.get("properties");
-        if (properties != null) {
-            for (var fieldWithValue : properties.properties()) {
-                JsonNode fieldProperties = fieldWithValue.getValue();
-                if (fieldProperties != null) {
-                    JsonNode typeNode = fieldProperties.get("type");
-                    if (typeNode != null && typeNode.asText().equals("exponential_histogram")) {
-                        return true;
-                    }
-                }
-            }
-        }
-        return false;
+        return containsFieldWithProperties(dataset, Map.of("type", "exponential_histogram"));
     }
 
     private static boolean containsTDigestFields(TestDataset dataset) throws IOException {
-        if (dataset.mappingFileName() == null) {
+        return containsFieldWithProperties(dataset, Map.of("type", "tdigest"));
+    }
+
+    private static boolean containsBFloat16ElementType(TestDataset dataset) throws IOException {
+        return containsFieldWithProperties(dataset, Map.of("element_type", "bfloat16"));
+    }
+
+    private static boolean containsFieldWithProperties(TestDataset dataset, Map<String, Object> properties) throws IOException {
+        if (dataset.mappingFileName() == null || properties.isEmpty()) {
             return false;
         }
+
         String mappingJsonText = readTextFile(getResource("/" + dataset.mappingFileName()));
-        JsonNode mappingNode = new ObjectMapper().readTree(mappingJsonText);
-        JsonNode properties = mappingNode.get("properties");
-        if (properties != null) {
-            for (var fieldWithValue : properties.properties()) {
-                JsonNode fieldProperties = fieldWithValue.getValue();
-                if (fieldProperties != null) {
-                    JsonNode typeNode = fieldProperties.get("type");
-                    if (typeNode != null && typeNode.asText().equals("tdigest")) {
-                        return true;
-                    }
+        Map<?, ?> mappingNode = new ObjectMapper().readValue(mappingJsonText, Map.class);
+        Object mappingProperties = mappingNode.get("properties");
+        if (mappingProperties instanceof Map<?, ?> mappingPropertiesMap) {
+            for (Object field : mappingPropertiesMap.values()) {
+                if (field instanceof Map<?, ?> fieldMap && fieldMap.entrySet().containsAll(properties.entrySet())) {
+                    return true;
                 }
             }
         }
+
         return false;
     }
 
@@ -461,7 +454,7 @@ public static void loadDataSetIntoEs(
         boolean supportsSourceFieldMapping,
         boolean inferenceEnabled
     ) throws IOException {
-        loadDataSetIntoEs(client, supportsIndexModeLookup, supportsSourceFieldMapping, inferenceEnabled, false, false, false);
+        loadDataSetIntoEs(client, supportsIndexModeLookup, supportsSourceFieldMapping, inferenceEnabled, false, false, false, false);
     }
 
     public static void loadDataSetIntoEs(
@@ -471,7 +464,8 @@ public static void loadDataSetIntoEs(
         boolean inferenceEnabled,
         boolean timeSeriesOnly,
         boolean exponentialHistogramFieldSupported,
-        boolean tDigestFieldSupported
+        boolean tDigestFieldSupported,
+        boolean bFloat16ElementTypeSupported
     ) throws IOException {
         loadDataSetIntoEs(
             client,
@@ -481,6 +475,7 @@ public static void loadDataSetIntoEs(
             timeSeriesOnly,
             exponentialHistogramFieldSupported,
             tDigestFieldSupported,
+            bFloat16ElementTypeSupported,
             (restClient, indexName, indexMapping, indexSettings) -> {
                 ESRestTestCase.createIndex(restClient, indexName, indexSettings, indexMapping, null);
             }
@@ -495,6 +490,7 @@ private static void loadDataSetIntoEs(
         boolean timeSeriesOnly,
         boolean exponentialHistogramFieldSupported,
         boolean tDigestFieldSupported,
+        boolean bFloat16ElementTypeSupported,
         IndexCreator indexCreator
     ) throws IOException {
         Logger logger = LogManager.getLogger(CsvTestsDataLoader.class);
@@ -507,7 +503,8 @@ private static void loadDataSetIntoEs(
             inferenceEnabled,
             timeSeriesOnly,
             exponentialHistogramFieldSupported,
-            tDigestFieldSupported
+            tDigestFieldSupported,
+            bFloat16ElementTypeSupported
         )) {
             load(client, dataset, logger, indexCreator);
             loadedDatasets.add(dataset.indexName);

x-pack/plugin/esql/qa/testFixtures/src/main/resources/data/dense_vector.csv

@@ -1,5 +1,5 @@
 id:l, float_vector:dense_vector, byte_vector:dense_vector, bit_vector:dense_vector
-0,    [1.0, 2.0, 3.0],          [10, 20, 30],          [13, 112]
-1,    [4.0, 5.0, 6.0],          [40, 50, 60],          [45, 9]
-2,    [9.0, 8.0, 7.0],          [90, 80, 70],          [127, 0]
-3,    [0.054, 0.032, 0.012],    [100, 110, 120],       [88, 53]
+0,    [1.0, 2.0, 3.0],           [10, 20, 30],            [13, 112]
+1,    [4.0, 5.0, 6.0],           [40, 50, 60],            [45, 9]
+2,    [9.0, 8.0, 7.0],           [90, 80, 70],            [127, 0]
+3,    [0.054, 0.032, 0.012],     [100, 110, 120],         [88, 53]

x-pack/plugin/esql/qa/testFixtures/src/main/resources/data/dense_vector_bfloat16.csv

@@ -0,0 +1,5 @@
+id:l, bfloat16_vector:dense_vector
+0,    [1.0, 2.0, 3.0]
+1,    [4.0, 5.0, 6.0]
+2,    [9.0, 8.0, 7.0]
+3,    [0.5390625, 0.3203125, 0.01202392578125]

x-pack/plugin/esql/qa/testFixtures/src/main/resources/dense_vector-bfloat16.csv-spec

@@ -0,0 +1,54 @@
+retrieveDenseBFloat16VectorData
+required_capability: dense_vector_field_type_released
+required_capability: dense_vector_agg_metric_double_if_version
+required_capability: l2_norm_vector_similarity_function
+required_capability: generic_vector_format
+
+FROM dense_vector_bfloat16
+| KEEP id, bfloat16_vector
+| SORT id
+;
+
+id:l | bfloat16_vector:dense_vector
+0    | [1.0, 2.0, 3.0]
+1    | [4.0, 5.0, 6.0]
+2    | [9.0, 8.0, 7.0]
+3    | [0.5390625, 0.3203125, 0.01202392578125]
+;
+
+denseBFloat16VectorWithEval
+required_capability: dense_vector_agg_metric_double_if_version
+required_capability: l2_norm_vector_similarity_function
+required_capability: generic_vector_format
+
+FROM dense_vector_bfloat16
+| EVAL v = bfloat16_vector
+| KEEP id, v
+| SORT id
+;
+
+id:l | v:dense_vector
+0    | [1.0, 2.0, 3.0]
+1    | [4.0, 5.0, 6.0]
+2    | [9.0, 8.0, 7.0]
+3    | [0.5390625, 0.3203125, 0.01202392578125]
+;
+
+denseBFloat16VectorWithRenameAndDrop
+required_capability: dense_vector_agg_metric_double_if_version
+required_capability: l2_norm_vector_similarity_function
+required_capability: generic_vector_format
+
+FROM dense_vector_bfloat16
+| EVAL v = bfloat16_vector 
+| RENAME v AS new_vector 
+| DROP bfloat16_vector
+| SORT id
+;
+
+id:l | new_vector:dense_vector
+0    | [1.0, 2.0, 3.0]
+1    | [4.0, 5.0, 6.0]
+2    | [9.0, 8.0, 7.0]
+3    | [0.5390625, 0.3203125, 0.01202392578125]
+;

x-pack/plugin/esql/qa/testFixtures/src/main/resources/mapping-dense_vector-all_element_types.json

@@ -0,0 +1,48 @@
+{
+  "properties": {
+    "id": {
+      "type": "long"
+    },
+    "float_vector": {
+      "type": "dense_vector",
+      "similarity": "l2_norm",
+      "index_options": {
+        "type": "hnsw",
+        "m": 16,
+        "ef_construction": 100
+      }
+    },
+    "byte_vector": {
+      "type": "dense_vector",
+      "similarity": "l2_norm",
+      "element_type": "byte",
+      "index_options": {
+        "type": "hnsw",
+        "m": 16,
+        "ef_construction": 100
+      }
+    },
+    "bit_vector": {
+      "type": "dense_vector",
+      "dims": 16,
+      "similarity": "l2_norm",
+      "element_type": "bit",
+      "index_options": {
+        "type": "hnsw",
+        "m": 16,
+        "ef_construction": 100
+      }
+    },
+    "bfloat16_vector": {
+      "type": "dense_vector",
+      "dims": 16,
+      "similarity": "l2_norm",
+      "element_type": "bfloat16",
+      "index_options": {
+        "type": "hnsw",
+        "m": 16,
+        "ef_construction": 100
+      }
+    }
+  }
+}

x-pack/plugin/esql/qa/testFixtures/src/main/resources/mapping-dense_vector_bfloat16.json

@@ -0,0 +1,17 @@
+{
+  "properties": {
+    "id": {
+      "type": "long"
+    },
+    "bfloat16_vector": {
+      "type": "dense_vector",
+      "similarity": "l2_norm",
+      "element_type": "bfloat16",
+      "index_options": {
+        "type": "hnsw",
+        "m": 16,
+        "ef_construction": 100
+      }
+    }
+  }
+}

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/vector/VectorSimilarityFunctionsIT.java

@@ -35,7 +35,6 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Locale;
-import java.util.Set;
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.hamcrest.CoreMatchers.containsString;
@@ -49,12 +48,12 @@ public class VectorSimilarityFunctionsIT extends AbstractEsqlIntegTestCase {
     public static Iterable<Object[]> parameters() throws Exception {
         List<Object[]> params = new ArrayList<>();
 
-        for (ElementType elementType : Set.of(ElementType.FLOAT, ElementType.BYTE, ElementType.BIT)) {
+        for (ElementType elementType : ElementType.values()) {
             params.add(new Object[] { "v_cosine", CosineSimilarity.SIMILARITY_FUNCTION, elementType });
             params.add(new Object[] { "v_dot_product", DotProduct.SIMILARITY_FUNCTION, elementType });
             params.add(new Object[] { "v_l1_norm", L1Norm.SIMILARITY_FUNCTION, elementType });
             params.add(new Object[] { "v_l2_norm", L2Norm.SIMILARITY_FUNCTION, elementType });
-            if (elementType != ElementType.FLOAT) {
+            if (elementType != ElementType.FLOAT && elementType != ElementType.BFLOAT16) {
                 params.add(new Object[] { "v_hamming", Hamming.EVALUATOR_SIMILARITY_FUNCTION, elementType });
             }
         }
@@ -236,7 +235,7 @@ private Double calculateSimilarity(
             case BYTE, BIT -> {
                 return (double) similarityFunction.calculateSimilarity(asByteArray(randomVector), asByteArray(vector));
             }
-            case FLOAT -> {
+            case FLOAT, BFLOAT16 -> {
                 return (double) similarityFunction.calculateSimilarity(asFloatArray(randomVector), asFloatArray(vector));
             }
             default -> throw new IllegalArgumentException("Unexpected element type: " + elementType);
@@ -335,7 +334,7 @@ private List<Number> randomVector(int numDims, boolean allowNull) {
         List<Number> vector = new ArrayList<>(dimensions);
         for (int j = 0; j < dimensions; j++) {
             switch (elementType) {
-                case FLOAT -> {
+                case FLOAT, BFLOAT16 -> {
                     if (dimensions == 1) {
                         vector.add(randomValueOtherThan(0f, () -> randomFloat()));
                     } else {