From 422833e31a2db79fcb2374ccb99903a5b9d83cbf Mon Sep 17 00:00:00 2001 From: "chenweiguo.vc" Date: Mon, 23 Sep 2024 16:11:14 +0800 Subject: [PATCH] [Java] Implement Holder-based functions for ViewVarCharVector --- .../arrow/vector/ViewVarCharVector.java | 61 ++++++++++++++---- .../arrow/vector/TestVarCharViewVector.java | 63 +++++++++++++++++++ 2 files changed, 111 insertions(+), 13 deletions(-) diff --git a/java/vector/src/main/java/org/apache/arrow/vector/ViewVarCharVector.java b/java/vector/src/main/java/org/apache/arrow/vector/ViewVarCharVector.java index 900e0493fe804..55d7b78442fc0 100644 --- a/java/vector/src/main/java/org/apache/arrow/vector/ViewVarCharVector.java +++ b/java/vector/src/main/java/org/apache/arrow/vector/ViewVarCharVector.java @@ -18,6 +18,7 @@ import static org.apache.arrow.vector.NullCheckingForGet.NULL_CHECKING_ENABLED; +import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.ReusableBuffer; import org.apache.arrow.vector.complex.impl.ViewVarCharReaderImpl; @@ -143,9 +144,31 @@ public void read(int index, ReusableBuffer buffer) { * @param holder data holder to be populated by this function */ public void get(int index, NullableViewVarCharHolder holder) { - // TODO: https://github.com/apache/arrow/issues/40937 - throw new UnsupportedOperationException( - "NullableViewVarCharHolder get operation not supported"); + final int dataLength = getValueLength(index); + if (isSet(index) == 0) { + holder.isSet = 0; + return; + } + holder.isSet = 1; + if (dataLength > INLINE_SIZE) { + // data is in the data buffer + // get buffer index + final int bufferIndex = + viewBuffer.getInt(((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH); + // get data offset + final int dataOffset = + viewBuffer.getInt( + ((long) index * ELEMENT_SIZE) + LENGTH_WIDTH + PREFIX_WIDTH + BUF_INDEX_WIDTH); + holder.buffer = dataBuffers.get(bufferIndex); + holder.start = dataOffset; + holder.end = dataOffset + dataLength; + } else { + final long dataOffset = ((long) index * ELEMENT_SIZE) + LENGTH_WIDTH; + // data is in the value buffer + holder.buffer = viewBuffer; + holder.start = (int) dataOffset; + holder.end = (int) dataOffset + dataLength; + } } /*----------------------------------------------------------------* @@ -162,8 +185,11 @@ public void get(int index, NullableViewVarCharHolder holder) { * @param holder holder that carries data buffer. */ public void set(int index, ViewVarCharHolder holder) { - // TODO: https://github.com/apache/arrow/issues/40937 - throw new UnsupportedOperationException("ViewVarCharHolder set operation not supported"); + int start = holder.start; + int length = holder.end - start; + ArrowBuf dataBuf = holder.buffer; + setBytes(index, dataBuf, start, length); + lastSet = index; } /** @@ -174,8 +200,9 @@ public void set(int index, ViewVarCharHolder holder) { * @param holder holder that carries data buffer. */ public void setSafe(int index, ViewVarCharHolder holder) { - // TODO: https://github.com/apache/arrow/issues/40937 - throw new UnsupportedOperationException("ViewVarCharHolder setSafe operation not supported"); + int length = holder.end - holder.start; + handleSafe(index, length); + set(index, holder); } /** @@ -186,9 +213,16 @@ public void setSafe(int index, ViewVarCharHolder holder) { * @param holder holder that carries data buffer. */ public void set(int index, NullableViewVarCharHolder holder) { - // TODO: https://github.com/apache/arrow/issues/40937 - throw new UnsupportedOperationException( - "NullableViewVarCharHolder set operation not supported"); + if (holder.isSet == 0) { + setNull(index); + } else { + BitVectorHelper.setBit(validityBuffer, index); + int start = holder.start; + int length = holder.end - start; + ArrowBuf dataBuf = holder.buffer; + setBytes(index, dataBuf, start, length); + } + lastSet = index; } /** @@ -199,9 +233,9 @@ public void set(int index, NullableViewVarCharHolder holder) { * @param holder holder that carries data buffer. */ public void setSafe(int index, NullableViewVarCharHolder holder) { - // TODO: https://github.com/apache/arrow/issues/40937 - throw new UnsupportedOperationException( - "NullableViewVarCharHolder setSafe operation not supported"); + int length = holder.end - holder.start; + handleSafe(index, length); + set(index, holder); } /** @@ -279,6 +313,7 @@ public TransferPair makeTransferPair(ValueVector to) { } private class TransferImpl implements TransferPair { + ViewVarCharVector to; public TransferImpl(String ref, BufferAllocator allocator) { diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestVarCharViewVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestVarCharViewVector.java index 232eec9ef1b23..30fe993c7a783 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/TestVarCharViewVector.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/TestVarCharViewVector.java @@ -48,6 +48,8 @@ import org.apache.arrow.memory.rounding.DefaultRoundingPolicy; import org.apache.arrow.memory.util.ArrowBufPointer; import org.apache.arrow.memory.util.CommonUtil; +import org.apache.arrow.util.AutoCloseables; +import org.apache.arrow.vector.holders.NullableViewVarCharHolder; import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.testing.ValueVectorDataPopulator; import org.apache.arrow.vector.types.Types; @@ -371,6 +373,67 @@ public void testMixedAllocation() { } } + @Test + public void testSetNullableViewVarCharHolder() { + try (final ViewVarCharVector viewVarCharVector = new ViewVarCharVector("myvector", allocator)) { + viewVarCharVector.allocateNew(0, 0); + final List strings = List.of(STR0, STR1, STR2, STR3, STR4, STR5); + + NullableViewVarCharHolder stringHolder = new NullableViewVarCharHolder(); + + setAndCheck(viewVarCharVector, 1, strings.get(0), stringHolder); + + // set not null + for (int i = 0; i < strings.size(); i++) { + setAndCheck(viewVarCharVector, i, strings.get(i), stringHolder); + } + + // set null + setAndCheck(viewVarCharVector, 6, null, stringHolder); + + // copy by holder + // len < 12 + copyAndCheck(viewVarCharVector, 0, stringHolder, 7); + // len > 12 + copyAndCheck(viewVarCharVector, 2, stringHolder, 8); + // null + copyAndCheck(viewVarCharVector, 6, stringHolder, 9); + } + } + + private static void copyAndCheck( + ViewVarCharVector viewVarCharVector, + int index, + NullableViewVarCharHolder stringHolder, + int index1) { + viewVarCharVector.get(index, stringHolder); + viewVarCharVector.setSafe(index1, stringHolder); + assertArrayEquals(viewVarCharVector.get(index), viewVarCharVector.get(index1)); + } + + private void setAndCheck( + ViewVarCharVector viewVarCharVector, + int index, + byte[] str, + NullableViewVarCharHolder stringHolder) { + ArrowBuf buf = null; + if (null == str) { + stringHolder.isSet = 0; + } else { + buf = allocator.buffer(str.length); + buf.setBytes(0, str); + stringHolder.isSet = 1; + stringHolder.start = 0; + stringHolder.end = str.length; + stringHolder.buffer = buf; + } + + viewVarCharVector.setSafe(index, stringHolder); + // verify results + assertArrayEquals(str, viewVarCharVector.get(index)); + AutoCloseables.closeNoChecked(buf); + } + @Test public void testAllocationIndexOutOfBounds() { assertThrows(