ARROW-1867: [Java] Add missing methods to BitVector from legacy vector class

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);
+      }
+    }
+  }
+
 
   /******************************************************************
    *                                                                *

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);
+  }
 }

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));
       }