[YAML]: add optional schema config for all transforms (#35952)

* rebase - fix changes confilct

* fix some errors

* add comment

* fix whitespace

* Update website/www/site/content/en/documentation/sdks/yaml-schema.md

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* Update website/www/site/content/en/documentation/sdks/yaml-schema.md

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* address gemeni review comment and lint issue

* address comments

* revert to filtering out None fields

* update yaml files based on combined error_handling output

* update readmes and docs

* update tests based on new logic

* update code logic based on new output design

* remove changes to text.yaml

* remove old logic

* fix wording

* update comments etc

* fix github action test

* add unit test for get_main_output_key

* update comments and logic

* add another yaml transform test to cover code change

* add comments and warnings about non-validated outputs

* updated get_main_output_keys test

* update again ut

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: derrickaw

CHANGES.md

@@ -121,7 +121,6 @@
   - Use the `--element_processing_timeout_minutes` option to reduce the chance of having stalled pipelines due to unexpected cases of slow processing, where slowness might not happen again if processing of the same element is retried.
 * (Python) Adding GCP Spanner Change Stream support for Python (apache_beam.io.gcp.spanner) ([#24103](https://github.com/apache/beam/issues/24103)).
 
-
 ## Breaking Changes
 
 * X behavior was changed ([#X](https://github.com/apache/beam/issues/X)).

sdks/python/apache_beam/yaml/README.md

@@ -62,9 +62,9 @@ or
 pytest -v integration_tests.py::<yaml_file_name_without_extension>Test
 ```
 
-To run the postcommit tests:
+To run some of the postcommit tests, for example:
 
 ```bash
-pytest -v integration_tests.py --test_files_dir="extended_tests"
+pytest -v integration_tests.py --test_files_dir="extended_tests/messaging"
 ```
 

sdks/python/apache_beam/yaml/json_utils.py

@@ -287,8 +287,10 @@ def row_to_json(beam_type: schema_pb2.FieldType) -> Callable[[Any], Any]:
         for field in beam_type.row_type.schema.fields
     }
     return lambda row: {
-        name: convert(getattr(row, name))
+        name: converted
         for (name, convert) in converters.items()
+        # To filter out nullable fields in rows
+        if (converted := convert(getattr(row, name, None))) is not None
     }
   elif type_info == "logical_type":
     return lambda value: value
@@ -348,6 +350,9 @@ def validate(row):
     nonlocal validator
     if validator is None:
       validator = jsonschema.validators.validator_for(json_schema)(json_schema)
+    # NOTE: A row like BeamSchema_...(name='Bob', score=None, age=25) needs to
+    # have any fields that are None to be filtered out or the validator will
+    # fail (e.g. {'age': 25, 'name': 'Bob'}).
     validator.validate(convert(row))
 
   return validate

sdks/python/apache_beam/yaml/pipeline.schema.yaml

@@ -45,6 +45,8 @@ $defs:
             properties: { __line__: {}}
             additionalProperties:
               type: string
+      output_schema:
+        type: object
     additionalProperties: true
     required:
       - type
@@ -129,6 +131,7 @@ $defs:
             name: {}
             input: {}
             output: {}
+            output_schema: { type: object }
             windowing: {}
             resource_hints: {}
             config: { type: object }

sdks/python/apache_beam/yaml/tests/assign_timestamps.yaml

@@ -82,3 +82,90 @@ pipelines:
           config:
             elements:
               - {user: bob, timestamp: 3}
+
+# Assign timestamp to beam row element with error handling and output schema
+# check.
+  - pipeline:
+      type: composite
+      transforms:
+        - type: Create
+          name: CreateVisits
+          config:
+            elements:
+              - {user: alice, timestamp: "not-valid"}
+              - {user: bob, timestamp: 3}
+        - type: AssignTimestamps
+          input: CreateVisits
+          config:
+            timestamp: timestamp
+            error_handling:
+              output: invalid_rows
+            output_schema:
+              type: object
+              properties:
+                user:
+                  type: string
+                timestamp:
+                  type: integer
+        - type: MapToFields
+          name: ExtractInvalidTimestamp
+          input: AssignTimestamps.invalid_rows
+          config:
+            language: python
+            fields:
+              user: "element.user"
+              timestamp: "element.timestamp"
+        - type: AssertEqual
+          input: ExtractInvalidTimestamp
+          config:
+            elements:
+              - {user: "alice", timestamp: "not-valid"}
+        - type: AssertEqual
+          input: AssignTimestamps
+          config:
+            elements: 
+              - {user: bob, timestamp: 3}
+
+# Assign timestamp to beam row element with error handling and output schema
+# check with more error handling.
+  - pipeline:
+      type: composite
+      transforms:
+        - type: Create
+          name: CreateVisits
+          config:
+            elements:
+              - {user: alice, timestamp: "not-valid"}
+              - {user: bob, timestamp: 3}
+        - type: AssignTimestamps
+          input: CreateVisits
+          config:
+            timestamp: timestamp
+            error_handling:
+              output: invalid_rows
+            output_schema:
+              type: object
+              properties:
+                user:
+                  type: string
+                timestamp:
+                  type: boolean
+        - type: MapToFields
+          name: ExtractInvalidTimestamp
+          input: AssignTimestamps.invalid_rows
+          config:
+            language: python
+            fields:
+              user: "element.user"
+              timestamp: "element.timestamp"
+        - type: AssertEqual
+          input: ExtractInvalidTimestamp
+          config:
+            elements:
+              - {user: "alice", timestamp: "not-valid"}
+              - {user: bob, timestamp: 3}
+        - type: AssertEqual
+          input: AssignTimestamps
+          config:
+            elements: []
+

sdks/python/apache_beam/yaml/tests/create.yaml

@@ -115,3 +115,26 @@ pipelines:
               - {sdk: MapReduce, year: 2004}
               - {}
               - {sdk: MillWheel, year: 2008}
+
+  # Simple Create with output schema check
+  - pipeline:
+      type: chain
+      transforms:
+        - type: Create
+          config:
+            elements:
+              - {sdk: MapReduce, year: 2004}
+              - {sdk: MillWheel, year: 2008}
+            output_schema:
+              type: object
+              properties:
+                sdk:
+                  type: string
+                year:
+                  type: integer
+        - type: AssertEqual
+          config:
+            elements:
+              - {sdk: MapReduce, year: 2004}
+              - {sdk: MillWheel, year: 2008}
+

sdks/python/apache_beam/yaml/tests/validate_with_schema.yaml

@@ -92,4 +92,48 @@ pipelines:
               - {name: "Alice", age: 30, score: 95.5}
               - {name: "Bob", age: 25, score: 88.0}
 
-
+  # Validate a Beam Row with a predefined schema, nulls, and error handling
+  - pipeline:
+      type: composite
+      transforms:
+        - type: Create
+          config:
+            elements:
+              - {name: "Alice", age: 30, score: 95.5}
+              - {name: "Bob", age: 25}
+              - {name: "Charlie", age: 27, score: "apple"}
+        - type: ValidateWithSchema
+          input: Create
+          config:
+            schema:
+              type: object
+              properties:
+                name:
+                  type: string
+                age:
+                  type: integer
+                score:
+                  type: number
+            error_handling:
+              output: invalid_rows
+        # ValidateWithSchema outputs the element, error msg, and traceback, so
+        # MapToFields is needed to easily assert downstream.
+        - type: MapToFields
+          input: ValidateWithSchema.invalid_rows
+          config:
+            language: python
+            fields:
+              name: "element.name"
+              age: "element.age"
+              score: "element.score"
+        - type: AssertEqual
+          input: ValidateWithSchema
+          config:
+            elements:
+              - {name: "Alice", age: 30, score: 95.5}
+              - {name: "Bob", age: 25}
+        - type: AssertEqual
+          input: MapToFields
+          config:
+            elements:
+              - {name: "Charlie", age: 27, score: "apple"}

sdks/python/apache_beam/yaml/yaml_errors.py

@@ -21,6 +21,7 @@
 
 import apache_beam as beam
 from apache_beam.typehints.row_type import RowTypeConstraint
+from apache_beam.yaml.yaml_utils import SafeLineLoader
 
 
 class ErrorHandlingConfig(NamedTuple):
@@ -35,9 +36,11 @@ class ErrorHandlingConfig(NamedTuple):
 
 def exception_handling_args(error_handling_spec):
   if error_handling_spec:
+    # error_handling_spec may have come from a yaml file and have metadata.
+    clean_spec = SafeLineLoader.strip_metadata(error_handling_spec)
     return {
         'dead_letter_tag' if k == 'output' else k: v
-        for (k, v) in error_handling_spec.items()
+        for (k, v) in clean_spec.items()
     }
   else:
     return None

sdks/python/apache_beam/yaml/yaml_mapping.py

@@ -487,7 +487,7 @@ def expand(self, pcoll):
           typing_from_runner_api(existing_fields[fld]))
 
 
-class _Validate(beam.PTransform):
+class Validate(beam.PTransform):
   """Validates each element of a PCollection against a json schema.
 
   Args:
@@ -982,7 +982,7 @@ def create_mapping_providers():
           'Partition-javascript': _Partition,
           'Partition-generic': _Partition,
           'StripErrorMetadata': _StripErrorMetadata,
-          'ValidateWithSchema': _Validate,
+          'ValidateWithSchema': Validate,
       }),
       yaml_provider.SqlBackedProvider({
           'Filter-sql': _SqlFilterTransform,

sdks/python/apache_beam/yaml/yaml_provider.py

@@ -759,9 +759,12 @@ def _unify_element_with_schema(element, target_schema):
   elif isinstance(element, dict):
     element_dict = element
   else:
-    # This element is not a row, so it can't be unified with a
-    # row schema.
-    return element
+    # This element is not a row-like object. If the target schema has a single
+    # field, assume this element is the value for that field.
+    if len(target_schema._fields) == 1:
+      return target_schema(**{target_schema._fields[0]: element})
+    else:
+      return element
 
   # Create new element with only the fields that exist in the original
   # element plus None for fields that are expected but missing