def test_read_do_write_with_start_bundle(self): input_path = self.create_temp_file('01234567890123456789\n0123456789') output_path = '%s.out' % input_path finish_path = '%s.finish' % input_path executor.MapTaskExecutor().execute(make_map_task([ maptask.WorkerRead( fileio.TextFileSource(file_path=input_path, start_offset=0, end_offset=15, strip_trailing_newlines=True, coder=coders.StrUtf8Coder()), output_coders=[self.OUTPUT_CODER]), maptask.WorkerDoFn(serialized_fn=pickle_with_side_inputs( DoFnUsingStartBundle(finish_path)), output_tags=['out'], output_coders=[self.OUTPUT_CODER], input=(0, 0), side_inputs=None), maptask.WorkerWrite( fileio.TextFileSink(file_path_prefix=output_path, append_trailing_newlines=True, coder=coders.ToStringCoder()), input=(1, 0), output_coders=(coders.ToStringCoder(),)) ])) with open(output_path) as f: self.assertEqual('XYZ: 01234567890123456789\n', f.read()) # Check that the finish_bundle method of the custom DoFn object left the # expected side-effect by writing a file with a specific content. with open(finish_path) as f: self.assertEqual('finish called.', f.read())
def test_create_do_write(self): output_path = self.create_temp_file('n/a') elements = ['abc', 'def', 'ghi'] executor.MapTaskExecutor().execute(make_map_task([ maptask.WorkerRead( inmemory.InMemorySource( elements=[pickler.dumps(e) for e in elements], # Start at the last element. start_index=2, # Go beyond the end to test that case is handled. end_index=15), output_coders=[coders.ToStringCoder()]), maptask.WorkerDoFn(serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn(lambda x: ['XYZ: %s' % x])), output_tags=['out'], output_coders=[self.OUTPUT_CODER], input=(0, 0), side_inputs=None), maptask.WorkerWrite( fileio.TextFileSink(file_path_prefix=output_path, append_trailing_newlines=True, coder=coders.ToStringCoder()), input=(1, 0), output_coders=(coders.ToStringCoder(),)) ])) with open(output_path) as f: self.assertEqual('XYZ: ghi\n', f.read())
def test_shuffle_read_do_write(self): output_path = self.create_temp_file('n/a') work_spec = [ maptask.WorkerGroupingShuffleRead(shuffle_reader_config='none', start_shuffle_position='aaa', end_shuffle_position='zzz', coder=self.SHUFFLE_CODER, output_coders=[self.SHUFFLE_CODER]), maptask.WorkerDoFn(serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn( lambda (k, vs): [str((k, v)) for v in vs])), output_tags=['out'], output_coders=[self.OUTPUT_CODER], input=(0, 0), side_inputs=None), maptask.WorkerWrite( fileio.TextFileSink(file_path_prefix=output_path, append_trailing_newlines=True, coder=coders.ToStringCoder()), input=(1, 0), output_coders=(coders.ToStringCoder(),)) ] shuffle_source_mock = mock.MagicMock() shuffle_source_mock.reader().__enter__().__iter__.return_value = [ (10, [1, 2]), (20, [3])] executor.MapTaskExecutor().execute( make_map_task(work_spec), test_shuffle_source=shuffle_source_mock) with open(output_path) as f: self.assertEqual('(10, 1)\n(10, 2)\n(20, 3)\n', f.read())
def test_read_do_write(self): input_path = self.create_temp_file('01234567890123456789\n0123456789') output_path = '%s.out' % input_path executor.MapTaskExecutor().execute(make_map_task([ maptask.WorkerRead( fileio.TextFileSource(file_path=input_path, start_offset=0, end_offset=15, strip_trailing_newlines=True, coder=coders.StrUtf8Coder()), output_coders=[self.OUTPUT_CODER]), maptask.WorkerDoFn(serialized_fn=pickle_with_side_inputs( ptransform.CallableWrapperDoFn(lambda x: ['XYZ: %s' % x])), output_tags=['out'], output_coders=[self.OUTPUT_CODER], input=(0, 0), side_inputs=None), maptask.WorkerWrite( fileio.TextFileSink(file_path_prefix=output_path, append_trailing_newlines=True, coder=coders.ToStringCoder()), input=(1, 0), output_coders=(coders.ToStringCoder(),)) ])) with open(output_path) as f: self.assertEqual('XYZ: 01234567890123456789\n', f.read())
def make_text_sink(output_path, input, coder=coders.ToStringCoder()): return maptask.WorkerWrite(fileio.NativeTextFileSink( file_path_prefix=output_path, append_trailing_newlines=True, coder=coder), input=input, output_coders=(coder, ))
def test_file_sink_writing(self): temp_path = tempfile.NamedTemporaryFile().name sink = MyFileSink(temp_path, file_name_suffix='.foo', coder=coders.ToStringCoder()) # Manually invoke the generic Sink API. init_token = sink.initialize_write() writer1 = sink.open_writer(init_token, '1') writer1.write('a') writer1.write('b') res1 = writer1.close() writer2 = sink.open_writer(init_token, '2') writer2.write('x') writer2.write('y') writer2.write('z') res2 = writer2.close() res = list(sink.finalize_write(init_token, [res1, res2])) # Retry the finalize operation (as if the first attempt was lost). res = list(sink.finalize_write(init_token, [res1, res2])) # Check the results. shard1 = temp_path + '-00000-of-00002.foo' shard2 = temp_path + '-00001-of-00002.foo' self.assertEqual(res, [shard1, shard2]) self.assertEqual(open(shard1).read(), '[start][a][b][end]') self.assertEqual(open(shard2).read(), '[start][x][y][z][end]') # Check that any temp files are deleted. self.assertItemsEqual([shard1, shard2], glob.glob(temp_path + '*'))
def test_file_sink_multi_shards(self): temp_path = tempfile.NamedTemporaryFile().name sink = MyFileSink(temp_path, file_name_suffix='.foo', coder=coders.ToStringCoder()) # Manually invoke the generic Sink API. init_token = sink.initialize_write() num_shards = 1000 writer_results = [] for i in range(num_shards): uuid = 'uuid-%05d' % i writer = sink.open_writer(init_token, uuid) writer.write('a') writer.write('b') writer.write(uuid) writer_results.append(writer.close()) res_first = list(sink.finalize_write(init_token, writer_results)) # Retry the finalize operation (as if the first attempt was lost). res_second = list(sink.finalize_write(init_token, writer_results)) self.assertItemsEqual(res_first, res_second) res = sorted(res_second) for i in range(num_shards): shard_name = '%s-%05d-of-%05d.foo' % (temp_path, i, num_shards) uuid = 'uuid-%05d' % i self.assertEqual(res[i], shard_name) self.assertEqual( open(shard_name).read(), ('[start][a][b][%s][end]' % uuid)) # Check that any temp files are deleted. self.assertItemsEqual(res, glob.glob(temp_path + '*'))
def __init__(self, file_path_prefix, append_trailing_newlines=True, file_name_suffix='', num_shards=0, shard_name_template=None, validate=True, coder=coders.ToStringCoder()): # We initialize a file_path attribute containing just the prefix part for # local runner environment. For now, sharding is not supported in the local # runner and sharding options (template, num, suffix) are ignored. # The attribute is also used in the worker environment when we just write # to a specific file. # TODO(silviuc): Add support for file sharding in the local runner. self.file_path = file_path_prefix self.append_trailing_newlines = append_trailing_newlines self.coder = coder self.is_gcs_sink = self.file_path.startswith('gs://') self.file_name_prefix = file_path_prefix self.file_name_suffix = file_name_suffix self.num_shards = num_shards # TODO(silviuc): Update this when the service supports more patterns. self.shard_name_template = ('-SSSSS-of-NNNNN' if shard_name_template is None else shard_name_template) # TODO(silviuc): Implement sink validation. self.validate = validate
def test_empty_write(self): temp_path = tempfile.NamedTemporaryFile().name sink = MyFileSink(temp_path, file_name_suffix='.foo', coder=coders.ToStringCoder()) p = df.Pipeline('DirectPipelineRunner') p | df.Create([]) | df.io.Write(sink) # pylint: disable=expression-not-assigned p.run() self.assertEqual( open(temp_path + '-00000-of-00001.foo').read(), '[start][end]')
def __init__(self, file_path_prefix, file_name_suffix='', coder=coders.ToStringCoder(), append_trailing_newlines=True): super(PureTextFileSink, self).__init__(file_path_prefix, file_name_suffix=file_name_suffix, coder=coder, mime_type='text/plain') self.append_trailing_newlines = append_trailing_newlines
def test_ungrouped_shuffle_read_and_write(self): output_path = self.create_temp_file('n/a') work_spec = [ maptask.WorkerUngroupedShuffleRead(shuffle_reader_config='none', start_shuffle_position='aaa', end_shuffle_position='zzz', coder=self.SHUFFLE_CODER, output_coders=[self.SHUFFLE_CODER]), maptask.WorkerWrite( fileio.TextFileSink(file_path_prefix=output_path, append_trailing_newlines=True, coder=coders.ToStringCoder()), input=(0, 0), output_coders=(coders.ToStringCoder(),)) ] shuffle_source_mock = mock.MagicMock() shuffle_source_mock.reader().__enter__().__iter__.return_value = [1, 2, 3] executor.MapTaskExecutor().execute( make_map_task(work_spec), test_shuffle_source=shuffle_source_mock) with open(output_path) as f: self.assertEqual('1\n2\n3\n', f.read())
def test_fixed_shard_write(self): temp_path = tempfile.NamedTemporaryFile().name sink = MyFileSink(temp_path, file_name_suffix='.foo', num_shards=3, shard_name_template='_NN_SSS_', coder=coders.ToStringCoder()) p = df.Pipeline('DirectPipelineRunner') p | df.Create(['a', 'b']) | df.io.Write(sink) # pylint: disable=expression-not-assigned p.run() concat = ''.join( open(temp_path + '_03_%03d_.foo' % shard_num).read() for shard_num in range(3)) self.assertTrue('][a][' in concat, concat) self.assertTrue('][b][' in concat, concat)
def test_file_sink_io_error(self): temp_path = tempfile.NamedTemporaryFile().name sink = MyFileSink(temp_path, file_name_suffix='.foo', coder=coders.ToStringCoder()) # Manually invoke the generic Sink API. init_token = sink.initialize_write() writer1 = sink.open_writer(init_token, '1') writer1.write('a') writer1.write('b') res1 = writer1.close() writer2 = sink.open_writer(init_token, '2') writer2.write('x') writer2.write('y') writer2.write('z') res2 = writer2.close() os.remove(res2) with self.assertRaises(IOError): list(sink.finalize_write(init_token, [res1, res2]))
def __init__( self, file_path_prefix, file_name_suffix='', append_trailing_newlines=True, num_shards=0, shard_name_template=None, coder=coders.ToStringCoder(), compression_type=CompressionTypes.NO_COMPRESSION, ): """Initialize a TextFileSink. Args: file_path_prefix: The file path to write to. The files written will begin with this prefix, followed by a shard identifier (see num_shards), and end in a common extension, if given by file_name_suffix. In most cases, only this argument is specified and num_shards, shard_name_template, and file_name_suffix use default values. file_name_suffix: Suffix for the files written. append_trailing_newlines: indicate whether this sink should write an additional newline char after writing each element. num_shards: The number of files (shards) used for output. If not set, the service will decide on the optimal number of shards. 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. shard_name_template: A template string containing placeholders for the shard number and shard count. Currently only '' and '-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'. coder: Coder used to encode each line. compression_type: Type of compression to use for this sink. Raises: TypeError: if file path parameters are not a string or if compression_type is not member of CompressionTypes. ValueError: if shard_name_template is not of expected format. Returns: A TextFileSink object usable for writing. """ if not isinstance(file_path_prefix, basestring): raise TypeError( 'TextFileSink: file_path_prefix must be a string; got %r instead' % file_path_prefix) if not isinstance(file_name_suffix, basestring): raise TypeError( 'TextFileSink: file_name_suffix must be a string; got %r instead' % file_name_suffix) if not CompressionTypes.valid_compression_type(compression_type): raise TypeError( 'compression_type must be CompressionType object but ' 'was %s' % type(compression_type)) if compression_type == CompressionTypes.DEFLATE: mime_type = 'application/x-gzip' else: mime_type = 'text/plain' super(TextFileSink, self).__init__(file_path_prefix, file_name_suffix=file_name_suffix, num_shards=num_shards, shard_name_template=shard_name_template, coder=coder, mime_type=mime_type) self.compression_type = compression_type self.append_trailing_newlines = append_trailing_newlines
def __init__(self, file_path_prefix, append_trailing_newlines=True, file_name_suffix='', num_shards=0, shard_name_template=None, validate=True, coder=coders.ToStringCoder()): """Initialize a TextSink. Args: file_path_prefix: The file path to write to. The files written will begin with this prefix, followed by a shard identifier (see num_shards), and end in a common extension, if given by file_name_suffix. In most cases, only this argument is specified and num_shards, shard_name_template, and file_name_suffix use default values. append_trailing_newlines: indicate whether this sink should write an additional newline char after writing each element. file_name_suffix: Suffix for the files written. num_shards: The number of files (shards) used for output. If not set, the service will decide on the optimal number of shards. 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. shard_name_template: A template string containing placeholders for the shard number and shard count. Currently only '' and '-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'. validate: Enable path validation on pipeline creation. coder: Coder used to encode each line. Raises: TypeError: if file_path is not a string. ValueError: if shard_name_template is not of expected format. """ if not isinstance(file_path_prefix, basestring): raise TypeError( '%s: file_path_prefix must be a string; got %r instead' % (self.__class__.__name__, file_path_prefix)) if not isinstance(file_name_suffix, basestring): raise TypeError( '%s: file_name_suffix must be a string; got %r instead' % (self.__class__.__name__, file_name_suffix)) # We initialize a file_path attribute containing just the prefix part for # local runner environment. For now, sharding is not supported in the local # runner and sharding options (template, num, suffix) are ignored. # The attribute is also used in the worker environment when we just write # to a specific file. # TODO(silviuc): Add support for file sharding in the local runner. self.file_path = file_path_prefix self.append_trailing_newlines = append_trailing_newlines self.coder = coder self.is_gcs_sink = self.file_path.startswith('gs://') self.file_name_prefix = file_path_prefix self.file_name_suffix = file_name_suffix self.num_shards = num_shards # TODO(silviuc): Update this when the service supports more patterns. if shard_name_template not in (None, '', '-SSSSS-of-NNNNN'): raise ValueError( 'The shard_name_template argument must be an empty string or the ' 'pattern -SSSSS-of-NNNNN instead of %s' % shard_name_template) self.shard_name_template = ( shard_name_template if shard_name_template is not None else '-SSSSS-of-NNNNN') # TODO(silviuc): Implement sink validation. self.validate = validate
def TextFileSink(file_path_prefix, # pylint: disable=invalid-name append_trailing_newlines=True, file_name_suffix='', num_shards=0, shard_name_template=None, validate=True, coder=coders.ToStringCoder()): """Initialize a TextSink. Args: file_path_prefix: The file path to write to. The files written will begin with this prefix, followed by a shard identifier (see num_shards), and end in a common extension, if given by file_name_suffix. In most cases, only this argument is specified and num_shards, shard_name_template, and file_name_suffix use default values. append_trailing_newlines: indicate whether this sink should write an additional newline char after writing each element. file_name_suffix: Suffix for the files written. num_shards: The number of files (shards) used for output. If not set, the service will decide on the optimal number of shards. 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. shard_name_template: A template string containing placeholders for the shard number and shard count. Currently only '' and '-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'. validate: Enable path validation on pipeline creation. coder: Coder used to encode each line. Raises: TypeError: if file_path is not a string. ValueError: if shard_name_template is not of expected format. Returns: A TextFileSink object usable for writing. """ if not isinstance(file_path_prefix, basestring): raise TypeError( 'TextFileSink: file_path_prefix must be a string; got %r instead' % file_path_prefix) if not isinstance(file_name_suffix, basestring): raise TypeError( 'TextFileSink: file_name_suffix must be a string; got %r instead' % file_name_suffix) if shard_name_template not in (None, '', '-SSSSS-of-NNNNN'): raise ValueError( 'The shard_name_template argument must be an empty string or the ' 'pattern -SSSSS-of-NNNNN instead of %s' % shard_name_template) if shard_name_template == '': # pylint: disable=g-explicit-bool-comparison num_shards = 1 if num_shards: return NativeTextFileSink(file_path_prefix, append_trailing_newlines=append_trailing_newlines, file_name_suffix=file_name_suffix, num_shards=num_shards, shard_name_template=shard_name_template, validate=validate, coder=coder) else: return PureTextFileSink(file_path_prefix, append_trailing_newlines=append_trailing_newlines, file_name_suffix=file_name_suffix, coder=coder)