Skip to content

Commit 9d5c078

Browse files
julienledemjulienledem
committed
fix read-write bug
1 parent 7f20b36 commit 9d5c078

12 files changed

Lines changed: 321 additions & 122 deletions

File tree

‎java/memory/src/main/java/io/netty/buffer/ArrowBuf.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -179,7 +179,10 @@ public ArrowBuf retain(BufferAllocator target) {
179179
historicalLog.recordEvent("retain(%s)", target.getName());
180180
}
181181
final BufferLedger otherLedger = this.ledger.getLedgerForAllocator(target);
182-
return otherLedger.newArrowBuf(offset, length, null);
182+
ArrowBuf newArrowBuf = otherLedger.newArrowBuf(offset, length, null);
183+
newArrowBuf.readerIndex(this.readerIndex);
184+
newArrowBuf.writerIndex(this.writerIndex);
185+
return newArrowBuf;
183186
}
184187

185188
/**
@@ -214,6 +217,8 @@ public TransferResult transferOwnership(BufferAllocator target) {
214217

215218
final BufferLedger otherLedger = this.ledger.getLedgerForAllocator(target);
216219
final ArrowBuf newBuf = otherLedger.newArrowBuf(offset, length, null);
220+
newBuf.readerIndex(this.readerIndex);
221+
newBuf.writerIndex(this.writerIndex);
217222
final boolean allocationFit = this.ledger.transferBalance(otherLedger);
218223
return new TransferResult(allocationFit, newBuf);
219224
}

‎java/memory/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java‎

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -22,16 +22,13 @@
2222
import static org.junit.Assert.assertNotNull;
2323
import static org.junit.Assert.assertTrue;
2424
import static org.junit.Assert.fail;
25-
import io.netty.buffer.ArrowBuf;
26-
import io.netty.buffer.ArrowBuf.TransferResult;
2725

28-
import org.apache.arrow.memory.AllocationReservation;
29-
import org.apache.arrow.memory.BufferAllocator;
30-
import org.apache.arrow.memory.OutOfMemoryException;
31-
import org.apache.arrow.memory.RootAllocator;
3226
import org.junit.Ignore;
3327
import org.junit.Test;
3428

29+
import io.netty.buffer.ArrowBuf;
30+
import io.netty.buffer.ArrowBuf.TransferResult;
31+
3532
public class TestBaseAllocator {
3633
// private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(TestBaseAllocator.class);
3734

@@ -134,6 +131,7 @@ public void testAllocator_transferOwnership() throws Exception {
134131
final ArrowBuf arrowBuf1 = childAllocator1.buffer(MAX_ALLOCATION / 4);
135132
rootAllocator.verify();
136133
TransferResult transferOwnership = arrowBuf1.transferOwnership(childAllocator2);
134+
assertEquiv(arrowBuf1, transferOwnership.buffer);
137135
final boolean allocationFit = transferOwnership.allocationFit;
138136
rootAllocator.verify();
139137
assertTrue(allocationFit);
@@ -160,6 +158,7 @@ public void testAllocator_shareOwnership() throws Exception {
160158
rootAllocator.verify();
161159
assertNotNull(arrowBuf2);
162160
assertNotEquals(arrowBuf2, arrowBuf1);
161+
assertEquiv(arrowBuf1, arrowBuf2);
163162

164163
// release original buffer (thus transferring ownership to allocator 2. (should leave allocator 1 in empty state)
165164
arrowBuf1.release();
@@ -172,6 +171,7 @@ public void testAllocator_shareOwnership() throws Exception {
172171
assertNotNull(arrowBuf3);
173172
assertNotEquals(arrowBuf3, arrowBuf1);
174173
assertNotEquals(arrowBuf3, arrowBuf2);
174+
assertEquiv(arrowBuf1, arrowBuf3);
175175
rootAllocator.verify();
176176

177177
arrowBuf2.release();
@@ -452,8 +452,10 @@ public void testAllocator_transferSliced() throws Exception {
452452
rootAllocator.verify();
453453

454454
TransferResult result1 = arrowBuf2s.transferOwnership(childAllocator1);
455+
assertEquiv(arrowBuf2s, result1.buffer);
455456
rootAllocator.verify();
456457
TransferResult result2 = arrowBuf1s.transferOwnership(childAllocator2);
458+
assertEquiv(arrowBuf1s, result2.buffer);
457459
rootAllocator.verify();
458460

459461
result1.buffer.release();
@@ -482,7 +484,9 @@ public void testAllocator_shareSliced() throws Exception {
482484
rootAllocator.verify();
483485

484486
final ArrowBuf arrowBuf2s1 = arrowBuf2s.retain(childAllocator1);
487+
assertEquiv(arrowBuf2s, arrowBuf2s1);
485488
final ArrowBuf arrowBuf1s2 = arrowBuf1s.retain(childAllocator2);
489+
assertEquiv(arrowBuf1s, arrowBuf1s2);
486490
rootAllocator.verify();
487491

488492
arrowBuf1s.release(); // releases arrowBuf1
@@ -512,11 +516,13 @@ public void testAllocator_transferShared() throws Exception {
512516
rootAllocator.verify();
513517
assertNotNull(arrowBuf2);
514518
assertNotEquals(arrowBuf2, arrowBuf1);
519+
assertEquiv(arrowBuf1, arrowBuf2);
515520

516521
TransferResult result = arrowBuf1.transferOwnership(childAllocator3);
517522
allocationFit = result.allocationFit;
518523
final ArrowBuf arrowBuf3 = result.buffer;
519524
assertTrue(allocationFit);
525+
assertEquiv(arrowBuf1, arrowBuf3);
520526
rootAllocator.verify();
521527

522528
// Since childAllocator3 now has childAllocator1's buffer, 1, can close
@@ -533,6 +539,7 @@ public void testAllocator_transferShared() throws Exception {
533539
allocationFit = result.allocationFit;
534540
final ArrowBuf arrowBuf4 = result2.buffer;
535541
assertTrue(allocationFit);
542+
assertEquiv(arrowBuf3, arrowBuf4);
536543
rootAllocator.verify();
537544

538545
arrowBuf3.release();
@@ -645,4 +652,9 @@ public void multiple() throws Exception {
645652

646653
}
647654
}
655+
656+
public void assertEquiv(ArrowBuf origBuf, ArrowBuf newBuf) {
657+
assertEquals(origBuf.readerIndex(), newBuf.readerIndex());
658+
assertEquals(origBuf.writerIndex(), newBuf.writerIndex());
659+
}
648660
}

‎java/tools/pom.xml‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,26 @@
4848
</dependencies>
4949

5050
<build>
51+
<plugins>
52+
<plugin>
53+
<artifactId>maven-assembly-plugin</artifactId>
54+
<version>2.6</version>
55+
<configuration>
56+
<descriptorRefs>
57+
<descriptorRef>jar-with-dependencies</descriptorRef>
58+
</descriptorRefs>
59+
</configuration>
60+
<executions>
61+
<execution>
62+
<id>make-assembly</id>
63+
<phase>package</phase>
64+
<goals>
65+
<goal>single</goal>
66+
</goals>
67+
</execution>
68+
</executions>
69+
</plugin>
70+
</plugins>
5171
</build>
5272

5373
</project>

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

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,15 +28,13 @@
2828
import org.apache.arrow.memory.BufferAllocator;
2929
import org.apache.arrow.memory.RootAllocator;
3030
import org.apache.arrow.vector.VectorLoader;
31+
import org.apache.arrow.vector.VectorSchemaRoot;
3132
import org.apache.arrow.vector.VectorUnloader;
32-
import org.apache.arrow.vector.complex.MapVector;
33-
import org.apache.arrow.vector.complex.NullableMapVector;
3433
import org.apache.arrow.vector.file.ArrowBlock;
3534
import org.apache.arrow.vector.file.ArrowFooter;
3635
import org.apache.arrow.vector.file.ArrowReader;
3736
import org.apache.arrow.vector.file.ArrowWriter;
3837
import org.apache.arrow.vector.schema.ArrowRecordBatch;
39-
import org.apache.arrow.vector.types.Types;
4038
import org.apache.arrow.vector.types.pojo.Schema;
4139
import org.apache.commons.cli.CommandLine;
4240
import org.apache.commons.cli.CommandLineParser;
@@ -107,13 +105,11 @@ int run(String[] args) {
107105
List<ArrowBlock> recordBatches = footer.getRecordBatches();
108106
for (ArrowBlock rbBlock : recordBatches) {
109107
try (ArrowRecordBatch inRecordBatch = arrowReader.readRecordBatch(rbBlock);
110-
MapVector parent = new MapVector("parent", allocator, null);) {
108+
VectorSchemaRoot root = new VectorSchemaRoot(schema, allocator);) {
111109

112-
NullableMapVector root = parent.addOrGet("root", Types.MinorType.MAP, NullableMapVector.class);
113-
VectorLoader vectorLoader = new VectorLoader(schema, root);
110+
VectorLoader vectorLoader = new VectorLoader(root);
114111
vectorLoader.load(inRecordBatch);
115112

116-
// NullableMapVector outParent = new NullableMapVector("parent", allocator, null);
117113
VectorUnloader vectorUnloader = new VectorUnloader(root);
118114
ArrowRecordBatch recordBatch = vectorUnloader.getRecordBatch();
119115
arrowWriter.writeRecordBatch(recordBatch);

‎java/vector/src/main/codegen/templates/NullableValueVectors.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,7 @@ public List<FieldVector> getChildrenFromFields() {
145145
@Override
146146
public void loadFieldBuffers(ArrowFieldNode fieldNode, List<ArrowBuf> ownBuffers) {
147147
org.apache.arrow.vector.BaseDataValueVector.load(getFieldInnerVectors(), ownBuffers);
148-
// TODO: do something with the sizes in fieldNode?
148+
bits.valueCount = fieldNode.getLength();
149149
}
150150

151151
public List<ArrowBuf> getFieldBuffers() {

‎java/vector/src/main/java/org/apache/arrow/vector/VectorLoader.java‎

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
import org.apache.arrow.vector.schema.ArrowRecordBatch;
2828
import org.apache.arrow.vector.schema.VectorLayout;
2929
import org.apache.arrow.vector.types.pojo.Field;
30-
import org.apache.arrow.vector.types.pojo.Schema;
3130

3231
import com.google.common.collect.Iterators;
3332

@@ -37,22 +36,16 @@
3736
* Loads buffers into vectors
3837
*/
3938
public class VectorLoader {
40-
private final List<FieldVector> fieldVectors;
41-
private final List<Field> fields;
39+
private final VectorSchemaRoot root;
4240

4341
/**
4442
* will create children in root based on schema
4543
* @param schema the expected schema
4644
* @param root the root to add vectors to based on schema
4745
*/
48-
public VectorLoader(Schema schema, FieldVector root) {
46+
public VectorLoader(VectorSchemaRoot root) {
4947
super();
50-
this.fields = schema.getFields();
51-
root.initializeChildrenFromFields(fields);
52-
this.fieldVectors = root.getChildrenFromFields();
53-
if (this.fieldVectors.size() != fields.size()) {
54-
throw new IllegalArgumentException("The root vector did not create the right number of children. found " + fieldVectors.size() + " expected " + fields.size());
55-
}
48+
this.root = root;
5649
}
5750

5851
/**
@@ -63,17 +56,19 @@ public VectorLoader(Schema schema, FieldVector root) {
6356
public void load(ArrowRecordBatch recordBatch) {
6457
Iterator<ArrowBuf> buffers = recordBatch.getBuffers().iterator();
6558
Iterator<ArrowFieldNode> nodes = recordBatch.getNodes().iterator();
59+
List<Field> fields = root.getSchema().getFields();
6660
for (int i = 0; i < fields.size(); ++i) {
6761
Field field = fields.get(i);
68-
FieldVector fieldVector = fieldVectors.get(i);
62+
FieldVector fieldVector = root.getVector(field.getName());
6963
loadBuffers(fieldVector, field, buffers, nodes);
70-
fieldVector.getMutator().setValueCount(recordBatch.getLength());
7164
}
65+
root.setRowCount(recordBatch.getLength());
7266
if (nodes.hasNext() || buffers.hasNext()) {
7367
throw new IllegalArgumentException("not all nodes and buffers where consumed. nodes: " + Iterators.toString(nodes) + " buffers: " + Iterators.toString(buffers));
7468
}
7569
}
7670

71+
7772
private void loadBuffers(FieldVector vector, Field field, Iterator<ArrowBuf> buffers, Iterator<ArrowFieldNode> nodes) {
7873
checkArgument(nodes.hasNext(),
7974
"no more field nodes for for field " + field + " and vector " + vector);
@@ -86,7 +81,7 @@ private void loadBuffers(FieldVector vector, Field field, Iterator<ArrowBuf> buf
8681
try {
8782
vector.loadFieldBuffers(fieldNode, ownBuffers);
8883
} catch (RuntimeException e) {
89-
throw new IllegalArgumentException("Could not load buffers for field " + field);
84+
throw new IllegalArgumentException("Could not load buffers for field " + field, e);
9085
}
9186
List<Field> children = field.getChildren();
9287
if (children.size() > 0) {
@@ -99,4 +94,5 @@ private void loadBuffers(FieldVector vector, Field field, Iterator<ArrowBuf> buf
9994
}
10095
}
10196
}
97+
10298
}

‎java/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
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, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.arrow.vector;
19+
20+
import java.util.ArrayList;
21+
import java.util.Collections;
22+
import java.util.HashMap;
23+
import java.util.List;
24+
import java.util.Map;
25+
26+
import org.apache.arrow.memory.BufferAllocator;
27+
import org.apache.arrow.vector.types.Types;
28+
import org.apache.arrow.vector.types.Types.MinorType;
29+
import org.apache.arrow.vector.types.pojo.Field;
30+
import org.apache.arrow.vector.types.pojo.Schema;
31+
32+
public class VectorSchemaRoot implements AutoCloseable {
33+
34+
private final Schema schema;
35+
private int rowCount;
36+
private final List<FieldVector> fieldVectors;
37+
private final Map<String, FieldVector> fieldVectorsMap = new HashMap<>();
38+
39+
public VectorSchemaRoot(FieldVector parent) {
40+
this.schema = new Schema(parent.getField().getChildren());
41+
this.rowCount = parent.getAccessor().getValueCount();
42+
this.fieldVectors = parent.getChildrenFromFields();
43+
for (int i = 0; i < schema.getFields().size(); ++i) {
44+
Field field = schema.getFields().get(i);
45+
FieldVector vector = fieldVectors.get(i);
46+
fieldVectorsMap.put(field.getName(), vector);
47+
}
48+
}
49+
50+
public VectorSchemaRoot(Schema schema, BufferAllocator allocator) {
51+
super();
52+
this.schema = schema;
53+
List<FieldVector> fieldVectors = new ArrayList<>();
54+
for (Field field : schema.getFields()) {
55+
MinorType minorType = Types.getMinorTypeForArrowType(field.getType());
56+
FieldVector vector = minorType.getNewVector(field.getName(), allocator, null);
57+
vector.initializeChildrenFromFields(field.getChildren());
58+
fieldVectors.add(vector);
59+
fieldVectorsMap.put(field.getName(), vector);
60+
}
61+
this.fieldVectors = Collections.unmodifiableList(fieldVectors);
62+
if (this.fieldVectors.size() != schema.getFields().size()) {
63+
throw new IllegalArgumentException("The root vector did not create the right number of children. found " + fieldVectors.size() + " expected " + schema.getFields().size());
64+
}
65+
}
66+
67+
public List<FieldVector> getFieldVectors() {
68+
return fieldVectors;
69+
}
70+
71+
public FieldVector getVector(String name) {
72+
return fieldVectorsMap.get(name);
73+
}
74+
75+
public Schema getSchema() {
76+
return schema;
77+
}
78+
79+
public int getRowCount() {
80+
return rowCount;
81+
}
82+
83+
public void setRowCount(int rowCount) {
84+
this.rowCount = rowCount;
85+
}
86+
87+
@Override
88+
public void close() {
89+
RuntimeException ex = null;
90+
for (FieldVector fieldVector : fieldVectors) {
91+
try {
92+
fieldVector.close();
93+
} catch (RuntimeException e) {
94+
ex = chain(ex, e);
95+
}
96+
}
97+
if (ex!= null) {
98+
throw ex;
99+
}
100+
}
101+
102+
private RuntimeException chain(RuntimeException root, RuntimeException e) {
103+
if (root == null) {
104+
root = e;
105+
} else {
106+
root.addSuppressed(e);
107+
}
108+
return root;
109+
}
110+
111+
private void printRow(StringBuilder sb, List<Object> row) {
112+
boolean first = true;
113+
for (Object v : row) {
114+
if (first) {
115+
first = false;
116+
} else {
117+
sb.append("\t");
118+
}
119+
sb.append(v);
120+
}
121+
sb.append("\n");
122+
}
123+
124+
public String contentToTSVString() {
125+
StringBuilder sb = new StringBuilder();
126+
List<Object> row = new ArrayList<>(schema.getFields().size());
127+
for (Field field : schema.getFields()) {
128+
row.add(field.getName());
129+
}
130+
printRow(sb, row);
131+
for (int i = 0; i < rowCount; i++) {
132+
row.clear();
133+
for (FieldVector v : fieldVectors) {
134+
row.add(v.getAccessor().getObject(i));
135+
}
136+
printRow(sb, row);
137+
}
138+
return sb.toString();
139+
}
140+
}

0 commit comments

Comments
 (0)