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 26a64ed0dd..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,13 +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); - // 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; 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 78a4e92d51..72dce5d1d5 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestVectorSchemaRoot.java @@ -25,17 +25,15 @@ 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; @@ -43,12 +41,9 @@ 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")); } } } 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");