ES|QL enable FloatBlock serialization (elastic#109858)

This commit enables ES|QL FloatBlock serialization.
jimczi · Jun 19, 2024 · d5488ad · d5488ad
1 parent b60d77e
commit d5488ad
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 6 deletions.
diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/Block.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/Block.java
@@ -276,6 +276,7 @@ static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         return List.of(
             IntBlock.ENTRY,
             LongBlock.ENTRY,
+            FloatBlock.ENTRY,
             DoubleBlock.ENTRY,
             BytesRefBlock.ENTRY,
             BooleanBlock.ENTRY,

diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BasicPageTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BasicPageTests.java
@@ -136,14 +136,15 @@ public void testEqualityAndHashCode() throws IOException {
         int blockCount = randomIntBetween(1, 256);
         Block[] blocks = new Block[blockCount];
         for (int blockIndex = 0; blockIndex < blockCount; blockIndex++) {
-            blocks[blockIndex] = switch (randomInt(6)) {
+            blocks[blockIndex] = switch (randomInt(7)) {
                 case 0 -> blockFactory.newIntArrayVector(randomInts(positions).toArray(), positions).asBlock();
                 case 1 -> blockFactory.newLongArrayVector(randomLongs(positions).toArray(), positions).asBlock();
-                case 2 -> blockFactory.newDoubleArrayVector(randomDoubles(positions).toArray(), positions).asBlock();
-                case 3 -> blockFactory.newConstantIntBlockWith(randomInt(), positions);
-                case 4 -> blockFactory.newConstantLongBlockWith(randomLong(), positions);
-                case 5 -> blockFactory.newConstantDoubleBlockWith(randomDouble(), positions);
-                case 6 -> blockFactory.newConstantBytesRefBlockWith(new BytesRef(Integer.toHexString(randomInt())), positions);
+                case 2 -> blockFactory.newFloatArrayVector(randomFloats(positions), positions).asBlock();
+                case 3 -> blockFactory.newDoubleArrayVector(randomDoubles(positions).toArray(), positions).asBlock();
+                case 4 -> blockFactory.newConstantIntBlockWith(randomInt(), positions);
+                case 5 -> blockFactory.newConstantLongBlockWith(randomLong(), positions);
+                case 6 -> blockFactory.newConstantDoubleBlockWith(randomDouble(), positions);
+                case 7 -> blockFactory.newConstantBytesRefBlockWith(new BytesRef(Integer.toHexString(randomInt())), positions);
                 default -> throw new AssertionError();
             };
         }
@@ -184,6 +185,7 @@ public void testPageSerializationSimple() throws IOException {
         Page origPage = new Page(
             blockFactory.newIntArrayVector(IntStream.range(0, 10).toArray(), 10).asBlock(),
             blockFactory.newLongArrayVector(LongStream.range(10, 20).toArray(), 10).asBlock(),
+            blockFactory.newFloatArrayVector(randomFloats(10), 10).asBlock(),
             blockFactory.newDoubleArrayVector(LongStream.range(30, 40).mapToDouble(i -> i).toArray(), 10).asBlock(),
             blockFactory.newBytesRefArrayVector(bytesRefArrayOf("0a", "1b", "2c", "3d", "4e", "5f", "6g", "7h", "8i", "9j"), 10).asBlock(),
             blockFactory.newConstantIntBlockWith(randomInt(), 10),
@@ -248,4 +250,10 @@ BytesRefArray bytesRefArrayOf(String... values) {
         Arrays.stream(values).map(BytesRef::new).forEach(array::append);
         return array;
     }
+
+    float[] randomFloats(int size) {
+        float[] fa = new float[size];
+        IntStream.range(0, size).forEach(i -> fa[i] = randomFloat());
+        return fa;
+    }
 }
diff --git a/...in/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockSerializationTests.java b/...in/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockSerializationTests.java
@@ -39,6 +39,10 @@ public void testConstantLongBlockLong() throws IOException {
         assertConstantBlockImpl(blockFactory.newConstantLongBlockWith(randomLong(), randomIntBetween(1, 8192)));
     }
 
+    public void testConstantFloatBlock() throws IOException {
+        assertConstantBlockImpl(blockFactory.newConstantFloatBlockWith(randomFloat(), randomIntBetween(1, 8192)));
+    }
+
     public void testConstantDoubleBlock() throws IOException {
         assertConstantBlockImpl(blockFactory.newConstantDoubleBlockWith(randomDouble(), randomIntBetween(1, 8192)));
     }
@@ -81,6 +85,17 @@ public void testEmptyLongBlock() throws IOException {
         }
     }
 
+    public void testEmptyFloatBlock() throws IOException {
+        assertEmptyBlock(blockFactory.newFloatBlockBuilder(0).build());
+        try (FloatBlock toFilter = blockFactory.newFloatBlockBuilder(0).appendNull().build()) {
+            assertEmptyBlock(toFilter.filter());
+        }
+        assertEmptyBlock(blockFactory.newFloatVectorBuilder(0).build().asBlock());
+        try (FloatVector toFilter = blockFactory.newFloatVectorBuilder(0).appendFloat(randomFloat()).build()) {
+            assertEmptyBlock(toFilter.filter().asBlock());
+        }
+    }
+
     public void testEmptyDoubleBlock() throws IOException {
         assertEmptyBlock(blockFactory.newDoubleBlockBuilder(0).build());
         try (DoubleBlock toFilter = blockFactory.newDoubleBlockBuilder(0).appendNull().build()) {
@@ -140,6 +155,22 @@ public void testFilterLongBlock() throws IOException {
         }
     }
 
+    public void testFilterFloatBlock() throws IOException {
+        try (FloatBlock toFilter = blockFactory.newFloatBlockBuilder(0).appendFloat(1).appendFloat(2).build()) {
+            assertFilterBlock(toFilter.filter(1));
+        }
+        try (FloatBlock toFilter = blockFactory.newFloatBlockBuilder(1).appendFloat(randomFloat()).appendNull().build()) {
+            assertFilterBlock(toFilter.filter(0));
+        }
+        try (FloatVector toFilter = blockFactory.newFloatVectorBuilder(1).appendFloat(randomFloat()).build()) {
+            assertFilterBlock(toFilter.filter(0).asBlock());
+
+        }
+        try (FloatVector toFilter = blockFactory.newFloatVectorBuilder(1).appendFloat(randomFloat()).appendFloat(randomFloat()).build()) {
+            assertFilterBlock(toFilter.filter(0).asBlock());
+        }
+    }
+
     public void testFilterDoubleBlock() throws IOException {
         try (DoubleBlock toFilter = blockFactory.newDoubleBlockBuilder(0).appendDouble(1).appendDouble(2).build()) {
             assertFilterBlock(toFilter.filter(1));