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,13 +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);
-      // 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));
-      }
+      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

@@ -25,30 +25,25 @@
 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.Collections;
 import java.util.List;
 import java.util.stream.Collectors;
 import org.apache.arrow.memory.BufferAllocator;
 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
@@ -363,11 +358,10 @@ public void testSchemaSyncRefreshesFieldVectorLookup() {
       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.
+      // 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);
 
@@ -379,9 +373,8 @@ public void testSchemaSyncRefreshesFieldVectorLookup() {
       // 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.
+      // Syncing the schema leaves the row count 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());
@@ -390,39 +383,33 @@ public void testSchemaSyncRefreshesFieldVectorLookup() {
   }
 
   @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);
+  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());
-      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);
-      }
+      assertEquals(schema.getCustomMetadata(), root.getSchema().getCustomMetadata());
+    }
+  }
 
-      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());
-      }
+  @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");