Skip to content

Commit e1c6ec7

Browse files
xxlaykxxxxlaykxx
authored
apacheGH-36375: [Java] Added creating MapWriter in ComplexWriter. (apache#36351) (#32) (#39)
Added new method rootAsMap() to ComplexWriter and implement it in ComplexWriterImpl for supporting map type. Previously in dremio side: When i trying to return map like output ComplexWrite with this code: org.apache.arrow.vector.complex.writer.BaseWriter.MapWriter mapWriter = out.rootAsList().map(false); mapWriter.startMap(); for (java.util.Map.Entry<java.lang.Integer, java.lang.Integer> element : map.entrySet()) { mapWriter.startEntry(); mapWriter.key().integer().writeInt((Integer) element.getKey()); mapWriter.value().integer().writeInt((Integer) element.getValue()); mapWriter.endEntry(); } mapWriter.endMap(); It use UnionMapWriter and generate schema like: EXPR$0: Map(false)<$data$: Union(Sparse, [1, 39])<struct: Struct<key: Int(32, true) not null, value: Int(32, true) not null> not null, map: Map(false)<entries: Struct<key: Int(32, true) not null, value: Int(32, true)> not null>>> But in OutputDerivation impl class where i should create output Complete type List<Field> children = Arrays.asList( CompleteType.INT.toField("key", false), CompleteType.INT.toField("value", false)); return new CompleteType(CompleteType.MAP.getType(), CompleteType.struct(children).toField(MapVector.DATA_VECTOR_NAME, false)); (This is only one valid case, because MapVector.initializeChildrenFromFields()) return EXPR$0::map<key::int32, value::int32> I found a place where it start using union - PromotableWriter.promoteToUnion. And in the end i have SCHEMA_CHANGE ERROR: Schema changed during projection. Schema was schema(EXPR$0::map<key::int32, value::int32>) but then changed to schema(EXPR$0::map<struct::struct<key::int32, value::int32>, map::map<key::int32, value::int32>>) * Closes: apache#36375 Authored-by: Ivan Chesnov <ivan.chesnov@dremio.com> Signed-off-by: David Li <li.davidm96@gmail.com>
1 parent 24636e3 commit e1c6ec7

File tree

3 files changed

+141
-1
lines changed

3 files changed

+141
-1
lines changed

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

+1
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@ public interface ComplexWriter {
106106
void copyReader(FieldReader reader);
107107
StructWriter rootAsStruct();
108108
ListWriter rootAsList();
109+
MapWriter rootAsMap(boolean keysSorted);
109110

110111
void setPosition(int index);
111112
void setValueCount(int count);

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

+36-1
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
import org.apache.arrow.util.Preconditions;
2121
import org.apache.arrow.vector.complex.ListVector;
22+
import org.apache.arrow.vector.complex.MapVector;
2223
import org.apache.arrow.vector.complex.NonNullableStructVector;
2324
import org.apache.arrow.vector.complex.StateTool;
2425
import org.apache.arrow.vector.complex.StructVector;
@@ -32,14 +33,15 @@ public class ComplexWriterImpl extends AbstractFieldWriter implements ComplexWri
3233

3334
private NullableStructWriter structRoot;
3435
private UnionListWriter listRoot;
36+
private UnionMapWriter mapRoot;
3537
private final NonNullableStructVector container;
3638

3739
Mode mode = Mode.INIT;
3840
private final String name;
3941
private final boolean unionEnabled;
4042
private final NullableStructWriterFactory nullableStructWriterFactory;
4143

42-
private enum Mode { INIT, STRUCT, LIST }
44+
private enum Mode { INIT, STRUCT, LIST, MAP }
4345

4446
/**
4547
* Constructs a new instance.
@@ -107,6 +109,9 @@ public void clear() {
107109
case LIST:
108110
listRoot.clear();
109111
break;
112+
case MAP:
113+
mapRoot.clear();
114+
break;
110115
default:
111116
break;
112117
}
@@ -121,6 +126,9 @@ public void setValueCount(int count) {
121126
case LIST:
122127
listRoot.setValueCount(count);
123128
break;
129+
case MAP:
130+
mapRoot.setValueCount(count);
131+
break;
124132
default:
125133
break;
126134
}
@@ -136,6 +144,9 @@ public void setPosition(int index) {
136144
case LIST:
137145
listRoot.setPosition(index);
138146
break;
147+
case MAP:
148+
mapRoot.setPosition(index);
149+
break;
139150
default:
140151
break;
141152
}
@@ -223,5 +234,29 @@ public ListWriter rootAsList() {
223234
return listRoot;
224235
}
225236

237+
@Override
238+
public MapWriter rootAsMap(boolean keysSorted) {
239+
switch (mode) {
226240

241+
case INIT:
242+
int vectorCount = container.size();
243+
// TODO allow dictionaries in complex types
244+
MapVector mapVector = container.addOrGetMap(name, keysSorted);
245+
if (container.size() > vectorCount) {
246+
mapVector.allocateNew();
247+
}
248+
mapRoot = new UnionMapWriter(mapVector);
249+
mapRoot.setPosition(idx());
250+
mode = Mode.MAP;
251+
break;
252+
253+
case MAP:
254+
break;
255+
256+
default:
257+
check(Mode.INIT, Mode.STRUCT);
258+
}
259+
260+
return mapRoot;
261+
}
227262
}

‎java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java

+104
Original file line numberDiff line numberDiff line change
@@ -1401,4 +1401,108 @@ public void testStructOfList() {
14011401
Assert.assertEquals(0, size);
14021402
}
14031403
}
1404+
1405+
@Test
1406+
public void testMap() {
1407+
try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
1408+
ComplexWriter writer = new ComplexWriterImpl("root", parent);
1409+
MapWriter mapWriter = writer.rootAsMap(false);
1410+
for (int i = 0; i < COUNT; i++) {
1411+
mapWriter.startMap();
1412+
for (int j = 0; j < i % 7; j++) {
1413+
mapWriter.startEntry();
1414+
if (j % 2 == 0) {
1415+
mapWriter.key().integer().writeInt(j);
1416+
mapWriter.value().integer().writeInt(j + 1);
1417+
} else {
1418+
IntHolder keyHolder = new IntHolder();
1419+
keyHolder.value = j;
1420+
IntHolder valueHolder = new IntHolder();
1421+
valueHolder.value = j + 1;
1422+
mapWriter.key().integer().write(keyHolder);
1423+
mapWriter.value().integer().write(valueHolder);
1424+
}
1425+
mapWriter.endEntry();
1426+
}
1427+
mapWriter.endMap();
1428+
}
1429+
writer.setValueCount(COUNT);
1430+
UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root");
1431+
for (int i = 0; i < COUNT; i++) {
1432+
mapReader.setPosition(i);
1433+
for (int j = 0; j < i % 7; j++) {
1434+
mapReader.next();
1435+
assertEquals(j, mapReader.key().readInteger().intValue());
1436+
assertEquals(j + 1, mapReader.value().readInteger().intValue());
1437+
}
1438+
}
1439+
}
1440+
}
1441+
1442+
@Test
1443+
public void testMapWithNulls() {
1444+
try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
1445+
ComplexWriter writer = new ComplexWriterImpl("root", parent);
1446+
MapWriter mapWriter = writer.rootAsMap(false);
1447+
mapWriter.startMap();
1448+
mapWriter.startEntry();
1449+
mapWriter.key().integer().writeNull();
1450+
mapWriter.value().integer().writeInt(1);
1451+
mapWriter.endEntry();
1452+
mapWriter.endMap();
1453+
writer.setValueCount(1);
1454+
UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root");
1455+
Assert.assertNull(mapReader.key().readInteger());
1456+
assertEquals(1, mapReader.value().readInteger().intValue());
1457+
}
1458+
}
1459+
1460+
@Test
1461+
public void testMapWithListKey() {
1462+
try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
1463+
ComplexWriter writer = new ComplexWriterImpl("root", parent);
1464+
MapWriter mapWriter = writer.rootAsMap(false);
1465+
mapWriter.startMap();
1466+
mapWriter.startEntry();
1467+
mapWriter.key().list().startList();
1468+
for (int i = 0; i < 3; i++) {
1469+
mapWriter.key().list().integer().writeInt(i);
1470+
}
1471+
mapWriter.key().list().endList();
1472+
mapWriter.value().integer().writeInt(1);
1473+
mapWriter.endEntry();
1474+
mapWriter.endMap();
1475+
writer.setValueCount(1);
1476+
UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root");
1477+
mapReader.key().next();
1478+
assertEquals(0, mapReader.key().reader().readInteger().intValue());
1479+
mapReader.key().next();
1480+
assertEquals(1, mapReader.key().reader().readInteger().intValue());
1481+
mapReader.key().next();
1482+
assertEquals(2, mapReader.key().reader().readInteger().intValue());
1483+
assertEquals(1, mapReader.value().readInteger().intValue());
1484+
}
1485+
}
1486+
1487+
@Test
1488+
public void testMapWithStructKey() {
1489+
try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
1490+
ComplexWriter writer = new ComplexWriterImpl("root", parent);
1491+
MapWriter mapWriter = writer.rootAsMap(false);
1492+
mapWriter.startMap();
1493+
mapWriter.startEntry();
1494+
mapWriter.key().struct().start();
1495+
mapWriter.key().struct().integer("value1").writeInt(1);
1496+
mapWriter.key().struct().integer("value2").writeInt(2);
1497+
mapWriter.key().struct().end();
1498+
mapWriter.value().integer().writeInt(1);
1499+
mapWriter.endEntry();
1500+
mapWriter.endMap();
1501+
writer.setValueCount(1);
1502+
UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root");
1503+
assertEquals(1, mapReader.key().reader("value1").readInteger().intValue());
1504+
assertEquals(2, mapReader.key().reader("value2").readInteger().intValue());
1505+
assertEquals(1, mapReader.value().readInteger().intValue());
1506+
}
1507+
}
14041508
}

0 commit comments

Comments
 (0)