Skip to content

Commit 545496c

Browse files
BryanCutlerwesm
BryanCutler
authored and
wesm
committed
ARROW-1619: [Java] Set lastSet in JsonFileReader
When reading a vector in JsonFileReader, lastSet should be set in VariableWidthVectors after reading inner vectors or else subsequent operations could corrupt the offsets. This also allows to simplify some of the related code. Additionally, ListVector.lastSet should be explicitly initialized to 0, which is it's starting offset. Author: Bryan Cutler <cutlerb@gmail.com> Closes apache#1140 from BryanCutler/java-JsonReader-setLast-ARROW-1619 and squashes the following commits: 8f97a3d [Bryan Cutler] added test for VarBinaryVector that checks lastSet after reading 70df0cc [Bryan Cutler] set lastSet in JsonFileReader and initialize lastSet for ListVector
1 parent ac997fb commit 545496c

5 files changed

Lines changed: 149 additions & 18 deletions

File tree

‎java/vector/src/main/java/org/apache/arrow/vector/complex/ListVector.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -368,7 +368,7 @@ public UnionVector promoteToUnion() {
368368
return vector;
369369
}
370370

371-
private int lastSet;
371+
private int lastSet = 0;
372372

373373
public class Accessor extends BaseRepeatedAccessor {
374374

‎java/vector/src/main/java/org/apache/arrow/vector/file/json/JsonFileReader.java‎

Lines changed: 29 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,8 @@
4444
import org.apache.arrow.vector.Float4Vector;
4545
import org.apache.arrow.vector.Float8Vector;
4646
import org.apache.arrow.vector.IntVector;
47+
import org.apache.arrow.vector.NullableVarBinaryVector;
48+
import org.apache.arrow.vector.NullableVarCharVector;
4749
import org.apache.arrow.vector.SmallIntVector;
4850
import org.apache.arrow.vector.TimeMicroVector;
4951
import org.apache.arrow.vector.TimeMilliVector;
@@ -63,10 +65,10 @@
6365
import org.apache.arrow.vector.UInt4Vector;
6466
import org.apache.arrow.vector.UInt8Vector;
6567
import org.apache.arrow.vector.ValueVector;
66-
import org.apache.arrow.vector.ValueVector.Mutator;
6768
import org.apache.arrow.vector.VarBinaryVector;
6869
import org.apache.arrow.vector.VarCharVector;
6970
import org.apache.arrow.vector.VectorSchemaRoot;
71+
import org.apache.arrow.vector.complex.ListVector;
7072
import org.apache.arrow.vector.dictionary.Dictionary;
7173
import org.apache.arrow.vector.dictionary.DictionaryProvider;
7274
import org.apache.arrow.vector.schema.ArrowVectorType;
@@ -84,7 +86,6 @@
8486
import com.google.common.base.Objects;
8587

8688
public class JsonFileReader implements AutoCloseable, DictionaryProvider {
87-
private final File inputFile;
8889
private final JsonParser parser;
8990
private final BufferAllocator allocator;
9091
private Schema schema;
@@ -93,7 +94,6 @@ public class JsonFileReader implements AutoCloseable, DictionaryProvider {
9394

9495
public JsonFileReader(File inputFile, BufferAllocator allocator) throws JsonParseException, IOException {
9596
super();
96-
this.inputFile = inputFile;
9797
this.allocator = allocator;
9898
MappingJsonFactory jsonFactory = new MappingJsonFactory();
9999
this.parser = jsonFactory.createParser(inputFile);
@@ -216,10 +216,9 @@ public VectorSchemaRoot read() throws IOException {
216216
}
217217
}
218218

219-
/*
220-
* TODO: This method doesn't load some vectors correctly. For instance, it doesn't initialize
221-
* `lastSet` in ListVector, VarCharVector, NullableVarBinaryVector A better way of implementing
222-
* this function is to use `loadFieldBuffers` methods in FieldVector.
219+
/**
220+
* TODO: A better way of implementing this function is to use `loadFieldBuffers` methods in
221+
* FieldVector to set the inner-vector data as done in `ArrowFileReader`.
223222
*/
224223
private void readVector(Field field, FieldVector vector) throws JsonParseException, IOException {
225224
List<ArrowVectorType> vectorTypes = field.getTypeLayout().getVectorTypes();
@@ -234,29 +233,42 @@ private void readVector(Field field, FieldVector vector) throws JsonParseExcepti
234233
if (started && !Objects.equal(field.getName(), name)) {
235234
throw new IllegalArgumentException("Expected field " + field.getName() + " but got " + name);
236235
}
236+
237+
// Initialize the vector with required capacity
237238
int count = readNextField("count", Integer.class);
239+
vector.setInitialCapacity(count);
238240
vector.allocateNew();
239-
vector.getMutator().setValueCount(count);
241+
242+
// Read inner vectors
240243
for (int v = 0; v < vectorTypes.size(); v++) {
241244
ArrowVectorType vectorType = vectorTypes.get(v);
242-
BufferBacked innerVector = fieldInnerVectors.get(v);
245+
ValueVector valueVector = (ValueVector) fieldInnerVectors.get(v);
243246
nextFieldIs(vectorType.getName());
244247
readToken(START_ARRAY);
245-
ValueVector valueVector = (ValueVector) innerVector;
246-
247248
int innerVectorCount = vectorType.equals(OFFSET) ? count + 1 : count;
248-
valueVector.setInitialCapacity(innerVectorCount);
249-
valueVector.allocateNew();
250-
251249
for (int i = 0; i < innerVectorCount; i++) {
252250
parser.nextToken();
253251
setValueFromParser(valueVector, i);
254252
}
255-
Mutator mutator = valueVector.getMutator();
256-
mutator.setValueCount(innerVectorCount);
257253
readToken(END_ARRAY);
258254
}
259-
// if children
255+
256+
// Set lastSet before valueCount to prevent setValueCount from filling empty values
257+
switch (vector.getMinorType()) {
258+
case LIST:
259+
// ListVector starts lastSet from index 0, so lastSet value is always last index written + 1
260+
((ListVector) vector).getMutator().setLastSet(count);
261+
break;
262+
case VARBINARY:
263+
((NullableVarBinaryVector) vector).getMutator().setLastSet(count - 1);
264+
break;
265+
case VARCHAR:
266+
((NullableVarCharVector) vector).getMutator().setLastSet(count - 1);
267+
break;
268+
}
269+
vector.getMutator().setValueCount(count);
270+
271+
// read child vectors, if any
260272
List<Field> fields = field.getChildren();
261273
if (!fields.isEmpty()) {
262274
List<FieldVector> vectorChildren = vector.getChildrenFromFields();

‎java/vector/src/test/java/org/apache/arrow/vector/file/BaseFileTest.java‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import org.apache.arrow.vector.NullableDecimalVector;
3333
import org.apache.arrow.vector.NullableIntVector;
3434
import org.apache.arrow.vector.NullableTimeMilliVector;
35+
import org.apache.arrow.vector.NullableVarBinaryVector;
3536
import org.apache.arrow.vector.NullableVarCharVector;
3637
import org.apache.arrow.vector.ValueVector.Accessor;
3738
import org.apache.arrow.vector.VectorSchemaRoot;
@@ -541,4 +542,52 @@ public void writeUnionData(int count, NullableMapVector parent) {
541542
writer.setValueCount(count);
542543
varchar.release();
543544
}
545+
546+
protected void writeVarBinaryData(int count, NullableMapVector parent) {
547+
Assert.assertTrue(count < 100);
548+
ComplexWriter writer = new ComplexWriterImpl("root", parent);
549+
MapWriter rootWriter = writer.rootAsMap();
550+
ListWriter listWriter = rootWriter.list("list");
551+
ArrowBuf varbin = allocator.buffer(count);
552+
for (int i = 0; i < count; i++) {
553+
varbin.setByte(i, i);
554+
listWriter.setPosition(i);
555+
listWriter.startList();
556+
for (int j = 0; j < i % 3; j++) {
557+
listWriter.varBinary().writeVarBinary(0, i + 1, varbin);
558+
}
559+
listWriter.endList();
560+
}
561+
writer.setValueCount(count);
562+
varbin.release();
563+
}
564+
565+
protected void validateVarBinary(int count, VectorSchemaRoot root) {
566+
Assert.assertEquals(count, root.getRowCount());
567+
ListVector listVector = (ListVector) root.getVector("list");
568+
byte[] expectedArray = new byte[count];
569+
int numVarBinaryValues = 0;
570+
for (int i = 0; i < count; i++) {
571+
expectedArray[i] = (byte) i;
572+
Object obj = listVector.getAccessor().getObject(i);
573+
List<?> objList = (List) obj;
574+
if (i % 3 == 0) {
575+
Assert.assertTrue(objList.isEmpty());
576+
} else {
577+
byte[] expected = Arrays.copyOfRange(expectedArray, 0, i + 1);
578+
for (int j = 0; j < i % 3; j++) {
579+
byte[] result = (byte[]) objList.get(j);
580+
Assert.assertArrayEquals(result, expected);
581+
numVarBinaryValues++;
582+
}
583+
}
584+
}
585+
586+
// ListVector lastSet should be the index of last value + 1
587+
Assert.assertEquals(listVector.getMutator().getLastSet(), count);
588+
589+
// NullableVarBinaryVector lastSet should be the index of last value
590+
NullableVarBinaryVector binaryVector = (NullableVarBinaryVector) listVector.getChildrenFromFields().get(0);
591+
Assert.assertEquals(binaryVector.getMutator().getLastSet(), numVarBinaryValues - 1);
592+
}
544593
}

‎java/vector/src/test/java/org/apache/arrow/vector/file/TestArrowFile.java‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -622,6 +622,47 @@ public void testWriteReadFixedSizeList() throws IOException {
622622
}
623623
}
624624

625+
@Test
626+
public void testWriteReadVarBin() throws IOException {
627+
File file = new File("target/mytest_varbin.arrow");
628+
ByteArrayOutputStream stream = new ByteArrayOutputStream();
629+
int count = COUNT;
630+
631+
// write
632+
try (
633+
BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE);
634+
NullableMapVector parent = NullableMapVector.empty("parent", vectorAllocator)) {
635+
writeVarBinaryData(count, parent);
636+
VectorSchemaRoot root = new VectorSchemaRoot(parent.getChild("root"));
637+
validateVarBinary(count, root);
638+
write(parent.getChild("root"), file, stream);
639+
}
640+
641+
// read from file
642+
try (
643+
BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE);
644+
FileInputStream fileInputStream = new FileInputStream(file);
645+
ArrowFileReader arrowReader = new ArrowFileReader(fileInputStream.getChannel(), readerAllocator)) {
646+
VectorSchemaRoot root = arrowReader.getVectorSchemaRoot();
647+
Schema schema = root.getSchema();
648+
LOGGER.debug("reading schema: " + schema);
649+
Assert.assertTrue(arrowReader.loadNextBatch());
650+
validateVarBinary(count, root);
651+
}
652+
653+
// read from stream
654+
try (
655+
BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE);
656+
ByteArrayInputStream input = new ByteArrayInputStream(stream.toByteArray());
657+
ArrowStreamReader arrowReader = new ArrowStreamReader(input, readerAllocator)) {
658+
VectorSchemaRoot root = arrowReader.getVectorSchemaRoot();
659+
Schema schema = root.getSchema();
660+
LOGGER.debug("reading schema: " + schema);
661+
Assert.assertTrue(arrowReader.loadNextBatch());
662+
validateVarBinary(count, root);
663+
}
664+
}
665+
625666

626667
/**
627668
* Writes the contents of parents to file. If outStream is non-null, also writes it

‎java/vector/src/test/java/org/apache/arrow/vector/file/json/TestJSONFile.java‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -285,4 +285,33 @@ public void testSetStructLength() throws IOException {
285285
}
286286
}
287287

288+
@Test
289+
public void testWriteReadVarBinJSON() throws IOException {
290+
File file = new File("target/mytest_varbin.json");
291+
int count = COUNT;
292+
293+
// write
294+
try (
295+
BufferAllocator vectorAllocator = allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE);
296+
NullableMapVector parent = NullableMapVector.empty("parent", vectorAllocator)) {
297+
writeVarBinaryData(count, parent);
298+
VectorSchemaRoot root = new VectorSchemaRoot(parent.getChild("root"));
299+
validateVarBinary(count, root);
300+
writeJSON(file, new VectorSchemaRoot(parent.getChild("root")), null);
301+
}
302+
303+
// read
304+
try (
305+
BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE)) {
306+
JsonFileReader reader = new JsonFileReader(file, readerAllocator);
307+
Schema schema = reader.start();
308+
LOGGER.debug("reading schema: " + schema);
309+
310+
// initialize vectors
311+
try (VectorSchemaRoot root = reader.read();) {
312+
validateVarBinary(count, root);
313+
}
314+
reader.close();
315+
}
316+
}
288317
}

0 commit comments

Comments
 (0)