Skip to content

Commit 4a121b6

Browse files
nbrunoemkornfield
nbruno
authored and
emkornfield
committed
ARROW-11173: [Java] Add map type in complex reader / writer
This pull request adds support for Map types in FieldReader and FieldWriter. Initial unit tests show how to add the following nested types: - A list of maps: `List<Map<Integer, Long>>` - Nesting a map as the value of another map: `Map<Long, Map<Long, Long>>` - Nesting maps as both the key and value of another map: `Map<Map<Integer, Integer>, Map<Long, Long>>` Appreciate any feedback or suggestions for improvement. Closes apache#9151 from nbruno/ARROW-11173 Authored-by: Nick Bruno <nbruno@users.noreply.github.com> Signed-off-by: Micah Kornfield <emkornfield@gmail.com>
1 parent e8a9615 commit 4a121b6

23 files changed

Lines changed: 1003 additions & 16 deletions

‎java/vector/src/main/codegen/includes/vv_imports.ftl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ import org.apache.arrow.vector.complex.impl.*;
3636
import org.apache.arrow.vector.complex.writer.*;
3737
import org.apache.arrow.vector.complex.writer.BaseWriter.StructWriter;
3838
import org.apache.arrow.vector.complex.writer.BaseWriter.ListWriter;
39+
import org.apache.arrow.vector.complex.writer.BaseWriter.MapWriter;
3940
import org.apache.arrow.vector.util.JsonStringArrayList;
4041

4142
import java.util.Arrays;

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,10 @@ public void copyAsField(String name, ListWriter writer) {
7777
fail("CopyAsFieldList");
7878
}
7979

80+
public void copyAsField(String name, MapWriter writer) {
81+
fail("CopyAsFieldMap");
82+
}
83+
8084
<#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
8185
<#assign boxedType = (minor.boxedType!type.boxedType) />
8286
public void read(${name}Holder holder) {

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

Lines changed: 55 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,12 +56,42 @@ public void end() {
5656

5757
@Override
5858
public void startList() {
59-
throw new IllegalStateException(String.format("You tried to start when you are using a ValueWriter of type %s.", this.getClass().getSimpleName()));
59+
throw new IllegalStateException(String.format("You tried to start a list when you are using a ValueWriter of type %s.", this.getClass().getSimpleName()));
6060
}
6161

6262
@Override
6363
public void endList() {
64-
throw new IllegalStateException(String.format("You tried to end when you are using a ValueWriter of type %s.", this.getClass().getSimpleName()));
64+
throw new IllegalStateException(String.format("You tried to end a list when you are using a ValueWriter of type %s.", this.getClass().getSimpleName()));
65+
}
66+
67+
@Override
68+
public void startMap() {
69+
throw new IllegalStateException(String.format("You tried to start a map when you are using a ValueWriter of type %s.", this.getClass().getSimpleName()));
70+
}
71+
72+
@Override
73+
public void endMap() {
74+
throw new IllegalStateException(String.format("You tried to end a map when you are using a ValueWriter of type %s.", this.getClass().getSimpleName()));
75+
}
76+
77+
@Override
78+
public void startEntry() {
79+
throw new IllegalStateException(String.format("You tried to start a map entry when you are using a ValueWriter of type %s.", this.getClass().getSimpleName()));
80+
}
81+
82+
@Override
83+
public MapWriter key() {
84+
throw new IllegalStateException(String.format("You tried to start a map key when you are using a ValueWriter of type %s.", this.getClass().getSimpleName()));
85+
}
86+
87+
@Override
88+
public MapWriter value() {
89+
throw new IllegalStateException(String.format("You tried to start a map value when you are using a ValueWriter of type %s.", this.getClass().getSimpleName()));
90+
}
91+
92+
@Override
93+
public void endEntry() {
94+
throw new IllegalStateException(String.format("You tried to end a map entry when you are using a ValueWriter of type %s.", this.getClass().getSimpleName()));
6595
}
6696

6797
<#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
@@ -123,6 +153,12 @@ public ListWriter list() {
123153
return null;
124154
}
125155

156+
@Override
157+
public MapWriter map() {
158+
fail("Map");
159+
return null;
160+
}
161+
126162
@Override
127163
public StructWriter struct(String name) {
128164
fail("Struct");
@@ -135,6 +171,23 @@ public ListWriter list(String name) {
135171
return null;
136172
}
137173

174+
@Override
175+
public MapWriter map(String name) {
176+
fail("Map");
177+
return null;
178+
}
179+
180+
@Override
181+
public MapWriter map(boolean keysSorted) {
182+
fail("Map");
183+
return null;
184+
}
185+
186+
@Override
187+
public MapWriter map(String name, boolean keysSorted) {
188+
fail("Map");
189+
return null;
190+
}
138191
<#list vv.types as type><#list type.minor as minor>
139192
<#assign lowerName = minor.class?uncap_first />
140193
<#if lowerName == "int" ><#assign lowerName = "integer" /></#if>

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

Lines changed: 55 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,6 @@
1515
* limitations under the License.
1616
*/
1717

18-
import org.apache.arrow.memory.ArrowBuf;
19-
import org.apache.arrow.vector.types.Types;
20-
import org.apache.arrow.vector.types.pojo.ArrowType;
21-
import org.apache.drill.common.types.TypeProtos.MinorType;
22-
2318
<@pp.dropOutputFile />
2419
<@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/AbstractPromotableFieldWriter.java" />
2520

@@ -44,7 +39,11 @@ abstract class AbstractPromotableFieldWriter extends AbstractFieldWriter {
4439
* @param type the type of the values we want to write
4540
* @return the corresponding field writer
4641
*/
47-
abstract protected FieldWriter getWriter(MinorType type);
42+
protected FieldWriter getWriter(MinorType type) {
43+
return getWriter(type, null);
44+
}
45+
46+
abstract protected FieldWriter getWriter(MinorType type, ArrowType arrowType);
4847

4948
/**
5049
* @return the current FieldWriter
@@ -73,6 +72,37 @@ public void endList() {
7372
setPosition(idx() + 1);
7473
}
7574

75+
@Override
76+
public void startMap() {
77+
getWriter(MinorType.MAP).startMap();
78+
}
79+
80+
@Override
81+
public void endMap() {
82+
getWriter(MinorType.MAP).endMap();
83+
setPosition(idx() + 1);
84+
}
85+
86+
@Override
87+
public void startEntry() {
88+
getWriter(MinorType.MAP).startEntry();
89+
}
90+
91+
@Override
92+
public MapWriter key() {
93+
return getWriter(MinorType.MAP).key();
94+
}
95+
96+
@Override
97+
public MapWriter value() {
98+
return getWriter(MinorType.MAP).value();
99+
}
100+
101+
@Override
102+
public void endEntry() {
103+
getWriter(MinorType.MAP).endEntry();
104+
}
105+
76106
<#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
77107
<#assign fields = minor.fields!type.fields />
78108
<#if minor.class != "Decimal" && minor.class != "Decimal256">
@@ -144,6 +174,16 @@ public ListWriter list() {
144174
return getWriter(MinorType.LIST).list();
145175
}
146176
177+
@Override
178+
public MapWriter map() {
179+
return getWriter(MinorType.LIST).map();
180+
}
181+
182+
@Override
183+
public MapWriter map(boolean keysSorted) {
184+
return getWriter(MinorType.MAP, new ArrowType.Map(keysSorted));
185+
}
186+
147187
@Override
148188
public StructWriter struct(String name) {
149189
return getWriter(MinorType.STRUCT).struct(name);
@@ -154,6 +194,15 @@ public ListWriter list(String name) {
154194
return getWriter(MinorType.STRUCT).list(name);
155195
}
156196
197+
@Override
198+
public MapWriter map(String name) {
199+
return getWriter(MinorType.STRUCT).map(name);
200+
}
201+
202+
@Override
203+
public MapWriter map(String name, boolean keysSorted) {
204+
return getWriter(MinorType.STRUCT).map(name, keysSorted);
205+
}
157206
<#list vv.types as type><#list type.minor as minor>
158207
<#assign lowerName = minor.class?uncap_first />
159208
<#if lowerName == "int" ><#assign lowerName = "integer" /></#if>

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,16 @@ public interface RepeatedListReader extends ListReader{
6060
int size();
6161
void copyAsValue(ListWriter writer);
6262
}
63+
64+
public interface MapReader extends BaseReader{
65+
FieldReader reader();
66+
}
67+
68+
public interface RepeatedMapReader extends MapReader{
69+
boolean next();
70+
int size();
71+
void copyAsValue(MapWriter writer);
72+
}
6373

6474
public interface ScalarReader extends
6575
<#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> ${name}Reader, </#list></#list>

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,8 @@ public interface StructWriter extends BaseWriter {
6262
void copyReaderToField(String name, FieldReader reader);
6363
StructWriter struct(String name);
6464
ListWriter list(String name);
65+
MapWriter map(String name);
66+
MapWriter map(String name, boolean keysSorted);
6567
void start();
6668
void end();
6769
}
@@ -71,6 +73,8 @@ public interface ListWriter extends BaseWriter {
7173
void endList();
7274
StructWriter struct();
7375
ListWriter list();
76+
MapWriter map();
77+
MapWriter map(boolean keysSorted);
7478
void copyReader(FieldReader reader);
7579
7680
<#list vv.types as type><#list type.minor as minor>
@@ -82,6 +86,17 @@ public interface ListWriter extends BaseWriter {
8286
</#list></#list>
8387
}
8488

89+
public interface MapWriter extends ListWriter {
90+
void startMap();
91+
void endMap();
92+
93+
void startEntry();
94+
void endEntry();
95+
96+
MapWriter key();
97+
MapWriter value();
98+
}
99+
85100
public interface ScalarWriter extends
86101
<#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> ${name}Writer, </#list></#list> BaseWriter {}
87102

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,18 @@ private FieldReader getList(byte typeId) {
129129
return listReader;
130130
}
131131

132+
private UnionMapReader mapReader;
133+
134+
private FieldReader getMap(byte typeId) {
135+
UnionMapReader mapReader = (UnionMapReader) readers[typeId];
136+
if (mapReader == null) {
137+
mapReader = new UnionMapReader((MapVector) data.getVectorByType(typeId));
138+
mapReader.setPosition(idx());
139+
readers[typeId] = mapReader;
140+
}
141+
return mapReader;
142+
}
143+
132144
@Override
133145
public java.util.Iterator<String> iterator() {
134146
throw new UnsupportedOperationException();

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -342,6 +342,22 @@ public ListVector getList(byte typeId) {
342342
return listVector;
343343
}
344344
345+
public MapVector getMap(byte typeId) {
346+
MapVector mapVector = typeId < 0 ? null : (MapVector) childVectors[typeId];
347+
if (mapVector == null) {
348+
int vectorCount = internalStruct.size();
349+
mapVector = addOrGet(typeId, MinorType.MAP, MapVector.class);
350+
if (internalStruct.size() > vectorCount) {
351+
mapVector.allocateNew();
352+
childVectors[typeId] = mapVector;
353+
if (callBack != null) {
354+
callBack.doWork();
355+
}
356+
}
357+
}
358+
return mapVector;
359+
}
360+
345361
public byte getTypeId(int index) {
346362
return typeBuffer.getByte(index * TYPE_WIDTH);
347363
}

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

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,13 +111,29 @@ public ListWriter asList(byte typeId) {
111111
return getListWriter(typeId);
112112
}
113113

114+
private MapWriter getMapWriter(byte typeId) {
115+
MapWriter mapWriter = (MapWriter) writers[typeId];
116+
if (mapWriter == null) {
117+
mapWriter = new UnionMapWriter((MapVector) data.getVectorByType(typeId));
118+
writers[typeId] = mapWriter;
119+
}
120+
return mapWriter;
121+
}
122+
123+
public MapWriter asMap(byte typeId) {
124+
data.setTypeId(idx(), typeId);
125+
return getMapWriter(typeId);
126+
}
127+
114128
BaseWriter getWriter(byte typeId) {
115129
MinorType minorType = data.getVectorByType(typeId).getMinorType();
116130
switch (minorType) {
117131
case STRUCT:
118132
return getStructWriter(typeId);
119133
case LIST:
120134
return getListWriter(typeId);
135+
case MAP:
136+
return getMapWriter(typeId);
121137
<#list vv.types as type>
122138
<#list type.minor as minor>
123139
<#assign name = minor.class?cap_first />
@@ -195,6 +211,30 @@ public ListWriter list(String name) {
195211
return getStructWriter(typeId).list(name);
196212
}
197213
214+
@Override
215+
public MapWriter map() {
216+
byte typeId = data.getTypeId(idx());
217+
data.setTypeId(idx(), typeId);
218+
getListWriter(typeId).setPosition(data.getOffset(idx()));
219+
return getMapWriter(typeId).map();
220+
}
221+
222+
@Override
223+
public MapWriter map(String name) {
224+
byte typeId = data.getTypeId(idx());
225+
data.setTypeId(idx(), typeId);
226+
getStructWriter(typeId).setPosition(data.getOffset(idx()));
227+
return getStructWriter(typeId).map(name);
228+
}
229+
230+
@Override
231+
public MapWriter map(String name, boolean keysSorted) {
232+
byte typeId = data.getTypeId(idx());
233+
data.setTypeId(idx(), typeId);
234+
getStructWriter(typeId).setPosition(data.getOffset(idx()));
235+
return getStructWriter(typeId).map(name, keysSorted);
236+
}
237+
198238
@Override
199239
public StructWriter struct(String name) {
200240
byte typeId = data.getTypeId(idx());

0 commit comments

Comments
 (0)