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..26a64ed0dd 100644 --- a/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java +++ b/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; 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..78a4e92d51 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java +++ b/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,6 +34,8 @@ 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; @@ -37,9 +43,12 @@ 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()); + } + } + } }