learning/prompts/code-generation/java/06_io_tfrecord.md

Prompt:
Write a sample Java code snippet that writes data to a TFRecord file using Apache Beam.

Response:
Your Apache Beam pipeline can write data to a TFRecord file using the Apache Beam TFRecordIO connector.

Here is an example of how to use the TFRecordIO connector to accomplish this:

```java
package tfrecord;

import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.Compression;
import org.apache.beam.sdk.io.TFRecordIO;
import org.apache.beam.sdk.options.Description;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.Validation;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Charsets;

import java.util.Arrays;
import java.util.List;

// Pipeline to write data to a TFRecord file using Apache Beam.
public class WriteTFRecordFile {

    // Pipeline options for writing data to a TFRecord file.
    public interface WriteTFRecordFileOptions extends PipelineOptions {

        @Description("File path prefix to write TFRecords files to")
        @Validation.Required
        String getFilePathPrefix();

        void setFilePathPrefix(String filePathPrefix);
    }

    // Main method to run the pipeline
    public static void main(String[] args) {

        // Parse the pipeline options from the command line
        WriteTFRecordFileOptions options =
                PipelineOptionsFactory.fromArgs(args)
                        .withValidation().as(WriteTFRecordFileOptions.class);

        // Create the pipeline
        Pipeline p = Pipeline.create(options);

        // Create a list of sample records
        List<String> records = Arrays.asList(
            "Record 1", "Record 2", "Record 3", "Record 4", "Record 5"
        );
        // Create a PCollection of sample records
        p.apply("Create collection of records", Create.of(records))
            // Convert each record to bytes using a DoFn
            .apply(
                "Convert records to bytes",
                ParDo.of(
                    new DoFn<String, byte[]>() {
                        @DoFn.ProcessElement
                        public void processElement(ProcessContext c) {
                            c.output(c.element().getBytes(Charsets.UTF_8));
                        }
                    }))
            // Write the records to a TFRecord file
            .apply(
                "Write to TFRecord file",
                    TFRecordIO.write()
                        .to(options.getFilePathPrefix())
                        .withCompression(Compression.UNCOMPRESSED)
                        .withNumShards(1));
        // Run the pipeline
        p.run();
    }
}
```

This code snippet utilizes the pipeline options pattern to parse command-line arguments.