Skip to content

Commit eaffcf4

Browse files
jhrotkojhrotko
committed
wip: remove lastExtentionType
1 parent 5bd2e06 commit eaffcf4

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
@@ -1281,12 +1281,12 @@ public void testMapVectorWithExtensionType() throws Exception {
12811281
writer.startEntry();
12821282
writer.key().bigInt().writeBigInt(0);
12831283
ExtensionWriter extensionWriter = writer.value().extension(UuidType.INSTANCE);
1284-
extensionWriter.writeExtension(u1);
1284+
extensionWriter.writeExtension(u1, UuidType.INSTANCE);
12851285
writer.endEntry();
12861286
writer.startEntry();
12871287
writer.key().bigInt().writeBigInt(1);
12881288
extensionWriter = writer.value().extension(UuidType.INSTANCE);
1289-
extensionWriter.writeExtension(u2);
1289+
extensionWriter.writeExtension(u2, UuidType.INSTANCE);
12901290
writer.endEntry();
12911291
writer.endMap();
12921292

@@ -1323,12 +1323,11 @@ public void testCopyFromForExtensionType() throws Exception {
13231323
writer.startEntry();
13241324
writer.key().bigInt().writeBigInt(0);
13251325
ExtensionWriter extensionWriter = writer.value().extension(UuidType.INSTANCE);
1326-
extensionWriter.writeExtension(u1);
1326+
extensionWriter.writeExtension(u1, UuidType.INSTANCE);
13271327
writer.endEntry();
13281328
writer.startEntry();
13291329
writer.key().bigInt().writeBigInt(1);
1330-
extensionWriter = writer.value().extension(UuidType.INSTANCE);
1331-
extensionWriter.writeExtension(u2);
1330+
extensionWriter.writeExtension(u2, UuidType.INSTANCE);
13321331
writer.endEntry();
13331332
writer.endMap();
13341333

@@ -1354,4 +1353,105 @@ public void testCopyFromForExtensionType() throws Exception {
13541353
assertEquals(u2, actualUuid);
13551354
}
13561355
}
1356+
1357+
@Test
1358+
public void testMapWithUuidKeyAndListUuidValue() throws Exception {
1359+
try (final MapVector mapVector = MapVector.empty("map", allocator, false)) {
1360+
mapVector.allocateNew();
1361+
UnionMapWriter writer = mapVector.getWriter();
1362+
1363+
// Create test UUIDs
1364+
UUID key1 = UUID.randomUUID();
1365+
UUID key2 = UUID.randomUUID();
1366+
UUID value1a = UUID.randomUUID();
1367+
UUID value1b = UUID.randomUUID();
1368+
UUID value2a = UUID.randomUUID();
1369+
UUID value2b = UUID.randomUUID();
1370+
UUID value2c = UUID.randomUUID();
1371+
1372+
// Write first map entry: {key1 -> [value1a, value1b]}
1373+
writer.setPosition(0);
1374+
writer.startMap();
1375+
1376+
writer.startEntry();
1377+
ExtensionWriter keyWriter = writer.key().extension(UuidType.INSTANCE);
1378+
keyWriter.writeExtension(key1, UuidType.INSTANCE);
1379+
ListWriter valueWriter = writer.value().list();
1380+
valueWriter.startList();
1381+
ExtensionWriter listItemWriter = valueWriter.extension(UuidType.INSTANCE);
1382+
listItemWriter.writeExtension(value1a, UuidType.INSTANCE);
1383+
listItemWriter = valueWriter.extension(UuidType.INSTANCE);
1384+
listItemWriter.writeExtension(value1b, UuidType.INSTANCE);
1385+
valueWriter.endList();
1386+
writer.endEntry();
1387+
1388+
writer.startEntry();
1389+
keyWriter = writer.key().extension(UuidType.INSTANCE);
1390+
keyWriter.writeExtension(key2, UuidType.INSTANCE);
1391+
valueWriter = writer.value().list();
1392+
valueWriter.startList();
1393+
listItemWriter = valueWriter.extension(UuidType.INSTANCE);
1394+
listItemWriter.writeExtension(value2a, UuidType.INSTANCE);
1395+
listItemWriter = valueWriter.extension(UuidType.INSTANCE);
1396+
listItemWriter.writeExtension(value2b, UuidType.INSTANCE);
1397+
listItemWriter = valueWriter.extension(UuidType.INSTANCE);
1398+
listItemWriter.writeExtension(value2c, UuidType.INSTANCE);
1399+
valueWriter.endList();
1400+
writer.endEntry();
1401+
1402+
writer.endMap();
1403+
writer.setValueCount(1);
1404+
1405+
// Read and verify the data
1406+
UnionMapReader mapReader = mapVector.getReader();
1407+
mapReader.setPosition(0);
1408+
1409+
// Read first entry
1410+
mapReader.next();
1411+
FieldReader keyReader = mapReader.key();
1412+
UuidHolder keyHolder = new UuidHolder();
1413+
keyReader.read(keyHolder);
1414+
ByteBuffer bb = ByteBuffer.wrap(keyHolder.value);
1415+
UUID actualKey = new UUID(bb.getLong(), bb.getLong());
1416+
assertEquals(key1, actualKey);
1417+
1418+
FieldReader valueReader = mapReader.value();
1419+
assertTrue(valueReader.isSet());
1420+
List<?> listValue = (List<?>) valueReader.readObject();
1421+
assertEquals(2, listValue.size());
1422+
1423+
// Verify first list item - readObject() returns UUID objects for extension types
1424+
UUID actualValue1a = (UUID) listValue.get(0);
1425+
assertEquals(value1a, actualValue1a);
1426+
1427+
// Verify second list item
1428+
UUID actualValue1b = (UUID) listValue.get(1);
1429+
assertEquals(value1b, actualValue1b);
1430+
1431+
// Read second entry
1432+
mapReader.next();
1433+
keyReader = mapReader.key();
1434+
keyReader.read(keyHolder);
1435+
bb = ByteBuffer.wrap(keyHolder.value);
1436+
actualKey = new UUID(bb.getLong(), bb.getLong());
1437+
assertEquals(key2, actualKey);
1438+
1439+
valueReader = mapReader.value();
1440+
assertTrue(valueReader.isSet());
1441+
listValue = (List<?>) valueReader.readObject();
1442+
assertEquals(3, listValue.size());
1443+
1444+
// Verify first list item - readObject() returns UUID objects for extension types
1445+
UUID actualValue2a = (UUID) listValue.get(0);
1446+
assertEquals(value2a, actualValue2a);
1447+
1448+
// Verify second list item
1449+
UUID actualValue2b = (UUID) listValue.get(1);
1450+
assertEquals(value2b, actualValue2b);
1451+
1452+
// Verify third list item
1453+
UUID actualValue2c = (UUID) listValue.get(2);
1454+
assertEquals(value2c, actualValue2c);
1455+
}
1456+
}
13571457
}

‎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)