AssertionError: DataFrame are different

DataFrame shape mismatch
[left]:  (3, 2)
[right]: (1, 2)
self = <apache_beam.yaml.yaml_transform_test.YamlTransformE2ETest testMethod=test_csv_to_json>

    def test_csv_to_json(self):
      try:
        import pandas as pd
      except ImportError:
        raise unittest.SkipTest('Pandas not available.')
    
      with tempfile.TemporaryDirectory() as tmpdir:
        data = pd.DataFrame([
            {
                'label': '11a', 'rank': 0
            },
            {
                'label': '37a', 'rank': 1
            },
            {
                'label': '389a', 'rank': 2
            },
        ])
        input = os.path.join(tmpdir, 'input.csv')
        output = os.path.join(tmpdir, 'output.json')
        data.to_csv(input, index=False)
    
        with beam.Pipeline() as p:
          result = p | YamlTransform(
              '''
              type: chain
              transforms:
                - type: ReadFromCsv
                  config:
                      path: %s
                - type: WriteToJson
                  config:
                      path: %s
                  num_shards: 1
              ''' % (repr(input), repr(output)))
    
        output_shard = list(glob.glob(output + "*"))[0]
        result = pd.read_json(
            output_shard, orient='records',
            lines=True).sort_values('rank').reindex()
>       pd.testing.assert_frame_equal(data, result)
E       AssertionError: DataFrame are different
E       
E       DataFrame shape mismatch
E       [left]:  (3, 2)
E       [right]: (1, 2)

apache_beam/yaml/yaml_transform_test.py:266: AssertionError