[DO NOT MERGE] Large export with Databoost temp fix - #4223
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request implements temporary performance optimizations for large Spanner exports. By introducing custom partitioning logic for both writing Avro files and reading partitions, the changes aim to improve stability and efficiency when handling massive datasets, avoiding common bottlenecks associated with default shuffle operations. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the Spanner export pipeline by introducing a custom AvroWriter that partitions rows into 128 lanes to write Avro files, and replaces Reshuffle.viaRandomKey() in BatchSpannerRead with a custom 400-bucket grouping and flattening mechanism. Feedback on these changes highlights a potential ArrayIndexOutOfBoundsException in the partitioning logic when dealing with Integer.MIN_VALUE hashes, performance concerns regarding the 128-branch Beam DAG size, and potential serialization issues due to anonymous inner DoFn classes capturing the outer class instance.
| PCollectionList<Struct> partitionedRows = | ||
| input.apply( | ||
| "Partition by table", | ||
| org.apache.beam.sdk.transforms.Partition.of( | ||
| 128, | ||
| new org.apache.beam.sdk.transforms.Partition.PartitionFn<Struct>() { | ||
| @Override | ||
| public int partitionFor(Struct elem, int numPartitions) { | ||
| return Math.abs(elem.getString(0).hashCode()) % numPartitions; | ||
| } | ||
| })); |
There was a problem hiding this comment.
Using Math.abs(hash) % numPartitions can result in a negative partition index if the hash code is Integer.MIN_VALUE (since Math.abs(Integer.MIN_VALUE) is still Integer.MIN_VALUE). If numPartitions is not a power of two, or if the code is modified in the future, this will throw an ArrayIndexOutOfBoundsException.
To prevent this, use Math.floorMod to ensure a non-negative index.
| PCollectionList<Struct> partitionedRows = | |
| input.apply( | |
| "Partition by table", | |
| org.apache.beam.sdk.transforms.Partition.of( | |
| 128, | |
| new org.apache.beam.sdk.transforms.Partition.PartitionFn<Struct>() { | |
| @Override | |
| public int partitionFor(Struct elem, int numPartitions) { | |
| return Math.abs(elem.getString(0).hashCode()) % numPartitions; | |
| } | |
| })); | |
| PCollectionList<Struct> partitionedRows = | |
| input.apply( | |
| "Partition by table", | |
| org.apache.beam.sdk.transforms.Partition.of( | |
| 128, | |
| new org.apache.beam.sdk.transforms.Partition.PartitionFn<Struct>() { | |
| @Override | |
| public int partitionFor(Struct elem, int numPartitions) { | |
| return Math.floorMod(elem.getString(0).hashCode(), numPartitions); | |
| } | |
| })); |
| PCollectionList<KV<String, String>> allResults = PCollectionList.empty(input.getPipeline()); | ||
| for (int i = 0; i < 128; i++) { | ||
| org.apache.beam.sdk.io.WriteFilesResult<String> res = | ||
| partitionedRows | ||
| .get(i) | ||
| .apply( | ||
| "Write lane " + i, | ||
| AvroIO.<Struct>writeCustomTypeToGenericRecords() | ||
| .to( | ||
| new SchemaBasedDynamicDestinations( | ||
| avroSchemas, uniqueIdView, dialectView, baseDir)) | ||
| .withTempDirectory(tempDir) | ||
| .withNoSpilling()); | ||
| allResults = allResults.and(res.getPerDestinationOutputFilenames()); | ||
| } | ||
| return allResults.apply("Flatten filenames", Flatten.pCollections()); |
There was a problem hiding this comment.
Creating 128 separate branches in the Beam DAG (each with its own AvroIO.write() transform) can significantly increase the pipeline graph size. This can lead to slow pipeline submission times, high memory usage on the job submitter, and potential graph size limit issues in runners like Dataflow.
While this is a temporary fix to avoid OOM/file limits by grouping tables into lanes, consider making the number of lanes configurable or using a single AvroIO.write() with optimized runner settings if possible.
| .apply( | ||
| "Assign Fixed Bucket", | ||
| ParDo.of( | ||
| new DoFn< | ||
| PartitionedReadOperation, | ||
| org.apache.beam.sdk.values.KV<Integer, PartitionedReadOperation>>() { | ||
| @ProcessElement | ||
| public void processElement(ProcessContext c) { | ||
| c.output( | ||
| org.apache.beam.sdk.values.KV.of( | ||
| java.util.concurrent.ThreadLocalRandom.current().nextInt(400), | ||
| c.element())); | ||
| } | ||
| })) | ||
| .apply("Group by Fixed Bucket", org.apache.beam.sdk.transforms.GroupByKey.create()) | ||
| .apply( | ||
| "Flatten Fixed Buckets", | ||
| ParDo.of( | ||
| new DoFn< | ||
| org.apache.beam.sdk.values.KV<Integer, Iterable<PartitionedReadOperation>>, | ||
| PartitionedReadOperation>() { | ||
| @ProcessElement | ||
| public void processElement(ProcessContext c) { | ||
| for (PartitionedReadOperation op : c.element().getValue()) { | ||
| c.output(op); | ||
| } | ||
| } | ||
| })) |
There was a problem hiding this comment.
The anonymous inner DoFn classes used here are defined within the instance method expand of BatchSpannerRead. In Java, non-static anonymous inner classes capture a reference to their enclosing outer class instance (BatchSpannerRead). This can lead to serialization overhead or NotSerializableException if the outer class or its fields are not serializable or change in the future.
To avoid capturing the outer instance, you can use MapElements and FlatMapElements with lambda expressions, which are cleaner and completely inline.
.apply(
"Assign Fixed Bucket",
org.apache.beam.sdk.transforms.MapElements.into(
org.apache.beam.sdk.values.TypeDescriptors.kvs(
org.apache.beam.sdk.values.TypeDescriptors.integers(),
org.apache.beam.sdk.values.TypeDescriptor.of(PartitionedReadOperation.class)))
.via(
op ->
org.apache.beam.sdk.values.KV.of(
java.util.concurrent.ThreadLocalRandom.current().nextInt(400),
op)))
.apply("Group by Fixed Bucket", org.apache.beam.sdk.transforms.GroupByKey.create())
.apply(
"Flatten Fixed Buckets",
org.apache.beam.sdk.transforms.FlatMapElements.into(
org.apache.beam.sdk.values.TypeDescriptor.of(PartitionedReadOperation.class))
.via(kv -> kv.getValue()))
Codecov Report❌ Patch coverage is
❌ Your patch check has failed because the patch coverage (38.09%) is below the target coverage (80.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## main #4223 +/- ##
============================================
- Coverage 55.94% 55.93% -0.01%
+ Complexity 7484 7031 -453
============================================
Files 1134 1134
Lines 70175 70210 +35
Branches 8023 8025 +2
============================================
+ Hits 39258 39271 +13
- Misses 28370 28388 +18
- Partials 2547 2551 +4
🚀 New features to boost your workflow:
|
1 participant
No description provided.