Skip to content

GH-1142: Fix buffer ownership in VectorSchemaRoot addVector/removeVector #1145

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
John-W-Lewis wants to merge 1 commit into
base: main
Choose a base branch
from
+119 −19
Open

GH-1142: Fix buffer ownership in VectorSchemaRoot addVector/removeVector #1145

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 27 additions & 10 deletions vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,10 @@ public FieldVector getVector(int index) {
/**
* Add vector to the record batch, producing a new VectorSchemaRoot.
*
* <p>Buffer ownership is transferred to the returned root via {@link TransferPair}. After this
* operation, the vectors in this root and the added vector are left in a transferred (empty)
* state. This root can be reused by calling {@link #allocateNew()}.
*
* @param index field index
* @param vector vector to be added.
* @return out VectorSchemaRoot with vector added
Expand All @@ -201,23 +205,33 @@ public VectorSchemaRoot addVector(int index, FieldVector vector) {
Preconditions.checkNotNull(vector);
Preconditions.checkArgument(index >= 0 && index <= fieldVectors.size());
List<FieldVector> newVectors = new ArrayList<>();
if (index == fieldVectors.size()) {
newVectors.addAll(fieldVectors);
newVectors.add(vector);
} else {
for (int i = 0; i < fieldVectors.size(); i++) {
if (i == index) {
newVectors.add(vector);
}
newVectors.add(fieldVectors.get(i));
for (int i = 0; i < fieldVectors.size(); i++) {
if (i == index) {
TransferPair addPair = vector.getTransferPair(vector.getAllocator());
addPair.transfer();
newVectors.add((FieldVector) addPair.getTo());
}
FieldVector v = fieldVectors.get(i);
TransferPair transferPair = v.getTransferPair(v.getAllocator());
transferPair.transfer();
newVectors.add((FieldVector) transferPair.getTo());
}
if (index == fieldVectors.size()) {
TransferPair addPair = vector.getTransferPair(vector.getAllocator());
addPair.transfer();
newVectors.add((FieldVector) addPair.getTo());
}
return new VectorSchemaRoot(newVectors);
}

/**
* Remove vector from the record batch, producing a new VectorSchemaRoot.
*
* <p>Buffer ownership is transferred to the returned root via {@link TransferPair}. After this
* operation, the vectors in this root are left in a transferred (empty) state. The removed
* vector's data is not transferred and is released. This root can be reused by calling {@link
* #allocateNew()}.
*
* @param index field index
* @return out VectorSchemaRoot with vector removed
*/
Expand All @@ -226,7 +240,10 @@ public VectorSchemaRoot removeVector(int index) {
List<FieldVector> newVectors = new ArrayList<>();
for (int i = 0; i < fieldVectors.size(); i++) {
if (i != index) {
newVectors.add(fieldVectors.get(i));
FieldVector v = fieldVectors.get(i);
TransferPair transferPair = v.getTransferPair(v.getAllocator());
transferPair.transfer();
newVectors.add((FieldVector) transferPair.getTo());
}
}
return new VectorSchemaRoot(newVectors);
Expand Down
101 changes: 92 additions & 9 deletions vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,14 +157,14 @@ private VectorSchemaRoot createBatch() {
public void testAddVector() {
try (final IntVector intVector1 = new IntVector("intVector1", allocator);
final IntVector intVector2 = new IntVector("intVector2", allocator);
final IntVector intVector3 = new IntVector("intVector3", allocator); ) {
final IntVector intVector3 = new IntVector("intVector3", allocator)) {

VectorSchemaRoot original = new VectorSchemaRoot(Arrays.asList(intVector1, intVector2));
assertEquals(2, original.getFieldVectors().size());

VectorSchemaRoot newRecordBatch = original.addVector(1, intVector3);
assertEquals(3, newRecordBatch.getFieldVectors().size());
assertEquals(intVector3, newRecordBatch.getFieldVectors().get(1));
assertEquals(intVector3.getField(), newRecordBatch.getFieldVectors().get(1).getField());

original.close();
newRecordBatch.close();
Expand All @@ -175,16 +175,16 @@ public void testAddVector() {
public void testAddVectorAtEnd() {
try (final IntVector intVector1 = new IntVector("intVector1", allocator);
final IntVector intVector2 = new IntVector("intVector2", allocator);
final IntVector intVector3 = new IntVector("intVector3", allocator); ) {
final IntVector intVector3 = new IntVector("intVector3", allocator)) {

VectorSchemaRoot original = new VectorSchemaRoot(Arrays.asList(intVector1, intVector2));
assertEquals(2, original.getFieldVectors().size());

VectorSchemaRoot newRecordBatch = original.addVector(2, intVector3);
assertEquals(3, newRecordBatch.getFieldVectors().size());
assertEquals(intVector1, newRecordBatch.getFieldVectors().get(0));
assertEquals(intVector2, newRecordBatch.getFieldVectors().get(1));
assertEquals(intVector3, newRecordBatch.getFieldVectors().get(2));
assertEquals(intVector1.getField(), newRecordBatch.getFieldVectors().get(0).getField());
assertEquals(intVector2.getField(), newRecordBatch.getFieldVectors().get(1).getField());
assertEquals(intVector3.getField(), newRecordBatch.getFieldVectors().get(2).getField());

original.close();
newRecordBatch.close();
Expand All @@ -195,16 +195,16 @@ public void testAddVectorAtEnd() {
public void testRemoveVector() {
try (final IntVector intVector1 = new IntVector("intVector1", allocator);
final IntVector intVector2 = new IntVector("intVector2", allocator);
final IntVector intVector3 = new IntVector("intVector3", allocator); ) {
final IntVector intVector3 = new IntVector("intVector3", allocator)) {

VectorSchemaRoot original =
new VectorSchemaRoot(Arrays.asList(intVector1, intVector2, intVector3));
assertEquals(3, original.getFieldVectors().size());

VectorSchemaRoot newRecordBatch = original.removeVector(0);
assertEquals(2, newRecordBatch.getFieldVectors().size());
assertEquals(intVector2, newRecordBatch.getFieldVectors().get(0));
assertEquals(intVector3, newRecordBatch.getFieldVectors().get(1));
assertEquals(intVector2.getField(), newRecordBatch.getFieldVectors().get(0).getField());
assertEquals(intVector3.getField(), newRecordBatch.getFieldVectors().get(1).getField());

original.close();
newRecordBatch.close();
Expand Down Expand Up @@ -344,4 +344,87 @@ public void testSchemaSync() {
assertFalse(schemaRoot.syncSchema());
}
}

@Test
public void testAddVectorOwnership() {
try (final IntVector intVector1 = new IntVector("intVector1", allocator);
final IntVector intVector2 = new IntVector("intVector2", allocator);
final IntVector intVector3 = new IntVector("intVector3", allocator)) {

intVector1.allocateNew();
intVector2.allocateNew();
intVector3.allocateNew();
for (int i = 0; i < 5; i++) {
intVector1.setSafe(i, i * 10);
intVector2.setSafe(i, i * 20);
intVector3.setSafe(i, i * 30);
}
intVector1.setValueCount(5);
intVector2.setValueCount(5);
intVector3.setValueCount(5);

VectorSchemaRoot original =
new VectorSchemaRoot(Arrays.asList(intVector1, intVector2));
original.setRowCount(5);

VectorSchemaRoot result = original.addVector(1, intVector3);

// Close the original root and the added vector -- the result should still have valid data
original.close();
intVector3.close();

assertEquals(3, result.getFieldVectors().size());
assertEquals(5, result.getRowCount());
IntVector resultVec0 = (IntVector) result.getVector(0);
IntVector resultVec1 = (IntVector) result.getVector(1);
IntVector resultVec2 = (IntVector) result.getVector(2);
for (int i = 0; i < 5; i++) {
assertEquals(i * 10, resultVec0.get(i));
assertEquals(i * 30, resultVec1.get(i));
assertEquals(i * 20, resultVec2.get(i));
}

result.close();
}
}

@Test
public void testRemoveVectorOwnership() {
try (final IntVector intVector1 = new IntVector("intVector1", allocator);
final IntVector intVector2 = new IntVector("intVector2", allocator);
final IntVector intVector3 = new IntVector("intVector3", allocator)) {

intVector1.allocateNew();
intVector2.allocateNew();
intVector3.allocateNew();
for (int i = 0; i < 5; i++) {
intVector1.setSafe(i, i * 10);
intVector2.setSafe(i, i * 20);
intVector3.setSafe(i, i * 30);
}
intVector1.setValueCount(5);
intVector2.setValueCount(5);
intVector3.setValueCount(5);

VectorSchemaRoot original =
new VectorSchemaRoot(Arrays.asList(intVector1, intVector2, intVector3));
original.setRowCount(5);

VectorSchemaRoot result = original.removeVector(1);

// Close the original root -- the result should still have valid data
original.close();

assertEquals(2, result.getFieldVectors().size());
assertEquals(5, result.getRowCount());
IntVector resultVec0 = (IntVector) result.getVector(0);
IntVector resultVec1 = (IntVector) result.getVector(1);
for (int i = 0; i < 5; i++) {
assertEquals(i * 10, resultVec0.get(i));
assertEquals(i * 30, resultVec1.get(i));
}

result.close();
}
}
}
Loading