ESQL: Fix serialization during can_match (elastic#111779)

This fixes that serialization of the `SingleValueQuery` during the `can_match` phase. Closes elastic#111726 Closes elastic#111701
cbuescher · Sep 4, 2024 · 1e8ac2f · 1e8ac2f
1 parent a64c6cf
commit 1e8ac2f
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 1 deletion.
diff --git a/docs/changelog/111779.yaml b/docs/changelog/111779.yaml
@@ -0,0 +1,7 @@
+pr: 111779
+summary: "ESQL: Fix serialization during `can_match`"
+area: ES|QL
+type: bug
+issues:
+ - 111701
+ - 111726
diff --git a/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/tree/Source.java b/x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/tree/Source.java
@@ -51,6 +51,21 @@ public static <S extends StreamInput & PlanStreamInput> Source readFrom(S in) th
         return new Source(new Location(line, charPositionInLine), text);
     }
 
+    /**
+     * Read the components of a {@link Source} and throw it away, returning
+     * {@link Source#EMPTY}. Use this when you will never use the {@link Source}
+     * and there is no chance of getting a {@link PlanStreamInput}.
+     */
+    public static Source readEmpty(StreamInput in) throws IOException {
+        if (in.readBoolean() == false) {
+            return EMPTY;
+        }
+        in.readInt();
+        in.readInt();
+        in.readInt();
+        return Source.EMPTY;
+    }
+
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         if (this == EMPTY) {

diff --git a/...ck/plugin/esql-core/src/test/java/org/elasticsearch/xpack/esql/core/tree/SourceTests.java b/...ck/plugin/esql-core/src/test/java/org/elasticsearch/xpack/esql/core/tree/SourceTests.java
@@ -6,13 +6,17 @@
  */
 package org.elasticsearch.xpack.esql.core.tree;
 
+import org.elasticsearch.common.io.stream.BytesStreamOutput;
+import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.test.ESTestCase;
 
+import java.io.IOException;
 import java.util.Arrays;
 import java.util.List;
 import java.util.function.Function;
 
 import static org.elasticsearch.test.EqualsHashCodeTestUtils.checkEqualsAndHashCode;
+import static org.hamcrest.Matchers.equalTo;
 
 public class SourceTests extends ESTestCase {
     public static Source randomSource() {
@@ -42,4 +46,16 @@ public void testEqualsAndHashCode() {
             SourceTests::mutate
         );
     }
+
+    public void testReadEmpty() throws IOException {
+        try (BytesStreamOutput out = new BytesStreamOutput()) {
+            randomSource().writeTo(out);
+            int test = randomInt();
+            out.writeInt(test);
+            try (StreamInput in = out.bytes().streamInput()) {
+                assertThat(Source.readEmpty(in), equalTo(Source.EMPTY));
+                assertThat(in.readInt(), equalTo(test));
+            }
+        }
+    }
 }
diff --git a/...ugin/esql/src/main/java/org/elasticsearch/xpack/esql/querydsl/query/SingleValueQuery.java b/...ugin/esql/src/main/java/org/elasticsearch/xpack/esql/querydsl/query/SingleValueQuery.java
@@ -107,7 +107,15 @@ public static class Builder extends AbstractQueryBuilder<Builder> {
             this.next = in.readNamedWriteable(QueryBuilder.class);
             this.field = in.readString();
             if (in.getTransportVersion().onOrAfter(TransportVersions.ESQL_SINGLE_VALUE_QUERY_SOURCE)) {
-                this.source = Source.readFrom((PlanStreamInput) in);
+                if (in instanceof PlanStreamInput psi) {
+                    this.source = Source.readFrom(psi);
+                } else {
+                    /*
+                     * For things like CanMatch we serialize without the Source. But we
+                     * don't use it, so that's ok.
+                     */
+                    this.source = Source.readEmpty(in);
+                }
             } else if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_12_0)) {
                 this.source = readOldSource(in);
             } else {