Skip to content

Commit d3338c7

Browse files
author
tianchen
committed
ARROW-6079: [Java] Implement/test UnionFixedSizeListWriter for FixedSizeListVector
1 parent 803dd89 commit d3338c7

4 files changed

Lines changed: 338 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)