Skip to content
This repository was archived by the owner on Oct 21, 2024. It is now read-only.

Commit a22d58a

Browse files
tianchen92tianchen92
committed
ARROW-6871: [Java] Enhance TransferPair related parameters check and tests
1 parent 45ad346 commit a22d58a

File tree

11 files changed

+331
-18
lines changed

11 files changed

+331
-18
lines changed

‎java/algorithm/src/main/java/org/apache/arrow/algorithm/sort/FixedWidthInPlaceVectorSorter.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ public void sortInPlace(V vec, VectorValueComparator<V> comparator) {
4646
this.comparator = comparator;
4747
this.pivotBuffer = (V) vec.getField().createVector(vec.getAllocator());
4848
this.pivotBuffer.allocateNew(1);
49+
this.pivotBuffer.setValueCount(1);
4950

5051
comparator.attachVectors(vec, pivotBuffer);
5152
quickSort();

‎java/performance/src/test/java/org/apache/arrow/vector/util/TransferPairBenchmarks.java‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
package org.apache.arrow.vector.util;
19+
20+
import java.nio.charset.StandardCharsets;
21+
import java.util.concurrent.TimeUnit;
22+
23+
import org.apache.arrow.memory.BufferAllocator;
24+
import org.apache.arrow.memory.RootAllocator;
25+
import org.apache.arrow.vector.IntVector;
26+
import org.apache.arrow.vector.VarCharVector;
27+
28+
import org.openjdk.jmh.annotations.Benchmark;
29+
import org.openjdk.jmh.annotations.BenchmarkMode;
30+
import org.openjdk.jmh.annotations.Mode;
31+
import org.openjdk.jmh.annotations.OutputTimeUnit;
32+
import org.openjdk.jmh.annotations.Scope;
33+
import org.openjdk.jmh.annotations.Setup;
34+
import org.openjdk.jmh.annotations.State;
35+
import org.openjdk.jmh.annotations.TearDown;
36+
import org.openjdk.jmh.runner.Runner;
37+
import org.openjdk.jmh.runner.RunnerException;
38+
import org.openjdk.jmh.runner.options.Options;
39+
import org.openjdk.jmh.runner.options.OptionsBuilder;
40+
41+
/**
42+
* Benchmarks for {@link TransferPair}.
43+
*/
44+
@State(Scope.Benchmark)
45+
public class TransferPairBenchmarks {
46+
47+
private static final int VECTOR_LENGTH = 1024;
48+
49+
private static final int ALLOCATOR_CAPACITY = 1024 * 1024;
50+
51+
private BufferAllocator allocator;
52+
53+
private IntVector intVector;
54+
55+
private VarCharVector varCharVector;
56+
57+
/**
58+
* Setup benchmarks.
59+
*/
60+
@Setup
61+
public void prepare() {
62+
allocator = new RootAllocator(ALLOCATOR_CAPACITY);
63+
intVector = new IntVector("intVector", allocator);
64+
varCharVector = new VarCharVector("varcharVector", allocator);
65+
66+
intVector.allocateNew(VECTOR_LENGTH);
67+
varCharVector.allocateNew(VECTOR_LENGTH);
68+
69+
for (int i = 0; i < VECTOR_LENGTH; i++) {
70+
if (i % 3 == 0) {
71+
intVector.setNull(i);
72+
varCharVector.setNull(i);
73+
} else {
74+
intVector.setSafe(i, i * i);
75+
varCharVector.setSafe(i, ("teststring" + i).getBytes(StandardCharsets.UTF_8));
76+
}
77+
}
78+
intVector.setValueCount(VECTOR_LENGTH);
79+
varCharVector.setValueCount(VECTOR_LENGTH);
80+
}
81+
82+
/**
83+
* Tear down benchmarks.
84+
*/
85+
@TearDown
86+
public void tearDown() {
87+
intVector.close();
88+
varCharVector.close();;
89+
allocator.close();
90+
}
91+
92+
@Benchmark
93+
@BenchmarkMode(Mode.AverageTime)
94+
@OutputTimeUnit(TimeUnit.MICROSECONDS)
95+
public int copyValueSafeForIntVector() {
96+
IntVector toVector = new IntVector("intVector", allocator);
97+
toVector.setValueCount(VECTOR_LENGTH);
98+
TransferPair transferPair = intVector.makeTransferPair(toVector);
99+
for (int i = 0; i < VECTOR_LENGTH; i++) {
100+
transferPair.copyValueSafe(i, i);
101+
}
102+
toVector.close();
103+
return 0;
104+
}
105+
106+
@Benchmark
107+
@BenchmarkMode(Mode.AverageTime)
108+
@OutputTimeUnit(TimeUnit.MICROSECONDS)
109+
public int copyValueSafeForVarcharVector() {
110+
VarCharVector toVector = new VarCharVector("varcharVector", allocator);
111+
toVector.setValueCount(VECTOR_LENGTH);
112+
TransferPair transferPair = varCharVector.makeTransferPair(toVector);
113+
for (int i = 0; i < VECTOR_LENGTH; i++) {
114+
transferPair.copyValueSafe(i, i);
115+
}
116+
toVector.close();
117+
return 0;
118+
}
119+
120+
public static void main(String [] args) throws RunnerException {
121+
Options opt = new OptionsBuilder()
122+
.include(TransferPairBenchmarks.class.getSimpleName())
123+
.forks(1)
124+
.build();
125+
126+
new Runner(opt).run();
127+
}
128+
}
129+

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -408,6 +408,8 @@ public TransferPair makeTransferPair(ValueVector target) {
408408
@Override
409409
public void copyFrom(int inIndex, int outIndex, ValueVector from) {
410410
Preconditions.checkArgument(this.getMinorType() == from.getMinorType());
411+
Preconditions.checkArgument(inIndex >= 0 && inIndex < from.getValueCount(),
412+
"Invalid inIndex: %s", inIndex);
411413
UnionVector fromCast = (UnionVector) from;
412414
fromCast.getReader().setPosition(inIndex);
413415
getWriter().setPosition(outIndex);
@@ -469,7 +471,10 @@ public void transfer() {
469471

470472
@Override
471473
public void splitAndTransfer(int startIndex, int length) {
472-
Preconditions.checkArgument(startIndex + length <= valueCount);
474+
Preconditions.checkArgument(startIndex >= 0 && startIndex < valueCount,
475+
"Invalid startIndex: %s", startIndex);
476+
Preconditions.checkArgument(startIndex + length <= valueCount,
477+
"Invalid length: %s", length);
473478
to.clear();
474479
internalStructVectorTransferPair.splitAndTransfer(startIndex, length);
475480
final int startPoint = startIndex * TYPE_WIDTH;

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -584,7 +584,10 @@ public void transferTo(BaseFixedWidthVector target) {
584584
*/
585585
public void splitAndTransferTo(int startIndex, int length,
586586
BaseFixedWidthVector target) {
587-
Preconditions.checkArgument(startIndex + length <= valueCount);
587+
Preconditions.checkArgument(startIndex >= 0 && startIndex < valueCount,
588+
"Invalid startIndex: %s", startIndex);
589+
Preconditions.checkArgument(startIndex + length <= valueCount,
590+
"Invalid length: %s", length);
588591
compareTypes(target, "splitAndTransferTo");
589592
target.clear();
590593
splitAndTransferValidityBuffer(startIndex, length, target);
@@ -830,6 +833,8 @@ protected void handleSafe(int index) {
830833
@Override
831834
public void copyFrom(int fromIndex, int thisIndex, ValueVector from) {
832835
Preconditions.checkArgument(this.getMinorType() == from.getMinorType());
836+
Preconditions.checkArgument(fromIndex >= 0 && fromIndex < from.getValueCount(),
837+
"Invalid fromIndex: %s", fromIndex);
833838
if (from.isNull(fromIndex)) {
834839
BitVectorHelper.unsetBit(this.getValidityBuffer(), thisIndex);
835840
} else {
@@ -851,6 +856,8 @@ public void copyFrom(int fromIndex, int thisIndex, ValueVector from) {
851856
@Override
852857
public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) {
853858
Preconditions.checkArgument(this.getMinorType() == from.getMinorType());
859+
Preconditions.checkArgument(fromIndex >= 0 && fromIndex < from.getValueCount(),
860+
"Invalid fromIndex: %s", fromIndex);
854861
handleSafe(thisIndex);
855862
copyFrom(fromIndex, thisIndex, from);
856863
}

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -714,6 +714,10 @@ public void transferTo(BaseVariableWidthVector target) {
714714
*/
715715
public void splitAndTransferTo(int startIndex, int length,
716716
BaseVariableWidthVector target) {
717+
Preconditions.checkArgument(startIndex >= 0 && startIndex < valueCount,
718+
"Invalid startIndex: %s", startIndex);
719+
Preconditions.checkArgument(startIndex + length <= valueCount,
720+
"Invalid length: %s", length);
717721
compareTypes(target, "splitAndTransferTo");
718722
target.clear();
719723
splitAndTransferValidityBuffer(startIndex, length, target);
@@ -750,7 +754,6 @@ private void splitAndTransferOffsetBuffer(int startIndex, int length, BaseVariab
750754
*/
751755
private void splitAndTransferValidityBuffer(int startIndex, int length,
752756
BaseVariableWidthVector target) {
753-
Preconditions.checkArgument(startIndex + length <= valueCount);
754757
int firstByteSource = BitVectorHelper.byteIndex(startIndex);
755758
int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1);
756759
int byteSizeTarget = getValidityBufferSizeFromCount(length);
@@ -1336,6 +1339,8 @@ public void copyFrom(int fromIndex, int thisIndex, ValueVector from) {
13361339
@Override
13371340
public void copyFromSafe(int fromIndex, int thisIndex, ValueVector from) {
13381341
Preconditions.checkArgument(this.getMinorType() == from.getMinorType());
1342+
Preconditions.checkArgument(fromIndex >= 0 && fromIndex < from.getValueCount(),
1343+
"Invalid fromIndex: %s", fromIndex);
13391344
if (from.isNull(fromIndex)) {
13401345
handleSafe(thisIndex, 0);
13411346
fillHoles(thisIndex);

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,10 @@ public int getBufferSize() {
158158
* @param target destination vector
159159
*/
160160
public void splitAndTransferTo(int startIndex, int length, BaseFixedWidthVector target) {
161+
Preconditions.checkArgument(startIndex >= 0 && startIndex < valueCount,
162+
"Invalid startIndex: %s", startIndex);
163+
Preconditions.checkArgument(startIndex + length <= valueCount,
164+
"Invalid length: %s", length);
161165
compareTypes(target, "splitAndTransferTo");
162166
target.clear();
163167
target.validityBuffer = splitAndTransferBuffer(startIndex, length, target,
@@ -174,7 +178,6 @@ private ArrowBuf splitAndTransferBuffer(
174178
BaseFixedWidthVector target,
175179
ArrowBuf sourceBuffer,
176180
ArrowBuf destBuffer) {
177-
assert startIndex + length <= valueCount;
178181
int firstByteSource = BitVectorHelper.byteIndex(startIndex);
179182
int lastByteSource = BitVectorHelper.byteIndex(valueCount - 1);
180183
int byteSizeTarget = getValidityBufferSizeFromCount(length);

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -590,7 +590,10 @@ public void transfer() {
590590

591591
@Override
592592
public void splitAndTransfer(int startIndex, int length) {
593-
Preconditions.checkArgument(startIndex + length <= valueCount);
593+
Preconditions.checkArgument(startIndex >= 0 && startIndex < valueCount,
594+
"Invalid startIndex: %s", startIndex);
595+
Preconditions.checkArgument(startIndex + length <= valueCount,
596+
"Invalid length: %s", length);
594597
final int startPoint = listSize * startIndex;
595598
final int sliceLength = listSize * length;
596599
to.clear();

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -364,6 +364,8 @@ public void copyFromSafe(int inIndex, int outIndex, ValueVector from) {
364364
@Override
365365
public void copyFrom(int inIndex, int outIndex, ValueVector from) {
366366
Preconditions.checkArgument(this.getMinorType() == from.getMinorType());
367+
Preconditions.checkArgument(inIndex >= 0 && inIndex < from.getValueCount(),
368+
"Invalid inIndex: %s", inIndex);
367369
FieldReader in = from.getReader();
368370
in.setPosition(inIndex);
369371
FieldWriter out = getWriter();
@@ -493,7 +495,10 @@ public void transfer() {
493495
*/
494496
@Override
495497
public void splitAndTransfer(int startIndex, int length) {
496-
Preconditions.checkArgument(startIndex + length <= valueCount);
498+
Preconditions.checkArgument(startIndex >= 0 && startIndex < valueCount,
499+
"Invalid startIndex: %s", startIndex);
500+
Preconditions.checkArgument(startIndex + length <= valueCount,
501+
"Invalid length: %s", length);
497502
final int startPoint = offsetBuffer.getInt(startIndex * OFFSET_WIDTH);
498503
final int sliceLength = offsetBuffer.getInt((startIndex + length) * OFFSET_WIDTH) - startPoint;
499504
to.clear();

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,8 @@ public FieldReader getReader() {
9999
@Override
100100
public void copyFrom(int fromIndex, int thisIndex, ValueVector from) {
101101
Preconditions.checkArgument(this.getMinorType() == from.getMinorType());
102+
Preconditions.checkArgument(fromIndex >= 0 && fromIndex < from.getValueCount(),
103+
"Invalid fromIndex: %s", fromIndex);
102104
if (ephPair == null || ephPair.from != from) {
103105
ephPair = (StructTransferPair) from.makeTransferPair(this);
104106
}

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -195,7 +195,10 @@ public void copyValueSafe(int fromIndex, int toIndex) {
195195

196196
@Override
197197
public void splitAndTransfer(int startIndex, int length) {
198-
Preconditions.checkArgument(startIndex + length <= valueCount);
198+
Preconditions.checkArgument(startIndex >= 0 && startIndex < valueCount,
199+
"Invalid startIndex: %s", startIndex);
200+
Preconditions.checkArgument(startIndex + length <= valueCount,
201+
"Invalid length: %s", length);
199202
target.clear();
200203
splitAndTransferValidityBuffer(startIndex, length, target);
201204
super.splitAndTransfer(startIndex, length);

0 commit comments

Comments
 (0)