GH-390: Refresh vector lookup when synchronizing the schema

vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java

@@ -312,6 +312,12 @@ public boolean syncSchema() {
         this.fieldVectors.stream().map(ValueVector::getField).collect(Collectors.toList());
     if (!oldFields.equals(newFields)) {
       this.schema = new Schema(newFields);
+      // Field equality is structural, and promotion (e.g. a promoted Struct/List child)
+      // produces new Field instances, so the existing keys would no longer match.
+      fieldVectorsMap.clear();
+      for (int i = 0; i < newFields.size(); i++) {
+        fieldVectorsMap.put(newFields.get(i), fieldVectors.get(i));
+      }
       return true;
     }
     return false;

vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java

@@ -19,9 +19,13 @@
 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.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -30,16 +34,21 @@
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.complex.ListVector;
 import org.apache.arrow.vector.complex.impl.UnionListWriter;
+import org.apache.arrow.vector.ipc.JsonFileReader;
+import org.apache.arrow.vector.ipc.JsonFileWriter;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.FieldType;
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
 
 public class TestVectorSchemaRoot {
 
+  @TempDir public Path tempDir;
+
   private BufferAllocator allocator;
 
   @BeforeEach
@@ -344,4 +353,76 @@ public void testSchemaSync() {
       assertFalse(schemaRoot.syncSchema());
     }
   }
+
+  @Test
+  public void testSchemaSyncRefreshesFieldVectorLookup() {
+    try (VectorSchemaRoot root = createBatch()) {
+      final FieldVector listVector = root.getVector("listCol");
+      // createBatch() writes two list rows but leaves the root row count unset.
+      root.setRowCount(listVector.getValueCount());
+      final int rowCount = root.getRowCount();
+      assertEquals(2, rowCount);
+      final Field staleField = root.getSchema().getFields().get(0);
+      final Object oldValue = listVector.getObject(0);
+
+      assertTrue(root.syncSchema());
+
+      // Promotion produced a new Field instance for the same, unchanged vector.
+      final Field syncedField = root.getSchema().getFields().get(0);
+      assertNotEquals(staleField, syncedField);
+
+      // Looking up by the synced Field must resolve to the very same vector instance.
+      assertNotNull(root.getVector(syncedField));
+      assertSame(listVector, root.getVector(syncedField));
+      assertSame(listVector, root.getVector("listCol"));
+
+      // The stale structural key is gone because the Field itself changed.
+      assertNull(root.getVector(staleField));
+
+      // Syncing the schema leaves the row count and the vector data untouched.
+      assertEquals(rowCount, root.getRowCount());
+      assertEquals(oldValue, listVector.getObject(0));
+
+      // A second sync is a no-op and the refreshed lookup stays correct.
+      assertFalse(root.syncSchema());
+      assertSame(listVector, root.getVector(syncedField));
+    }
+  }
+
+  @Test
+  public void testJsonWriteAndReadAfterSchemaSync() throws Exception {
+    final Path jsonFile = tempDir.resolve("promoted-schema.json");
+    try (VectorSchemaRoot root = createBatch()) {
+      final FieldVector expectedVector = root.getVector("listCol");
+      // createBatch() writes two list rows but leaves the root row count unset.
+      root.setRowCount(expectedVector.getValueCount());
+      final int rowCount = root.getRowCount();
+      assertEquals(2, rowCount);
+      assertTrue(root.syncSchema());
+      final Schema syncedSchema = root.getSchema();
+
+      // Writing resolves every vector through VectorSchemaRoot#getVector(Field), so a
+      // stale Field-keyed lookup makes this write fail.
+      try (JsonFileWriter writer =
+          new JsonFileWriter(jsonFile.toFile(), JsonFileWriter.config().pretty(true))) {
+        writer.start(syncedSchema, null);
+        writer.write(root);
+      }
+
+      try (BufferAllocator readerAllocator =
+              allocator.newChildAllocator("json-reader", 0, Long.MAX_VALUE);
+          JsonFileReader reader = new JsonFileReader(jsonFile.toFile(), readerAllocator)) {
+        assertEquals(syncedSchema, reader.start());
+        try (VectorSchemaRoot readRoot = reader.read()) {
+          assertNotNull(readRoot);
+          assertEquals(rowCount, readRoot.getRowCount());
+          final FieldVector actualVector = readRoot.getVector("listCol");
+          // The JSON schema round-trips exactly, so the nested rows compare structurally.
+          assertEquals(expectedVector.getObject(0), actualVector.getObject(0));
+          assertEquals(expectedVector.getObject(1), actualVector.getObject(1));
+        }
+        assertNull(reader.read());
+      }
+    }
+  }
 }