addressing PR comments - simplifying toXContent for VectorData

Author: pmpailis

server/src/main/java/org/elasticsearch/search/vectors/ExactKnnQueryBuilder.java

@@ -22,7 +22,6 @@
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
-import java.util.Map;
 import java.util.Objects;
 
 /**
@@ -31,7 +30,6 @@
  */
 public class ExactKnnQueryBuilder extends AbstractQueryBuilder<ExactKnnQueryBuilder> {
     public static final String NAME = "exact_knn";
-    private static final MapParams QUERY_VECTOR_PARAMS = new MapParams(Map.of(VectorData.XCONTENT_PARAM_NAME, "query"));
     private final String field;
     private final VectorData query;
 
@@ -92,7 +90,7 @@ protected void doWriteTo(StreamOutput out) throws IOException {
     @Override
     protected void doXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject(NAME);
-        query.toXContent(builder, QUERY_VECTOR_PARAMS);
+        builder.field("query", query);
         builder.field("field", field);
         boostAndQueryNameToXContent(builder);
         builder.endObject();

server/src/main/java/org/elasticsearch/search/vectors/KnnScoreDocQueryBuilder.java

@@ -25,7 +25,6 @@
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Map;
 import java.util.Objects;
 
 /**
@@ -35,7 +34,6 @@
  */
 public class KnnScoreDocQueryBuilder extends AbstractQueryBuilder<KnnScoreDocQueryBuilder> {
     public static final String NAME = "knn_score_doc";
-    private static final MapParams QUERY_VECTOR_PARAMS = new MapParams(Map.of(VectorData.XCONTENT_PARAM_NAME, "query"));
     private final ScoreDoc[] scoreDocs;
     private final String fieldName;
     private final VectorData queryVector;
@@ -129,7 +127,7 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
             builder.field("field", fieldName);
         }
         if (queryVector != null) {
-            queryVector.toXContent(builder, QUERY_VECTOR_PARAMS);
+            builder.field("query", queryVector);
         }
         boostAndQueryNameToXContent(builder);
         builder.endObject();

server/src/main/java/org/elasticsearch/search/vectors/KnnSearchBuilder.java

@@ -28,7 +28,6 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 import java.util.Objects;
 import java.util.function.Supplier;
 
@@ -55,9 +54,6 @@ public class KnnSearchBuilder implements Writeable, ToXContentFragment, Rewritea
     public static final ParseField FILTER_FIELD = new ParseField("filter");
     public static final ParseField BOOST_FIELD = AbstractQueryBuilder.BOOST_FIELD;
     public static final ParseField INNER_HITS_FIELD = new ParseField("inner_hits");
-    private static final MapParams QUERY_VECTOR_PARAMS = new MapParams(
-        Map.of(VectorData.XCONTENT_PARAM_NAME, QUERY_VECTOR_FIELD.getPreferredName())
-    );
 
     @SuppressWarnings("unchecked")
     private static final ConstructingObjectParser<KnnSearchBuilder.Builder, Void> PARSER = new ConstructingObjectParser<>("knn", args -> {
@@ -423,7 +419,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.field(queryVectorBuilder.getWriteableName(), queryVectorBuilder);
             builder.endObject();
         } else {
-            queryVector.toXContent(builder, QUERY_VECTOR_PARAMS);
+            builder.field(QUERY_VECTOR_FIELD.getPreferredName(), queryVector);
         }
         if (similarity != null) {
             builder.field(VECTOR_SIMILARITY.getPreferredName(), similarity);

server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java

@@ -37,7 +37,6 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 import java.util.Objects;
 
 import static org.elasticsearch.search.SearchService.DEFAULT_SIZE;
@@ -59,9 +58,6 @@ public class KnnVectorQueryBuilder extends AbstractQueryBuilder<KnnVectorQueryBu
     public static final ParseField QUERY_VECTOR_FIELD = new ParseField("query_vector");
     public static final ParseField VECTOR_SIMILARITY_FIELD = new ParseField("similarity");
     public static final ParseField FILTER_FIELD = new ParseField("filter");
-    private static final MapParams QUERY_VECTOR_PARAMS = new MapParams(
-        Map.of(VectorData.XCONTENT_PARAM_NAME, QUERY_VECTOR_FIELD.getPreferredName())
-    );
 
     @SuppressWarnings("unchecked")
     public static final ConstructingObjectParser<KnnVectorQueryBuilder, Void> PARSER = new ConstructingObjectParser<>(
@@ -225,7 +221,7 @@ protected void doWriteTo(StreamOutput out) throws IOException {
     protected void doXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject(NAME);
         builder.field(FIELD_FIELD.getPreferredName(), fieldName);
-        queryVector.toXContent(builder, QUERY_VECTOR_PARAMS);
+        builder.field(QUERY_VECTOR_FIELD.getPreferredName(), queryVector);
         if (numCands != null) {
             builder.field(NUM_CANDS_FIELD.getPreferredName(), numCands);
         }

server/src/main/java/org/elasticsearch/search/vectors/VectorData.java

@@ -28,9 +28,6 @@
 
 public record VectorData(float[] floatVector, byte[] byteVector) implements Writeable, ToXContentFragment {
 
-    public static final String XCONTENT_PARAM_NAME = "field";
-    private static final String DEFAULT_XCONTENT_NAME = "vector";
-
     private VectorData(float[] floatVector) {
         this(floatVector, null);
     }
@@ -90,11 +87,17 @@ public void writeTo(StreamOutput out) throws IOException {
 
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startArray();
         if (floatVector != null) {
-            builder.array(params.param(XCONTENT_PARAM_NAME, DEFAULT_XCONTENT_NAME), floatVector);
+            for (float v : floatVector) {
+                builder.value(v);
+            }
         } else {
-            builder.array(params.param(XCONTENT_PARAM_NAME, DEFAULT_XCONTENT_NAME), byteVector);
+            for (byte b : byteVector) {
+                builder.value(b);
+            }
         }
+        builder.endArray();
         return builder;
     }
 
@@ -157,12 +160,10 @@ private static VectorData parseNumberVector(XContentParser parser) throws IOExce
     }
 
     public static VectorData fromFloats(float[] vec) {
-        if (vec == null) return null;
-        return new VectorData(vec);
+        return vec == null ? null : new VectorData(vec);
     }
 
     public static VectorData fromBytes(byte[] vec) {
-        if (vec == null) return null;
-        return new VectorData(vec);
+        return vec == null ? null : new VectorData(vec);
     }
 }

server/src/test/java/org/elasticsearch/search/vectors/VectorDataTests.java

@@ -28,7 +28,10 @@ public void testThrowsIfBothVectorsAreNull() {
     }
 
     public void testThrowsIfBothVectorsAreNonNull() {
-        IllegalArgumentException ex = expectThrows(IllegalArgumentException.class, () -> new VectorData(new float[] {0f}, new byte[] {1}));
+        IllegalArgumentException ex = expectThrows(
+            IllegalArgumentException.class,
+            () -> new VectorData(new float[] { 0f }, new byte[] { 1 })
+        );
         assertThat(ex.getMessage(), containsString("please supply exactly either a float or a byte vector"));
     }