apacheGH-37863: [Java] Add typed getters for StructVector (apache#37916)

### Rationale for this change Add methods for getting child vectors as specific vector subtypes for convenience. ### What changes are included in this PR? Add child getters which let the caller specify the target vector type to StructVector. ### Are these changes tested? Yes. ### Are there any user-facing changes? No. * Closes: apache#37863 Authored-by: James Duong <[email protected]> Signed-off-by: David Li <[email protected]>
dgreiss · Feb 17, 2024 · df04687 · df04687
1 parent e4438e8
commit df04687
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/java/vector/src/main/java/org/apache/arrow/vector/complex/NonNullableStructVector.java b/java/vector/src/main/java/org/apache/arrow/vector/complex/NonNullableStructVector.java
@@ -374,6 +374,18 @@ public ValueVector getVectorById(int id) {
     return getChildByOrdinal(id);
   }
 
+  /**
+   * Gets a child vector by ordinal position and casts to the specified class.
+   */
+  public <V extends ValueVector> V getVectorById(int id, Class<V> clazz) {
+    ValueVector untyped = getVectorById(id);
+    if (clazz.isInstance(untyped)) {
+      return clazz.cast(untyped);
+    }
+    throw new ClassCastException("Id " + id + " had the wrong type. Expected " + clazz.getCanonicalName() +
+        " but was " + untyped.getClass().getCanonicalName());
+  }
+
   @Override
   public void setValueCount(int valueCount) {
     for (final ValueVector v : getChildren()) {

diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestStructVector.java
@@ -282,4 +282,12 @@ public void testAddChildVectorsWithDuplicatedFieldNamesForConflictPolicyReplace(
     }
   }
 
+  @Test
+  public void testTypedGetters() {
+    try (final StructVector s1 = StructVector.empty("s1", allocator)) {
+      s1.addOrGet("struct_child", FieldType.nullable(MinorType.INT.getType()), IntVector.class);
+      assertEquals(IntVector.class, s1.getChild("struct_child", IntVector.class).getClass());
+      assertEquals(IntVector.class, s1.getVectorById(0, IntVector.class).getClass());
+    }
+  }
 }