You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by "adityaRaj369 (via GitHub)" <gi...@apache.org> on 2026/06/20 14:34:23 UTC

[PR] Improve WithKeys coder inference context [beam]

adityaRaj369 opened a new pull request, #39043:
URL: https://github.com/apache/beam/pull/39043

   Fixes #18571
   
   This updates WithKeys to build the output KV type descriptor from the explicit key type and the input PCollection type when both are available. That lets coder inference use the actual input coder context instead of only inferring the key coder independently.
   
   The existing fallback behavior, including schema coder fallback, is preserved.
   
   Tested:
   ./gradlew :sdks:java:core:test --tests org.apache.beam.sdk.transforms.WithKeysTest --no-daemon


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Improve WithKeys coder inference context [beam]

Posted by "Abacn (via GitHub)" <gi...@apache.org>.
Abacn commented on PR #39043:
URL: https://github.com/apache/beam/pull/39043#issuecomment-5064077312

   It seems this breaks coder inferences in tests: #35514
   
   ```
   java.lang.IllegalStateException: Unable to return a default Coder for BeamSortRel_28989/WithDummyKey/AddKeys/Map/ParMultiDo(Anonymous).output [PCollection@426526029]. Correct one of the following root causes:
     No Coder has been manually specified;  you may do so using .setCoder().
     Inferring a Coder from the CoderRegistry failed: Cannot provide coder for parameterized type org.apache.beam.sdk.values.KV<java.lang.String, org.apache.beam.sdk.values.Row>: Cannot provide a coder for a Beam Row. Please provide a schema instead using PCollection.setRowSchema.
     Using the default output Coder from the producing PTransform failed: PTransform.getOutputCoder called.
   	at org.apache.beam.sdk.util.Preconditions.checkStateNotNull(Preconditions.java:499)
   	at org.apache.beam.sdk.values.PCollection.getCoder(PCollection.java:282)
   	at org.apache.beam.sdk.values.PCollection.finishSpecifying(PCollection.java:113)
   	at org.apache.beam.sdk.runners.TransformHierarchy.finishSpecifyingInput(TransformHierarchy.java:154)
   	at org.apache.beam.sdk.Pipeline.applyInternal(Pipeline.java:558)
   	at org.apache.beam.sdk.Pipeline.applyTransform(Pipeline.java:507)
   ```


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Improve WithKeys coder inference context [beam]

Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39043:
URL: https://github.com/apache/beam/pull/39043#issuecomment-4855277966

   Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment `assign to next reviewer`:
   
   R: @Abacn for label java.
   
   Available commands:
   - `stop reviewer notifications` - opt out of the automated review tooling
   - `remind me after tests pass` - tag the comment author after tests pass
   - `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Improve WithKeys coder inference context [beam]

Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39043:
URL: https://github.com/apache/beam/pull/39043#issuecomment-5034130614

   Reminder, please take a look at this pr: @kennknowles 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Improve WithKeys coder inference context [beam]

Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39043:
URL: https://github.com/apache/beam/pull/39043#issuecomment-4759617370

   Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment `assign set of reviewers`


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Improve WithKeys coder inference context [beam]

Posted by "adityaRaj369 (via GitHub)" <gi...@apache.org>.
adityaRaj369 commented on PR #39043:
URL: https://github.com/apache/beam/pull/39043#issuecomment-4759767147

   assign set of reviewers
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Improve WithKeys coder inference context [beam]

Posted by "adityaRaj369 (via GitHub)" <gi...@apache.org>.
adityaRaj369 commented on PR #39043:
URL: https://github.com/apache/beam/pull/39043#issuecomment-5036971830

   The PR is approved now, but `beam_PreCommit_Java_IOs_Direct` is still failing while the Java/core-related checks are green.
   
   This PR only changes `WithKeys.java` and `WithKeysTest.java`, so the IOs Direct failure appears unrelated to the current diff. Could a maintainer please rerun the failed check when convenient?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Improve WithKeys coder inference context [beam]

Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39043:
URL: https://github.com/apache/beam/pull/39043#issuecomment-4826127681

   Reminder, please take a look at this pr: @kennknowles 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Improve WithKeys coder inference context [beam]

Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on PR #39043:
URL: https://github.com/apache/beam/pull/39043#issuecomment-4758519140

   ## Summary of Changes
   
   Hello, I'm Gemini Code Assist[^1]! 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 improves the coder inference mechanism within the WithKeys transform. By leveraging the input PCollection's type descriptor alongside the explicit key type, the transform can now more effectively determine the appropriate coder for the resulting KV pairs. This change ensures better type safety and consistency in pipeline construction while preserving backward compatibility for existing use cases.
   
   ### Highlights
   
   * **Improved Coder Inference**: Updated WithKeys to build the output KV type descriptor using both the explicit key type and the input PCollection type, allowing for more accurate coder inference.
   * **Refined Coder Registration**: Modified the expansion logic to use the derived output type descriptor for coder registration, while maintaining support for existing fallback behaviors.
   * **Added Validation Test**: Included a new test case in WithKeysTest to verify that the output type descriptor is correctly set when a key type is provided.
   
   <details>
   <summary><b>New Features</b></summary>
   
   🧠 You can now enable **Memory** (public preview) 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](https://codeassist.google/code-review/login) to enable Memory in your admin console.**
   </details>
   
   <details>
   <summary><b>Using Gemini Code Assist</b></summary>
   <br>
   
   The full guide for Gemini Code Assist can be found on our [documentation page](https://developers.google.com/gemini-code-assist/docs/review-github-code), here are some quick tips.
   
   <b>Invoking Gemini</b>
   
   You can request assistance from Gemini at any point by creating a comment using either `/gemini <command>` or `@gemini-code-assist <command>`. Below is a summary of the supported commands on the current page.
   
   Feature | Command | Description
   --- | --- | ---
   Code Review | `/gemini review` | Performs a code review for the current pull request in its current state.
   Pull Request Summary | `/gemini summary` | Provides a summary of the current pull request in its current state.
   Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in pull request comments and review comments.
   Help | `/gemini help` | Displays a list of available commands.
   
   <b>Customization</b>
   
   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 `.gemini/` folder in the base of the repository. Detailed instructions can be found [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github).
   
   <b>Limitations & Feedback</b>
   
   Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with :thumbsup: and :thumbsdown: 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](https://google.qualtrics.com/jfe/form/SV_2cyuGuTWsEw84yG).
   </details>
   
   
   [^1]: Review the [Privacy Notices](https://policies.google.com/privacy), [Generative AI Prohibited Use Policy](https://policies.google.com/terms/generative-ai/use-policy), [Terms of Service](https://policies.google.com/terms), and learn how to configure Gemini Code Assist in GitHub [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github). Gemini can make mistakes, so double check it and [use code with caution](https://support.google.com/legal/answer/13505487).
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Improve WithKeys coder inference context [beam]

Posted by "adityaRaj369 (via GitHub)" <gi...@apache.org>.
adityaRaj369 commented on code in PR #39043:
URL: https://github.com/apache/beam/pull/39043#discussion_r3446937257


##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/WithKeys.java:
##########
@@ -110,27 +110,32 @@ public WithKeys<K, V> withKeyType(TypeDescriptor<K> keyType) {
 
   @Override
   public PCollection<KV<K, V>> expand(PCollection<V> in) {
+    TypeDescriptor<V> inputType = in.getTypeDescriptor();
+    TypeDescriptor<KV<K, V>> outputType = getOutputTypeDescriptor(inputType);
     PCollection<KV<K, V>> result =
         in.apply(
             "AddKeys",
-            MapElements.via(
-                new SimpleFunction<V, KV<K, V>>() {
-                  @Override
-                  public KV<K, V> apply(V element) {
-                    return KV.of(fn.apply(element), element);
-                  }
-                }));
+            outputType == null
+                ? MapElements.via(
+                    new SimpleFunction<V, KV<K, V>>() {
+                      @Override
+                      public KV<K, V> apply(V element) {
+                        return KV.of(fn.apply(element), element);
+                      }
+                    })
+                : MapElements.into(outputType)
+                    .via(
+                        (SerializableFunction<V, KV<K, V>>)
+                            element -> KV.of(fn.apply(element), element)));

Review Comment:
   Fixed. The mapping function now uses a local copy of fn so it does not capture the outer WithKeys instance. I also applied it to both mapping branches and reran WithKeysTest successfully.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Improve WithKeys coder inference context [beam]

Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39043:
URL: https://github.com/apache/beam/pull/39043#issuecomment-4759770415

   Assigning reviewers:
   
   R: @kennknowles for label java.
   
   
   Note: If you would like to opt out of this review, comment `assign to next reviewer`.
   
   Available commands:
   - `stop reviewer notifications` - opt out of the automated review tooling
   - `remind me after tests pass` - tag the comment author after tests pass
   - `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)
   
   The PR bot will only process comments in the main thread (not review comments).


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Improve WithKeys coder inference context [beam]

Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #39043:
URL: https://github.com/apache/beam/pull/39043#discussion_r3446666769


##########
sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/WithKeys.java:
##########
@@ -110,27 +110,32 @@ public WithKeys<K, V> withKeyType(TypeDescriptor<K> keyType) {
 
   @Override
   public PCollection<KV<K, V>> expand(PCollection<V> in) {
+    TypeDescriptor<V> inputType = in.getTypeDescriptor();
+    TypeDescriptor<KV<K, V>> outputType = getOutputTypeDescriptor(inputType);
     PCollection<KV<K, V>> result =
         in.apply(
             "AddKeys",
-            MapElements.via(
-                new SimpleFunction<V, KV<K, V>>() {
-                  @Override
-                  public KV<K, V> apply(V element) {
-                    return KV.of(fn.apply(element), element);
-                  }
-                }));
+            outputType == null
+                ? MapElements.via(
+                    new SimpleFunction<V, KV<K, V>>() {
+                      @Override
+                      public KV<K, V> apply(V element) {
+                        return KV.of(fn.apply(element), element);
+                      }
+                    })
+                : MapElements.into(outputType)
+                    .via(
+                        (SerializableFunction<V, KV<K, V>>)
+                            element -> KV.of(fn.apply(element), element)));

Review Comment:
   ![high](https://www.gstatic.com/codereviewagent/high-priority.svg)
   
   The lambda expression `element -> KV.of(fn.apply(element), element)` captures the instance field `fn`, which implicitly captures the outer `WithKeys` instance (`this`). In Apache Beam, capturing the `PTransform` instance in a user-defined function/lambda can lead to `NotSerializableException` at runtime if the transform contains non-serializable fields, and it unnecessarily increases the serialized payload size sent to workers.
   
   To avoid capturing `this`, copy `this.fn` to a local variable (e.g., `localFn`) before the lambda and reference that local variable instead.
   
   ```suggestion
       SerializableFunction<V, K> localFn = this.fn;
       TypeDescriptor<V> inputType = in.getTypeDescriptor();
       TypeDescriptor<KV<K, V>> outputType = getOutputTypeDescriptor(inputType);
       PCollection<KV<K, V>> result =
           in.apply(
               "AddKeys",
               outputType == null
                   ? MapElements.via(
                       new SimpleFunction<V, KV<K, V>>() {
                         @Override
                         public KV<K, V> apply(V element) {
                           return KV.of(localFn.apply(element), element);
                         }
                       })
                   : MapElements.into(outputType)
                       .via(
                           (SerializableFunction<V, KV<K, V>>)
                               element -> KV.of(localFn.apply(element), element)));
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Improve WithKeys coder inference context [beam]

Posted by "kennknowles (via GitHub)" <gi...@apache.org>.
kennknowles commented on PR #39043:
URL: https://github.com/apache/beam/pull/39043#issuecomment-5048568369

   Yea, I think the test runs are too old for GitHub to give me the "rerun" button, or something has changed in the GitHub UI. Going against my usual procedure of ensuring fully green, since RRIO is a known flake and seems to be the only issue.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Improve WithKeys coder inference context [beam]

Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39043:
URL: https://github.com/apache/beam/pull/39043#issuecomment-4925483155

   Reminder, please take a look at this pr: @Abacn 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Improve WithKeys coder inference context [beam]

Posted by "adityaRaj369 (via GitHub)" <gi...@apache.org>.
adityaRaj369 commented on PR #39043:
URL: https://github.com/apache/beam/pull/39043#issuecomment-4759769912

   The failing PreCommit Java IOs Direct job appears unrelated to this PR. It failed in :sdks:java:io:rrio:test at CallTest > givenTeardownTimeout_throwsError, while this PR only changes WithKeys and WithKeysTest under sdks/java/core. The focused WithKeysTest passes locally, and the main Java precommit is green.
   
   Could a maintainer please rerun the failed IOs Direct check?
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Improve WithKeys coder inference context [beam]

Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39043:
URL: https://github.com/apache/beam/pull/39043#issuecomment-4969228636

   Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment `assign to next reviewer`:
   
   R: @kennknowles for label java.
   
   Available commands:
   - `stop reviewer notifications` - opt out of the automated review tooling
   - `remind me after tests pass` - tag the comment author after tests pass
   - `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Improve WithKeys coder inference context [beam]

Posted by "kennknowles (via GitHub)" <gi...@apache.org>.
kennknowles merged PR #39043:
URL: https://github.com/apache/beam/pull/39043


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org