| # |
| # 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 binascii |
| import glob |
| import gzip |
| import io |
| import json |
| import logging |
| import os |
| import pickle |
| import random |
| 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 Create |
| from apache_beam import coders |
| from apache_beam.io.filesystem import CompressionTypes |
| from apache_beam.io.tfrecordio import ReadAllFromTFRecord |
| from apache_beam.io.tfrecordio import ReadFromTFRecord |
| from apache_beam.io.tfrecordio import WriteToTFRecord |
| from apache_beam.io.tfrecordio import _TFRecordSink |
| from apache_beam.io.tfrecordio import _TFRecordUtil |
| 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.util import LogElements |
| |
| try: |
| import tensorflow.compat.v1 as tf # pylint: disable=import-error |
| except ImportError: |
| try: |
| import tensorflow as tf # pylint: disable=import-error |
| except ImportError: |
| tf = None # pylint: disable=invalid-name |
| logging.warning('Tensorflow is not installed, so skipping some tests.') |
| |
| try: |
| import crcmod |
| except ImportError: |
| crcmod = None |
| |
| # Created by running following code in python: |
| # >>> import tensorflow as tf |
| # >>> import base64 |
| # >>> writer = tf.python_io.TFRecordWriter('/tmp/python_foo.tfrecord') |
| # >>> writer.write(b'foo') |
| # >>> writer.close() |
| # >>> with open('/tmp/python_foo.tfrecord', 'rb') as f: |
| # ... data = base64.b64encode(f.read()) |
| # ... print(data) |
| FOO_RECORD_BASE64 = b'AwAAAAAAAACwmUkOZm9vYYq+/g==' |
| |
| # Same as above but containing two records [b'foo', b'bar'] |
| FOO_BAR_RECORD_BASE64 = b'AwAAAAAAAACwmUkOZm9vYYq+/gMAAAAAAAAAsJlJDmJhckYA5cg=' |
| |
| |
| def _write_file(path, base64_records): |
| record = binascii.a2b_base64(base64_records) |
| with open(path, 'wb') as f: |
| f.write(record) |
| |
| |
| def _write_file_deflate(path, base64_records): |
| record = binascii.a2b_base64(base64_records) |
| with open(path, 'wb') as f: |
| f.write(zlib.compress(record)) |
| |
| |
| def _write_file_gzip(path, base64_records): |
| record = binascii.a2b_base64(base64_records) |
| with gzip.GzipFile(path, 'wb') as f: |
| f.write(record) |
| |
| |
| class TestTFRecordUtil(unittest.TestCase): |
| def setUp(self): |
| self.record = binascii.a2b_base64(FOO_RECORD_BASE64) |
| |
| def _as_file_handle(self, contents): |
| result = io.BytesIO() |
| result.write(contents) |
| result.seek(0) |
| return result |
| |
| def _increment_value_at_index(self, value, index): |
| l = list(value) |
| l[index] = l[index] + 1 |
| return bytes(l) |
| |
| def _test_error(self, record, error_text): |
| with self.assertRaisesRegex(ValueError, re.escape(error_text)): |
| _TFRecordUtil.read_record(self._as_file_handle(record)) |
| |
| def test_masked_crc32c(self): |
| self.assertEqual(0xfd7fffa, _TFRecordUtil._masked_crc32c(b'\x00' * 32)) |
| self.assertEqual(0xf909b029, _TFRecordUtil._masked_crc32c(b'\xff' * 32)) |
| self.assertEqual(0xfebe8a61, _TFRecordUtil._masked_crc32c(b'foo')) |
| self.assertEqual( |
| 0xe4999b0, |
| _TFRecordUtil._masked_crc32c(b'\x03\x00\x00\x00\x00\x00\x00\x00')) |
| |
| @unittest.skipIf(crcmod is None, 'crcmod not installed.') |
| def test_masked_crc32c_crcmod(self): |
| crc32c_fn = crcmod.predefined.mkPredefinedCrcFun('crc-32c') |
| self.assertEqual( |
| 0xfd7fffa, |
| _TFRecordUtil._masked_crc32c(b'\x00' * 32, crc32c_fn=crc32c_fn)) |
| self.assertEqual( |
| 0xf909b029, |
| _TFRecordUtil._masked_crc32c(b'\xff' * 32, crc32c_fn=crc32c_fn)) |
| self.assertEqual( |
| 0xfebe8a61, _TFRecordUtil._masked_crc32c(b'foo', crc32c_fn=crc32c_fn)) |
| self.assertEqual( |
| 0xe4999b0, |
| _TFRecordUtil._masked_crc32c( |
| b'\x03\x00\x00\x00\x00\x00\x00\x00', crc32c_fn=crc32c_fn)) |
| |
| def test_write_record(self): |
| file_handle = io.BytesIO() |
| _TFRecordUtil.write_record(file_handle, b'foo') |
| self.assertEqual(self.record, file_handle.getvalue()) |
| |
| def test_read_record(self): |
| actual = _TFRecordUtil.read_record(self._as_file_handle(self.record)) |
| self.assertEqual(b'foo', actual) |
| |
| def test_read_record_invalid_record(self): |
| self._test_error(b'bar', 'Not a valid TFRecord. Fewer than 12 bytes') |
| |
| def test_read_record_invalid_length_mask(self): |
| record = self._increment_value_at_index(self.record, 9) |
| self._test_error(record, 'Mismatch of length mask') |
| |
| def test_read_record_invalid_data_mask(self): |
| record = self._increment_value_at_index(self.record, 16) |
| self._test_error(record, 'Mismatch of data mask') |
| |
| def test_compatibility_read_write(self): |
| for record in [b'', b'blah', b'another blah']: |
| file_handle = io.BytesIO() |
| _TFRecordUtil.write_record(file_handle, record) |
| file_handle.seek(0) |
| actual = _TFRecordUtil.read_record(file_handle) |
| self.assertEqual(record, actual) |
| |
| |
| class TestTFRecordSink(unittest.TestCase): |
| def _write_lines(self, sink, path, lines): |
| f = sink.open(path) |
| for l in lines: |
| sink.write_record(f, l) |
| sink.close(f) |
| |
| def test_write_record_single(self): |
| with TempDir() as temp_dir: |
| path = temp_dir.create_temp_file('result') |
| record = binascii.a2b_base64(FOO_RECORD_BASE64) |
| sink = _TFRecordSink( |
| path, |
| coder=coders.BytesCoder(), |
| file_name_suffix='', |
| num_shards=0, |
| shard_name_template=None, |
| compression_type=CompressionTypes.UNCOMPRESSED) |
| self._write_lines(sink, path, [b'foo']) |
| |
| with open(path, 'rb') as f: |
| self.assertEqual(f.read(), record) |
| |
| def test_write_record_multiple(self): |
| with TempDir() as temp_dir: |
| path = temp_dir.create_temp_file('result') |
| record = binascii.a2b_base64(FOO_BAR_RECORD_BASE64) |
| sink = _TFRecordSink( |
| path, |
| coder=coders.BytesCoder(), |
| file_name_suffix='', |
| num_shards=0, |
| shard_name_template=None, |
| compression_type=CompressionTypes.UNCOMPRESSED) |
| self._write_lines(sink, path, [b'foo', b'bar']) |
| |
| with open(path, 'rb') as f: |
| self.assertEqual(f.read(), record) |
| |
| |
| @unittest.skipIf(tf is None, 'tensorflow not installed.') |
| class TestWriteToTFRecord(TestTFRecordSink): |
| def test_write_record_gzip(self): |
| with TempDir() as temp_dir: |
| file_path_prefix = temp_dir.create_temp_file('result') |
| with TestPipeline() as p: |
| input_data = [b'foo', b'bar'] |
| _ = p | beam.Create(input_data) | WriteToTFRecord( |
| file_path_prefix, compression_type=CompressionTypes.GZIP) |
| |
| actual = [] |
| file_name = glob.glob(file_path_prefix + '-*')[0] |
| for r in tf.python_io.tf_record_iterator( |
| file_name, |
| options=tf.python_io.TFRecordOptions( |
| tf.python_io.TFRecordCompressionType.GZIP)): |
| actual.append(r) |
| self.assertEqual(sorted(actual), sorted(input_data)) |
| |
| def test_write_record_auto(self): |
| with TempDir() as temp_dir: |
| file_path_prefix = temp_dir.create_temp_file('result') |
| with TestPipeline() as p: |
| input_data = [b'foo', b'bar'] |
| _ = p | beam.Create(input_data) | WriteToTFRecord( |
| file_path_prefix, file_name_suffix='.gz') |
| |
| actual = [] |
| file_name = glob.glob(file_path_prefix + '-*.gz')[0] |
| for r in tf.python_io.tf_record_iterator( |
| file_name, |
| options=tf.python_io.TFRecordOptions( |
| tf.python_io.TFRecordCompressionType.GZIP)): |
| actual.append(r) |
| self.assertEqual(sorted(actual), sorted(input_data)) |
| |
| |
| class TestReadFromTFRecord(unittest.TestCase): |
| def test_process_single(self): |
| with TempDir() as temp_dir: |
| path = temp_dir.create_temp_file('result') |
| _write_file(path, FOO_RECORD_BASE64) |
| with TestPipeline() as p: |
| result = ( |
| p |
| | ReadFromTFRecord( |
| path, |
| coder=coders.BytesCoder(), |
| compression_type=CompressionTypes.AUTO, |
| validate=True)) |
| assert_that(result, equal_to([b'foo'])) |
| |
| def test_process_multiple(self): |
| with TempDir() as temp_dir: |
| path = temp_dir.create_temp_file('result') |
| _write_file(path, FOO_BAR_RECORD_BASE64) |
| with TestPipeline() as p: |
| result = ( |
| p |
| | ReadFromTFRecord( |
| path, |
| coder=coders.BytesCoder(), |
| compression_type=CompressionTypes.AUTO, |
| validate=True)) |
| assert_that(result, equal_to([b'foo', b'bar'])) |
| |
| def test_process_deflate(self): |
| with TempDir() as temp_dir: |
| path = temp_dir.create_temp_file('result') |
| _write_file_deflate(path, FOO_BAR_RECORD_BASE64) |
| with TestPipeline() as p: |
| result = ( |
| p |
| | ReadFromTFRecord( |
| path, |
| coder=coders.BytesCoder(), |
| compression_type=CompressionTypes.DEFLATE, |
| validate=True)) |
| assert_that(result, equal_to([b'foo', b'bar'])) |
| |
| def test_process_gzip_with_coder(self): |
| with TempDir() as temp_dir: |
| path = temp_dir.create_temp_file('result') |
| _write_file_gzip(path, FOO_BAR_RECORD_BASE64) |
| with TestPipeline() as p: |
| result = ( |
| p |
| | ReadFromTFRecord( |
| path, |
| coder=coders.BytesCoder(), |
| compression_type=CompressionTypes.GZIP, |
| validate=True)) |
| assert_that(result, equal_to([b'foo', b'bar'])) |
| |
| def test_process_gzip_without_coder(self): |
| with TempDir() as temp_dir: |
| path = temp_dir.create_temp_file('result') |
| _write_file_gzip(path, FOO_BAR_RECORD_BASE64) |
| with TestPipeline() as p: |
| result = ( |
| p |
| | ReadFromTFRecord(path, compression_type=CompressionTypes.GZIP)) |
| assert_that(result, equal_to([b'foo', b'bar'])) |
| |
| def test_process_auto(self): |
| with TempDir() as temp_dir: |
| path = temp_dir.create_temp_file('result.gz') |
| _write_file_gzip(path, FOO_BAR_RECORD_BASE64) |
| with TestPipeline() as p: |
| result = ( |
| p |
| | ReadFromTFRecord( |
| path, |
| coder=coders.BytesCoder(), |
| compression_type=CompressionTypes.AUTO, |
| validate=True)) |
| assert_that(result, equal_to([b'foo', b'bar'])) |
| |
| def test_process_gzip_auto(self): |
| with TempDir() as temp_dir: |
| path = temp_dir.create_temp_file('result.gz') |
| _write_file_gzip(path, FOO_BAR_RECORD_BASE64) |
| with TestPipeline() as p: |
| result = ( |
| p |
| | ReadFromTFRecord(path, compression_type=CompressionTypes.AUTO)) |
| assert_that(result, equal_to([b'foo', b'bar'])) |
| |
| |
| class TestReadAllFromTFRecord(unittest.TestCase): |
| def _write_glob(self, temp_dir, suffix, include_empty=False): |
| for _ in range(3): |
| path = temp_dir.create_temp_file(suffix) |
| _write_file(path, FOO_BAR_RECORD_BASE64) |
| if include_empty: |
| path = temp_dir.create_temp_file(suffix) |
| _write_file(path, '') |
| |
| def test_process_single(self): |
| with TempDir() as temp_dir: |
| path = temp_dir.create_temp_file('result') |
| _write_file(path, FOO_RECORD_BASE64) |
| with TestPipeline() as p: |
| result = ( |
| p |
| | Create([path]) |
| | ReadAllFromTFRecord( |
| coder=coders.BytesCoder(), |
| compression_type=CompressionTypes.AUTO)) |
| assert_that(result, equal_to([b'foo'])) |
| |
| def test_process_multiple(self): |
| with TempDir() as temp_dir: |
| path = temp_dir.create_temp_file('result') |
| _write_file(path, FOO_BAR_RECORD_BASE64) |
| with TestPipeline() as p: |
| result = ( |
| p |
| | Create([path]) |
| | ReadAllFromTFRecord( |
| coder=coders.BytesCoder(), |
| compression_type=CompressionTypes.AUTO)) |
| assert_that(result, equal_to([b'foo', b'bar'])) |
| |
| def test_process_with_filename(self): |
| with TempDir() as temp_dir: |
| num_files = 3 |
| files = [] |
| for i in range(num_files): |
| path = temp_dir.create_temp_file('result%s' % i) |
| _write_file(path, FOO_BAR_RECORD_BASE64) |
| files.append(path) |
| content = [b'foo', b'bar'] |
| expected = [(file, line) for file in files for line in content] |
| pattern = temp_dir.get_path() + os.path.sep + '*' |
| |
| with TestPipeline() as p: |
| result = ( |
| p |
| | Create([pattern]) |
| | ReadAllFromTFRecord( |
| coder=coders.BytesCoder(), |
| compression_type=CompressionTypes.AUTO, |
| with_filename=True)) |
| assert_that(result, equal_to(expected)) |
| |
| def test_process_glob(self): |
| with TempDir() as temp_dir: |
| self._write_glob(temp_dir, 'result') |
| glob = temp_dir.get_path() + os.path.sep + '*result' |
| with TestPipeline() as p: |
| result = ( |
| p |
| | Create([glob]) |
| | ReadAllFromTFRecord( |
| coder=coders.BytesCoder(), |
| compression_type=CompressionTypes.AUTO)) |
| assert_that(result, equal_to([b'foo', b'bar'] * 3)) |
| |
| def test_process_glob_with_empty_file(self): |
| with TempDir() as temp_dir: |
| self._write_glob(temp_dir, 'result', include_empty=True) |
| glob = temp_dir.get_path() + os.path.sep + '*result' |
| with TestPipeline() as p: |
| result = ( |
| p |
| | Create([glob]) |
| | ReadAllFromTFRecord( |
| coder=coders.BytesCoder(), |
| compression_type=CompressionTypes.AUTO)) |
| assert_that(result, equal_to([b'foo', b'bar'] * 3)) |
| |
| def test_process_multiple_globs(self): |
| with TempDir() as temp_dir: |
| globs = [] |
| for i in range(3): |
| suffix = 'result' + str(i) |
| self._write_glob(temp_dir, suffix) |
| globs.append(temp_dir.get_path() + os.path.sep + '*' + suffix) |
| |
| with TestPipeline() as p: |
| result = ( |
| p |
| | Create(globs) |
| | ReadAllFromTFRecord( |
| coder=coders.BytesCoder(), |
| compression_type=CompressionTypes.AUTO)) |
| assert_that(result, equal_to([b'foo', b'bar'] * 9)) |
| |
| def test_process_deflate(self): |
| with TempDir() as temp_dir: |
| path = temp_dir.create_temp_file('result') |
| _write_file_deflate(path, FOO_BAR_RECORD_BASE64) |
| with TestPipeline() as p: |
| result = ( |
| p |
| | Create([path]) |
| | ReadAllFromTFRecord( |
| coder=coders.BytesCoder(), |
| compression_type=CompressionTypes.DEFLATE)) |
| assert_that(result, equal_to([b'foo', b'bar'])) |
| |
| def test_process_gzip(self): |
| with TempDir() as temp_dir: |
| path = temp_dir.create_temp_file('result') |
| _write_file_gzip(path, FOO_BAR_RECORD_BASE64) |
| with TestPipeline() as p: |
| result = ( |
| p |
| | Create([path]) |
| | ReadAllFromTFRecord( |
| coder=coders.BytesCoder(), |
| compression_type=CompressionTypes.GZIP)) |
| assert_that(result, equal_to([b'foo', b'bar'])) |
| |
| def test_process_auto(self): |
| with TempDir() as temp_dir: |
| path = temp_dir.create_temp_file('result.gz') |
| _write_file_gzip(path, FOO_BAR_RECORD_BASE64) |
| with TestPipeline() as p: |
| result = ( |
| p |
| | Create([path]) |
| | ReadAllFromTFRecord( |
| coder=coders.BytesCoder(), |
| compression_type=CompressionTypes.AUTO)) |
| assert_that(result, equal_to([b'foo', b'bar'])) |
| |
| |
| class TestEnd2EndWriteAndRead(unittest.TestCase): |
| def create_inputs(self): |
| input_array = [[random.random() - 0.5 for _ in range(15)] |
| for _ in range(12)] |
| memfile = io.BytesIO() |
| pickle.dump(input_array, memfile) |
| return memfile.getvalue() |
| |
| def test_end2end(self): |
| with TempDir() as temp_dir: |
| file_path_prefix = temp_dir.create_temp_file('result') |
| |
| # Generate a TFRecord file. |
| with TestPipeline() as p: |
| expected_data = [self.create_inputs() for _ in range(0, 10)] |
| _ = p | beam.Create(expected_data) | WriteToTFRecord(file_path_prefix) |
| |
| # Read the file back and compare. |
| with TestPipeline() as p: |
| actual_data = p | ReadFromTFRecord(file_path_prefix + '-*') |
| assert_that(actual_data, equal_to(expected_data)) |
| |
| def test_end2end_auto_compression(self): |
| with TempDir() as temp_dir: |
| file_path_prefix = temp_dir.create_temp_file('result') |
| |
| # Generate a TFRecord file. |
| with TestPipeline() as p: |
| expected_data = [self.create_inputs() for _ in range(0, 10)] |
| _ = p | beam.Create(expected_data) | WriteToTFRecord( |
| file_path_prefix, file_name_suffix='.gz') |
| |
| # Read the file back and compare. |
| with TestPipeline() as p: |
| actual_data = p | ReadFromTFRecord(file_path_prefix + '-*') |
| assert_that(actual_data, equal_to(expected_data)) |
| |
| def test_end2end_auto_compression_unsharded(self): |
| with TempDir() as temp_dir: |
| file_path_prefix = temp_dir.create_temp_file('result') |
| |
| # Generate a TFRecord file. |
| with TestPipeline() as p: |
| expected_data = [self.create_inputs() for _ in range(0, 10)] |
| _ = p | beam.Create(expected_data) | WriteToTFRecord( |
| file_path_prefix + '.gz', shard_name_template='') |
| |
| # Read the file back and compare. |
| with TestPipeline() as p: |
| actual_data = p | ReadFromTFRecord(file_path_prefix + '.gz') |
| assert_that(actual_data, equal_to(expected_data)) |
| |
| @unittest.skipIf(tf is None, 'tensorflow not installed.') |
| def test_end2end_example_proto(self): |
| with TempDir() as temp_dir: |
| file_path_prefix = temp_dir.create_temp_file('result') |
| |
| example = tf.train.Example() |
| example.features.feature['int'].int64_list.value.extend(list(range(3))) |
| example.features.feature['bytes'].bytes_list.value.extend( |
| [b'foo', b'bar']) |
| |
| with TestPipeline() as p: |
| _ = p | beam.Create([example]) | WriteToTFRecord( |
| file_path_prefix, coder=beam.coders.ProtoCoder(example.__class__)) |
| |
| # Read the file back and compare. |
| with TestPipeline() as p: |
| actual_data = ( |
| p | ReadFromTFRecord( |
| file_path_prefix + '-*', |
| coder=beam.coders.ProtoCoder(example.__class__))) |
| assert_that(actual_data, equal_to([example])) |
| |
| def test_end2end_read_write_read(self): |
| with TempDir() as temp_dir: |
| path = temp_dir.create_temp_file('result') |
| with TestPipeline() as p: |
| # Initial read to validate the pipeline doesn't fail before the file is |
| # created. |
| _ = p | ReadFromTFRecord(path + '-*', validate=False) |
| expected_data = [self.create_inputs() for _ in range(0, 10)] |
| _ = p | beam.Create(expected_data) | WriteToTFRecord( |
| path, file_name_suffix='.gz') |
| |
| # Read the file back and compare. |
| with TestPipeline() as p: |
| actual_data = p | ReadFromTFRecord(path + '-*', validate=True) |
| assert_that(actual_data, equal_to(expected_data)) |
| |
| |
| 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)) |
| | "encode" >> beam.Map(lambda s: json.dumps(s).encode('utf-8'))) |
| #TFrecordIO |
| output2 = output | 'WriteToTFRecord' >> beam.io.WriteToTFRecord( |
| file_path_prefix=self.tempdir + "/ouput_WriteToTFRecord", |
| file_name_suffix=".tfrecord", |
| num_shards=num_shards, |
| ) |
| _ = output2 | 'LogElements after WriteToTFRecord' >> LogElements( |
| prefix='after WriteToTFRecord ', with_window=True, level=logging.INFO) |
| |
| # Regex to match the expected windowed file pattern |
| # Example: |
| # ouput_WriteToTFRecord-[1614556800.0, 1614556805.0)-00000-of-00002.tfrecord |
| # 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})\.tfrecord$') |
| pattern = re.compile(pattern_string) |
| file_names = [] |
| for file_name in glob.glob(self.tempdir + '/ouput_WriteToTFRecord*'): |
| 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() |
| | "encode" >> beam.Map(lambda s: json.dumps(s).encode('utf-8'))) |
| #TFrecordIO |
| output2 = output | 'WriteToTFRecord' >> beam.io.WriteToTFRecord( |
| file_path_prefix=self.tempdir + "/ouput_WriteToTFRecord", |
| file_name_suffix=".tfrecord", |
| shard_name_template=shard_name_template, |
| num_shards=num_shards, |
| triggering_frequency=60, |
| ) |
| _ = output2 | 'LogElements after WriteToTFRecord' >> LogElements( |
| prefix='after WriteToTFRecord ', with_window=True, level=logging.INFO) |
| |
| # Regex to match the expected windowed file pattern |
| # Example: |
| # ouput_WriteToTFRecord-[2021-03-01T00-00-00, 2021-03-01T00-01-00)- |
| # 00000-of-00002.tfrecord |
| # 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})\.tfrecord$') |
| pattern = re.compile(pattern_string) |
| file_names = [] |
| for file_name in glob.glob(self.tempdir + '/ouput_WriteToTFRecord*'): |
| 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() |
| | "encode" >> beam.Map(lambda s: json.dumps(s).encode('utf-8'))) |
| #TFrecordIO |
| output2 = output | 'WriteToTFRecord' >> beam.io.WriteToTFRecord( |
| file_path_prefix=self.tempdir + "/ouput_WriteToTFRecord", |
| file_name_suffix=".tfrecord", |
| shard_name_template=shard_name_template, |
| num_shards=num_shards, |
| triggering_frequency=triggering_frequency, |
| ) |
| _ = output2 | 'LogElements after WriteToTFRecord' >> LogElements( |
| prefix='after WriteToTFRecord ', with_window=True, level=logging.INFO) |
| |
| # Regex to match the expected windowed file pattern |
| # Example: |
| # ouput_WriteToTFRecord-[2021-03-01T00-00-00, 2021-03-01T00-01-00)- |
| # 00000-of-00002.tfrecord |
| # 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})\.tfrecord$') |
| pattern = re.compile(pattern_string) |
| file_names = [] |
| for file_name in glob.glob(self.tempdir + '/ouput_WriteToTFRecord*'): |
| 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() |
