This is an automated email from the ASF dual-hosted git repository.
xqhu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new b60ce5be96b Add type checking to confirm that the flatten inputs are
actually pco… (#35874)
b60ce5be96b is described below
commit b60ce5be96b2ea97084c1330430c8556ca59a8d9
Author: Hai Joey Tran <[email protected]>
AuthorDate: Fri Aug 15 19:32:47 2025 -0400
Add type checking to confirm that the flatten inputs are actually pco…
(#35874)
* Add type checking to confirm that the flatten inputs are actually
pcollections
* Allow iterables of non-pcollections to flatten again
* Update sdks/python/apache_beam/transforms/core.py
Co-authored-by: gemini-code-assist[bot]
<176961590+gemini-code-assist[bot]@users.noreply.github.com>
---------
Co-authored-by: gemini-code-assist[bot]
<176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
sdks/python/apache_beam/transforms/core.py | 9 +++++++++
sdks/python/apache_beam/typehints/typecheck_test.py | 9 +++++++++
2 files changed, 18 insertions(+)
diff --git a/sdks/python/apache_beam/transforms/core.py
b/sdks/python/apache_beam/transforms/core.py
index 1de1506159e..11c09702e9a 100644
--- a/sdks/python/apache_beam/transforms/core.py
+++ b/sdks/python/apache_beam/transforms/core.py
@@ -3866,6 +3866,15 @@ class Flatten(PTransform):
raise ValueError(
'Input to Flatten must be an iterable. '
'Got a value of type %s instead.' % type(pvalueish))
+
+ # Spot check to see if any of the items are iterables of PCollections
+ # and raise an error if so. This is always a user-error
+ for idx, item in enumerate(pvalueish):
+ if isinstance(item, (list, tuple)) and any(
+ isinstance(sub_item, pvalue.PCollection) for sub_item in item):
+ raise TypeError(
+ 'Inputs to Flatten cannot include an iterable of PCollections. '
+ f'(input at index {idx}: "{item}")')
return pvalueish, pvalueish
def expand(self, pcolls):
diff --git a/sdks/python/apache_beam/typehints/typecheck_test.py
b/sdks/python/apache_beam/typehints/typecheck_test.py
index 15e822b6f2d..c2eaa0f6f9f 100644
--- a/sdks/python/apache_beam/typehints/typecheck_test.py
+++ b/sdks/python/apache_beam/typehints/typecheck_test.py
@@ -179,6 +179,15 @@ class PerformanceRuntimeTypeCheckTest(unittest.TestCase):
(self.p | beam.Create(['1', '1']) | beam.ParDo(ToInt()))
self.p.run().wait_until_finish()
+ def test_bad_flatten_input(self):
+ with self.assertRaisesRegex(
+ TypeError,
+ "Inputs to Flatten cannot include an iterable of PCollections. "):
+ with beam.Pipeline() as p:
+ pc = p | beam.Create([1, 1])
+ flatten_inputs = [pc, (pc, )]
+ flatten_inputs | beam.Flatten()
+
def test_do_fn_returning_non_iterable_throws_error(self):
# This function is incorrect because it returns a non-iterable object
def incorrect_par_do_fn(x):
