Skip to content
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

Jump to bottom

DX-66740: GH-36375 [Java] Added creating MapWriter in ComplexWriter. #32

Merged
merged 1 commit into from
Jul 7, 2023
Merged

DX-66740: GH-36375 [Java] Added creating MapWriter in ComplexWriter. #32

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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;
}
}
104 changes: 104 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 @@ -1401,4 +1401,108 @@ public void testStructOfList() {
Assert.assertEquals(0, size);
}
}

@Test
public void testMap() {
try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
ComplexWriter writer = new ComplexWriterImpl("root", parent);
MapWriter mapWriter = writer.rootAsMap(false);
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();
}
writer.setValueCount(COUNT);
UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root");
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 (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
ComplexWriter writer = new ComplexWriterImpl("root", parent);
MapWriter mapWriter = writer.rootAsMap(false);
mapWriter.startMap();
mapWriter.startEntry();
mapWriter.key().integer().writeNull();
mapWriter.value().integer().writeInt(1);
mapWriter.endEntry();
mapWriter.endMap();
writer.setValueCount(1);
UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root");
Assert.assertNull(mapReader.key().readInteger());
assertEquals(1, mapReader.value().readInteger().intValue());
}
}

@Test
public void testMapWithListKey() {
try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
ComplexWriter writer = new ComplexWriterImpl("root", parent);
MapWriter mapWriter = writer.rootAsMap(false);
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();
writer.setValueCount(1);
UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root");
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 testMapWithStructKey() {
try (NonNullableStructVector parent = NonNullableStructVector.empty("parent", allocator)) {
ComplexWriter writer = new ComplexWriterImpl("root", parent);
MapWriter mapWriter = writer.rootAsMap(false);
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();
writer.setValueCount(1);
UnionMapReader mapReader = (UnionMapReader) new SingleStructReaderImpl(parent).reader("root");
assertEquals(1, mapReader.key().reader("value1").readInteger().intValue());
assertEquals(2, mapReader.key().reader("value2").readInteger().intValue());
assertEquals(1, mapReader.value().readInteger().intValue());
}
}
}