How do I add headers for the output csv for apache beam dataflow?

Question

I noticed in the java SDK, there is a function that allows you to write the headers of a csv file. https://cloud.google.com/dataflow/java-sdk/JavaDoc/com/google/cloud/dataflow/sdk/io/TextIO.Write.html#withHeader-java.lang.String-

Is this feature mirrored on the python SDK?

Answer 1

You can now write to a text and specify a header using the text sink.

From the documentation:

class apache_beam.io.textio.WriteToText(
    file_path_prefix, file_name_suffix='',
    append_trailing_newlines=True, num_shards=0, 
    shard_name_template=None, coder=ToStringCoder, 
    compression_type='auto', header=None)

So you can use the following piece of code:

beam.io.WriteToText(bucket_name, file_name_suffix='.csv', 
    header='colname1, colname2')

The complete documentation is available here if you want details or check how it is implemented: https://beam.apache.org/documentation/sdks/pydoc/2.0.0/_modules/apache_beam/io/textio.html#WriteToText

Answer 2

This is not implemented at this moment. However you can implement/extend it yourself (see attached notebook for an example+test with my version of apache_beam).

This is based on a note in the docstring of the superclass FileSink, mentioning that you should overwrite the open function:

The new class that works for my version of apache_beam ('0.3.0-incubating.dev'):

import apache_beam as beam
from apache_beam.io import TextFileSink
from apache_beam.io.fileio import ChannelFactory,CompressionTypes
from apache_beam import coders


class TextFileSinkWithHeader(TextFileSink):
    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,
               header=None):
        super(TextFileSinkWithHeader, self).__init__(
            file_path_prefix,
            file_name_suffix=file_name_suffix,
            num_shards=num_shards,
            shard_name_template=shard_name_template,
            coder=coder,

            compression_type=compression_type,
            append_trailing_newlines=append_trailing_newlines)
        self.header = header

    def open(self, temp_path):
        channel_factory = ChannelFactory.open(
            temp_path,
            'wb',
            mime_type=self.mime_type)
        channel_factory.write(self.header+"\n")
        return channel_factory

You can subsequently use it as follows:

beam.io.Write(TextFileSinkWithHeader('./names_w_headers',header="names"))

See the notebook for the complete overview.

Answer 3

This feature does not yet exist in the Python SDK

Answer 4

For the Python SDK:

beam.io.Write(beam.io.WriteToText(
                     file_path_prefix=os.path.join(OUTPUT_DIR),
                     file_name_suffix='.csv', header='colname1, colname2')
              )