Add validation to check for non-ints from partitionfns #35881

hjtran · 2025-08-15T16:22:17Z

A user ran into a confusing scenario where their partition function that returned Trues/Falses ended up "filtering" everything out of their input pcollection.

The Partition transform checks for indexes that are outside of the partition range but bools end up looking like 0s and 1s.

I've added a type check to make sure bools and other non-ints aren't returned from the partitionfn

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

Mention the appropriate issue in your description (for example: 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, comment fixes #<ISSUE NUMBER> instead.
Update CHANGES.md with noteworthy changes.
If this contribution is large, please file an Apache Individual Contributor License Agreement.

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.

gemini-code-assist

Summary of Changes

Hello @hjtran, I'm Gemini Code Assist¹! 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 a crucial validation to the Partition transform in Apache Beam's Python SDK. It addresses a scenario where user-defined partition functions might return non-integer values, such as booleans, leading to unexpected and confusing behavior. The change ensures that partition functions strictly return integer indices, preventing data loss or incorrect partitioning by raising a ValueError for invalid types.

Highlights

Enhanced Type Validation for Partition Functions: The Partition transform's process method now includes a check to ensure that the output of the user-provided partition function is an integer.
Improved Error Messaging: If a partition function yields a bool or any other non-integer type, a ValueError is raised with a clear message indicating the expected integer type.
New Test Case for Boolean Partitions: A dedicated test test_partition_with_bools has been added to core_test.py to specifically verify that passing boolean values to the partition function correctly triggers the new validation error.

Using Gemini Code Assist

The 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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

github-actions · 2025-08-16T12:10:35Z

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

hjtran · 2025-08-16T12:18:46Z

R: @liferoad
The failures look unrelated

github-actions · 2025-08-16T12:19:52Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

liferoad · 2025-08-16T13:46:22Z

/gemini review

liferoad · 2025-08-16T13:46:55Z

R: @liferoad The failures look unrelated

LGTM. Asked Gemini to review in case we could improve the PR.

gemini-code-assist

Code Review

This pull request adds a validation to the Partition transform to ensure the partition function returns an integer, preventing confusing behavior with boolean return values. The implementation is correct, but the condition can be simplified for better readability. The added test case is good, but it could be more specific about the exception type and could be extended to cover other non-integer types like floats to be more comprehensive. I've also corrected a minor grammatical error in the error message.

gemini-code-assist · 2025-08-16T13:47:38Z

sdks/python/apache_beam/transforms/core.py

+      if isinstance(partition, bool) or not isinstance(partition, int):
+        raise ValueError(
+            f"PartitionFn yielded a '{type(partition).__name__}' "
+            "when it should only yields integers")


The condition isinstance(partition, bool) or not isinstance(partition, int) is correct, but it can be simplified. Since bool is a subclass of int, isinstance(True, int) evaluates to True. A more direct way to check if a value is strictly an integer (and not a subclass like bool) is to use type(partition) is not int. This is more readable and achieves the same result. I've also corrected a small grammatical error in the error message ('yields' to 'yield').

Suggested change

if isinstance(partition, bool) or not isinstance(partition, int):

raise ValueError(

f"PartitionFn yielded a '{type(partition).__name__}' "

"when it should only yields integers")

if type(partition) is not int:

raise ValueError(

f"PartitionFn yielded a '{type(partition).__name__}' "

"when it should only yield integers")

I thought about this. I am wondering whether we could do partition = int(partition) although partition is bool, this could make the changes less aggressive.

gemini-code-assist · 2025-08-16T13:47:38Z

sdks/python/apache_beam/transforms/core_test.py

+    with pytest.raises(
+        Exception,
+        match="PartitionFn yielded a 'bool' when it should only yields integers"
+    ):


The test correctly checks for the exception, but it's good practice to be more specific. The code raises a ValueError, so the test should expect a ValueError instead of a generic Exception. This makes the test more robust. Also, there's a small grammatical error in the match string; it should be 'yield' instead of 'yields'.

Suggested change

with pytest.raises(

Exception,

match="PartitionFn yielded a 'bool' when it should only yields integers"

):

with pytest.raises(

ValueError,

match="PartitionFn yielded a 'bool' when it should only yield integers"

):

Please fix this.

I initially checked for ValueError and the test works with the directrunner but the portable runner will cast all exceptions to RuntimeError. (https://github.com/apache/beam/runs/48183396870)

I can narrow this to check for either RuntimeError or ValueError though

gemini-code-assist · 2025-08-16T13:47:39Z

sdks/python/apache_beam/transforms/core_test.py

+  def test_partition_with_bools(self):
+    with pytest.raises(
+        Exception,
+        match="PartitionFn yielded a 'bool' when it should only yields integers"
+    ):
+      with beam.testing.test_pipeline.TestPipeline() as p:
+        _ = (p | beam.Create([True]) | beam.Partition(lambda x, _: x, 2))


This test covers the case where the partition function returns a boolean. To make the tests more comprehensive, it would be beneficial to add another test case for other non-integer types that should be rejected, such as floats. This would ensure the validation logic is robust against various invalid types.

For example, you could add a test like this:

def test_partition_with_floats(self): with pytest.raises( ValueError, match="PartitionFn yielded a 'float' when it should only yield integers" ): with beam.testing.test_pipeline.TestPipeline() as p: _ = (p | beam.Create([1.0]) | beam.Partition(lambda x, _: x, 2))

let us add another type like floats to test.

added floats and a couple of other types

hjtran · 2025-08-16T19:53:51Z

I feel that it'd be better for users to be explicit. Since the documentation states that partitions should be integers between 0 and max partitions, I think it's a little surprising if bools worked (even if they are technically ints, which I actually didn't know until I wrote this PR)

…

On Sat, Aug 16, 2025, 10:13 AM liferoad ***@***.***> wrote: ***@***.**** commented on this pull request. ------------------------------ In sdks/python/apache_beam/transforms/core.py <#35881 (comment)>: > + if isinstance(partition, bool) or not isinstance(partition, int): + raise ValueError( + f"PartitionFn yielded a '{type(partition).__name__}' " + "when it should only yields integers") I thought about this. I am wondering whether we could do partition = int(partition) although partition is bool, this could make the changes less aggressive. — Reply to this email directly, view it on GitHub <#35881 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ACQJSFD5MQKXE3DOJWDFBKD3N4375AVCNFSM6AAAAACD73EIWKVHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZTCMRVHA3TQMZRGU> . You are receiving this because you were mentioned.Message ID: ***@***.***>

liferoad · 2025-08-17T13:33:29Z

I feel that it'd be better for users to be explicit. Since the documentation states that partitions should be integers between 0 and max partitions, I think it's a little surprising if bools worked (even if they are technically ints, which I actually didn't know until I wrote this PR)
…
On Sat, Aug 16, 2025, 10:13 AM liferoad @.> wrote: @.* commented on this pull request. ------------------------------ In sdks/python/apache_beam/transforms/core.py <#35881 (comment)>: > + if isinstance(partition, bool) or not isinstance(partition, int): + raise ValueError( + f"PartitionFn yielded a '{type(partition).name}' " + "when it should only yields integers") I thought about this. I am wondering whether we could do partition = int(partition) although partition is bool, this could make the changes less aggressive. — Reply to this email directly, view it on GitHub <#35881 (comment)>, or unsubscribe https://github.com/notifications/unsubscribe-auth/ACQJSFD5MQKXE3DOJWDFBKD3N4375AVCNFSM6AAAAACD73EIWKVHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZTCMRVHA3TQMZRGU . You are receiving this because you were mentioned.Message ID: @.***>

ok.

ahmedabu98 · 2025-09-02T14:50:40Z

@liferoad @hjtran any further steps needed here?

liferoad · 2025-09-02T14:53:14Z

We can merge this. I thought @hjtran wants to do this by himself since he is a committer now. :)

hjtran · 2025-09-02T14:55:53Z

Whoops, sorry I was away from computer for the past couple of weeks. Thanks for merging!

claudevdm · 2025-09-04T00:24:14Z

This is causing some internal failures e.g. ValueError: PartitionFn yielded a 'int32' when it should only yield integers @liferoad

hjtran · 2025-09-04T01:29:48Z

This is causing some internal failures e.g. ValueError: PartitionFn yielded a 'int32' when it should only yield integers @liferoad

Created #36043 to allow for nump ints too

hjtran added 2 commits August 15, 2025 12:19

add validation to check for non-ints from partitions

200b753

move type check earlier

549a1f4

gemini-code-assist bot reviewed Aug 15, 2025

View reviewed changes

github-actions bot added the python label Aug 15, 2025

hjtran added 2 commits August 15, 2025 14:49

just use regular pipeline

f21e23e

just check for exception in case of portable runner

22b8343

hjtran marked this pull request as ready for review August 16, 2025 11:55

liferoad approved these changes Aug 16, 2025

View reviewed changes

gemini-code-assist bot reviewed Aug 16, 2025

View reviewed changes

hjtran added 2 commits August 20, 2025 08:05

check for valueerror or runtimeerror

622bc06

test a few more types

ebcd037

liferoad merged commit 9f3b160 into apache:master Sep 2, 2025
90 checks passed

hjtran mentioned this pull request Sep 4, 2025

allow numpy integers for partition functions for backwards compatibility #36043

Closed

3 tasks

Add validation to check for non-ints from partitionfns #35881

Add validation to check for non-ints from partitionfns #35881

Uh oh!

Conversation

hjtran commented Aug 15, 2025

GitHub Actions Tests Status (on master branch)

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

github-actions bot commented Aug 16, 2025

Uh oh!

hjtran commented Aug 16, 2025

Uh oh!

github-actions bot commented Aug 16, 2025

Uh oh!

liferoad commented Aug 16, 2025

Uh oh!

liferoad commented Aug 16, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Aug 16, 2025

Choose a reason for hiding this comment

Uh oh!

liferoad Aug 16, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Aug 16, 2025

Choose a reason for hiding this comment

Uh oh!

liferoad Aug 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hjtran Aug 20, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Aug 16, 2025

Choose a reason for hiding this comment

Uh oh!

liferoad Aug 17, 2025

Choose a reason for hiding this comment

Uh oh!

hjtran Aug 20, 2025

Choose a reason for hiding this comment

Uh oh!

hjtran commented Aug 16, 2025 via email

Uh oh!

liferoad commented Aug 17, 2025

Uh oh!

ahmedabu98 commented Sep 2, 2025

Uh oh!

liferoad commented Sep 2, 2025

Uh oh!

Uh oh!

hjtran commented Sep 2, 2025

Uh oh!

claudevdm commented Sep 4, 2025

Uh oh!

hjtran commented Sep 4, 2025

Uh oh!

Uh oh!

liferoad Aug 17, 2025 •

edited

Loading