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

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_V2.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);
             }
         }
@@ -449,6 +453,27 @@ private static boolean containsTDigestFields(TestDataset dataset) throws IOExcep
         return false;
     }
 
+    private static boolean containsBFloat16ElementType(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("element_type");
+                    if (typeNode != null && typeNode.asText().equals("bfloat16")) {
+                        return true;
+                    }
+                }
+            }
+        }
+        return false;
+    }
+
     private static boolean isTimeSeries(TestDataset dataset) throws IOException {
         Settings settings = dataset.readSettingsFile();
         String mode = settings.get("index.mode");
@@ -461,7 +486,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 +496,8 @@ public static void loadDataSetIntoEs(
         boolean inferenceEnabled,
         boolean timeSeriesOnly,
         boolean exponentialHistogramFieldSupported,
-        boolean tDigestFieldSupported
+        boolean tDigestFieldSupported,
+        boolean bFloat16ElementTypeSupported
     ) throws IOException {
         loadDataSetIntoEs(
             client,
@@ -481,6 +507,7 @@ public static void loadDataSetIntoEs(
             timeSeriesOnly,
             exponentialHistogramFieldSupported,
             tDigestFieldSupported,
+            bFloat16ElementTypeSupported,
             (restClient, indexName, indexMapping, indexSettings) -> {
                 ESRestTestCase.createIndex(restClient, indexName, indexSettings, indexMapping, null);
             }
@@ -495,6 +522,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 +535,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/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -1418,6 +1418,11 @@ public enum Cap {
          */
         DENSE_VECTOR_FIELD_TYPE_BIT_ELEMENTS,
 
+        /**
+         * Support directIO rescoring and `bfloat16` for `bbq_hnsw` and `bbq_disk`, and `bfloat16` for `hnsw` ans `bbq_flat` index types.
+         */
+        GENERIC_VECTOR_FORMAT,
+
         /**
          * Support null elements on vector similarity functions
          */