diff --git a/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java b/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java index 4c1fbf761a..16eb5303e4 100644 --- a/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java +++ b/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java @@ -54,6 +54,7 @@ public class VectorSchemaRoot implements AutoCloseable { private int rowCount; private final List fieldVectors; private final Map fieldVectorsMap = new LinkedHashMap<>(); + private final Map fieldVectorsByName = new LinkedHashMap<>(); /** Constructs new instance containing each of the vectors. */ public VectorSchemaRoot(Iterable vectors) { @@ -116,6 +117,7 @@ public VectorSchemaRoot(Schema schema, List fieldVectors, int rowCo Field field = schema.getFields().get(i); FieldVector vector = fieldVectors.get(i); fieldVectorsMap.put(field, vector); + fieldVectorsByName.putIfAbsent(field.getName(), vector); } } @@ -173,12 +175,7 @@ public List getFieldVectors() { *

if name occurs multiple times this returns the first inserted entry for name */ public FieldVector getVector(String name) { - for (Map.Entry entry : fieldVectorsMap.entrySet()) { - if (entry.getKey().getName().equals(name)) { - return entry.getValue(); - } - } - return null; + return fieldVectorsByName.get(name); } public FieldVector getVector(Field field) { diff --git a/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java b/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java index bd3113f8bc..1736b8b6b8 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java @@ -19,6 +19,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -90,6 +91,17 @@ private void checkCount(BitVector vec1, IntVector vec2, VectorSchemaRoot vsr, in assertEquals(vsr.getRowCount(), count); } + @Test + public void testGetVectorByName() { + try (IntVector intVector = new IntVector("intVector", allocator); + BitVector bitVector = new BitVector("bitVector", allocator); + VectorSchemaRoot root = VectorSchemaRoot.of(intVector, bitVector)) { + assertEquals(intVector, root.getVector("intVector")); + assertEquals(bitVector, root.getVector("bitVector")); + assertNull(root.getVector("missingVector")); + } + } + private VectorSchemaRoot createBatch() { FieldType varCharType = new FieldType(true, new ArrowType.Utf8(), /* dictionary= */ null); FieldType listType = new FieldType(true, new ArrowType.List(), /* dictionary= */ null);