pickle fix

apache · Sep 18, 2023 · 0a91bd2 · 0a91bd2
1 parent 7ff7f85
commit 0a91bd2
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/sdks/python/apache_beam/yaml/yaml_udf_test.py b/sdks/python/apache_beam/yaml/yaml_udf_test.py
@@ -27,7 +27,9 @@
 from apache_beam.testing.util import equal_to
 from apache_beam.yaml.yaml_transform import YamlTransform
 
-AsRows = beam.Map(lambda named_tuple: beam.Row(**named_tuple._asdict()))
+
+def AsRows():
+  return beam.Map(lambda named_tuple: beam.Row(**named_tuple._asdict()))
 
 
 class YamlUDFMappingTest(unittest.TestCase):
@@ -108,7 +110,7 @@ def test_filter_inline_js(self):
           callable: "function filter(x) {return x.rank > 0}"
       ''')
       assert_that(
-          result | AsRows,
+          result | AsRows(),
           equal_to([
               beam.Row(label='37a', conductor=37, rank=1),
               beam.Row(label='389a', conductor=389, rank=2),
@@ -128,7 +130,7 @@ def test_filter_inline_py(self):
           callable: "lambda x: x.rank > 0"
       ''')
       assert_that(
-          result | AsRows,
+          result | AsRows(),
           equal_to([
               beam.Row(label='37a', conductor=37, rank=1),
               beam.Row(label='389a', conductor=389, rank=2),
@@ -148,7 +150,7 @@ def test_filter_expression_js(self):
           expression: "label.toUpperCase().indexOf('3') == -1 && conductor"
       ''')
       assert_that(
-          result | AsRows,
+          result | AsRows(),
           equal_to([
               beam.Row(label='11a', conductor=11, rank=0),
           ]))
@@ -167,7 +169,7 @@ def test_filter_expression_py(self):
           expression: "'3' not in label"
       ''')
       assert_that(
-          result | AsRows,
+          result | AsRows(),
           equal_to([
               beam.Row(label='11a', conductor=11, rank=0),
           ]))
@@ -200,7 +202,7 @@ def test_filter_inline_js_file(self):
             name: "f"
         ''')
       assert_that(
-          result | AsRows,
+          result | AsRows(),
           equal_to([
               beam.Row(label='37a', conductor=37, rank=1),
               beam.Row(label='389a', conductor=389, rank=2),
@@ -232,7 +234,7 @@ def g(x):
             name: "f"
         ''')
       assert_that(
-          result | AsRows,
+          result | AsRows(),
           equal_to([
               beam.Row(label='37a', conductor=37, rank=1),
               beam.Row(label='389a', conductor=389, rank=2),