YAML tests for bfloat16 rank_vectors

x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/rank_vectors/rank_vectors_dv_fields_api.yml

@@ -18,6 +18,10 @@ setup:
               vector:
                  type: rank_vectors
                  dims: 5
+              bfloat16_vector:
+                type: rank_vectors
+                dims: 5
+                element_type: bfloat16
               byte_vector:
                  type: rank_vectors
                  dims: 5
@@ -33,6 +37,7 @@ setup:
         body:
           id: "1"
           vector: [[230.0, 300.33, -34.8988, 15.555, -200.0],  [-0.5, 100.0, -13, 14.8, -156.0]]
+          bfloat16_vector: [[230.0, 300.33, -34.8988, 15.555, -200.0],  [-0.5, 100.0, -13, 14.8, -156.0]]
           byte_vector: [[8, 5, -15, 1, -7], [-1, 115, -3, 4, -128]]
           bit_vector: [[8, 5, -15, 1, -7], [-1, 115, -3, 4, -128]]
 
@@ -43,6 +48,7 @@ setup:
         body:
           id: "3"
           vector: [[0.5, 111.3, -13.0, 14.8, -156.0]]
+          bfloat16_vector: [[0.5, 111.3, -13.0, 14.8, -156.0]]
           byte_vector: [[2, 18, -5, 0, -124]]
           bit_vector: [[2, 18, -5, 0, -124]]
 
@@ -73,6 +79,26 @@ setup:
   - match: {hits.hits.1._id: "3"}
   - close_to: {hits.hits.1._score: {value: 192.6447, error: 0.01}}
 
+  - do:
+      headers:
+        Content-Type: application/json
+      search:
+        rest_total_hits_as_int: true
+        body:
+          query:
+            script_score:
+              query: {match_all: {} }
+              script:
+                source: "doc['bfloat16_vector'].magnitudes[0]"
+
+  - match: {hits.total: 2}
+
+  - match: {hits.hits.0._id: "1"}
+  - close_to: {hits.hits.0._score: {value: 429.6021, error: 0.5}}
+
+  - match: {hits.hits.1._id: "3"}
+  - close_to: {hits.hits.1._score: {value: 192.6447, error: 0.5}}
+
   - do:
       headers:
         Content-Type: application/json
@@ -137,6 +163,26 @@ setup:
   - match: {hits.hits.1._id: "3"}
   - close_to: {hits.hits.1._score: {value: 0.5, error: 0.01}}
 
+  - do:
+      headers:
+        Content-Type: application/json
+      search:
+        rest_total_hits_as_int: true
+        body:
+          query:
+            script_score:
+              query: {match_all: {} }
+              script:
+                source: "doc['bfloat16_vector'].vectorValues.next()[0];"
+
+  - match: {hits.total: 2}
+
+  - match: {hits.hits.0._id: "1"}
+  - close_to: {hits.hits.0._score: {value: 230, error: 0.5}}
+
+  - match: {hits.hits.1._id: "3"}
+  - close_to: {hits.hits.1._score: {value: 0.5, error: 0.5}}
+
   - do:
       headers:
         Content-Type: application/json
@@ -186,7 +232,7 @@ setup:
           _source: false
           index: test-index
           body:
-            docvalue_fields: [name, bit_vector, byte_vector, vector]
+            docvalue_fields: [name, bit_vector, byte_vector, bfloat16_vector, vector]
             sort: [{id: "asc"}]
 
     - match: {hits.hits.0._id: "1"}
@@ -202,6 +248,17 @@ setup:
     - close_to: {hits.hits.0.fields.vector.0.1.3: {value: 14.8, error: 0.0001}}
     - close_to: {hits.hits.0.fields.vector.0.1.4: {value: -156, error: 0.0001}}
 
+    - close_to: {hits.hits.0.fields.bfloat16_vector.0.0.0: {value: 230, error: 0.5}}
+    - close_to: {hits.hits.0.fields.bfloat16_vector.0.0.1: {value: 300.33, error: 0.5}}
+    - close_to: {hits.hits.0.fields.bfloat16_vector.0.0.2: {value: -34.8988, error: 0.2}}
+    - close_to: {hits.hits.0.fields.bfloat16_vector.0.0.3: {value: 15.555, error: 0.5}}
+    - close_to: {hits.hits.0.fields.bfloat16_vector.0.0.4: {value: -200, error: 0.5}}
+    - close_to: {hits.hits.0.fields.bfloat16_vector.0.1.0: {value: -0.5, error: 0.1}}
+    - close_to: {hits.hits.0.fields.bfloat16_vector.0.1.1: {value: 100, error: 0.5}}
+    - close_to: {hits.hits.0.fields.bfloat16_vector.0.1.2: {value: -13, error: 0.1}}
+    - close_to: {hits.hits.0.fields.bfloat16_vector.0.1.3: {value: 14.8, error: 0.1}}
+    - close_to: {hits.hits.0.fields.bfloat16_vector.0.1.4: {value: -156, error: 0.5}}
+
     # byte_vector: [[8, 5, -15, 1, -7], [-1, 115, -3, 4, -128]]
     - match: {hits.hits.0.fields.byte_vector.0.0: [8, 5, -15, 1, -7]}
     - match: {hits.hits.0.fields.bit_vector.0.0: [8, 5, -15, 1, -7]}

x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/rank_vectors/rank_vectors_max_sim.yml

@@ -16,6 +16,10 @@ setup:
               vector:
                 type: rank_vectors
                 dims: 5
+              bfloat16_vector:
+                type: rank_vectors
+                dims: 5
+                element_type: bfloat16
               byte_vector:
                 type: rank_vectors
                 dims: 5
@@ -30,6 +34,7 @@ setup:
         id: "1"
         body:
           vector: [[230.0, 300.33, -34.8988, 15.555, -200.0],  [-0.5, 100.0, -13, 14.8, -156.0]]
+          bfloat16_vector: [[230.0, 300.33, -34.8988, 15.555, -200.0],  [-0.5, 100.0, -13, 14.8, -156.0]]
           byte_vector: [[8, 5, -15, 1, -7], [-1, 115, -3, 4, -128]]
           bit_vector: [[8, 5, -15, 1, -7], [-1, 115, -3, 4, -128]]
 
@@ -39,6 +44,7 @@ setup:
         id: "3"
         body:
           vector: [[0.5, 111.3, -13.0, 14.8, -156.0]]
+          bfloat16_vector: [[0.5, 111.3, -13.0, 14.8, -156.0]]
           byte_vector: ["0212fb0084"]
           bit_vector: ["0212fb0084"]
 
@@ -71,6 +77,28 @@ setup:
   - match: {hits.hits.1._id: "3"}
   - close_to: {hits.hits.1._score: {value: 68.90001, error: 0.01}}
 
+  - do:
+      headers:
+        Content-Type: application/json
+      search:
+        rest_total_hits_as_int: true
+        body:
+          query:
+            script_score:
+              query: {match_all: {} }
+              script:
+                source: "maxSimDotProduct(params.query_vector, 'bfloat16_vector')"
+                params:
+                  query_vector: [[1, 2, 1, 1, 1]]
+
+  - match: {hits.total: 2}
+
+  - match: {hits.hits.0._id: "1"}
+  - close_to: {hits.hits.0._score: {value: 611.316, error: 1.0}}
+
+  - match: {hits.hits.1._id: "3"}
+  - close_to: {hits.hits.1._score: {value: 68.90001, error: 0.5}}
+
   - do:
       headers:
         Content-Type: application/json
@@ -156,6 +184,21 @@ setup:
                 source: "maxSimInvHamming(params.query_vector, 'vector')"
                 params:
                   query_vector: [[1, 2, 1, 1, 1]]
+  - do:
+      catch: bad_request
+      headers:
+        Content-Type: application/json
+      search:
+        rest_total_hits_as_int: true
+        body:
+          query:
+            script_score:
+              query: {match_all: {} }
+              script:
+                source: "maxSimInvHamming(params.query_vector, 'bfloat16_vector')"
+                params:
+                  query_vector: [[1, 2, 1, 1, 1]]
+
 
   - do:
       headers: