Simplify merging enrich output (#107018)

The merge logic in MergePositionsOperator is excessively complex and lacks flexibility. It relies on the source operator emitting pages with ascending positions. Additionally, this merge logic introduced an unusual method, appendAllValuesToCurrentPosition, to the Block.Builder. We should replace this with a simpler and more flexible approach. This PR uses a mechanism similar to the grouping aggregation. In fact, it is very close to the values aggregation. Initially, I considered using the GroupingState from ValuesAggregator. However, unlike in the values aggregation, we don't expect many multi-values in enrich. Hence, I introduced the new EnrichResultBuilders instead.
elastic · Apr 9, 2024 · 24aed5c · 24aed5c
1 parent aba7566
commit 24aed5c
Show file tree

Hide file tree

Showing 30 changed files with 910 additions and 710 deletions.
diff --git a/x-pack/plugin/esql/build.gradle b/x-pack/plugin/esql/build.gradle
@@ -2,6 +2,7 @@ import org.elasticsearch.gradle.internal.info.BuildParams
 
 apply plugin: 'elasticsearch.internal-es-plugin'
 apply plugin: 'elasticsearch.internal-cluster-test'
+apply plugin: 'elasticsearch.string-templates'
 esplugin {
   name 'x-pack-esql'
   description 'The plugin that powers ESQL for Elasticsearch'
@@ -222,3 +223,59 @@ tasks.register("regen") {
     }
   }
 }
+
+tasks.named("spotlessJava") { dependsOn stringTemplates }
+tasks.named('checkstyleMain').configure {
+  excludes = [ "**/*.java.st" ]
+}
+
+def prop(Type, type, TYPE, BYTES, Array) {
+  return [
+    "Type" : Type,
+    "type" : type,
+    "TYPE" : TYPE,
+    "BYTES" : BYTES,
+    "Array" : Array,
+
+    "int" : type == "int" ? "true" : "",
+    "long" : type == "long" ? "true" : "",
+    "double" : type == "double" ? "true" : "",
+    "BytesRef" : type == "BytesRef" ? "true" : "",
+    "boolean" : type == "boolean" ? "true" : "",
+  ]
+}
+
+tasks.named('stringTemplates').configure {
+  var intProperties = prop("Int", "int", "INT", "Integer.BYTES", "IntArray")
+  var longProperties = prop("Long", "long", "LONG", "Long.BYTES", "LongArray")
+  var doubleProperties = prop("Double", "double", "DOUBLE", "Double.BYTES", "DoubleArray")
+  var bytesRefProperties = prop("BytesRef", "BytesRef", "BYTES_REF", "org.apache.lucene.util.RamUsageEstimator.NUM_BYTES_OBJECT_REF", "")
+  var booleanProperties = prop("Boolean", "boolean", "BOOLEAN", "Byte.BYTES", "BitArray")
+  // enrich
+  File enrichResultBuilderInput = file("src/main/java/org/elasticsearch/xpack/esql/enrich/X-EnrichResultBuilder.java.st")
+  template {
+    it.properties = intProperties
+    it.inputFile = enrichResultBuilderInput
+    it.outputFile = "org/elasticsearch/xpack/esql/enrich/EnrichResultBuilderForInt.java"
+  }
+  template {
+    it.properties = longProperties
+    it.inputFile = enrichResultBuilderInput
+    it.outputFile = "org/elasticsearch/xpack/esql/enrich/EnrichResultBuilderForLong.java"
+  }
+  template {
+    it.properties = doubleProperties
+    it.inputFile = enrichResultBuilderInput
+    it.outputFile = "org/elasticsearch/xpack/esql/enrich/EnrichResultBuilderForDouble.java"
+  }
+  template {
+    it.properties = bytesRefProperties
+    it.inputFile = enrichResultBuilderInput
+    it.outputFile = "org/elasticsearch/xpack/esql/enrich/EnrichResultBuilderForBytesRef.java"
+  }
+  template {
+    it.properties = booleanProperties
+    it.inputFile = enrichResultBuilderInput
+    it.outputFile = "org/elasticsearch/xpack/esql/enrich/EnrichResultBuilderForBoolean.java"
+  }
+}
diff --git a/...ugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlock.java b/...ugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlock.java
@@ -223,19 +223,6 @@ sealed interface Builder extends Block.Builder, BlockLoader.BooleanBuilder permi
         @Override
         Builder mvOrdering(Block.MvOrdering mvOrdering);
 
-        /**
-         * Appends the all values of the given block into a the current position
-         * in this builder.
-         */
-        @Override
-        Builder appendAllValuesToCurrentPosition(Block block);
-
-        /**
-         * Appends the all values of the given block into a the current position
-         * in this builder.
-         */
-        Builder appendAllValuesToCurrentPosition(BooleanBlock block);
-
         @Override
         BooleanBlock build();
     }

diff --git a/...ql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java b/...ql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java
@@ -71,55 +71,6 @@ public BooleanBlockBuilder endPositionEntry() {
         return this;
     }
 
-    /**
-     * Appends the all values of the given block into a the current position
-     * in this builder.
-     */
-    @Override
-    public BooleanBlockBuilder appendAllValuesToCurrentPosition(Block block) {
-        if (block.areAllValuesNull()) {
-            return appendNull();
-        }
-        return appendAllValuesToCurrentPosition((BooleanBlock) block);
-    }
-
-    /**
-     * Appends the all values of the given block into a the current position
-     * in this builder.
-     */
-    @Override
-    public BooleanBlockBuilder appendAllValuesToCurrentPosition(BooleanBlock block) {
-        final int positionCount = block.getPositionCount();
-        if (positionCount == 0) {
-            return appendNull();
-        }
-        final int totalValueCount = block.getTotalValueCount();
-        if (totalValueCount == 0) {
-            return appendNull();
-        }
-        if (totalValueCount > 1) {
-            beginPositionEntry();
-        }
-        final BooleanVector vector = block.asVector();
-        if (vector != null) {
-            for (int p = 0; p < positionCount; p++) {
-                appendBoolean(vector.getBoolean(p));
-            }
-        } else {
-            for (int p = 0; p < positionCount; p++) {
-                int count = block.getValueCount(p);
-                int i = block.getFirstValueIndex(p);
-                for (int v = 0; v < count; v++) {
-                    appendBoolean(block.getBoolean(i++));
-                }
-            }
-        }
-        if (totalValueCount > 1) {
-            endPositionEntry();
-        }
-        return this;
-    }
-
     @Override
     public BooleanBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusive) {
         if (block.areAllValuesNull()) {

diff --git a/...gin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlock.java b/...gin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlock.java
@@ -229,19 +229,6 @@ sealed interface Builder extends Block.Builder, BlockLoader.BytesRefBuilder perm
         @Override
         Builder mvOrdering(Block.MvOrdering mvOrdering);
 
-        /**
-         * Appends the all values of the given block into a the current position
-         * in this builder.
-         */
-        @Override
-        Builder appendAllValuesToCurrentPosition(Block block);
-
-        /**
-         * Appends the all values of the given block into a the current position
-         * in this builder.
-         */
-        Builder appendAllValuesToCurrentPosition(BytesRefBlock block);
-
         @Override
         BytesRefBlock build();
     }

diff --git a/...l/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java b/...l/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java
@@ -78,56 +78,6 @@ protected void writeNullValue() {
         values.append(BytesRefBlock.NULL_VALUE);
     }
 
-    /**
-     * Appends the all values of the given block into a the current position
-     * in this builder.
-     */
-    @Override
-    public BytesRefBlockBuilder appendAllValuesToCurrentPosition(Block block) {
-        if (block.areAllValuesNull()) {
-            return appendNull();
-        }
-        return appendAllValuesToCurrentPosition((BytesRefBlock) block);
-    }
-
-    /**
-     * Appends the all values of the given block into a the current position
-     * in this builder.
-     */
-    @Override
-    public BytesRefBlockBuilder appendAllValuesToCurrentPosition(BytesRefBlock block) {
-        final int positionCount = block.getPositionCount();
-        if (positionCount == 0) {
-            return appendNull();
-        }
-        final int totalValueCount = block.getTotalValueCount();
-        if (totalValueCount == 0) {
-            return appendNull();
-        }
-        if (totalValueCount > 1) {
-            beginPositionEntry();
-        }
-        BytesRef scratch = new BytesRef();
-        final BytesRefVector vector = block.asVector();
-        if (vector != null) {
-            for (int p = 0; p < positionCount; p++) {
-                appendBytesRef(vector.getBytesRef(p, scratch));
-            }
-        } else {
-            for (int p = 0; p < positionCount; p++) {
-                int count = block.getValueCount(p);
-                int i = block.getFirstValueIndex(p);
-                for (int v = 0; v < count; v++) {
-                    appendBytesRef(block.getBytesRef(i++, scratch));
-                }
-            }
-        }
-        if (totalValueCount > 1) {
-            endPositionEntry();
-        }
-        return this;
-    }
-
     @Override
     public BytesRefBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusive) {
         if (block.areAllValuesNull()) {

diff --git a/...lugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlock.java b/...lugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlock.java
@@ -224,19 +224,6 @@ sealed interface Builder extends Block.Builder, BlockLoader.DoubleBuilder permit
         @Override
         Builder mvOrdering(Block.MvOrdering mvOrdering);
 
-        /**
-         * Appends the all values of the given block into a the current position
-         * in this builder.
-         */
-        @Override
-        Builder appendAllValuesToCurrentPosition(Block block);
-
-        /**
-         * Appends the all values of the given block into a the current position
-         * in this builder.
-         */
-        Builder appendAllValuesToCurrentPosition(DoubleBlock block);
-
         @Override
         DoubleBlock build();
     }

diff --git a/...sql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java b/...sql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java
@@ -71,55 +71,6 @@ public DoubleBlockBuilder endPositionEntry() {
         return this;
     }
 
-    /**
-     * Appends the all values of the given block into a the current position
-     * in this builder.
-     */
-    @Override
-    public DoubleBlockBuilder appendAllValuesToCurrentPosition(Block block) {
-        if (block.areAllValuesNull()) {
-            return appendNull();
-        }
-        return appendAllValuesToCurrentPosition((DoubleBlock) block);
-    }
-
-    /**
-     * Appends the all values of the given block into a the current position
-     * in this builder.
-     */
-    @Override
-    public DoubleBlockBuilder appendAllValuesToCurrentPosition(DoubleBlock block) {
-        final int positionCount = block.getPositionCount();
-        if (positionCount == 0) {
-            return appendNull();
-        }
-        final int totalValueCount = block.getTotalValueCount();
-        if (totalValueCount == 0) {
-            return appendNull();
-        }
-        if (totalValueCount > 1) {
-            beginPositionEntry();
-        }
-        final DoubleVector vector = block.asVector();
-        if (vector != null) {
-            for (int p = 0; p < positionCount; p++) {
-                appendDouble(vector.getDouble(p));
-            }
-        } else {
-            for (int p = 0; p < positionCount; p++) {
-                int count = block.getValueCount(p);
-                int i = block.getFirstValueIndex(p);
-                for (int v = 0; v < count; v++) {
-                    appendDouble(block.getDouble(i++));
-                }
-            }
-        }
-        if (totalValueCount > 1) {
-            endPositionEntry();
-        }
-        return this;
-    }
-
     @Override
     public DoubleBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusive) {
         if (block.areAllValuesNull()) {

diff --git a/...k/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlock.java b/...k/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlock.java
@@ -223,19 +223,6 @@ sealed interface Builder extends Block.Builder, BlockLoader.IntBuilder permits I
         @Override
         Builder mvOrdering(Block.MvOrdering mvOrdering);
 
-        /**
-         * Appends the all values of the given block into a the current position
-         * in this builder.
-         */
-        @Override
-        Builder appendAllValuesToCurrentPosition(Block block);
-
-        /**
-         * Appends the all values of the given block into a the current position
-         * in this builder.
-         */
-        Builder appendAllValuesToCurrentPosition(IntBlock block);
-
         @Override
         IntBlock build();
     }

diff --git a/...n/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java b/...n/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java
@@ -71,55 +71,6 @@ public IntBlockBuilder endPositionEntry() {
         return this;
     }
 
-    /**
-     * Appends the all values of the given block into a the current position
-     * in this builder.
-     */
-    @Override
-    public IntBlockBuilder appendAllValuesToCurrentPosition(Block block) {
-        if (block.areAllValuesNull()) {
-            return appendNull();
-        }
-        return appendAllValuesToCurrentPosition((IntBlock) block);
-    }
-
-    /**
-     * Appends the all values of the given block into a the current position
-     * in this builder.
-     */
-    @Override
-    public IntBlockBuilder appendAllValuesToCurrentPosition(IntBlock block) {
-        final int positionCount = block.getPositionCount();
-        if (positionCount == 0) {
-            return appendNull();
-        }
-        final int totalValueCount = block.getTotalValueCount();
-        if (totalValueCount == 0) {
-            return appendNull();
-        }
-        if (totalValueCount > 1) {
-            beginPositionEntry();
-        }
-        final IntVector vector = block.asVector();
-        if (vector != null) {
-            for (int p = 0; p < positionCount; p++) {
-                appendInt(vector.getInt(p));
-            }
-        } else {
-            for (int p = 0; p < positionCount; p++) {
-                int count = block.getValueCount(p);
-                int i = block.getFirstValueIndex(p);
-                for (int v = 0; v < count; v++) {
-                    appendInt(block.getInt(i++));
-                }
-            }
-        }
-        if (totalValueCount > 1) {
-            endPositionEntry();
-        }
-        return this;
-    }
-
     @Override
     public IntBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusive) {
         if (block.areAllValuesNull()) {

diff --git a/.../plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlock.java b/.../plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlock.java
@@ -224,19 +224,6 @@ sealed interface Builder extends Block.Builder, BlockLoader.LongBuilder permits
         @Override
         Builder mvOrdering(Block.MvOrdering mvOrdering);
 
-        /**
-         * Appends the all values of the given block into a the current position
-         * in this builder.
-         */
-        @Override
-        Builder appendAllValuesToCurrentPosition(Block block);
-
-        /**
-         * Appends the all values of the given block into a the current position
-         * in this builder.
-         */
-        Builder appendAllValuesToCurrentPosition(LongBlock block);
-
         @Override
         LongBlock build();
     }