This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 544704c87b86 [SPARK-57929][CORE][SQL] Support Binary/List view types
in ArrowColumnVector
544704c87b86 is described below
commit 544704c87b86d16bade9b1f0acaa55b3578f8fcb
Author: Robert Kruszewski <[email protected]>
AuthorDate: Thu Jul 16 22:15:25 2026 +0800
[SPARK-57929][CORE][SQL] Support Binary/List view types in ArrowColumnVector
### What changes were proposed in this pull request?
Add support for constructing ArrowColumnVector from String/BinaryView and
ListView
### Why are the changes needed?
Arrow in spec 1.4 has added these as more efficient way to represent list
and string types in memory
### Does this PR introduce _any_ user-facing change?
No, this is internal developer api
### How was this patch tested?
Added Tests
### Was this patch authored or co-authored using generative AI tooling?
Yes, this pr was authored with use of Fable 5
Closes #57210 from robert3005/arrow-view-column-vectors.
Authored-by: Robert Kruszewski <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
.../org/apache/spark/sql/util/ArrowUtils.scala | 4 +-
.../spark/sql/vectorized/ArrowColumnVector.java | 63 +++++++--
.../spark/sql/execution/arrow/ArrowWriter.scala | 6 +
.../connect/client/arrow/ArrowEncoderSuite.scala | 54 +++++++-
.../connect/client/arrow/ArrowVectorReader.scala | 13 ++
.../sql/execution/arrow/ArrowFileReadWrite.scala | 30 ++++-
.../execution/arrow/ArrowFileReadWriteSuite.scala | 40 +++++-
.../sql/execution/arrow/ArrowWriterSuite.scala | 28 +++-
.../sql/vectorized/ArrowColumnVectorSuite.scala | 145 +++++++++++++++++++++
9 files changed, 363 insertions(+), 20 deletions(-)
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 e69a0fa7f415..b8dab164cf78 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 f44de5ffa9df..d6e2f3f2c12e 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 @@ public class ArrowColumnVector extends ColumnVector {
} 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 @@ public class ArrowColumnVector extends ColumnVector {
} 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 @@ public class ArrowColumnVector extends ColumnVector {
}
}
+ // 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 @@ public class ArrowColumnVector extends ColumnVector {
}
}
- 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<T extends BaseListVector & RepeatedValueVector>
+ 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 b96e57ce49af..f4c3bc08f403 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 9f171cda2ea2..efb6a7a2304f 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 54311cecc162..1f71ed8f8432 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 e7ec2d2b7984..516b01b25fdc 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 8accda825403..dda9827e80a6 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 2584193008ff..6e8e2deeba67 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 9180ce1aee19..ea6bc4a78f02 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)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
