#

# 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()