lang/java/tools/src/main/java/org/apache/avro/tool/CatTool.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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
 *
 *     https://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.
 */
package org.apache.avro.tool;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.PrintStream;
import java.util.List;

import joptsimple.OptionParser;
import joptsimple.OptionSet;
import joptsimple.OptionSpec;

import org.apache.avro.Schema;
import org.apache.avro.file.CodecFactory;
import org.apache.avro.file.DataFileConstants;
import org.apache.avro.file.DataFileStream;
import org.apache.avro.file.DataFileWriter;
import org.apache.avro.generic.GenericDatumReader;
import org.apache.avro.generic.GenericDatumWriter;
import org.apache.avro.generic.GenericRecord;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.Path;

/** Tool to extract samples from an Avro data file. */
public class CatTool implements Tool {

  private long totalCopied;
  private double sampleCounter;

  private GenericRecord reuse;
  private DataFileStream<GenericRecord> reader;
  private DataFileWriter<GenericRecord> writer;
  private Schema schema;
  private List<Path> inFiles;
  private int currentInput;

  @Override
  public int run(InputStream in, PrintStream out, PrintStream err, List<String> args) throws Exception {
    OptionParser optParser = new OptionParser();
    OptionSpec<Long> offsetOpt = optParser.accepts("offset", "offset for reading input").withRequiredArg()
        .ofType(Long.class).defaultsTo(Long.valueOf(0));
    OptionSpec<Long> limitOpt = optParser.accepts("limit", "maximum number of records in the outputfile")
        .withRequiredArg().ofType(Long.class).defaultsTo(Long.MAX_VALUE);
    OptionSpec<Double> fracOpt = optParser.accepts("samplerate", "rate at which records will be collected")
        .withRequiredArg().ofType(Double.class).defaultsTo(Double.valueOf(1));

    OptionSet opts = optParser.parse(args.toArray(new String[0]));
    List<String> nargs = (List<String>) opts.nonOptionArguments();
    if (nargs.size() < 2) {
      printHelp(out);
      return 0;
    }

    inFiles = Util.getFiles(nargs.subList(0, nargs.size() - 1));

    System.out.println("List of input files:");
    for (Path p : inFiles) {
      System.out.println(p);
    }
    currentInput = -1;
    nextInput();

    OutputStream output = out;
    String lastArg = nargs.get(nargs.size() - 1);
    if (nargs.size() > 1 && !lastArg.equals("-")) {
      output = Util.createFromFS(lastArg);
    }
    writer = new DataFileWriter<>(new GenericDatumWriter<>());

    String codecName = reader.getMetaString(DataFileConstants.CODEC);
    CodecFactory codec = (codecName == null) ? CodecFactory.fromString(DataFileConstants.NULL_CODEC)
        : CodecFactory.fromString(codecName);
    writer.setCodec(codec);
    for (String key : reader.getMetaKeys()) {
      if (!DataFileWriter.isReservedMeta(key)) {
        writer.setMeta(key, reader.getMeta(key));
      }
    }
    writer.create(schema, output);

    long offset = opts.valueOf(offsetOpt);
    long limit = opts.valueOf(limitOpt);
    double samplerate = opts.valueOf(fracOpt);
    sampleCounter = 1;
    totalCopied = 0;
    reuse = null;

    if (limit < 0) {
      System.out.println("limit has to be non-negative");
      this.printHelp(out);
      return 1;
    }
    if (offset < 0) {
      System.out.println("offset has to be non-negative");
      this.printHelp(out);
      return 1;
    }
    if (samplerate < 0 || samplerate > 1) {
      System.out.println("samplerate has to be a number between 0 and 1");
      this.printHelp(out);
      return 1;
    }

    skip(offset);
    writeRecords(limit, samplerate);
    System.out.println(totalCopied + " records written.");

    writer.flush();
    writer.close();
    Util.close(out);
    return 0;
  }

  private void nextInput() throws IOException {
    currentInput++;
    Path path = inFiles.get(currentInput);
    FSDataInputStream input = new FSDataInputStream(Util.openFromFS(path));
    reader = new DataFileStream<>(input, new GenericDatumReader<>());
    if (schema == null) { // if this is the first file, the schema gets saved
      schema = reader.getSchema();
    } else if (!schema.equals(reader.getSchema())) { // subsequent files have to have equal schemas
      throw new IOException("schemas dont match");
    }
  }

  private boolean hasNextInput() {
    return inFiles.size() > (currentInput + 1);
  }

  /** skips a number of records from the input */
  private long skip(long skip) throws IOException {
    long skipped = 0;
    while (0 < skip && reader.hasNext()) {
      reader.next(reuse);
      skip--;
      skipped++;
    }
    if ((0 < skip) && hasNextInput()) { // goto next file
      nextInput();
      skipped = skipped + skip(skip);
    }
    return skipped;
  }

  /**
   * writes records with the given samplerate The record at position offset is
   * guaranteed to be taken
   */
  private long writeRecords(long count, double samplerate) throws IOException {
    long written = 0;
    while (written < count && reader.hasNext()) {
      reuse = reader.next(reuse);
      sampleCounter = sampleCounter + samplerate;
      if (sampleCounter >= 1) {
        writer.append(reuse);
        written++;
        sampleCounter--;
      }
    }
    totalCopied = totalCopied + written;
    if (written < count && hasNextInput()) { // goto next file
      nextInput();
      written = written + writeRecords(count - written, samplerate);
    }
    return written;
  }

  private void printHelp(PrintStream out) {
    out.println("cat --offset <offset> --limit <limit> --samplerate <samplerate> [input-files...] output-file");
    out.println();
    out.println("extracts records from a list of input files into a new file.");
    out.println("--offset      start of the extract");
    out.println("--limit       maximum number of records in the output file.");
    out.println("--samplerate  rate at which records will be collected");
    out.println("A dash ('-') can be given to direct output to stdout");
  }

  @Override
  public String getName() {
    return "cat";
  }

  @Override
  public String getShortDescription() {
    return "Extracts samples from files";
  }

}