diff --git a/CHANGES.md b/CHANGES.md index dd3eff8eccf5..a36597854440 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -32,7 +32,6 @@ ## I/Os * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). -* Add support for streaming writes in IOBase (Python) ## New Features / Improvements @@ -107,6 +106,9 @@ * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). * Debezium IO upgraded to 3.1.1 requires Java 17 (Java) ([#34747](https://github.com/apache/beam/issues/34747)). +* Add support for streaming writes in IOBase (Python) +* Implement support for streaming writes in FileBasedSink (Python) +* Expose support for streaming writes in TextIO (Python) ## New Features / Improvements diff --git a/sdks/python/apache_beam/io/filebasedsink.py b/sdks/python/apache_beam/io/filebasedsink.py index 8bb0f7e2171e..510d253c7376 100644 --- a/sdks/python/apache_beam/io/filebasedsink.py +++ b/sdks/python/apache_beam/io/filebasedsink.py @@ -33,9 +33,12 @@ from apache_beam.options.value_provider import StaticValueProvider from apache_beam.options.value_provider import ValueProvider from apache_beam.options.value_provider import check_accessible +from apache_beam.transforms import window from apache_beam.transforms.display import DisplayDataItem DEFAULT_SHARD_NAME_TEMPLATE = '-SSSSS-of-NNNNN' +DEFAULT_WINDOW_SHARD_NAME_TEMPLATE = '-W-SSSSS-of-NNNNN' +DEFAULT_TRIGGERING_FREQUENCY = 0 __all__ = ['FileBasedSink'] @@ -71,7 +74,9 @@ def __init__( *, max_records_per_shard=None, max_bytes_per_shard=None, - skip_if_empty=False): + skip_if_empty=False, + convert_fn=None, + triggering_frequency=None): """ Raises: TypeError: if file path parameters are not a :class:`str` or @@ -98,6 +103,8 @@ def __init__( shard_name_template = DEFAULT_SHARD_NAME_TEMPLATE elif shard_name_template == '': num_shards = 1 + if triggering_frequency is None: + triggering_frequency = DEFAULT_TRIGGERING_FREQUENCY if isinstance(file_path_prefix, str): file_path_prefix = StaticValueProvider(str, file_path_prefix) if isinstance(file_name_suffix, str): @@ -106,6 +113,7 @@ def __init__( self.file_name_suffix = file_name_suffix self.num_shards = num_shards self.coder = coder + self.shard_name_template = shard_name_template self.shard_name_format = self._template_to_format(shard_name_template) self.shard_name_glob_format = self._template_to_glob_format( shard_name_template) @@ -114,6 +122,8 @@ def __init__( self.max_records_per_shard = max_records_per_shard self.max_bytes_per_shard = max_bytes_per_shard self.skip_if_empty = skip_if_empty + self.convert_fn = convert_fn + self.triggering_frequency = triggering_frequency def display_data(self): return { @@ -202,19 +212,39 @@ def open_writer(self, init_result, uid): return FileBasedSinkWriter(self, writer_path) @check_accessible(['file_path_prefix', 'file_name_suffix']) - def _get_final_name(self, shard_num, num_shards): + def _get_final_name(self, shard_num, num_shards, w=None): + if w is None or isinstance(w, window.GlobalWindow): + window_utc = None + else: + window_utc = ( + '[' + w.start.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S") + ', ' + + w.end.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S") + ')') return ''.join([ self.file_path_prefix.get(), - self.shard_name_format % - dict(shard_num=shard_num, num_shards=num_shards), + self.shard_name_format % dict( + shard_num=shard_num, + num_shards=num_shards, + uuid=(uuid.uuid4()), + window=w, + window_utc=window_utc), self.file_name_suffix.get() ]) @check_accessible(['file_path_prefix', 'file_name_suffix']) - def _get_final_name_glob(self, num_shards): + def _get_final_name_glob(self, num_shards, w=None): + if w is None or isinstance(w, window.GlobalWindow): + window_utc = None + else: + window_utc = ( + '[' + w.start.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S") + ', ' + + w.end.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S") + ')') return ''.join([ self.file_path_prefix.get(), - self.shard_name_glob_format % dict(num_shards=num_shards), + self.shard_name_glob_format % dict( + num_shards=num_shards, + uuid=(uuid.uuid4()), + window=w, + window_utc=window_utc), self.file_name_suffix.get() ]) @@ -233,7 +263,23 @@ def pre_finalize(self, init_result, writer_results): self.shard_name_glob_format) FileSystems.delete(dst_glob_files) - def _check_state_for_finalize_write(self, writer_results, num_shards): + def pre_finalize_windowed(self, init_result, writer_results, window=None): + num_shards = len(list(writer_results)) + dst_glob = self._get_final_name_glob(num_shards, window) + dst_glob_files = [ + file_metadata.path for mr in FileSystems.match([dst_glob]) + for file_metadata in mr.metadata_list + ] + + if dst_glob_files: + _LOGGER.warning( + 'Deleting %d existing files in target path matching: %s', + len(dst_glob_files), + self.shard_name_glob_format) + FileSystems.delete(dst_glob_files) + + def _check_state_for_finalize_write( + self, writer_results, num_shards, window=None): """Checks writer output files' states. Returns: @@ -248,7 +294,7 @@ def _check_state_for_finalize_write(self, writer_results, num_shards): return [], [], [], 0 src_glob = FileSystems.join(FileSystems.split(writer_results[0])[0], '*') - dst_glob = self._get_final_name_glob(num_shards) + dst_glob = self._get_final_name_glob(num_shards, window) src_glob_files = set( file_metadata.path for mr in FileSystems.match([src_glob]) for file_metadata in mr.metadata_list) @@ -261,7 +307,7 @@ def _check_state_for_finalize_write(self, writer_results, num_shards): delete_files = [] num_skipped = 0 for shard_num, src in enumerate(writer_results): - final_name = self._get_final_name(shard_num, num_shards) + final_name = self._get_final_name(shard_num, num_shards, window) dst = final_name src_exists = src in src_glob_files dst_exists = dst in dst_glob_files @@ -300,11 +346,19 @@ def _report_sink_lineage(self, dst_glob, dst_files): @check_accessible(['file_path_prefix']) def finalize_write( self, init_result, writer_results, unused_pre_finalize_results): + #Legacy finalize_write now has shares the implementation with + #finalize_windowed_write when window is None. + return self.finalize_windowed_write( + init_result, writer_results, unused_pre_finalize_results, None) + + @check_accessible(['file_path_prefix']) + def finalize_windowed_write( + self, init_result, writer_results, unused_pre_finalize_results, w=None): writer_results = sorted(writer_results) num_shards = len(writer_results) src_files, dst_files, delete_files, num_skipped = ( - self._check_state_for_finalize_write(writer_results, num_shards)) + self._check_state_for_finalize_write(writer_results, num_shards, w)) num_skipped += len(delete_files) FileSystems.delete(delete_files) num_shards_to_finalize = len(src_files) @@ -322,16 +376,8 @@ def finalize_write( ] if num_shards_to_finalize: - _LOGGER.info( - 'Starting finalize_write threads with num_shards: %d (skipped: %d), ' - 'batches: %d, num_threads: %d', - num_shards_to_finalize, - num_skipped, - len(source_file_batch), - num_threads) start_time = time.time() - # Use a thread pool for renaming operations. def _rename_batch(batch): """_rename_batch executes batch rename operations.""" source_files, destination_files = batch @@ -355,19 +401,36 @@ def _rename_batch(batch): _LOGGER.debug('Rename successful: %s -> %s', src, dst) return exceptions - exception_batches = util.run_using_threadpool( - _rename_batch, - list(zip(source_file_batch, destination_file_batch)), - num_threads) - - all_exceptions = [ - e for exception_batch in exception_batches for e in exception_batch - ] - if all_exceptions: - raise Exception( - 'Encountered exceptions in finalize_write: %s' % all_exceptions) - - yield from dst_files + if w is None or isinstance(w, window.GlobalWindow): + # bounded input was handled by finalize_write legacy method + # the implementation here should be called by finalize_write + # Use a thread pool for renaming operations. + exception_batches = util.run_using_threadpool( + _rename_batch, + list(zip(source_file_batch, destination_file_batch)), + num_threads) + + all_exceptions = [ + e for exception_batch in exception_batches for e in exception_batch + ] + if all_exceptions: + raise Exception( + 'Encountered exceptions in finalize_write: %s' % all_exceptions) + + yield from dst_files + else: + # unbounded input + batch = list([src_files, dst_files]) + exception_batches = _rename_batch(batch) + + all_exceptions = [ + e for exception_batch in exception_batches for e in exception_batch + ] + if all_exceptions: + raise Exception( + 'Encountered exceptions in finalize_write: %s' % all_exceptions) + + yield from dst_files _LOGGER.info( 'Renamed %d shards in %.2f seconds.', @@ -385,6 +448,26 @@ def _rename_batch(batch): # This error is not serious, we simply log it. _LOGGER.info('Unable to delete file: %s', init_result) + @staticmethod + def _template_replace_window(shard_name_template): + match = re.search('W+', shard_name_template) + if match: + shard_name_template = shard_name_template.replace( + match.group(0), '%%(window)0%ds' % len(match.group(0))) + match = re.search('V+', shard_name_template) + if match: + shard_name_template = shard_name_template.replace( + match.group(0), '%%(window_utc)0%ds' % len(match.group(0))) + return shard_name_template + + @staticmethod + def _template_replace_uuid(shard_name_template): + match = re.search('U+', shard_name_template) + if match: + shard_name_template = shard_name_template.replace( + match.group(0), '%%(uuid)0%dd' % len(match.group(0))) + return shard_name_template + @staticmethod def _template_replace_num_shards(shard_name_template): match = re.search('N+', shard_name_template) @@ -394,17 +477,30 @@ def _template_replace_num_shards(shard_name_template): return shard_name_template @staticmethod - def _template_to_format(shard_name_template): - if not shard_name_template: - return '' + def _template_replace_shard_num(shard_name_template): match = re.search('S+', shard_name_template) if match is None: + # shard name is required in the template. raise ValueError( "Shard number pattern S+ not found in shard_name_template: %s" % shard_name_template) - shard_name_format = shard_name_template.replace( + return shard_name_template.replace( match.group(0), '%%(shard_num)0%dd' % len(match.group(0))) - return FileBasedSink._template_replace_num_shards(shard_name_format) + + @staticmethod + def _template_to_format(shard_name_template): + if not shard_name_template: + return '' + # shard_num is required in the template, while others are optional. + replace_funcs = [ + FileBasedSink._template_replace_shard_num, + FileBasedSink._template_replace_num_shards, + FileBasedSink._template_replace_uuid, + FileBasedSink._template_replace_window + ] + for func in replace_funcs: + shard_name_template = func(shard_name_template) + return shard_name_template @staticmethod def _template_to_glob_format(shard_name_template): diff --git a/sdks/python/apache_beam/io/textio.py b/sdks/python/apache_beam/io/textio.py index d817463cfef6..ad7cbe6ea765 100644 --- a/sdks/python/apache_beam/io/textio.py +++ b/sdks/python/apache_beam/io/textio.py @@ -451,7 +451,8 @@ def __init__( *, max_records_per_shard=None, max_bytes_per_shard=None, - skip_if_empty=False): + skip_if_empty=False, + triggering_frequency=None): """Initialize a _TextSink. Args: @@ -468,13 +469,23 @@ def __init__( Constraining the number of shards is likely to reduce the performance of a pipeline. Setting this value is not recommended unless you require a specific number of output files. + In streaming if not set, the service will write a file per bundle. shard_name_template: A template string containing placeholders for - the shard number and shard count. When constructing a filename for a - particular shard number, the upper-case letters 'S' and 'N' are - replaced with the 0-padded shard number and shard count respectively. - This argument can be '' in which case it behaves as if num_shards was - set to 1 and only one file will be generated. The default pattern used - is '-SSSSS-of-NNNNN' if None is passed as the shard_name_template. + the shard number and shard count. Currently only ``''``, + ``'-SSSSS-of-NNNNN'``, ``'-W-SSSSS-of-NNNNN'`` and + ``'-V-SSSSS-of-NNNNN'`` are patterns accepted by the service. + When constructing a filename for a particular shard number, the + upper-case letters ``S`` and ``N`` are replaced with the ``0``-padded + shard number and shard count respectively. This argument can be ``''`` + in which case it behaves as if num_shards was set to 1 and only one file + will be generated. The default pattern used is ``'-SSSSS-of-NNNNN'`` for + bounded PCollections and for ``'-W-SSSSS-of-NNNNN'`` unbounded + PCollections. + W is used for windowed shard naming and is replaced with + ``[window.start, window.end)`` + V is used for windowed shard naming and is replaced with + ``[window.start.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S"), + window.end.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S")`` coder: Coder used to encode each line. compression_type: Used to handle compressed output files. Typical value is CompressionTypes.AUTO, in which case the final file path's @@ -494,6 +505,10 @@ def __init__( to exceed this value. This also tracks the uncompressed, not compressed, size of the shard. skip_if_empty: Don't write any shards if the PCollection is empty. + triggering_frequency: (int) Every triggering_frequency duration, a window + will be triggered and all bundles in the window will be written. + If set it overrides user windowing. Mandatory for GlobalWindow. + Returns: A _TextSink object usable for writing. @@ -508,7 +523,8 @@ def __init__( compression_type=compression_type, max_records_per_shard=max_records_per_shard, max_bytes_per_shard=max_bytes_per_shard, - skip_if_empty=skip_if_empty) + skip_if_empty=skip_if_empty, + triggering_frequency=triggering_frequency) self._append_trailing_newlines = append_trailing_newlines self._header = header self._footer = footer @@ -833,7 +849,8 @@ def __init__( *, max_records_per_shard=None, max_bytes_per_shard=None, - skip_if_empty=False): + skip_if_empty=False, + triggering_frequency=None): r"""Initialize a :class:`WriteToText` transform. Args: @@ -852,13 +869,21 @@ def __init__( the performance of a pipeline. Setting this value is not recommended unless you require a specific number of output files. shard_name_template (str): A template string containing placeholders for - the shard number and shard count. Currently only ``''`` and - ``'-SSSSS-of-NNNNN'`` are patterns accepted by the service. + the shard number and shard count. Currently only ``''``, + ``'-SSSSS-of-NNNNN'``, ``'-W-SSSSS-of-NNNNN'`` and + ``'-V-SSSSS-of-NNNNN'`` are patterns accepted by the service. When constructing a filename for a particular shard number, the upper-case letters ``S`` and ``N`` are replaced with the ``0``-padded shard number and shard count respectively. This argument can be ``''`` in which case it behaves as if num_shards was set to 1 and only one file - will be generated. The default pattern used is ``'-SSSSS-of-NNNNN'``. + will be generated. The default pattern used is ``'-SSSSS-of-NNNNN'`` for + bounded PCollections and for ``'-W-SSSSS-of-NNNNN'`` unbounded + PCollections. + W is used for windowed shard naming and is replaced with + ``[window.start, window.end)`` + V is used for windowed shard naming and is replaced with + ``[window.start.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S"), + window.end.to_utc_datetime().strftime("%Y-%m-%dT%H-%M-%S")`` coder (~apache_beam.coders.coders.Coder): Coder used to encode each line. compression_type (str): Used to handle compressed output files. Typical value is :class:`CompressionTypes.AUTO @@ -883,6 +908,8 @@ def __init__( skip_if_empty: Don't write any shards if the PCollection is empty. In case of an empty PCollection, this will still delete existing files having same file path and not create new ones. + triggering_frequency: (int) Every triggering_frequency duration, a window + will be triggered and all bundles in the window will be written. """ self._sink = _TextSink( @@ -897,9 +924,18 @@ def __init__( footer, max_records_per_shard=max_records_per_shard, max_bytes_per_shard=max_bytes_per_shard, - skip_if_empty=skip_if_empty) + skip_if_empty=skip_if_empty, + triggering_frequency=triggering_frequency) def expand(self, pcoll): + if (not pcoll.is_bounded and self._sink.shard_name_template + == filebasedsink.DEFAULT_SHARD_NAME_TEMPLATE): + self._sink.shard_name_template = ( + filebasedsink.DEFAULT_WINDOW_SHARD_NAME_TEMPLATE) + self._sink.shard_name_format = self._sink._template_to_format( + self._sink.shard_name_template) + self._sink.shard_name_glob_format = self._sink._template_to_glob_format( + self._sink.shard_name_template) return pcoll | Write(self._sink) diff --git a/sdks/python/apache_beam/io/textio_test.py b/sdks/python/apache_beam/io/textio_test.py index 30ddc5d62e07..192ef3c6220f 100644 --- a/sdks/python/apache_beam/io/textio_test.py +++ b/sdks/python/apache_beam/io/textio_test.py @@ -24,10 +24,14 @@ import logging import os import platform +import re import shutil import tempfile import unittest import zlib +from datetime import datetime + +import pytz import apache_beam as beam from apache_beam import coders @@ -45,11 +49,13 @@ from apache_beam.io.textio import WriteToText from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.test_stream import TestStream from apache_beam.testing.test_utils import TempDir from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to from apache_beam.transforms.core import Create from apache_beam.transforms.userstate import CombiningValueStateSpec +from apache_beam.transforms.util import LogElements from apache_beam.utils.timestamp import Timestamp @@ -1849,6 +1855,406 @@ def check_types(element): _ = pcoll | beam.Map(check_types) +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()) + #TextIO + output2 = output | 'TextIO WriteToText' >> beam.io.WriteToText( + file_path_prefix=self.tempdir + "/ouput_WriteToText", + file_name_suffix=".txt", + num_shards=num_shards, + triggering_frequency=60) + _ = output2 | 'LogElements after WriteToText' >> LogElements( + prefix='after WriteToText ', with_window=True, level=logging.INFO) + + # Regex to match the expected windowed file pattern + # Example: + # ouput_WriteToText-[1614556800.0, 1614556805.0)-00000-of-00002.txt + # It captures: window_interval, shard_num, total_shards + pattern_string = ( + r'.*-\[(?P[\d\.]+), ' + r'(?P[\d\.]+|Infinity)\)-' + r'(?P\d{5})-of-(?P\d{5})\.txt$') + pattern = re.compile(pattern_string) + file_names = [] + for file_name in glob.glob(self.tempdir + '/ouput_WriteToText*'): + 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_default_shard_name_template_windowed_pcoll( + self, num_shards=2): + with TestPipeline() as p: + output = ( + p | GenerateEvent.sample_data() + | 'User windowing' >> beam.transforms.core.WindowInto( + beam.transforms.window.FixedWindows(10), + trigger=beam.transforms.trigger.AfterWatermark(), + accumulation_mode=beam.transforms.trigger.AccumulationMode. + DISCARDING, + allowed_lateness=beam.utils.timestamp.Duration(seconds=0))) + #TextIO + output2 = output | 'TextIO WriteToText' >> beam.io.WriteToText( + file_path_prefix=self.tempdir + "/ouput_WriteToText", + file_name_suffix=".txt", + num_shards=num_shards, + ) + _ = output2 | 'LogElements after WriteToText' >> LogElements( + prefix='after WriteToText ', with_window=True, level=logging.INFO) + + # Regex to match the expected windowed file pattern + # Example: + # ouput_WriteToText-[1614556800.0, 1614556805.0)-00000-of-00002.txt + # It captures: window_interval, shard_num, total_shards + pattern_string = ( + r'.*-\[(?P[\d\.]+), ' + r'(?P[\d\.]+|Infinity)\)-' + r'(?P\d{5})-of-(?P\d{5})\.txt$') + pattern = re.compile(pattern_string) + file_names = [] + for file_name in glob.glob(self.tempdir + '/ouput_WriteToText*'): + 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 * 3, #25s of data covered by 3 10s windows + "expected %d files, but got: %d" % (num_shards * 3, len(file_names))) + + def test_write_streaming_undef_shards_default_shard_name_template_windowed_pcoll( # pylint: disable=line-too-long + self): + with TestPipeline() as p: + output = ( + p | GenerateEvent.sample_data() + | 'User windowing' >> beam.transforms.core.WindowInto( + beam.transforms.window.FixedWindows(10), + trigger=beam.transforms.trigger.AfterWatermark(), + accumulation_mode=beam.transforms.trigger.AccumulationMode. + DISCARDING, + allowed_lateness=beam.utils.timestamp.Duration(seconds=0))) + #TextIO + output2 = output | 'TextIO WriteToText' >> beam.io.WriteToText( + file_path_prefix=self.tempdir + "/ouput_WriteToText", + file_name_suffix=".txt", + num_shards=0, + ) + _ = output2 | 'LogElements after WriteToText' >> LogElements( + prefix='after WriteToText ', with_window=True, level=logging.INFO) + + # Regex to match the expected windowed file pattern + # Example: + # ouput_WriteToText-[1614556800.0, 1614556805.0)-00000-of-00002.txt + # It captures: window_interval, shard_num, total_shards + pattern_string = ( + r'.*-\[(?P[\d\.]+), ' + r'(?P[\d\.]+|Infinity)\)-' + r'(?P\d{5})-of-(?P\d{5})\.txt$') + pattern = re.compile(pattern_string) + file_names = [] + for file_name in glob.glob(self.tempdir + '/ouput_WriteToText*'): + 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.assertGreaterEqual( + len(file_names), + 1 * 3, #25s of data covered by 3 10s windows + "expected %d files, but got: %d" % (1 * 3, len(file_names))) + + def test_write_streaming_undef_shards_default_shard_name_template_windowed_pcoll_and_trig_freq( # pylint: disable=line-too-long + self): + 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))) + #TextIO + output2 = output | 'TextIO WriteToText' >> beam.io.WriteToText( + file_path_prefix=self.tempdir + "/ouput_WriteToText", + file_name_suffix=".txt", + num_shards=0, + triggering_frequency=10, + ) + _ = output2 | 'LogElements after WriteToText' >> LogElements( + prefix='after WriteToText ', with_window=True, level=logging.INFO) + + # Regex to match the expected windowed file pattern + # Example: + # ouput_WriteToText-[1614556800.0, 1614556805.0)-00000-of-00002.txt + # It captures: window_interval, shard_num, total_shards + pattern_string = ( + r'.*-\[(?P[\d\.]+), ' + r'(?P[\d\.]+|Infinity)\)-' + r'(?P\d{5})-of-(?P\d{5})\.txt$') + pattern = re.compile(pattern_string) + file_names = [] + for file_name in glob.glob(self.tempdir + '/ouput_WriteToText*'): + 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.assertGreaterEqual( + len(file_names), + 1 * 3, #25s of data covered by 3 10s windows + "expected %d files, but got: %d" % (1 * 3, len(file_names))) + + def test_write_streaming_undef_shards_default_shard_name_template_global_window_pcoll( # pylint: disable=line-too-long + self): + with TestPipeline() as p: + output = (p | GenerateEvent.sample_data()) + #TextIO + output2 = output | 'TextIO WriteToText' >> beam.io.WriteToText( + file_path_prefix=self.tempdir + "/ouput_WriteToText", + file_name_suffix=".txt", + num_shards=0, #0 means undef nb of shards, same as omitted/default + triggering_frequency=60, + ) + _ = output2 | 'LogElements after WriteToText' >> LogElements( + prefix='after WriteToText ', with_window=True, level=logging.INFO) + + # Regex to match the expected windowed file pattern + # Example: + # ouput_WriteToText-[1614556800.0, 1614556805.0)-00000-of-00002.txt + # It captures: window_interval, shard_num, total_shards + pattern_string = ( + r'.*-\[(?P[\d\.]+), ' + r'(?P[\d\.]+|Infinity)\)-' + r'(?P\d{5})-of-(?P\d{5})\.txt$') + pattern = re.compile(pattern_string) + file_names = [] + for file_name in glob.glob(self.tempdir + '/ouput_WriteToText*'): + 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.assertGreaterEqual( + len(file_names), + 1, #25s of data covered by 60s windows + "expected %d files, but got: %d" % (1, 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()) + #TextIO + output2 = output | 'TextIO WriteToText' >> beam.io.WriteToText( + file_path_prefix=self.tempdir + "/ouput_WriteToText", + file_name_suffix=".txt", + shard_name_template=shard_name_template, + num_shards=num_shards, + triggering_frequency=60, + ) + _ = output2 | 'LogElements after WriteToText' >> LogElements( + prefix='after WriteToText ', with_window=True, level=logging.INFO) + + # Regex to match the expected windowed file pattern + # Example: + # ouput_WriteToText-[2021-03-01T00-00-00, 2021-03-01T00-01-00)- + # 00000-of-00002.txt + # It captures: window_interval, shard_num, total_shards + pattern_string = ( + r'.*-\[(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}), ' + r'(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}|Infinity)\)-' + r'(?P\d{5})-of-(?P\d{5})\.txt$') + pattern = re.compile(pattern_string) + file_names = [] + for file_name in glob.glob(self.tempdir + '/ouput_WriteToText*'): + 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()) + #TextIO + output2 = output | 'TextIO WriteToText' >> beam.io.WriteToText( + file_path_prefix=self.tempdir + "/ouput_WriteToText", + file_name_suffix=".txt", + shard_name_template=shard_name_template, + num_shards=num_shards, + triggering_frequency=triggering_frequency, + ) + _ = output2 | 'LogElements after WriteToText' >> LogElements( + prefix='after WriteToText ', with_window=True, level=logging.INFO) + + # Regex to match the expected windowed file pattern + # Example: + # ouput_WriteToText-[2021-03-01T00-00-00, 2021-03-01T00-01-00)- + # 00000-of-00002.txt + # It captures: window_interval, shard_num, total_shards + pattern_string = ( + r'.*-\[(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}), ' + r'(?P\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}|Infinity)\)-' + r'(?P\d{5})-of-(?P\d{5})\.txt$') + pattern = re.compile(pattern_string) + file_names = [] + for file_name in glob.glob(self.tempdir + '/ouput_WriteToText*'): + 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()