浏览代码

Dry up o.e.c.lucene.Lucene transport logic (#126646)

We've accumulated quited a bit of duplication here, drying this logic up a little.
Armin Braun 6 月之前
父节点
当前提交
55453aeef6

+ 44 - 76
server/src/main/java/org/elasticsearch/common/lucene/Lucene.java

@@ -313,18 +313,15 @@ public class Lucene {
         } else if (type == 1) {
             TotalHits totalHits = readTotalHits(in);
             float maxScore = in.readFloat();
-            SortField[] fields = in.readArray(Lucene::readSortField, SortField[]::new);
-            FieldDoc[] fieldDocs = new FieldDoc[in.readVInt()];
-            for (int i = 0; i < fieldDocs.length; i++) {
-                fieldDocs[i] = readFieldDoc(in);
-            }
+            SortField[] fields = readSortFieldArray(in);
+            FieldDoc[] fieldDocs = in.readArray(Lucene::readFieldDoc, FieldDoc[]::new);
             return new TopDocsAndMaxScore(new TopFieldDocs(totalHits, fieldDocs, fields), maxScore);
         } else if (type == 2) {
             TotalHits totalHits = readTotalHits(in);
             float maxScore = in.readFloat();
 
             String field = in.readString();
-            SortField[] fields = in.readArray(Lucene::readSortField, SortField[]::new);
+            SortField[] fields = readSortFieldArray(in);
             int size = in.readVInt();
             Object[] collapseValues = new Object[size];
             FieldDoc[] fieldDocs = new FieldDoc[size];
@@ -339,65 +336,30 @@ public class Lucene {
     }
 
     public static FieldDoc readFieldDoc(StreamInput in) throws IOException {
-        Comparable<?>[] cFields = new Comparable<?>[in.readVInt()];
-        for (int j = 0; j < cFields.length; j++) {
-            byte type = in.readByte();
-            if (type == 0) {
-                cFields[j] = null;
-            } else if (type == 1) {
-                cFields[j] = in.readString();
-            } else if (type == 2) {
-                cFields[j] = in.readInt();
-            } else if (type == 3) {
-                cFields[j] = in.readLong();
-            } else if (type == 4) {
-                cFields[j] = in.readFloat();
-            } else if (type == 5) {
-                cFields[j] = in.readDouble();
-            } else if (type == 6) {
-                cFields[j] = in.readByte();
-            } else if (type == 7) {
-                cFields[j] = in.readShort();
-            } else if (type == 8) {
-                cFields[j] = in.readBoolean();
-            } else if (type == 9) {
-                cFields[j] = in.readBytesRef();
-            } else if (type == 10) {
-                cFields[j] = new BigInteger(in.readString());
-            } else {
-                throw new IOException("Can't match type [" + type + "]");
-            }
-        }
-        return new FieldDoc(in.readVInt(), in.readFloat(), cFields);
+        var sortValues = readSortValues(in);
+        return new FieldDoc(in.readVInt(), in.readFloat(), sortValues);
+    }
+
+    public static Object[] readSortValues(StreamInput in) throws IOException {
+        return in.readArray(Lucene::readSortValue, Object[]::new);
     }
 
     public static Comparable<?> readSortValue(StreamInput in) throws IOException {
         byte type = in.readByte();
-        if (type == 0) {
-            return null;
-        } else if (type == 1) {
-            return in.readString();
-        } else if (type == 2) {
-            return in.readInt();
-        } else if (type == 3) {
-            return in.readLong();
-        } else if (type == 4) {
-            return in.readFloat();
-        } else if (type == 5) {
-            return in.readDouble();
-        } else if (type == 6) {
-            return in.readByte();
-        } else if (type == 7) {
-            return in.readShort();
-        } else if (type == 8) {
-            return in.readBoolean();
-        } else if (type == 9) {
-            return in.readBytesRef();
-        } else if (type == 10) {
-            return new BigInteger(in.readString());
-        } else {
-            throw new IOException("Can't match type [" + type + "]");
-        }
+        return switch (type) {
+            case 0 -> null;
+            case 1 -> in.readString();
+            case 2 -> in.readInt();
+            case 3 -> in.readLong();
+            case 4 -> in.readFloat();
+            case 5 -> in.readDouble();
+            case 6 -> in.readByte();
+            case 7 -> in.readShort();
+            case 8 -> in.readBoolean();
+            case 9 -> in.readBytesRef();
+            case 10 -> new BigInteger(in.readString());
+            default -> throw new IOException("Can't match type [" + type + "]");
+        };
     }
 
     public static ScoreDoc readScoreDoc(StreamInput in) throws IOException {
@@ -426,7 +388,7 @@ public class Lucene {
             out.writeByte((byte) 2);
             writeTotalHits(out, topDocs.totalHits);
             out.writeString(topFieldGroups.field);
-            out.writeArray(Lucene::writeSortField, topFieldGroups.fields);
+            writeSortFieldArray(out, topFieldGroups.fields);
             out.writeVInt(topDocs.scoreDocs.length);
             for (int i = 0; i < topDocs.scoreDocs.length; i++) {
                 ScoreDoc doc = topFieldGroups.scoreDocs[i];
@@ -437,7 +399,7 @@ public class Lucene {
         } else if (topDocs instanceof TopFieldDocs topFieldDocs) {
             out.writeByte((byte) 1);
             writeTotalHits(out, topDocs.totalHits);
-            out.writeArray(Lucene::writeSortField, topFieldDocs.fields);
+            writeSortFieldArray(out, topFieldDocs.fields);
             out.writeArray((o, doc) -> {
                 writeFieldDoc(o, (FieldDoc) doc);
                 o.writeVInt(doc.shardIndex);
@@ -452,6 +414,10 @@ public class Lucene {
         }
     }
 
+    public static void writeSortFieldArray(StreamOutput out, SortField[] sortFields) throws IOException {
+        out.writeArray(Lucene::writeSortField, sortFields);
+    }
+
     /**
      * Read side counterpart to {@link #writeTopDocsIncludingShardIndex} and the same as {@link #readTopDocs(StreamInput)} but for the
      * added shard index values that are read.
@@ -474,7 +440,7 @@ public class Lucene {
             return new TopDocs(totalHits, scoreDocs);
         } else if (type == 1) {
             TotalHits totalHits = readTotalHits(in);
-            SortField[] fields = in.readArray(Lucene::readSortField, SortField[]::new);
+            SortField[] fields = readSortFieldArray(in);
             FieldDoc[] fieldDocs = new FieldDoc[in.readVInt()];
             for (int i = 0; i < fieldDocs.length; i++) {
                 var fieldDoc = readFieldDoc(in);
@@ -485,7 +451,7 @@ public class Lucene {
         } else if (type == 2) {
             TotalHits totalHits = readTotalHits(in);
             String field = in.readString();
-            SortField[] fields = in.readArray(Lucene::readSortField, SortField[]::new);
+            SortField[] fields = readSortFieldArray(in);
             int size = in.readVInt();
             Object[] collapseValues = new Object[size];
             FieldDoc[] fieldDocs = new FieldDoc[size];
@@ -501,6 +467,10 @@ public class Lucene {
         }
     }
 
+    public static SortField[] readSortFieldArray(StreamInput in) throws IOException {
+        return in.readArray(Lucene::readSortField, SortField[]::new);
+    }
+
     public static void writeTopDocs(StreamOutput out, TopDocsAndMaxScore topDocs) throws IOException {
         if (topDocs.topDocs instanceof TopFieldGroups topFieldGroups) {
             out.writeByte((byte) 2);
@@ -509,7 +479,7 @@ public class Lucene {
             out.writeFloat(topDocs.maxScore);
 
             out.writeString(topFieldGroups.field);
-            out.writeArray(Lucene::writeSortField, topFieldGroups.fields);
+            writeSortFieldArray(out, topFieldGroups.fields);
 
             out.writeVInt(topFieldGroups.scoreDocs.length);
             for (int i = 0; i < topFieldGroups.scoreDocs.length; i++) {
@@ -523,7 +493,7 @@ public class Lucene {
             writeTotalHits(out, topFieldDocs.totalHits);
             out.writeFloat(topDocs.maxScore);
 
-            out.writeArray(Lucene::writeSortField, topFieldDocs.fields);
+            writeSortFieldArray(out, topFieldDocs.fields);
             out.writeArray((o, doc) -> writeFieldDoc(o, (FieldDoc) doc), topFieldDocs.scoreDocs);
         } else {
             out.writeByte((byte) 0);
@@ -598,14 +568,17 @@ public class Lucene {
 
     public static void writeFieldDoc(StreamOutput out, FieldDoc fieldDoc) throws IOException {
         out.writeArray(Lucene::writeSortValue, fieldDoc.fields);
-        out.writeVInt(fieldDoc.doc);
-        out.writeFloat(fieldDoc.score);
+        doWriteScoreDoc(out, fieldDoc);
     }
 
     public static void writeScoreDoc(StreamOutput out, ScoreDoc scoreDoc) throws IOException {
         if (scoreDoc.getClass().equals(ScoreDoc.class) == false) {
             throw new IllegalArgumentException("This method can only be used to serialize a ScoreDoc, not a " + scoreDoc.getClass());
         }
+        doWriteScoreDoc(out, scoreDoc);
+    }
+
+    private static void doWriteScoreDoc(StreamOutput out, ScoreDoc scoreDoc) throws IOException {
         out.writeVInt(scoreDoc.doc);
         out.writeFloat(scoreDoc.score);
     }
@@ -662,17 +635,12 @@ public class Lucene {
         }
     }
 
-    public static void writeSortField(StreamOutput out, SortField sortField) throws IOException {
+    static void writeSortField(StreamOutput out, SortField sortField) throws IOException {
         sortField = rewriteMergeSortField(sortField);
         if (sortField.getClass() != SortField.class) {
             throw new IllegalArgumentException("Cannot serialize SortField impl [" + sortField + "]");
         }
-        if (sortField.getField() == null) {
-            out.writeBoolean(false);
-        } else {
-            out.writeBoolean(true);
-            out.writeString(sortField.getField());
-        }
+        out.writeOptionalString(sortField.getField());
         if (sortField.getComparatorSource() != null) {
             IndexFieldData.XFieldComparatorSource comparatorSource = (IndexFieldData.XFieldComparatorSource) sortField
                 .getComparatorSource();

+ 3 - 3
server/src/main/java/org/elasticsearch/search/SearchHits.java

@@ -140,9 +140,9 @@ public final class SearchHits implements Writeable, ChunkedToXContent, RefCounte
                 isPooled = isPooled || hit.isPooled();
             }
         }
-        var sortFields = in.readOptionalArray(Lucene::readSortField, SortField[]::new);
+        var sortFields = in.readOptional(Lucene::readSortFieldArray);
         var collapseField = in.readOptionalString();
-        var collapseValues = in.readOptionalArray(Lucene::readSortValue, Object[]::new);
+        var collapseValues = in.readOptional(Lucene::readSortValues);
         if (isPooled) {
             return new SearchHits(hits, totalHits, maxScore, sortFields, collapseField, collapseValues);
         } else {
@@ -164,7 +164,7 @@ public final class SearchHits implements Writeable, ChunkedToXContent, RefCounte
         }
         out.writeFloat(maxScore);
         out.writeArray(hits);
-        out.writeOptionalArray(Lucene::writeSortField, sortFields);
+        out.writeOptional(Lucene::writeSortFieldArray, sortFields);
         out.writeOptionalString(collapseField);
         out.writeOptionalArray(Lucene::writeSortValue, collapseValues);
     }

+ 2 - 2
server/src/main/java/org/elasticsearch/search/SearchSortValues.java

@@ -51,8 +51,8 @@ public class SearchSortValues implements ToXContentFragment, Writeable {
     }
 
     public static SearchSortValues readFrom(StreamInput in) throws IOException {
-        Object[] formattedSortValues = in.readArray(Lucene::readSortValue, Object[]::new);
-        Object[] rawSortValues = in.readArray(Lucene::readSortValue, Object[]::new);
+        Object[] formattedSortValues = Lucene.readSortValues(in);
+        Object[] rawSortValues = Lucene.readSortValues(in);
         if (formattedSortValues.length == 0 && rawSortValues.length == 0) {
             return EMPTY;
         }

+ 2 - 2
server/src/main/java/org/elasticsearch/search/SearchSortValuesAndFormats.java

@@ -48,8 +48,8 @@ public class SearchSortValuesAndFormats implements Writeable {
     }
 
     public SearchSortValuesAndFormats(StreamInput in) throws IOException {
-        this.rawSortValues = in.readArray(Lucene::readSortValue, Object[]::new);
-        this.formattedSortValues = in.readArray(Lucene::readSortValue, Object[]::new);
+        this.rawSortValues = Lucene.readSortValues(in);
+        this.formattedSortValues = Lucene.readSortValues(in);
         this.sortValueFormats = new DocValueFormat[formattedSortValues.length];
         for (int i = 0; i < sortValueFormats.length; ++i) {
             sortValueFormats[i] = in.readNamedWriteable(DocValueFormat.class);