BatchElements transform for Java SDK#38369
BatchElements transform for Java SDK#38369ganesh-skumar wants to merge 8 commits intoapache:masterfrom
Conversation
|
Continuing work from my previous PR under a new account due to access issues with @Ganeshsivakumar Original PR: #38280 No major changes, this PR just picks up where the previous one left off. |
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 introduces the BatchElements transform to the Apache Beam Java SDK. This transform is designed to optimize processing performance for operations with significant fixed costs by grouping elements into batches. It features a sophisticated, self-tuning mechanism that dynamically adjusts batch sizes based on observed processing times, allowing it to adapt to varying workloads while respecting user-defined constraints. The implementation handles both global and windowed PCollections, ensuring correct window isolation for batched elements. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. 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 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 counter productive. 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 introduces the BatchElements transform, which dynamically adjusts batch sizes by profiling downstream processing time using linear regression. The implementation provides a flexible configuration and handles both global and windowed data. Several critical issues were identified: a unit mismatch in the batch size estimator (seconds vs. milliseconds), a window contamination bug in the eviction logic of WindowAwareBatchingDoFn, and the need for higher precision timing using System.nanoTime(). Additionally, it is recommended to use the existing Random instance for jitter and to increase the MAX_LIVE_WINDOWS limit to improve performance in complex windowing scenarios.
| target = Math.min(target, (config.targetBatchDurationSecsWithFixedCost - a) / b); | ||
| } | ||
|
|
||
| // 2: b*x = targetDurationSecs | ||
| if (config.targetBatchDurationSecs > 0) { | ||
| target = Math.min(target, config.targetBatchDurationSecs / b); |
There was a problem hiding this comment.
There is a unit mismatch in the batch size calculation. The configuration parameters targetBatchDurationSecs and targetBatchDurationSecsWithFixedCost are defined in seconds, but the regression coefficients a (fixed cost) and b (per-element cost) are derived from System.currentTimeMillis(), meaning they are in milliseconds. This results in the estimator calculating a target batch size that is 1000x smaller than intended. You should convert the target durations to milliseconds before performing the calculation.
There was a problem hiding this comment.
converts both to milliseconds in new commit
| Map.Entry<BoundedWindow, SizedBatch<T>> largest = | ||
| batches.entrySet().stream().max(Comparator.comparingInt(e -> e.getValue().size)).get(); | ||
|
|
||
| try (BatchSizeEstimator.Stopwatch sw = estimator.recordTime(largest.getValue().size)) { |
There was a problem hiding this comment.
This eviction logic causes cross-window data contamination. In Beam, receiver.output() called within processElement automatically associates the output with the current element's window. By evicting and outputting a batch from a different window (largest.getKey()), those elements will be incorrectly re-assigned to the current window. To correctly output elements to a different window, you must do so within finishBundle using the FinishBundleContext, or use a more advanced API that allows specifying the window for each output.
There was a problem hiding this comment.
Fixed. Attaches correct window value before emitting
| this.startTime = System.currentTimeMillis(); | ||
| } | ||
|
|
||
| @Override | ||
| public void close() { | ||
| long elapsed = System.currentTimeMillis() - startTime; |
There was a problem hiding this comment.
System.currentTimeMillis() lacks the precision required for profiling fast operations like element batching, as it often has a resolution of 1ms or more (depending on the OS). For small batches that process quickly, this will frequently result in zero-duration measurements, which negatively impacts the accuracy of the linear regression. It is highly recommended to use System.nanoTime() for more precise micro-benchmarking.
| // add jitter to avoid any single batch size | ||
| int jitter = data.size() % 2; | ||
| if (data.size() > 10) { | ||
| target += (int) (target * config.variance * 2 * (Math.random() - 0.5)); |
There was a problem hiding this comment.
The class already maintains a Random instance. It is better to use random.nextDouble() for consistency and to allow for potential reproducibility in tests if a seed is provided.
| target += (int) (target * config.variance * 2 * (Math.random() - 0.5)); | |
| target += (int) (target * config.variance * 2 * (random.nextDouble() - 0.5)); |
| private Map<BoundedWindow, SizedBatch<T>> batches; | ||
| private int targetBatchSize; | ||
|
|
||
| private static final int MAX_LIVE_WINDOWS = 10; |
There was a problem hiding this comment.
The MAX_LIVE_WINDOWS limit of 10 is extremely low for many Beam use cases, such as pipelines with sliding windows or significant data out-of-orderness. When this limit is reached, the transform will aggressively evict batches, leading to very small batch sizes and poor efficiency. Consider increasing this default significantly (e.g., to 1000) or making it a configurable parameter in BatchConfig.
There was a problem hiding this comment.
Goal is to match python's BatchElements implementation.
|
Hi Jack, could you please review this PR when you have some time. |
|
Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment |
2 participants
Introduces BatchElements transform for Java SDK with dynamic batching
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.
To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md
GitHub Actions Tests Status (on master branch)
See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.