Compact serialization of blocks

Author: dnhatn

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

@@ -184,6 +184,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion INCLUDE_INDEX_MODE_IN_GET_DATA_STREAM = def(9_023_0_00);
     public static final TransportVersion MAX_OPERATION_SIZE_REJECTIONS_ADDED = def(9_024_0_00);
     public static final TransportVersion RETRY_ILM_ASYNC_ACTION_REQUIRE_ERROR = def(9_025_0_00);
+    public static final TransportVersion ESQL_SERIALIZE_BLOCK_TYPE_CODE = def(9_026_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlock.java

@@ -9,14 +9,17 @@
 
 import org.apache.lucene.util.Accountable;
 import org.apache.lucene.util.RamUsageEstimator;
-import org.elasticsearch.common.io.stream.NamedWriteable;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.ReleasableIterator;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.index.mapper.BlockLoader;
 
+import java.io.IOException;
+
 /**
  * A Block is a columnar representation of homogenous data. It has a position (row) count, and
  * various data retrieval methods for accessing the underlying data that is stored at a given
@@ -35,7 +38,7 @@
  * <p> Block are immutable and can be passed between threads as long as no two threads hold a reference to
  * the same block at the same time.
  */
-public interface Block extends Accountable, BlockLoader.Block, NamedWriteable, RefCounted, Releasable {
+public interface Block extends Accountable, BlockLoader.Block, Writeable, RefCounted, Releasable {
     /**
      * The maximum number of values that can be added to one position via lookup.
      * TODO maybe make this everywhere?
@@ -329,6 +332,30 @@ static Block[] buildAll(Block.Builder... builders) {
         }
     }
 
+    /**
+     * Writes only the data of the block to a stream output.
+     * This method should be used when the type of the block is known during reading.
+     */
+    void writeTo(StreamOutput out) throws IOException;
+
+    /**
+     * Writes the type of the block followed by the block data to a stream output.
+     * This should be paired with {@link #readTypedBlock(BlockStreamInput)}
+     */
+    static void writeTypedBlock(Block block, StreamOutput out) throws IOException {
+        block.elementType().writeTo(out);
+        block.writeTo(out);
+    }
+
+    /**
+     * Reads the block type and then the block data from a stream input
+     * This should be paired with {@link #writeTypedBlock(Block, StreamOutput)}
+     */
+    static Block readTypedBlock(BlockStreamInput in) throws IOException {
+        ElementType elementType = ElementType.readFrom(in);
+        return elementType.reader.readBlock(in);
+    }
+
     /**
      * Serialization type for blocks: 0 and 1 replace false/true used in pre-8.14
      */

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlock.java

@@ -8,7 +8,6 @@
 package org.elasticsearch.compute.data;
 
 import org.apache.lucene.util.Accountable;
-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.unit.ByteSizeValue;
@@ -42,8 +41,6 @@ public CompositeBlock(Block[] blocks) {
         }
     }
 
-    static NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Block.class, "CompositeBlock", CompositeBlock::readFrom);
-
     @Override
     public Vector asVector() {
         return null;
@@ -208,18 +205,14 @@ public long ramBytesUsed() {
         return Arrays.stream(blocks).mapToLong(Accountable::ramBytesUsed).sum();
     }
 
-    @Override
-    public String getWriteableName() {
-        return "CompositeBlock";
-    }
-
     static Block readFrom(StreamInput in) throws IOException {
         final int numBlocks = in.readVInt();
         boolean success = false;
         final Block[] blocks = new Block[numBlocks];
+        BlockStreamInput blockStreamInput = (BlockStreamInput) in;
         try {
             for (int b = 0; b < numBlocks; b++) {
-                blocks[b] = in.readNamedWriteable(Block.class);
+                blocks[b] = Block.readTypedBlock(blockStreamInput);
             }
             CompositeBlock result = new CompositeBlock(blocks);
             success = true;
@@ -235,7 +228,7 @@ static Block readFrom(StreamInput in) throws IOException {
     public void writeTo(StreamOutput out) throws IOException {
         out.writeVInt(blocks.length);
         for (Block block : blocks) {
-            out.writeNamedWriteable(block);
+            Block.writeTypedBlock(block, out);
         }
     }