Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
2 changes: 2 additions & 0 deletions vector/src/main/codegen/includes/vv_imports.ftl
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,8 @@ import org.apache.arrow.vector.util.JsonStringArrayList;
import java.util.Arrays;
import java.util.Random;
import java.util.List;
import java.util.HashMap;
import java.util.Map;

import java.io.Closeable;
import java.io.InputStream;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ public void copyAsField(String name, ${name}Writer writer) {

</#list></#list>

public void copyAsValue(StructWriter writer, ExtensionTypeWriterFactory writerFactory) {
public void copyAsValue(StructWriter writer, ExtensionTypeFactory writerFactory) {
fail("CopyAsValue StructWriter");
}

Expand Down
4 changes: 2 additions & 2 deletions vector/src/main/codegen/templates/AbstractFieldWriter.java
Original file line number Diff line number Diff line change
Expand Up @@ -110,10 +110,10 @@ public void endEntry() {
public void write(ExtensionHolder var1) {
this.fail("ExtensionType");
}
public void writeExtension(Object var1) {
public void writeExtension(Object var1, ExtensionType var2) {
this.fail("ExtensionType");
}
public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory var1) {
public void addExtensionTypeWriterFactory(ExtensionTypeFactory var1, ExtensionType var2) {
this.fail("ExtensionType");
}

Expand Down
6 changes: 3 additions & 3 deletions vector/src/main/codegen/templates/BaseReader.java
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ public interface RepeatedStructReader extends StructReader{
boolean next();
int size();
void copyAsValue(StructWriter writer);
void copyAsValue(StructWriter writer, ExtensionTypeWriterFactory writerFactory);
void copyAsValue(StructWriter writer, ExtensionTypeFactory writerFactory);
}

public interface ListReader extends BaseReader{
Expand All @@ -60,7 +60,7 @@ public interface RepeatedListReader extends ListReader{
boolean next();
int size();
void copyAsValue(ListWriter writer);
void copyAsValue(ListWriter writer, ExtensionTypeWriterFactory writerFactory);
void copyAsValue(ListWriter writer, ExtensionTypeFactory writerFactory);
}

public interface MapReader extends BaseReader{
Expand All @@ -71,7 +71,7 @@ public interface RepeatedMapReader extends MapReader{
boolean next();
int size();
void copyAsValue(MapWriter writer);
void copyAsValue(MapWriter writer, ExtensionTypeWriterFactory writerFactory);
void copyAsValue(MapWriter writer, ExtensionTypeFactory writerFactory);
}

public interface ScalarReader extends
Expand Down
4 changes: 2 additions & 2 deletions vector/src/main/codegen/templates/BaseWriter.java
Original file line number Diff line number Diff line change
Expand Up @@ -122,14 +122,14 @@ public interface ExtensionWriter extends BaseWriter {
*
* @param value the extension type value to write
*/
void writeExtension(Object value);
void writeExtension(Object value, ExtensionType extensionType);

/**
* Adds the given extension type factory. This factory allows configuring writer implementations for specific ExtensionTypeVector.
*
* @param factory the extension type factory to add
*/
void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory);
void addExtensionTypeWriterFactory(ExtensionTypeFactory factory, ExtensionType extensionType);
}

public interface ScalarWriter extends
Expand Down
10 changes: 6 additions & 4 deletions vector/src/main/codegen/templates/ComplexCopier.java
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import org.apache.arrow.vector.complex.reader.FieldReader;
import org.apache.arrow.vector.complex.writer.FieldWriter;
import org.apache.arrow.vector.types.Types;
import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;

<@pp.dropOutputFile />
<@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/ComplexCopier.java" />
Expand All @@ -45,11 +46,11 @@ public static void copy(FieldReader input, FieldWriter output) {
writeValue(input, output, null);
}

public static void copy(FieldReader input, FieldWriter output, ExtensionTypeWriterFactory extensionTypeWriterFactory) {
public static void copy(FieldReader input, FieldWriter output, ExtensionTypeFactory extensionTypeWriterFactory) {
writeValue(input, output, extensionTypeWriterFactory);
}

private static void writeValue(FieldReader reader, FieldWriter writer, ExtensionTypeWriterFactory extensionTypeWriterFactory) {
private static void writeValue(FieldReader reader, FieldWriter writer, ExtensionTypeFactory extensionTypeWriterFactory) {
final MinorType mt = reader.getMinorType();

switch (mt) {
Expand Down Expand Up @@ -120,9 +121,10 @@ private static void writeValue(FieldReader reader, FieldWriter writer, Extension
}
if (reader.isSet()) {
Object value = reader.readObject();
ExtensionType extensionType = (ExtensionType) reader.getField().getType();
if (value != null) {
writer.addExtensionTypeWriterFactory(extensionTypeWriterFactory);
writer.writeExtension(value);
writer.addExtensionTypeWriterFactory(extensionTypeWriterFactory, extensionType);
writer.writeExtension(value, extensionType);
}
} else {
writer.writeNull();
Expand Down
2 changes: 1 addition & 1 deletion vector/src/main/codegen/templates/NullReader.java
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,7 @@ public void read(int arrayIndex, Nullable${name}Holder holder){
}
</#list></#list>

public void copyAsValue(StructWriter writer, ExtensionTypeWriterFactory writerFactory){}
public void copyAsValue(StructWriter writer, ExtensionTypeFactory writerFactory){}
public void read(ExtensionHolder holder) {
holder.isSet = 0;
}
Expand Down
15 changes: 6 additions & 9 deletions vector/src/main/codegen/templates/PromotableWriter.java
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,8 @@
* limitations under the License.
*/

<@pp.dropOutputFile />
import org.apache.arrow.vector.complex.impl.ExtensionTypeFactory;
import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;<@pp.dropOutputFile />
<@pp.changeOutputFile name="/org/apache/arrow/vector/complex/impl/PromotableWriter.java" />

<#include "/@includes/license.ftl" />
Expand Down Expand Up @@ -541,17 +542,13 @@ public void writeLargeVarChar(String value) {
}

@Override
public void writeExtension(Object value) {
getWriter(MinorType.EXTENSIONTYPE).writeExtension(value);
public void writeExtension(Object value, ExtensionType extensionType) {
getWriter(MinorType.EXTENSIONTYPE, extensionType).writeExtension(value, extensionType);
}

@Override
public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory) {
getWriter(MinorType.EXTENSIONTYPE).addExtensionTypeWriterFactory(factory);
}

public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory factory, ArrowType arrowType) {
getWriter(MinorType.EXTENSIONTYPE, arrowType).addExtensionTypeWriterFactory(factory);
public void addExtensionTypeWriterFactory(ExtensionTypeFactory var1, ExtensionType var2) {
getWriter(MinorType.EXTENSIONTYPE, var2).addExtensionTypeWriterFactory(var1, var2);
}

@Override
Expand Down
10 changes: 6 additions & 4 deletions vector/src/main/codegen/templates/UnionListWriter.java
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
*/

import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.vector.complex.impl.ExtensionTypeFactory;
import org.apache.arrow.vector.complex.writer.Decimal256Writer;
import org.apache.arrow.vector.complex.writer.DecimalWriter;
import org.apache.arrow.vector.holders.Decimal256Holder;
Expand All @@ -24,6 +25,7 @@

import java.lang.UnsupportedOperationException;
import java.math.BigDecimal;
import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;

<@pp.dropOutputFile />
<#list ["List", "ListView", "LargeList", "LargeListView"] as listName>
Expand Down Expand Up @@ -336,14 +338,14 @@ public void writeNull() {
}

@Override
public void writeExtension(Object value) {
writer.writeExtension(value);
public void writeExtension(Object value, ExtensionType extensionType) {
writer.writeExtension(value, extensionType);
writer.setPosition(writer.idx() + 1);
}

@Override
public void addExtensionTypeWriterFactory(ExtensionTypeWriterFactory var1) {
writer.addExtensionTypeWriterFactory(var1, extensionType);
public void addExtensionTypeWriterFactory(ExtensionTypeFactory var1, ExtensionType var2) {
writer.addExtensionTypeWriterFactory(var1, var2);
}

public void write(ExtensionHolder var1) {
Expand Down
30 changes: 29 additions & 1 deletion vector/src/main/codegen/templates/UnionReader.java
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@
*/


import java.util.HashMap;
import java.util.Map;
import org.apache.arrow.vector.types.Types.MinorType;
import org.apache.arrow.vector.types.pojo.Field;

Expand Down Expand Up @@ -43,9 +45,12 @@ public class UnionReader extends AbstractFieldReader {

private BaseReader[] readers = new BaseReader[NUM_SUPPORTED_TYPES];
public UnionVector data;

private ExtensionTypeFactory extensionTypeFactory;
private Map<ExtensionType, ExtensionReader> extensionTypeReaders = new HashMap<>();

public UnionReader(UnionVector data) {
this.data = data;
this.extensionTypeFactory = data.getExtensionTypeFactory();
}

public MinorType getMinorType() {
Expand Down Expand Up @@ -79,6 +84,10 @@ public void read(int index, UnionHolder holder) {
}

private FieldReader getReaderForIndex(int index) {
return getReaderForIndex(index, null);
}

private FieldReader getReaderForIndex(int index, ExtensionType extensionType) {
int typeValue = data.getTypeValue(index);
FieldReader reader = (FieldReader) readers[typeValue];
if (reader != null) {
Expand All @@ -95,6 +104,11 @@ private FieldReader getReaderForIndex(int index) {
return (FieldReader) getListView();
case MAP:
return (FieldReader) getMap();
case EXTENSIONTYPE:
if(extensionType == null) {
throw new IllegalStateException("Cannot read extension type without extensionType");
}
return (FieldReader) getExtension(extensionType);
<#list vv.types as type>
<#list type.minor as minor>
<#assign name = minor.class?cap_first />
Expand Down Expand Up @@ -214,6 +228,20 @@ public void copyAsValue(${name}Writer writer){
</#list>
</#list>

public void read(ExtensionHolder holder){
getReaderForIndex(idx(), extensionTypeFactory.getExtensionTypeByHolder(holder)).read(holder);
}

private ExtensionReader getExtension(ExtensionType arrowType) {
ExtensionReader extensionReader = extensionTypeReaders.get(arrowType);
if (extensionReader == null) {
extensionReader = extensionTypeFactory.getReaderImpl(data.getExtensionTypeVector(arrowType));
extensionReader.setPosition(idx());
extensionTypeReaders.put(arrowType, extensionReader);
}
return extensionReader;
}

@Override
public void copyAsValue(ListWriter writer) {
ComplexCopier.copy(this, (FieldWriter) writer);
Expand Down
37 changes: 37 additions & 0 deletions vector/src/main/codegen/templates/UnionVector.java
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,8 @@ public class UnionVector extends AbstractContainerVector implements FieldVector
private ValueVector singleVector;

private int typeBufferAllocationSizeInBytes;

private ExtensionTypeFactory extensionTypeFactory;

private final FieldType fieldType;
private final Field[] typeIds = new Field[Byte.MAX_VALUE + 1];
Expand Down Expand Up @@ -325,6 +327,26 @@ public StructVector getStruct() {
</#if>
</#list>
</#list>

private ExtensionTypeVector extensionTypeVector;

public <T extends ExtensionTypeVector> T getExtensionTypeVector(ExtensionType type) {
return getExtensionTypeVector(null, type);
}

public <T extends ExtensionTypeVector> T getExtensionTypeVector(String name, ExtensionType type) {
if (extensionTypeVector == null) {
int vectorCount = internalStruct.size();
extensionTypeVector = addOrGet(name, MinorType.EXTENSIONTYPE, type, extensionTypeFactory.getVectorClass(type));
if (internalStruct.size() > vectorCount) {
extensionTypeVector.allocateNew();
if (callBack != null) {
callBack.doWork();
}
}
}
return (T) extensionTypeVector;
}

public ListVector getList() {
if (listVector == null) {
Expand Down Expand Up @@ -725,6 +747,8 @@ public ValueVector getVectorByType(int typeId, ArrowType arrowType) {
return getListView();
case MAP:
return getMap(name, arrowType);
case EXTENSIONTYPE:
return getExtensionTypeVector(name, (ExtensionType) arrowType);
default:
throw new UnsupportedOperationException("Cannot support type: " + MinorType.values()[typeId]);
}
Expand Down Expand Up @@ -847,6 +871,11 @@ public void setSafe(int index, Nullable${name}Holder holder) {
</#list>
</#list>

public void setSafe(int index, ExtensionHolder holder) {
setType(index, MinorType.EXTENSIONTYPE);
getExtensionTypeVector(null).setSafe(index, holder);
}

public void setType(int index, MinorType type) {
while (index >= getTypeBufferValueCapacity()) {
reallocTypeBuffer();
Expand Down Expand Up @@ -929,4 +958,12 @@ public void setInitialCapacity(int valueCount, double density) {
public void setNull(int index) {
throw new UnsupportedOperationException("The method setNull() is not supported on UnionVector.");
}

public void setExtensionTypeFactory(ExtensionTypeFactory extensionTypeFactory) {
this.extensionTypeFactory = extensionTypeFactory;
}

public ExtensionTypeFactory getExtensionTypeFactory() {
return extensionTypeFactory;
}
}
Loading
Loading