apacheGH-40999: [Java] Fix AIOOBE trying to splitAndTransfer DUV with…

…in nullable struct (apache#41000) We add a `typeId >= 0` guard to `DUV.TransferImpl.splitAndTransfer` to fix apache#40999. ### Are these changes tested? Yes ### Are there any user-facing changes? No * GitHub Issue: apache#40999 Authored-by: James Henderson <[email protected]> Signed-off-by: David Li <[email protected]>
jorisvandenbossche · Apr 7, 2024 · 84f6ede · 84f6ede
1 parent 805327e
commit 84f6ede
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 4 deletions.
diff --git a/java/vector/src/main/codegen/templates/DenseUnionVector.java b/java/vector/src/main/codegen/templates/DenseUnionVector.java
@@ -676,10 +676,12 @@ public void splitAndTransfer(int startIndex, int length) {
 
       for (int i = startIndex; i < startIndex + length; i++) {
         byte typeId = typeBuffer.getByte(i);
-        to.offsetBuffer.setInt((long) (i - startIndex) * OFFSET_WIDTH, typeCounts[typeId]);
-        typeCounts[typeId] += 1;
-        if (typeStarts[typeId] == -1) {
-          typeStarts[typeId] = offsetBuffer.getInt((long) i * OFFSET_WIDTH);
+        if (typeId >= 0) {
+          to.offsetBuffer.setInt((long) (i - startIndex) * OFFSET_WIDTH, typeCounts[typeId]);
+          typeCounts[typeId] += 1;
+          if (typeStarts[typeId] == -1) {
+            typeStarts[typeId] = offsetBuffer.getInt((long) i * 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
@@ -363,6 +363,34 @@ public void testSplitAndTransferWithMixedVectors() throws Exception {
     }
   }
 
+  @Test
+  public void testSplitAndTransferDuvInStruct() {
+    try (StructVector struct = StructVector.empty("struct", allocator)) {
+      DenseUnionVector duv = struct.addOrGet("duv",
+          FieldType.notNullable(MinorType.DENSEUNION.getType()),
+          DenseUnionVector.class);
+      byte i32TypeId = duv.registerNewTypeId(Field.notNullable("i32", MinorType.INT.getType()));
+      duv.addVector(i32TypeId, new IntVector("i32", allocator));
+
+      struct.setIndexDefined(0);
+      duv.setTypeId(0, i32TypeId);
+      duv.setSafe(0, newIntHolder(42));
+
+      struct.setNull(1);
+      struct.setValueCount(2);
+
+      try (StructVector dest = StructVector.empty("dest", allocator)) {
+        TransferPair pair = struct.makeTransferPair(dest);
+        pair.splitAndTransfer(0, 2);
+
+        assertEquals(2, dest.getValueCount());
+        assertFalse(dest.isNull(0));
+        assertEquals(42, dest.getObject(0).get("duv"));
+        assertTrue(dest.isNull(1));
+      }
+    }
+  }
+
   @Test
   public void testGetFieldTypeInfo() throws Exception {
     Map<String, String> metadata = new HashMap<>();