[Bug]: YAML Flatten incorrectly drops fields when input PCollections' schema are different

[charlespnh]

What happened?

Run the following YAML pipeline:

pipeline:
  transforms:
    - type: Create
      name: Create1
      config:
        elements:
          - {'ride_id': '1',
             'passenger_count': 1}
          - {'ride_id': '2',
             'passenger_count': 2}

    - type: Create
      name: Create2
      config:
          elements:
            - {'ride_id': '3'}
            - {'ride_id': '4'}

    - type: Flatten
      name: Flatten1
      input:
        - Create1
        - Create2

    - type: LogForTesting
      name: LogForTesting
      input: Flatten1

... gives the following result:

INFO:root:{"ride_id": "3"}
INFO:root:{"ride_id": "4"}
INFO:root:{"ride_id": "1"}
INFO:root:{"ride_id": "2"}

... But I'm expecting:

INFO:root:{"ride_id": "3"}
INFO:root:{"ride_id": "4"}
INFO:root:{"ride_id": "1", "passenger_count": 1}
INFO:root:{"ride_id": "2", "passenger_count": 2}

This is found on apache_beam 2.66.0. Tested with Python SDK on Beam Playground and didn't see this issue of fields being dropped.

Issue Priority

Priority: 2 (default / most bugs should be filed as P2)

Issue Components

• Component: Python SDK
• Component: Java SDK
• Component: Go SDK
• Component: Typescript SDK
• Component: IO connector
• Component: Beam YAML
• Component: Beam examples
• Component: Beam playground
• Component: Beam katas
• Component: Website
• Component: Infrastructure
• Component: Spark Runner
• Component: Flink Runner
• Component: Samza Runner
• Component: Twister2 Runner
• Component: Hazelcast Jet Runner
• Component: Google Cloud Dataflow Runner

[charlespnh]

CC @chamikaramj

[chamikaramj]

Making this a P1 and a release blocker.

cc: @derrickaw @liferoad @damccorm @tvalentyn

[liferoad]

I fixed the create issue here: #35438. Have you checked the dev version?

[liferoad]

I fixed the create issue here: #35438. Have you checked the dev version?

confirmed the issue still exists. Probably caused by the schema inferred by Flatten.

[liferoad]

The workaround is to use the same schema:

pipeline:
  transforms:
    - type: Create
      name: Create1
      config:
        elements:
          - {'ride_id': '1',
             'passenger_count': 1}
          - {'ride_id': '2',
             'passenger_count': 2}

    - type: Create
      name: Create2
      config:
          elements:
            - {'ride_id': '3', 'passenger_count':null}
            - {'ride_id': '4', 'passenger_count':null}

    - type: Flatten
      name: Flatten1
      input:
        - Create2
        - Create1

    - type: LogForTesting
      name: LogForTesting
      input: Flatten1

[derrickaw]

When trying this:

  - pipeline:
      transforms:
        - type: Create
          name: Create1
          config:
            elements:
              - {'ride_id': '1', 'passenger_count': 1}
              - {'ride_id': '2', 'passenger_count': 2}
        - type: Create
          name: Create2
          config:
              elements:
                - {'ride_id': '3'}
                - {'ride_id': '4'}
        - type: Flatten
          name: Flatten1
          input:
            - Create1
            - Create2
        - type: AssertEqual
          input: Flatten1
          config:
            elements:
              - {'ride_id': '1','passenger_count': 1}
              - {'ride_id': '2','passenger_count': 2}
              - {'ride_id': '3'}
              - {'ride_id': '4'}

apache_beam.testing.util.BeamAssertException:
Failed assert:
[Row(ride_id='1', passenger_count=1), Row(ride_id='2', passenger_count=2), Row(ride_id='3'), Row(ride_id='4')] == [Row(ride_id='1'), Row(ride_id='2'), Row(ride_id='3'), Row(ride_id='4')],
unexpected elements [Row(ride_id='1'), Row(ride_id='2')],
missing elements [Row(ride_id='1', passenger_count=1), Row(ride_id='2', passenger_count=2)]

[liferoad]

Using the Python SDK works well:

import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions

def run():
    """
    This function defines and runs the Apache Beam pipeline.
    """
    # Define the pipeline options.
    # Using DirectRunner to run the pipeline locally.
    options = PipelineOptions()

    # The pipeline is created using a 'with' block, which ensures that the
    # pipeline's run() method is called and any necessary cleanup is performed.
    with beam.Pipeline(options=options) as p:
        # Transform 1: Create the first PCollection from a list of dictionaries.
        # This corresponds to the 'Create1' transform in the YAML file.
        create1 = p | 'Create1' >> beam.Create([
            {'ride_id': '1', 'passenger_count': 1},
            {'ride_id': '2', 'passenger_count': 2}
        ])

        # Transform 2: Create the second PCollection.
        # This corresponds to the 'Create2' transform in the YAML file.
        create2 = p | 'Create2' >> beam.Create([
            {'ride_id': '3'},
            {'ride_id': '4'}
        ])

        # Transform 3: Flatten the two PCollections into a single one.
        # The input PCollections are provided as a tuple to the Flatten transform.
        # This corresponds to the 'Flatten1' transform in the YAML file.
        flattened_pcollection = (create1, create2) | 'Flatten1' >> beam.Flatten()

        # Transform 4: Log the elements of the flattened PCollection.
        # This is a simple way to inspect the contents of a PCollection,
        # similar to the 'LogForTesting' transform in the YAML.
        flattened_pcollection | 'LogForTesting' >> beam.Map(print)

if __name__ == '__main__':
    # This is the entry point of the script.
    # It calls the run() function to execute the pipeline.
    run()

Output:

{'ride_id': '3'}
{'ride_id': '4'}
{'ride_id': '1', 'passenger_count': 1}
{'ride_id': '2', 'passenger_count': 2}