Use built-in batch grouping transform

Fixes GoogleCloudPlatform#1416
fbiville · Apr 8, 2024 · 6409955 · 6409955
1 parent 980e7ab
commit 6409955
Show file tree

Hide file tree

Showing 5 changed files with 2 additions and 247 deletions.
diff --git a/.../src/main/java/com/google/cloud/teleport/v2/neo4j/transforms/Neo4jRowWriterTransform.java b/.../src/main/java/com/google/cloud/teleport/v2/neo4j/transforms/Neo4jRowWriterTransform.java
@@ -31,7 +31,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import java.util.Map;
 import java.util.Set;
-import org.apache.beam.sdk.transforms.GroupByKey;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.SerializableFunction;
@@ -104,8 +104,7 @@ public PCollection<Row> expand(@NonNull PCollection<Row> input) {
 
     return input
         .apply("Create KV pairs", CreateKvTransform.of(parallelism))
-        .apply("Group by keys", GroupByKey.create())
-        .apply("Split into batches", ParDo.of(SplitIntoBatches.of(batchSize)))
+        .apply("Group into batches", GroupIntoBatches.ofSize(batchSize))
         .apply(target.getSequence() + ": Neo4j write " + target.getName(), ParDo.of(neo4jUnwindFn))
         .setRowSchema(input.getSchema());
   }

diff --git a/...oud-to-neo4j/src/main/java/com/google/cloud/teleport/v2/neo4j/transforms/Partitioner.java b/...oud-to-neo4j/src/main/java/com/google/cloud/teleport/v2/neo4j/transforms/Partitioner.java
diff --git a/...o-neo4j/src/main/java/com/google/cloud/teleport/v2/neo4j/transforms/SplitIntoBatches.java b/...o-neo4j/src/main/java/com/google/cloud/teleport/v2/neo4j/transforms/SplitIntoBatches.java
diff --git a/...to-neo4j/src/test/java/com/google/cloud/teleport/v2/neo4j/transforms/PartitionerTest.java b/...to-neo4j/src/test/java/com/google/cloud/teleport/v2/neo4j/transforms/PartitionerTest.java
diff --git a/...o4j/src/test/java/com/google/cloud/teleport/v2/neo4j/transforms/SplitIntoBatchesTest.java b/...o4j/src/test/java/com/google/cloud/teleport/v2/neo4j/transforms/SplitIntoBatchesTest.java