From c7696df8648f2719763af363e1324fe4cc49cb69 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 15 Jul 2026 16:38:03 +0100 Subject: [PATCH] typytype Signed-off-by: Robert Kruszewski --- .../test/java/dev/vortex/api/TestMinimal.java | 9 +- .../java/dev/vortex/jni/JNIWriterTest.java | 3 +- .../java/dev/vortex/spark/ArrowUtils.java | 111 ++-- .../spark/read/VortexArrowColumnVector.java | 302 +++++++--- .../java/dev/vortex/spark/ArrowUtilsTest.java | 129 ++++- .../read/VortexArrowColumnVectorTest.java | 523 ++++++++++++++++++ vortex-jni/src/dtype.rs | 55 +- vortex-jni/src/scan.rs | 13 +- 8 files changed, 947 insertions(+), 198 deletions(-) create mode 100644 java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexArrowColumnVectorTest.java diff --git a/java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java b/java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java index aeee8febcf4..3a6f9ca0883 100644 --- a/java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java +++ b/java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java @@ -23,7 +23,6 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.FieldVector; -import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.ipc.ArrowReader; @@ -155,8 +154,8 @@ public void testProjectedScan() throws Exception { List people = readAll(ds, options, allocator, batch -> { List results = new ArrayList<>(); - VarCharVector names = (VarCharVector) batch.getVector("Name"); - VarCharVector states = (VarCharVector) batch.getVector("State"); + ViewVarCharVector names = (ViewVarCharVector) batch.getVector("Name"); + ViewVarCharVector states = (ViewVarCharVector) batch.getVector("State"); for (int i = 0; i < batch.getRowCount(); i++) { String name = names.isNull(i) ? null : new String(names.get(i), UTF_8); String state = states.isNull(i) ? null : new String(states.get(i), UTF_8); @@ -272,9 +271,9 @@ private static List readAll( private static List readFullBatch(VectorSchemaRoot root) { List result = new ArrayList<>(); - VarCharVector names = (VarCharVector) root.getVector("Name"); + ViewVarCharVector names = (ViewVarCharVector) root.getVector("Name"); FieldVector salaries = root.getVector("Salary"); - VarCharVector states = (VarCharVector) root.getVector("State"); + ViewVarCharVector states = (ViewVarCharVector) root.getVector("State"); for (int i = 0; i < root.getRowCount(); i++) { String name = names.isNull(i) ? null : new String(names.get(i), UTF_8); diff --git a/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java b/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java index 03869ae8598..efc8c6064f5 100644 --- a/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java +++ b/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java @@ -30,6 +30,7 @@ import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.types.pojo.ArrowType; @@ -192,7 +193,7 @@ public void testWriteBatch() throws IOException { try (ArrowReader reader = p.scanArrow(allocator)) { reader.loadNextBatch(); VectorSchemaRoot resultRoot = reader.getVectorSchemaRoot(); - VarCharVector nameOut = (VarCharVector) resultRoot.getVector("name"); + ViewVarCharVector nameOut = (ViewVarCharVector) resultRoot.getVector("name"); IntVector ageOut = (IntVector) resultRoot.getVector("age"); assertEquals("Alice", nameOut.getObject(0).toString()); assertEquals("Bob", nameOut.getObject(1).toString()); diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/ArrowUtils.java b/java/vortex-spark/src/main/java/dev/vortex/spark/ArrowUtils.java index 6b3c130c467..6e0166a2873 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/ArrowUtils.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/ArrowUtils.java @@ -4,7 +4,6 @@ package dev.vortex.spark; import dev.vortex.relocated.org.apache.arrow.vector.types.DateUnit; -import dev.vortex.relocated.org.apache.arrow.vector.types.FloatingPointPrecision; import dev.vortex.relocated.org.apache.arrow.vector.types.TimeUnit; import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.ArrowType; import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.Field; @@ -18,8 +17,18 @@ * Utility class for converting Arrow types to Spark SQL data types. * *

This class provides static methods to convert Arrow field definitions and type definitions into their - * corresponding Spark SQL DataType representations. It handles the mapping between Arrow's type system and Spark's type - * system, including complex types like structs and arrays. + * corresponding Spark SQL DataType representations. The mapping matches Spark 4.1's own {@code ArrowUtils}, extended + * where Vortex produces types Spark's mapping does not cover: + * + *

    + *
  • the Arrow view types map to their logical Spark types: {@code Utf8View} to {@code StringType}, + * {@code BinaryView} to {@code BinaryType}, and {@code ListView} to {@code ArrayType} + *
  • timestamps of every unit map to {@code TimestampType}/{@code TimestampNTZType}; + * {@link dev.vortex.spark.read.VortexArrowColumnVector} normalizes non-microsecond values on read + *
+ * + *

The one intentional divergence from Spark 4.1 is Arrow's Time type: Spark's {@code TimeType} only exists in Spark + * 4.1+, and this module compiles a single source set against both Spark 3.5 and 4.1, so Time is rejected. */ public final class ArrowUtils { private ArrowUtils() {} @@ -27,8 +36,8 @@ private ArrowUtils() {} /** * Converts an Arrow Field to a Spark SQL DataType. * - *

This method handles complex types like structs and arrays by recursively converting their child fields. For - * primitive types, it delegates to {@link #fromArrowType(ArrowType)}. + *

This method handles nested types like structs, lists and maps by recursively converting their child fields. + * For non-nested types, it delegates to {@link #fromArrowType(ArrowType)}. * * @param field the Arrow field to convert * @return the corresponding Spark SQL DataType @@ -43,11 +52,19 @@ public static DataType fromArrowField(Field field) { return new StructField(child.getName(), dt, child.isNullable(), Metadata.empty()); }) .collect(Collectors.toList())); - case List: { + case List: + case ListView: { Field elementField = field.getChildren().get(0); DataType elementType = fromArrowField(elementField); return DataTypes.createArrayType(elementType, elementField.isNullable()); } + case Map: { + Field entries = field.getChildren().get(0); + Field keyField = entries.getChildren().get(0); + Field valueField = entries.getChildren().get(1); + return DataTypes.createMapType( + fromArrowField(keyField), fromArrowField(valueField), valueField.isNullable()); + } default: return fromArrowType(field.getType()); } @@ -56,13 +73,12 @@ public static DataType fromArrowField(Field field) { /** * Converts an Arrow type to a Spark SQL DataType. * - *

This method maps primitive Arrow types to their corresponding Spark SQL types. It supports most common Arrow - * types including integers, floating point numbers, strings, binary data, dates, timestamps, decimals, and nulls. + *

This method maps non-nested Arrow types to their corresponding Spark SQL types, following Spark 4.1's own + * Arrow type mapping plus the view types (see the class documentation). * * @param dt the Arrow type to convert * @return the corresponding Spark SQL DataType - * @throws UnsupportedOperationException if the Arrow type configuration is not supported - * @throws RuntimeException if the Arrow type is not recognized + * @throws UnsupportedOperationException if the Arrow type has no Spark representation */ public static DataType fromArrowType(ArrowType dt) { switch (dt.getTypeID()) { @@ -70,26 +86,31 @@ public static DataType fromArrowType(ArrowType dt) { return DataTypes.BooleanType; case Int: { ArrowType.Int intType = (ArrowType.Int) dt; - if (intType.getIsSigned() && intType.getBitWidth() == 8) { - return DataTypes.ByteType; - } else if (intType.getIsSigned() && intType.getBitWidth() == 16) { - return DataTypes.ShortType; - } else if (intType.getIsSigned() && intType.getBitWidth() == 32) { - return DataTypes.IntegerType; - } else if (intType.getIsSigned() && intType.getBitWidth() == 64) { - return DataTypes.LongType; - } else { - throw new UnsupportedOperationException("Unsupported Arrow type: " + dt); + if (!intType.getIsSigned()) { + throw new UnsupportedOperationException("Unsupported Arrow unsigned integer type: " + dt); + } + switch (intType.getBitWidth()) { + case 8: + return DataTypes.ByteType; + case 16: + return DataTypes.ShortType; + case 32: + return DataTypes.IntegerType; + case 64: + return DataTypes.LongType; + default: + throw new UnsupportedOperationException("Unsupported Arrow integer bit width: " + dt); } } case FloatingPoint: { ArrowType.FloatingPoint floatType = (ArrowType.FloatingPoint) dt; - if (floatType.getPrecision() == FloatingPointPrecision.SINGLE) { - return DataTypes.FloatType; - } else if (floatType.getPrecision() == FloatingPointPrecision.DOUBLE) { - return DataTypes.DoubleType; - } else { - throw new UnsupportedOperationException("Unsupported Arrow type: " + dt); + switch (floatType.getPrecision()) { + case SINGLE: + return DataTypes.FloatType; + case DOUBLE: + return DataTypes.DoubleType; + default: + throw new UnsupportedOperationException("Unsupported Arrow float precision: " + dt); } } case Decimal: { @@ -98,30 +119,52 @@ public static DataType fromArrowType(ArrowType dt) { } case Utf8: case LargeUtf8: + case Utf8View: return DataTypes.StringType; case Binary: case LargeBinary: + case BinaryView: return DataTypes.BinaryType; case Date: { ArrowType.Date dateType = (ArrowType.Date) dt; if (dateType.getUnit() == DateUnit.DAY) { return DataTypes.DateType; - } else { - throw new UnsupportedOperationException("Unsupported Arrow type: " + dt); } + throw new UnsupportedOperationException("Unsupported Arrow date unit: " + dt); } case Timestamp: { + // Spark timestamps are physically microseconds; VortexArrowColumnVector's accessor + // normalizes second/millisecond/nanosecond values on read. ArrowType.Timestamp ts = (ArrowType.Timestamp) dt; - if (ts.getUnit() == TimeUnit.MICROSECOND) { - return ts.getTimezone() != null ? DataTypes.TimestampType : DataTypes.TimestampNTZType; - } else { - throw new UnsupportedOperationException("Unsupported Arrow type: " + dt); - } + return ts.getTimezone() != null ? DataTypes.TimestampType : DataTypes.TimestampNTZType; } case Null: return DataTypes.NullType; + case Interval: { + ArrowType.Interval interval = (ArrowType.Interval) dt; + switch (interval.getUnit()) { + case YEAR_MONTH: + return DataTypes.createYearMonthIntervalType(); + case MONTH_DAY_NANO: + return DataTypes.CalendarIntervalType; + default: + throw new UnsupportedOperationException("Unsupported Arrow interval unit: " + dt); + } + } + case Duration: { + ArrowType.Duration duration = (ArrowType.Duration) dt; + if (duration.getUnit() != TimeUnit.MICROSECOND) { + throw new UnsupportedOperationException("Unsupported Arrow duration unit: " + dt); + } + return DataTypes.createDayTimeIntervalType(); + } + case Time: + // Spark 3.5 has no TIME type (TimeType only exists in Spark 4.1+), and this + // module compiles a single source set against both Spark versions. + throw new UnsupportedOperationException("Arrow Time type has no Spark 3.5 representation: " + dt); default: - throw new IllegalArgumentException("Unsupported Arrow type: " + dt); + throw new UnsupportedOperationException( + "Unsupported Arrow type: " + dt + " (type id: " + dt.getTypeID() + ")"); } } } diff --git a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexArrowColumnVector.java b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexArrowColumnVector.java index 698707cfdee..be191a0ed43 100644 --- a/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexArrowColumnVector.java +++ b/java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexArrowColumnVector.java @@ -7,12 +7,16 @@ import com.jakewharton.nopen.annotation.Open; import dev.vortex.relocated.org.apache.arrow.vector.*; import dev.vortex.relocated.org.apache.arrow.vector.complex.ListVector; +import dev.vortex.relocated.org.apache.arrow.vector.complex.ListViewVector; import dev.vortex.relocated.org.apache.arrow.vector.complex.MapVector; import dev.vortex.relocated.org.apache.arrow.vector.complex.StructVector; import dev.vortex.relocated.org.apache.arrow.vector.holders.NullableLargeVarCharHolder; import dev.vortex.relocated.org.apache.arrow.vector.holders.NullableVarCharHolder; +import dev.vortex.relocated.org.apache.arrow.vector.types.TimeUnit; +import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.ArrowType; import dev.vortex.spark.ArrowUtils; import org.apache.spark.sql.types.DataType; +import org.apache.spark.sql.types.DataTypes; import org.apache.spark.sql.types.Decimal; import org.apache.spark.sql.vectorized.ColumnVector; import org.apache.spark.sql.vectorized.ColumnarArray; @@ -23,12 +27,17 @@ * Spark ColumnVector implementation that wraps Apache Arrow vectors from Vortex data. *

* This class provides a bridge between Vortex's Arrow-based data representation and Spark's - * ColumnVector interface. It supports all major Arrow data types including primitives, strings, - * binary data, decimals, dates, timestamps, arrays, maps, and structs. + * ColumnVector interface. It supports the same Arrow vector types as Spark 4.1's own + * {@code ArrowColumnVector} — booleans, signed integers, single/double floats, 128-bit decimals, + * strings and binary (regular and large), day dates, timestamps of every unit with and without + * timezone (normalized to microseconds), microsecond durations, year-month and month-day-nano + * intervals, nulls, lists, maps, and structs — plus the Arrow view types (Utf8View, BinaryView, + * ListView) that Vortex produces natively. *

- * The implementation uses type-specific accessors to efficiently retrieve values from the - * underlying Arrow vectors while maintaining Spark's expected API contract. - * + * Arrow types outside that set (unsigned integers, half floats, 256-bit decimals, + * non-microsecond durations, Time, fixed-size binary and lists, large lists, unions, run-end and + * dictionary encodings) are rejected with a descriptive {@link UnsupportedOperationException}. + * * @see ColumnVector * @see ValueVector */ @@ -39,7 +48,7 @@ public class VortexArrowColumnVector extends ColumnVector { /** * Returns the underlying Apache Arrow ValueVector wrapped by this column vector. - * + * * @return the Arrow ValueVector containing the actual data */ public ValueVector getValueVector() { @@ -48,7 +57,7 @@ public ValueVector getValueVector() { /** * Returns whether this column contains any null values. - * + * * @return true if the column contains at least one null value, false otherwise */ @Override @@ -58,7 +67,7 @@ public boolean hasNull() { /** * Returns the total number of null values in this column. - * + * * @return the count of null values */ @Override @@ -76,7 +85,7 @@ public void close() {} /** * Returns whether the value at the specified row is null. - * + * * @param rowId the row index to check * @return true if the value at rowId is null, false otherwise */ @@ -87,7 +96,7 @@ public boolean isNullAt(int rowId) { /** * Returns the boolean value at the specified row. - * + * * @param rowId the row index * @return the boolean value at rowId * @throws UnsupportedOperationException if this column is not of boolean type @@ -99,7 +108,7 @@ public boolean getBoolean(int rowId) { /** * Returns the byte value at the specified row. - * + * * @param rowId the row index * @return the byte value at rowId * @throws UnsupportedOperationException if this column is not of byte type @@ -111,7 +120,7 @@ public byte getByte(int rowId) { /** * Returns the short value at the specified row. - * + * * @param rowId the row index * @return the short value at rowId * @throws UnsupportedOperationException if this column is not of short type @@ -123,7 +132,7 @@ public short getShort(int rowId) { /** * Returns the int value at the specified row. - * + * * @param rowId the row index * @return the int value at rowId * @throws UnsupportedOperationException if this column is not of int type @@ -135,7 +144,7 @@ public int getInt(int rowId) { /** * Returns the long value at the specified row. - * + * * @param rowId the row index * @return the long value at rowId * @throws UnsupportedOperationException if this column is not of long type @@ -147,7 +156,7 @@ public long getLong(int rowId) { /** * Returns the float value at the specified row. - * + * * @param rowId the row index * @return the float value at rowId * @throws UnsupportedOperationException if this column is not of float type @@ -159,7 +168,7 @@ public float getFloat(int rowId) { /** * Returns the double value at the specified row. - * + * * @param rowId the row index * @return the double value at rowId * @throws UnsupportedOperationException if this column is not of double type @@ -171,7 +180,7 @@ public double getDouble(int rowId) { /** * Returns the decimal value at the specified row with the given precision and scale. - * + * * @param rowId the row index * @param precision the precision of the decimal * @param scale the scale of the decimal @@ -186,7 +195,7 @@ public Decimal getDecimal(int rowId, int precision, int scale) { /** * Returns the UTF8String value at the specified row. - * + * * @param rowId the row index * @return the UTF8String value at rowId, or null if the value is null * @throws UnsupportedOperationException if this column is not of string type @@ -199,7 +208,7 @@ public UTF8String getUTF8String(int rowId) { /** * Returns the binary data (byte array) at the specified row. - * + * * @param rowId the row index * @return the byte array at rowId, or null if the value is null * @throws UnsupportedOperationException if this column is not of binary type @@ -212,7 +221,7 @@ public byte[] getBinary(int rowId) { /** * Returns the array value at the specified row. - * + * * @param rowId the row index * @return the ColumnarArray at rowId, or null if the value is null * @throws UnsupportedOperationException if this column is not of array type @@ -225,7 +234,7 @@ public ColumnarArray getArray(int rowId) { /** * Returns the map value at the specified row. - * + * * @param rowId the row index * @return the ColumnarMap at rowId, or null if the value is null * @throws UnsupportedOperationException if this column is not of map type @@ -240,8 +249,9 @@ public ColumnarMap getMap(int rowId) { * Returns the child column at the specified ordinal. *

* This is used for complex types like structs where each field is represented - * as a child column. - * + * as a child column. For month-day-nano interval columns, the children follow the + * months/days/microseconds protocol documented on {@link ColumnVector#getInterval(int)}. + * * @param ordinal the index of the child column * @return the child VortexArrowColumnVector at the specified ordinal * @throws ArrayIndexOutOfBoundsException if ordinal is out of bounds @@ -256,7 +266,7 @@ public VortexArrowColumnVector getChild(int ordinal) { *

* This constructor is used internally for creating column vectors before * the underlying Arrow vector is available. - * + * * @param type the Spark DataType for this column */ VortexArrowColumnVector(DataType type) { @@ -268,7 +278,7 @@ public VortexArrowColumnVector getChild(int ordinal) { *

* This constructor automatically determines the appropriate Spark DataType from * the Arrow field and initializes the type-specific accessor. - * + * * @param vector the Arrow ValueVector to wrap * @throws UnsupportedOperationException if the vector type is not supported */ @@ -277,60 +287,78 @@ public VortexArrowColumnVector(ValueVector vector) { initAccessor(vector); } + private static VortexArrowColumnVector withAccessor(DataType type, ArrowVectorAccessor accessor) { + VortexArrowColumnVector column = new VortexArrowColumnVector(type); + column.accessor = accessor; + return column; + } + void initAccessor(ValueVector vector) { - if (vector instanceof BitVector) { - accessor = new VortexArrowColumnVector.BooleanAccessor((BitVector) vector); - } else if (vector instanceof TinyIntVector) { - accessor = new VortexArrowColumnVector.ByteAccessor((TinyIntVector) vector); - } else if (vector instanceof SmallIntVector) { - accessor = new VortexArrowColumnVector.ShortAccessor((SmallIntVector) vector); - } else if (vector instanceof IntVector) { - accessor = new VortexArrowColumnVector.IntAccessor((IntVector) vector); - } else if (vector instanceof BigIntVector) { - accessor = new VortexArrowColumnVector.LongAccessor((BigIntVector) vector); - } else if (vector instanceof Float4Vector) { - accessor = new VortexArrowColumnVector.FloatAccessor((Float4Vector) vector); - } else if (vector instanceof Float8Vector) { - accessor = new VortexArrowColumnVector.DoubleAccessor((Float8Vector) vector); - } else if (vector instanceof DecimalVector) { - accessor = new VortexArrowColumnVector.DecimalAccessor((DecimalVector) vector); - } else if (vector instanceof VarCharVector) { - accessor = new VortexArrowColumnVector.StringAccessor((VarCharVector) vector); - } else if (vector instanceof LargeVarCharVector) { - accessor = new VortexArrowColumnVector.LargeStringAccessor((LargeVarCharVector) vector); - } else if (vector instanceof VarBinaryVector) { - accessor = new VortexArrowColumnVector.BinaryAccessor((VarBinaryVector) vector); - } else if (vector instanceof LargeVarBinaryVector) { - accessor = new VortexArrowColumnVector.LargeBinaryAccessor((LargeVarBinaryVector) vector); - } else if (vector instanceof DateDayVector) { - accessor = new VortexArrowColumnVector.DateAccessor((DateDayVector) vector); - } else if (vector instanceof TimeStampMicroTZVector) { - accessor = new VortexArrowColumnVector.TimestampAccessor((TimeStampMicroTZVector) vector); - } else if (vector instanceof TimeStampMicroVector) { - accessor = new VortexArrowColumnVector.TimestampNTZAccessor((TimeStampMicroVector) vector); - } else if (vector instanceof MapVector) { - MapVector mapVector = (MapVector) vector; + if (vector instanceof BitVector bitVector) { + accessor = new VortexArrowColumnVector.BooleanAccessor(bitVector); + } else if (vector instanceof TinyIntVector tinyIntVector) { + accessor = new VortexArrowColumnVector.ByteAccessor(tinyIntVector); + } else if (vector instanceof SmallIntVector smallIntVector) { + accessor = new VortexArrowColumnVector.ShortAccessor(smallIntVector); + } else if (vector instanceof IntVector intVector) { + accessor = new VortexArrowColumnVector.IntAccessor(intVector); + } else if (vector instanceof BigIntVector bigIntVector) { + accessor = new VortexArrowColumnVector.LongAccessor(bigIntVector); + } else if (vector instanceof Float4Vector float4Vector) { + accessor = new VortexArrowColumnVector.FloatAccessor(float4Vector); + } else if (vector instanceof Float8Vector float8Vector) { + accessor = new VortexArrowColumnVector.DoubleAccessor(float8Vector); + } else if (vector instanceof DecimalVector decimalVector) { + accessor = new VortexArrowColumnVector.DecimalAccessor(decimalVector); + } else if (vector instanceof VarCharVector varCharVector) { + accessor = new VortexArrowColumnVector.StringAccessor(varCharVector); + } else if (vector instanceof LargeVarCharVector largeVarCharVector) { + accessor = new VortexArrowColumnVector.LargeStringAccessor(largeVarCharVector); + } else if (vector instanceof ViewVarCharVector viewVarCharVector) { + accessor = new VortexArrowColumnVector.StringViewAccessor(viewVarCharVector); + } else if (vector instanceof VarBinaryVector varBinaryVector) { + accessor = new VortexArrowColumnVector.BinaryAccessor(varBinaryVector); + } else if (vector instanceof LargeVarBinaryVector largeVarBinaryVector) { + accessor = new VortexArrowColumnVector.LargeBinaryAccessor(largeVarBinaryVector); + } else if (vector instanceof ViewVarBinaryVector viewVarBinaryVector) { + accessor = new VortexArrowColumnVector.BinaryViewAccessor(viewVarBinaryVector); + } else if (vector instanceof DateDayVector dateDayVector) { + accessor = new VortexArrowColumnVector.DateAccessor(dateDayVector); + } else if (vector instanceof TimeStampVector timeStampVector) { + // Covers all eight unit/timezone variants; values are normalized to microseconds. + accessor = new VortexArrowColumnVector.TimestampAccessor(timeStampVector); + } else if (vector instanceof MapVector mapVector) { + // MapVector extends ListVector, so this check must come first. accessor = new VortexArrowColumnVector.MapAccessor(mapVector); - } else if (vector instanceof ListVector) { - ListVector listVector = (ListVector) vector; + } else if (vector instanceof ListVector listVector) { accessor = new VortexArrowColumnVector.ArrayAccessor(listVector); - } else if (vector instanceof StructVector) { - StructVector structVector = (StructVector) vector; + } else if (vector instanceof ListViewVector listViewVector) { + accessor = new VortexArrowColumnVector.ListViewAccessor(listViewVector); + } else if (vector instanceof StructVector structVector) { accessor = new VortexArrowColumnVector.StructAccessor(structVector); childColumns = new VortexArrowColumnVector[structVector.size()]; for (int i = 0; i < childColumns.length; ++i) { childColumns[i] = new VortexArrowColumnVector(structVector.getVectorById(i)); } - } else if (vector instanceof NullVector) { - accessor = new VortexArrowColumnVector.NullAccessor((NullVector) vector); - } else if (vector instanceof IntervalYearVector) { - accessor = new VortexArrowColumnVector.IntervalYearAccessor((IntervalYearVector) vector); - } else if (vector instanceof DurationVector) { - accessor = new VortexArrowColumnVector.DurationAccessor((DurationVector) vector); + } else if (vector instanceof NullVector nullVector) { + accessor = new VortexArrowColumnVector.NullAccessor(nullVector); + } else if (vector instanceof IntervalYearVector intervalYearVector) { + accessor = new VortexArrowColumnVector.IntervalYearAccessor(intervalYearVector); + } else if (vector instanceof IntervalMonthDayNanoVector intervalMonthDayNanoVector) { + accessor = new VortexArrowColumnVector.IntervalMonthDayNanoAccessor(intervalMonthDayNanoVector); + // CalendarInterval values are read through ColumnVector.getInterval, which uses the + // months/days/microseconds child column protocol. + childColumns = new VortexArrowColumnVector[] { + withAccessor(DataTypes.IntegerType, new IntervalMonthsAccessor(intervalMonthDayNanoVector)), + withAccessor(DataTypes.IntegerType, new IntervalDaysAccessor(intervalMonthDayNanoVector)), + withAccessor(DataTypes.LongType, new IntervalMicrosAccessor(intervalMonthDayNanoVector)), + }; + } else if (vector instanceof DurationVector durationVector) { + accessor = new VortexArrowColumnVector.DurationAccessor(durationVector); } else { - throw new UnsupportedOperationException( - "Unsupported Arrow vector type: " + vector.getClass().getName()); + throw new UnsupportedOperationException("Unsupported Arrow vector type: " + + vector.getClass().getSimpleName() + " for field " + vector.getField()); } } @@ -579,6 +607,24 @@ final UTF8String getUTF8String(int rowId) { } } + @Open + static class StringViewAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { + + private final ViewVarCharVector accessor; + + StringViewAccessor(ViewVarCharVector vector) { + super(vector); + this.accessor = vector; + } + + @Override + final UTF8String getUTF8String(int rowId) { + // View values may be inlined in the view buffer rather than contiguous in a data + // buffer, so copy out rather than aliasing vector memory. + return UTF8String.fromBytes(accessor.get(rowId)); + } + } + @Open static class BinaryAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { @@ -611,6 +657,22 @@ final byte[] getBinary(int rowId) { } } + @Open + static class BinaryViewAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { + + private final ViewVarBinaryVector accessor; + + BinaryViewAccessor(ViewVarBinaryVector vector) { + super(vector); + this.accessor = vector; + } + + @Override + final byte[] getBinary(int rowId) { + return accessor.getObject(rowId); + } + } + @Open static class DateAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { @@ -627,45 +689,63 @@ final int getInt(int rowId) { } } + /** + * Reads any of the eight timestamp vector variants (four units, with or without timezone), + * normalizing values to the microseconds Spark expects for TimestampType and TimestampNTZType. + * Seconds and milliseconds fail on overflow rather than silently wrapping; nanoseconds floor + * towards negative infinity, matching Spark's nanosecond-to-microsecond conversions. + */ @Open static class TimestampAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { - private final TimeStampMicroTZVector accessor; + private final TimeStampVector accessor; + private final TimeUnit unit; - TimestampAccessor(TimeStampMicroTZVector vector) { + TimestampAccessor(TimeStampVector vector) { super(vector); this.accessor = vector; + this.unit = ((ArrowType.Timestamp) vector.getField().getType()).getUnit(); } @Override final long getLong(int rowId) { - return accessor.get(rowId); + long value = accessor.get(rowId); + return switch (unit) { + case SECOND -> Math.multiplyExact(value, 1_000_000L); + case MILLISECOND -> Math.multiplyExact(value, 1_000L); + case MICROSECOND -> value; + case NANOSECOND -> Math.floorDiv(value, 1_000L); + }; } } @Open - static class TimestampNTZAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { + static class ArrayAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { - private final TimeStampMicroVector accessor; + private final ListVector accessor; + private final VortexArrowColumnVector arrayData; - TimestampNTZAccessor(TimeStampMicroVector vector) { + ArrayAccessor(ListVector vector) { super(vector); this.accessor = vector; + this.arrayData = new VortexArrowColumnVector(vector.getDataVector()); } @Override - final long getLong(int rowId) { - return accessor.get(rowId); + final ColumnarArray getArray(int rowId) { + int start = accessor.getElementStartIndex(rowId); + int end = accessor.getElementEndIndex(rowId); + return new ColumnarArray(arrayData, start, end - start); } } @Open - static class ArrayAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { + static class ListViewAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { - private final ListVector accessor; + private final ListViewVector accessor; private final VortexArrowColumnVector arrayData; - ArrayAccessor(ListVector vector) { + ListViewAccessor(ListViewVector vector) { super(vector); this.accessor = vector; this.arrayData = new VortexArrowColumnVector(vector.getDataVector()); @@ -741,6 +821,68 @@ int getInt(int rowId) { } } + /** + * Accessor for the {@code MONTH_DAY_NANO} interval vector itself. Values are read through + * {@link ColumnVector#getInterval(int)}, which uses the months/days/microseconds child + * columns; this accessor only provides null tracking. + */ + @Open + static class IntervalMonthDayNanoAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { + + IntervalMonthDayNanoAccessor(IntervalMonthDayNanoVector vector) { + super(vector); + } + } + + @Open + static class IntervalMonthsAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { + + private final IntervalMonthDayNanoVector accessor; + + IntervalMonthsAccessor(IntervalMonthDayNanoVector vector) { + super(vector); + this.accessor = vector; + } + + @Override + final int getInt(int rowId) { + return IntervalMonthDayNanoVector.getMonths(accessor.getDataBuffer(), rowId); + } + } + + @Open + static class IntervalDaysAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { + + private final IntervalMonthDayNanoVector accessor; + + IntervalDaysAccessor(IntervalMonthDayNanoVector vector) { + super(vector); + this.accessor = vector; + } + + @Override + final int getInt(int rowId) { + return IntervalMonthDayNanoVector.getDays(accessor.getDataBuffer(), rowId); + } + } + + @Open + static class IntervalMicrosAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { + + private final IntervalMonthDayNanoVector accessor; + + IntervalMicrosAccessor(IntervalMonthDayNanoVector vector) { + super(vector); + this.accessor = vector; + } + + @Override + final long getLong(int rowId) { + // Truncating division matches Spark's ArrowColumnVector month-day-nano handling. + return IntervalMonthDayNanoVector.getNanoseconds(accessor.getDataBuffer(), rowId) / 1_000L; + } + } + @Open static class DurationAccessor extends VortexArrowColumnVector.ArrowVectorAccessor { diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/ArrowUtilsTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/ArrowUtilsTest.java index 881634e6630..50c74eb41f8 100644 --- a/java/vortex-spark/src/test/java/dev/vortex/spark/ArrowUtilsTest.java +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/ArrowUtilsTest.java @@ -8,7 +8,9 @@ import dev.vortex.relocated.org.apache.arrow.vector.types.DateUnit; import dev.vortex.relocated.org.apache.arrow.vector.types.FloatingPointPrecision; +import dev.vortex.relocated.org.apache.arrow.vector.types.IntervalUnit; import dev.vortex.relocated.org.apache.arrow.vector.types.TimeUnit; +import dev.vortex.relocated.org.apache.arrow.vector.types.UnionMode; import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.ArrowType; import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.Field; import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.FieldType; @@ -52,23 +54,26 @@ void floatingPointMapsByPrecision() { } @Test - @DisplayName("Decimal preserves precision and scale") + @DisplayName("Decimal preserves precision and scale regardless of bit width") void decimalPreservesPrecisionAndScale() { assertEquals(DataTypes.createDecimalType(20, 4), ArrowUtils.fromArrowType(new ArrowType.Decimal(20, 4, 128))); + assertEquals(DataTypes.createDecimalType(20, 4), ArrowUtils.fromArrowType(new ArrowType.Decimal(20, 4, 256))); } @Test - @DisplayName("Utf8 and LargeUtf8 map to StringType") + @DisplayName("Utf8, LargeUtf8 and Utf8View map to StringType") void utf8MapsToString() { assertEquals(DataTypes.StringType, ArrowUtils.fromArrowType(new ArrowType.Utf8())); assertEquals(DataTypes.StringType, ArrowUtils.fromArrowType(new ArrowType.LargeUtf8())); + assertEquals(DataTypes.StringType, ArrowUtils.fromArrowType(new ArrowType.Utf8View())); } @Test - @DisplayName("Binary and LargeBinary map to BinaryType") + @DisplayName("Binary, LargeBinary and BinaryView map to BinaryType") void binaryMapsToBinary() { assertEquals(DataTypes.BinaryType, ArrowUtils.fromArrowType(new ArrowType.Binary())); assertEquals(DataTypes.BinaryType, ArrowUtils.fromArrowType(new ArrowType.LargeBinary())); + assertEquals(DataTypes.BinaryType, ArrowUtils.fromArrowType(new ArrowType.BinaryView())); } @Test @@ -78,14 +83,13 @@ void dateDayMapsToDate() { } @Test - @DisplayName("Timestamp(MICROSECOND) maps to Timestamp with tz, TimestampNTZ without") + @DisplayName("Timestamps of every unit map to Timestamp with tz, TimestampNTZ without") void timestampMapsByTimezonePresence() { - assertEquals( - DataTypes.TimestampType, - ArrowUtils.fromArrowType(new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"))); - assertEquals( - DataTypes.TimestampNTZType, - ArrowUtils.fromArrowType(new ArrowType.Timestamp(TimeUnit.MICROSECOND, null))); + for (TimeUnit unit : + new TimeUnit[] {TimeUnit.SECOND, TimeUnit.MILLISECOND, TimeUnit.MICROSECOND, TimeUnit.NANOSECOND}) { + assertEquals(DataTypes.TimestampType, ArrowUtils.fromArrowType(new ArrowType.Timestamp(unit, "UTC"))); + assertEquals(DataTypes.TimestampNTZType, ArrowUtils.fromArrowType(new ArrowType.Timestamp(unit, null))); + } } @Test @@ -94,6 +98,30 @@ void nullMapsToNull() { assertEquals(DataTypes.NullType, ArrowUtils.fromArrowType(new ArrowType.Null())); } + @Test + @DisplayName("Interval(YEAR_MONTH) maps to YearMonthIntervalType") + void yearMonthIntervalMapsToYearMonthIntervalType() { + assertEquals( + DataTypes.createYearMonthIntervalType(), + ArrowUtils.fromArrowType(new ArrowType.Interval(IntervalUnit.YEAR_MONTH))); + } + + @Test + @DisplayName("Interval(MONTH_DAY_NANO) maps to CalendarIntervalType") + void monthDayNanoIntervalMapsToCalendarIntervalType() { + assertEquals( + DataTypes.CalendarIntervalType, + ArrowUtils.fromArrowType(new ArrowType.Interval(IntervalUnit.MONTH_DAY_NANO))); + } + + @Test + @DisplayName("Duration(MICROSECOND) maps to DayTimeIntervalType") + void microsecondDurationMapsToDayTimeIntervalType() { + assertEquals( + DataTypes.createDayTimeIntervalType(), + ArrowUtils.fromArrowType(new ArrowType.Duration(TimeUnit.MICROSECOND))); + } + @Test @DisplayName("fromArrowField builds a StructType from nested children") void structFieldBuildsStructType() { @@ -122,6 +150,35 @@ void listFieldBuildsArrayType() { assertEquals(DataTypes.createArrayType(DataTypes.IntegerType, true), ArrowUtils.fromArrowField(list)); } + @Test + @DisplayName("fromArrowField builds an ArrayType from a ListView field") + void listViewFieldBuildsArrayType() { + Field listView = new Field( + "lv", + FieldType.nullable(new ArrowType.ListView()), + List.of(new Field("element", FieldType.notNullable(new ArrowType.Utf8View()), null))); + + assertEquals(DataTypes.createArrayType(DataTypes.StringType, false), ArrowUtils.fromArrowField(listView)); + } + + @Test + @DisplayName("fromArrowField builds a MapType carrying the value's nullability") + void mapFieldBuildsMapType() { + assertEquals( + DataTypes.createMapType(DataTypes.StringType, DataTypes.LongType, true), + ArrowUtils.fromArrowField(mapField(FieldType.nullable(new ArrowType.Int(64, true))))); + assertEquals( + DataTypes.createMapType(DataTypes.StringType, DataTypes.LongType, false), + ArrowUtils.fromArrowField(mapField(FieldType.notNullable(new ArrowType.Int(64, true))))); + } + + private static Field mapField(FieldType valueType) { + Field key = new Field("key", FieldType.notNullable(new ArrowType.Utf8()), null); + Field value = new Field("value", valueType, null); + Field entries = new Field("entries", FieldType.notNullable(new ArrowType.Struct()), List.of(key, value)); + return new Field("m", FieldType.nullable(new ArrowType.Map(false)), List.of(entries)); + } + @Test @DisplayName("Unsigned integers are unsupported") void unsignedIntegerIsUnsupported() { @@ -145,18 +202,56 @@ void millisecondDateIsUnsupported() { } @Test - @DisplayName("Non-microsecond timestamp units are unsupported") - void secondTimestampIsUnsupported() { + @DisplayName("Non-microsecond duration units are unsupported") + void nonMicrosecondDurationsAreUnsupported() { + for (TimeUnit unit : new TimeUnit[] {TimeUnit.SECOND, TimeUnit.MILLISECOND, TimeUnit.NANOSECOND}) { + assertThrows( + UnsupportedOperationException.class, () -> ArrowUtils.fromArrowType(new ArrowType.Duration(unit))); + } + } + + @Test + @DisplayName("Time is unsupported: Spark's TimeType only exists in Spark 4.1+") + void timeIsUnsupported() { + assertThrows( + UnsupportedOperationException.class, + () -> ArrowUtils.fromArrowType(new ArrowType.Time(TimeUnit.NANOSECOND, 64))); + assertThrows( + UnsupportedOperationException.class, + () -> ArrowUtils.fromArrowType(new ArrowType.Time(TimeUnit.MILLISECOND, 32))); + } + + @Test + @DisplayName("Interval(DAY_TIME) is unsupported") + void dayTimeIntervalIsUnsupported() { assertThrows( UnsupportedOperationException.class, - () -> ArrowUtils.fromArrowType(new ArrowType.Timestamp(TimeUnit.SECOND, "UTC"))); + () -> ArrowUtils.fromArrowType(new ArrowType.Interval(IntervalUnit.DAY_TIME))); + } + + @Test + @DisplayName("FixedSizeBinary is unsupported") + void fixedSizeBinaryIsUnsupported() { + assertThrows( + UnsupportedOperationException.class, () -> ArrowUtils.fromArrowType(new ArrowType.FixedSizeBinary(16))); } @Test - @DisplayName("Unrecognized Arrow types raise IllegalArgumentException") - void unrecognizedTypeIsIllegalArgument() { + @DisplayName("Union is unsupported") + void unionIsUnsupported() { assertThrows( - IllegalArgumentException.class, - () -> ArrowUtils.fromArrowType(new ArrowType.Duration(TimeUnit.SECOND))); + UnsupportedOperationException.class, + () -> ArrowUtils.fromArrowType(new ArrowType.Union(UnionMode.Sparse, new int[0]))); + } + + @Test + @DisplayName("LargeList and FixedSizeList fields are unsupported") + void largeAndFixedSizeListFieldsAreUnsupported() { + Field element = new Field("element", FieldType.nullable(new ArrowType.Int(32, true)), null); + Field largeList = new Field("ll", FieldType.nullable(new ArrowType.LargeList()), List.of(element)); + Field fixedSizeList = new Field("fsl", FieldType.nullable(new ArrowType.FixedSizeList(2)), List.of(element)); + + assertThrows(UnsupportedOperationException.class, () -> ArrowUtils.fromArrowField(largeList)); + assertThrows(UnsupportedOperationException.class, () -> ArrowUtils.fromArrowField(fixedSizeList)); } } diff --git a/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexArrowColumnVectorTest.java b/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexArrowColumnVectorTest.java new file mode 100644 index 00000000000..9ae0ba86fd8 --- /dev/null +++ b/java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexArrowColumnVectorTest.java @@ -0,0 +1,523 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +package dev.vortex.spark.read; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import dev.vortex.relocated.org.apache.arrow.memory.BufferAllocator; +import dev.vortex.relocated.org.apache.arrow.memory.RootAllocator; +import dev.vortex.relocated.org.apache.arrow.vector.BigIntVector; +import dev.vortex.relocated.org.apache.arrow.vector.BitVector; +import dev.vortex.relocated.org.apache.arrow.vector.DateDayVector; +import dev.vortex.relocated.org.apache.arrow.vector.DateMilliVector; +import dev.vortex.relocated.org.apache.arrow.vector.Decimal256Vector; +import dev.vortex.relocated.org.apache.arrow.vector.DecimalVector; +import dev.vortex.relocated.org.apache.arrow.vector.DurationVector; +import dev.vortex.relocated.org.apache.arrow.vector.Float2Vector; +import dev.vortex.relocated.org.apache.arrow.vector.Float4Vector; +import dev.vortex.relocated.org.apache.arrow.vector.Float8Vector; +import dev.vortex.relocated.org.apache.arrow.vector.IntVector; +import dev.vortex.relocated.org.apache.arrow.vector.IntervalMonthDayNanoVector; +import dev.vortex.relocated.org.apache.arrow.vector.IntervalYearVector; +import dev.vortex.relocated.org.apache.arrow.vector.LargeVarBinaryVector; +import dev.vortex.relocated.org.apache.arrow.vector.LargeVarCharVector; +import dev.vortex.relocated.org.apache.arrow.vector.NullVector; +import dev.vortex.relocated.org.apache.arrow.vector.SmallIntVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeMicroVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampMicroTZVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampMicroVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampMilliTZVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampMilliVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampNanoTZVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampNanoVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampSecTZVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampSecVector; +import dev.vortex.relocated.org.apache.arrow.vector.TimeStampVector; +import dev.vortex.relocated.org.apache.arrow.vector.TinyIntVector; +import dev.vortex.relocated.org.apache.arrow.vector.UInt4Vector; +import dev.vortex.relocated.org.apache.arrow.vector.VarBinaryVector; +import dev.vortex.relocated.org.apache.arrow.vector.VarCharVector; +import dev.vortex.relocated.org.apache.arrow.vector.ViewVarBinaryVector; +import dev.vortex.relocated.org.apache.arrow.vector.ViewVarCharVector; +import dev.vortex.relocated.org.apache.arrow.vector.complex.ListVector; +import dev.vortex.relocated.org.apache.arrow.vector.complex.ListViewVector; +import dev.vortex.relocated.org.apache.arrow.vector.complex.MapVector; +import dev.vortex.relocated.org.apache.arrow.vector.complex.StructVector; +import dev.vortex.relocated.org.apache.arrow.vector.complex.impl.NullableStructWriter; +import dev.vortex.relocated.org.apache.arrow.vector.complex.impl.UnionListViewWriter; +import dev.vortex.relocated.org.apache.arrow.vector.complex.impl.UnionListWriter; +import dev.vortex.relocated.org.apache.arrow.vector.complex.impl.UnionMapWriter; +import dev.vortex.relocated.org.apache.arrow.vector.types.TimeUnit; +import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.ArrowType; +import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.FieldType; +import dev.vortex.relocated.org.apache.arrow.vector.util.Text; +import java.math.BigDecimal; +import java.nio.charset.StandardCharsets; +import org.apache.spark.sql.types.ArrayType; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.MapType; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.vectorized.ColumnarArray; +import org.apache.spark.sql.vectorized.ColumnarMap; +import org.apache.spark.unsafe.types.CalendarInterval; +import org.apache.spark.unsafe.types.UTF8String; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for {@link VortexArrowColumnVector} covering every Arrow vector type it supports (Spark 4.1's own + * ArrowColumnVector set plus the Arrow view types): the Spark type mapping, the value conversion, and null handling. + */ +final class VortexArrowColumnVectorTest { + + private static final BufferAllocator ALLOCATOR = new RootAllocator(); + + @AfterAll + static void closeAllocator() { + ALLOCATOR.close(); + } + + @Test + @DisplayName("BitVector maps to BooleanType") + void booleanVector() { + try (BitVector vector = new BitVector("bool", ALLOCATOR)) { + vector.allocateNew(3); + vector.set(0, 1); + vector.set(2, 0); + vector.setValueCount(3); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(DataTypes.BooleanType, column.dataType()); + assertTrue(column.getBoolean(0)); + assertTrue(column.isNullAt(1)); + assertFalse(column.getBoolean(2)); + assertTrue(column.hasNull()); + assertEquals(1, column.numNulls()); + } + } + + @Test + @DisplayName("Signed integer vectors map to Byte/Short/Integer/LongType") + void signedIntegers() { + try (TinyIntVector i8 = new TinyIntVector("i8", ALLOCATOR); + SmallIntVector i16 = new SmallIntVector("i16", ALLOCATOR); + IntVector i32 = new IntVector("i32", ALLOCATOR); + BigIntVector i64 = new BigIntVector("i64", ALLOCATOR)) { + i8.allocateNew(2); + i8.set(0, Byte.MIN_VALUE); + i8.setValueCount(2); + i16.allocateNew(2); + i16.set(0, Short.MIN_VALUE); + i16.setValueCount(2); + i32.allocateNew(2); + i32.set(0, Integer.MIN_VALUE); + i32.setValueCount(2); + i64.allocateNew(2); + i64.set(0, Long.MIN_VALUE); + i64.setValueCount(2); + + VortexArrowColumnVector byteColumn = new VortexArrowColumnVector(i8); + assertEquals(DataTypes.ByteType, byteColumn.dataType()); + assertEquals(Byte.MIN_VALUE, byteColumn.getByte(0)); + assertTrue(byteColumn.isNullAt(1)); + + VortexArrowColumnVector shortColumn = new VortexArrowColumnVector(i16); + assertEquals(DataTypes.ShortType, shortColumn.dataType()); + assertEquals(Short.MIN_VALUE, shortColumn.getShort(0)); + + VortexArrowColumnVector intColumn = new VortexArrowColumnVector(i32); + assertEquals(DataTypes.IntegerType, intColumn.dataType()); + assertEquals(Integer.MIN_VALUE, intColumn.getInt(0)); + + VortexArrowColumnVector longColumn = new VortexArrowColumnVector(i64); + assertEquals(DataTypes.LongType, longColumn.dataType()); + assertEquals(Long.MIN_VALUE, longColumn.getLong(0)); + } + } + + @Test + @DisplayName("Float vectors map to Float/DoubleType") + void floats() { + try (Float4Vector f32 = new Float4Vector("f32", ALLOCATOR); + Float8Vector f64 = new Float8Vector("f64", ALLOCATOR)) { + f32.allocateNew(2); + f32.set(0, 2.5f); + f32.setValueCount(2); + f64.allocateNew(2); + f64.set(0, 3.5d); + f64.setValueCount(2); + + VortexArrowColumnVector floatColumn = new VortexArrowColumnVector(f32); + assertEquals(DataTypes.FloatType, floatColumn.dataType()); + assertEquals(2.5f, floatColumn.getFloat(0)); + assertTrue(floatColumn.isNullAt(1)); + + VortexArrowColumnVector doubleColumn = new VortexArrowColumnVector(f64); + assertEquals(DataTypes.DoubleType, doubleColumn.dataType()); + assertEquals(3.5d, doubleColumn.getDouble(0)); + } + } + + @Test + @DisplayName("Decimal128 vectors map to DecimalType") + void decimal() { + try (DecimalVector d128 = new DecimalVector("d128", ALLOCATOR, 10, 2)) { + d128.allocateNew(2); + d128.set(0, new BigDecimal("12345678.90")); + d128.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(d128); + assertEquals(DataTypes.createDecimalType(10, 2), column.dataType()); + assertEquals( + new BigDecimal("12345678.90"), column.getDecimal(0, 10, 2).toJavaBigDecimal()); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("String vectors (regular, large, view) map to StringType") + void strings() { + try (VarCharVector utf8 = new VarCharVector("utf8", ALLOCATOR); + LargeVarCharVector largeUtf8 = new LargeVarCharVector("large_utf8", ALLOCATOR); + ViewVarCharVector utf8View = new ViewVarCharVector("utf8_view", ALLOCATOR)) { + utf8.allocateNew(2); + utf8.setSafe(0, "hello".getBytes(StandardCharsets.UTF_8)); + utf8.setValueCount(2); + largeUtf8.allocateNew(2); + largeUtf8.setSafe(0, "world".getBytes(StandardCharsets.UTF_8)); + largeUtf8.setValueCount(2); + utf8View.allocateNew(2); + utf8View.setSafe(0, new Text("a string long enough to not be inlined")); + utf8View.setValueCount(2); + + VortexArrowColumnVector utf8Column = new VortexArrowColumnVector(utf8); + assertEquals(DataTypes.StringType, utf8Column.dataType()); + assertEquals(UTF8String.fromString("hello"), utf8Column.getUTF8String(0)); + assertTrue(utf8Column.isNullAt(1)); + + VortexArrowColumnVector largeUtf8Column = new VortexArrowColumnVector(largeUtf8); + assertEquals(DataTypes.StringType, largeUtf8Column.dataType()); + assertEquals(UTF8String.fromString("world"), largeUtf8Column.getUTF8String(0)); + + VortexArrowColumnVector utf8ViewColumn = new VortexArrowColumnVector(utf8View); + assertEquals(DataTypes.StringType, utf8ViewColumn.dataType()); + assertEquals( + UTF8String.fromString("a string long enough to not be inlined"), utf8ViewColumn.getUTF8String(0)); + assertTrue(utf8ViewColumn.isNullAt(1)); + } + } + + @Test + @DisplayName("Binary vectors (regular, large, view) map to BinaryType") + void binary() { + byte[] payload = new byte[] {1, 2, 3, 4}; + try (VarBinaryVector bin = new VarBinaryVector("bin", ALLOCATOR); + LargeVarBinaryVector largeBin = new LargeVarBinaryVector("large_bin", ALLOCATOR); + ViewVarBinaryVector binView = new ViewVarBinaryVector("bin_view", ALLOCATOR)) { + bin.allocateNew(2); + bin.setSafe(0, payload); + bin.setValueCount(2); + largeBin.allocateNew(2); + largeBin.setSafe(0, payload); + largeBin.setValueCount(2); + binView.allocateNew(2); + binView.setSafe(0, payload); + binView.setValueCount(2); + + for (VortexArrowColumnVector column : new VortexArrowColumnVector[] { + new VortexArrowColumnVector(bin), + new VortexArrowColumnVector(largeBin), + new VortexArrowColumnVector(binView), + }) { + assertEquals(DataTypes.BinaryType, column.dataType()); + assertArrayEquals(payload, column.getBinary(0)); + assertTrue(column.isNullAt(1)); + } + } + } + + @Test + @DisplayName("Day-unit date vectors map to DateType") + void date() { + try (DateDayVector vector = new DateDayVector("date_day", ALLOCATOR)) { + vector.allocateNew(2); + vector.set(0, 19000); + vector.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(DataTypes.DateType, column.dataType()); + assertEquals(19000, column.getInt(0)); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("Timestamp vectors of every unit normalize to microseconds") + void timestampsWithoutTimezone() { + try (TimeStampSecVector sec = new TimeStampSecVector("ts_s", ALLOCATOR); + TimeStampMilliVector milli = new TimeStampMilliVector("ts_ms", ALLOCATOR); + TimeStampMicroVector micro = new TimeStampMicroVector("ts_us", ALLOCATOR); + TimeStampNanoVector nano = new TimeStampNanoVector("ts_ns", ALLOCATOR)) { + sec.allocateNew(3); + sec.set(0, 1_700_000_000L); + sec.setValueCount(3); + milli.allocateNew(3); + milli.set(0, 1_700_000_000_123L); + milli.setValueCount(3); + micro.allocateNew(3); + micro.set(0, 1_700_000_000_123_456L); + micro.setValueCount(3); + nano.allocateNew(3); + nano.set(0, 1_700_000_000_123_456_789L); + // Negative nanos floor towards negative infinity when reduced to micros. + nano.set(2, -1_500L); + nano.setValueCount(3); + + assertTimestamp(sec, DataTypes.TimestampNTZType, 1_700_000_000_000_000L); + assertTimestamp(milli, DataTypes.TimestampNTZType, 1_700_000_000_123_000L); + assertTimestamp(micro, DataTypes.TimestampNTZType, 1_700_000_000_123_456L); + VortexArrowColumnVector nanoColumn = + assertTimestamp(nano, DataTypes.TimestampNTZType, 1_700_000_000_123_456L); + assertEquals(-2L, nanoColumn.getLong(2)); + } + } + + @Test + @DisplayName("Timezone-aware timestamp vectors of every unit map to TimestampType") + void timestampsWithTimezone() { + try (TimeStampSecTZVector sec = new TimeStampSecTZVector("ts_s", ALLOCATOR, "UTC"); + TimeStampMilliTZVector milli = new TimeStampMilliTZVector("ts_ms", ALLOCATOR, "UTC"); + TimeStampMicroTZVector micro = new TimeStampMicroTZVector("ts_us", ALLOCATOR, "UTC"); + TimeStampNanoTZVector nano = new TimeStampNanoTZVector("ts_ns", ALLOCATOR, "UTC")) { + sec.allocateNew(1); + sec.set(0, 1_700_000_000L); + sec.setValueCount(1); + milli.allocateNew(1); + milli.set(0, 1_700_000_000_123L); + milli.setValueCount(1); + micro.allocateNew(1); + micro.set(0, 1_700_000_000_123_456L); + micro.setValueCount(1); + nano.allocateNew(1); + nano.set(0, 1_700_000_000_123_456_789L); + nano.setValueCount(1); + + assertTimestamp(sec, DataTypes.TimestampType, 1_700_000_000_000_000L); + assertTimestamp(milli, DataTypes.TimestampType, 1_700_000_000_123_000L); + assertTimestamp(micro, DataTypes.TimestampType, 1_700_000_000_123_456L); + assertTimestamp(nano, DataTypes.TimestampType, 1_700_000_000_123_456L); + } + } + + private static VortexArrowColumnVector assertTimestamp( + TimeStampVector vector, org.apache.spark.sql.types.DataType expectedType, long expectedMicros) { + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(expectedType, column.dataType()); + assertEquals(expectedMicros, column.getLong(0)); + return column; + } + + @Test + @DisplayName("Microsecond duration vectors map to DayTimeIntervalType") + void duration() { + try (DurationVector vector = new DurationVector( + "dur_us", FieldType.nullable(new ArrowType.Duration(TimeUnit.MICROSECOND)), ALLOCATOR)) { + vector.allocateNew(2); + vector.set(0, 1_500L); + vector.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(DataTypes.createDayTimeIntervalType(), column.dataType()); + assertEquals(1_500L, column.getLong(0)); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("Year-month interval vectors map to YearMonthIntervalType as months") + void intervalYear() { + try (IntervalYearVector vector = new IntervalYearVector("interval_ym", ALLOCATOR)) { + vector.allocateNew(2); + vector.set(0, 14); + vector.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(DataTypes.createYearMonthIntervalType(), column.dataType()); + assertEquals(14, column.getInt(0)); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("Month-day-nano interval vectors map to CalendarIntervalType") + void intervalMonthDayNano() { + try (IntervalMonthDayNanoVector vector = new IntervalMonthDayNanoVector("interval_mdn", ALLOCATOR)) { + vector.allocateNew(2); + vector.set(0, 1, 2, 3_500L); + vector.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(DataTypes.CalendarIntervalType, column.dataType()); + // Nanoseconds truncate to microseconds, matching Spark's ArrowColumnVector. + assertEquals(new CalendarInterval(1, 2, 3L), column.getInterval(0)); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("Null vectors map to NullType") + void nullVector() { + try (NullVector vector = new NullVector("null_col")) { + vector.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(DataTypes.NullType, column.dataType()); + assertTrue(column.isNullAt(0)); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("List vectors map to ArrayType") + void list() { + try (ListVector vector = ListVector.empty("list", ALLOCATOR)) { + UnionListWriter writer = vector.getWriter(); + writer.allocate(); + writer.setPosition(0); + writer.startList(); + writer.writeInt(1); + writer.writeInt(2); + writer.writeInt(3); + writer.endList(); + writer.setPosition(2); + writer.startList(); + writer.endList(); + vector.setValueCount(3); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(DataTypes.IntegerType, ((ArrayType) column.dataType()).elementType()); + ColumnarArray array = column.getArray(0); + assertEquals(3, array.numElements()); + assertEquals(1, array.getInt(0)); + assertEquals(3, array.getInt(2)); + assertTrue(column.isNullAt(1)); + assertEquals(0, column.getArray(2).numElements()); + } + } + + @Test + @DisplayName("List view vectors map to ArrayType") + void listView() { + try (ListViewVector vector = ListViewVector.empty("list_view", ALLOCATOR)) { + UnionListViewWriter writer = vector.getWriter(); + writer.allocate(); + writer.setPosition(0); + writer.startListView(); + writer.writeInt(7); + writer.writeInt(8); + writer.endListView(); + vector.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + assertEquals(DataTypes.IntegerType, ((ArrayType) column.dataType()).elementType()); + ColumnarArray array = column.getArray(0); + assertEquals(2, array.numElements()); + assertEquals(7, array.getInt(0)); + assertEquals(8, array.getInt(1)); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("Map vectors map to MapType") + void map() { + try (MapVector vector = MapVector.empty("map", ALLOCATOR, false)) { + UnionMapWriter writer = vector.getWriter(); + writer.allocate(); + writer.setPosition(0); + writer.startMap(); + writer.startEntry(); + writer.key().integer().writeInt(1); + writer.value().bigInt().writeBigInt(100L); + writer.endEntry(); + writer.startEntry(); + writer.key().integer().writeInt(2); + writer.value().bigInt().writeBigInt(200L); + writer.endEntry(); + writer.endMap(); + vector.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + MapType mapType = (MapType) column.dataType(); + assertEquals(DataTypes.IntegerType, mapType.keyType()); + assertEquals(DataTypes.LongType, mapType.valueType()); + + ColumnarMap columnarMap = column.getMap(0); + assertEquals(2, columnarMap.numElements()); + assertEquals(1, columnarMap.keyArray().getInt(0)); + assertEquals(100L, columnarMap.valueArray().getLong(0)); + assertEquals(2, columnarMap.keyArray().getInt(1)); + assertEquals(200L, columnarMap.valueArray().getLong(1)); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("Struct vectors map to StructType with child columns") + void struct() { + try (StructVector vector = StructVector.empty("struct", ALLOCATOR)) { + NullableStructWriter writer = vector.getWriter(); + writer.allocate(); + writer.setPosition(0); + writer.start(); + writer.integer("a").writeInt(5); + writer.bigInt("b").writeBigInt(7L); + writer.end(); + vector.setValueCount(2); + + VortexArrowColumnVector column = new VortexArrowColumnVector(vector); + StructType structType = (StructType) column.dataType(); + assertEquals(2, structType.fields().length); + assertEquals(DataTypes.IntegerType, structType.fields()[0].dataType()); + assertEquals(DataTypes.LongType, structType.fields()[1].dataType()); + assertEquals(5, column.getChild(0).getInt(0)); + assertEquals(7L, column.getChild(1).getLong(0)); + assertTrue(column.isNullAt(1)); + } + } + + @Test + @DisplayName("Arrow types outside the supported set are rejected with descriptive errors") + void unsupportedTypes() { + try (UInt4Vector unsignedInt = new UInt4Vector("u32", ALLOCATOR); + Float2Vector halfFloat = new Float2Vector("f16", ALLOCATOR); + Decimal256Vector decimal256 = new Decimal256Vector("d256", ALLOCATOR, 38, 2); + DateMilliVector dateMilli = new DateMilliVector("date_ms", ALLOCATOR); + TimeMicroVector time = new TimeMicroVector("time_us", ALLOCATOR); + DurationVector durationSec = new DurationVector( + "dur_s", FieldType.nullable(new ArrowType.Duration(TimeUnit.SECOND)), ALLOCATOR)) { + assertUnsupported(unsignedInt, "unsigned"); + assertUnsupported(halfFloat, "float precision"); + // Decimal256 maps to DecimalType but has no accessor, matching Spark's ArrowColumnVector. + assertUnsupported(decimal256, "Decimal256Vector"); + assertUnsupported(dateMilli, "date unit"); + assertUnsupported(time, "Time"); + assertUnsupported(durationSec, "duration unit"); + } + } + + private static void assertUnsupported( + dev.vortex.relocated.org.apache.arrow.vector.ValueVector vector, String expectedMessagePart) { + UnsupportedOperationException e = + assertThrows(UnsupportedOperationException.class, () -> new VortexArrowColumnVector(vector)); + assertTrue( + e.getMessage().contains(expectedMessagePart), + "expected error message to contain \"" + expectedMessagePart + "\" but was: " + e.getMessage()); + } +} diff --git a/vortex-jni/src/dtype.rs b/vortex-jni/src/dtype.rs index 8969d7b573a..1014ae07613 100644 --- a/vortex-jni/src/dtype.rs +++ b/vortex-jni/src/dtype.rs @@ -7,70 +7,23 @@ use std::ptr; use arrow_array::ffi::FFI_ArrowSchema; -use arrow_schema::DataType; -use arrow_schema::FieldRef; -use arrow_schema::Fields; use arrow_schema::Schema; use vortex::dtype::DType; use vortex::error::VortexResult; use vortex_arrow::ToArrowType; -/// Export a Vortex [`DType`] to the Arrow C Data Interface struct at `schema_addr`. Views -/// (Utf8View/BinaryView) are downgraded to regular Utf8/Binary so Spark and other consumers -/// without view support can read them. +/// Export a Vortex [`DType`] to the Arrow C Data Interface struct at `schema_addr`. String and +/// binary columns are exported as their native view types (Utf8View/BinaryView); consumers are +/// expected to handle them. pub(crate) fn export_dtype_to_arrow(dtype: &DType, schema_addr: i64) -> VortexResult<()> { let arrow_schema = dtype.to_arrow_schema()?; - let viewless = strip_views(DataType::Struct(arrow_schema.fields().clone())); - let fields = match viewless { - DataType::Struct(fields) => fields, - _ => unreachable!("Vortex DType always exports as a struct"), - }; - let schema = Schema::new(fields); - let ffi_schema = FFI_ArrowSchema::try_from(&schema)?; + let ffi_schema = FFI_ArrowSchema::try_from(&arrow_schema)?; unsafe { ptr::write(schema_addr as *mut FFI_ArrowSchema, ffi_schema); } Ok(()) } -/// Replace view-based Arrow types with their non-view counterparts throughout the tree. -pub(crate) fn strip_views(data_type: DataType) -> DataType { - match data_type { - DataType::BinaryView => DataType::Binary, - DataType::Utf8View => DataType::Utf8, - DataType::List(inner) | DataType::ListView(inner) => { - let new_inner = (*inner) - .clone() - .with_data_type(strip_views(inner.data_type().clone())); - DataType::List(FieldRef::new(new_inner)) - } - DataType::LargeList(inner) | DataType::LargeListView(inner) => { - let new_inner = (*inner) - .clone() - .with_data_type(strip_views(inner.data_type().clone())); - DataType::LargeList(FieldRef::new(new_inner)) - } - DataType::Struct(fields) => { - let viewless_fields: Vec = fields - .iter() - .map(|field_ref| { - let field = (**field_ref).clone(); - let data_type = field.data_type().clone(); - FieldRef::new(field.with_data_type(strip_views(data_type))) - }) - .collect(); - DataType::Struct(Fields::from(viewless_fields)) - } - DataType::FixedSizeList(inner, size) => { - let new_inner = (*inner) - .clone() - .with_data_type(strip_views(inner.data_type().clone())); - DataType::FixedSizeList(FieldRef::new(new_inner), size) - } - dt => dt, - } -} - /// Decode an [`FFI_ArrowSchema`] pointed to by `schema_addr` into an Arrow [`Schema`]. pub(crate) fn import_arrow_schema(schema_addr: i64) -> VortexResult { let ffi_schema = unsafe { &*(schema_addr as *const FFI_ArrowSchema) }; diff --git a/vortex-jni/src/scan.rs b/vortex-jni/src/scan.rs index 87819681686..4b067706eba 100644 --- a/vortex-jni/src/scan.rs +++ b/vortex-jni/src/scan.rs @@ -19,7 +19,6 @@ use arrow_array::RecordBatch; use arrow_array::cast::AsArray; use arrow_array::ffi_stream::FFI_ArrowArrayStream; use arrow_schema::ArrowError; -use arrow_schema::DataType; use arrow_schema::Field; use futures::StreamExt; use jni::EnvUnowned; @@ -49,7 +48,7 @@ use vortex_arrow::ToArrowType; use crate::POOL; use crate::RUNTIME; use crate::data_source::NativeDataSource; -use crate::dtype::strip_views; +use crate::dtype::export_dtype_to_arrow; use crate::errors::try_or_throw; use crate::session::session_ref; @@ -216,7 +215,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeScan_arrowSchema( let NativeScan::Pending(scan) = scan else { throw_runtime!("schema unavailable: scan already started"); }; - crate::dtype::export_dtype_to_arrow(scan.dtype(), schema_addr)?; + export_dtype_to_arrow(scan.dtype(), schema_addr)?; Ok(()) }); } @@ -344,13 +343,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativePartition_scanArrow( let array_stream = partition.execute()?; let dtype = array_stream.dtype().clone(); - let raw_schema = dtype.to_arrow_schema()?; - let viewless = strip_views(DataType::Struct(raw_schema.fields().clone())); - let fields = match viewless { - DataType::Struct(fields) => fields, - _ => unreachable!("Vortex DType always exports as a struct"), - }; - let schema = Arc::new(arrow_schema::Schema::new(fields)); + let schema = Arc::new(dtype.to_arrow_schema()?); let target = Arc::new(Field::new_struct("", schema.fields().clone(), false)); let session = unsafe { session_ref(session_ptr) };