Dedupe

Author: nik9000

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/BreakingBytesRefBuilder.java

@@ -12,6 +12,7 @@
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.util.ByteUtils;
 import org.elasticsearch.core.Releasable;
 
 /**
@@ -131,6 +132,12 @@ public void append(BytesRef bytes) {
         append(bytes.bytes, bytes.offset, bytes.length);
     }
 
+    public void append(long l) {
+        grow(bytes.length + Long.BYTES);
+        ByteUtils.writeLongBE(l, bytes.bytes, bytes.length);
+        bytes.length += Long.BYTES;
+    }
+
     /**
      * Set the content of the builder to the given bytes.
      */

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/CollectOperator.java

@@ -14,7 +14,9 @@
 import org.elasticsearch.action.bulk.BulkResponse;
 import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.action.support.SubscribableListener;
+import org.elasticsearch.action.update.UpdateRequest;
 import org.elasticsearch.client.internal.Client;
+import org.elasticsearch.common.hash.MurmurHash3;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.core.Releasables;
@@ -23,18 +25,30 @@
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.util.Base64;
 import java.util.List;
 import java.util.function.Function;
 
 public class CollectOperator implements Operator {
-    public interface Writer {
-        XContentBuilder write(XContentBuilder builder, int valueIndex) throws IOException;
+    public static final long ID_SEED = 0;
+
+    public interface FieldWriter {
+        XContentBuilder write(XContentBuilder builder, int position) throws IOException;
+    }
+
+    public interface IdWriter {
+        void write(BreakingBytesRefBuilder builder, int position);
     }
 
-    public record Factory(Client client, String index, List<Function<Block, Writer>> writers) implements OperatorFactory {
+    public record Factory(
+        Client client,
+        String index,
+        List<Function<Block, FieldWriter>> fieldWriters,
+        List<Function<Page, IdWriter>> idWriters
+    ) implements OperatorFactory {
         @Override
         public CollectOperator get(DriverContext driverContext) {
-            return new CollectOperator(client, driverContext, index, writers);
+            return new CollectOperator(client, driverContext, index, fieldWriters, idWriters);
         }
 
         @Override
@@ -48,7 +62,9 @@ public String describe() {
     private final Client client;
     private final DriverContext driverContext;
     private final String index;
-    private final List<Function<Block, Writer>> writers;
+    private final List<Function<Block, FieldWriter>> fieldWriters;
+    private final List<Function<Page, IdWriter>> idWriters;
+    private final BreakingBytesRefBuilder idBuilder;
 
     private volatile Phase phase = Phase.COLLECTING;
     private volatile IsBlockedResult blocked = NOT_BLOCKED;
@@ -59,11 +75,19 @@ public String describe() {
     private long rowsEmitted;
     private long bulkBytesSent;
 
-    public CollectOperator(Client client, DriverContext driverContext, String index, List<Function<Block, Writer>> writers) {
+    public CollectOperator(
+        Client client,
+        DriverContext driverContext,
+        String index,
+        List<Function<Block, FieldWriter>> fieldWriters,
+        List<Function<Page, IdWriter>> idWriters
+    ) {
         this.client = client;
         this.driverContext = driverContext;
         this.index = index;
-        this.writers = writers;
+        this.fieldWriters = fieldWriters;
+        this.idWriters = idWriters;
+        this.idBuilder = new BreakingBytesRefBuilder(driverContext.breaker(), "id");
     }
 
     @Override
@@ -77,7 +101,6 @@ public void addInput(Page page) {
         checkFailure();
         pagesReceived++;
         rowsReceived += page.getPositionCount();
-
         try {
             BulkRequest request = request(page);
             bulkBytesSent += request.estimatedSizeInBytes();
@@ -96,14 +119,29 @@ private BulkRequest request(Page page) throws IOException {
         }
         for (int b = 0; b < page.getBlockCount(); b++) {
             Block block = page.getBlock(b);
-            Writer writer = writers.get(b).apply(block);
+            FieldWriter writer = fieldWriters.get(b).apply(block);
             for (int p = 0; p < block.getPositionCount(); p++) {
                 writer.write(source[p], p);
             }
         }
+        List<IdWriter> idWriters = this.idWriters.stream().map(w -> w.apply(page)).toList();
         BulkRequest request = new BulkRequest();
+        MurmurHash3.Hash128 hash = new MurmurHash3.Hash128();
+        byte[] hashBytes = hash.getBytes();
         for (int p = 0; p < page.getPositionCount(); p++) {
-            request.add(new IndexRequest(index).source(source[p].endObject()));
+            source[p].endObject();
+            if (idWriters.isEmpty()) {
+                request.add(new IndexRequest(index).source(source[p]));
+            } else {
+                for (IdWriter idWriter : idWriters) {
+                    idWriter.write(idBuilder, p);
+                }
+                MurmurHash3.hash128(idBuilder.bytes(), 0, idBuilder.length(), ID_SEED, hash);
+                hash.getBytes(hashBytes, 0);
+                String id = Base64.getUrlEncoder().encodeToString(hashBytes);
+                request.add(new UpdateRequest().index(index).id(id).doc(source[p]).docAsUpsert(true));
+                idBuilder.clear();
+            }
         }
         return request;
     }