missing file

Author: julienledem

java/tools/src/test/java/org/apache/arrow/tools/TestFileRoundtrip.java

@@ -0,0 +1,159 @@
+/*
+ * 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
+ *
+ *   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.
+ */
+package org.apache.arrow.tools;
+
+import static org.junit.Assert.assertEquals;
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileNotFoundException;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.util.List;
+
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.VectorLoader;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.VectorUnloader;
+import org.apache.arrow.vector.complex.MapVector;
+import org.apache.arrow.vector.complex.impl.ComplexWriterImpl;
+import org.apache.arrow.vector.complex.writer.BaseWriter.ComplexWriter;
+import org.apache.arrow.vector.complex.writer.BaseWriter.MapWriter;
+import org.apache.arrow.vector.complex.writer.BigIntWriter;
+import org.apache.arrow.vector.complex.writer.IntWriter;
+import org.apache.arrow.vector.file.ArrowBlock;
+import org.apache.arrow.vector.file.ArrowFooter;
+import org.apache.arrow.vector.file.ArrowReader;
+import org.apache.arrow.vector.file.ArrowWriter;
+import org.apache.arrow.vector.schema.ArrowRecordBatch;
+import org.apache.arrow.vector.types.pojo.Schema;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+public class TestFileRoundtrip {
+  private static final int COUNT = 10;
+
+  @Rule
+  public TemporaryFolder testFolder = new TemporaryFolder();
+
+  private BufferAllocator allocator;
+
+  @Before
+  public void init() {
+    allocator = new RootAllocator(Integer.MAX_VALUE);
+  }
+
+  @After
+  public void tearDown() {
+    allocator.close();
+  }
+
+  private void writeData(int count, MapVector parent) {
+    ComplexWriter writer = new ComplexWriterImpl("root", parent);
+    MapWriter rootWriter = writer.rootAsMap();
+    IntWriter intWriter = rootWriter.integer("int");
+    BigIntWriter bigIntWriter = rootWriter.bigInt("bigInt");
+    for (int i = 0; i < count; i++) {
+      intWriter.setPosition(i);
+      intWriter.writeInt(i);
+      bigIntWriter.setPosition(i);
+      bigIntWriter.writeBigInt(i);
+    }
+    writer.setValueCount(count);
+  }
+
+  @Test
+  public void test() throws Exception {
+    File testInFile = testFolder.newFile("testIn.arrow");
+    File testOutFile = testFolder.newFile("testOut.arrow");
+
+    writeInput(testInFile);
+
+    String[] args = { "-i", testInFile.getAbsolutePath(), "-o",  testOutFile.getAbsolutePath()};
+    int result = new FileRoundtrip(System.out, System.err).run(args);
+    assertEquals(0, result);
+
+    validateOutput(testOutFile);
+  }
+
+  private void validateOutput(File testOutFile) throws Exception {
+    // read
+    try (
+        BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE);
+        FileInputStream fileInputStream = new FileInputStream(testOutFile);
+        ArrowReader arrowReader = new ArrowReader(fileInputStream.getChannel(), readerAllocator);
+        BufferAllocator vectorAllocator = allocator.newChildAllocator("final vectors", 0, Integer.MAX_VALUE);
+        ) {
+      ArrowFooter footer = arrowReader.readFooter();
+      Schema schema = footer.getSchema();
+
+      // initialize vectors
+      try (VectorSchemaRoot root = new VectorSchemaRoot(schema, readerAllocator)) {
+        VectorLoader vectorLoader = new VectorLoader(root);
+
+        List<ArrowBlock> recordBatches = footer.getRecordBatches();
+        for (ArrowBlock rbBlock : recordBatches) {
+          try (ArrowRecordBatch recordBatch = arrowReader.readRecordBatch(rbBlock)) {
+            vectorLoader.load(recordBatch);
+          }
+          validateContent(COUNT, root);
+        }
+      }
+    }
+  }
+
+  private void validateContent(int count, VectorSchemaRoot root) {
+    Assert.assertEquals(count, root.getRowCount());
+    for (int i = 0; i < count; i++) {
+      Assert.assertEquals(i, root.getVector("int").getAccessor().getObject(i));
+      Assert.assertEquals(Long.valueOf(i), root.getVector("bigInt").getAccessor().getObject(i));
+    }
+  }
+
+  public void writeInput(File testInFile) throws FileNotFoundException, IOException {
+    int count = COUNT;
+    try (
+        BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE);
+        MapVector parent = new MapVector("parent", vectorAllocator, null)) {
+      writeData(count, parent);
+      write(parent.getChild("root"), testInFile);
+    }
+  }
+
+  private void write(FieldVector parent, File file) throws FileNotFoundException, IOException {
+    Schema schema = new Schema(parent.getField().getChildren());
+    int valueCount = parent.getAccessor().getValueCount();
+    List<FieldVector> fields = parent.getChildrenFromFields();
+    VectorUnloader vectorUnloader = new VectorUnloader(schema, valueCount, fields);
+    try (
+        FileOutputStream fileOutputStream = new FileOutputStream(file);
+        ArrowWriter arrowWriter = new ArrowWriter(fileOutputStream.getChannel(), schema);
+        ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch();
+            ) {
+      arrowWriter.writeRecordBatch(recordBatch);
+    }
+  }
+
+}