| # |
| # 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. |
| # |
| # pytype: skip-file |
| |
| import glob |
| import json |
| import logging |
| import math |
| import os |
| import pytz |
| import pytest |
| import re |
| import shutil |
| import tempfile |
| import unittest |
| from typing import List, Any |
| |
| import fastavro |
| import hamcrest as hc |
| |
| from fastavro.schema import parse_schema |
| from fastavro import writer |
| |
| import apache_beam as beam |
| from apache_beam import Create, schema_pb2 |
| from apache_beam.io import avroio |
| from apache_beam.io import filebasedsource |
| from apache_beam.io import iobase |
| from apache_beam.io import source_test_utils |
| from apache_beam.io.avroio import _FastAvroSource # For testing |
| from apache_beam.io.avroio import avro_schema_to_beam_schema # For testing |
| from apache_beam.io.avroio import beam_schema_to_avro_schema # For testing |
| from apache_beam.io.avroio import avro_union_type_to_beam_type # For testing |
| from apache_beam.io.avroio import avro_dict_to_beam_row # For testing |
| from apache_beam.io.avroio import beam_row_to_avro_dict # For testing |
| from apache_beam.io.avroio import _create_avro_sink # For testing |
| from apache_beam.io.filesystems import FileSystems |
| from apache_beam.options.pipeline_options import StandardOptions |
| from apache_beam.testing.test_pipeline import TestPipeline |
| from apache_beam.testing.test_stream import TestStream |
| from apache_beam.testing.util import assert_that |
| from apache_beam.testing.util import equal_to |
| from apache_beam.transforms.display import DisplayData |
| from apache_beam.transforms.display_test import DisplayDataItemMatcher |
| from apache_beam.transforms.sql import SqlTransform |
| from apache_beam.transforms.userstate import CombiningValueStateSpec |
| from apache_beam.transforms.util import LogElements |
| from apache_beam.utils.timestamp import Timestamp |
| from apache_beam.typehints import schemas |
| from datetime import datetime |
| |
| # Import snappy optionally; some tests will be skipped when import fails. |
| try: |
| import snappy # pylint: disable=import-error |
| except ImportError: |
| snappy = None # pylint: disable=invalid-name |
| logging.warning('python-snappy is not installed; some tests will be skipped.') |
| |
| RECORDS = [{ |
| 'name': 'Thomas', 'favorite_number': 1, 'favorite_color': 'blue' |
| }, { |
| 'name': 'Henry', 'favorite_number': 3, 'favorite_color': 'green' |
| }, { |
| 'name': 'Toby', 'favorite_number': 7, 'favorite_color': 'brown' |
| }, { |
| 'name': 'Gordon', 'favorite_number': 4, 'favorite_color': 'blue' |
| }, { |
| 'name': 'Emily', 'favorite_number': -1, 'favorite_color': 'Red' |
| }, { |
| 'name': 'Percy', 'favorite_number': 6, 'favorite_color': 'Green' |
| }] |
| |
| |
| class AvroBase(object): |
| |
| _temp_files: List[str] = [] |
| |
| def __init__(self, methodName='runTest'): |
| super().__init__(methodName) |
| self.RECORDS = RECORDS |
| self.SCHEMA_STRING = ''' |
| {"namespace": "example.avro", |
| "type": "record", |
| "name": "User", |
| "fields": [ |
| {"name": "name", "type": "string"}, |
| {"name": "favorite_number", "type": ["int", "null"]}, |
| {"name": "favorite_color", "type": ["string", "null"]} |
| ] |
| } |
| ''' |
| |
| def setUp(self): |
| # Reducing the size of thread pools. Without this test execution may fail in |
| # environments with limited amount of resources. |
| filebasedsource.MAX_NUM_THREADS_FOR_SIZE_ESTIMATION = 2 |
| |
| def tearDown(self): |
| for path in self._temp_files: |
| if os.path.exists(path): |
| os.remove(path) |
| self._temp_files = [] |
| |
| def _write_data( |
| self, |
| directory=None, |
| prefix=None, |
| codec=None, |
| count=None, |
| sync_interval=None): |
| raise NotImplementedError |
| |
| def _write_pattern(self, num_files, return_filenames=False): |
| assert num_files > 0 |
| temp_dir = tempfile.mkdtemp() |
| |
| file_name = None |
| file_list = [] |
| for _ in range(num_files): |
| file_name = self._write_data(directory=temp_dir, prefix='mytemp') |
| file_list.append(file_name) |
| |
| assert file_name |
| file_name_prefix = file_name[:file_name.rfind(os.path.sep)] |
| if return_filenames: |
| return (file_name_prefix + os.path.sep + 'mytemp*', file_list) |
| return file_name_prefix + os.path.sep + 'mytemp*' |
| |
| def _run_avro_test( |
| self, pattern, desired_bundle_size, perform_splitting, expected_result): |
| source = _FastAvroSource(pattern) |
| |
| if perform_splitting: |
| assert desired_bundle_size |
| splits = [ |
| split |
| for split in source.split(desired_bundle_size=desired_bundle_size) |
| ] |
| if len(splits) < 2: |
| raise ValueError( |
| 'Test is trivial. Please adjust it so that at least ' |
| 'two splits get generated') |
| |
| sources_info = [(split.source, split.start_position, split.stop_position) |
| for split in splits] |
| source_test_utils.assert_sources_equal_reference_source( |
| (source, None, None), sources_info) |
| else: |
| read_records = source_test_utils.read_from_source(source, None, None) |
| self.assertCountEqual(expected_result, read_records) |
| |
| def test_schema_read_write(self): |
| with tempfile.TemporaryDirectory() as tmp_dirname: |
| path = os.path.join(tmp_dirname, 'tmp_filename') |
| rows = [beam.Row(a=1, b=['x', 'y']), beam.Row(a=2, b=['t', 'u'])] |
| stable_repr = lambda row: json.dumps(row._asdict()) |
| with TestPipeline() as p: |
| _ = p | Create(rows) | avroio.WriteToAvro(path) | beam.Map(print) |
| with TestPipeline() as p: |
| readback = ( |
| p |
| | avroio.ReadFromAvro(path + '*', as_rows=True) |
| | beam.Map(stable_repr)) |
| assert_that(readback, equal_to([stable_repr(r) for r in rows])) |
| |
| @pytest.mark.xlang_sql_expansion_service |
| @unittest.skipIf( |
| TestPipeline().get_pipeline_options().view_as(StandardOptions).runner |
| is None, |
| "Must be run with a runner that supports staging java artifacts.") |
| def test_avro_schema_to_beam_schema_with_nullable_atomic_fields(self): |
| records = [] |
| records.extend(self.RECORDS) |
| records.append({ |
| 'name': 'Bruce', 'favorite_number': None, 'favorite_color': None |
| }) |
| avro_schema = fastavro.parse_schema(json.loads(self.SCHEMA_STRING)) |
| beam_schema = avro_schema_to_beam_schema(avro_schema) |
| |
| with TestPipeline() as p: |
| readback = ( |
| p |
| | Create(records) |
| | beam.Map(avro_dict_to_beam_row(avro_schema, beam_schema)) |
| | SqlTransform("SELECT * FROM PCOLLECTION") |
| | beam.Map(beam_row_to_avro_dict(avro_schema, beam_schema))) |
| assert_that(readback, equal_to(records)) |
| |
| def test_avro_union_type_to_beam_type_with_nullable_long(self): |
| union_type = ['null', 'long'] |
| beam_type = avro_union_type_to_beam_type(union_type) |
| expected_beam_type = schema_pb2.FieldType( |
| atomic_type=schema_pb2.INT64, nullable=True) |
| hc.assert_that(beam_type, hc.equal_to(expected_beam_type)) |
| |
| def test_avro_union_type_to_beam_type_with_string_long(self): |
| union_type = ['string', 'long'] |
| beam_type = avro_union_type_to_beam_type(union_type) |
| expected_beam_type = schemas.typing_to_runner_api(Any) |
| hc.assert_that(beam_type, hc.equal_to(expected_beam_type)) |
| |
| def test_avro_union_type_to_beam_type_with_record_and_null(self): |
| record_type = { |
| 'type': 'record', |
| 'name': 'TestRecord', |
| 'fields': [{ |
| 'name': 'field1', 'type': 'string' |
| }, { |
| 'name': 'field2', 'type': 'int' |
| }] |
| } |
| union_type = [record_type, 'null'] |
| beam_type = avro_union_type_to_beam_type(union_type) |
| expected_beam_type = schema_pb2.FieldType( |
| row_type=schema_pb2.RowType( |
| schema=schema_pb2.Schema( |
| fields=[ |
| schemas.schema_field( |
| 'field1', |
| schema_pb2.FieldType(atomic_type=schema_pb2.STRING)), |
| schemas.schema_field( |
| 'field2', |
| schema_pb2.FieldType(atomic_type=schema_pb2.INT32)) |
| ])), |
| nullable=True) |
| hc.assert_that(beam_type, hc.equal_to(expected_beam_type)) |
| |
| def test_avro_union_type_to_beam_type_with_nullable_annotated_string(self): |
| annotated_string_type = {"avro.java.string": "String", "type": "string"} |
| union_type = ['null', annotated_string_type] |
| |
| beam_type = avro_union_type_to_beam_type(union_type) |
| |
| expected_beam_type = schema_pb2.FieldType( |
| atomic_type=schema_pb2.STRING, nullable=True) |
| hc.assert_that(beam_type, hc.equal_to(expected_beam_type)) |
| |
| def test_avro_union_type_to_beam_type_with_only_null(self): |
| union_type = ['null'] |
| beam_type = avro_union_type_to_beam_type(union_type) |
| expected_beam_type = schemas.typing_to_runner_api(Any) |
| hc.assert_that(beam_type, hc.equal_to(expected_beam_type)) |
| |
| def test_avro_union_type_to_beam_type_with_multiple_types(self): |
| union_type = ['null', 'string', 'int'] |
| beam_type = avro_union_type_to_beam_type(union_type) |
| expected_beam_type = schemas.typing_to_runner_api(Any) |
| hc.assert_that(beam_type, hc.equal_to(expected_beam_type)) |
| |
| def test_avro_schema_to_beam_and_back(self): |
| avro_schema = fastavro.parse_schema(json.loads(self.SCHEMA_STRING)) |
| beam_schema = avro_schema_to_beam_schema(avro_schema) |
| converted_avro_schema = beam_schema_to_avro_schema(beam_schema) |
| expected_fields = json.loads(self.SCHEMA_STRING)["fields"] |
| hc.assert_that( |
| converted_avro_schema["fields"], hc.equal_to(expected_fields)) |
| |
| def test_read_without_splitting(self): |
| file_name = self._write_data() |
| expected_result = self.RECORDS |
| self._run_avro_test(file_name, None, False, expected_result) |
| |
| def test_read_with_splitting(self): |
| file_name = self._write_data() |
| expected_result = self.RECORDS |
| self._run_avro_test(file_name, 100, True, expected_result) |
| |
| def test_source_display_data(self): |
| file_name = 'some_avro_source' |
| source = \ |
| _FastAvroSource( |
| file_name, |
| validate=False, |
| ) |
| dd = DisplayData.create_from(source) |
| |
| # No extra avro parameters for AvroSource. |
| expected_items = [ |
| DisplayDataItemMatcher('compression', 'auto'), |
| DisplayDataItemMatcher('file_pattern', file_name) |
| ] |
| hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) |
| |
| def test_read_display_data(self): |
| file_name = 'some_avro_source' |
| read = \ |
| avroio.ReadFromAvro( |
| file_name, |
| validate=False) |
| dd = DisplayData.create_from(read) |
| |
| # No extra avro parameters for AvroSource. |
| expected_items = [ |
| DisplayDataItemMatcher('compression', 'auto'), |
| DisplayDataItemMatcher('file_pattern', file_name) |
| ] |
| hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) |
| |
| def test_sink_display_data(self): |
| file_name = 'some_avro_sink' |
| sink = _create_avro_sink( |
| file_name, self.SCHEMA, 'null', '.end', 0, None, 'application/x-avro') |
| dd = DisplayData.create_from(sink) |
| |
| expected_items = [ |
| DisplayDataItemMatcher('schema', str(self.SCHEMA)), |
| DisplayDataItemMatcher( |
| 'file_pattern', |
| 'some_avro_sink-%(shard_num)05d-of-%(num_shards)05d.end'), |
| DisplayDataItemMatcher('codec', 'null'), |
| DisplayDataItemMatcher('compression', 'uncompressed') |
| ] |
| |
| hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) |
| |
| def test_write_display_data(self): |
| file_name = 'some_avro_sink' |
| write = avroio.WriteToAvro(file_name, self.SCHEMA) |
| write.expand(beam.PCollection(beam.Pipeline())) |
| dd = DisplayData.create_from(write) |
| expected_items = [ |
| DisplayDataItemMatcher('schema', str(self.SCHEMA)), |
| DisplayDataItemMatcher( |
| 'file_pattern', |
| 'some_avro_sink-%(shard_num)05d-of-%(num_shards)05d'), |
| DisplayDataItemMatcher('codec', 'deflate'), |
| DisplayDataItemMatcher('compression', 'uncompressed') |
| ] |
| hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) |
| |
| def test_read_reentrant_without_splitting(self): |
| file_name = self._write_data() |
| source = _FastAvroSource(file_name) |
| source_test_utils.assert_reentrant_reads_succeed((source, None, None)) |
| |
| def test_read_reantrant_with_splitting(self): |
| file_name = self._write_data() |
| source = _FastAvroSource(file_name) |
| splits = [split for split in source.split(desired_bundle_size=100000)] |
| assert len(splits) == 1 |
| source_test_utils.assert_reentrant_reads_succeed( |
| (splits[0].source, splits[0].start_position, splits[0].stop_position)) |
| |
| def test_read_without_splitting_multiple_blocks(self): |
| file_name = self._write_data(count=12000) |
| expected_result = self.RECORDS * 2000 |
| self._run_avro_test(file_name, None, False, expected_result) |
| |
| def test_read_with_splitting_multiple_blocks(self): |
| file_name = self._write_data(count=12000) |
| expected_result = self.RECORDS * 2000 |
| self._run_avro_test(file_name, 10000, True, expected_result) |
| |
| def test_split_points(self): |
| num_records = 12000 |
| sync_interval = 16000 |
| file_name = self._write_data(count=num_records, sync_interval=sync_interval) |
| |
| source = _FastAvroSource(file_name) |
| |
| splits = [split for split in source.split(desired_bundle_size=float('inf'))] |
| assert len(splits) == 1 |
| range_tracker = splits[0].source.get_range_tracker( |
| splits[0].start_position, splits[0].stop_position) |
| |
| split_points_report = [] |
| |
| for _ in splits[0].source.read(range_tracker): |
| split_points_report.append(range_tracker.split_points()) |
| # There will be a total of num_blocks in the generated test file, |
| # proportional to number of records in the file divided by syncronization |
| # interval used by avro during write. Each block has more than 10 records. |
| num_blocks = int(math.ceil(14.5 * num_records / sync_interval)) |
| assert num_blocks > 1 |
| # When reading records of the first block, range_tracker.split_points() |
| # should return (0, iobase.RangeTracker.SPLIT_POINTS_UNKNOWN) |
| self.assertEqual( |
| split_points_report[:10], |
| [(0, iobase.RangeTracker.SPLIT_POINTS_UNKNOWN)] * 10) |
| |
| # When reading records of last block, range_tracker.split_points() should |
| # return (num_blocks - 1, 1) |
| self.assertEqual(split_points_report[-10:], [(num_blocks - 1, 1)] * 10) |
| |
| def test_read_without_splitting_compressed_deflate(self): |
| file_name = self._write_data(codec='deflate') |
| expected_result = self.RECORDS |
| self._run_avro_test(file_name, None, False, expected_result) |
| |
| def test_read_with_splitting_compressed_deflate(self): |
| file_name = self._write_data(codec='deflate') |
| expected_result = self.RECORDS |
| self._run_avro_test(file_name, 100, True, expected_result) |
| |
| @unittest.skipIf(snappy is None, 'python-snappy not installed.') |
| def test_read_without_splitting_compressed_snappy(self): |
| file_name = self._write_data(codec='snappy') |
| expected_result = self.RECORDS |
| self._run_avro_test(file_name, None, False, expected_result) |
| |
| @unittest.skipIf(snappy is None, 'python-snappy not installed.') |
| def test_read_with_splitting_compressed_snappy(self): |
| file_name = self._write_data(codec='snappy') |
| expected_result = self.RECORDS |
| self._run_avro_test(file_name, 100, True, expected_result) |
| |
| def test_read_without_splitting_pattern(self): |
| pattern = self._write_pattern(3) |
| expected_result = self.RECORDS * 3 |
| self._run_avro_test(pattern, None, False, expected_result) |
| |
| def test_read_with_splitting_pattern(self): |
| pattern = self._write_pattern(3) |
| expected_result = self.RECORDS * 3 |
| self._run_avro_test(pattern, 100, True, expected_result) |
| |
| def test_dynamic_work_rebalancing_exhaustive(self): |
| def compare_split_points(file_name): |
| source = _FastAvroSource(file_name) |
| splits = [ |
| split for split in source.split(desired_bundle_size=float('inf')) |
| ] |
| assert len(splits) == 1 |
| source_test_utils.assert_split_at_fraction_exhaustive(splits[0].source) |
| |
| # Adjusting block size so that we can perform a exhaustive dynamic |
| # work rebalancing test that completes within an acceptable amount of time. |
| file_name = self._write_data(count=5, sync_interval=2) |
| |
| compare_split_points(file_name) |
| |
| def test_corrupted_file(self): |
| file_name = self._write_data() |
| with open(file_name, 'rb') as f: |
| data = f.read() |
| |
| # Corrupt the last character of the file which is also the last character of |
| # the last sync_marker. |
| # https://avro.apache.org/docs/current/spec.html#Object+Container+Files |
| corrupted_data = bytearray(data) |
| corrupted_data[-1] = (corrupted_data[-1] + 1) % 256 |
| with tempfile.NamedTemporaryFile(delete=False, |
| prefix=tempfile.template) as f: |
| f.write(corrupted_data) |
| corrupted_file_name = f.name |
| |
| source = _FastAvroSource(corrupted_file_name) |
| with self.assertRaisesRegex(ValueError, r'expected sync marker'): |
| source_test_utils.read_from_source(source, None, None) |
| |
| def test_read_from_avro(self): |
| path = self._write_data() |
| with TestPipeline() as p: |
| assert_that(p | avroio.ReadFromAvro(path), equal_to(self.RECORDS)) |
| |
| def test_read_all_from_avro_single_file(self): |
| path = self._write_data() |
| with TestPipeline() as p: |
| assert_that( |
| p \ |
| | Create([path]) \ |
| | avroio.ReadAllFromAvro(), |
| equal_to(self.RECORDS)) |
| |
| def test_read_all_from_avro_many_single_files(self): |
| path1 = self._write_data() |
| path2 = self._write_data() |
| path3 = self._write_data() |
| with TestPipeline() as p: |
| assert_that( |
| p \ |
| | Create([path1, path2, path3]) \ |
| | avroio.ReadAllFromAvro(), |
| equal_to(self.RECORDS * 3)) |
| |
| def test_read_all_from_avro_file_pattern(self): |
| file_pattern = self._write_pattern(5) |
| with TestPipeline() as p: |
| assert_that( |
| p \ |
| | Create([file_pattern]) \ |
| | avroio.ReadAllFromAvro(), |
| equal_to(self.RECORDS * 5)) |
| |
| def test_read_all_from_avro_many_file_patterns(self): |
| file_pattern1 = self._write_pattern(5) |
| file_pattern2 = self._write_pattern(2) |
| file_pattern3 = self._write_pattern(3) |
| with TestPipeline() as p: |
| assert_that( |
| p \ |
| | Create([file_pattern1, file_pattern2, file_pattern3]) \ |
| | avroio.ReadAllFromAvro(), |
| equal_to(self.RECORDS * 10)) |
| |
| def test_read_all_from_avro_with_filename(self): |
| file_pattern, file_paths = self._write_pattern(3, return_filenames=True) |
| result = [(path, record) for path in file_paths for record in self.RECORDS] |
| with TestPipeline() as p: |
| assert_that( |
| p \ |
| | Create([file_pattern]) \ |
| | avroio.ReadAllFromAvro(with_filename=True), |
| equal_to(result)) |
| |
| class _WriteFilesFn(beam.DoFn): |
| """writes a couple of files with deferral.""" |
| |
| COUNT_STATE = CombiningValueStateSpec('count', combine_fn=sum) |
| |
| def __init__(self, SCHEMA, RECORDS, tempdir): |
| self._thread = None |
| self.SCHEMA = SCHEMA |
| self.RECORDS = RECORDS |
| self.tempdir = tempdir |
| |
| def get_expect(self, match_updated_files): |
| results_file1 = [('file1', x) for x in self.gen_records(1)] |
| results_file2 = [('file2', x) for x in self.gen_records(3)] |
| if match_updated_files: |
| results_file1 += [('file1', x) for x in self.gen_records(2)] |
| return results_file1 + results_file2 |
| |
| def gen_records(self, count): |
| return self.RECORDS * (count // len(self.RECORDS)) + self.RECORDS[:( |
| count % len(self.RECORDS))] |
| |
| def process(self, element, count_state=beam.DoFn.StateParam(COUNT_STATE)): |
| counter = count_state.read() |
| if counter == 0: |
| count_state.add(1) |
| with open(FileSystems.join(self.tempdir, 'file1'), 'wb') as f: |
| writer(f, self.SCHEMA, self.gen_records(2)) |
| with open(FileSystems.join(self.tempdir, 'file2'), 'wb') as f: |
| writer(f, self.SCHEMA, self.gen_records(3)) |
| # convert dumb key to basename in output |
| basename = FileSystems.split(element[1][0])[1] |
| content = element[1][1] |
| yield basename, content |
| |
| def test_read_all_continuously_new(self): |
| with TestPipeline() as pipeline: |
| tempdir = tempfile.mkdtemp() |
| writer_fn = self._WriteFilesFn(self.SCHEMA, self.RECORDS, tempdir) |
| with open(FileSystems.join(tempdir, 'file1'), 'wb') as f: |
| writer(f, writer_fn.SCHEMA, writer_fn.gen_records(1)) |
| match_pattern = FileSystems.join(tempdir, '*') |
| interval = 0.5 |
| last = 2 |
| |
| p_read_once = ( |
| pipeline |
| | 'Continuously read new files' >> avroio.ReadAllFromAvroContinuously( |
| match_pattern, |
| with_filename=True, |
| start_timestamp=Timestamp.now(), |
| interval=interval, |
| stop_timestamp=Timestamp.now() + last, |
| match_updated_files=False) |
| | 'add dumb key' >> beam.Map(lambda x: (0, x)) |
| | 'Write files on-the-fly' >> beam.ParDo(writer_fn)) |
| assert_that( |
| p_read_once, |
| equal_to(writer_fn.get_expect(match_updated_files=False)), |
| label='assert read new files results') |
| |
| def test_read_all_continuously_update(self): |
| with TestPipeline() as pipeline: |
| tempdir = tempfile.mkdtemp() |
| writer_fn = self._WriteFilesFn(self.SCHEMA, self.RECORDS, tempdir) |
| with open(FileSystems.join(tempdir, 'file1'), 'wb') as f: |
| writer(f, writer_fn.SCHEMA, writer_fn.gen_records(1)) |
| match_pattern = FileSystems.join(tempdir, '*') |
| interval = 0.5 |
| last = 2 |
| |
| p_read_upd = ( |
| pipeline |
| | 'Continuously read updated files' >> |
| avroio.ReadAllFromAvroContinuously( |
| match_pattern, |
| with_filename=True, |
| start_timestamp=Timestamp.now(), |
| interval=interval, |
| stop_timestamp=Timestamp.now() + last, |
| match_updated_files=True) |
| | 'add dumb key' >> beam.Map(lambda x: (0, x)) |
| | 'Write files on-the-fly' >> beam.ParDo(writer_fn)) |
| assert_that( |
| p_read_upd, |
| equal_to(writer_fn.get_expect(match_updated_files=True)), |
| label='assert read updated files results') |
| |
| def test_sink_transform(self): |
| with tempfile.NamedTemporaryFile() as dst: |
| path = dst.name |
| with TestPipeline() as p: |
| # pylint: disable=expression-not-assigned |
| p \ |
| | beam.Create(self.RECORDS) \ |
| | avroio.WriteToAvro(path, self.SCHEMA,) |
| with TestPipeline() as p: |
| # json used for stable sortability |
| readback = \ |
| p \ |
| | avroio.ReadFromAvro(path + '*', ) \ |
| | beam.Map(json.dumps) |
| assert_that(readback, equal_to([json.dumps(r) for r in self.RECORDS])) |
| |
| @unittest.skipIf(snappy is None, 'python-snappy not installed.') |
| def test_sink_transform_snappy(self): |
| with tempfile.NamedTemporaryFile() as dst: |
| path = dst.name |
| with TestPipeline() as p: |
| # pylint: disable=expression-not-assigned |
| p \ |
| | beam.Create(self.RECORDS) \ |
| | avroio.WriteToAvro( |
| path, |
| self.SCHEMA, |
| codec='snappy') |
| with TestPipeline() as p: |
| # json used for stable sortability |
| readback = \ |
| p \ |
| | avroio.ReadFromAvro(path + '*') \ |
| | beam.Map(json.dumps) |
| assert_that(readback, equal_to([json.dumps(r) for r in self.RECORDS])) |
| |
| def test_writer_open_and_close(self): |
| # Create and then close a temp file so we can manually open it later |
| dst = tempfile.NamedTemporaryFile(delete=False) |
| dst.close() |
| |
| schema = parse_schema(json.loads(self.SCHEMA_STRING)) |
| sink = _create_avro_sink( |
| 'some_avro_sink', schema, 'null', '.end', 0, None, 'application/x-avro') |
| |
| w = sink.open(dst.name) |
| |
| sink.close(w) |
| |
| os.unlink(dst.name) |
| |
| |
| class TestFastAvro(AvroBase, unittest.TestCase): |
| def __init__(self, methodName='runTest'): |
| super().__init__(methodName) |
| self.SCHEMA = parse_schema(json.loads(self.SCHEMA_STRING)) |
| |
| def _write_data( |
| self, |
| directory=None, |
| prefix=tempfile.template, |
| codec='null', |
| count=len(RECORDS), |
| **kwargs): |
| all_records = self.RECORDS * \ |
| (count // len(self.RECORDS)) + self.RECORDS[:(count % len(self.RECORDS))] |
| with tempfile.NamedTemporaryFile(delete=False, |
| dir=directory, |
| prefix=prefix, |
| mode='w+b') as f: |
| writer(f, self.SCHEMA, all_records, codec=codec, **kwargs) |
| self._temp_files.append(f.name) |
| return f.name |
| |
| |
| class GenerateEvent(beam.PTransform): |
| @staticmethod |
| def sample_data(): |
| return GenerateEvent() |
| |
| def expand(self, input): |
| elemlist = [{'age': 10}, {'age': 20}, {'age': 30}] |
| elem = elemlist |
| return ( |
| input |
| | TestStream().add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 1, 0, |
| tzinfo=pytz.UTC).timestamp()).add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 2, 0, |
| tzinfo=pytz.UTC).timestamp()).add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 3, 0, |
| tzinfo=pytz.UTC).timestamp()).add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 4, 0, |
| tzinfo=pytz.UTC).timestamp()). |
| advance_watermark_to( |
| datetime(2021, 3, 1, 0, 0, 5, 0, |
| tzinfo=pytz.UTC).timestamp()).add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 5, 0, |
| tzinfo=pytz.UTC).timestamp()). |
| add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 6, |
| 0, tzinfo=pytz.UTC).timestamp()).add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 7, 0, |
| tzinfo=pytz.UTC).timestamp()).add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 8, 0, |
| tzinfo=pytz.UTC).timestamp()).add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 9, 0, |
| tzinfo=pytz.UTC).timestamp()). |
| advance_watermark_to( |
| datetime(2021, 3, 1, 0, 0, 10, 0, |
| tzinfo=pytz.UTC).timestamp()).add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 10, 0, |
| tzinfo=pytz.UTC).timestamp()).add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 11, 0, |
| tzinfo=pytz.UTC).timestamp()). |
| add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 12, 0, |
| tzinfo=pytz.UTC).timestamp()).add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 13, 0, |
| tzinfo=pytz.UTC).timestamp()).add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 14, 0, |
| tzinfo=pytz.UTC).timestamp()). |
| advance_watermark_to( |
| datetime(2021, 3, 1, 0, 0, 15, 0, |
| tzinfo=pytz.UTC).timestamp()).add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 15, 0, |
| tzinfo=pytz.UTC).timestamp()).add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 16, 0, |
| tzinfo=pytz.UTC).timestamp()). |
| add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 17, 0, |
| tzinfo=pytz.UTC).timestamp()).add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 18, 0, |
| tzinfo=pytz.UTC).timestamp()).add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 19, 0, |
| tzinfo=pytz.UTC).timestamp()). |
| advance_watermark_to( |
| datetime(2021, 3, 1, 0, 0, 20, 0, |
| tzinfo=pytz.UTC).timestamp()).add_elements( |
| elements=elem, |
| event_timestamp=datetime( |
| 2021, 3, 1, 0, 0, 20, 0, |
| tzinfo=pytz.UTC).timestamp()).advance_watermark_to( |
| datetime( |
| 2021, 3, 1, 0, 0, 25, 0, tzinfo=pytz.UTC). |
| timestamp()).advance_watermark_to_infinity()) |
| |
| |
| class WriteStreamingTest(unittest.TestCase): |
| def setUp(self): |
| super().setUp() |
| self.tempdir = tempfile.mkdtemp() |
| |
| def tearDown(self): |
| if os.path.exists(self.tempdir): |
| shutil.rmtree(self.tempdir) |
| |
| def test_write_streaming_2_shards_default_shard_name_template( |
| self, num_shards=2): |
| with TestPipeline() as p: |
| output = ( |
| p |
| | GenerateEvent.sample_data() |
| | 'User windowing' >> beam.transforms.core.WindowInto( |
| beam.transforms.window.FixedWindows(60), |
| trigger=beam.transforms.trigger.AfterWatermark(), |
| accumulation_mode=beam.transforms.trigger.AccumulationMode. |
| DISCARDING, |
| allowed_lateness=beam.utils.timestamp.Duration(seconds=0))) |
| #AvroIO |
| avroschema = { |
| 'name': 'dummy', # your supposed to be file name with .avro extension |
| 'type': 'record', # type of avro serilazation, there are more (see |
| # above docs) |
| 'fields': [ # this defines actual keys & their types |
| {'name': 'age', 'type': 'int'}, |
| ], |
| } |
| output2 = output | 'WriteToAvro' >> beam.io.WriteToAvro( |
| file_path_prefix=self.tempdir + "/ouput_WriteToAvro", |
| file_name_suffix=".avro", |
| num_shards=num_shards, |
| schema=avroschema) |
| _ = output2 | 'LogElements after WriteToAvro' >> LogElements( |
| prefix='after WriteToAvro ', with_window=True, level=logging.INFO) |
| |
| # Regex to match the expected windowed file pattern |
| # Example: |
| # ouput_WriteToAvro-[1614556800.0, 1614556805.0)-00000-of-00002.avro |
| # It captures: window_interval, shard_num, total_shards |
| pattern_string = ( |
| r'.*-\[(?P<window_start>[\d\.]+), ' |
| r'(?P<window_end>[\d\.]+|Infinity)\)-' |
| r'(?P<shard_num>\d{5})-of-(?P<total_shards>\d{5})\.avro$') |
| pattern = re.compile(pattern_string) |
| file_names = [] |
| for file_name in glob.glob(self.tempdir + '/ouput_WriteToAvro*'): |
| match = pattern.match(file_name) |
| self.assertIsNotNone( |
| match, f"File name {file_name} did not match expected pattern.") |
| if match: |
| file_names.append(file_name) |
| print("Found files matching expected pattern:", file_names) |
| self.assertEqual( |
| len(file_names), |
| num_shards, |
| "expected %d files, but got: %d" % (num_shards, len(file_names))) |
| |
| def test_write_streaming_2_shards_custom_shard_name_template( |
| self, num_shards=2, shard_name_template='-V-SSSSS-of-NNNNN'): |
| with TestPipeline() as p: |
| output = (p | GenerateEvent.sample_data()) |
| #AvroIO |
| avroschema = { |
| 'name': 'dummy', # your supposed to be file name with .avro extension |
| 'type': 'record', # type of avro serilazation |
| 'fields': [ # this defines actual keys & their types |
| {'name': 'age', 'type': 'int'}, |
| ], |
| } |
| output2 = output | 'WriteToAvro' >> beam.io.WriteToAvro( |
| file_path_prefix=self.tempdir + "/ouput_WriteToAvro", |
| file_name_suffix=".avro", |
| shard_name_template=shard_name_template, |
| num_shards=num_shards, |
| triggering_frequency=60, |
| schema=avroschema) |
| _ = output2 | 'LogElements after WriteToAvro' >> LogElements( |
| prefix='after WriteToAvro ', with_window=True, level=logging.INFO) |
| |
| # Regex to match the expected windowed file pattern |
| # Example: |
| # ouput_WriteToAvro-[2021-03-01T00-00-00, 2021-03-01T00-01-00)- |
| # 00000-of-00002.avro |
| # It captures: window_interval, shard_num, total_shards |
| pattern_string = ( |
| r'.*-\[(?P<window_start>\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}), ' |
| r'(?P<window_end>\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}|Infinity)\)-' |
| r'(?P<shard_num>\d{5})-of-(?P<total_shards>\d{5})\.avro$') |
| pattern = re.compile(pattern_string) |
| file_names = [] |
| for file_name in glob.glob(self.tempdir + '/ouput_WriteToAvro*'): |
| match = pattern.match(file_name) |
| self.assertIsNotNone( |
| match, f"File name {file_name} did not match expected pattern.") |
| if match: |
| file_names.append(file_name) |
| print("Found files matching expected pattern:", file_names) |
| self.assertEqual( |
| len(file_names), |
| num_shards, |
| "expected %d files, but got: %d" % (num_shards, len(file_names))) |
| |
| def test_write_streaming_2_shards_custom_shard_name_template_5s_window( |
| self, |
| num_shards=2, |
| shard_name_template='-V-SSSSS-of-NNNNN', |
| triggering_frequency=5): |
| with TestPipeline() as p: |
| output = (p | GenerateEvent.sample_data()) |
| #AvroIO |
| avroschema = { |
| 'name': 'dummy', # your supposed to be file name with .avro extension |
| 'type': 'record', # type of avro serilazation |
| 'fields': [ # this defines actual keys & their types |
| {'name': 'age', 'type': 'int'}, |
| ], |
| } |
| output2 = output | 'WriteToAvro' >> beam.io.WriteToAvro( |
| file_path_prefix=self.tempdir + "/ouput_WriteToAvro", |
| file_name_suffix=".txt", |
| shard_name_template=shard_name_template, |
| num_shards=num_shards, |
| triggering_frequency=triggering_frequency, |
| schema=avroschema) |
| _ = output2 | 'LogElements after WriteToAvro' >> LogElements( |
| prefix='after WriteToAvro ', with_window=True, level=logging.INFO) |
| |
| # Regex to match the expected windowed file pattern |
| # Example: |
| # ouput_WriteToAvro-[2021-03-01T00-00-00, 2021-03-01T00-01-00)- |
| # 00000-of-00002.avro |
| # It captures: window_interval, shard_num, total_shards |
| pattern_string = ( |
| r'.*-\[(?P<window_start>\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}), ' |
| r'(?P<window_end>\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}|Infinity)\)-' |
| r'(?P<shard_num>\d{5})-of-(?P<total_shards>\d{5})\.txt$') |
| pattern = re.compile(pattern_string) |
| file_names = [] |
| for file_name in glob.glob(self.tempdir + '/ouput_WriteToAvro*'): |
| match = pattern.match(file_name) |
| self.assertIsNotNone( |
| match, f"File name {file_name} did not match expected pattern.") |
| if match: |
| file_names.append(file_name) |
| print("Found files matching expected pattern:", file_names) |
| # for 5s window size, the input should be processed by 5 windows with |
| # 2 shards per window |
| self.assertEqual( |
| len(file_names), |
| 10, |
| "expected %d files, but got: %d" % (num_shards, len(file_names))) |
| |
| |
| if __name__ == '__main__': |
| logging.getLogger().setLevel(logging.INFO) |
| unittest.main() |
