Skip to content

Commit

Permalink
[Java] Implement Holder-based functions for ViewVarCharVector
Browse files Browse the repository at this point in the history
  • Loading branch information
@ViggoC
ViggoC committed Sep 23, 2024
1 parent 6730c1a commit 422833e
Show file tree
Hide file tree
Showing 2 changed files with 111 additions and 13 deletions.
61 changes: 48 additions & 13 deletions java/vector/src/main/java/org/apache/arrow/vector/ViewVarCharVector.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
}

/*----------------------------------------------------------------*
Expand All @@ -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;
}

/**
Expand All @@ -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);
}

/**
Expand All @@ -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;
}

/**
Expand All @@ -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);
}

/**
Expand Down Expand Up @@ -279,6 +313,7 @@ public TransferPair makeTransferPair(ValueVector to) {
}

private class TransferImpl implements TransferPair {

ViewVarCharVector to;

public TransferImpl(String ref, BufferAllocator allocator) {
Expand Down
63 changes: 63 additions & 0 deletions java/vector/src/test/java/org/apache/arrow/vector/TestVarCharViewVector.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<byte[]> 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(
Expand Down

0 comments on commit 422833e

Please sign in to comment.