From d5488ad78dd615fb46f791ae09355c8ed768c022 Mon Sep 17 00:00:00 2001 From: Chris Hegarty <62058229+ChrisHegarty@users.noreply.github.com> Date: Wed, 19 Jun 2024 09:37:15 +0100 Subject: [PATCH] ES|QL enable FloatBlock serialization (#109858) This commit enables ES|QL FloatBlock serialization. --- .../org/elasticsearch/compute/data/Block.java | 1 + .../compute/data/BasicPageTests.java | 20 ++++++++---- .../compute/data/BlockSerializationTests.java | 31 +++++++++++++++++++ 3 files changed, 46 insertions(+), 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 index ca3ce1349c47f..282bc9064b308 100644 --- 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 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 index f76ff0708120b..e0cf277e99967 100644 --- 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/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockSerializationTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockSerializationTests.java index 2daf7755841f7..8ca02b64f01ff 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockSerializationTests.java +++ b/x-pack/plugin/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));