(arrow) branch main updated: GH-24826: [Java] Add DUV.setOffset method (#40985)

[lidavidm]

This is an automated email from the ASF dual-hosted git repository.

lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/main by this push:
     new 36ed0328b4 GH-24826: [Java] Add DUV.setOffset method (#40985)
36ed0328b4 is described below

commit 36ed0328b43ca39533e58a889f8e091d1f1ca7dc
Author: James Henderson <ja...@jarohen.dev>
AuthorDate: Thu Apr 4 05:39:37 2024 +0100

    GH-24826: [Java] Add DUV.setOffset method (#40985)
    
    ### Are these changes tested?
    
    Yes
    
    ### Are there any user-facing changes?
    
    Yes - the addition of a public DUV.setOffset method
    * GitHub Issue: #24826
    
    Authored-by: James Henderson <ja...@jarohen.dev>
    Signed-off-by: David Li <li.david...@gmail.com>
---
 .../main/codegen/templates/DenseUnionVector.java   |  8 +++++
 .../apache/arrow/vector/TestDenseUnionVector.java  | 38 ++++++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/java/vector/src/main/codegen/templates/DenseUnionVector.java 
b/java/vector/src/main/codegen/templates/DenseUnionVector.java
index c23caf3bb5..8edd167152 100644
--- a/java/vector/src/main/codegen/templates/DenseUnionVector.java
+++ b/java/vector/src/main/codegen/templates/DenseUnionVector.java
@@ -908,6 +908,14 @@ public class DenseUnionVector extends 
AbstractContainerVector implements FieldVe
     return (int) typeBuffer.capacity() / TYPE_WIDTH;
   }
 
+  public void setOffset(int index, int offset) {
+    while (index >= getOffsetBufferValueCapacity()) {
+      reallocOffsetBuffer();
+    }
+
+    offsetBuffer.setInt((long) index * OFFSET_WIDTH, offset);
+  }
+
   private long getOffsetBufferValueCapacity() {
     return offsetBuffer.capacity() / OFFSET_WIDTH;
   }
diff --git 
a/java/vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java 
b/java/vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java
index 8fd33eb5a8..2c29861561 100644
--- 
a/java/vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java
+++ 
b/java/vector/src/test/java/org/apache/arrow/vector/TestDenseUnionVector.java
@@ -99,6 +99,44 @@ public class TestDenseUnionVector {
     }
   }
 
+  @Test
+  public void testSetOffset() {
+    try (DenseUnionVector duv = DenseUnionVector.empty("foo", allocator)) {
+      duv.allocateNew();
+      byte i32TypeId = duv.registerNewTypeId(Field.notNullable("i32", 
MinorType.INT.getType()));
+      byte f64TypeId = duv.registerNewTypeId(Field.notNullable("f64", 
MinorType.FLOAT8.getType()));
+
+      IntVector i32Vector = ((IntVector) duv.addVector(i32TypeId, new 
IntVector("i32", allocator)));
+      Float8Vector f64Vector = ((Float8Vector) duv.addVector(f64TypeId, new 
Float8Vector("f64", allocator)));
+
+      i32Vector.allocateNew(3);
+      f64Vector.allocateNew(1);
+
+      duv.setTypeId(0, i32TypeId);
+      duv.setOffset(0, 0);
+      i32Vector.set(0, 42);
+
+      duv.setTypeId(1, i32TypeId);
+      duv.setOffset(1, 1);
+      i32Vector.set(1, 43);
+
+      duv.setTypeId(2, f64TypeId);
+      duv.setOffset(2, 0);
+      f64Vector.set(0, 3.14);
+
+      duv.setTypeId(3, i32TypeId);
+      duv.setOffset(3, 2);
+      i32Vector.set(2, 44);
+
+      duv.setValueCount(4);
+
+      assertEquals(42, duv.getObject(0));
+      assertEquals(43, duv.getObject(1));
+      assertEquals(3.14, duv.getObject(2));
+      assertEquals(44, duv.getObject(3));
+    }
+  }
+
   @Test
   public void testTransfer() throws Exception {
     try (DenseUnionVector srcVector = new DenseUnionVector(EMPTY_SCHEMA_PATH, 
allocator, null, null)) {