LineIO Example

LineIO.java

/*
 * Copyright (C) 2015 Google Inc.
 *
 * Licensed 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.
 */

import com.google.cloud.dataflow.sdk.coders.Coder;
import com.google.cloud.dataflow.sdk.coders.StringUtf8Coder;
import com.google.cloud.dataflow.sdk.io.FileBasedSource;
import com.google.cloud.dataflow.sdk.io.Read;
import com.google.cloud.dataflow.sdk.options.PipelineOptions;
import com.google.cloud.dataflow.sdk.util.CoderUtils;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.ReadableByteChannel;
import java.nio.channels.SeekableByteChannel;
import java.util.NoSuchElementException;

/**
 * A source that reads lines from text files.
 *
 * <p>As of writing, this is basically a copy from {@code FileBasedSourceTest}
 * intended for testing dynamic work rebalancing with custom sources.
 */
public class LineIO {
  // Match TextIO.
  public static Read.Bounded<String> readFilepattern(String filepattern) {
    return Read.from(new LineSource(filepattern, 1));
  }

  private static class LineSource extends FileBasedSource<String> {
    public LineSource(String fileOrPattern, long minBundleSize) {
      super(fileOrPattern, minBundleSize);
    }

    public LineSource(String filename, long minBundleSize, long startOffset, long endOffset) {
      super(filename, minBundleSize, startOffset, endOffset);
    }

    @Override
    public boolean producesSortedKeys(PipelineOptions options) throws Exception {
      return false;
    }

    @Override
    public void validate() {}

    @Override
    public Coder<String> getDefaultOutputCoder() {
      return StringUtf8Coder.of();
    }

    @Override
    public FileBasedSource<String> createForSubrangeOfFile(String fileName, long start, long end) {
      return new LineSource(fileName, getMinBundleSize(), start, end);
    }

    @Override
    public FileBasedReader<String> createSingleFileReader(PipelineOptions options) {
      return new LineReader(this);
    }
  }

  /**
   * A reader that can read lines of text from a {@link LineSource}.
   */
  private static class LineReader extends FileBasedSource.FileBasedReader<String> {
    private static final Logger LOG = LoggerFactory.getLogger(LineReader.class);
    private ReadableByteChannel channel = null;
    private final byte boundary;
    private long nextOffset = 0;
    private long currentOffset = 0;
    private boolean isAtSplitPoint = false;
    private final ByteBuffer buf;
    private static final int BUF_SIZE = 1024;
    private String currentValue = null;

    public LineReader(LineSource source) {
      super(source);
      boundary = '\n';
      buf = ByteBuffer.allocate(BUF_SIZE);
      buf.flip();
    }

    private int readNextLine(ByteArrayOutputStream out) throws IOException {
      int byteCount = 0;
      while (true) {
        if (!buf.hasRemaining()) {
          buf.clear();
          int read = channel.read(buf);
          if (read < 0) {
            break;
          }
          buf.flip();
        }
        byte b = buf.get();
        byteCount++;
        if (b == boundary) {
          break;
        }
        out.write(b);
      }
      return byteCount;
    }

    @Override
    protected void startReading(ReadableByteChannel channel) throws IOException {
      this.channel = channel;
      if (getCurrentSource().getMode() == FileBasedSource.Mode.SINGLE_FILE_OR_SUBRANGE) {
        // If we are not at the beginning of a line, we should ignore the current line.
        if (getCurrentSource().getStartOffset() > 0) {
          SeekableByteChannel seekChannel = (SeekableByteChannel) channel;
          // Start from one character back and read till we find a new line.
          seekChannel.position(seekChannel.position() - 1);
          nextOffset = seekChannel.position() + readNextLine(new ByteArrayOutputStream());
        }
      }
    }

    @Override
    protected boolean readNextRecord() throws IOException {
      currentOffset = nextOffset;

      ByteArrayOutputStream buf = new ByteArrayOutputStream();
      int offsetAdjustment = readNextLine(buf);
      if (offsetAdjustment == 0) {
        // EOF
        return false;
      }
      nextOffset += offsetAdjustment;
      isAtSplitPoint = true;
      currentValue = CoderUtils.decodeFromByteArray(StringUtf8Coder.of(), buf.toByteArray());
      return true;
    }

    @Override
    protected boolean isAtSplitPoint() {
      return isAtSplitPoint;
    }

    @Override
    protected long getCurrentOffset() {
      return currentOffset;
    }

    @Override
    public String getCurrent() throws NoSuchElementException {
      return currentValue;
    }
  }

}

Fine to share for now, obviously needs commenting and cleanup before we import it to SDK.