ceph/src/arrow/java/tools/src/main/java/org/apache/arrow/tools/FileToStream.java

   1 /*
   2  * Licensed to the Apache Software Foundation (ASF) under one or more
   3  * contributor license agreements.  See the NOTICE file distributed with
   4  * this work for additional information regarding copyright ownership.
   5  * The ASF licenses this file to You under the Apache License, Version 2.0
   6  * (the "License"); you may not use this file except in compliance with
   7  * the License.  You may obtain a copy of the License at
   8  *
   9  *    http://www.apache.org/licenses/LICENSE-2.0
  10  *
  11  * Unless required by applicable law or agreed to in writing, software
  12  * distributed under the License is distributed on an "AS IS" BASIS,
  13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14  * See the License for the specific language governing permissions and
  15  * limitations under the License.
  16  */
  17
  18 package org.apache.arrow.tools;
  19
  20 import java.io.File;
  21 import java.io.FileInputStream;
  22 import java.io.FileOutputStream;
  23 import java.io.IOException;
  24 import java.io.OutputStream;
  25
  26 import org.apache.arrow.memory.BufferAllocator;
  27 import org.apache.arrow.memory.RootAllocator;
  28 import org.apache.arrow.vector.VectorSchemaRoot;
  29 import org.apache.arrow.vector.ipc.ArrowFileReader;
  30 import org.apache.arrow.vector.ipc.ArrowStreamWriter;
  31
  32 /**
  33  * Converts an Arrow file to an Arrow stream. The file should be specified as the
  34  * first argument and the output is written to standard out.
  35  */
  36 public class FileToStream {
  37   private FileToStream() {}
  38
  39   /**
  40    * Reads an Arrow file from in and writes it back to out.
  41    */
  42   public static void convert(FileInputStream in, OutputStream out) throws IOException {
  43     BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
  44     try (ArrowFileReader reader = new ArrowFileReader(in.getChannel(), allocator)) {
  45       VectorSchemaRoot root = reader.getVectorSchemaRoot();
  46       // load the first batch before instantiating the writer so that we have any dictionaries
  47       // only writeBatches if we loaded one in the first place.
  48       boolean writeBatches = reader.loadNextBatch();
  49       try (ArrowStreamWriter writer = new ArrowStreamWriter(root, reader, out)) {
  50         writer.start();
  51         while (writeBatches) {
  52           writer.writeBatch();
  53           if (!reader.loadNextBatch()) {
  54             break;
  55           }
  56         }
  57         writer.end();
  58       }
  59     }
  60   }
  61
  62   /**
  63    * Main method.  The first arg is the file path.  The second, optional argument,
  64    * is an output file location (defaults to standard out).
  65    */
  66   public static void main(String[] args) throws IOException {
  67     if (args.length != 1 && args.length != 2) {
  68       System.err.println("Usage: FileToStream <input file> [output file]");
  69       System.exit(1);
  70     }
  71
  72     FileInputStream in = new FileInputStream(new File(args[0]));
  73     OutputStream out = args.length == 1 ?
  74         System.out : new FileOutputStream(new File(args[1]));
  75
  76     convert(in, out);
  77   }
  78 }