Skip to content

Commit 858b1b1

Browse files
julienledemkou
julienledem
authored and
kou
committed
ARROW-367: converter json <=> Arrow file format for Integration tests
Author: Julien Le Dem <julien@dremio.com> Closes apache#203 from julienledem/integration and squashes the following commits: b3cd326 [Julien Le Dem] add license fdbe03f [Julien Le Dem] ARROW-367: converter json <=> Arrow file format for Integration tests
1 parent 0c4c153 commit 858b1b1

6 files changed

Lines changed: 554 additions & 114 deletions

File tree

‎tools/src/main/java/org/apache/arrow/tools/Integration.java‎

Lines changed: 262 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,262 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.apache.arrow.tools;
20+
21+
import java.io.File;
22+
import java.io.FileInputStream;
23+
import java.io.FileOutputStream;
24+
import java.io.IOException;
25+
import java.util.Arrays;
26+
import java.util.Iterator;
27+
import java.util.List;
28+
29+
import org.apache.arrow.memory.BufferAllocator;
30+
import org.apache.arrow.memory.RootAllocator;
31+
import org.apache.arrow.vector.FieldVector;
32+
import org.apache.arrow.vector.VectorLoader;
33+
import org.apache.arrow.vector.VectorSchemaRoot;
34+
import org.apache.arrow.vector.VectorUnloader;
35+
import org.apache.arrow.vector.file.ArrowBlock;
36+
import org.apache.arrow.vector.file.ArrowFooter;
37+
import org.apache.arrow.vector.file.ArrowReader;
38+
import org.apache.arrow.vector.file.ArrowWriter;
39+
import org.apache.arrow.vector.file.json.JsonFileReader;
40+
import org.apache.arrow.vector.file.json.JsonFileWriter;
41+
import org.apache.arrow.vector.schema.ArrowRecordBatch;
42+
import org.apache.arrow.vector.types.pojo.Field;
43+
import org.apache.arrow.vector.types.pojo.Schema;
44+
import org.apache.commons.cli.CommandLine;
45+
import org.apache.commons.cli.CommandLineParser;
46+
import org.apache.commons.cli.Options;
47+
import org.apache.commons.cli.ParseException;
48+
import org.apache.commons.cli.PosixParser;
49+
import org.slf4j.Logger;
50+
import org.slf4j.LoggerFactory;
51+
52+
import com.google.common.base.Objects;
53+
54+
public class Integration {
55+
private static final Logger LOGGER = LoggerFactory.getLogger(Integration.class);
56+
57+
public static void main(String[] args) {
58+
try {
59+
new Integration().run(args);
60+
} catch (ParseException e) {
61+
fatalError("Invalid parameters", e);
62+
} catch (IOException e) {
63+
fatalError("Error accessing files", e);
64+
} catch (RuntimeException e) {
65+
fatalError("Incompatible files", e);
66+
}
67+
}
68+
69+
private final Options options;
70+
71+
enum Command {
72+
ARROW_TO_JSON(true, false) {
73+
@Override
74+
public void execute(File arrowFile, File jsonFile) throws IOException {
75+
try(
76+
BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
77+
FileInputStream fileInputStream = new FileInputStream(arrowFile);
78+
ArrowReader arrowReader = new ArrowReader(fileInputStream.getChannel(), allocator);) {
79+
ArrowFooter footer = arrowReader.readFooter();
80+
Schema schema = footer.getSchema();
81+
LOGGER.debug("Input file size: " + arrowFile.length());
82+
LOGGER.debug("Found schema: " + schema);
83+
try (JsonFileWriter writer = new JsonFileWriter(jsonFile);) {
84+
writer.start(schema);
85+
List<ArrowBlock> recordBatches = footer.getRecordBatches();
86+
for (ArrowBlock rbBlock : recordBatches) {
87+
try (ArrowRecordBatch inRecordBatch = arrowReader.readRecordBatch(rbBlock);
88+
VectorSchemaRoot root = new VectorSchemaRoot(schema, allocator);) {
89+
VectorLoader vectorLoader = new VectorLoader(root);
90+
vectorLoader.load(inRecordBatch);
91+
writer.write(root);
92+
}
93+
}
94+
}
95+
LOGGER.debug("Output file size: " + jsonFile.length());
96+
}
97+
}
98+
},
99+
JSON_TO_ARROW(false, true) {
100+
@Override
101+
public void execute(File arrowFile, File jsonFile) throws IOException {
102+
try (
103+
BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
104+
JsonFileReader reader = new JsonFileReader(jsonFile, allocator);
105+
) {
106+
Schema schema = reader.start();
107+
LOGGER.debug("Input file size: " + jsonFile.length());
108+
LOGGER.debug("Found schema: " + schema);
109+
try (
110+
FileOutputStream fileOutputStream = new FileOutputStream(arrowFile);
111+
ArrowWriter arrowWriter = new ArrowWriter(fileOutputStream.getChannel(), schema);
112+
) {
113+
114+
// initialize vectors
115+
VectorSchemaRoot root;
116+
while ((root = reader.read()) != null) {
117+
VectorUnloader vectorUnloader = new VectorUnloader(root);
118+
try (ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch();) {
119+
arrowWriter.writeRecordBatch(recordBatch);
120+
}
121+
root.close();
122+
}
123+
}
124+
LOGGER.debug("Output file size: " + arrowFile.length());
125+
}
126+
}
127+
},
128+
VALIDATE(true, true) {
129+
@Override
130+
public void execute(File arrowFile, File jsonFile) throws IOException {
131+
try (
132+
BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
133+
JsonFileReader jsonReader = new JsonFileReader(jsonFile, allocator);
134+
FileInputStream fileInputStream = new FileInputStream(arrowFile);
135+
ArrowReader arrowReader = new ArrowReader(fileInputStream.getChannel(), allocator);
136+
) {
137+
Schema jsonSchema = jsonReader.start();
138+
ArrowFooter footer = arrowReader.readFooter();
139+
Schema arrowSchema = footer.getSchema();
140+
LOGGER.debug("Arrow Input file size: " + arrowFile.length());
141+
LOGGER.debug("ARROW schema: " + arrowSchema);
142+
LOGGER.debug("JSON Input file size: " + jsonFile.length());
143+
LOGGER.debug("JSON schema: " + jsonSchema);
144+
compareSchemas(jsonSchema, arrowSchema);
145+
146+
List<ArrowBlock> recordBatches = footer.getRecordBatches();
147+
Iterator<ArrowBlock> iterator = recordBatches.iterator();
148+
VectorSchemaRoot jsonRoot;
149+
while ((jsonRoot = jsonReader.read()) != null && iterator.hasNext()) {
150+
ArrowBlock rbBlock = iterator.next();
151+
try (ArrowRecordBatch inRecordBatch = arrowReader.readRecordBatch(rbBlock);
152+
VectorSchemaRoot arrowRoot = new VectorSchemaRoot(arrowSchema, allocator);) {
153+
VectorLoader vectorLoader = new VectorLoader(arrowRoot);
154+
vectorLoader.load(inRecordBatch);
155+
// TODO: compare
156+
compare(arrowRoot, jsonRoot);
157+
}
158+
jsonRoot.close();
159+
}
160+
boolean hasMoreJSON = jsonRoot != null;
161+
boolean hasMoreArrow = iterator.hasNext();
162+
if (hasMoreJSON || hasMoreArrow) {
163+
throw new IllegalArgumentException("Unexpected RecordBatches. J:" + hasMoreJSON + " A:" + hasMoreArrow);
164+
}
165+
}
166+
}
167+
};
168+
169+
public final boolean arrowExists;
170+
public final boolean jsonExists;
171+
172+
Command(boolean arrowExists, boolean jsonExists) {
173+
this.arrowExists = arrowExists;
174+
this.jsonExists = jsonExists;
175+
}
176+
177+
abstract public void execute(File arrowFile, File jsonFile) throws IOException;
178+
179+
}
180+
181+
Integration() {
182+
this.options = new Options();
183+
this.options.addOption("a", "arrow", true, "arrow file");
184+
this.options.addOption("j", "json", true, "json file");
185+
this.options.addOption("c", "command", true, "command to execute: " + Arrays.toString(Command.values()));
186+
}
187+
188+
private File validateFile(String type, String fileName, boolean shouldExist) {
189+
if (fileName == null) {
190+
throw new IllegalArgumentException("missing " + type + " file parameter");
191+
}
192+
File f = new File(fileName);
193+
if (shouldExist && (!f.exists() || f.isDirectory())) {
194+
throw new IllegalArgumentException(type + " file not found: " + f.getAbsolutePath());
195+
}
196+
if (!shouldExist && f.exists()) {
197+
throw new IllegalArgumentException(type + " file already exists: " + f.getAbsolutePath());
198+
}
199+
return f;
200+
}
201+
202+
void run(String[] args) throws ParseException, IOException {
203+
CommandLineParser parser = new PosixParser();
204+
CommandLine cmd = parser.parse(options, args, false);
205+
206+
207+
Command command = toCommand(cmd.getOptionValue("command"));
208+
File arrowFile = validateFile("arrow", cmd.getOptionValue("arrow"), command.arrowExists);
209+
File jsonFile = validateFile("json", cmd.getOptionValue("json"), command.jsonExists);
210+
command.execute(arrowFile, jsonFile);
211+
}
212+
213+
private Command toCommand(String commandName) {
214+
try {
215+
return Command.valueOf(commandName);
216+
} catch (IllegalArgumentException e) {
217+
throw new IllegalArgumentException("Unknown command: " + commandName + " expected one of " + Arrays.toString(Command.values()));
218+
}
219+
}
220+
221+
private static void fatalError(String message, Throwable e) {
222+
System.err.println(message);
223+
LOGGER.error(message, e);
224+
System.exit(1);
225+
}
226+
227+
228+
private static void compare(VectorSchemaRoot arrowRoot, VectorSchemaRoot jsonRoot) {
229+
compareSchemas(jsonRoot.getSchema(), arrowRoot.getSchema());
230+
if (arrowRoot.getRowCount() != jsonRoot.getRowCount()) {
231+
throw new IllegalArgumentException("Different row count:\n" + arrowRoot.getRowCount() + "\n" + jsonRoot.getRowCount());
232+
}
233+
List<FieldVector> arrowVectors = arrowRoot.getFieldVectors();
234+
List<FieldVector> jsonVectors = jsonRoot.getFieldVectors();
235+
if (arrowVectors.size() != jsonVectors.size()) {
236+
throw new IllegalArgumentException("Different column count:\n" + arrowVectors.size() + "\n" + jsonVectors.size());
237+
}
238+
for (int i = 0; i < arrowVectors.size(); i++) {
239+
Field field = arrowRoot.getSchema().getFields().get(i);
240+
FieldVector arrowVector = arrowVectors.get(i);
241+
FieldVector jsonVector = jsonVectors.get(i);
242+
int valueCount = arrowVector.getAccessor().getValueCount();
243+
if (valueCount != jsonVector.getAccessor().getValueCount()) {
244+
throw new IllegalArgumentException("Different value count for field " + field + " : " + valueCount + " != " + jsonVector.getAccessor().getValueCount());
245+
}
246+
for (int j = 0; j < valueCount; j++) {
247+
Object arrow = arrowVector.getAccessor().getObject(j);
248+
Object json = jsonVector.getAccessor().getObject(j);
249+
if (!Objects.equal(arrow, json)) {
250+
throw new IllegalArgumentException(
251+
"Different values in column:\n" + field + " at index " + j + ": " + arrow + " != " + json);
252+
}
253+
}
254+
}
255+
}
256+
257+
private static void compareSchemas(Schema jsonSchema, Schema arrowSchema) {
258+
if (!arrowSchema.equals(jsonSchema)) {
259+
throw new IllegalArgumentException("Different schemas:\n" + arrowSchema + "\n" + jsonSchema);
260+
}
261+
}
262+
}

‎tools/src/test/java/org/apache/arrow/tools/ArrowFileTestFixtures.java‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.apache.arrow.tools;
20+
21+
import java.io.File;
22+
import java.io.FileInputStream;
23+
import java.io.FileNotFoundException;
24+
import java.io.FileOutputStream;
25+
import java.io.IOException;
26+
import java.util.List;
27+
28+
import org.apache.arrow.memory.BufferAllocator;
29+
import org.apache.arrow.vector.FieldVector;
30+
import org.apache.arrow.vector.VectorLoader;
31+
import org.apache.arrow.vector.VectorSchemaRoot;
32+
import org.apache.arrow.vector.VectorUnloader;
33+
import org.apache.arrow.vector.complex.MapVector;
34+
import org.apache.arrow.vector.complex.impl.ComplexWriterImpl;
35+
import org.apache.arrow.vector.complex.writer.BaseWriter.ComplexWriter;
36+
import org.apache.arrow.vector.complex.writer.BaseWriter.MapWriter;
37+
import org.apache.arrow.vector.complex.writer.BigIntWriter;
38+
import org.apache.arrow.vector.complex.writer.IntWriter;
39+
import org.apache.arrow.vector.file.ArrowBlock;
40+
import org.apache.arrow.vector.file.ArrowFooter;
41+
import org.apache.arrow.vector.file.ArrowReader;
42+
import org.apache.arrow.vector.file.ArrowWriter;
43+
import org.apache.arrow.vector.schema.ArrowRecordBatch;
44+
import org.apache.arrow.vector.types.pojo.Schema;
45+
import org.junit.Assert;
46+
47+
public class ArrowFileTestFixtures {
48+
static final int COUNT = 10;
49+
50+
static void writeData(int count, MapVector parent) {
51+
ComplexWriter writer = new ComplexWriterImpl("root", parent);
52+
MapWriter rootWriter = writer.rootAsMap();
53+
IntWriter intWriter = rootWriter.integer("int");
54+
BigIntWriter bigIntWriter = rootWriter.bigInt("bigInt");
55+
for (int i = 0; i < count; i++) {
56+
intWriter.setPosition(i);
57+
intWriter.writeInt(i);
58+
bigIntWriter.setPosition(i);
59+
bigIntWriter.writeBigInt(i);
60+
}
61+
writer.setValueCount(count);
62+
}
63+
64+
static void validateOutput(File testOutFile, BufferAllocator allocator) throws Exception {
65+
// read
66+
try (
67+
BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE);
68+
FileInputStream fileInputStream = new FileInputStream(testOutFile);
69+
ArrowReader arrowReader = new ArrowReader(fileInputStream.getChannel(), readerAllocator);
70+
BufferAllocator vectorAllocator = allocator.newChildAllocator("final vectors", 0, Integer.MAX_VALUE);
71+
) {
72+
ArrowFooter footer = arrowReader.readFooter();
73+
Schema schema = footer.getSchema();
74+
75+
// initialize vectors
76+
try (VectorSchemaRoot root = new VectorSchemaRoot(schema, readerAllocator)) {
77+
VectorLoader vectorLoader = new VectorLoader(root);
78+
79+
List<ArrowBlock> recordBatches = footer.getRecordBatches();
80+
for (ArrowBlock rbBlock : recordBatches) {
81+
try (ArrowRecordBatch recordBatch = arrowReader.readRecordBatch(rbBlock)) {
82+
vectorLoader.load(recordBatch);
83+
}
84+
validateContent(COUNT, root);
85+
}
86+
}
87+
}
88+
}
89+
90+
static void validateContent(int count, VectorSchemaRoot root) {
91+
Assert.assertEquals(count, root.getRowCount());
92+
for (int i = 0; i < count; i++) {
93+
Assert.assertEquals(i, root.getVector("int").getAccessor().getObject(i));
94+
Assert.assertEquals(Long.valueOf(i), root.getVector("bigInt").getAccessor().getObject(i));
95+
}
96+
}
97+
98+
static void write(FieldVector parent, File file) throws FileNotFoundException, IOException {
99+
Schema schema = new Schema(parent.getField().getChildren());
100+
int valueCount = parent.getAccessor().getValueCount();
101+
List<FieldVector> fields = parent.getChildrenFromFields();
102+
VectorUnloader vectorUnloader = new VectorUnloader(schema, valueCount, fields);
103+
try (
104+
FileOutputStream fileOutputStream = new FileOutputStream(file);
105+
ArrowWriter arrowWriter = new ArrowWriter(fileOutputStream.getChannel(), schema);
106+
ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch();
107+
) {
108+
arrowWriter.writeRecordBatch(recordBatch);
109+
}
110+
}
111+
112+
113+
static void writeInput(File testInFile, BufferAllocator allocator) throws FileNotFoundException, IOException {
114+
int count = ArrowFileTestFixtures.COUNT;
115+
try (
116+
BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE);
117+
MapVector parent = new MapVector("parent", vectorAllocator, null)) {
118+
writeData(count, parent);
119+
write(parent.getChild("root"), testInFile);
120+
}
121+
}
122+
}

0 commit comments

Comments
 (0)