[YAML] Schema-producing Create.

sdks/python/apache_beam/yaml/yaml_provider.py

@@ -460,6 +460,13 @@ def create_transform(self, type, args, yaml_create_transform):
 }
 
 
+def element_to_rows(e):
+  if isinstance(e, dict):
+    return dicts_to_rows(e)
+  else:
+    return beam.Row(element=dicts_to_rows(e))
+
+
 def dicts_to_rows(o):
   if isinstance(o, dict):
     return beam.Row(**{k: dicts_to_rows(v) for k, v in o.items()})
@@ -487,7 +494,7 @@ def create(elements: Iterable[Any], reshuffle: bool = True):
         reshuffle (optional): Whether to introduce a reshuffle if there is more
             than one element in the collection. Defaults to True.
     """
-    return beam.Create(dicts_to_rows(elements), reshuffle)
+    return beam.Create([element_to_rows(e) for e in elements], reshuffle)
 
   def with_schema(**args):
     # TODO: This is preliminary.

sdks/python/apache_beam/yaml/yaml_transform_test.py

@@ -54,11 +54,11 @@ def __init__(self, limit, error_handling):
     self._error_handling = error_handling
 
   def expand(self, pcoll):
-    def raise_on_big(element):
-      if len(element) > self._limit:
-        raise ValueError(element)
+    def raise_on_big(row):
+      if len(row.element) > self._limit:
+        raise ValueError(row.element)
       else:
-        return element
+        return row.element
 
     good, bad = pcoll | beam.Map(raise_on_big).with_exception_handling()
     return {'small_elements': good, self._error_handling['output']: bad}
@@ -201,7 +201,7 @@ def test_implicit_flatten(self):
             - type: PyMap
               input: [CreateBig, CreateSmall]
               config:
-                  fn: "lambda x: x * x"
+                  fn: "lambda x: x.element * x.element"
           output: PyMap
           ''',
           providers=TEST_PROVIDERS)
@@ -263,7 +263,7 @@ def test_name_is_not_ambiguous(self):
               - type: PyMap
                 name: PyMap
                 config:
-                    fn: "lambda elem: elem * elem"
+                    fn: "lambda row: row.element * row.element"
                 input: Create
             output: PyMap
             ''',
@@ -421,11 +421,14 @@ def test_mapping_errors(self):
             - type: Create
               config:
                   elements: [0, 1, 2, 4]
-            - type: PyMap
+            - type: MapToFields
               name: ToRow
               input: Create
               config:
-                  fn: "lambda x: beam.Row(num=x, str='a' * x or 'bbb')"
+                  language: python
+                  fields:
+                      num: element
+                      str: "'a' * element or 'bbb'"
             - type: Filter
               input: ToRow
               config:
@@ -585,7 +588,8 @@ class AnnotatingProvider(yaml_provider.InlineProvider):
   """
   def __init__(self, name, transform_names):
     super().__init__({
-        transform_name: lambda: beam.Map(lambda x: (x or ()) + (name, ))
+        transform_name:
+        lambda: beam.Map(lambda x: (x if type(x) == tuple else ()) + (name, ))
         for transform_name in transform_names.strip().split()
     })
     self._name = name
@@ -718,7 +722,7 @@ def __init__(self, a, b):
   def expand(self, pcoll):
     a = self._a
     b = self._b
-    return pcoll | beam.Map(lambda x: a * x + b)
+    return pcoll | beam.Map(lambda x: a * x.element + b)
 
 
 if __name__ == '__main__':