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..270e8630ab 100644 --- a/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java +++ b/vector/src/main/java/org/apache/arrow/vector/VectorSchemaRoot.java @@ -112,10 +112,18 @@ public VectorSchemaRoot(Schema schema, List 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 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 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; diff --git a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java b/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java index 85dd501a1b..0cdc65de8a 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java +++ b/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 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 fields = schema.getFields(); + for (int i = 0; i < fields.size(); i++) { + readFromJsonIntoVector(fields.get(i), recordBatch.getVector(i)); } } readToken(END_ARRAY); diff --git a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileWriter.java b/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileWriter.java index 68700fe6af..bc16f3725f 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileWriter.java +++ b/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 fields = recordBatch.getSchema().getFields(); + for (int i = 0; i < fields.size(); i++) { + writeFromVectorIntoJson(fields.get(i), recordBatch.getVector(i)); } generator.writeEndArray(); } 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..72dce5d1d5 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java +++ b/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")); + } + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java b/vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java index 0a6fade9a1..1a5c4c682a 100644 --- a/vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java +++ b/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");