This is an automated email from the ASF dual-hosted git repository.
zhouyuan pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new d64a51dfd0 [GLUTEN-13487][VL] Offload Lance scans with vector columns
(#13488)
d64a51dfd0 is described below
commit d64a51dfd01bbe1b8e4adb6c1f74331fb75f531f
Author: EJ Song <[email protected]>
AuthorDate: Mon Nov 2 22:33:51 2026 -0800
[GLUTEN-13487][VL] Offload Lance scans with vector columns (#13488)
---
.../gluten/execution/LanceScanTransformer.scala | 129 ++++++++-
.../execution/LanceFixedSizeListsSuite.scala | 322 +++++++++++++++++++++
.../apache/gluten/execution/VeloxLanceSuite.scala | 196 ++++++++++++-
docs/backends/velox/get-started/VeloxLance.md | 6 +
4 files changed, 632 insertions(+), 21 deletions(-)
diff --git
a/backends-velox/src-lance/main/scala/org/apache/gluten/execution/LanceScanTransformer.scala
b/backends-velox/src-lance/main/scala/org/apache/gluten/execution/LanceScanTransformer.scala
index 7a3c8bf002..e603e2bfab 100644
---
a/backends-velox/src-lance/main/scala/org/apache/gluten/execution/LanceScanTransformer.scala
+++
b/backends-velox/src-lance/main/scala/org/apache/gluten/execution/LanceScanTransformer.scala
@@ -34,8 +34,11 @@ import org.apache.spark.sql.vectorized.{ColumnarBatch,
ColumnVector}
import org.apache.arrow.c.{ArrowArrayStream, Data}
import org.apache.arrow.memory.BufferAllocator
-import org.apache.arrow.vector.FieldVector
+import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{BaseRepeatedValueVector,
FixedSizeListVector}
import org.apache.arrow.vector.ipc.ArrowReader
+import org.apache.arrow.vector.ipc.message.ArrowFieldNode
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
import org.lance.spark.internal.LanceArrowStreamScanner
import org.lance.spark.read.{LanceInputPartition, LanceScan}
@@ -202,17 +205,7 @@ private class LanceFragmentIterator(
advanced = false
val root = reader.getVectorSchemaRoot
val rowCount = root.getRowCount
- // Transfer (zero-copy move) each column out of the reader's reused root
so the emitted batch
- // owns its buffers: it stays valid after the reader is closed, and is
safe for the offload
- // transition to close.
- val transferred = new
java.util.ArrayList[FieldVector](root.getFieldVectors.size())
- root.getFieldVectors.asScala.foreach {
- fv =>
- val pair = fv.getTransferPair(allocator)
- pair.transfer()
- transferred.add(pair.getTo.asInstanceOf[FieldVector])
- }
- val vectors = ArrowColumnVectors.wrap(transferred)
+ val vectors =
ArrowColumnVectors.wrap(LanceFragmentIterator.takeColumns(root, allocator))
new ColumnarBatch(vectors.map(_.asInstanceOf[ColumnVector]), rowCount)
}
@@ -228,3 +221,115 @@ private class LanceFragmentIterator(
}
}
}
+
+private object LanceFragmentIterator {
+
+ /**
+ * Transfers (zero-copy moves) each column out of the reader's reused `root`
so the emitted batch
+ * owns its buffers: it stays valid after the reader is closed, and is safe
for the offload
+ * transition to close. Columns holding a FixedSizeList (Lance vectors) are
rebuilt as Lists on
+ * the way out, still without copying their values. If a column fails, the
columns already taken
+ * are released.
+ */
+ def takeColumns(
+ root: VectorSchemaRoot,
+ allocator: BufferAllocator): java.util.List[FieldVector] = {
+ val taken = new
java.util.ArrayList[FieldVector](root.getFieldVectors.size())
+ try {
+ root.getFieldVectors.asScala.foreach {
+ fv =>
+ if (FixedSizeLists.contains(fv.getField)) {
+ taken.add(FixedSizeLists.toLists(fv, allocator))
+ } else {
+ val pair = fv.getTransferPair(allocator)
+ pair.transfer()
+ taken.add(pair.getTo.asInstanceOf[FieldVector])
+ }
+ }
+ taken
+ } catch {
+ case t: Throwable =>
+ taken.asScala.foreach(_.close())
+ throw t
+ }
+ }
+}
+
+/**
+ * Rewrites FixedSizeList vectors, which is how Lance stores vector
(embedding) columns, as
+ * variable-size Lists. Neither Spark's
[[org.apache.spark.sql.vectorized.ArrowColumnVector]], which
+ * wraps the emitted batches, nor Velox's Arrow import supports FixedSizeList,
while a List holds
+ * the same values and matches the ArrayType that lance-spark declares for the
column. Element `i`
+ * of a FixedSizeList of size `n` is child range `[i * n, (i + 1) * n)` also
for null entries, so
+ * the offsets are `i * n` throughout and the child values are reused as is:
only the offsets buffer
+ * is new. Nested FixedSizeLists, e.g. a vector inside a struct or a list of
vectors, are rewritten
+ * too.
+ */
+private object FixedSizeLists {
+
+ /** Whether `field` is a FixedSizeList or has one among its descendants. */
+ def contains(field: Field): Boolean =
+ field.getType.isInstanceOf[ArrowType.FixedSizeList] ||
+ field.getChildren.asScala.exists(contains)
+
+ /**
+ * Builds a vector in `allocator` holding the data of `source`, with every
FixedSizeList rewritten
+ * as a List. Buffers are shared or moved rather than copied, so `source`
may be left without its
+ * child data; its owner still closes it.
+ */
+ def toLists(source: FieldVector, allocator: BufferAllocator): FieldVector = {
+ val target = toListField(source.getField).createVector(allocator)
+ try {
+ load(source, target, allocator)
+ target
+ } catch {
+ case t: Throwable =>
+ target.close()
+ throw t
+ }
+ }
+
+ private def toListField(field: Field): Field =
+ if (!contains(field)) {
+ field
+ } else {
+ val fieldType = field.getType match {
+ case _: ArrowType.FixedSizeList =>
+ new FieldType(
+ field.isNullable,
+ ArrowType.List.INSTANCE,
+ field.getDictionary,
+ field.getMetadata)
+ case _ => field.getFieldType
+ }
+ new Field(field.getName, fieldType,
field.getChildren.asScala.map(toListField).asJava)
+ }
+
+ /** Loads `source` into `target`, whose field is
`toListField(source.getField)`. */
+ private def load(source: FieldVector, target: FieldVector, allocator:
BufferAllocator): Unit =
+ if (!contains(source.getField)) {
+ // Same type on both sides: move the buffers.
+ source.makeTransferPair(target).transfer()
+ } else {
+ val node = new ArrowFieldNode(source.getValueCount, source.getNullCount)
+ source match {
+ case fixed: FixedSizeListVector =>
+ val width = BaseRepeatedValueVector.OFFSET_WIDTH
+ val offsets = allocator.buffer((fixed.getValueCount + 1).toLong *
width)
+ try {
+ (0 to fixed.getValueCount).foreach {
+ i => offsets.setInt(i.toLong * width, i * fixed.getListSize)
+ }
+ target.loadFieldBuffers(node,
java.util.Arrays.asList(fixed.getValidityBuffer, offsets))
+ } finally {
+ // loadFieldBuffers took its own reference.
+ offsets.close()
+ }
+ case _ =>
+ target.loadFieldBuffers(node, source.getFieldBuffers)
+ }
+
source.getChildrenFromFields.asScala.zip(target.getChildrenFromFields.asScala).foreach
{
+ case (s, t) => load(s, t, allocator)
+ }
+ }
+}
diff --git
a/backends-velox/src-lance/test/scala/org/apache/gluten/execution/LanceFixedSizeListsSuite.scala
b/backends-velox/src-lance/test/scala/org/apache/gluten/execution/LanceFixedSizeListsSuite.scala
new file mode 100644
index 0000000000..1f40068ccf
--- /dev/null
+++
b/backends-velox/src-lance/test/scala/org/apache/gluten/execution/LanceFixedSizeListsSuite.scala
@@ -0,0 +1,322 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.execution
+
+import org.apache.spark.sql.vectorized.ArrowColumnVector
+
+import org.apache.arrow.memory.{BufferAllocator, OutOfMemoryException,
RootAllocator}
+import org.apache.arrow.vector.{FieldVector, Float4Vector, IntVector,
VarCharVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.{FixedSizeListVector, ListVector,
StructVector}
+import org.apache.arrow.vector.types.FloatingPointPrecision
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.scalatest.funsuite.AnyFunSuite
+
+import java.nio.charset.StandardCharsets
+
+import scala.collection.JavaConverters._
+
+/**
+ * Unit tests for how the Lance scan takes the columns of an imported batch
+ * ([[LanceFragmentIterator.takeColumns]]) and rewrites its FixedSizeList
vectors as Lists
+ * ([[FixedSizeLists]]). Every test runs on its own [[RootAllocator]], whose
`close()` fails if any
+ * buffer leaked.
+ */
+class LanceFixedSizeListsSuite extends AnyFunSuite {
+
+ private val floatType = new
ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)
+
+ private def withAllocator(f: BufferAllocator => Unit): Unit = {
+ val allocator = new RootAllocator(Long.MaxValue)
+ try f(allocator)
+ finally allocator.close()
+ }
+
+ private def vectorField(name: String, dim: Int): Field =
+ new Field(
+ name,
+ FieldType.nullable(new ArrowType.FixedSizeList(dim)),
+ Seq(new Field("item", FieldType.notNullable(floatType), null)).asJava)
+
+ /** Fills the FixedSizeList `vector` with `rows`; null slots get junk child
values. */
+ private def fill(vector: FixedSizeListVector, rows:
Seq[Option[Seq[Float]]]): Unit = {
+ val dim = vector.getListSize
+ val values = vector.getDataVector.asInstanceOf[Float4Vector]
+ rows.zipWithIndex.foreach {
+ case (row, i) =>
+ row match {
+ case Some(_) => vector.setNotNull(i)
+ case None => vector.setNull(i)
+ }
+ (0 until dim).foreach(j => values.setSafe(i * dim + j,
row.map(_(j)).getOrElse(-1f)))
+ }
+ }
+
+ private def floatVectors(
+ allocator: BufferAllocator,
+ dim: Int,
+ rows: Seq[Option[Seq[Float]]]): FixedSizeListVector = {
+ val vector = vectorField("vec",
dim).createVector(allocator).asInstanceOf[FixedSizeListVector]
+ vector.allocateNew()
+ fill(vector, rows)
+ vector.setValueCount(rows.size)
+ vector
+ }
+
+ /** Reads the arrays of `column` through Spark's ArrowColumnVector. */
+ private def read(column: ArrowColumnVector, rows: Int):
Seq[Option[Seq[Float]]] =
+ (0 until rows).map {
+ i => if (column.isNullAt(i)) None else
Some(column.getArray(i).toFloatArray.toSeq)
+ }
+
+ /**
+ * Rewrites the `dim`-sized vectors in `source`, checks the result reads as
`expected`, and
+ * releases both vectors.
+ */
+ private def checkRewrite(
+ source: FieldVector,
+ allocator: BufferAllocator,
+ dim: Int,
+ expected: Seq[Option[Seq[Float]]]): Unit = {
+ val converted = FixedSizeLists.toLists(source, allocator)
+ source.close()
+ val column = new ArrowColumnVector(converted)
+ try {
+ assert(converted.getField.getType == ArrowType.List.INSTANCE)
+ assert(converted.getValueCount == expected.size)
+ val offsets = converted.asInstanceOf[ListVector].getOffsetBuffer
+ (0 to expected.size).foreach(i => assert(offsets.getInt(i * 4L) == i *
dim))
+ assert(read(column, expected.size) == expected)
+ } finally {
+ column.close()
+ }
+ }
+
+ test("contains finds FixedSizeLists at any depth") {
+ val plain = new Field("id", FieldType.nullable(new ArrowType.Int(32,
true)), null)
+ val list = new Field(
+ "values",
+ FieldType.nullable(ArrowType.List.INSTANCE),
+ Seq(new Field("item", FieldType.nullable(floatType), null)).asJava)
+ def struct(children: Field*): Field =
+ new Field("s", FieldType.nullable(ArrowType.Struct.INSTANCE),
children.asJava)
+ def listOf(child: Field): Field =
+ new Field("l", FieldType.nullable(ArrowType.List.INSTANCE),
Seq(child).asJava)
+
+ assert(!FixedSizeLists.contains(plain))
+ assert(!FixedSizeLists.contains(list))
+ assert(!FixedSizeLists.contains(struct(plain, list)))
+ assert(FixedSizeLists.contains(vectorField("vec", 4)))
+ assert(FixedSizeLists.contains(struct(plain, vectorField("vec", 4))))
+ assert(FixedSizeLists.contains(listOf(vectorField("vec", 4))))
+ assert(FixedSizeLists.contains(struct(listOf(struct(vectorField("vec",
4))))))
+ }
+
+ test("rewrites a FixedSizeList with null rows as a List with the same
values") {
+ withAllocator {
+ allocator =>
+ // The null row's child slots hold junk, which must not leak into the
neighboring rows.
+ val rows = Seq(
+ Some(Seq(0f, 1f, 2f, 3f)),
+ None,
+ Some(Seq(8f, 9f, 10f, 11f)),
+ None,
+ Some(Seq(16f, 17f, 18f, 19f)))
+ checkRewrite(floatVectors(allocator, 4, rows), allocator, 4, rows)
+ }
+ }
+
+ test("rewrites FixedSizeLists with no nulls, only nulls, and no rows") {
+ withAllocator {
+ allocator =>
+ // Uniform validity, and an empty batch whose offsets buffer holds a
single 0.
+ val noNulls = Seq(Some(Seq(1f, 2f)), Some(Seq(3f, 4f)), Some(Seq(5f,
6f)))
+ checkRewrite(floatVectors(allocator, 2, noNulls), allocator, 2,
noNulls)
+ val onlyNulls = Seq(None, None, None)
+ checkRewrite(floatVectors(allocator, 2, onlyNulls), allocator, 2,
onlyNulls)
+ checkRewrite(floatVectors(allocator, 2, Seq.empty), allocator, 2,
Seq.empty)
+ }
+ }
+
+ test("moves the vector values without copying them") {
+ withAllocator {
+ allocator =>
+ val source = floatVectors(allocator, 3, Seq(Some(Seq(1f, 2f, 3f)),
None))
+ val values =
source.getDataVector.asInstanceOf[Float4Vector].getDataBufferAddress
+ val converted = FixedSizeLists.toLists(source, allocator)
+ try {
+ val list = converted.asInstanceOf[ListVector]
+
assert(list.getDataVector.asInstanceOf[Float4Vector].getDataBufferAddress ==
values)
+ assert(list.getValidityBufferAddress ==
source.getValidityBufferAddress)
+ } finally {
+ converted.close()
+ source.close()
+ }
+ }
+ }
+
+ test("rewrites vectors nested in a struct and in a list") {
+ withAllocator {
+ allocator =>
+ val docField = new Field(
+ "doc",
+ FieldType.nullable(ArrowType.Struct.INSTANCE),
+ Seq(
+ new Field("label", FieldType.nullable(ArrowType.Utf8.INSTANCE),
null),
+ vectorField("emb", 2)).asJava
+ )
+ val doc = docField.createVector(allocator).asInstanceOf[StructVector]
+ doc.allocateNew()
+ val label = doc.getChild("label", classOf[VarCharVector])
+ val emb = doc.getChild("emb", classOf[FixedSizeListVector])
+ // Rows: a full struct, a struct with a null vector, and a null struct.
+ Seq(0, 1).foreach {
+ i =>
+ doc.setIndexDefined(i)
+ label.setSafe(i, s"d$i".getBytes(StandardCharsets.UTF_8))
+ }
+ doc.setNull(2)
+ fill(emb, Seq(Some(Seq(1f, 2f)), None, None))
+ doc.setValueCount(3)
+
+ val multiField =
+ new Field(
+ "multi",
+ FieldType.nullable(ArrowType.List.INSTANCE),
+ Seq(vectorField("e", 2)).asJava)
+ val multi = multiField.createVector(allocator).asInstanceOf[ListVector]
+ multi.allocateNew()
+ val elements = multi.getDataVector.asInstanceOf[FixedSizeListVector]
+ // Rows: two vectors, an empty list, and a null list.
+ multi.startNewValue(0)
+ multi.endValue(0, 2)
+ multi.startNewValue(1)
+ multi.endValue(1, 0)
+ multi.setNull(2)
+ fill(elements, Seq(Some(Seq(1f, 2f)), Some(Seq(3f, 4f))))
+ multi.setValueCount(3)
+
+ val convertedDoc = FixedSizeLists.toLists(doc, allocator)
+ val convertedMulti = FixedSizeLists.toLists(multi, allocator)
+ doc.close()
+ multi.close()
+ val docColumn = new ArrowColumnVector(convertedDoc)
+ val multiColumn = new ArrowColumnVector(convertedMulti)
+ try {
+ assert(
+ convertedDoc.getField.getChildren.asScala.map(_.getType) ==
+ Seq(ArrowType.Utf8.INSTANCE, ArrowType.List.INSTANCE))
+ assert((0 until 3).map(docColumn.isNullAt) == Seq(false, false,
true))
+ assert((0 until
2).map(docColumn.getChild(0).getUTF8String(_).toString) == Seq(
+ "d0",
+ "d1"))
+ assert(read(docColumn.getChild(1), 2) == Seq(Some(Seq(1f, 2f)),
None))
+
+ assert(convertedMulti.getField.getChildren.get(0).getType ==
ArrowType.List.INSTANCE)
+ assert((0 until 3).map(multiColumn.isNullAt) == Seq(false, false,
true))
+ val first = multiColumn.getArray(0)
+ assert(
+ (0 until
first.numElements).map(first.getArray(_).toFloatArray.toSeq) ==
+ Seq(Seq(1f, 2f), Seq(3f, 4f)))
+ assert(multiColumn.getArray(1).numElements == 0)
+ } finally {
+ docColumn.close()
+ multiColumn.close()
+ }
+ }
+ }
+
+ test("releases the partial result when the rewrite fails") {
+ withAllocator {
+ allocator =>
+ val docField = new Field(
+ "doc",
+ FieldType.nullable(ArrowType.Struct.INSTANCE),
+ Seq(
+ new Field("id", FieldType.nullable(new ArrowType.Int(32, true)),
null),
+ vectorField("emb", 2)).asJava
+ )
+ val doc = docField.createVector(allocator).asInstanceOf[StructVector]
+ doc.allocateNew()
+ val id = doc.getChild("id", classOf[IntVector])
+ doc.setIndexDefined(0)
+ id.setSafe(0, 7)
+ doc.setNull(1)
+ fill(doc.getChild("emb", classOf[FixedSizeListVector]),
Seq(Some(Seq(1f, 2f)), None))
+ doc.setValueCount(2)
+
+ // `id` moves into the rewrite before the offsets of `emb` fail to
allocate, so the failed
+ // rewrite has to release what it already took.
+ val limited = allocator.newChildAllocator("limited", 0, 0)
+ try {
+ intercept[OutOfMemoryException](FixedSizeLists.toLists(doc, limited))
+ assert(limited.getAllocatedMemory == 0)
+ } finally {
+ limited.close()
+ doc.close()
+ }
+ }
+ }
+
+ /** A batch with an int column `id` and a vector column `vec`. */
+ private def batch(allocator: BufferAllocator): VectorSchemaRoot = {
+ val id = new IntVector("id", allocator)
+ id.allocateNew(2)
+ id.set(0, 7)
+ id.set(1, 8)
+ id.setValueCount(2)
+ val vec = floatVectors(allocator, 2, Seq(Some(Seq(1f, 2f)), None))
+ new VectorSchemaRoot(Seq[FieldVector](id, vec).asJava)
+ }
+
+ test("takes the columns of a batch, rewriting only the vector columns") {
+ withAllocator {
+ allocator =>
+ val root = batch(allocator)
+ try {
+ val columns =
+ LanceFragmentIterator.takeColumns(root, allocator).asScala.map(new
ArrowColumnVector(_))
+ try {
+ assert(
+ columns.map(_.getValueVector.getField.getType) ==
+ Seq(new ArrowType.Int(32, true), ArrowType.List.INSTANCE))
+ assert((0 until 2).map(columns(0).getInt) == Seq(7, 8))
+ assert(read(columns(1), 2) == Seq(Some(Seq(1f, 2f)), None))
+ } finally {
+ columns.foreach(_.close())
+ }
+ } finally {
+ root.close()
+ }
+ }
+ }
+
+ test("releases the columns already taken when a later column fails") {
+ withAllocator {
+ allocator =>
+ val root = batch(allocator)
+ // `id` moves into `limited` before the offsets of `vec` fail to
allocate there.
+ val limited = allocator.newChildAllocator("limited", 0, 0)
+ try {
+
intercept[OutOfMemoryException](LanceFragmentIterator.takeColumns(root,
limited))
+ assert(limited.getAllocatedMemory == 0)
+ } finally {
+ limited.close()
+ root.close()
+ }
+ }
+ }
+}
diff --git
a/backends-velox/src-lance/test/scala/org/apache/gluten/execution/VeloxLanceSuite.scala
b/backends-velox/src-lance/test/scala/org/apache/gluten/execution/VeloxLanceSuite.scala
index 16dcc4b478..8909c87b1f 100644
---
a/backends-velox/src-lance/test/scala/org/apache/gluten/execution/VeloxLanceSuite.scala
+++
b/backends-velox/src-lance/test/scala/org/apache/gluten/execution/VeloxLanceSuite.scala
@@ -18,8 +18,11 @@ package org.apache.gluten.execution
import org.apache.spark.SparkConf
import org.apache.spark.sql.{DataFrame, Row}
-import org.apache.spark.sql.execution.{ColumnarToRowExec, InputAdapter,
SparkPlan}
+import org.apache.spark.sql.execution.{ColumnarShuffleExchangeExec,
ColumnarToRowExec, InputAdapter, SparkPlan}
import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
+import org.apache.spark.sql.types._
+
+import scala.collection.JavaConverters._
/**
* Exercises the read-only Lance offload: a columnar Lance scan is handed to
Velox through the Arrow
@@ -52,27 +55,94 @@ class VeloxLanceSuite extends
VeloxWholeStageTransformerSuite {
view: String,
rows: Int,
writeOptions: Map[String, String] = Map.empty)(f: => Unit): Unit = {
+ val data = spark
+ .range(rows)
+ .selectExpr(
+ "cast(id as int) as id",
+ "cast(id * 2 as int) as v",
+ "concat('n', cast(id as string)) as name")
+ withLanceView(view, data, writeOptions, Map.empty)(f)
+ }
+
+ /** Writes `data` as a Lance dataset and registers it as a temp view. */
+ private def withLanceView(
+ view: String,
+ data: DataFrame,
+ writeOptions: Map[String, String],
+ readOptions: Map[String, String])(f: => Unit): Unit = {
withTempPath {
path =>
val uri = s"${path.getCanonicalPath}/$view.lance"
- spark
- .range(rows)
- .selectExpr(
- "cast(id as int) as id",
- "cast(id * 2 as int) as v",
- "concat('n', cast(id as string)) as name")
- .write
+ data.write
.format("lance")
.options(writeOptions)
.option("path", uri)
.save()
- spark.read.format("lance").option("path",
uri).load().createOrReplaceTempView(view)
+ spark.read
+ .format("lance")
+ .options(readOptions)
+ .option("path", uri)
+ .load()
+ .createOrReplaceTempView(view)
try f
finally spark.catalog.dropTempView(view)
}
}
+ /** Field metadata that makes lance-spark write an array column as a
FixedSizeList (a vector). */
+ private def vectorMetadata(dim: Int): Metadata =
+ new MetadataBuilder().putLong("arrow.fixed-size-list.size", dim).build()
+
+ /**
+ * Row `i` of the vector dataset (see [[withLanceVectorView]]). Every tenth
row, from row 3 on,
+ * has a null `vec`, a null `doc`, a null vector inside `doc`, and a null
`multi`, respectively.
+ */
+ private def vectorRow(i: Int): Row = {
+ val f = i.toFloat
+ Row(
+ i,
+ if (i % 10 == 3) null else Seq(f, f + 1, f + 2, f + 3),
+ Seq(i * 10.0, i * 10.0 + 1),
+ if (i % 10 == 5) null else Row(s"d$i", if (i % 10 == 6) null else
Seq(-f, -f - 1, -f - 2)),
+ if (i % 10 == 7) null else Seq.tabulate(i % 3)(j => Seq(f * 100 + j, f *
100 + j + 0.5f))
+ )
+ }
+
+ /**
+ * Writes `rows` rows of FixedSizeList vector columns as a Lance dataset:
`vec` (float, size 4),
+ * `dvec` (double, size 2), a vector inside the struct `doc`, and the list
of vectors `multi`.
+ */
+ private def withLanceVectorView(
+ view: String,
+ rows: Int = 10,
+ writeOptions: Map[String, String] = Map.empty,
+ readOptions: Map[String, String] = Map.empty)(f: => Unit): Unit = {
+ // The element metadata of `multi` makes each element a FixedSizeList of
size 2.
+ val multiMetadata =
+ new MetadataBuilder().putString("_lance.element",
vectorMetadata(2).json).build()
+ val schema = new StructType()
+ .add("id", IntegerType)
+ .add("vec", ArrayType(FloatType, containsNull = false), nullable = true,
vectorMetadata(4))
+ .add("dvec", ArrayType(DoubleType, containsNull = false), nullable =
true, vectorMetadata(2))
+ .add(
+ "doc",
+ new StructType()
+ .add("label", StringType)
+ .add(
+ "emb",
+ ArrayType(FloatType, containsNull = false),
+ nullable = true,
+ vectorMetadata(3)))
+ .add(
+ "multi",
+ ArrayType(ArrayType(FloatType, containsNull = false), containsNull =
false),
+ nullable = true,
+ multiMetadata)
+ val data = spark.createDataFrame((0 until rows).map(vectorRow).asJava,
schema)
+ withLanceView(view, data, writeOptions, readOptions)(f)
+ }
+
private def isLanceScan(plan: SparkPlan): Boolean = plan match {
case _: LanceScanTransformer => true
case InputAdapter(child) => isLanceScan(child)
@@ -85,6 +155,15 @@ class VeloxLanceSuite extends
VeloxWholeStageTransformerSuite {
.filterNot(_.isInstanceOf[InputAdapter])
.filter(_.children.exists(isLanceScan))
+ /** Asserts that Velox reads the Lance scan's batches directly, with no row
conversion between. */
+ private def assertScanFeedsVelox(df: DataFrame): Unit = {
+ val parents = lanceScanParents(df)
+ assert(
+ parents.nonEmpty && parents.forall(_.isInstanceOf[OffloadArrowDataExec]),
+ s"Expected OffloadArrowDataExec over the Lance scan,
got:\n${df.queryExecution.executedPlan}"
+ )
+ }
+
test("lance scan offloads to LanceScanTransformer") {
withLanceView("lance_basic", 100) {
runQueryAndCompare("select id, v, name from lance_basic") {
@@ -170,4 +249,103 @@ class VeloxLanceSuite extends
VeloxWholeStageTransformerSuite {
checkAnswer(df, Seq(Row(50L)))
}
}
+
+ test("lance scan offloads vector columns") {
+ withLanceVectorView("lance_vectors") {
+ runQueryAndCompare("select id, vec, dvec from lance_vectors") {
+ df =>
+ checkGlutenPlan[LanceScanTransformer](df)
+ checkAnswer(df, (0 until 10).map(vectorRow).map(r => Row(r(0), r(1),
r(2))))
+ }
+ }
+ }
+
+ test("lance scan offloads nested vector columns") {
+ withLanceVectorView("lance_nested_vectors") {
+ runQueryAndCompare("select id, doc, multi from lance_nested_vectors
where id >= 2") {
+ df =>
+ checkGlutenPlan[LanceScanTransformer](df)
+ checkAnswer(df, (2 until 10).map(vectorRow).map(r => Row(r(0), r(3),
r(4))))
+ }
+ }
+ }
+
+ test("lance vector columns feed Velox operators") {
+ withLanceVectorView("lance_vectors_velox") {
+ // Velox evaluates the element accesses and the aggregation over the
imported vectors; null
+ // vectors, structs, and lists must stay null rather than expose their
child slots.
+ runQueryAndCompare(
+ """select id % 2 as k, count(vec) as c, sum(vec[0]) as s, sum(dvec[1])
as d,
+ | count(doc) as docs, count(doc.emb) as embs, sum(doc.emb[2]) as e,
+ | count(multi) as m, sum(multi[0][1]) as mm
+ |from lance_vectors_velox group by id % 2""".stripMargin) {
+ df =>
+ checkGlutenPlan[LanceScanTransformer](df)
+ checkGlutenPlan[HashAggregateExecBaseTransformer](df)
+ assertScanFeedsVelox(df)
+ }
+ }
+ }
+
+ test("lance vector columns pass through a Velox filter, shuffle and sort") {
+ withLanceVectorView("lance_vectors_shuffle") {
+ // `id % 3` is not pushed into Lance, so Velox filters the imported
vectors, shuffles and
+ // sorts them, and converts them back to rows.
+ runQueryAndCompare(
+ "select id, vec, doc, multi from lance_vectors_shuffle where id % 3 !=
1 order by id") {
+ df =>
+ checkGlutenPlan[LanceScanTransformer](df)
+ assertScanFeedsVelox(df)
+ checkGlutenPlan[FilterExecTransformerBase](df)
+ checkGlutenPlan[ColumnarShuffleExchangeExec](df)
+ checkGlutenPlan[SortExecTransformer](df)
+ checkAnswer(
+ df,
+ (0 until 10).filter(_ % 3 != 1).map(vectorRow).map(r => Row(r(0),
r(1), r(3), r(4))))
+ }
+ }
+ }
+
+ test("lance scan reads vector batches that are all null or empty") {
+ withLanceVectorView("lance_vectors_edge") {
+ // Single-row batches whose vector, struct, or list is null, and a scan
with no rows.
+ Seq(
+ "select id, vec from lance_vectors_edge where id = 3" -> Seq(Row(3,
null)),
+ "select id, doc from lance_vectors_edge where id = 5" -> Seq(Row(5,
null)),
+ "select id, doc from lance_vectors_edge where id = 6" -> Seq(Row(6,
Row("d6", null))),
+ "select id, multi from lance_vectors_edge where id = 7" -> Seq(Row(7,
null)),
+ "select id, vec, doc, multi from lance_vectors_edge where id < 0" ->
Seq.empty[Row]
+ ).foreach {
+ case (query, expected) =>
+ runQueryAndCompare(query) {
+ df =>
+ checkGlutenPlan[LanceScanTransformer](df)
+ checkAnswer(df, expected)
+ }
+ }
+ }
+ }
+
+ test("lance scan reads vectors across fragments and batches") {
+ withLanceVectorView(
+ "lance_vectors_multi",
+ rows = 1000,
+ writeOptions = Map("max_row_per_file" -> "300"),
+ readOptions = Map("batch_size" -> "64")) {
+ runQueryAndCompare("select * from lance_vectors_multi") {
+ df =>
+ checkGlutenPlan[LanceScanTransformer](df)
+ assert(df.rdd.getNumPartitions > 1, "Expected one partition per
Lance fragment")
+ checkAnswer(df, (0 until 1000).map(vectorRow))
+ }
+ runQueryAndCompare(
+ """select count(vec), sum(vec[0]), sum(vec[3]), sum(doc.emb[1]),
sum(multi[1][0])
+ |from lance_vectors_multi""".stripMargin) {
+ df =>
+ checkGlutenPlan[LanceScanTransformer](df)
+ checkGlutenPlan[HashAggregateExecBaseTransformer](df)
+ assertScanFeedsVelox(df)
+ }
+ }
+ }
}
diff --git a/docs/backends/velox/get-started/VeloxLance.md
b/docs/backends/velox/get-started/VeloxLance.md
index 7a789101b8..256be3a6b4 100644
--- a/docs/backends/velox/get-started/VeloxLance.md
+++ b/docs/backends/velox/get-started/VeloxLance.md
@@ -58,9 +58,15 @@ serve it:
| Query shape | Behavior |
|--------------------------------|----------|
| Projection / filter | Offload |
+| Vector (`FixedSizeList`) columns | Offload |
| Pushed-down aggregation (e.g. `COUNT(*)`) | Fallback |
| Full-text query (`_score`) | Fallback |
+Vector columns are Arrow `FixedSizeList`s, which neither Spark's
`ArrowColumnVector` nor Velox
+reads, so the scan rewrites them as variable-size lists, the `ArrayType`
lance-spark already
+exposes. The vector values are not copied; each vector column only gets an
offsets buffer. This
+also applies to vectors nested in a struct or a list.
+
The offload decision is made at plan time
(`LanceScanTransformer.isExportable`), so a
non-offloadable scan is detected before execution rather than failing inside
Velox.
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
