Skip to content

Commit 1a9f205

Browse files
jhrotkojhrotko
committed
wip: remove lastExtentionType
1 parent a7b0dae commit 1a9f205

File tree

11 files changed

+143
-33
lines changed

11 files changed

+143
-33
lines changed

‎vector/src/main/codegen/templates/AbstractFieldWriter.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -109,11 +109,15 @@ public void endEntry() {
109109

110110
@Override
111111
public void write(ExtensionHolder var1) {
112-
this.fail("ExtensionType");
112+
this.fail("Cannot write ExtensionHolder");
113113
}
114114
@Override
115115
public void writeExtension(Object var1) {
116-
this.fail("ExtensionType");
116+
this.fail("Cannot write extension object");
117+
}
118+
@Override
119+
public void writeExtension(Object var1, ArrowType type) {
120+
this.fail("Cannot write extension with type " + type);
117121
}
118122

119123
<#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />

‎vector/src/main/codegen/templates/BaseWriter.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,14 @@ public interface ExtensionWriter extends BaseWriter {
123123
* @param value the extension type value to write
124124
*/
125125
void writeExtension(Object value);
126+
127+
/**
128+
* Writes the given extension type value.
129+
*
130+
* @param value the extension type value to write
131+
* @param type of the extension
132+
*/
133+
void writeExtension(Object value, ArrowType type);
126134
}
127135

128136
public interface ScalarWriter extends

‎vector/src/main/codegen/templates/ComplexCopier.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,7 @@ public static void copy(FieldReader reader, FieldWriter writer) {
111111
if (reader.isSet()) {
112112
Object value = reader.readObject();
113113
if (value != null) {
114-
writer.writeExtension(value);
114+
writer.writeExtension(value, reader.getField().getType());
115115
}
116116
} else {
117117
writer.writeNull();

‎vector/src/main/codegen/templates/PromotableWriter.java‎

Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -325,9 +325,6 @@ protected boolean requiresArrowType(MinorType type) {
325325

326326
@Override
327327
protected FieldWriter getWriter(MinorType type, ArrowType arrowType) {
328-
if(type == MinorType.EXTENSIONTYPE) {
329-
lastExtensionType = arrowType;
330-
}
331328
if (state == State.UNION) {
332329
if (requiresArrowType(type)) {
333330
((UnionWriter) writer).getWriter(type, arrowType);
@@ -543,20 +540,9 @@ public void writeLargeVarChar(String value) {
543540
getWriter(MinorType.LARGEVARCHAR).writeLargeVarChar(value);
544541
}
545542

546-
protected ArrowType lastExtensionType;
547-
548543
@Override
549-
public void writeExtension(Object value) {
550-
FieldWriter writer = getWriter(MinorType.EXTENSIONTYPE, lastExtensionType);
551-
if(writer instanceof UnionWriter) {
552-
((UnionWriter) writer).writeExtension(value, lastExtensionType);
553-
} else {
554-
writer.writeExtension(value);
555-
}
556-
}
557-
558544
public void writeExtension(Object value, ArrowType arrowType) {
559-
getWriter(MinorType.EXTENSIONTYPE, arrowType).writeExtension(value);
545+
getWriter(MinorType.EXTENSIONTYPE, arrowType).writeExtension(value, arrowType);
560546
}
561547

562548
@Override

‎vector/src/main/codegen/templates/UnionListWriter.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -341,6 +341,11 @@ public void writeExtension(Object value) {
341341
writer.setPosition(writer.idx() + 1);
342342
}
343343
344+
@Override
345+
public void writeExtension(Object value, ArrowType type) {
346+
writeExtension(value);
347+
}
348+
344349
public void write(ExtensionHolder var1) {
345350
writer.write(var1);
346351
writer.setPosition(writer.idx() + 1);

‎vector/src/main/java/org/apache/arrow/vector/complex/impl/UnionExtensionWriter.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ public void writeExtension(Object var1) {
5959
this.writer.writeExtension(var1);
6060
}
6161

62+
@Override
6263
public void write(ExtensionHolder holder) {
6364
this.writer.write(holder);
6465
}

‎vector/src/test/java/org/apache/arrow/vector/TestMapVector.java‎

Lines changed: 105 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1284,12 +1284,12 @@ public void testMapVectorWithExtensionType() throws Exception {
12841284
writer.startEntry();
12851285
writer.key().bigInt().writeBigInt(0);
12861286
ExtensionWriter extensionWriter = writer.value().extension(UuidType.INSTANCE);
1287-
extensionWriter.writeExtension(u1);
1287+
extensionWriter.writeExtension(u1, UuidType.INSTANCE);
12881288
writer.endEntry();
12891289
writer.startEntry();
12901290
writer.key().bigInt().writeBigInt(1);
12911291
extensionWriter = writer.value().extension(UuidType.INSTANCE);
1292-
extensionWriter.writeExtension(u2);
1292+
extensionWriter.writeExtension(u2, UuidType.INSTANCE);
12931293
writer.endEntry();
12941294
writer.endMap();
12951295

@@ -1326,12 +1326,11 @@ public void testCopyFromForExtensionType() throws Exception {
13261326
writer.startEntry();
13271327
writer.key().bigInt().writeBigInt(0);
13281328
ExtensionWriter extensionWriter = writer.value().extension(UuidType.INSTANCE);
1329-
extensionWriter.writeExtension(u1);
1329+
extensionWriter.writeExtension(u1, UuidType.INSTANCE);
13301330
writer.endEntry();
13311331
writer.startEntry();
13321332
writer.key().bigInt().writeBigInt(1);
1333-
extensionWriter = writer.value().extension(UuidType.INSTANCE);
1334-
extensionWriter.writeExtension(u2);
1333+
extensionWriter.writeExtension(u2, UuidType.INSTANCE);
13351334
writer.endEntry();
13361335
writer.endMap();
13371336

@@ -1575,4 +1574,105 @@ public void testFixedSizeBinaryFirstInitialization() {
15751574
assertArrayEquals(new byte[] {32, 21}, (byte[]) resultStruct.get(MapVector.VALUE_NAME));
15761575
}
15771576
}
1577+
1578+
@Test
1579+
public void testMapWithUuidKeyAndListUuidValue() throws Exception {
1580+
try (final MapVector mapVector = MapVector.empty("map", allocator, false)) {
1581+
mapVector.allocateNew();
1582+
UnionMapWriter writer = mapVector.getWriter();
1583+
1584+
// Create test UUIDs
1585+
UUID key1 = UUID.randomUUID();
1586+
UUID key2 = UUID.randomUUID();
1587+
UUID value1a = UUID.randomUUID();
1588+
UUID value1b = UUID.randomUUID();
1589+
UUID value2a = UUID.randomUUID();
1590+
UUID value2b = UUID.randomUUID();
1591+
UUID value2c = UUID.randomUUID();
1592+
1593+
// Write first map entry: {key1 -> [value1a, value1b]}
1594+
writer.setPosition(0);
1595+
writer.startMap();
1596+
1597+
writer.startEntry();
1598+
ExtensionWriter keyWriter = writer.key().extension(UuidType.INSTANCE);
1599+
keyWriter.writeExtension(key1, UuidType.INSTANCE);
1600+
ListWriter valueWriter = writer.value().list();
1601+
valueWriter.startList();
1602+
ExtensionWriter listItemWriter = valueWriter.extension(UuidType.INSTANCE);
1603+
listItemWriter.writeExtension(value1a, UuidType.INSTANCE);
1604+
listItemWriter = valueWriter.extension(UuidType.INSTANCE);
1605+
listItemWriter.writeExtension(value1b, UuidType.INSTANCE);
1606+
valueWriter.endList();
1607+
writer.endEntry();
1608+
1609+
writer.startEntry();
1610+
keyWriter = writer.key().extension(UuidType.INSTANCE);
1611+
keyWriter.writeExtension(key2, UuidType.INSTANCE);
1612+
valueWriter = writer.value().list();
1613+
valueWriter.startList();
1614+
listItemWriter = valueWriter.extension(UuidType.INSTANCE);
1615+
listItemWriter.writeExtension(value2a, UuidType.INSTANCE);
1616+
listItemWriter = valueWriter.extension(UuidType.INSTANCE);
1617+
listItemWriter.writeExtension(value2b, UuidType.INSTANCE);
1618+
listItemWriter = valueWriter.extension(UuidType.INSTANCE);
1619+
listItemWriter.writeExtension(value2c, UuidType.INSTANCE);
1620+
valueWriter.endList();
1621+
writer.endEntry();
1622+
1623+
writer.endMap();
1624+
writer.setValueCount(1);
1625+
1626+
// Read and verify the data
1627+
UnionMapReader mapReader = mapVector.getReader();
1628+
mapReader.setPosition(0);
1629+
1630+
// Read first entry
1631+
mapReader.next();
1632+
FieldReader keyReader = mapReader.key();
1633+
UuidHolder keyHolder = new UuidHolder();
1634+
keyReader.read(keyHolder);
1635+
ByteBuffer bb = ByteBuffer.wrap(keyHolder.value);
1636+
UUID actualKey = new UUID(bb.getLong(), bb.getLong());
1637+
assertEquals(key1, actualKey);
1638+
1639+
FieldReader valueReader = mapReader.value();
1640+
assertTrue(valueReader.isSet());
1641+
List<?> listValue = (List<?>) valueReader.readObject();
1642+
assertEquals(2, listValue.size());
1643+
1644+
// Verify first list item - readObject() returns UUID objects for extension types
1645+
UUID actualValue1a = (UUID) listValue.get(0);
1646+
assertEquals(value1a, actualValue1a);
1647+
1648+
// Verify second list item
1649+
UUID actualValue1b = (UUID) listValue.get(1);
1650+
assertEquals(value1b, actualValue1b);
1651+
1652+
// Read second entry
1653+
mapReader.next();
1654+
keyReader = mapReader.key();
1655+
keyReader.read(keyHolder);
1656+
bb = ByteBuffer.wrap(keyHolder.value);
1657+
actualKey = new UUID(bb.getLong(), bb.getLong());
1658+
assertEquals(key2, actualKey);
1659+
1660+
valueReader = mapReader.value();
1661+
assertTrue(valueReader.isSet());
1662+
listValue = (List<?>) valueReader.readObject();
1663+
assertEquals(3, listValue.size());
1664+
1665+
// Verify first list item - readObject() returns UUID objects for extension types
1666+
UUID actualValue2a = (UUID) listValue.get(0);
1667+
assertEquals(value2a, actualValue2a);
1668+
1669+
// Verify second list item
1670+
UUID actualValue2b = (UUID) listValue.get(1);
1671+
assertEquals(value2b, actualValue2b);
1672+
1673+
// Verify third list item
1674+
UUID actualValue2c = (UUID) listValue.get(2);
1675+
assertEquals(value2c, actualValue2c);
1676+
}
1677+
}
15781678
}

‎vector/src/test/java/org/apache/arrow/vector/complex/impl/TestComplexCopier.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -896,9 +896,9 @@ public void testCopyMapVectorWithExtensionType() {
896896
mapWriter.startMap();
897897
mapWriter.startEntry();
898898
ExtensionWriter extensionKeyWriter = mapWriter.key().extension(UuidType.INSTANCE);
899-
extensionKeyWriter.writeExtension(UUID.randomUUID());
899+
extensionKeyWriter.writeExtension(UUID.randomUUID(), UuidType.INSTANCE);
900900
ExtensionWriter extensionValueWriter = mapWriter.value().extension(UuidType.INSTANCE);
901-
extensionValueWriter.writeExtension(UUID.randomUUID());
901+
extensionValueWriter.writeExtension(UUID.randomUUID(), UuidType.INSTANCE);
902902
mapWriter.endEntry();
903903
mapWriter.endMap();
904904
}
@@ -932,9 +932,9 @@ public void testCopyStructVectorWithExtensionType() {
932932
structWriter.setPosition(i);
933933
structWriter.start();
934934
ExtensionWriter extensionWriter1 = structWriter.extension("uuid1", UuidType.INSTANCE);
935-
extensionWriter1.writeExtension(UUID.randomUUID());
935+
extensionWriter1.writeExtension(UUID.randomUUID(), UuidType.INSTANCE);
936936
ExtensionWriter extensionWriter2 = structWriter.extension("uuid2", UuidType.INSTANCE);
937-
extensionWriter2.writeExtension(UUID.randomUUID());
937+
extensionWriter2.writeExtension(UUID.randomUUID(), UuidType.INSTANCE);
938938
structWriter.end();
939939
}
940940

‎vector/src/test/java/org/apache/arrow/vector/complex/impl/TestPromotableWriter.java‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ public void testPromoteToUnion() throws Exception {
133133

134134
writer.setPosition(9);
135135
UUID uuid = UUID.randomUUID();
136-
writer.extension("A", UuidType.INSTANCE).writeExtension(uuid);
136+
writer.extension("A", UuidType.INSTANCE).writeExtension(uuid, UuidType.INSTANCE);
137137
writer.end();
138138

139139
writer.setPosition(10);
@@ -814,9 +814,9 @@ public void testExtensionType() throws Exception {
814814
container.setValueCount(1);
815815

816816
writer.setPosition(0);
817-
writer.writeExtension(u1);
817+
writer.writeExtension(u1, UuidType.INSTANCE);
818818
writer.setPosition(1);
819-
writer.writeExtension(u2);
819+
writer.writeExtension(u2, UuidType.INSTANCE);
820820

821821
container.setValueCount(2);
822822

@@ -839,9 +839,9 @@ public void testExtensionTypeForList() throws Exception {
839839
container.setValueCount(1);
840840

841841
writer.setPosition(0);
842-
writer.writeExtension(u1);
842+
writer.writeExtension(u1, UuidType.INSTANCE);
843843
writer.setPosition(1);
844-
writer.writeExtension(u2);
844+
writer.writeExtension(u2, UuidType.INSTANCE);
845845

846846
container.setValueCount(2);
847847

‎vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import org.apache.arrow.vector.UuidVector;
2222
import org.apache.arrow.vector.holder.UuidHolder;
2323
import org.apache.arrow.vector.holders.ExtensionHolder;
24+
import org.apache.arrow.vector.types.pojo.ArrowType;
2425

2526
public class UuidWriterImpl extends AbstractExtensionTypeWriter<UuidVector> {
2627

@@ -38,6 +39,11 @@ public void writeExtension(Object value) {
3839
vector.setValueCount(getPosition() + 1);
3940
}
4041

42+
@Override
43+
public void writeExtension(Object value, ArrowType type) {
44+
writeExtension(value);
45+
}
46+
4147
@Override
4248
public void write(ExtensionHolder holder) {
4349
UuidHolder uuidHolder = (UuidHolder) holder;

0 commit comments

Comments
 (0)