diff --git a/sql/api/src/main/scala/org/apache/spark/sql/util/ArrowUtils.scala b/sql/api/src/main/scala/org/apache/spark/sql/util/ArrowUtils.scala index e69a0fa7f4154..b8dab164cf780 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/util/ArrowUtils.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/util/ArrowUtils.scala @@ -94,6 +94,8 @@ private[sql] object ArrowUtils { case ArrowType.Binary.INSTANCE => BinaryType case ArrowType.LargeUtf8.INSTANCE => StringType case ArrowType.LargeBinary.INSTANCE => BinaryType + case ArrowType.Utf8View.INSTANCE => StringType + case ArrowType.BinaryView.INSTANCE => BinaryType case d: ArrowType.Decimal => DecimalType(d.getPrecision, d.getScale) case date: ArrowType.Date if date.getUnit == DateUnit.DAY => DateType case ts: ArrowType.Timestamp @@ -516,7 +518,7 @@ private[sql] object ArrowUtils { val keyType = fromArrowField(elementField.getChildren.get(0)) val valueType = fromArrowField(elementField.getChildren.get(1)) MapType(keyType, valueType, elementField.getChildren.get(1).isNullable) - case ArrowType.List.INSTANCE => + case ArrowType.List.INSTANCE | ArrowType.ListView.INSTANCE => val elementField = field.getChildren().get(0) val elementType = fromArrowField(elementField) ArrayType(elementType, containsNull = elementField.isNullable) diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ArrowColumnVector.java b/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ArrowColumnVector.java index f44de5ffa9df7..d6e2f3f2c12e6 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ArrowColumnVector.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ArrowColumnVector.java @@ -17,6 +17,7 @@ package org.apache.spark.sql.vectorized; +import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.*; import org.apache.arrow.vector.complex.*; import org.apache.arrow.vector.holders.NullableIntervalMonthDayNanoHolder; @@ -209,7 +210,11 @@ void initAccessor(ValueVector vector) { } else if (vector instanceof VarBinaryVector varBinaryVector) { accessor = new BinaryAccessor(varBinaryVector); } else if (vector instanceof LargeVarBinaryVector largeVarBinaryVector) { - accessor = new LargeBinaryAccessor(largeVarBinaryVector); + accessor = new BinaryAccessor(largeVarBinaryVector); + } else if (vector instanceof ViewVarCharVector viewVarCharVector) { + accessor = new StringViewAccessor(viewVarCharVector); + } else if (vector instanceof ViewVarBinaryVector viewVarBinaryVector) { + accessor = new BinaryAccessor(viewVarBinaryVector); } else if (vector instanceof DateDayVector dateDayVector) { accessor = new DateAccessor(dateDayVector); } else if (vector instanceof TimeStampMicroTZVector timeStampMicroTZVector) { @@ -225,7 +230,9 @@ void initAccessor(ValueVector vector) { } else if (vector instanceof MapVector mapVector) { accessor = new MapAccessor(mapVector); } else if (vector instanceof ListVector listVector) { - accessor = new ArrayAccessor(listVector); + accessor = new ArrayAccessor<>(listVector); + } else if (vector instanceof ListViewVector listViewVector) { + accessor = new ArrayAccessor<>(listViewVector); } else if (vector instanceof StructVector structVector) { if (ArrowUtils.isTimestampNanosStructField(structVector.getField())) { // Lossless struct representation of a nanosecond timestamp (ArrowUtils.toArrowField with @@ -500,33 +507,59 @@ final UTF8String getUTF8String(int rowId) { } } + // Covers VarBinaryVector, LargeVarBinaryVector and ViewVarBinaryVector: all of them return + // the value as a byte array from getObject, with a null check. static class BinaryAccessor extends ArrowVectorAccessor { - private final VarBinaryVector accessor; + private final VariableWidthFieldVector accessor; - BinaryAccessor(VarBinaryVector vector) { + BinaryAccessor(VariableWidthFieldVector vector) { super(vector); this.accessor = vector; } @Override final byte[] getBinary(int rowId) { - return accessor.getObject(rowId); + return (byte[]) accessor.getObject(rowId); } } - static class LargeBinaryAccessor extends ArrowVectorAccessor { + static class StringViewAccessor extends ArrowVectorAccessor { - private final LargeVarBinaryVector accessor; + private final ViewVarCharVector accessor; - LargeBinaryAccessor(LargeVarBinaryVector vector) { + StringViewAccessor(ViewVarCharVector vector) { super(vector); this.accessor = vector; } @Override - final byte[] getBinary(int rowId) { - return accessor.getObject(rowId); + final UTF8String getUTF8String(int rowId) { + if (accessor.isNull(rowId)) { + return null; + } + // Decode the 16-byte view struct directly rather than through Arrow's + // NullableViewVarCharHolder: the holder's int start/end fields truncate the inline-value + // offset (rowId * 16 + 4) once the views buffer grows past Integer.MAX_VALUE bytes, which + // would silently read the wrong memory. Both branches read the value with zero copy, like + // the non-view string accessors. + ArrowBuf views = accessor.getDataBuffer(); + long viewOffset = (long) rowId * BaseVariableWidthViewVector.ELEMENT_SIZE; + int length = views.getInt(viewOffset); + if (length <= BaseVariableWidthViewVector.INLINE_SIZE) { + // Short values are stored inline in the views buffer, right after the length. + long start = viewOffset + BaseVariableWidthViewVector.LENGTH_WIDTH; + return UTF8String.fromAddress(null, views.memoryAddress() + start, length); + } else { + // Long values live in one of the variadic data buffers; the view struct holds the buffer + // index and the offset within that buffer, after the length and a 4-byte prefix. + long bufferIndexOffset = viewOffset + BaseVariableWidthViewVector.LENGTH_WIDTH + + BaseVariableWidthViewVector.PREFIX_WIDTH; + int bufferIndex = views.getInt(bufferIndexOffset); + int start = views.getInt(bufferIndexOffset + BaseVariableWidthViewVector.BUF_INDEX_WIDTH); + ArrowBuf dataBuffer = accessor.getDataBuffers().get(bufferIndex); + return UTF8String.fromAddress(null, dataBuffer.memoryAddress() + start, length); + } } } @@ -679,12 +712,16 @@ final CalendarInterval getInterval(int rowId) { } } - static class ArrayAccessor extends ArrowVectorAccessor { + // Covers ListVector and ListViewVector. ListView stores an explicit per-value offset and size + // (rather than ListVector's contiguous offsets), but both expose a value's element range in the + // data vector as [getElementStartIndex, getElementEndIndex). + static class ArrayAccessor + extends ArrowVectorAccessor { - private final ListVector accessor; + private final T accessor; private final ArrowColumnVector arrayData; - ArrayAccessor(ListVector vector) { + ArrayAccessor(T vector) { super(vector); this.accessor = vector; this.arrayData = new ArrowColumnVector(vector.getDataVector()); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala index b96e57ce49af7..f4c3bc08f403f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala @@ -131,6 +131,12 @@ object ArrowWriter { createFieldWriter(vector.getChildByOrdinal(ordinal)) } new GeographyWriter(dt, vector, children.toArray) + // The Arrow view types are readable through ArrowColumnVector but have no field writers. + // Their Spark types resolve to the non-view vector cases above, so without this case they + // would fall through to UNSUPPORTED_DATATYPE naming a fully supported Spark type; report + // the offending Arrow type instead. + case (_, vector @ (_: ViewVarCharVector | _: ViewVarBinaryVector | _: ListViewVector)) => + throw ExecutionErrors.unsupportedArrowTypeError(vector.getField.getType) case (dt, _) => throw ExecutionErrors.unsupportedDataTypeError(dt) } diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/arrow/ArrowEncoderSuite.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/arrow/ArrowEncoderSuite.scala index 9f171cda2ea23..efb6a7a2304f0 100644 --- a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/arrow/ArrowEncoderSuite.scala +++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/arrow/ArrowEncoderSuite.scala @@ -16,7 +16,7 @@ */ package org.apache.spark.sql.connect.client.arrow -import java.io.File +import java.io.{ByteArrayOutputStream, File} import java.math.BigInteger import java.net.URLClassLoader import java.time.{Duration, Period, ZoneOffset} @@ -32,7 +32,8 @@ import scala.reflect.classTag import scala.reflect.runtime.{universe => ru} import org.apache.arrow.memory.{BufferAllocator, RootAllocator} -import org.apache.arrow.vector.VarBinaryVector +import org.apache.arrow.vector.{BaseVariableWidthViewVector, FieldVector, VarBinaryVector, VectorSchemaRoot, ViewVarBinaryVector, ViewVarCharVector} +import org.apache.arrow.vector.ipc.ArrowStreamWriter import org.apache.spark.{SparkRuntimeException, SparkUnsupportedOperationException} import org.apache.spark.sql.{Encoders, Row} @@ -259,6 +260,55 @@ class ArrowEncoderSuite extends ConnectFunSuite { } } + test("deserializing string and binary view vectors") { + // The client never produces view-encoded batches itself, but it can receive them, so the + // readers must handle them. Mix short (inline, <= 12 bytes) and long (stored in a data + // buffer) values to exercise both view-storage paths. + val values = Seq("a", "a-string-longer-than-twelve-bytes", null) + + def serializeViewVector(vector: BaseVariableWidthViewVector): Array[Byte] = { + vector.allocateNew() + values.zipWithIndex.foreach { + case (null, i) => vector.setNull(i) + case (s, i) => + val bytes = s.getBytes("utf8") + vector.setSafe(i, bytes, 0, bytes.length) + } + vector.setValueCount(values.size) + val root = new VectorSchemaRoot(Collections.singletonList[FieldVector](vector)) + try { + val out = new ByteArrayOutputStream() + val writer = new ArrowStreamWriter(root, null, out) + writer.start() + writer.writeBatch() + writer.end() + out.toByteArray + } finally { + root.close() + } + } + + withAllocator { allocator => + val strings = ArrowDeserializers.deserializeFromArrow( + Iterator.single(serializeViewVector(new ViewVarCharVector("s", allocator))), + StringEncoder, + allocator, + timeZoneId = "UTC") + compareIterators(values.iterator, strings) + strings.close() + + val binaries = ArrowDeserializers.deserializeFromArrow( + Iterator.single(serializeViewVector(new ViewVarBinaryVector("b", allocator))), + BinaryEncoder, + allocator, + timeZoneId = "UTC") + compareIterators( + values.iterator.map(Option(_).map(_.getBytes("utf8").toSeq)), + binaries.map(Option(_).map(_.toSeq))) + binaries.close() + } + } + test("single batch") { val inspector = new CountingBatchInspector roundTripAndCheckIdentical(singleIntEncoder, inspectBatch = inspector) { () => diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/arrow/ArrowVectorReader.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/arrow/ArrowVectorReader.scala index 54311cecc1627..1f71ed8f84322 100644 --- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/arrow/ArrowVectorReader.scala +++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/arrow/ArrowVectorReader.scala @@ -95,8 +95,10 @@ object ArrowVectorReader { case v: DecimalVector => new DecimalVectorReader(v) case v: VarCharVector => new VarCharVectorReader(v) case v: LargeVarCharVector => new LargeVarCharVectorReader(v) + case v: ViewVarCharVector => new ViewVarCharVectorReader(v) case v: VarBinaryVector => new VarBinaryVectorReader(v) case v: LargeVarBinaryVector => new LargeVarBinaryVectorReader(v) + case v: ViewVarBinaryVector => new ViewVarBinaryVectorReader(v) case v: DurationVector => new DurationVectorReader(v) case v: IntervalYearVector => new IntervalYearVectorReader(v) case v: DateDayVector => new DateDayVectorReader(v, timeZoneId) @@ -215,6 +217,11 @@ private[arrow] class LargeVarCharVectorReader(v: LargeVarCharVector) override def getString(i: Int): String = Text.decode(vector.get(i)) } +private[arrow] class ViewVarCharVectorReader(v: ViewVarCharVector) + extends TypedArrowVectorReader[ViewVarCharVector](v) { + override def getString(i: Int): String = Text.decode(vector.get(i)) +} + private[arrow] class VarBinaryVectorReader(v: VarBinaryVector) extends TypedArrowVectorReader[VarBinaryVector](v) { override def getBytes(i: Int): Array[Byte] = vector.get(i) @@ -227,6 +234,12 @@ private[arrow] class LargeVarBinaryVectorReader(v: LargeVarBinaryVector) override def getString(i: Int): String = SparkStringUtils.getHexString(getBytes(i)) } +private[arrow] class ViewVarBinaryVectorReader(v: ViewVarBinaryVector) + extends TypedArrowVectorReader[ViewVarBinaryVector](v) { + override def getBytes(i: Int): Array[Byte] = vector.get(i) + override def getString(i: Int): String = SparkStringUtils.getHexString(getBytes(i)) +} + private[arrow] class DurationVectorReader(v: DurationVector) extends TypedArrowVectorReader[DurationVector](v) { override def getDuration(i: Int): Duration = vector.getObject(i) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowFileReadWrite.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowFileReadWrite.scala index e7ec2d2b79841..516b01b25fdc8 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowFileReadWrite.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowFileReadWrite.scala @@ -24,9 +24,10 @@ import scala.jdk.CollectionConverters._ import org.apache.arrow.vector._ import org.apache.arrow.vector.ipc.{ArrowFileReader, ArrowFileWriter} -import org.apache.arrow.vector.types.pojo.Schema +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} import org.apache.spark.sql.classic.{DataFrame, SparkSession} +import org.apache.spark.sql.errors.ExecutionErrors import org.apache.spark.sql.util.ArrowUtils private[sql] class SparkArrowFileWriter(schema: Schema, path: Path) extends AutoCloseable { @@ -97,6 +98,33 @@ private[spark] object ArrowFileReadWrite { def load(spark: SparkSession, path: Path): DataFrame = { val reader = new SparkArrowFileReader(path) val schema = ArrowUtils.fromArrowSchema(reader.schema) + // `toDataFrame` rebuilds vectors from `schema` via `toArrowSchema` and loads the file's raw + // record batches into them, so the file's physical layout must match Spark's canonical Arrow + // encoding of that schema. An Arrow type that converts to a Spark type but is encoded + // differently (e.g. Utf8View or LargeUtf8 vs Utf8) would have its buffers reinterpreted as + // the wrong layout and read back as corrupt values, so reject it up front. + val canonicalSchema = ArrowUtils.toArrowSchema( + schema, "UTC", errorOnDuplicatedFieldNames = true, largeVarTypes = false) + reader.schema.getFields.asScala.zip(canonicalSchema.getFields.asScala).foreach { + case (actual, canonical) => checkLayoutMatch(actual, canonical) + } ArrowConverters.toDataFrame(reader.read(), schema, spark, "UTC", true, false) } + + private def checkLayoutMatch(actual: Field, canonical: Field): Unit = { + val compatible = (actual.getType, canonical.getType) match { + // The timezone label does not affect the physical encoding (values are epoch-based), and + // `fromArrowSchema` already rejects the units Spark cannot read. + case (a: ArrowType.Timestamp, c: ArrowType.Timestamp) => a.getUnit == c.getUnit + // Map equality includes the keysSorted flag, which has no layout impact. + case (_: ArrowType.Map, _: ArrowType.Map) => true + case (a, c) => a == c + } + if (!compatible) { + throw ExecutionErrors.unsupportedArrowTypeError(actual.getType) + } + actual.getChildren.asScala.zip(canonical.getChildren.asScala).foreach { + case (a, c) => checkLayoutMatch(a, c) + } + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowFileReadWriteSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowFileReadWriteSuite.scala index 8accda825403f..dda9827e80a68 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowFileReadWriteSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowFileReadWriteSuite.scala @@ -17,9 +17,19 @@ package org.apache.spark.sql.execution.arrow import java.io.File +import java.nio.channels.Channels +import java.nio.file.Files +import scala.jdk.CollectionConverters._ + +import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot, ViewVarCharVector} +import org.apache.arrow.vector.ipc.ArrowFileWriter +import org.apache.arrow.vector.types.pojo.ArrowType + +import org.apache.spark.SparkUnsupportedOperationException import org.apache.spark.sql.functions._ import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.util.ArrowUtils import org.apache.spark.util.Utils class ArrowFileReadWriteSuite extends SharedSparkSession { @@ -38,7 +48,8 @@ class ArrowFileReadWriteSuite extends SharedSparkSession { lit(2L).alias("long"), lit(3.0).alias("double"), lit("a string").alias("str"), - lit(Array(1.0, 2.0, Double.NaN, Double.NegativeInfinity)).alias("arr")) + lit(Array(1.0, 2.0, Double.NaN, Double.NegativeInfinity)).alias("arr"), + col("id").cast("timestamp").alias("ts")) val path = new File(tempDataPath, "simple.arrowfile").toPath ArrowFileReadWrite.save(df, path) @@ -57,4 +68,31 @@ class ArrowFileReadWriteSuite extends SharedSparkSession { val df2 = ArrowFileReadWrite.load(spark, path) checkAnswer(df, df2) } + + test("loading a file whose layout differs from the canonical Arrow encoding fails fast") { + // `load` rebuilds vectors from the Spark schema, so an Arrow type that converts to a Spark + // type but is encoded differently (here Utf8View vs Utf8) cannot be loaded; it must be + // rejected at the schema check instead of having its buffers misread as garbage values. + val allocator = ArrowUtils.rootAllocator.newChildAllocator("stringViewFile", 0, Long.MaxValue) + val vector = new ViewVarCharVector("v", allocator) + vector.allocateNew() + val bytes = "a-string-longer-than-twelve-bytes".getBytes("utf8") + vector.setSafe(0, bytes, 0, bytes.length) + vector.setValueCount(1) + val root = new VectorSchemaRoot(Seq[FieldVector](vector).asJava) + val path = new File(tempDataPath, "stringview.arrowfile").toPath + val writer = new ArrowFileWriter(root, null, Channels.newChannel(Files.newOutputStream(path))) + writer.start() + writer.writeBatch() + writer.close() + root.close() + allocator.close() + + checkError( + exception = intercept[SparkUnsupportedOperationException] { + ArrowFileReadWrite.load(spark, path) + }, + condition = "UNSUPPORTED_ARROWTYPE", + parameters = Map("typeName" -> ArrowType.Utf8View.INSTANCE.toString)) + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowWriterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowWriterSuite.scala index 2584193008ff6..6e8e2deeba67f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowWriterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowWriterSuite.scala @@ -19,9 +19,11 @@ package org.apache.spark.sql.execution.arrow import scala.jdk.CollectionConverters._ -import org.apache.arrow.vector.VectorSchemaRoot +import org.apache.arrow.vector.{VectorSchemaRoot, ViewVarBinaryVector, ViewVarCharVector} +import org.apache.arrow.vector.complex.ListViewVector +import org.apache.arrow.vector.types.pojo.{ArrowType, FieldType} -import org.apache.spark.{SparkArithmeticException, SparkFunSuite} +import org.apache.spark.{SparkArithmeticException, SparkFunSuite, SparkUnsupportedOperationException} import org.apache.spark.sql.Row import org.apache.spark.sql.YearUDT import org.apache.spark.sql.catalyst.InternalRow @@ -153,6 +155,28 @@ class ArrowWriterSuite extends SparkFunSuite { check(new YearUDT, Seq(2020, 2021, null, 2022)) } + test("view vectors are rejected with UNSUPPORTED_ARROWTYPE") { + // The view types are readable through ArrowColumnVector but have no field writers; the error + // must name the Arrow type, not the Spark type its schema maps to. + val allocator = ArrowUtils.rootAllocator.newChildAllocator("view", 0, Long.MaxValue) + val listView = ListViewVector.empty("arr", allocator) + listView.addOrGetVector(FieldType.nullable(new ArrowType.Int(8 * 4, true))) + val vectors = Seq( + new ViewVarCharVector("str", allocator), + new ViewVarBinaryVector("bin", allocator), + listView) + vectors.foreach { vector => + checkError( + exception = intercept[SparkUnsupportedOperationException] { + ArrowWriter.createFieldWriter(vector) + }, + condition = "UNSUPPORTED_ARROWTYPE", + parameters = Map("typeName" -> vector.getField.getType.toString)) + vector.close() + } + allocator.close() + } + test("timestamp nanos round-trip") { // Decompose an int64 epoch-nanoseconds value into the (epochMicros, nanosWithinMicro) pair, // matching how the Arrow reader reconstructs it. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/vectorized/ArrowColumnVectorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/vectorized/ArrowColumnVectorSuite.scala index 9180ce1aee198..ea6bc4a78f02d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/vectorized/ArrowColumnVectorSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/vectorized/ArrowColumnVectorSuite.scala @@ -19,6 +19,7 @@ package org.apache.spark.sql.vectorized import org.apache.arrow.vector._ import org.apache.arrow.vector.complex._ +import org.apache.arrow.vector.types.pojo.{ArrowType, FieldType} import org.apache.spark.SparkFunSuite import org.apache.spark.sql.types._ @@ -331,6 +332,96 @@ class ArrowColumnVectorSuite extends SparkFunSuite { allocator.close() } + test("string_view") { + val allocator = ArrowUtils.rootAllocator.newChildAllocator("string_view", 0, Long.MaxValue) + val vector = new ViewVarCharVector("stringView", allocator) + vector.allocateNew() + + // Mix short (inline, <= 12 bytes) and long (stored in a data buffer, > 12 bytes) values to + // exercise both view-storage paths. + val values = (0 until 10).map { i => + if (i % 2 == 0) s"str$i" else s"a-long-string-value-$i" + } + values.zipWithIndex.foreach { case (s, i) => + val utf8 = s.getBytes("utf8") + vector.setSafe(i, utf8, 0, utf8.length) + } + vector.setNull(10) + vector.setValueCount(11) + + val columnVector = new ArrowColumnVector(vector) + assert(columnVector.dataType === StringType) + assert(columnVector.hasNull) + assert(columnVector.numNulls === 1) + + values.zipWithIndex.foreach { case (s, i) => + assert(columnVector.getUTF8String(i) === UTF8String.fromString(s)) + } + assert(columnVector.isNullAt(10)) + + columnVector.close() + allocator.close() + } + + test("binary_view") { + val allocator = ArrowUtils.rootAllocator.newChildAllocator("binary_view", 0, Long.MaxValue) + val vector = new ViewVarBinaryVector("binaryView", allocator) + vector.allocateNew() + + // Mix short (inline, <= 12 bytes) and long (stored in a data buffer, > 12 bytes) values to + // exercise both view-storage paths. + val values = (0 until 10).map { i => + if (i % 2 == 0) s"str$i" else s"a-long-binary-value-$i" + } + values.zipWithIndex.foreach { case (s, i) => + val utf8 = s.getBytes("utf8") + vector.setSafe(i, utf8, 0, utf8.length) + } + vector.setNull(10) + vector.setValueCount(11) + + val columnVector = new ArrowColumnVector(vector) + assert(columnVector.dataType === BinaryType) + assert(columnVector.hasNull) + assert(columnVector.numNulls === 1) + + values.zipWithIndex.foreach { case (s, i) => + assert(columnVector.getBinary(i) === s.getBytes("utf8")) + } + assert(columnVector.isNullAt(10)) + + columnVector.close() + allocator.close() + } + + test("string_view with multiple data buffers") { + val allocator = ArrowUtils.rootAllocator.newChildAllocator("string_view", 0, Long.MaxValue) + val vector = new ViewVarCharVector("stringView", allocator) + // Keep the variadic data buffers small (16 * 8 = 128 bytes each) so the long values below + // spill into multiple buffers, exercising the non-zero buffer-index branch of the accessor. + vector.setInitialCapacity(16, 8) + vector.allocateNew() + + val values = (0 until 16).map(i => s"a-long-string-value-spilling-over-$i") + values.zipWithIndex.foreach { case (s, i) => + val utf8 = s.getBytes("utf8") + vector.setSafe(i, utf8, 0, utf8.length) + } + vector.setValueCount(16) + // The values must not fit in a single data buffer, otherwise this test exercises nothing + // beyond the plain string_view test. + assert(vector.getDataBuffers.size() > 1) + + val columnVector = new ArrowColumnVector(vector) + assert(columnVector.dataType === StringType) + values.zipWithIndex.foreach { case (s, i) => + assert(columnVector.getUTF8String(i) === UTF8String.fromString(s)) + } + + columnVector.close() + allocator.close() + } + test("array") { val allocator = ArrowUtils.rootAllocator.newChildAllocator("array", 0, Long.MaxValue) val vector = ArrowUtils.toArrowField("array", ArrayType(IntegerType), nullable = true, null) @@ -385,6 +476,60 @@ class ArrowColumnVectorSuite extends SparkFunSuite { allocator.close() } + test("array_view") { + val allocator = ArrowUtils.rootAllocator.newChildAllocator("array_view", 0, Long.MaxValue) + val vector = ListViewVector.empty("arrayView", allocator) + vector.addOrGetVector(FieldType.nullable(new ArrowType.Int(8 * 4, true))) + vector.allocateNew() + val elementVector = vector.getDataVector().asInstanceOf[IntVector] + + // [1, 2] + vector.startNewValue(0) + elementVector.setSafe(0, 1) + elementVector.setSafe(1, 2) + vector.endValue(0, 2) + + // [3, null, 5] + vector.startNewValue(1) + elementVector.setSafe(2, 3) + elementVector.setNull(3) + elementVector.setSafe(4, 5) + vector.endValue(1, 3) + + // null + + // [] + vector.startNewValue(3) + vector.endValue(3, 0) + + elementVector.setValueCount(5) + vector.setValueCount(4) + + val columnVector = new ArrowColumnVector(vector) + assert(columnVector.dataType === ArrayType(IntegerType)) + assert(columnVector.hasNull) + assert(columnVector.numNulls === 1) + + val array0 = columnVector.getArray(0) + assert(array0.numElements() === 2) + assert(array0.getInt(0) === 1) + assert(array0.getInt(1) === 2) + + val array1 = columnVector.getArray(1) + assert(array1.numElements() === 3) + assert(array1.getInt(0) === 3) + assert(array1.isNullAt(1)) + assert(array1.getInt(2) === 5) + + assert(columnVector.isNullAt(2)) + + val array3 = columnVector.getArray(3) + assert(array3.numElements() === 0) + + columnVector.close() + allocator.close() + } + test("non nullable struct") { val allocator = ArrowUtils.rootAllocator.newChildAllocator("struct", 0, Long.MaxValue) val schema = new StructType().add("int", IntegerType).add("long", LongType)