ES|QL: Add support for cached strings in plan serialization (#112929)

luigidellaquila · web-flow · commit a9b9172acc4d · 2024-10-11T10:57:13.000+02:00
diff --git a/docs/changelog/112929.yaml b/docs/changelog/112929.yaml
@@ -0,0 +1,5 @@
+pr: 112929
+summary: "ES|QL: Add support for cached strings in plan serialization"
+area: ES|QL
+type: enhancement
+issues: []
diff --git a/server/src/main/java/org/elasticsearch/TransportVersions.java b/server/src/main/java/org/elasticsearch/TransportVersions.java
@@ -239,6 +239,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion TEXT_SIMILARITY_RERANKER_QUERY_REWRITE = def(8_763_00_0);
     public static final TransportVersion SIMULATE_INDEX_TEMPLATES_SUBSTITUTIONS = def(8_764_00_0);
     public static final TransportVersion RETRIEVERS_TELEMETRY_ADDED = def(8_765_00_0);
+    public static final TransportVersion ESQL_CACHED_STRING_SERIALIZATION = def(8_766_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,
diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/FieldAttribute.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/FieldAttribute.java
@@ -115,7 +115,7 @@ private FieldAttribute(StreamInput in) throws IOException {
         this(
             Source.readFrom((StreamInput & PlanStreamInput) in),
             in.readOptionalWriteable(FieldAttribute::readFrom),
-            in.readString(),
+            ((PlanStreamInput) in).readCachedString(),
             DataType.readFrom(in),
             EsField.readFrom(in),
             in.readOptionalString(),
@@ -130,7 +130,7 @@ public void writeTo(StreamOutput out) throws IOException {
         if (((PlanStreamOutput) out).writeAttributeCacheHeader(this)) {
             Source.EMPTY.writeTo(out);
             out.writeOptionalWriteable(parent);
-            out.writeString(name());
+            ((PlanStreamOutput) out).writeCachedString(name());
             dataType().writeTo(out);
             field.writeTo(out);
             // We used to write the qualifier here. We can still do if needed in the future.
diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/DataType.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/DataType.java
@@ -14,6 +14,8 @@
 import org.elasticsearch.index.mapper.SourceFieldMapper;
 import org.elasticsearch.index.mapper.TimeSeriesIdFieldMapper;
 import org.elasticsearch.xpack.esql.core.plugin.EsqlCorePlugin;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamInput;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamOutput;
 
 import java.io.IOException;
 import java.math.BigInteger;
@@ -519,12 +521,12 @@ public DataType counter() {
     }
 
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeString(typeName);
+        ((PlanStreamOutput) out).writeCachedString(typeName);
     }
 
     public static DataType readFrom(StreamInput in) throws IOException {
         // TODO: Use our normal enum serialization pattern
-        return readFrom(in.readString());
+        return readFrom(((PlanStreamInput) in).readCachedString());
     }
 
     /**
diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/DateEsField.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/DateEsField.java
@@ -8,6 +8,8 @@
 
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamInput;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamOutput;
 
 import java.io.IOException;
 import java.util.Map;
@@ -26,12 +28,12 @@ private DateEsField(String name, DataType dataType, Map<String, EsField> propert
     }
 
     protected DateEsField(StreamInput in) throws IOException {
-        this(in.readString(), DataType.DATETIME, in.readImmutableMap(EsField::readFrom), in.readBoolean());
+        this(((PlanStreamInput) in).readCachedString(), DataType.DATETIME, in.readImmutableMap(EsField::readFrom), in.readBoolean());
     }
 
     @Override
     public void writeContent(StreamOutput out) throws IOException {
-        out.writeString(getName());
+        ((PlanStreamOutput) out).writeCachedString(getName());
         out.writeMap(getProperties(), (o, x) -> x.writeTo(out));
         out.writeBoolean(isAggregatable());
     }
diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/EsField.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/EsField.java
@@ -60,15 +60,15 @@ public EsField(String name, DataType esDataType, Map<String, EsField> properties
     }
 
     public EsField(StreamInput in) throws IOException {
-        this.name = in.readString();
+        this.name = ((PlanStreamInput) in).readCachedString();
         this.esDataType = readDataType(in);
         this.properties = in.readImmutableMap(EsField::readFrom);
         this.aggregatable = in.readBoolean();
         this.isAlias = in.readBoolean();
     }
 
     private DataType readDataType(StreamInput in) throws IOException {
-        String name = in.readString();
+        String name = ((PlanStreamInput) in).readCachedString();
         if (in.getTransportVersion().before(TransportVersions.ESQL_NESTED_UNSUPPORTED) && name.equalsIgnoreCase("NESTED")) {
             /*
              * The "nested" data type existed in older versions of ESQL but was
@@ -98,7 +98,7 @@ public void writeTo(StreamOutput out) throws IOException {
      * This needs to be overridden by subclasses for specific serialization
      */
     public void writeContent(StreamOutput out) throws IOException {
-        out.writeString(name);
+        ((PlanStreamOutput) out).writeCachedString(name);
         esDataType.writeTo(out);
         out.writeMap(properties, (o, x) -> x.writeTo(out));
         out.writeBoolean(aggregatable);
diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/InvalidMappedField.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/InvalidMappedField.java
@@ -10,6 +10,8 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.xpack.esql.core.QlIllegalArgumentException;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamInput;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamOutput;
 
 import java.io.IOException;
 import java.util.Map;
@@ -52,7 +54,7 @@ private InvalidMappedField(String name, String errorMessage, Map<String, EsField
     }
 
     protected InvalidMappedField(StreamInput in) throws IOException {
-        this(in.readString(), in.readString(), in.readImmutableMap(StreamInput::readString, EsField::readFrom));
+        this(((PlanStreamInput) in).readCachedString(), in.readString(), in.readImmutableMap(StreamInput::readString, EsField::readFrom));
     }
 
     public Set<DataType> types() {
@@ -61,7 +63,7 @@ public Set<DataType> types() {
 
     @Override
     public void writeContent(StreamOutput out) throws IOException {
-        out.writeString(getName());
+        ((PlanStreamOutput) out).writeCachedString(getName());
         out.writeString(errorMessage);
         out.writeMap(getProperties(), (o, x) -> x.writeTo(out));
     }
diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/KeywordEsField.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/KeywordEsField.java
@@ -8,6 +8,8 @@
 
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamInput;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamOutput;
 
 import java.io.IOException;
 import java.util.Collections;
@@ -59,7 +61,7 @@ protected KeywordEsField(
 
     public KeywordEsField(StreamInput in) throws IOException {
         this(
-            in.readString(),
+            ((PlanStreamInput) in).readCachedString(),
             KEYWORD,
             in.readImmutableMap(EsField::readFrom),
             in.readBoolean(),
@@ -71,7 +73,7 @@ public KeywordEsField(StreamInput in) throws IOException {
 
     @Override
     public void writeContent(StreamOutput out) throws IOException {
-        out.writeString(getName());
+        ((PlanStreamOutput) out).writeCachedString(getName());
         out.writeMap(getProperties(), (o, x) -> x.writeTo(out));
         out.writeBoolean(isAggregatable());
         out.writeInt(precision);
diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/MultiTypeEsField.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/MultiTypeEsField.java
@@ -10,6 +10,8 @@
 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.util.PlanStreamInput;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamOutput;
 
 import java.io.IOException;
 import java.util.HashMap;
@@ -36,13 +38,18 @@ public MultiTypeEsField(String name, DataType dataType, boolean aggregatable, Ma
     }
 
     protected MultiTypeEsField(StreamInput in) throws IOException {
-        this(in.readString(), DataType.readFrom(in), in.readBoolean(), in.readImmutableMap(i -> i.readNamedWriteable(Expression.class)));
+        this(
+            ((PlanStreamInput) in).readCachedString(),
+            DataType.readFrom(in),
+            in.readBoolean(),
+            in.readImmutableMap(i -> i.readNamedWriteable(Expression.class))
+        );
     }
 
     @Override
     public void writeContent(StreamOutput out) throws IOException {
-        out.writeString(getName());
-        out.writeString(getDataType().typeName());
+        ((PlanStreamOutput) out).writeCachedString(getName());
+        getDataType().writeTo(out);
         out.writeBoolean(isAggregatable());
         out.writeMap(getIndexToConversionExpressions(), (o, v) -> out.writeNamedWriteable(v));
     }
diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/TextEsField.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/TextEsField.java
@@ -10,6 +10,8 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.core.Tuple;
 import org.elasticsearch.xpack.esql.core.QlIllegalArgumentException;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamInput;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamOutput;
 
 import java.io.IOException;
 import java.util.Map;
@@ -32,12 +34,12 @@ public TextEsField(String name, Map<String, EsField> properties, boolean hasDocV
     }
 
     protected TextEsField(StreamInput in) throws IOException {
-        this(in.readString(), in.readImmutableMap(EsField::readFrom), in.readBoolean(), in.readBoolean());
+        this(((PlanStreamInput) in).readCachedString(), in.readImmutableMap(EsField::readFrom), in.readBoolean(), in.readBoolean());
     }
 
     @Override
     public void writeContent(StreamOutput out) throws IOException {
-        out.writeString(getName());
+        ((PlanStreamOutput) out).writeCachedString(getName());
         out.writeMap(getProperties(), (o, x) -> x.writeTo(out));
         out.writeBoolean(isAggregatable());
         out.writeBoolean(isAlias());
diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/UnsupportedEsField.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/UnsupportedEsField.java
@@ -8,6 +8,8 @@
 
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamInput;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamOutput;
 
 import java.io.IOException;
 import java.util.Map;
@@ -34,13 +36,18 @@ public UnsupportedEsField(String name, String originalType, String inherited, Ma
     }
 
     public UnsupportedEsField(StreamInput in) throws IOException {
-        this(in.readString(), in.readString(), in.readOptionalString(), in.readImmutableMap(EsField::readFrom));
+        this(
+            ((PlanStreamInput) in).readCachedString(),
+            ((PlanStreamInput) in).readCachedString(),
+            in.readOptionalString(),
+            in.readImmutableMap(EsField::readFrom)
+        );
     }
 
     @Override
     public void writeContent(StreamOutput out) throws IOException {
-        out.writeString(getName());
-        out.writeString(getOriginalType());
+        ((PlanStreamOutput) out).writeCachedString(getName());
+        ((PlanStreamOutput) out).writeCachedString(getOriginalType());
         out.writeOptionalString(getInherited());
         out.writeMap(getProperties(), (o, x) -> x.writeTo(out));
     }
diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/util/PlanStreamInput.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/util/PlanStreamInput.java
@@ -47,4 +47,6 @@ public interface PlanStreamInput {
     <A extends Attribute> A readAttributeWithCache(CheckedFunction<StreamInput, A, IOException> constructor) throws IOException;
 
     <A extends EsField> A readEsFieldWithCache() throws IOException;
+
+    String readCachedString() throws IOException;
 }
diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/util/PlanStreamOutput.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/util/PlanStreamOutput.java
@@ -31,4 +31,6 @@ public interface PlanStreamOutput {
      * @throws IOException
      */
     boolean writeEsFieldCacheHeader(EsField field) throws IOException;
+
+    void writeCachedString(String field) throws IOException;
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/UnsupportedAttribute.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/UnsupportedAttribute.java
@@ -76,7 +76,7 @@ public UnsupportedAttribute(Source source, String name, UnsupportedEsField field
     private UnsupportedAttribute(StreamInput in) throws IOException {
         this(
             Source.readFrom((PlanStreamInput) in),
-            in.readString(),
+            ((PlanStreamInput) in).readCachedString(),
             in.getTransportVersion().onOrAfter(TransportVersions.ESQL_ES_FIELD_CACHED_SERIALIZATION)
                 || in.getTransportVersion().isPatchFrom(TransportVersions.ESQL_ATTRIBUTE_CACHED_SERIALIZATION_8_15)
                     ? EsField.readFrom(in)
@@ -90,7 +90,7 @@ private UnsupportedAttribute(StreamInput in) throws IOException {
     public void writeTo(StreamOutput out) throws IOException {
         if (((PlanStreamOutput) out).writeAttributeCacheHeader(this)) {
             Source.EMPTY.writeTo(out);
-            out.writeString(name());
+            ((PlanStreamOutput) out).writeCachedString(name());
             if (out.getTransportVersion().onOrAfter(TransportVersions.ESQL_ES_FIELD_CACHED_SERIALIZATION)
                 || out.getTransportVersion().isPatchFrom(TransportVersions.ESQL_ATTRIBUTE_CACHED_SERIALIZATION_8_15)) {
                 field().writeTo(out);
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanStreamInput.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanStreamInput.java
@@ -62,9 +62,11 @@ public NameId apply(long streamNameId) {
 
     private final Map<Integer, Block> cachedBlocks = new HashMap<>();
 
-    private Attribute[] attributesCache = new Attribute[64];
+    private Attribute[] attributesCache = new Attribute[1024];
 
-    private EsField[] esFieldsCache = new EsField[64];
+    private EsField[] esFieldsCache = new EsField[1024];
+
+    private String[] stringCache = new String[1024];
 
     // hook for nameId, where can cache and map, for now just return a NameId of the same long value.
     private final LongFunction<NameId> nameIdFunction;
@@ -195,10 +197,11 @@ public <A extends Attribute> A readAttributeWithCache(CheckedFunction<StreamInpu
     }
 
     private Attribute attributeFromCache(int id) throws IOException {
-        if (attributesCache[id] == null) {
+        Attribute attribute = attributesCache[id];
+        if (attribute == null) {
             throw new IOException("Attribute ID not found in serialization cache [" + id + "]");
         }
-        return attributesCache[id];
+        return attribute;
     }
 
     /**
@@ -221,7 +224,7 @@ public <A extends EsField> A readEsFieldWithCache() throws IOException {
             // it's safe to cast to int, since the max value for this is {@link PlanStreamOutput#MAX_SERIALIZED_ATTRIBUTES}
             int cacheId = Math.toIntExact(readZLong());
             if (cacheId < 0) {
-                String className = readString();
+                String className = readCachedString();
                 Writeable.Reader<? extends EsField> reader = EsField.getReader(className);
                 cacheId = -1 - cacheId;
                 EsField result = reader.read(this);
@@ -231,17 +234,37 @@ public <A extends EsField> A readEsFieldWithCache() throws IOException {
                 return (A) esFieldFromCache(cacheId);
             }
         } else {
-            String className = readString();
+            String className = readCachedString();
             Writeable.Reader<? extends EsField> reader = EsField.getReader(className);
             return (A) reader.read(this);
         }
     }
 
+    /**
+     * Reads a cached string, serialized with {@link PlanStreamOutput#writeCachedString(String)}.
+     */
+    @Override
+    public String readCachedString() throws IOException {
+        if (getTransportVersion().before(TransportVersions.ESQL_CACHED_STRING_SERIALIZATION)) {
+            return readString();
+        }
+        int cacheId = Math.toIntExact(readZLong());
+        if (cacheId < 0) {
+            String string = readString();
+            cacheId = -1 - cacheId;
+            cacheString(cacheId, string);
+            return string;
+        } else {
+            return stringFromCache(cacheId);
+        }
+    }
+
     private EsField esFieldFromCache(int id) throws IOException {
-        if (esFieldsCache[id] == null) {
+        EsField field = esFieldsCache[id];
+        if (field == null) {
             throw new IOException("Attribute ID not found in serialization cache [" + id + "]");
         }
-        return esFieldsCache[id];
+        return field;
     }
 
     /**
@@ -257,4 +280,27 @@ private void cacheEsField(int id, EsField field) {
         esFieldsCache[id] = field;
     }
 
+    private String stringFromCache(int id) throws IOException {
+        String value = stringCache[id];
+        if (value == null) {
+            throw new IOException("String not found in serialization cache [" + id + "]");
+        }
+        return value;
+    }
+
+    private void cacheString(int id, String string) {
+        assert id >= 0;
+        if (id >= stringCache.length) {
+            stringCache = ArrayUtil.grow(stringCache, id + 1);
+        }
+        stringCache[id] = string;
+    }
+
+    @Override
+    public void close() throws IOException {
+        super.close();
+        this.stringCache = null;
+        this.attributesCache = null;
+        this.esFieldsCache = null;
+    }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanStreamOutput.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanStreamOutput.java
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/index/EsIndexSerializationTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/index/EsIndexSerializationTests.java
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/physical/ExchangeSinkExecSerializationTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/physical/ExchangeSinkExecSerializationTests.java

Original file line number	Diff line number	Diff line change
`@@ -14,6 +14,8 @@`
`14`	`14`	`import org.elasticsearch.index.mapper.SourceFieldMapper;`
`15`	`15`	`import org.elasticsearch.index.mapper.TimeSeriesIdFieldMapper;`
`16`	`16`	`import org.elasticsearch.xpack.esql.core.plugin.EsqlCorePlugin;`
	`17`	`+import org.elasticsearch.xpack.esql.core.util.PlanStreamInput;`
	`18`	`+import org.elasticsearch.xpack.esql.core.util.PlanStreamOutput;`
`17`	`19`
`18`	`20`	`import java.io.IOException;`
`19`	`21`	`import java.math.BigInteger;`
`@@ -519,12 +521,12 @@ public DataType counter() {`
`519`	`521`	`}`
`520`	`522`
`521`	`523`	`public void writeTo(StreamOutput out) throws IOException {`
`522`		`- out.writeString(typeName);`
	`524`	`+ ((PlanStreamOutput) out).writeCachedString(typeName);`
`523`	`525`	`}`
`524`	`526`
`525`	`527`	`public static DataType readFrom(StreamInput in) throws IOException {`
`526`	`528`	`// TODO: Use our normal enum serialization pattern`
`527`		`- return readFrom(in.readString());`
	`529`	`+ return readFrom(((PlanStreamInput) in).readCachedString());`
`528`	`530`	`}`
`529`	`531`
`530`	`532`	`/**`
Original file line number	Diff line number	Diff line change
`@@ -8,6 +8,8 @@`
`8`	`8`
`9`	`9`	`import org.elasticsearch.common.io.stream.StreamInput;`
`10`	`10`	`import org.elasticsearch.common.io.stream.StreamOutput;`
	`11`	`+import org.elasticsearch.xpack.esql.core.util.PlanStreamInput;`
	`12`	`+import org.elasticsearch.xpack.esql.core.util.PlanStreamOutput;`
`11`	`13`
`12`	`14`	`import java.io.IOException;`
`13`	`15`	`import java.util.Map;`
`@@ -26,12 +28,12 @@ private DateEsField(String name, DataType dataType, Map<String, EsField> propert`
`26`	`28`	`}`
`27`	`29`
`28`	`30`	`protected DateEsField(StreamInput in) throws IOException {`
`29`		`- this(in.readString(), DataType.DATETIME, in.readImmutableMap(EsField::readFrom), in.readBoolean());`
	`31`	`+ this(((PlanStreamInput) in).readCachedString(), DataType.DATETIME, in.readImmutableMap(EsField::readFrom), in.readBoolean());`
`30`	`32`	`}`
`31`	`33`
`32`	`34`	`@Override`
`33`	`35`	`public void writeContent(StreamOutput out) throws IOException {`
`34`		`- out.writeString(getName());`
	`36`	`+ ((PlanStreamOutput) out).writeCachedString(getName());`
`35`	`37`	`out.writeMap(getProperties(), (o, x) -> x.writeTo(out));`
`36`	`38`	`out.writeBoolean(isAggregatable());`
`37`	`39`	`}`