Skip to content

Commit a4b7059

Browse files
tianchenemkornfield
tianchen
authored and
emkornfield
committed
ARROW-6079: [Java] Implement/test UnionFixedSizeListWriter for FixedSizeListVector
Related to [ARROW-6079](https://issues.apache.org/jira/browse/ARROW-6079). Now we have two list vectors: ListVector and FixedSizeListVector. ListVector has already implemented UnionListWriter for writing data, however, FixedSizeListVector doesn't have this yet and seems the only way for users to write data is getting inner vector and set value manually. Implement a writer for FixedSizeListVector is useful in some cases. Closes apache#4973 from tianchen92/ARROW-6079 and squashes the following commits: ad18d1c <tianchen> fix test to avoid string comparision d3338c7 <tianchen> ARROW-6079: Implement/test UnionFixedSizeListWriter for FixedSizeListVector Authored-by: tianchen <niki.lj@alibaba-inc.com> Signed-off-by: Micah Kornfield <emkornfield@gmail.com>
1 parent bc27f17 commit a4b7059

4 files changed

Lines changed: 360 additions & 17 deletions

File tree

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

Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,217 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
import io.netty.buffer.ArrowBuf;
19+
import org.apache.arrow.vector.complex.writer.DecimalWriter;
20+
import org.apache.arrow.vector.holders.DecimalHolder;
21+
22+
import java.lang.UnsupportedOperationException;
23+
import java.math.BigDecimal;
24+
25+
<@pp.dropOutputFile />
26+
<@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/UnionFixedSizeListWriter.java" />
27+
28+
29+
<#include "/@includes/license.ftl" />
30+
31+
package org.apache.arrow.vector.complex.impl;
32+
33+
<#include "/@includes/vv_imports.ftl" />
34+
35+
/*
36+
* This class is generated using freemarker and the ${.template_name} template.
37+
*/
38+
39+
@SuppressWarnings("unused")
40+
public class UnionFixedSizeListWriter extends AbstractFieldWriter {
41+
42+
protected FixedSizeListVector vector;
43+
protected PromotableWriter writer;
44+
private boolean inStruct = false;
45+
private String structName;
46+
private int lastIndex = 0;
47+
private final int listSize;
48+
49+
public UnionFixedSizeListWriter(FixedSizeListVector vector) {
50+
this(vector, NullableStructWriterFactory.getNullableStructWriterFactoryInstance());
51+
}
52+
53+
public UnionFixedSizeListWriter(FixedSizeListVector vector, NullableStructWriterFactory nullableStructWriterFactory) {
54+
this.vector = vector;
55+
this.writer = new PromotableWriter(vector.getDataVector(), vector, nullableStructWriterFactory);
56+
this.listSize = vector.getListSize();
57+
}
58+
59+
public UnionFixedSizeListWriter(FixedSizeListVector vector, AbstractFieldWriter parent) {
60+
this(vector);
61+
}
62+
63+
@Override
64+
public void allocate() {
65+
vector.allocateNew();
66+
}
67+
68+
@Override
69+
public void clear() {
70+
vector.clear();
71+
}
72+
73+
@Override
74+
public Field getField() {
75+
return null;
76+
}
77+
78+
public void setValueCount(int count) {
79+
vector.setValueCount(count);
80+
}
81+
82+
@Override
83+
public int getValueCapacity() {
84+
return vector.getValueCapacity();
85+
}
86+
87+
@Override
88+
public void close() throws Exception {
89+
90+
}
91+
92+
@Override
93+
public void setPosition(int index) {
94+
super.setPosition(index);
95+
}
96+
<#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first />
97+
<#assign fields = minor.fields!type.fields />
98+
<#assign uncappedName = name?uncap_first/>
99+
<#if uncappedName == "int" ><#assign uncappedName = "integer" /></#if>
100+
<#if !minor.typeParams?? >
101+
102+
@Override
103+
public ${name}Writer ${uncappedName}() {
104+
return this;
105+
}
106+
107+
@Override
108+
public ${name}Writer ${uncappedName}(String name) {
109+
structName = name;
110+
return writer.${uncappedName}(name);
111+
}
112+
</#if>
113+
</#list></#list>
114+
115+
@Override
116+
public DecimalWriter decimal() {
117+
return this;
118+
}
119+
120+
@Override
121+
public DecimalWriter decimal(String name, int scale, int precision) {
122+
return writer.decimal(name, scale, precision);
123+
}
124+
125+
@Override
126+
public DecimalWriter decimal(String name) {
127+
return writer.decimal(name);
128+
}
129+
130+
@Override
131+
public StructWriter struct() {
132+
inStruct = true;
133+
return this;
134+
}
135+
136+
@Override
137+
public ListWriter list() {
138+
return writer;
139+
}
140+
141+
@Override
142+
public ListWriter list(String name) {
143+
ListWriter listWriter = writer.list(name);
144+
return listWriter;
145+
}
146+
147+
@Override
148+
public StructWriter struct(String name) {
149+
StructWriter structWriter = writer.struct(name);
150+
return structWriter;
151+
}
152+
153+
@Override
154+
public void startList() {
155+
int start = vector.startNewValue(idx());
156+
writer.setPosition(start);
157+
}
158+
159+
@Override
160+
public void endList() {
161+
setPosition(idx() + 1);
162+
}
163+
164+
@Override
165+
public void start() {
166+
writer.start();
167+
}
168+
169+
@Override
170+
public void end() {
171+
writer.end();
172+
inStruct = false;
173+
}
174+
175+
@Override
176+
public void write(DecimalHolder holder) {
177+
writer.write(holder);
178+
writer.setPosition(writer.idx() + 1);
179+
}
180+
181+
public void writeDecimal(int start, ArrowBuf buffer) {
182+
writer.writeDecimal(start, buffer);
183+
writer.setPosition(writer.idx() + 1);
184+
}
185+
186+
public void writeDecimal(BigDecimal value) {
187+
writer.writeDecimal(value);
188+
writer.setPosition(writer.idx() + 1);
189+
}
190+
191+
<#list vv.types as type>
192+
<#list type.minor as minor>
193+
<#assign name = minor.class?cap_first />
194+
<#assign fields = minor.fields!type.fields />
195+
<#assign uncappedName = name?uncap_first/>
196+
<#if !minor.typeParams?? >
197+
@Override
198+
public void write${name}(<#list fields as field>${field.type} ${field.name}<#if field_has_next>, </#if></#list>) {
199+
if (writer.idx() >= (idx() + 1) * listSize) {
200+
throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
201+
}
202+
writer.write${name}(<#list fields as field>${field.name}<#if field_has_next>, </#if></#list>);
203+
writer.setPosition(writer.idx() + 1);
204+
}
205+
206+
public void write(${name}Holder holder) {
207+
if (writer.idx() >= (idx() + 1) * listSize) {
208+
throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize));
209+
}
210+
writer.write${name}(<#list fields as field>holder.${field.name}<#if field_has_next>, </#if></#list>);
211+
writer.setPosition(writer.idx() + 1);
212+
}
213+
214+
</#if>
215+
</#list>
216+
</#list>
217+
}

‎java/vector/src/main/java/org/apache/arrow/vector/complex/FixedSizeListVector.java‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
import org.apache.arrow.vector.ValueVector;
4141
import org.apache.arrow.vector.ZeroVector;
4242
import org.apache.arrow.vector.complex.impl.UnionFixedSizeListReader;
43+
import org.apache.arrow.vector.complex.impl.UnionFixedSizeListWriter;
4344
import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
4445
import org.apache.arrow.vector.types.Types.MinorType;
4546
import org.apache.arrow.vector.types.pojo.ArrowType;
@@ -267,6 +268,24 @@ public FieldVector getDataVector() {
267268
return vector;
268269
}
269270

271+
/**
272+
* Start a new value in the list vector.
273+
*
274+
* @param index index of the value to start
275+
*/
276+
public int startNewValue(int index) {
277+
while (index >= getValidityBufferValueCapacity()) {
278+
reallocValidityBuffer();
279+
}
280+
281+
BitVectorHelper.setValidityBitToOne(validityBuffer, index);
282+
return index * listSize;
283+
}
284+
285+
public UnionFixedSizeListWriter getWriter() {
286+
return new UnionFixedSizeListWriter(this);
287+
}
288+
270289
@Override
271290
public void setInitialCapacity(int numRecords) {
272291
validityAllocationSizeInBytes = getValidityBufferSizeFromCount(numRecords);

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

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import org.apache.arrow.vector.ValueVector;
2424
import org.apache.arrow.vector.ZeroVector;
2525
import org.apache.arrow.vector.complex.AbstractStructVector;
26+
import org.apache.arrow.vector.complex.FixedSizeListVector;
2627
import org.apache.arrow.vector.complex.ListVector;
2728
import org.apache.arrow.vector.complex.StructVector;
2829
import org.apache.arrow.vector.complex.UnionVector;
@@ -48,6 +49,7 @@ public class PromotableWriter extends AbstractPromotableFieldWriter {
4849

4950
private final AbstractStructVector parentContainer;
5051
private final ListVector listVector;
52+
private final FixedSizeListVector fixedListVector;
5153
private final NullableStructWriterFactory nullableStructWriterFactory;
5254
private int position;
5355
private static final int MAX_DECIMAL_PRECISION = 38;
@@ -86,6 +88,7 @@ public PromotableWriter(
8688
NullableStructWriterFactory nullableStructWriterFactory) {
8789
this.parentContainer = parentContainer;
8890
this.listVector = null;
91+
this.fixedListVector = null;
8992
this.nullableStructWriterFactory = nullableStructWriterFactory;
9093
init(v);
9194
}
@@ -100,6 +103,16 @@ public PromotableWriter(ValueVector v, ListVector listVector) {
100103
this(v, listVector, NullableStructWriterFactory.getNullableStructWriterFactoryInstance());
101104
}
102105

106+
/**
107+
* Constructs a new instance.
108+
*
109+
* @param v The vector to initialize the writer with.
110+
* @param fixedListVector The vector that serves as a parent of v.
111+
*/
112+
public PromotableWriter(ValueVector v, FixedSizeListVector fixedListVector) {
113+
this(v, fixedListVector, NullableStructWriterFactory.getNullableStructWriterFactoryInstance());
114+
}
115+
103116
/**
104117
* Constructs a new instance.
105118
*
@@ -113,6 +126,25 @@ public PromotableWriter(
113126
NullableStructWriterFactory nullableStructWriterFactory) {
114127
this.listVector = listVector;
115128
this.parentContainer = null;
129+
this.fixedListVector = null;
130+
this.nullableStructWriterFactory = nullableStructWriterFactory;
131+
init(v);
132+
}
133+
134+
/**
135+
* Constructs a new instance.
136+
*
137+
* @param v The vector to initialize the writer with.
138+
* @param fixedListVector The vector that serves as a parent of v.
139+
* @param nullableStructWriterFactory The factory to create the delegate writer.
140+
*/
141+
public PromotableWriter(
142+
ValueVector v,
143+
FixedSizeListVector fixedListVector,
144+
NullableStructWriterFactory nullableStructWriterFactory) {
145+
this.fixedListVector = fixedListVector;
146+
this.parentContainer = null;
147+
this.listVector = null;
116148
this.nullableStructWriterFactory = nullableStructWriterFactory;
117149
init(v);
118150
}
@@ -189,7 +221,8 @@ protected FieldWriter getWriter(MinorType type, ArrowType arrowType) {
189221
arrowType = type.getType();
190222
}
191223
FieldType fieldType = new FieldType(addVectorAsNullable, arrowType, null, null);
192-
ValueVector v = listVector.addOrGetVector(fieldType).getVector();
224+
ValueVector v = listVector != null ? listVector.addOrGetVector(fieldType).getVector() :
225+
fixedListVector.addOrGetVector(fieldType).getVector();
193226
v.allocateNew();
194227
setWriter(v, arrowType);
195228
writer.setPosition(position);
@@ -219,6 +252,8 @@ private FieldWriter promoteToUnion() {
219252
unionVector.allocateNew();
220253
} else if (listVector != null) {
221254
unionVector = listVector.promoteToUnion();
255+
} else if (fixedListVector != null) {
256+
unionVector = fixedListVector.promoteToUnion();
222257
}
223258
unionVector.addVector((FieldVector) tp.getTo());
224259
writer = new UnionWriter(unionVector, nullableStructWriterFactory);

0 commit comments

Comments
 (0)