rdblue commented on a change in pull request #3533:
URL: https://github.com/apache/iceberg/pull/3533#discussion_r749712875
##########
File path:
arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedArrowReader.java
##########
@@ -407,35 +412,52 @@ public void setBatchSize(int batchSize) {
}
private static final class PositionVectorReader extends
VectorizedArrowReader {
+ private static final Field ROW_POSITION_ARROW_FIELD =
ArrowSchemaUtil.convert(MetadataColumns.ROW_POSITION);
+ private final BufferAllocator bufferAllocator =
ArrowAllocation.rootAllocator();
+ private final boolean setArrowValidityVector;
private long rowStart;
+ private int batchSize;
private NullabilityHolder nulls;
+ PositionVectorReader(boolean setArrowValidityVector) {
+ this.setArrowValidityVector = setArrowValidityVector;
+ }
+
@Override
public VectorHolder read(VectorHolder reuse, int numValsToRead) {
- Field arrowField = ArrowSchemaUtil.convert(MetadataColumns.ROW_POSITION);
- FieldVector vec =
arrowField.createVector(ArrowAllocation.rootAllocator());
-
- if (reuse != null) {
- vec.setValueCount(0);
- nulls.reset();
+ FieldVector vec;
+ if (reuse == null) {
+ vec = newVector(batchSize);
} else {
- ((BigIntVector) vec).allocateNew(numValsToRead);
- for (int i = 0; i < numValsToRead; i += 1) {
- vec.getDataBuffer().setLong(i * Long.BYTES, rowStart + i);
- }
+ FieldVector reusedVector = reuse.vector();
+ reusedVector.setValueCount(0);
+ vec = reusedVector;
+ }
+
+ ArrowBuf dataBuffer = vec.getDataBuffer();
+ for (int i = 0; i < numValsToRead; i += 1) {
+ dataBuffer.setLong((long) i * Long.BYTES, rowStart + i);
+ }
+
+ if (setArrowValidityVector) {
+ ArrowBuf validityBuffer = vec.getValidityBuffer();
for (int i = 0; i < numValsToRead; i += 1) {
- BitVectorHelper.setBit(vec.getValidityBuffer(), i);
+ BitVectorHelper.setBit(validityBuffer, i);
}
- nulls = new NullabilityHolder(numValsToRead);
}
rowStart += numValsToRead;
vec.setValueCount(numValsToRead);
- nulls.setNotNulls(0, numValsToRead);
return new VectorHolder.PositionVectorHolder(vec,
MetadataColumns.ROW_POSITION.type(), nulls);
}
+ private BigIntVector newVector(int valueCount) {
Review comment:
Could be `static` if we pass in the allocator as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
