[BEAM-12093] Overhaul ElasticsearchIO.Write (#14347)

Presently, the Write transform has 2 responsibilities: -Convert input documents into Bulk API entities, serializing based on user settings (partial update, delete, upsert, etc) -> DocToBulk -Batch the converted Bulk API entities together and interface with the target ES cluster -> BulkIO This change aims to separate these 2 responsibilities into discrete PTransforms to allow for greater flexibility while also maintaining the convenience of the Write transform to perform both document conversion and IO serially.
apache · May 27, 2021 · 4f1f1c1 · 4f1f1c1
1 parent 4cc279c
commit 4f1f1c1
Show file tree

Hide file tree

Showing 16 changed files with 1,691 additions and 324 deletions.
diff --git a/CHANGES.md b/CHANGES.md
@@ -93,6 +93,8 @@
 ## I/Os
 
 * Support for X source added (Java/Python) ([BEAM-X](https://issues.apache.org/jira/browse/BEAM-X)).
+* Allow splitting apart document serialization and IO for ElasticsearchIO
+* Support Bulk API request size optimization through addition of ElasticsearchIO.Write.withStatefulBatches
 
 ## New Features / Improvements
 

diff --git a/...csearch-tests-2/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java b/...csearch-tests-2/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java
@@ -103,6 +103,15 @@ public void testWriteVolume() throws Exception {
     elasticsearchIOTestCommonWrite.testWrite();
   }
 
+  @Test
+  public void testWriteVolumeStateful() throws Exception {
+    // cannot share elasticsearchIOTestCommon because tests run in parallel.
+    ElasticsearchIOTestCommon elasticsearchIOTestCommonWrite =
+        new ElasticsearchIOTestCommon(writeConnectionConfiguration, restClient, true);
+    elasticsearchIOTestCommonWrite.setPipeline(pipeline);
+    elasticsearchIOTestCommonWrite.testWriteStateful();
+  }
+
   @Test
   public void testSizesVolume() throws Exception {
     elasticsearchIOTestCommon.testSizes();
@@ -123,6 +132,23 @@ public void testWriteVolumeWithFullAddressing() throws Exception {
     elasticsearchIOTestCommonWrite.testWriteWithFullAddressing();
   }
 
+  @Test
+  public void testWriteWithAllowableErrors() throws Exception {
+    elasticsearchIOTestCommon.testWriteWithAllowedErrors();
+  }
+
+  @Test
+  public void testWriteWithRouting() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteWithRouting();
+  }
+
+  @Test
+  public void testWriteWithDocVersion() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteWithDocVersion();
+  }
+
   /**
    * This test verifies volume partial updates of Elasticsearch. The test dataset index is cloned
    * and then a new field is added to each document using a partial update. The test then asserts

diff --git a/...earch-tests-2/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java b/...earch-tests-2/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java
@@ -201,6 +201,29 @@ public void testWritePartialUpdate() throws Exception {
     elasticsearchIOTestCommon.testWritePartialUpdate();
   }
 
+  @Test
+  public void testWriteWithAllowedErrors() throws Exception {
+    elasticsearchIOTestCommon.testWriteWithAllowedErrors();
+  }
+
+  @Test
+  public void testWriteWithRouting() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteWithRouting();
+  }
+
+  @Test
+  public void testWriteWithDocVersion() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteWithDocVersion();
+  }
+
+  @Test
+  public void testMaxParallelRequestsPerWindow() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testMaxParallelRequestsPerWindow();
+  }
+
   @Test
   public void testReadWithMetadata() throws Exception {
     elasticsearchIOTestCommon.setPipeline(pipeline);

diff --git a/sdks/java/io/elasticsearch-tests/elasticsearch-tests-5/build.gradle b/sdks/java/io/elasticsearch-tests/elasticsearch-tests-5/build.gradle
@@ -48,6 +48,8 @@ dependencies {
   testCompile project(path: ":sdks:java:io:elasticsearch-tests:elasticsearch-tests-common", configuration: "testRuntime")
   testCompile "org.elasticsearch.test:framework:$elastic_search_version"
   testCompile "org.elasticsearch.plugin:transport-netty4-client:$elastic_search_version"
+  // TODO(egalpin): Remove painless plugin in favour of containerized tests
+  testImplementation group: 'org.codelibs.elasticsearch.module', name: 'lang-painless', version: "$elastic_search_version"
   testCompile "com.carrotsearch.randomizedtesting:randomizedtesting-runner:2.5.0"
   testCompile "org.elasticsearch:elasticsearch:$elastic_search_version"
 

diff --git a/...csearch-tests-5/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java b/...csearch-tests-5/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java
@@ -107,6 +107,15 @@ public void testWriteVolume() throws Exception {
     elasticsearchIOTestCommonWrite.testWrite();
   }
 
+  @Test
+  public void testWriteVolumeStateful() throws Exception {
+    // cannot share elasticsearchIOTestCommon because tests run in parallel.
+    ElasticsearchIOTestCommon elasticsearchIOTestCommonWrite =
+        new ElasticsearchIOTestCommon(writeConnectionConfiguration, restClient, true);
+    elasticsearchIOTestCommonWrite.setPipeline(pipeline);
+    elasticsearchIOTestCommonWrite.testWriteStateful();
+  }
+
   @Test
   public void testSizesVolume() throws Exception {
     elasticsearchIOTestCommon.testSizes();
@@ -127,6 +136,29 @@ public void testWriteWithFullAddressingVolume() throws Exception {
     elasticsearchIOTestCommonWrite.testWriteWithFullAddressing();
   }
 
+  @Test
+  public void testWriteWithAllowableErrors() throws Exception {
+    elasticsearchIOTestCommon.testWriteWithAllowedErrors();
+  }
+
+  @Test
+  public void testWriteWithRouting() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteWithRouting();
+  }
+
+  @Test
+  public void testWriteScriptedUpsert() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteScriptedUpsert();
+  }
+
+  @Test
+  public void testWriteWithDocVersion() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteWithDocVersion();
+  }
+
   /**
    * This test verifies volume partial updates of Elasticsearch. The test dataset index is cloned
    * and then a new field is added to each document using a partial update. The test then asserts

diff --git a/...earch-tests-5/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java b/...earch-tests-5/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java
@@ -30,6 +30,7 @@
 import java.util.Collection;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.painless.PainlessPlugin;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.test.ESIntegTestCase;
 import org.elasticsearch.transport.Netty4Plugin;
@@ -85,6 +86,7 @@ public Settings indexSettings() {
   protected Collection<Class<? extends Plugin>> nodePlugins() {
     ArrayList<Class<? extends Plugin>> plugins = new ArrayList<>();
     plugins.add(Netty4Plugin.class);
+    plugins.add(PainlessPlugin.class);
     return plugins;
   }
 
@@ -199,6 +201,35 @@ public void testWritePartialUpdate() throws Exception {
     elasticsearchIOTestCommon.testWritePartialUpdate();
   }
 
+  @Test
+  public void testWriteWithAllowableErrors() throws Exception {
+    elasticsearchIOTestCommon.testWriteWithAllowedErrors();
+  }
+
+  @Test
+  public void testWriteWithRouting() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteWithRouting();
+  }
+
+  @Test
+  public void testWriteScriptedUpsert() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteScriptedUpsert();
+  }
+
+  @Test
+  public void testWriteWithDocVersion() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteWithDocVersion();
+  }
+
+  @Test
+  public void testMaxParallelRequestsPerWindow() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testMaxParallelRequestsPerWindow();
+  }
+
   @Test
   public void testReadWithMetadata() throws Exception {
     elasticsearchIOTestCommon.setPipeline(pipeline);

diff --git a/sdks/java/io/elasticsearch-tests/elasticsearch-tests-6/build.gradle b/sdks/java/io/elasticsearch-tests/elasticsearch-tests-6/build.gradle
@@ -48,6 +48,8 @@ dependencies {
   testCompile project(path: ":sdks:java:io:elasticsearch-tests:elasticsearch-tests-common", configuration: "testRuntime")
   testCompile "org.elasticsearch.test:framework:$elastic_search_version"
   testCompile "org.elasticsearch.plugin:transport-netty4-client:$elastic_search_version"
+  // TODO(egalpin): Remove painless plugin in favour of containerized tests
+  testImplementation group: 'org.codelibs.elasticsearch.module', name: 'lang-painless', version: "$elastic_search_version"
   testCompile "com.carrotsearch.randomizedtesting:randomizedtesting-runner:2.5.2"
   testCompile "org.elasticsearch:elasticsearch:$elastic_search_version"
 

diff --git a/...csearch-tests-6/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java b/...csearch-tests-6/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java
@@ -107,6 +107,15 @@ public void testWriteVolume() throws Exception {
     elasticsearchIOTestCommonWrite.testWrite();
   }
 
+  @Test
+  public void testWriteVolumeStateful() throws Exception {
+    // cannot share elasticsearchIOTestCommon because tests run in parallel.
+    ElasticsearchIOTestCommon elasticsearchIOTestCommonWrite =
+        new ElasticsearchIOTestCommon(writeConnectionConfiguration, restClient, true);
+    elasticsearchIOTestCommonWrite.setPipeline(pipeline);
+    elasticsearchIOTestCommonWrite.testWriteStateful();
+  }
+
   @Test
   public void testSizesVolume() throws Exception {
     elasticsearchIOTestCommon.testSizes();
@@ -127,6 +136,29 @@ public void testWriteWithFullAddressingVolume() throws Exception {
     elasticsearchIOTestCommonWrite.testWriteWithFullAddressing();
   }
 
+  @Test
+  public void testWriteWithAllowableErrors() throws Exception {
+    elasticsearchIOTestCommon.testWriteWithAllowedErrors();
+  }
+
+  @Test
+  public void testWriteWithRouting() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteWithRouting();
+  }
+
+  @Test
+  public void testWriteScriptedUpsert() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteScriptedUpsert();
+  }
+
+  @Test
+  public void testWriteWithDocVersion() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteWithDocVersion();
+  }
+
   /**
    * This test verifies volume partial updates of Elasticsearch. The test dataset index is cloned
    * and then a new field is added to each document using a partial update. The test then asserts

diff --git a/...earch-tests-6/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java b/...earch-tests-6/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java
@@ -30,6 +30,7 @@
 import java.util.Collection;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.painless.PainlessPlugin;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.test.ESIntegTestCase;
 import org.elasticsearch.transport.Netty4Plugin;
@@ -85,6 +86,7 @@ public Settings indexSettings() {
   protected Collection<Class<? extends Plugin>> nodePlugins() {
     ArrayList<Class<? extends Plugin>> plugins = new ArrayList<>();
     plugins.add(Netty4Plugin.class);
+    plugins.add(PainlessPlugin.class);
     return plugins;
   }
 
@@ -193,6 +195,35 @@ public void testWritePartialUpdate() throws Exception {
     elasticsearchIOTestCommon.testWritePartialUpdate();
   }
 
+  @Test
+  public void testWriteWithAllowableErrors() throws Exception {
+    elasticsearchIOTestCommon.testWriteWithAllowedErrors();
+  }
+
+  @Test
+  public void testWriteWithRouting() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteWithRouting();
+  }
+
+  @Test
+  public void testWriteScriptedUpsert() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteScriptedUpsert();
+  }
+
+  @Test
+  public void testWriteWithDocVersion() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteWithDocVersion();
+  }
+
+  @Test
+  public void testMaxParallelRequestsPerWindow() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testMaxParallelRequestsPerWindow();
+  }
+
   @Test
   public void testReadWithMetadata() throws Exception {
     elasticsearchIOTestCommon.setPipeline(pipeline);

diff --git a/sdks/java/io/elasticsearch-tests/elasticsearch-tests-7/build.gradle b/sdks/java/io/elasticsearch-tests/elasticsearch-tests-7/build.gradle
@@ -48,6 +48,8 @@ dependencies {
   testCompile project(path: ":sdks:java:io:elasticsearch-tests:elasticsearch-tests-common", configuration: "testRuntime")
   testCompile "org.elasticsearch.test:framework:$elastic_search_version"
   testCompile "org.elasticsearch.plugin:transport-netty4-client:$elastic_search_version"
+  // TODO(egalpin): Remove painless plugin in favour of containerized tests
+  testImplementation group: 'org.codelibs.elasticsearch.module', name: 'lang-painless', version: "$elastic_search_version"
   testCompile "com.carrotsearch.randomizedtesting:randomizedtesting-runner:2.7.5"
   testCompile "org.elasticsearch:elasticsearch:$elastic_search_version"
 

diff --git a/...csearch-tests-7/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java b/...csearch-tests-7/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java
@@ -107,6 +107,15 @@ public void testWriteVolume() throws Exception {
     elasticsearchIOTestCommonWrite.testWrite();
   }
 
+  @Test
+  public void testWriteVolumeStateful() throws Exception {
+    // cannot share elasticsearchIOTestCommon because tests run in parallel.
+    ElasticsearchIOTestCommon elasticsearchIOTestCommonWrite =
+        new ElasticsearchIOTestCommon(writeConnectionConfiguration, restClient, true);
+    elasticsearchIOTestCommonWrite.setPipeline(pipeline);
+    elasticsearchIOTestCommonWrite.testWriteStateful();
+  }
+
   @Test
   public void testSizesVolume() throws Exception {
     elasticsearchIOTestCommon.testSizes();
@@ -127,6 +136,29 @@ public void testWriteWithFullAddressingVolume() throws Exception {
     elasticsearchIOTestCommonWrite.testWriteWithFullAddressing();
   }
 
+  @Test
+  public void testWriteWithAllowableErrors() throws Exception {
+    elasticsearchIOTestCommon.testWriteWithAllowedErrors();
+  }
+
+  @Test
+  public void testWriteWithRouting() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteWithRouting();
+  }
+
+  @Test
+  public void testWriteScriptedUpsert() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteScriptedUpsert();
+  }
+
+  @Test
+  public void testWriteWithDocVersion() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteWithDocVersion();
+  }
+
   /**
    * This test verifies volume partial updates of Elasticsearch. The test dataset index is cloned
    * and then a new field is added to each document using a partial update. The test then asserts

diff --git a/...earch-tests-7/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java b/...earch-tests-7/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java
@@ -30,6 +30,7 @@
 import java.util.Collection;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.painless.PainlessPlugin;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.test.ESIntegTestCase;
 import org.elasticsearch.transport.Netty4Plugin;
@@ -89,6 +90,7 @@ public Settings indexSettings() {
   protected Collection<Class<? extends Plugin>> nodePlugins() {
     ArrayList<Class<? extends Plugin>> plugins = new ArrayList<>();
     plugins.add(Netty4Plugin.class);
+    plugins.add(PainlessPlugin.class);
     return plugins;
   }
 
@@ -197,6 +199,35 @@ public void testWritePartialUpdate() throws Exception {
     elasticsearchIOTestCommon.testWritePartialUpdate();
   }
 
+  @Test
+  public void testWriteWithAllowableErrors() throws Exception {
+    elasticsearchIOTestCommon.testWriteWithAllowedErrors();
+  }
+
+  @Test
+  public void testWriteWithRouting() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteWithRouting();
+  }
+
+  @Test
+  public void testWriteScriptedUpsert() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteScriptedUpsert();
+  }
+
+  @Test
+  public void testWriteWithDocVersion() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testWriteWithDocVersion();
+  }
+
+  @Test
+  public void testMaxParallelRequestsPerWindow() throws Exception {
+    elasticsearchIOTestCommon.setPipeline(pipeline);
+    elasticsearchIOTestCommon.testMaxParallelRequestsPerWindow();
+  }
+
   @Test
   public void testReadWithMetadata() throws Exception {
     elasticsearchIOTestCommon.setPipeline(pipeline);