sdks/python/apache_beam/examples/fastavro_it_test.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""End-to-end test for Avro IO's fastavro support.

Writes a configurable number of records to a temporary location with fastavro,
then reads them back in from source, joins the generated records and records
that are read from the source, and verifies they have the same elements.



Usage:

  DataFlowRunner:
    pytest apache_beam/examples/fastavro_it_test.py \
        --test-pipeline-options="
          --runner=TestDataflowRunner
          --project=...
          --region=...
          --staging_location=gs://...
          --temp_location=gs://...
          --output=gs://...
          --sdk_location=...
        "

  DirectRunner:
    pytest apache_beam/examples/fastavro_it_test.py \
      --test-pipeline-options="
        --output=/tmp
        --records=5000
      "
"""

# pytype: skip-file

import json
import logging
import unittest
import uuid

import pytest
from fastavro import parse_schema

from apache_beam.io.avroio import ReadAllFromAvro
from apache_beam.io.avroio import WriteToAvro
from apache_beam.runners.runner import PipelineState
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.test_utils import delete_files
from apache_beam.testing.util import BeamAssertException
from apache_beam.transforms.core import Create
from apache_beam.transforms.core import FlatMap
from apache_beam.transforms.core import Map
from apache_beam.transforms.util import CoGroupByKey

LABELS = ['abc', 'def', 'ghi', 'jkl', 'mno', 'pqr', 'stu', 'vwx']
COLORS = ['RED', 'ORANGE', 'YELLOW', 'GREEN', 'BLUE', 'PURPLE', None]


def record(i):
  return {
      'label': LABELS[i % len(LABELS)],
      'number': i,
      'number_str': str(i),
      'color': COLORS[i % len(COLORS)]
  }


def assertEqual(l, r):
  if l != r:
    raise BeamAssertException('Assertion failed: %s == %s' % (l, r))


def check(element):
  assert element['color'] in COLORS
  assert element['label'] in LABELS
  assertEqual(
      sorted(element.keys()), ['color', 'label', 'number', 'number_str'])


class FastavroIT(unittest.TestCase):

  SCHEMA_STRING = '''
    {"namespace": "example.avro",
     "type": "record",
     "name": "User",
     "fields": [
         {"name": "label", "type": "string"},
         {"name": "number",  "type": ["int", "null"]},
         {"name": "number_str", "type": ["string", "null"]},
         {"name": "color", "type": ["string", "null"]}
     ]
    }
    '''

  def setUp(self):
    self.test_pipeline = TestPipeline(is_integration_test=True)
    self.uuid = str(uuid.uuid4())
    self.output = '/'.join([self.test_pipeline.get_option('output'), self.uuid])

  @pytest.mark.it_postcommit
  def test_avro_it(self):
    num_records = self.test_pipeline.get_option('records')
    num_records = int(num_records) if num_records else 1000000
    fastavro_output = '/'.join([self.output, 'fastavro'])

    # Seed a `PCollection` with indices that will each be FlatMap'd into
    # `batch_size` records, to avoid having a too-large list in memory at
    # the outset
    batch_size = self.test_pipeline.get_option('batch-size')
    batch_size = int(batch_size) if batch_size else 10000

    # pylint: disable=bad-option-value
    batches = range(int(num_records / batch_size))

    def batch_indices(start):
      # pylint: disable=bad-option-value
      return range(start * batch_size, (start + 1) * batch_size)

    # A `PCollection` with `num_records` avro records
    records_pcoll = \
        self.test_pipeline \
        | 'create-batches' >> Create(batches) \
        | 'expand-batches' >> FlatMap(batch_indices) \
        | 'create-records' >> Map(record)

    # pylint: disable=expression-not-assigned
    records_pcoll \
    | 'write_fastavro' >> WriteToAvro(
        fastavro_output,
        parse_schema(json.loads(self.SCHEMA_STRING)),
    )
    result = self.test_pipeline.run()
    result.wait_until_finish()
    fastavro_pcoll = self.test_pipeline \
                     | 'create-fastavro' >> Create(['%s*' % fastavro_output]) \
                     | 'read-fastavro' >> ReadAllFromAvro()

    mapped_fastavro_pcoll = fastavro_pcoll | "map_fastavro" >> Map(
        lambda x: (x['number'], x))
    mapped_record_pcoll = records_pcoll | "map_record" >> Map(
        lambda x: (x['number'], x))

    def validate_record(elem):
      v = elem[1]

      def assertEqual(l, r):
        if l != r:
          raise BeamAssertException('Assertion failed: %s == %s' % (l, r))

      assertEqual(sorted(v.keys()), ['fastavro', 'record_pcoll'])
      record_pcoll_values = v['record_pcoll']
      fastavro_values = v['fastavro']
      assertEqual(record_pcoll_values, fastavro_values)
      assertEqual(len(record_pcoll_values), 1)

    {
        "record_pcoll": mapped_record_pcoll, "fastavro": mapped_fastavro_pcoll
    } | CoGroupByKey() | Map(validate_record)

    result = self.test_pipeline.run()
    result.wait_until_finish()

    self.addCleanup(delete_files, [self.output])
    assert result.state == PipelineState.DONE


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.DEBUG)
  unittest.main()