com.wtanaka.beam.StdoutIO Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of beam Show documentation
Apache Beam Utility Code
There is a newer version: v0.0.0-128-g8bec6d4
/*
 * com.wtanaka.beam
 *
 * Copyright (C) 2017 Wesley Tanaka 
 *
 * This program is free software: you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see
 * .
 */
package com.wtanaka.beam;

import java.io.IOException;
import java.io.OutputStream;
import java.io.Serializable;

import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.io.Sink;
import org.apache.beam.sdk.io.Write;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PDone;

/**
 * Transform for outputting to System.out
 */
class StdoutIO
{
   /**
    * Sink implementation for Stdout
    */
   static class StdoutSink extends Sink
   {
      private static final long serialVersionUID = 1L;
      private final OutputStream m_serializableOutputStream;

      static class StdoutWriter extends Writer
      {
         private final WriteOperation m_writeOperation;
         private final OutputStream m_outputStream;

         StdoutWriter(
            final WriteOperation writeOperation,
            final OutputStream outputStream)
         {
            m_writeOperation = writeOperation;
            m_outputStream = outputStream;
         }

         @Override
         public Void close()
         {
            return null;
         }

         @Override
         public WriteOperation getWriteOperation()
         {
            return m_writeOperation;
         }

         @Override
         public void open(final String uId)
         {
         }

         @Override
         public void write(final byte[] value) throws IOException
         {
            m_outputStream.write(value);
            // System.out.write((int) '\n');
         }
      }

      static class StdoutWriteOperation extends Sink
         .WriteOperation
      {
         private static final long serialVersionUID = 1L;
         private final OutputStream m_serializableOutputStream;
         private final Sink m_sink;

         StdoutWriteOperation(final Sink sink,
                              final OutputStream serializableOutputStream)
         {
            m_sink = sink;
            assert serializableOutputStream == null ||
               (serializableOutputStream instanceof Serializable) :
               "Stream " + serializableOutputStream + " must be Serializable";
            m_serializableOutputStream = serializableOutputStream;
         }

         @Override
         public Writer createWriter(
            final PipelineOptions options)
         {
            return new StdoutSink.StdoutWriter(this,
               getStream());
         }

         @Override
         public void finalize(final Iterable writerResults,
                              final PipelineOptions options)
            throws IOException
         {
            getStream().flush();
         }

         @Override
         public Sink getSink()
         {
            return m_sink;
         }

         private OutputStream getStream()
         {
            if (m_serializableOutputStream == null)
            {
               return System.out;
            }
            else
            {
               return m_serializableOutputStream;
            }
         }

         @Override
         public Coder getWriterResultCoder()
         {
            return VoidCoder.of();
         }

         @Override
         public void initialize(final PipelineOptions options)
         {
         }
      }

      StdoutSink()
      {
         m_serializableOutputStream = null;
      }

      StdoutSink(
         final OutputStream serializableOutput)
      {
         m_serializableOutputStream = serializableOutput;
      }

      @Override
      public WriteOperation createWriteOperation(
         final PipelineOptions options)
      {
         return new StdoutSink.StdoutWriteOperation(this,
            m_serializableOutputStream);
      }

      @Override
      public void validate(final PipelineOptions options)
      {

      }
   }

   public static PTransform, PDone> write()
   {
      return Write.to(new StdoutSink());
   }
}