GH-1473: Keep columns with equal Fields and the schema metadata in VectorSchemaRoot

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

@@ -112,10 +112,18 @@ public VectorSchemaRoot(Schema schema, List<FieldVector> fieldVectors, int rowCo
     this.schema = schema;
     this.rowCount = rowCount;
     this.fieldVectors = fieldVectors;
-    for (int i = 0; i < schema.getFields().size(); ++i) {
-      Field field = schema.getFields().get(i);
-      FieldVector vector = fieldVectors.get(i);
-      fieldVectorsMap.put(field, vector);
+    indexFieldVectors();
+  }
+
+  /**
+   * Rebuilds the Field lookup from the current schema. Fields compare structurally, so columns with
+   * equal Fields share one entry, which keeps the first of them.
+   */
+  private void indexFieldVectors() {
+    fieldVectorsMap.clear();
+    List<Field> fields = schema.getFields();
+    for (int i = 0; i < fields.size(); ++i) {
+      fieldVectorsMap.putIfAbsent(fields.get(i), fieldVectors.get(i));
     }
   }
 
@@ -311,7 +319,10 @@ public boolean syncSchema() {
     List<Field> newFields =
         this.fieldVectors.stream().map(ValueVector::getField).collect(Collectors.toList());
     if (!oldFields.equals(newFields)) {
-      this.schema = new Schema(newFields);
+      this.schema = new Schema(newFields, this.schema.getCustomMetadata());
+      // A promoted vector (e.g. a Struct/List child promoted to a union) reports a structurally
+      // different Field, so the lookup keyed by the old Fields no longer matches.
+      indexFieldVectors();
       return true;
     }
     return false;

vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java

@@ -201,9 +201,10 @@ public boolean read(VectorSchemaRoot root) throws IOException {
         nextFieldIs("columns");
         readToken(START_ARRAY);
         {
-          for (Field field : root.getSchema().getFields()) {
-            FieldVector vector = root.getVector(field);
-            readFromJsonIntoVector(field, vector);
+          // Look vectors up by position: columns can have equal Fields, e.g. duplicate names
+          List<Field> fields = root.getSchema().getFields();
+          for (int i = 0; i < fields.size(); i++) {
+            readFromJsonIntoVector(fields.get(i), root.getVector(i));
           }
         }
         readToken(END_ARRAY);
@@ -229,9 +230,9 @@ public VectorSchemaRoot read() throws IOException {
         nextFieldIs("columns");
         readToken(START_ARRAY);
         {
-          for (Field field : schema.getFields()) {
-            FieldVector vector = recordBatch.getVector(field);
-            readFromJsonIntoVector(field, vector);
+          List<Field> fields = schema.getFields();
+          for (int i = 0; i < fields.size(); i++) {
+            readFromJsonIntoVector(fields.get(i), recordBatch.getVector(i));
           }
         }
         readToken(END_ARRAY);

vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileWriter.java

@@ -192,9 +192,10 @@ private void writeBatch(VectorSchemaRoot recordBatch) throws IOException {
     {
       generator.writeObjectField("count", recordBatch.getRowCount());
       generator.writeArrayFieldStart("columns");
-      for (Field field : recordBatch.getSchema().getFields()) {
-        FieldVector vector = recordBatch.getVector(field);
-        writeFromVectorIntoJson(field, vector);
+      // Look vectors up by position: columns can have equal Fields, e.g. duplicate names
+      List<Field> fields = recordBatch.getSchema().getFields();
+      for (int i = 0; i < fields.size(); i++) {
+        writeFromVectorIntoJson(fields.get(i), recordBatch.getVector(i));
       }
       generator.writeEndArray();
     }

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

@@ -19,11 +19,15 @@
 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.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.stream.Collectors;
 import org.apache.arrow.memory.BufferAllocator;
@@ -344,4 +348,68 @@ 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);
+
+      assertTrue(root.syncSchema());
+
+      // Promotion changed the vector's Field, so the synced Field differs from the stale one.
+      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 untouched.
+      assertEquals(rowCount, root.getRowCount());
+
+      // A second sync is a no-op and the refreshed lookup stays correct.
+      assertFalse(root.syncSchema());
+      assertSame(listVector, root.getVector(syncedField));
+    }
+  }
+
+  @Test
+  public void testSchemaSyncKeepsCustomMetadata() {
+    Field child = new Field("varCharCol", FieldType.nullable(new ArrowType.Utf8()), null);
+    Field listField =
+        new Field("listCol", FieldType.nullable(new ArrowType.List()), Arrays.asList(child));
+    Schema schema = new Schema(Arrays.asList(listField), Collections.singletonMap("key", "value"));
+    try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
+      ListVector vector = (ListVector) root.getVector(0);
+      vector.allocateNew();
+      UnionListWriter writer = vector.getWriter();
+      writer.startList();
+      // A big int does not fit the varchar child, so the writer promotes it to a union.
+      writer.writeBigInt(0L);
+      writer.endList();
+
+      assertTrue(root.syncSchema());
+      assertEquals(schema.getCustomMetadata(), root.getSchema().getCustomMetadata());
+    }
+  }
+
+  @Test
+  public void testGetVectorWithDuplicateFields() {
+    Field field = new Field("x", FieldType.nullable(new ArrowType.Int(32, true)), null);
+    try (VectorSchemaRoot root =
+        VectorSchemaRoot.create(new Schema(Arrays.asList(field, field)), allocator)) {
+      // Equal Fields share one lookup entry, which resolves to the first column.
+      assertSame(root.getVector(0), root.getVector(field));
+      assertSame(root.getVector(0), root.getVector("x"));
+    }
+  }
 }

vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java

@@ -20,6 +20,7 @@
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.io.File;
 import java.io.IOException;
@@ -31,9 +32,12 @@
 import java.util.List;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.IntVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.complex.ListVector;
 import org.apache.arrow.vector.complex.StructVector;
 import org.apache.arrow.vector.complex.impl.ComplexWriterImpl;
+import org.apache.arrow.vector.complex.impl.UnionListWriter;
 import org.apache.arrow.vector.complex.writer.BaseWriter;
 import org.apache.arrow.vector.dictionary.DictionaryProvider;
 import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider;
@@ -195,6 +199,73 @@ public void testWriteReadUnionJSON() throws IOException {
     }
   }
 
+  @Test
+  public void testWriteReadAfterSchemaSync() throws IOException {
+    File file = tempDir.resolve("promoted.json").toFile();
+    Field child = new Field("varchar", FieldType.nullable(new ArrowType.Utf8()), null);
+    Field listField =
+        new Field("list", FieldType.nullable(new ArrowType.List()), Arrays.asList(child));
+    try (BufferAllocator vectorAllocator =
+            allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE);
+        VectorSchemaRoot root =
+            VectorSchemaRoot.create(new Schema(Arrays.asList(listField)), vectorAllocator)) {
+      ListVector vector = (ListVector) root.getVector(0);
+      vector.allocateNew();
+      UnionListWriter writer = vector.getWriter();
+      writer.startList();
+      writer.writeVarChar("a");
+      // A big int does not fit the varchar child, so the writer promotes it to a union
+      writer.writeBigInt(1L);
+      writer.endList();
+      root.setRowCount(1);
+      assertTrue(root.syncSchema());
+
+      writeJSON(file, root, null);
+
+      try (BufferAllocator readerAllocator =
+              allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE);
+          JsonFileReader reader = new JsonFileReader(file, readerAllocator)) {
+        assertEquals(root.getSchema(), reader.start());
+        try (VectorSchemaRoot rootFromJson = reader.read()) {
+          Validator.compareVectorSchemaRoot(root, rootFromJson);
+        }
+      }
+    }
+  }
+
+  @Test
+  public void testWriteReadDuplicateFields() throws IOException {
+    File file = tempDir.resolve("duplicate-fields.json").toFile();
+    Field field = new Field("x", FieldType.nullable(new ArrowType.Int(32, true)), null);
+    try (BufferAllocator vectorAllocator =
+            allocator.newChildAllocator("original vectors", 0, Integer.MAX_VALUE);
+        VectorSchemaRoot root =
+            VectorSchemaRoot.create(new Schema(Arrays.asList(field, field)), vectorAllocator)) {
+      ((IntVector) root.getVector(0)).setSafe(0, 1);
+      ((IntVector) root.getVector(1)).setSafe(0, 2);
+      root.setRowCount(1);
+
+      writeJSON(file, root, null);
+
+      try (BufferAllocator readerAllocator =
+              allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE);
+          JsonFileReader reader = new JsonFileReader(file, readerAllocator)) {
+        reader.start();
+        try (VectorSchemaRoot rootFromJson = reader.read()) {
+          Validator.compareVectorSchemaRoot(root, rootFromJson);
+        }
+      }
+      try (BufferAllocator readerAllocator =
+              allocator.newChildAllocator("reader", 0, Integer.MAX_VALUE);
+          JsonFileReader reader = new JsonFileReader(file, readerAllocator);
+          VectorSchemaRoot rootFromJson =
+              VectorSchemaRoot.create(reader.start(), readerAllocator)) {
+        assertTrue(reader.read(rootFromJson));
+        Validator.compareVectorSchemaRoot(root, rootFromJson);
+      }
+    }
+  }
+
   @Test
   public void testWriteReadDateTimeJSON() throws IOException {
     File file = new File("target/mytest_datetime.json");