liyafan82 commented on a change in pull request #9744:
URL: https://github.com/apache/arrow/pull/9744#discussion_r597399889
##########
File path:
java/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/BinaryConsumer.java
##########
@@ -57,30 +55,44 @@ public BinaryConsumer(VarBinaryVector vector, int index) {
}
}
+ /**
+ * Set the variable length element at the specified index to the supplied
+ * byte array supplied in the input stream. Besides, it will handle the
+ * case where index and length of new element are beyond the existing
+ * capacity of the vector.
+ *
+ * @param vector binary vector
+ * @param index position of the element to set
+ * @param is input stream
+ */
+ public static void setSafe(VarBinaryVector vector, int index, InputStream
is) throws IOException {
+ assert index >= 0;
+ while (index >= vector.getValueCapacity()) {
+ vector.reallocValidityAndOffsetBuffers();
+ }
+ final int startOffset = vector.getStartOffset(index);
+ final ArrowBuf offsetBuffer = vector.getOffsetBuffer();
+ int dataLength = 0;
+ int read;
+ byte[] bytes = new byte[BUFFER_SIZE];
+ while ((read = is.read(bytes)) != -1) {
+ while (vector.getDataBuffer().capacity() < (startOffset + dataLength +
read)) {
+ vector.reallocDataBuffer();
+ }
+ vector.getDataBuffer().setBytes(startOffset + dataLength, bytes, 0,
read);
+ dataLength += read;
+ }
+ offsetBuffer.setInt((index + 1) * VarBinaryVector.OFFSET_WIDTH,
startOffset + dataLength);
+ BitVectorHelper.setBit(vector.getValidityBuffer(), index);
+ vector.setLastSet(index);
+ }
+
/**
* consume a InputStream.
*/
public void consume(InputStream is) throws IOException {
if (is != null) {
-
- int read;
- byte[] bytes = new byte[BUFFER_SIZE];
- int totalBytes = 0;
-
- ArrowBuf dataBuffer = vector.getDataBuffer();
- ArrowBuf offsetBuffer = vector.getOffsetBuffer();
- int startIndex = offsetBuffer.getInt(currentIndex * 4);
- while ((read = is.read(bytes)) != -1) {
- while ((dataBuffer.writerIndex() + read) > dataBuffer.capacity()) {
- vector.reallocDataBuffer();
- }
- PlatformDependent.copyMemory(bytes, 0,
- dataBuffer.memoryAddress() + startIndex + totalBytes, read);
- totalBytes += read;
- }
- offsetBuffer.setInt((currentIndex + 1) * 4, startIndex + totalBytes);
- BitVectorHelper.setBit(vector.getValidityBuffer(), currentIndex);
- vector.setLastSet(currentIndex);
+ setSafe(vector, currentIndex, is);
Review comment:
nit: it seems this method is simply enough. why extract another method?
--
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.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
