diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/BigIntVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/BigIntVectorConverter.java index 04e90e1a4..7497b7717 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/BigIntVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/BigIntVectorConverter.java @@ -8,6 +8,8 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.FieldType; @SnowflakeJdbcInternalApi public class BigIntVectorConverter extends SimpleArrowFullVectorConverter { @@ -21,6 +23,10 @@ public BigIntVectorConverter( super(allocator, vector, context, session, idx); } + private static FieldType getFieldType(boolean nullable) { + return new FieldType(nullable, new ArrowType.Int(64, true), null); + } + @Override protected boolean matchingType() { return (vector instanceof BigIntVector); @@ -28,7 +34,9 @@ protected boolean matchingType() { @Override protected BigIntVector initVector() { - BigIntVector resultVector = new BigIntVector(vector.getName(), allocator); + boolean nullable = vector.getField().isNullable(); + BigIntVector resultVector = + new BigIntVector(vector.getName(), getFieldType(nullable), allocator); resultVector.allocateNew(vector.getValueCount()); return resultVector; } diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/BinaryVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/BinaryVectorConverter.java index 8cee6d3f5..43c6e7276 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/BinaryVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/BinaryVectorConverter.java @@ -8,6 +8,8 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.FieldType; @SnowflakeJdbcInternalApi public class BinaryVectorConverter extends SimpleArrowFullVectorConverter { @@ -20,6 +22,10 @@ public BinaryVectorConverter( super(allocator, vector, context, session, idx); } + private static FieldType getFieldType(boolean nullable) { + return new FieldType(nullable, new ArrowType.Binary(), null); + } + @Override protected boolean matchingType() { return vector instanceof VarBinaryVector; @@ -27,7 +33,9 @@ protected boolean matchingType() { @Override protected VarBinaryVector initVector() { - VarBinaryVector resultVector = new VarBinaryVector(vector.getName(), allocator); + boolean nullable = vector.getField().isNullable(); + VarBinaryVector resultVector = + new VarBinaryVector(vector.getName(), getFieldType(nullable), allocator); resultVector.allocateNew(vector.getValueCount()); return resultVector; } diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/BitVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/BitVectorConverter.java index 76701800f..53056c4a6 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/BitVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/BitVectorConverter.java @@ -8,6 +8,8 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.FieldType; @SnowflakeJdbcInternalApi public class BitVectorConverter extends SimpleArrowFullVectorConverter { @@ -21,6 +23,10 @@ public BitVectorConverter( super(allocator, vector, context, session, idx); } + private static FieldType getFieldType(boolean nullable) { + return new FieldType(nullable, new ArrowType.Bool(), null); + } + @Override protected boolean matchingType() { return vector instanceof BitVector; @@ -28,7 +34,8 @@ protected boolean matchingType() { @Override protected BitVector initVector() { - BitVector resultVector = new BitVector(vector.getName(), allocator); + boolean nullable = vector.getField().isNullable(); + BitVector resultVector = new BitVector(vector.getName(), getFieldType(nullable), allocator); resultVector.allocateNew(vector.getValueCount()); return resultVector; } diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/DateVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/DateVectorConverter.java index c509af685..8920ba5c8 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/DateVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/DateVectorConverter.java @@ -9,6 +9,9 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.types.DateUnit; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.FieldType; @SnowflakeJdbcInternalApi public class DateVectorConverter extends SimpleArrowFullVectorConverter { @@ -25,6 +28,10 @@ public DateVectorConverter( this.timeZone = timeZone; } + private static FieldType getFieldType(boolean nullable) { + return new FieldType(nullable, new ArrowType.Date(DateUnit.DAY), null); + } + @Override protected boolean matchingType() { return vector instanceof DateDayVector; @@ -32,7 +39,9 @@ protected boolean matchingType() { @Override protected DateDayVector initVector() { - DateDayVector resultVector = new DateDayVector(vector.getName(), allocator); + boolean nullable = vector.getField().isNullable(); + DateDayVector resultVector = + new DateDayVector(vector.getName(), getFieldType(nullable), allocator); resultVector.allocateNew(vector.getValueCount()); return resultVector; } diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/DecimalVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/DecimalVectorConverter.java index d7421f858..d8e37eeb7 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/DecimalVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/DecimalVectorConverter.java @@ -8,6 +8,8 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.FieldType; @SnowflakeJdbcInternalApi public class DecimalVectorConverter extends SimpleArrowFullVectorConverter { @@ -21,6 +23,10 @@ public DecimalVectorConverter( super(allocator, vector, context, session, idx); } + private static FieldType getFieldType(boolean nullable, int scale, int precision) { + return new FieldType(nullable, new ArrowType.Decimal(precision, scale, 128), null); + } + @Override protected boolean matchingType() { return (vector instanceof DecimalVector); @@ -32,7 +38,9 @@ protected DecimalVector initVector() { String precisionString = vector.getField().getMetadata().get("precision"); int scale = Integer.parseInt(scaleString); int precision = Integer.parseInt(precisionString); - DecimalVector resultVector = new DecimalVector(vector.getName(), allocator, precision, scale); + boolean nullable = vector.getField().isNullable(); + DecimalVector resultVector = + new DecimalVector(vector.getName(), getFieldType(nullable, scale, precision), allocator); resultVector.allocateNew(vector.getValueCount()); return resultVector; } diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/FixedSizeListVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/FixedSizeListVectorConverter.java index 30d9f9f77..cc40d1c3e 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/FixedSizeListVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/FixedSizeListVectorConverter.java @@ -12,7 +12,9 @@ import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.FixedSizeListVector; +import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; @SnowflakeJdbcInternalApi public class FixedSizeListVectorConverter extends AbstractFullVectorConverter { @@ -41,6 +43,10 @@ public class FixedSizeListVectorConverter extends AbstractFullVectorConverter { this.valueTargetType = valueTargetType; } + private static FieldType getFieldType(boolean nullable, int size) { + return new FieldType(nullable, new ArrowType.FixedSizeList(size), null); + } + @Override protected FieldVector convertVector() throws SFException, SnowflakeSQLException, SFArrowException { @@ -50,8 +56,12 @@ protected FieldVector convertVector() FieldVector convertedDataVector = ArrowFullVectorConverterUtil.convert( allocator, dataVector, context, session, timeZoneToUse, 0, valueTargetType); + + boolean nullable = vector.getField().isNullable(); + int listSize = listVector.getListSize(); FixedSizeListVector convertedListVector = - FixedSizeListVector.empty(listVector.getName(), listVector.getListSize(), allocator); + new FixedSizeListVector( + listVector.getName(), allocator, getFieldType(nullable, listSize), null); ArrayList fields = new ArrayList<>(); fields.add(convertedDataVector.getField()); convertedListVector.initializeChildrenFromFields(fields); diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/FloatVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/FloatVectorConverter.java index e47079293..7c3a282a6 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/FloatVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/FloatVectorConverter.java @@ -8,6 +8,9 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.Float8Vector; import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.types.FloatingPointPrecision; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.FieldType; @SnowflakeJdbcInternalApi public class FloatVectorConverter extends SimpleArrowFullVectorConverter { @@ -26,9 +29,16 @@ protected boolean matchingType() { return vector instanceof Float8Vector; } + private static FieldType getFieldType(boolean nullable) { + return new FieldType( + nullable, new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), null); + } + @Override protected Float8Vector initVector() { - Float8Vector resultVector = new Float8Vector(vector.getName(), allocator); + boolean nullable = vector.getField().isNullable(); + Float8Vector resultVector = + new Float8Vector(vector.getName(), getFieldType(nullable), allocator); resultVector.allocateNew(vector.getValueCount()); return resultVector; } diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/IntVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/IntVectorConverter.java index db199e703..c8d9dd3b7 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/IntVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/IntVectorConverter.java @@ -8,6 +8,8 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.FieldType; @SnowflakeJdbcInternalApi public class IntVectorConverter extends SimpleArrowFullVectorConverter { @@ -26,9 +28,14 @@ protected boolean matchingType() { return (vector instanceof IntVector); } + private static FieldType getFieldType(boolean nullable) { + return new FieldType(nullable, new ArrowType.Int(32, true), null); + } + @Override protected IntVector initVector() { - IntVector resultVector = new IntVector(vector.getName(), allocator); + boolean nullable = vector.getField().isNullable(); + IntVector resultVector = new IntVector(vector.getName(), getFieldType(nullable), allocator); resultVector.allocateNew(vector.getValueCount()); return resultVector; } diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/ListVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/ListVectorConverter.java index bf61245b4..c4060afc9 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/ListVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/ListVectorConverter.java @@ -12,7 +12,9 @@ import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; @SnowflakeJdbcInternalApi public class ListVectorConverter extends AbstractFullVectorConverter { @@ -41,8 +43,13 @@ public class ListVectorConverter extends AbstractFullVectorConverter { this.valueTargetType = valueTargetType; } + private static FieldType getFieldType(boolean nullable) { + return new FieldType(nullable, new ArrowType.List(), null); + } + protected ListVector initVector(String name, Field field) { - ListVector convertedListVector = ListVector.empty(name, allocator); + boolean nullable = vector.getField().isNullable(); + ListVector convertedListVector = new ListVector(name, allocator, getFieldType(nullable), null); ArrayList fields = new ArrayList<>(); fields.add(field); convertedListVector.initializeChildrenFromFields(fields); @@ -58,11 +65,13 @@ protected FieldVector convertVector() FieldVector convertedDataVector = ArrowFullVectorConverterUtil.convert( allocator, dataVector, context, session, timeZoneToUse, 0, valueTargetType); - // TODO: change to convertedDataVector and make all necessary changes to make it work - ListVector convertedListVector = initVector(vector.getName(), dataVector.getField()); + ListVector convertedListVector = initVector(vector.getName(), convertedDataVector.getField()); convertedListVector.allocateNew(); convertedListVector.setValueCount(listVector.getValueCount()); - convertedListVector.getOffsetBuffer().setBytes(0, listVector.getOffsetBuffer()); + + ArrowBuf offsetBuffer = listVector.getOffsetBuffer(); + convertedListVector.getOffsetBuffer().setBytes(0L, offsetBuffer, 0L, offsetBuffer.capacity()); + ArrowBuf validityBuffer = listVector.getValidityBuffer(); convertedListVector .getValidityBuffer() diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/MapVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/MapVectorConverter.java index 0b8ec963e..abe77ba09 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/MapVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/MapVectorConverter.java @@ -9,7 +9,9 @@ import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.complex.MapVector; +import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; @SnowflakeJdbcInternalApi public class MapVectorConverter extends ListVectorConverter { @@ -25,9 +27,15 @@ public class MapVectorConverter extends ListVectorConverter { super(allocator, vector, context, session, timeZoneToUse, idx, valueTargetType); } + private static FieldType getFieldType(boolean nullable) { + return new FieldType(nullable, new ArrowType.Map(false), null); + } + @Override protected ListVector initVector(String name, Field field) { - MapVector convertedMapVector = MapVector.empty(name, allocator, false); + boolean nullable = vector.getField().isNullable(); + MapVector convertedMapVector = + new MapVector(vector.getName(), allocator, getFieldType(nullable), null); ArrayList fields = new ArrayList<>(); fields.add(field); convertedMapVector.initializeChildrenFromFields(fields); diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/SmallIntVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/SmallIntVectorConverter.java index f15a027ef..dbe96d85d 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/SmallIntVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/SmallIntVectorConverter.java @@ -8,6 +8,8 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.SmallIntVector; import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.FieldType; @SnowflakeJdbcInternalApi public class SmallIntVectorConverter extends SimpleArrowFullVectorConverter { @@ -21,6 +23,10 @@ public SmallIntVectorConverter( super(allocator, vector, context, session, idx); } + private static FieldType getFieldType(boolean nullable) { + return new FieldType(nullable, new ArrowType.Int(16, true), null); + } + @Override protected boolean matchingType() { return (vector instanceof SmallIntVector); @@ -28,7 +34,9 @@ protected boolean matchingType() { @Override protected SmallIntVector initVector() { - SmallIntVector resultVector = new SmallIntVector(vector.getName(), allocator); + boolean nullable = vector.getField().isNullable(); + SmallIntVector resultVector = + new SmallIntVector(vector.getName(), getFieldType(nullable), allocator); resultVector.allocateNew(vector.getValueCount()); return resultVector; } diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/StructVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/StructVectorConverter.java index 2dbbe4dec..09ad1a300 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/StructVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/StructVectorConverter.java @@ -15,7 +15,9 @@ import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.StructVector; +import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.TransferPair; @SnowflakeJdbcInternalApi @@ -45,6 +47,10 @@ public class StructVectorConverter extends AbstractFullVectorConverter { this.targetTypes = targetTypes; } + private static FieldType getFieldType(boolean nullable) { + return new FieldType(nullable, new ArrowType.Struct(), null); + } + protected FieldVector convertVector() throws SFException, SnowflakeSQLException, SFArrowException { try { @@ -63,7 +69,10 @@ protected FieldVector convertVector() List convertedFields = convertedVectors.stream().map(ValueVector::getField).collect(Collectors.toList()); - StructVector converted = StructVector.empty(vector.getName(), allocator); + + boolean nullable = vector.getField().isNullable(); + StructVector converted = + new StructVector(vector.getName(), allocator, getFieldType(nullable), null); converted.allocateNew(); converted.initializeChildrenFromFields(convertedFields); for (FieldVector convertedVector : convertedVectors) { diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeMicroVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeMicroVectorConverter.java index 93bc6318e..762523f12 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeMicroVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeMicroVectorConverter.java @@ -4,6 +4,7 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.TimeMicroVector; import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.types.TimeUnit; @SnowflakeJdbcInternalApi public class TimeMicroVectorConverter extends TimeVectorConverter { @@ -12,9 +13,20 @@ public TimeMicroVectorConverter(RootAllocator allocator, ValueVector vector) { super(allocator, vector); } + @Override + protected TimeUnit getTimeUnit() { + return TimeUnit.MICROSECOND; + } + + @Override + protected int getWidth() { + return 64; + } + @Override protected TimeMicroVector initVector() { - return new TimeMicroVector(vector.getName(), allocator); + return new TimeMicroVector( + vector.getName(), getFieldType(vector.getField().isNullable()), allocator); } @Override diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeMilliVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeMilliVectorConverter.java index 63a56c73c..392a16296 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeMilliVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeMilliVectorConverter.java @@ -4,6 +4,7 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.TimeMilliVector; import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.types.TimeUnit; @SnowflakeJdbcInternalApi public class TimeMilliVectorConverter extends TimeVectorConverter { @@ -11,9 +12,20 @@ public TimeMilliVectorConverter(RootAllocator allocator, ValueVector vector) { super(allocator, vector); } + @Override + protected TimeUnit getTimeUnit() { + return TimeUnit.MILLISECOND; + } + + @Override + protected int getWidth() { + return 32; + } + @Override protected TimeMilliVector initVector() { - return new TimeMilliVector(vector.getName(), allocator); + return new TimeMilliVector( + vector.getName(), getFieldType(vector.getField().isNullable()), allocator); } @Override diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeNanoVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeNanoVectorConverter.java index ad91e7a67..aea438a45 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeNanoVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeNanoVectorConverter.java @@ -4,6 +4,7 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.TimeNanoVector; import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.types.TimeUnit; @SnowflakeJdbcInternalApi public class TimeNanoVectorConverter extends TimeVectorConverter { @@ -12,9 +13,20 @@ public TimeNanoVectorConverter(RootAllocator allocator, ValueVector vector) { super(allocator, vector); } + @Override + protected TimeUnit getTimeUnit() { + return TimeUnit.NANOSECOND; + } + + @Override + protected int getWidth() { + return 64; + } + @Override protected TimeNanoVector initVector() { - return new TimeNanoVector(vector.getName(), allocator); + return new TimeNanoVector( + vector.getName(), getFieldType(vector.getField().isNullable()), allocator); } @Override diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeSecVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeSecVectorConverter.java index 64498c715..ff5c7cef9 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeSecVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeSecVectorConverter.java @@ -4,6 +4,7 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.TimeSecVector; import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.types.TimeUnit; @SnowflakeJdbcInternalApi public class TimeSecVectorConverter extends TimeVectorConverter { @@ -11,9 +12,20 @@ public TimeSecVectorConverter(RootAllocator allocator, ValueVector vector) { super(allocator, vector); } + @Override + protected TimeUnit getTimeUnit() { + return TimeUnit.SECOND; + } + + @Override + protected int getWidth() { + return 32; + } + @Override protected TimeSecVector initVector() { - return new TimeSecVector(vector.getName(), allocator); + return new TimeSecVector( + vector.getName(), getFieldType(vector.getField().isNullable()), allocator); } @Override diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeVectorConverter.java index 60d3bae3f..e514a73e6 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimeVectorConverter.java @@ -9,6 +9,9 @@ import org.apache.arrow.vector.BaseIntVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.types.TimeUnit; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.FieldType; @SnowflakeJdbcInternalApi public abstract class TimeVectorConverter @@ -21,6 +24,14 @@ public TimeVectorConverter(RootAllocator allocator, ValueVector vector) { this.vector = vector; } + protected FieldType getFieldType(boolean nullable) { + return new FieldType(nullable, new ArrowType.Time(getTimeUnit(), getWidth()), null); + } + + protected abstract TimeUnit getTimeUnit(); + + protected abstract int getWidth(); + protected abstract T initVector(); protected abstract void convertValue(T dstVector, int idx, long value); diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimestampVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimestampVectorConverter.java index 2d9f5b121..a35f8d59c 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimestampVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TimestampVectorConverter.java @@ -15,7 +15,9 @@ import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.complex.StructVector; +import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; @SnowflakeJdbcInternalApi public class TimestampVectorConverter implements ArrowFullVectorConverter { @@ -37,6 +39,9 @@ public class TimestampVectorConverter implements ArrowFullVectorConverter { private static final long NANOS_PER_MILLI = 1000000L; private static final int MILLIS_PER_SECOND = 1000; private static final int SECONDS_PER_MINUTE = 60; + private static final FieldType intType = new FieldType(false, new ArrowType.Int(32, true), null); + private static final FieldType bigIntType = + new FieldType(false, new ArrowType.Int(64, true), null); public TimestampVectorConverter( RootAllocator allocator, @@ -51,8 +56,12 @@ public TimestampVectorConverter( this.isNTZ = isNTZ; } + private static FieldType getFieldType(boolean nullable) { + return new FieldType(nullable, new ArrowType.Struct(), null); + } + private IntVector makeVectorOfZeroes(int length) { - IntVector vector = new IntVector(FIELD_NAME_FRACTION, allocator); + IntVector vector = new IntVector(FIELD_NAME_FRACTION, intType, allocator); vector.allocateNew(length); vector.zeroVector(); vector.setValueCount(length); @@ -60,7 +69,7 @@ private IntVector makeVectorOfZeroes(int length) { } private IntVector makeVectorOfUTCOffsets(int length) { - IntVector vector = new IntVector(FIELD_NAME_TIME_ZONE_INDEX, allocator); + IntVector vector = new IntVector(FIELD_NAME_TIME_ZONE_INDEX, intType, allocator); vector.allocateNew(length); vector.setValueCount(length); for (int i = 0; i < length; i++) { @@ -74,7 +83,7 @@ private SFPair normalizeTimeSinceEpoch(BigIntVector vec int scale = Integer.parseInt(vector.getField().getMetadata().get("scale")); if (scale == 0) { IntVector fractions = makeVectorOfZeroes(length); - BigIntVector epoch = new BigIntVector(FIELD_NAME_EPOCH, allocator); + BigIntVector epoch = new BigIntVector(FIELD_NAME_EPOCH, bigIntType, allocator); fractions .getValidityBuffer() .setBytes(0L, vector.getValidityBuffer(), 0L, vector.getValidityBuffer().capacity()); @@ -83,10 +92,10 @@ private SFPair normalizeTimeSinceEpoch(BigIntVector vec } long scaleFactor = ArrowResultUtil.powerOfTen(scale); long fractionScaleFactor = ArrowResultUtil.powerOfTen(9 - scale); - BigIntVector epoch = new BigIntVector(FIELD_NAME_EPOCH, allocator); + BigIntVector epoch = new BigIntVector(FIELD_NAME_EPOCH, bigIntType, allocator); epoch.allocateNew(length); epoch.setValueCount(length); - IntVector fractions = new IntVector(FIELD_NAME_FRACTION, allocator); + IntVector fractions = new IntVector(FIELD_NAME_FRACTION, intType, allocator); fractions.allocateNew(length); fractions.setValueCount(length); for (int i = 0; i < length; i++) { @@ -98,7 +107,7 @@ private SFPair normalizeTimeSinceEpoch(BigIntVector vec private IntVector makeTimeZoneOffsets( BigIntVector seconds, IntVector fractions, TimeZone timeZone) { - IntVector offsets = new IntVector(FIELD_NAME_TIME_ZONE_INDEX, allocator); + IntVector offsets = new IntVector(FIELD_NAME_TIME_ZONE_INDEX, intType, allocator); offsets.allocateNew(seconds.getValueCount()); offsets.setValueCount(seconds.getValueCount()); for (int i = 0; i < seconds.getValueCount(); i++) { @@ -113,7 +122,9 @@ private IntVector makeTimeZoneOffsets( } private StructVector pack(BigIntVector seconds, IntVector fractions, IntVector offsets) { - StructVector result = StructVector.empty(vector.getName(), allocator); + boolean nullable = vector.getField().isNullable(); + StructVector result = + new StructVector(vector.getName(), allocator, getFieldType(nullable), null); List fields = new ArrayList() { { diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TinyIntVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TinyIntVectorConverter.java index a4c7bdb22..93bd5dd0d 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TinyIntVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/TinyIntVectorConverter.java @@ -8,6 +8,8 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.TinyIntVector; import org.apache.arrow.vector.ValueVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.FieldType; @SnowflakeJdbcInternalApi public class TinyIntVectorConverter extends SimpleArrowFullVectorConverter { @@ -21,6 +23,10 @@ public TinyIntVectorConverter( super(allocator, vector, context, session, idx); } + private static FieldType getFieldType(boolean nullable) { + return new FieldType(nullable, new ArrowType.Int(8, true), null); + } + @Override protected boolean matchingType() { return (vector instanceof TinyIntVector); @@ -28,7 +34,9 @@ protected boolean matchingType() { @Override protected TinyIntVector initVector() { - TinyIntVector resultVector = new TinyIntVector(vector.getName(), allocator); + boolean nullable = vector.getField().isNullable(); + TinyIntVector resultVector = + new TinyIntVector(vector.getName(), getFieldType(nullable), allocator); resultVector.allocateNew(vector.getValueCount()); return resultVector; } diff --git a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/VarCharVectorConverter.java b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/VarCharVectorConverter.java index 8898d4498..decf46cab 100644 --- a/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/VarCharVectorConverter.java +++ b/src/main/java/net/snowflake/client/core/arrow/fullvectorconverters/VarCharVectorConverter.java @@ -8,6 +8,8 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.util.Text; @SnowflakeJdbcInternalApi @@ -21,6 +23,10 @@ public VarCharVectorConverter( super(allocator, vector, context, session, idx); } + private static FieldType getFieldType(boolean nullable) { + return new FieldType(nullable, new ArrowType.Utf8(), null); + } + @Override protected boolean matchingType() { return (vector instanceof VarCharVector); @@ -28,7 +34,9 @@ protected boolean matchingType() { @Override protected VarCharVector initVector() { - VarCharVector resultVector = new VarCharVector(vector.getName(), allocator); + boolean nullable = vector.getField().isNullable(); + VarCharVector resultVector = + new VarCharVector(vector.getName(), getFieldType(nullable), allocator); resultVector.allocateNew(vector.getValueCount()); return resultVector; } diff --git a/src/test/java/net/snowflake/client/core/arrow/ArrowBatchesTest.java b/src/test/java/net/snowflake/client/core/arrow/ArrowBatchesTest.java index dfb1fe598..cbda19f21 100644 --- a/src/test/java/net/snowflake/client/core/arrow/ArrowBatchesTest.java +++ b/src/test/java/net/snowflake/client/core/arrow/ArrowBatchesTest.java @@ -4,14 +4,29 @@ import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; +import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Map; import net.snowflake.client.core.SFException; import net.snowflake.client.core.arrow.fullvectorconverters.ArrowFullVectorConverterUtil; import net.snowflake.client.core.arrow.fullvectorconverters.IntVectorConverter; import net.snowflake.client.core.arrow.fullvectorconverters.SFArrowException; import net.snowflake.client.jdbc.SnowflakeSQLException; import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.IntervalDayVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.complex.BaseRepeatedValueVector; +import org.apache.arrow.vector.complex.MapVector; +import org.apache.arrow.vector.complex.StructVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.util.JsonStringArrayList; +import org.apache.arrow.vector.util.JsonStringHashMap; +import org.apache.arrow.vector.util.Text; import org.junit.Test; public class ArrowBatchesTest extends BaseConverterTest { @@ -52,4 +67,79 @@ public void testUnknownType() { } fail("Should throw on unsupported type"); } + + @Test + public void testMapVectorConverter() throws SFArrowException { + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + Map metadata = + new HashMap() { + { + put("logicalType", "FIXED"); + put("scale", "3"); + put("precision", "18"); + } + }; + FieldType valueVectorFieldType = + new FieldType(false, new ArrowType.Int(32, true), null, metadata); + IntVector decimalVector = new IntVector("value", valueVectorFieldType, allocator); + decimalVector.allocateNew(2); + decimalVector.set(0, 1); + decimalVector.set(1, 4); + decimalVector.setValueCount(2); + + FieldType keyVectorFieldType = new FieldType(false, new ArrowType.Utf8(), null); + VarCharVector keyVector = new VarCharVector("key", keyVectorFieldType, allocator); + keyVector.allocateNew(2); + keyVector.set(0, "a".getBytes()); + keyVector.set(1, "b".getBytes()); + keyVector.setValueCount(2); + + FieldType entriesVectorFieldType = new FieldType(false, new ArrowType.Struct(), null); + StructVector entryVector = new StructVector("entries", allocator, entriesVectorFieldType, null); + entryVector.initializeChildrenFromFields( + new ArrayList() { + { + add(decimalVector.getField()); + add(keyVector.getField()); + } + }); + entryVector.allocateNew(); + entryVector.setValueCount(2); + entryVector.getValidityBuffer().setByte(0, 3); + decimalVector.makeTransferPair(entryVector.getChild("value")).transfer(); + keyVector.makeTransferPair(entryVector.getChild("key")).transfer(); + + FieldType mapVectorFieldType = new FieldType(true, new ArrowType.Map(false), null); + MapVector mapVector = new MapVector("map", allocator, mapVectorFieldType, null); + mapVector.initializeChildrenFromFields( + new ArrayList() { + { + add(entryVector.getField()); + } + }); + mapVector.allocateNew(); + mapVector.setValueCount(1); + mapVector.getOffsetBuffer().setByte(0, 0); + mapVector.getOffsetBuffer().setByte(BaseRepeatedValueVector.OFFSET_WIDTH, 2); + mapVector.getValidityBuffer().setByte(0, 1); + entryVector.makeTransferPair(mapVector.getDataVector()).transfer(); + + FieldVector convertedVector = + ArrowFullVectorConverterUtil.convert(allocator, mapVector, this, null, null, 0, null); + assertEquals(convertedVector.getField().getChildren().size(), 1); + assertEquals(convertedVector.getField().getChildren().get(0).getType(), new ArrowType.Struct()); + assertEquals(convertedVector.getField().getChildren().get(0).getChildren().size(), 2); + assertEquals( + convertedVector.getField().getChildren().get(0).getChildren().get(1).getType(), + new ArrowType.Utf8()); + assertEquals( + convertedVector.getField().getChildren().get(0).getChildren().get(0).getType(), + new ArrowType.Decimal(18, 3, 128)); + JsonStringArrayList> result = + (JsonStringArrayList>) convertedVector.getObject(0); + assertEquals(result.get(0).get("key"), new Text("a")); + assertEquals(result.get(1).get("key"), new Text("b")); + assertEquals(result.get(0).get("value"), BigDecimal.valueOf(0.001)); + assertEquals(result.get(1).get("value"), BigDecimal.valueOf(0.004)); + } }