[beam] branch master updated: Fix FlatMap numpy array bug (#22006)

[bh...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

bhulette 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 bf2ac4119aa Fix FlatMap numpy array bug (#22006)
bf2ac4119aa is described below

commit bf2ac4119aa21d9a4e930f34c102b2eb1ea306cf
Author: Brian Hulette <bh...@google.com>
AuthorDate: Wed Jun 22 17:00:47 2022 -0700

    Fix FlatMap numpy array bug (#22006)
    
    * Add failing test
    
    * Fix flatmap bug
---
 sdks/python/apache_beam/runners/common.py                        | 7 +++++--
 .../runners/portability/fn_api_runner/fn_runner_test.py          | 9 +++++++++
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/sdks/python/apache_beam/runners/common.py b/sdks/python/apache_beam/runners/common.py
index 9d47c2ce6a2..66a7bfc38d0 100644
--- a/sdks/python/apache_beam/runners/common.py
+++ b/sdks/python/apache_beam/runners/common.py
@@ -1561,7 +1561,8 @@ class _OutputHandler(OutputHandler):
     A value wrapped in a TaggedOutput object will be unwrapped and
     then dispatched to the appropriate indexed output.
     """
-    results = results or []
+    if results is None:
+      results = []
 
     # TODO(https://github.com/apache/beam/issues/20404): Verify that the
     #  results object is a valid iterable type if
@@ -1614,7 +1615,9 @@ class _OutputHandler(OutputHandler):
     A value wrapped in a TaggedOutput object will be unwrapped and
     then dispatched to the appropriate indexed output.
     """
-    results = results or []
+    if results is None:
+      results = []
+
     output_element_count = 0
     for result in results:
       tag, result = self._handle_tagged_output(result)
diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py b/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py
index 2218f8d9018..fbc36f38011 100644
--- a/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py
+++ b/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py
@@ -454,6 +454,15 @@ class FnApiRunnerTest(unittest.TestCase):
                 ExpectingSideInputsFn(f'Do{k}'),
                 *[beam.pvalue.AsList(inputs[s]) for s in range(1, k)]))
 
+  def test_flatmap_numpy_array(self):
+    with self.create_pipeline() as p:
+      pc = (
+          p
+          | beam.Create([np.array(range(10))])
+          | beam.FlatMap(lambda arr: arr))
+
+      assert_that(pc, equal_to([np.int64(i) for i in range(10)]))
+
   @unittest.skip('https://github.com/apache/beam/issues/21228')
   def test_pardo_side_input_sparse_dependencies(self):
     with self.create_pipeline() as p: