diff --git a/java/vector/src/main/java/org/apache/arrow/vector/BitVector.java b/java/vector/src/main/java/org/apache/arrow/vector/BitVector.java index b6f7323a796..3887da4a618 100644 --- a/java/vector/src/main/java/org/apache/arrow/vector/BitVector.java +++ b/java/vector/src/main/java/org/apache/arrow/vector/BitVector.java @@ -444,6 +444,72 @@ public void setSafe(int index, int isSet, int value) { set(index, isSet, value); } + /** + * Set the element at the given index to one. + * + * @param index position of element + */ + public void setToOne(int index) { + BitVectorHelper.setValidityBitToOne(validityBuffer, index); + BitVectorHelper.setValidityBitToOne(valueBuffer, index); + } + + /** + * Same as {@link #setToOne(int)} except that it handles the case when + * index is greater than or equal to current value capacity of the vector. + * + * @param index position of the element + */ + public void setSafeToOne(int index) { + handleSafe(index); + setToOne(index); + } + + /** + * Set count bits to 1 in data starting at firstBitIndex + * + * @param firstBitIndex the index of the first bit to set + * @param count the number of bits to set + */ + public void setRangeToOne(int firstBitIndex, int count) { + int startByteIndex = BitVectorHelper.byteIndex(firstBitIndex); + final int lastBitIndex = firstBitIndex + count; + final int endByteIndex = BitVectorHelper.byteIndex(lastBitIndex); + final int startByteBitIndex = BitVectorHelper.bitIndex(firstBitIndex); + final int endBytebitIndex = BitVectorHelper.bitIndex(lastBitIndex); + if (count < 8 && startByteIndex == endByteIndex) { + // handles the case where we don't have a first and a last byte + byte bitMask = 0; + for (int i = startByteBitIndex; i < endBytebitIndex; ++i) { + bitMask |= (byte) (1L << i); + } + BitVectorHelper.setBitMaskedByte(validityBuffer, startByteIndex, bitMask); + BitVectorHelper.setBitMaskedByte(valueBuffer, startByteIndex, bitMask); + } else { + // fill in first byte (if it's not full) + if (startByteBitIndex != 0) { + final byte bitMask = (byte) (0xFFL << startByteBitIndex); + BitVectorHelper.setBitMaskedByte(validityBuffer, startByteIndex, bitMask); + BitVectorHelper.setBitMaskedByte(valueBuffer, startByteIndex, bitMask); + ++startByteIndex; + } + + // fill in one full byte at a time + for (int i = startByteIndex; i < endByteIndex; i++) { + validityBuffer.setByte(i, 0xFF); + valueBuffer.setByte(i, 0xFF); + } + + // fill in the last byte (if it's not full) + if (endBytebitIndex != 0) { + final int byteIndex = BitVectorHelper.byteIndex(lastBitIndex - endBytebitIndex); + final byte bitMask = (byte) (0xFFL >>> ((8 - endBytebitIndex) & 7)); + BitVectorHelper.setBitMaskedByte(validityBuffer, byteIndex, bitMask); + BitVectorHelper.setBitMaskedByte(valueBuffer, byteIndex, bitMask); + } + } + } + /****************************************************************** * * diff --git a/java/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java b/java/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java index 2d4db85c583..8322a1ac8fc 100644 --- a/java/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java +++ b/java/vector/src/main/java/org/apache/arrow/vector/BitVectorHelper.java @@ -198,4 +198,18 @@ public static ArrowBuf loadValidityBuffer(final ArrowFieldNode fieldNode, return newBuffer; } + + /** + * Set the byte of the given index in the data buffer by applying a bit mask to + * the current byte at that index. + * + * @param data + * @param byteIndex + * @param bitMask + */ + static void setBitMaskedByte(ArrowBuf data, int byteIndex, byte bitMask) { + byte currentByte = data.getByte(byteIndex); + currentByte |= bitMask; + data.setByte(byteIndex, currentByte); + } } diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestBitVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestBitVector.java index 36365fa9d40..a59e5cdd48f 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/TestBitVector.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/TestBitVector.java @@ -232,7 +232,7 @@ public void testReallocAfterVectorTransfer1() { for (int i = 0; i < valueCapacity; i++) { if ((i & 1) == 1) { - vector.set(i, 1); + vector.setToOne(i); } } @@ -246,12 +246,12 @@ public void testReallocAfterVectorTransfer1() { } /* trigger first realloc */ - vector.setSafe(valueCapacity, 1); + vector.setSafeToOne(valueCapacity); assertEquals(valueCapacity * 2, vector.getValueCapacity()); for (int i = valueCapacity; i < valueCapacity*2; i++) { if ((i & 1) == 1) { - vector.set(i, 1); + vector.setToOne(i); } } @@ -265,12 +265,12 @@ public void testReallocAfterVectorTransfer1() { } /* trigger second realloc */ - vector.setSafe(valueCapacity*2, 1); + vector.setSafeToOne(valueCapacity*2); assertEquals(valueCapacity * 4, vector.getValueCapacity()); for (int i = valueCapacity*2; i < valueCapacity*4; i++) { if ((i & 1) == 1) { - vector.set(i, 1); + vector.setToOne(i); } } @@ -291,7 +291,7 @@ public void testReallocAfterVectorTransfer1() { assertEquals(valueCapacity * 4, toVector.getValueCapacity()); /* realloc the toVector */ - toVector.setSafe(valueCapacity * 4, 1); + toVector.setSafeToOne(valueCapacity * 4); for (int i = 0; i < toVector.getValueCapacity(); i++) { if (i <= valueCapacity * 4) { @@ -505,9 +505,7 @@ private void validateRange(int length, int start, int count) { try (BitVector bitVector = new BitVector("bits", allocator)) { bitVector.reset(); bitVector.allocateNew(length); - for (int i = start; i < start + count; i++) { - bitVector.set(i, 1); - } + bitVector.setRangeToOne(start, count); for (int i = 0; i < start; i++) { Assert.assertTrue(desc + i, bitVector.isNull(i)); }