Skip to content

Commit 660583d

Browse files
Add producer for arrays type
1 parent db4c933 commit 660583d

File tree

2 files changed

+74
-0
lines changed

2 files changed

+74
-0
lines changed

adapter/avro/src/main/java/org/apache/arrow/adapter/avro/ArrowToAvroUtils.java

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
package org.apache.arrow.adapter.avro;
1919

20+
import org.apache.arrow.adapter.avro.producers.AvroArraysProducer;
2021
import org.apache.arrow.adapter.avro.producers.AvroBooleanProducer;
2122
import org.apache.arrow.adapter.avro.producers.AvroBytesProducer;
2223
import org.apache.arrow.adapter.avro.producers.AvroDoubleProducer;
@@ -54,6 +55,7 @@
5455
import org.apache.arrow.vector.TimeStampMilliVector;
5556
import org.apache.arrow.vector.VarBinaryVector;
5657
import org.apache.arrow.vector.VarCharVector;
58+
import org.apache.arrow.vector.complex.ListVector;
5759
import org.apache.arrow.vector.complex.StructVector;
5860
import org.apache.arrow.vector.types.Types;
5961

@@ -147,6 +149,13 @@ private static BaseAvroProducer<?> createProducer(FieldVector vector, boolean nu
147149
}
148150
return new AvroStructProducer(structVector, childProducers);
149151

152+
case LIST:
153+
154+
ListVector listVector = (ListVector) vector;
155+
FieldVector itemVector = listVector.getDataVector();
156+
Producer<?> itemProducer = createProducer(itemVector, itemVector.getField().isNullable());
157+
return new AvroArraysProducer(listVector, itemProducer);
158+
150159
// Not all Arrow types are supported for encoding (yet)!
151160

152161
default:
Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
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+
package org.apache.arrow.adapter.avro.producers;
19+
20+
import java.io.IOException;
21+
22+
import org.apache.arrow.vector.FieldVector;
23+
import org.apache.arrow.vector.complex.ListVector;
24+
import org.apache.avro.io.Encoder;
25+
26+
/**
27+
* Producer which produces array type values to an Avro encoder.
28+
* Writes the data from a {@link ListVector}.
29+
*/
30+
public class AvroArraysProducer extends BaseAvroProducer<ListVector> {
31+
32+
private final Producer<? extends FieldVector> delegate;
33+
34+
/** Instantiate an ArraysProducer. */
35+
public AvroArraysProducer(ListVector vector, Producer<? extends FieldVector> delegate) {
36+
super(vector);
37+
this.delegate = delegate;
38+
}
39+
40+
@Override
41+
public void produce(Encoder encoder) throws IOException {
42+
43+
int startOffset = vector.getOffsetBuffer().getInt(currentIndex * (long) Integer.BYTES);
44+
int endOffset = vector.getOffsetBuffer().getInt((currentIndex + 1) * (long) Integer.BYTES);
45+
int nItems = endOffset - startOffset;
46+
47+
encoder.writeArrayStart();
48+
encoder.setItemCount(nItems);
49+
50+
for (int i = 0; i < nItems; i++) {
51+
encoder.startItem();
52+
delegate.produce(encoder);
53+
}
54+
55+
encoder.writeArrayEnd();
56+
currentIndex++;
57+
}
58+
59+
@Override
60+
@SuppressWarnings("unchecked")
61+
public boolean resetValueVector(ListVector vector) {
62+
((Producer<FieldVector>) delegate).resetValueVector(vector.getDataVector());
63+
return super.resetValueVector(vector);
64+
}
65+
}

0 commit comments

Comments
 (0)