ES|QL: cache EsField on serialization (#112008)

As a follow-up to #111447,
with this change we also cache `EsFields`.

This gives us an additional 30-40% reduction on the size of serialized
plan, according to [these
tests](#111980)

Related to #111358

Author: luigidellaquila

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -197,6 +197,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion LTR_SERVERLESS_RELEASE = def(8_727_00_0);
     public static final TransportVersion ALLOW_PARTIAL_SEARCH_RESULTS_IN_PIT = def(8_728_00_0);
     public static final TransportVersion RANK_DOCS_RETRIEVER = def(8_729_00_0);
+    public static final TransportVersion ESQL_ES_FIELD_CACHED_SERIALIZATION = def(8_730_00_0);
     /*
      * STOP! READ THIS FIRST! No, really,
      *        ____ _____ ___  ____  _        ____  _____    _    ____    _____ _   _ ___ ____    _____ ___ ____  ____ _____ _

test/framework/src/main/java/org/elasticsearch/test/AbstractWireSerializingTestCase.java

@@ -25,11 +25,18 @@ public abstract class AbstractWireSerializingTestCase<T extends Writeable> exten
      */
     protected abstract Writeable.Reader<T> instanceReader();
 
+    /**
+     * Returns a {@link Writeable.Writer} that will be used to serialize the instance
+     */
+    protected Writeable.Writer<T> instanceWriter() {
+        return StreamOutput::writeWriteable;
+    }
+
     /**
      * Copy the {@link Writeable} by round tripping it through {@linkplain StreamInput} and {@linkplain StreamOutput}.
      */
     @Override
     protected final T copyInstance(T instance, TransportVersion version) throws IOException {
-        return copyWriteable(instance, getNamedWriteableRegistry(), instanceReader(), version);
+        return copyInstance(instance, getNamedWriteableRegistry(), instanceWriter(), instanceReader(), version);
     }
 }

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/FieldAttribute.java

@@ -112,7 +112,7 @@ private FieldAttribute(StreamInput in) throws IOException {
             in.readOptionalWriteable(FieldAttribute::readFrom),
             in.readString(),
             DataType.readFrom(in),
-            in.readNamedWriteable(EsField.class),
+            EsField.readFrom(in),
             in.readOptionalString(),
             in.readEnum(Nullability.class),
             NameId.readFrom((StreamInput & PlanStreamInput) in),
@@ -127,7 +127,7 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeOptionalWriteable(parent);
             out.writeString(name());
             dataType().writeTo(out);
-            out.writeNamedWriteable(field);
+            field.writeTo(out);
             // We used to write the qualifier here. We can still do if needed in the future.
             out.writeOptionalString(null);
             out.writeEnum(nullable());

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/DateEsField.java

@@ -6,7 +6,6 @@
  */
 package org.elasticsearch.xpack.esql.core.type;
 
-import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 
@@ -17,7 +16,6 @@
  * Information about a field in an ES index with the {@code date} type
  */
 public class DateEsField extends EsField {
-    static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(EsField.class, "DateEsField", DateEsField::new);
 
     public static DateEsField dateEsField(String name, Map<String, EsField> properties, boolean hasDocValues) {
         return new DateEsField(name, DataType.DATETIME, properties, hasDocValues);
@@ -27,19 +25,19 @@ private DateEsField(String name, DataType dataType, Map<String, EsField> propert
         super(name, dataType, properties, hasDocValues);
     }
 
-    private DateEsField(StreamInput in) throws IOException {
-        this(in.readString(), DataType.DATETIME, in.readMap(i -> i.readNamedWriteable(EsField.class)), in.readBoolean());
+    protected DateEsField(StreamInput in) throws IOException {
+        this(in.readString(), DataType.DATETIME, in.readImmutableMap(EsField::readFrom), in.readBoolean());
     }
 
     @Override
-    public void writeTo(StreamOutput out) throws IOException {
+    protected void writeContent(StreamOutput out) throws IOException {
         out.writeString(getName());
-        out.writeMap(getProperties(), StreamOutput::writeNamedWriteable);
+        out.writeMap(getProperties(), (o, x) -> x.writeTo(out));
         out.writeBoolean(isAggregatable());
     }
 
-    @Override
     public String getWriteableName() {
-        return ENTRY.name;
+        return "DateEsField";
     }
+
 }

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/EsField.java

@@ -7,34 +7,40 @@
 package org.elasticsearch.xpack.esql.core.type;
 
 import org.elasticsearch.TransportVersions;
-import org.elasticsearch.common.io.stream.NamedWriteable;
-import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.core.Nullable;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamInput;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamOutput;
 
 import java.io.IOException;
-import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 
 /**
  * Information about a field in an ES index.
  */
-public class EsField implements NamedWriteable {
-    public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
-        return List.of(
-            EsField.ENTRY,
-            DateEsField.ENTRY,
-            InvalidMappedField.ENTRY,
-            KeywordEsField.ENTRY,
-            TextEsField.ENTRY,
-            UnsupportedEsField.ENTRY
-        );
+public class EsField implements Writeable {
+
+    private static Map<String, Writeable.Reader<? extends EsField>> readers = Map.ofEntries(
+        Map.entry("EsField", EsField::new),
+        Map.entry("DateEsField", DateEsField::new),
+        Map.entry("InvalidMappedField", InvalidMappedField::new),
+        Map.entry("KeywordEsField", KeywordEsField::new),
+        Map.entry("MultiTypeEsField", MultiTypeEsField::new),
+        Map.entry("TextEsField", TextEsField::new),
+        Map.entry("UnsupportedEsField", UnsupportedEsField::new)
+    );
+
+    public static Writeable.Reader<? extends EsField> getReader(String name) {
+        Reader<? extends EsField> result = readers.get(name);
+        if (result == null) {
+            throw new IllegalArgumentException("Invalid EsField type [" + name + "]");
+        }
+        return result;
     }
 
-    static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(EsField.class, "EsField", EsField::new);
-
     private final DataType esDataType;
     private final boolean aggregatable;
     private final Map<String, EsField> properties;
@@ -53,10 +59,10 @@ public EsField(String name, DataType esDataType, Map<String, EsField> properties
         this.isAlias = isAlias;
     }
 
-    public EsField(StreamInput in) throws IOException {
+    protected EsField(StreamInput in) throws IOException {
         this.name = in.readString();
         this.esDataType = readDataType(in);
-        this.properties = in.readImmutableMap(i -> i.readNamedWriteable(EsField.class));
+        this.properties = in.readImmutableMap(EsField::readFrom);
         this.aggregatable = in.readBoolean();
         this.isAlias = in.readBoolean();
     }
@@ -77,18 +83,33 @@ private DataType readDataType(StreamInput in) throws IOException {
         return DataType.readFrom(name);
     }
 
+    public static <A extends EsField> A readFrom(StreamInput in) throws IOException {
+        return ((PlanStreamInput) in).readEsFieldWithCache();
+    }
+
     @Override
     public void writeTo(StreamOutput out) throws IOException {
+        if (((PlanStreamOutput) out).writeEsFieldCacheHeader(this)) {
+            writeContent(out);
+        }
+    }
+
+    /**
+     * This needs to be overridden by subclasses for specific serialization
+     */
+    protected void writeContent(StreamOutput out) throws IOException {
         out.writeString(name);
         esDataType.writeTo(out);
-        out.writeMap(properties, StreamOutput::writeNamedWriteable);
+        out.writeMap(properties, (o, x) -> x.writeTo(out));
         out.writeBoolean(aggregatable);
         out.writeBoolean(isAlias);
     }
 
-    @Override
+    /**
+     * This needs to be overridden by subclasses for specific serialization
+     */
     public String getWriteableName() {
-        return ENTRY.name;
+        return "EsField";
     }
 
     /**

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/InvalidMappedField.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.esql.core.type;
 
-import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.xpack.esql.core.QlIllegalArgumentException;
@@ -27,11 +26,6 @@
  * It is used specifically for the 'union types' feature in ES|QL.
  */
 public class InvalidMappedField extends EsField {
-    static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
-        EsField.class,
-        "InvalidMappedField",
-        InvalidMappedField::new
-    );
 
     private final String errorMessage;
     private final Map<String, Set<String>> typesToIndices;
@@ -44,10 +38,6 @@ public InvalidMappedField(String name, String errorMessage) {
         this(name, errorMessage, new TreeMap<>());
     }
 
-    public InvalidMappedField(String name) {
-        this(name, StringUtils.EMPTY, new TreeMap<>());
-    }
-
     /**
      * Constructor supporting union types, used in ES|QL.
      */
@@ -61,24 +51,23 @@ private InvalidMappedField(String name, String errorMessage, Map<String, EsField
         this.typesToIndices = typesToIndices;
     }
 
-    private InvalidMappedField(StreamInput in) throws IOException {
-        this(in.readString(), in.readString(), in.readImmutableMap(StreamInput::readString, i -> i.readNamedWriteable(EsField.class)));
+    protected InvalidMappedField(StreamInput in) throws IOException {
+        this(in.readString(), in.readString(), in.readImmutableMap(StreamInput::readString, EsField::readFrom));
     }
 
     public Set<DataType> types() {
         return typesToIndices.keySet().stream().map(DataType::fromTypeName).collect(Collectors.toSet());
     }
 
     @Override
-    public void writeTo(StreamOutput out) throws IOException {
+    protected void writeContent(StreamOutput out) throws IOException {
         out.writeString(getName());
         out.writeString(errorMessage);
-        out.writeMap(getProperties(), StreamOutput::writeNamedWriteable);
+        out.writeMap(getProperties(), (o, x) -> x.writeTo(out));
     }
 
-    @Override
     public String getWriteableName() {
-        return ENTRY.name;
+        return "InvalidMappedField";
     }
 
     public String errorMessage() {

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/KeywordEsField.java

@@ -6,7 +6,6 @@
  */
 package org.elasticsearch.xpack.esql.core.type;
 
-import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 
@@ -21,11 +20,6 @@
  * Information about a field in an ES index with the {@code keyword} type.
  */
 public class KeywordEsField extends EsField {
-    static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
-        EsField.class,
-        "KeywordEsField",
-        KeywordEsField::new
-    );
 
     private final int precision;
     private final boolean normalized;
@@ -63,11 +57,11 @@ protected KeywordEsField(
         this.normalized = normalized;
     }
 
-    private KeywordEsField(StreamInput in) throws IOException {
+    public KeywordEsField(StreamInput in) throws IOException {
         this(
             in.readString(),
             KEYWORD,
-            in.readMap(i -> i.readNamedWriteable(EsField.class)),
+            in.readImmutableMap(EsField::readFrom),
             in.readBoolean(),
             in.readInt(),
             in.readBoolean(),
@@ -76,18 +70,17 @@ private KeywordEsField(StreamInput in) throws IOException {
     }
 
     @Override
-    public void writeTo(StreamOutput out) throws IOException {
+    protected void writeContent(StreamOutput out) throws IOException {
         out.writeString(getName());
-        out.writeMap(getProperties(), StreamOutput::writeNamedWriteable);
+        out.writeMap(getProperties(), (o, x) -> x.writeTo(out));
         out.writeBoolean(isAggregatable());
         out.writeInt(precision);
         out.writeBoolean(normalized);
         out.writeBoolean(isAlias());
     }
 
-    @Override
     public String getWriteableName() {
-        return ENTRY.name;
+        return "KeywordEsField";
     }
 
     public int getPrecision() {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/type/MultiTypeEsField.java renamed to x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/MultiTypeEsField.java

@@ -5,15 +5,11 @@
  * 2.0.
  */
 
-package org.elasticsearch.xpack.esql.type;
+package org.elasticsearch.xpack.esql.core.type;
 
-import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
-import org.elasticsearch.xpack.esql.core.type.DataType;
-import org.elasticsearch.xpack.esql.core.type.EsField;
-import org.elasticsearch.xpack.esql.core.type.InvalidMappedField;
 
 import java.io.IOException;
 import java.util.HashMap;
@@ -31,11 +27,6 @@
  * type conversion is done at the data node level.
  */
 public class MultiTypeEsField extends EsField {
-    public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
-        EsField.class,
-        "MultiTypeEsField",
-        MultiTypeEsField::new
-    );
 
     private final Map<String, Expression> indexToConversionExpressions;
 
@@ -44,21 +35,20 @@ public MultiTypeEsField(String name, DataType dataType, boolean aggregatable, Ma
         this.indexToConversionExpressions = indexToConversionExpressions;
     }
 
-    public MultiTypeEsField(StreamInput in) throws IOException {
+    protected MultiTypeEsField(StreamInput in) throws IOException {
         this(in.readString(), DataType.readFrom(in), in.readBoolean(), in.readImmutableMap(i -> i.readNamedWriteable(Expression.class)));
     }
 
     @Override
-    public void writeTo(StreamOutput out) throws IOException {
+    public void writeContent(StreamOutput out) throws IOException {
         out.writeString(getName());
         out.writeString(getDataType().typeName());
         out.writeBoolean(isAggregatable());
         out.writeMap(getIndexToConversionExpressions(), (o, v) -> out.writeNamedWriteable(v));
     }
 
-    @Override
     public String getWriteableName() {
-        return ENTRY.name;
+        return "MultiTypeEsField";
     }
 
     public Map<String, Expression> getIndexToConversionExpressions() {