Copilot commented on code in PR #12958: URL: https://github.com/apache/gluten/pull/12958#discussion_r4229332854
########## backends-velox/src-lance/main/scala/org/apache/gluten/execution/LanceScanTransformer.scala: ########## @@ -0,0 +1,198 @@ +/* + * 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.gluten.backendsapi.arrow.ArrowBatchTypes.ArrowJavaBatchType +import org.apache.gluten.extension.columnar.transition.Convention +import org.apache.gluten.iterator.Iterators +import org.apache.gluten.memory.arrow.alloc.ArrowBufferAllocators +import org.apache.gluten.vectorized.ArrowWritableColumnVector + +import org.apache.spark.{Partition, SparkContext, TaskContext} +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartitioning} +import org.apache.spark.sql.connector.read.{InputPartition, Scan} +import org.apache.spark.sql.execution.LeafExecNode +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} + +import org.apache.arrow.c.{ArrowArrayStream, Data} +import org.apache.arrow.vector.FieldVector +import org.lance.spark.internal.LanceArrowStreamScanner +import org.lance.spark.read.{LanceInputPartition, LanceScan} + +import scala.collection.JavaConverters._ + +/** + * A Gluten leaf that reads a Lance table through the Arrow C Data Interface and produces + * [[ArrowJavaBatchType]] columnar batches, which Gluten then offloads to Velox via its automatic + * ArrowJava -> ArrowNative transition. + * + * This replaces a vanilla [[BatchScanExec]] over a [[LanceScan]]. Each Spark partition maps to one + * Lance fragment. For every fragment the leaf asks lance-spark to export the planned scan as an + * [[ArrowArrayStream]] ([[LanceArrowStreamScanner.export]]) and hands back only the C-struct + * address. The address is re-wrapped with Gluten's own Arrow build and imported here, so + * lance-spark and Gluten never share Arrow Java objects across their classloaders -- only the + * version-stable C ABI struct crosses the boundary. This is the consumer side of lance-core's + * {@code LanceScanner#exportArrowStream(long)} (lance#7259). + * + * Each imported batch is transferred out of the reader's root into independent Arrow buffers so the + * emitted [[ColumnarBatch]] owns its data: Gluten's offload transition C-exports then closes each + * batch, and the recycle callback closes it again, exactly as for [[ColumnarRangeExec]]. The reader + * and the lance-spark handle are closed on task completion, which releases the native scan and the + * caller-owned stream struct. + */ +case class LanceScanTransformer(@transient batchScan: BatchScanExec) + extends LeafExecNode + with ValidatablePlan { + + @transient private lazy val lanceScan: LanceScan = + batchScan.scan.asInstanceOf[LanceScan] + + @transient private lazy val inputPartitions: Array[InputPartition] = + lanceScan.toBatch.planInputPartitions() + + // Captured eagerly: `batchScan` is transient, but `output` is also read on executors (e.g. by + // DeserializeToObjectExec) after this node has been serialized. + override val output: Seq[Attribute] = batchScan.output + + override def outputPartitioning: Partitioning = UnknownPartitioning(inputPartitions.length) + + override def batchType(): Convention.BatchType = ArrowJavaBatchType + + override def rowType(): Convention.RowType = Convention.RowType.None + + override protected def doExecute(): RDD[InternalRow] = + throw new UnsupportedOperationException(s"$nodeName does not support row-based execution.") + + override protected def doExecuteColumnar(): RDD[ColumnarBatch] = + new LanceColumnarRDD(sparkContext, inputPartitions) +} + +object LanceScanTransformer { + + /** Whether this DSv2 scan is a Lance scan that can be offloaded to Velox through Arrow C. */ + def supportsBatchScan(scan: Scan): Boolean = scan.isInstanceOf[LanceScan] + + /** + * Whether the planned Lance scan can actually be exported through the Arrow C stream. + * + * A pushed aggregation (e.g. {@code COUNT(*)}) or a full-text query makes the native fragment + * scan's output diverge from the partition's declared Spark schema: for an aggregation the + * partition's output is the aggregate result while the fragment scan returns data rows, and a + * full-text query auto-projects a {@code _score} column. [[LanceArrowStreamScanner.export]] + * rejects both at runtime (checkNoPushedAggregation / checkNativeSchemaMatchesPartition), but by + * then the node is already offloaded to Velox, so the throw fails the query instead of falling + * back. Detecting them here keeps such scans on vanilla Spark (e.g. the fast COUNT(*) metadata + * reader) rather than offloading a scan that can only crash. + */ + def isExportable(batchScan: BatchScanExec): Boolean = + batchScan.inputPartitions.forall { + case p: LanceInputPartition => + !p.getPushedAggregation.isPresent && p.getReadOptions.getFullTextQuery == null + case _ => false + } +} + +/** One RDD partition per planned Lance [[InputPartition]] (one Lance fragment each). */ +private class LancePartition(override val index: Int, val inputPartition: LanceInputPartition) + extends Partition + +/** + * Executes the planned Lance partitions, importing each fragment's exported Arrow C stream with + * Gluten's Arrow build and yielding offloadable [[ArrowWritableColumnVector]] batches. + */ +private class LanceColumnarRDD( + @transient private val sc: SparkContext, + private val inputPartitions: Array[InputPartition]) + extends RDD[ColumnarBatch](sc, Nil) { + + override protected def getPartitions: Array[Partition] = + inputPartitions.zipWithIndex.map { + case (p, i) => new LancePartition(i, p.asInstanceOf[LanceInputPartition]) + } + + override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { + val lancePartition = split.asInstanceOf[LancePartition].inputPartition + val allocator = ArrowBufferAllocators.contextInstance() + + val batches: Iterator[ColumnarBatch] = + lancePartition.getLanceSplit.getFragments.asScala.iterator.flatMap { + fragId => + // Plan + export the fragment scan on the lance-spark side; only the C-struct address + // crosses over. wrap() views that struct with Gluten's Arrow, and importArrayStream() + // moves it into a reader that owns and drains the native scan. + val handle = LanceArrowStreamScanner.export(fragId.intValue(), lancePartition) + val stream = ArrowArrayStream.wrap(handle.streamAddress()) + val reader = Data.importArrayStream(allocator, stream) + context.addTaskCompletionListener[Unit] { + _ => + try { + // Drains + runs the C release callback, tearing down the native scan. + reader.close() + } finally { + // Frees the caller-owned stream struct and closes the scanner + dataset. + handle.close() + } Review Comment: Resource cleanup is not exception-safe during fragment initialization: if `wrap(...)` or `importArrayStream(...)` throws after `export(...)` succeeds, `handle` (and its native resources / stream struct) can leak. Wrap initialization in a construct that guarantees `handle.close()` on failure (and `reader.close()` if created) before propagating the exception. ########## pom.xml: ########## @@ -1962,6 +1964,90 @@ </plugins> </build> </profile> + <profile> + <id>lance</id> + <activation> + <activeByDefault>false</activeByDefault> + </activation> + <build> + <plugins> + <plugin> + <groupId>org.codehaus.mojo</groupId> + <artifactId>build-helper-maven-plugin</artifactId> + <executions> + <execution> + <id>add-lance-sources</id> + <goals> + <goal>add-source</goal> + </goals> + <phase>generate-sources</phase> + <configuration> + <sources> + <source>${project.basedir}/src-lance/main/scala</source> + </sources> + </configuration> + </execution> + <execution> + <id>add-lance-resources</id> + <goals> + <goal>add-resource</goal> + </goals> + <phase>generate-resources</phase> + <configuration> + <resources> + <resource> + <directory>${project.basedir}/src-lance/main/resources</directory> + </resource> + </resources> + </configuration> + </execution> + </executions> + </plugin> + </plugins> + </build> + </profile> + <profile> + <id>lance-test</id> + <activation> + <activeByDefault>false</activeByDefault> + </activation> + <build> + <plugins> + <plugin> + <groupId>org.codehaus.mojo</groupId> + <artifactId>build-helper-maven-plugin</artifactId> + <executions> + <execution> + <id>add-lance-test-sources</id> + <goals> + <goal>add-test-source</goal> + </goals> + <phase>generate-test-sources</phase> + <configuration> + <sources> + <source>${project.basedir}/src-lance/test/scala</source> + </sources> + </configuration> + </execution> + <execution> + <id>add-lance-test-resources</id> + <goals> + <goal>add-test-resource</goal> + </goals> + <phase>generate-test-resources</phase> + <configuration> + <resources> + <resource> + <directory>${project.basedir}/src-lance/test/resources</directory> + </resource> + </resources> + </configuration> + </execution> + </executions> + </plugin> + </plugins> + </build> Review Comment: The parent POM profile adds extra source/resource roots for every module when `-Plance` is enabled. While `${project.basedir}` will resolve per-module, this still broadens build behavior across the entire reactor (and can complicate IDE import / incremental builds). Consider scoping the build-helper configuration to the `backends-velox` module POM (where the `src-lance` directories actually live), keeping the parent profile limited to shared properties/versioning. ########## docs/get-started/VeloxLance.md: ########## @@ -0,0 +1,79 @@ +--- +layout: page +title: Lance Support in Velox Backend +nav_order: 9 +parent: Getting-Started +--- + +# Lance Support in Velox Backend + +Gluten can offload reads of [Lance](https://lancedb.github.io/lance/) datasets (via +[lance-spark](https://github.com/lance-format/lance-spark)) to the Velox backend. The scan is +handed to Velox through the Arrow C Data Interface: lance-spark exports each fragment as an Arrow C +stream and Gluten imports it, so only the Arrow C-struct address crosses the boundary and Gluten's +Arrow version stays decoupled from lance-spark's. + +Only reads are supported. Writes fall back to vanilla Spark. + +## Requirements + +- **lance-spark** 0.8.0 or later on the classpath (the first release with the Arrow C stream + export, `LanceArrowStreamScanner`), e.g. `org.lance:lance-spark-bundle-3.5_2.12:0.8.0`. The + pinned version is the `lance.version` property in the root `pom.xml`. +- **Platform:** lance-core ships its native library for `linux-x86-64`, `linux-aarch64` and + `darwin-aarch64`. +- **Spark:** tested with Spark 3.5. lance-spark also publishes Spark 3.4 and 4.x artifacts. + +## Building + +Enable the `lance` profile alongside the Velox backend: + +``` +mvn clean package -Pbackends-velox -Pspark-3.5 -Plance -DskipTests +``` + +The read-only offload is entirely Velox-specific, so it lives under `backends-velox/src-lance` +(mirroring `backends-velox/src-iceberg`) rather than in a top-level module. Add `-Plance-test` to +also compile and run the Lance test suite: + +``` +mvn test -pl backends-velox -Pbackends-velox -Pspark-3.5 -Plance -Plance-test \ + -Dtest=none -DwildcardSuites=org.apache.gluten.execution.VeloxLanceSuite +``` + +## Reading + +Offload / Fallback. Review Comment: This reads like a placeholder and doesn’t add information. Consider turning it into a proper section heading (e.g., `## Offload and fallback behavior`) or removing the line to keep the doc polished. ########## backends-velox/src-lance/main/scala/org/apache/gluten/execution/LanceScanTransformer.scala: ########## @@ -0,0 +1,198 @@ +/* + * 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.gluten.backendsapi.arrow.ArrowBatchTypes.ArrowJavaBatchType +import org.apache.gluten.extension.columnar.transition.Convention +import org.apache.gluten.iterator.Iterators +import org.apache.gluten.memory.arrow.alloc.ArrowBufferAllocators +import org.apache.gluten.vectorized.ArrowWritableColumnVector + +import org.apache.spark.{Partition, SparkContext, TaskContext} +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartitioning} +import org.apache.spark.sql.connector.read.{InputPartition, Scan} +import org.apache.spark.sql.execution.LeafExecNode +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} + +import org.apache.arrow.c.{ArrowArrayStream, Data} +import org.apache.arrow.vector.FieldVector +import org.lance.spark.internal.LanceArrowStreamScanner +import org.lance.spark.read.{LanceInputPartition, LanceScan} + +import scala.collection.JavaConverters._ + +/** + * A Gluten leaf that reads a Lance table through the Arrow C Data Interface and produces + * [[ArrowJavaBatchType]] columnar batches, which Gluten then offloads to Velox via its automatic + * ArrowJava -> ArrowNative transition. + * + * This replaces a vanilla [[BatchScanExec]] over a [[LanceScan]]. Each Spark partition maps to one + * Lance fragment. For every fragment the leaf asks lance-spark to export the planned scan as an + * [[ArrowArrayStream]] ([[LanceArrowStreamScanner.export]]) and hands back only the C-struct + * address. The address is re-wrapped with Gluten's own Arrow build and imported here, so + * lance-spark and Gluten never share Arrow Java objects across their classloaders -- only the + * version-stable C ABI struct crosses the boundary. This is the consumer side of lance-core's + * {@code LanceScanner#exportArrowStream(long)} (lance#7259). + * + * Each imported batch is transferred out of the reader's root into independent Arrow buffers so the + * emitted [[ColumnarBatch]] owns its data: Gluten's offload transition C-exports then closes each + * batch, and the recycle callback closes it again, exactly as for [[ColumnarRangeExec]]. The reader + * and the lance-spark handle are closed on task completion, which releases the native scan and the + * caller-owned stream struct. + */ +case class LanceScanTransformer(@transient batchScan: BatchScanExec) + extends LeafExecNode + with ValidatablePlan { + + @transient private lazy val lanceScan: LanceScan = + batchScan.scan.asInstanceOf[LanceScan] + + @transient private lazy val inputPartitions: Array[InputPartition] = + lanceScan.toBatch.planInputPartitions() + + // Captured eagerly: `batchScan` is transient, but `output` is also read on executors (e.g. by + // DeserializeToObjectExec) after this node has been serialized. + override val output: Seq[Attribute] = batchScan.output + + override def outputPartitioning: Partitioning = UnknownPartitioning(inputPartitions.length) + + override def batchType(): Convention.BatchType = ArrowJavaBatchType + + override def rowType(): Convention.RowType = Convention.RowType.None + + override protected def doExecute(): RDD[InternalRow] = + throw new UnsupportedOperationException(s"$nodeName does not support row-based execution.") + + override protected def doExecuteColumnar(): RDD[ColumnarBatch] = + new LanceColumnarRDD(sparkContext, inputPartitions) +} + +object LanceScanTransformer { + + /** Whether this DSv2 scan is a Lance scan that can be offloaded to Velox through Arrow C. */ + def supportsBatchScan(scan: Scan): Boolean = scan.isInstanceOf[LanceScan] + + /** + * Whether the planned Lance scan can actually be exported through the Arrow C stream. + * + * A pushed aggregation (e.g. {@code COUNT(*)}) or a full-text query makes the native fragment + * scan's output diverge from the partition's declared Spark schema: for an aggregation the + * partition's output is the aggregate result while the fragment scan returns data rows, and a + * full-text query auto-projects a {@code _score} column. [[LanceArrowStreamScanner.export]] + * rejects both at runtime (checkNoPushedAggregation / checkNativeSchemaMatchesPartition), but by + * then the node is already offloaded to Velox, so the throw fails the query instead of falling + * back. Detecting them here keeps such scans on vanilla Spark (e.g. the fast COUNT(*) metadata + * reader) rather than offloading a scan that can only crash. + */ + def isExportable(batchScan: BatchScanExec): Boolean = + batchScan.inputPartitions.forall { + case p: LanceInputPartition => + !p.getPushedAggregation.isPresent && p.getReadOptions.getFullTextQuery == null + case _ => false + } +} + +/** One RDD partition per planned Lance [[InputPartition]] (one Lance fragment each). */ +private class LancePartition(override val index: Int, val inputPartition: LanceInputPartition) + extends Partition + +/** + * Executes the planned Lance partitions, importing each fragment's exported Arrow C stream with + * Gluten's Arrow build and yielding offloadable [[ArrowWritableColumnVector]] batches. + */ +private class LanceColumnarRDD( + @transient private val sc: SparkContext, + private val inputPartitions: Array[InputPartition]) + extends RDD[ColumnarBatch](sc, Nil) { + + override protected def getPartitions: Array[Partition] = + inputPartitions.zipWithIndex.map { + case (p, i) => new LancePartition(i, p.asInstanceOf[LanceInputPartition]) + } + + override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { + val lancePartition = split.asInstanceOf[LancePartition].inputPartition + val allocator = ArrowBufferAllocators.contextInstance() + + val batches: Iterator[ColumnarBatch] = + lancePartition.getLanceSplit.getFragments.asScala.iterator.flatMap { + fragId => + // Plan + export the fragment scan on the lance-spark side; only the C-struct address + // crosses over. wrap() views that struct with Gluten's Arrow, and importArrayStream() + // moves it into a reader that owns and drains the native scan. + val handle = LanceArrowStreamScanner.export(fragId.intValue(), lancePartition) + val stream = ArrowArrayStream.wrap(handle.streamAddress()) + val reader = Data.importArrayStream(allocator, stream) + context.addTaskCompletionListener[Unit] { + _ => + try { + // Drains + runs the C release callback, tearing down the native scan. + reader.close() + } finally { + // Frees the caller-owned stream struct and closes the scanner + dataset. + handle.close() + } + } + + new Iterator[ColumnarBatch] { + private var advanced = false + private var hasMore = false + + private def advance(): Unit = + if (!advanced) { + hasMore = reader.loadNextBatch() + advanced = true + } + + override def hasNext: Boolean = { + advance() + hasMore + } Review Comment: Each fragment’s `reader`/`handle` is only closed on task completion. If a Spark partition contains many fragments, this retains native scanners/streams longer than necessary and can increase peak memory/native resource usage. Prefer closing `reader` + `handle` as soon as a fragment iterator is exhausted (e.g., when `loadNextBatch()` returns false, or in a per-fragment `try/finally`/close-on-finish iterator), while keeping the task-completion listener only as a safety net for early task termination. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
