Skip to content

GH-36375: [Java] Added creating MapWriter in ComplexWriter. #36351

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lidavidm merged 4 commits into from
Jul 5, 2023
Merged

GH-36375: [Java] Added creating MapWriter in ComplexWriter. #36351

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
847a4c3
Added creating MapWriter in ComplexWriter.
xxlaykxx Jun 28, 2023
7176860
Revert accidentally removed tests
xxlaykxx Jun 29, 2023
b5f1333
fixed formating
xxlaykxx Jun 29, 2023
88cd2c2
Updated test for using ComplexWriterImpl
xxlaykxx Jul 4, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions java/vector/src/main/codegen/templates/BaseWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,7 @@ public interface ComplexWriter {
void copyReader(FieldReader reader);
StructWriter rootAsStruct();
ListWriter rootAsList();
MapWriter rootAsMap(boolean keysSorted);

void setPosition(int index);
void setValueCount(int count);
Expand Down
37 changes: 36 additions & 1 deletion java/vector/src/main/java/org/apache/arrow/vector/complex/impl/ComplexWriterImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

import org.apache.arrow.util.Preconditions;
import org.apache.arrow.vector.complex.ListVector;
import org.apache.arrow.vector.complex.MapVector;
import org.apache.arrow.vector.complex.NonNullableStructVector;
import org.apache.arrow.vector.complex.StateTool;
import org.apache.arrow.vector.complex.StructVector;
Expand All @@ -32,14 +33,15 @@ public class ComplexWriterImpl extends AbstractFieldWriter implements ComplexWri

private NullableStructWriter structRoot;
private UnionListWriter listRoot;
private UnionMapWriter mapRoot;
private final NonNullableStructVector container;

Mode mode = Mode.INIT;
private final String name;
private final boolean unionEnabled;
private final NullableStructWriterFactory nullableStructWriterFactory;

private enum Mode { INIT, STRUCT, LIST }
private enum Mode { INIT, STRUCT, LIST, MAP }

/**
* Constructs a new instance.
Expand Down Expand Up @@ -107,6 +109,9 @@ public void clear() {
case LIST:
listRoot.clear();
break;
case MAP:
mapRoot.clear();
break;
default:
break;
}
Expand All @@ -121,6 +126,9 @@ public void setValueCount(int count) {
case LIST:
listRoot.setValueCount(count);
break;
case MAP:
mapRoot.setValueCount(count);
break;
default:
break;
}
Expand All @@ -136,6 +144,9 @@ public void setPosition(int index) {
case LIST:
listRoot.setPosition(index);
break;
case MAP:
mapRoot.setPosition(index);
break;
default:
break;
}
Expand Down Expand Up @@ -223,5 +234,29 @@ public ListWriter rootAsList() {
return listRoot;
}

@Override
public MapWriter rootAsMap(boolean keysSorted) {
switch (mode) {

case INIT:
int vectorCount = container.size();
// TODO allow dictionaries in complex types
MapVector mapVector = container.addOrGetMap(name, keysSorted);
if (container.size() > vectorCount) {
mapVector.allocateNew();
}
mapRoot = new UnionMapWriter(mapVector);
mapRoot.setPosition(idx());
mode = Mode.MAP;
break;

case MAP:
break;

default:
check(Mode.INIT, Mode.STRUCT);
}

return mapRoot;
}
}
105 changes: 105 additions & 0 deletions java/vector/src/test/java/org/apache/arrow/vector/complex/writer/TestComplexWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@
import org.apache.arrow.vector.complex.impl.UnionListReader;
import org.apache.arrow.vector.complex.impl.UnionListWriter;
import org.apache.arrow.vector.complex.impl.UnionMapReader;
import org.apache.arrow.vector.complex.impl.UnionMapWriter;
import org.apache.arrow.vector.complex.impl.UnionReader;
import org.apache.arrow.vector.complex.impl.UnionWriter;
import org.apache.arrow.vector.complex.reader.BaseReader.StructReader;
Expand Down Expand Up @@ -1563,4 +1564,108 @@ public void testStructOfList() {
Assert.assertEquals(0, size);
}
}

@Test
public void testMap() {
try (MapVector mapVector = MapVector.empty("map", allocator, false)) {
mapVector.allocateNew();
UnionMapWriter mapWriter = new UnionMapWriter(mapVector);
for (int i = 0; i < COUNT; i++) {
mapWriter.startMap();
for (int j = 0; j < i % 7; j++) {
mapWriter.startEntry();
if (j % 2 == 0) {
mapWriter.key().integer().writeInt(j);
mapWriter.value().integer().writeInt(j + 1);
} else {
IntHolder keyHolder = new IntHolder();
keyHolder.value = j;
IntHolder valueHolder = new IntHolder();
valueHolder.value = j + 1;
mapWriter.key().integer().write(keyHolder);
mapWriter.value().integer().write(valueHolder);
}
mapWriter.endEntry();
}
mapWriter.endMap();
}
mapWriter.setValueCount(COUNT);
UnionMapReader mapReader = new UnionMapReader(mapVector);
for (int i = 0; i < COUNT; i++) {
mapReader.setPosition(i);
for (int j = 0; j < i % 7; j++) {
mapReader.next();
assertEquals(j, mapReader.key().readInteger().intValue());
assertEquals(j + 1, mapReader.value().readInteger().intValue());
}
}
}
}

@Test
public void testMapWithNulls() {
try (MapVector mapVector = MapVector.empty("map", allocator, false)) {
mapVector.allocateNew();
UnionMapWriter mapWriter = new UnionMapWriter(mapVector);
mapWriter.startMap();
mapWriter.startEntry();
mapWriter.key().integer().writeNull();
mapWriter.value().integer().writeInt(1);
mapWriter.endEntry();
mapWriter.endMap();
mapWriter.setValueCount(1);
UnionMapReader mapReader = new UnionMapReader(mapVector);
Assert.assertNull(mapReader.key().readInteger());
assertEquals(1, mapReader.value().readInteger().intValue());
}
}

@Test
public void testMapWithListKey() {
try (MapVector mapVector = MapVector.empty("map", allocator, false)) {
mapVector.allocateNew();
UnionMapWriter mapWriter = new UnionMapWriter(mapVector);
mapWriter.startMap();
mapWriter.startEntry();
mapWriter.key().list().startList();
for (int i = 0; i < 3; i++) {
mapWriter.key().list().integer().writeInt(i);
}
mapWriter.key().list().endList();
mapWriter.value().integer().writeInt(1);
mapWriter.endEntry();
mapWriter.endMap();
mapWriter.setValueCount(1);
UnionMapReader mapReader = new UnionMapReader(mapVector);
mapReader.key().next();
assertEquals(0, mapReader.key().reader().readInteger().intValue());
mapReader.key().next();
assertEquals(1, mapReader.key().reader().readInteger().intValue());
mapReader.key().next();
assertEquals(2, mapReader.key().reader().readInteger().intValue());
assertEquals(1, mapReader.value().readInteger().intValue());
}
}

@Test
public void testMapWithStruckKey() {
try (MapVector mapVector = MapVector.empty("map", allocator, false)) {
mapVector.allocateNew();
UnionMapWriter mapWriter = new UnionMapWriter(mapVector);
mapWriter.startMap();
mapWriter.startEntry();
mapWriter.key().struct().start();
mapWriter.key().struct().integer("value1").writeInt(1);
mapWriter.key().struct().integer("value2").writeInt(2);
mapWriter.key().struct().end();
mapWriter.value().integer().writeInt(1);
mapWriter.endEntry();
mapWriter.endMap();
mapWriter.setValueCount(1);
UnionMapReader mapReader = new UnionMapReader(mapVector);
assertEquals(1, mapReader.key().reader("value1").readInteger().intValue());
assertEquals(2, mapReader.key().reader("value2").readInteger().intValue());
assertEquals(1, mapReader.value().readInteger().intValue());
}
}
}