This is an automated email from the ASF dual-hosted git repository. taiyang-li pushed a commit to branch fake_add_bolt_backend in repository https://gitbox.apache.org/repos/asf/gluten.git
commit f9f85ee0f466c0df35e24041ec84b34efebaf777 Author: liyang.127 <[email protected]> AuthorDate: Tue Jun 30 21:02:51 2026 +0800 [GLUTEN][CORE] Share RowToColumnarExec under backends-core Extract the common Row -> Arrow ColumnarBatch conversion logic from RowToVeloxColumnarExec into a new SharedRowToColumnarExec under backends-core. Backends only need to subclass it and supply two hooks: - preferredBatchBytes (the backend's preferred Arrow batch byte size) - sparkToBackendUnsafe (the backend's `BroadcastUtils.sparkToXxxUnsafe`) VeloxRowToColumnarExec becomes a thin subclass that fills the two backend-specific hooks. The `object RowToVeloxColumnarExec` keeps existing call sites by forwarding to SharedRowToColumnarExec. Generated-by: TraeCli openrouter-3o --- .../gluten/execution/SharedRowToColumnarExec.scala | 42 ++-- .../gluten/execution/RowToVeloxColumnarExec.scala | 224 +++------------------ 2 files changed, 52 insertions(+), 214 deletions(-) diff --git a/backends-velox/src/main/scala/org/apache/gluten/execution/RowToVeloxColumnarExec.scala b/backends-core/src/main/scala/org/apache/gluten/execution/SharedRowToColumnarExec.scala similarity index 87% copy from backends-velox/src/main/scala/org/apache/gluten/execution/RowToVeloxColumnarExec.scala copy to backends-core/src/main/scala/org/apache/gluten/execution/SharedRowToColumnarExec.scala index f11de5894a..5f375f53ae 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/execution/RowToVeloxColumnarExec.scala +++ b/backends-core/src/main/scala/org/apache/gluten/execution/SharedRowToColumnarExec.scala @@ -18,7 +18,7 @@ package org.apache.gluten.execution import org.apache.gluten.backendsapi.BackendsApiManager import org.apache.gluten.columnarbatch.ColumnarBatches -import org.apache.gluten.config.{GlutenConfig, VeloxConfig} +import org.apache.gluten.config.GlutenConfig import org.apache.gluten.iterator.Iterators import org.apache.gluten.memory.arrow.alloc.ArrowBufferAllocators import org.apache.gluten.runtime.Runtimes @@ -43,19 +43,39 @@ import org.apache.arrow.memory.ArrowBuf import scala.collection.mutable.ListBuffer -case class RowToVeloxColumnarExec(child: SparkPlan) extends RowToColumnarExecBase(child = child) { +/** + * Backend-agnostic implementation of `Row -> ColumnarBatch` conversion. Backends only need to + * subclass it and supply the per-backend preferred batch byte size and the broadcast unsafe + * relation builder. + */ +abstract class SharedRowToColumnarExec(child: SparkPlan) extends RowToColumnarExecBase(child) { + + /** Backend-specific preferred batch byte size used to flush an Arrow batch. */ + protected def preferredBatchBytes: Long + + /** + * Backend-specific `Spark -> Unsafe Arrow` broadcast builder. Equivalent to + * `BroadcastUtils.sparkTo{Bolt,Velox}Unsafe`. + */ + protected def sparkToBackendUnsafe[F, T]( + sc: org.apache.spark.SparkContext, + mode: org.apache.spark.sql.catalyst.plans.physical.BroadcastMode, + schema: StructType, + relation: Broadcast[F], + itrTransformer: Iterator[InternalRow] => Iterator[ColumnarBatch]): Broadcast[T] + override def doExecuteColumnarInternal(): RDD[ColumnarBatch] = { val numInputRows = longMetric("numInputRows") val numOutputBatches = longMetric("numOutputBatches") val convertTime = longMetric("convertTime") val numRows = GlutenConfig.get.maxBatchSize - val numBytes = VeloxConfig.get.veloxPreferredBatchBytes + val numBytes = preferredBatchBytes // This avoids calling `schema` in the RDD closure, so that we don't need to include the entire // plan (this) in the closure. val localSchema = schema child.execute().mapPartitions { rowIterator => - RowToVeloxColumnarExec.toColumnarBatchIterator( + SharedRowToColumnarExec.toColumnarBatchIterator( rowIterator, localSchema, numInputRows, @@ -71,16 +91,16 @@ case class RowToVeloxColumnarExec(child: SparkPlan) extends RowToColumnarExecBas val numOutputBatches = longMetric("numOutputBatches") val convertTime = longMetric("convertTime") val numRows = GlutenConfig.get.maxBatchSize - val numBytes = VeloxConfig.get.veloxPreferredBatchBytes + val numBytes = preferredBatchBytes val mode = BroadcastUtils.getBroadcastMode(outputPartitioning) val relation = child.executeBroadcast() - BroadcastUtils.sparkToVeloxUnsafe( + sparkToBackendUnsafe( sparkContext, mode, schema, relation, itr => - RowToVeloxColumnarExec.toColumnarBatchIterator( + SharedRowToColumnarExec.toColumnarBatchIterator( itr, schema, numInputRows, @@ -90,13 +110,9 @@ case class RowToVeloxColumnarExec(child: SparkPlan) extends RowToColumnarExecBas numBytes) ) } - - // For spark 3.2. - protected def withNewChildInternal(newChild: SparkPlan): RowToVeloxColumnarExec = - copy(child = newChild) } -object RowToVeloxColumnarExec { +object SharedRowToColumnarExec { def toColumnarBatchIterator( it: Iterator[InternalRow], @@ -106,7 +122,7 @@ object RowToVeloxColumnarExec { val numInputRows = new SQLMetric("numInputRows") val numOutputBatches = new SQLMetric("numOutputBatches") val convertTime = new SQLMetric("convertTime") - RowToVeloxColumnarExec.toColumnarBatchIterator( + toColumnarBatchIterator( it, schema, numInputRows, diff --git a/backends-velox/src/main/scala/org/apache/gluten/execution/RowToVeloxColumnarExec.scala b/backends-velox/src/main/scala/org/apache/gluten/execution/RowToVeloxColumnarExec.scala index f11de5894a..c8300e8d36 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/execution/RowToVeloxColumnarExec.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/execution/RowToVeloxColumnarExec.scala @@ -16,79 +16,28 @@ */ package org.apache.gluten.execution -import org.apache.gluten.backendsapi.BackendsApiManager -import org.apache.gluten.columnarbatch.ColumnarBatches -import org.apache.gluten.config.{GlutenConfig, VeloxConfig} -import org.apache.gluten.iterator.Iterators -import org.apache.gluten.memory.arrow.alloc.ArrowBufferAllocators -import org.apache.gluten.runtime.Runtimes -import org.apache.gluten.utils.ArrowAbiUtil -import org.apache.gluten.vectorized._ +import org.apache.gluten.config.VeloxConfig +import org.apache.spark.SparkContext import org.apache.spark.broadcast.Broadcast -import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{UnsafeProjection, UnsafeRow} +import org.apache.spark.sql.catalyst.plans.physical.BroadcastMode import org.apache.spark.sql.execution.{BroadcastUtils, SparkPlan} import org.apache.spark.sql.execution.metric.SQLMetric -import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StructType -import org.apache.spark.sql.utils.SparkArrowUtil import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.spark.task.TaskResources -import org.apache.spark.unsafe.Platform -import org.apache.arrow.c.ArrowSchema -import org.apache.arrow.memory.ArrowBuf +case class RowToVeloxColumnarExec(child: SparkPlan) extends SharedRowToColumnarExec(child) { -import scala.collection.mutable.ListBuffer + override protected def preferredBatchBytes: Long = VeloxConfig.get.veloxPreferredBatchBytes -case class RowToVeloxColumnarExec(child: SparkPlan) extends RowToColumnarExecBase(child = child) { - override def doExecuteColumnarInternal(): RDD[ColumnarBatch] = { - val numInputRows = longMetric("numInputRows") - val numOutputBatches = longMetric("numOutputBatches") - val convertTime = longMetric("convertTime") - val numRows = GlutenConfig.get.maxBatchSize - val numBytes = VeloxConfig.get.veloxPreferredBatchBytes - // This avoids calling `schema` in the RDD closure, so that we don't need to include the entire - // plan (this) in the closure. - val localSchema = schema - child.execute().mapPartitions { - rowIterator => - RowToVeloxColumnarExec.toColumnarBatchIterator( - rowIterator, - localSchema, - numInputRows, - numOutputBatches, - convertTime, - numRows, - numBytes) - } - } - - override def doExecuteBroadcast[T](): Broadcast[T] = { - val numInputRows = longMetric("numInputRows") - val numOutputBatches = longMetric("numOutputBatches") - val convertTime = longMetric("convertTime") - val numRows = GlutenConfig.get.maxBatchSize - val numBytes = VeloxConfig.get.veloxPreferredBatchBytes - val mode = BroadcastUtils.getBroadcastMode(outputPartitioning) - val relation = child.executeBroadcast() - BroadcastUtils.sparkToVeloxUnsafe( - sparkContext, - mode, - schema, - relation, - itr => - RowToVeloxColumnarExec.toColumnarBatchIterator( - itr, - schema, - numInputRows, - numOutputBatches, - convertTime, - numRows, - numBytes) - ) + override protected def sparkToBackendUnsafe[F, T]( + sc: SparkContext, + mode: BroadcastMode, + schema: StructType, + relation: Broadcast[F], + itrTransformer: Iterator[InternalRow] => Iterator[ColumnarBatch]): Broadcast[T] = { + BroadcastUtils.sparkToVeloxUnsafe(sc, mode, schema, relation, itrTransformer) } // For spark 3.2. @@ -102,19 +51,8 @@ object RowToVeloxColumnarExec { it: Iterator[InternalRow], schema: StructType, columnBatchSize: Int, - columnBatchBytes: Long): Iterator[ColumnarBatch] = { - val numInputRows = new SQLMetric("numInputRows") - val numOutputBatches = new SQLMetric("numOutputBatches") - val convertTime = new SQLMetric("convertTime") - RowToVeloxColumnarExec.toColumnarBatchIterator( - it, - schema, - numInputRows, - numOutputBatches, - convertTime, - columnBatchSize, - columnBatchBytes) - } + columnBatchBytes: Long): Iterator[ColumnarBatch] = + SharedRowToColumnarExec.toColumnarBatchIterator(it, schema, columnBatchSize, columnBatchBytes) def toColumnarBatchIterator( it: Iterator[InternalRow], @@ -123,129 +61,13 @@ object RowToVeloxColumnarExec { numOutputBatches: SQLMetric, convertTime: SQLMetric, columnBatchSize: Int, - columnBatchBytes: Long): Iterator[ColumnarBatch] = { - if (it.isEmpty) { - return Iterator.empty - } - - val arrowSchema = - SparkArrowUtil.toArrowSchema(schema, SQLConf.get.sessionLocalTimeZone) - val runtime = Runtimes.contextInstance(BackendsApiManager.getBackendName, "RowToColumnar") - val jniWrapper = NativeRowToColumnarJniWrapper.create(runtime) - val arrowAllocator = ArrowBufferAllocators.contextInstance() - val cSchema = ArrowSchema.allocateNew(arrowAllocator) - val factory = UnsafeProjection - val converter = factory.create(schema) - val r2cHandle = - try { - ArrowAbiUtil.exportSchema(arrowAllocator, arrowSchema, cSchema) - jniWrapper.init(cSchema.memoryAddress()) - } finally { - cSchema.close() - } - - val res: Iterator[ColumnarBatch] = new Iterator[ColumnarBatch] { - var finished = false - - override def hasNext: Boolean = { - if (finished) { - false - } else { - it.hasNext - } - } - - def convertToUnsafeRow(row: InternalRow): UnsafeRow = { - row match { - case unsafeRow: UnsafeRow => unsafeRow - case _ => - converter.apply(row) - } - } - - override def next(): ColumnarBatch = { - var arrowBuf: ArrowBuf = null - TaskResources.addRecycler("RowToColumnar_arrowBuf", 100) { - if (arrowBuf != null && arrowBuf.refCnt() != 0) { - arrowBuf.close() - } - } - val rowLength = new ListBuffer[Long]() - var rowCount = 0 - var offset = 0L - while (rowCount < columnBatchSize && offset < columnBatchBytes && !finished) { - if (!it.hasNext) { - finished = true - } else { - val row = it.next() - val start = System.currentTimeMillis() - val unsafeRow = convertToUnsafeRow(row) - val sizeInBytes = unsafeRow.getSizeInBytes - - // allocate buffer based on first row - if (rowCount == 0) { - // allocate buffer based on 1st row, but if first row is very big, this will cause OOM - // maybe we should optimize to list ArrayBuf to native to avoid buf close and allocate - // 31760L origins from BaseVariableWidthVector.lastValueAllocationSizeInBytes - // experimental value - val estimatedBufSize = Math.min( - Math.max( - Math.min(sizeInBytes.toDouble * columnBatchSize * 1.2, 31760L * columnBatchSize), - sizeInBytes.toDouble * 10), - // Limit the size of the buffer to columnBatchBytes or the size of the first row, - // whichever is greater so we always have enough space for the first row. - Math.max(columnBatchBytes, sizeInBytes) - ) - arrowBuf = arrowAllocator.buffer(estimatedBufSize.toLong) - } - - if ((offset + sizeInBytes) > arrowBuf.capacity()) { - val bufSize = if (offset + sizeInBytes > columnBatchBytes) { - // If adding the current row causes the batch size to exceed columnBatchBytes add - // just enough space to add the current row. - offset + sizeInBytes - } else { - Math.min((offset + sizeInBytes * 2), columnBatchBytes) - } - val tmpBuf = arrowAllocator.buffer(bufSize) - tmpBuf.setBytes(0, arrowBuf, 0, offset) - arrowBuf.close() - arrowBuf = tmpBuf - } - Platform.copyMemory( - unsafeRow.getBaseObject, - unsafeRow.getBaseOffset, - null, - arrowBuf.memoryAddress() + offset, - sizeInBytes) - offset += sizeInBytes - rowLength += sizeInBytes.toLong - rowCount += 1 - convertTime += System.currentTimeMillis() - start - } - } - numInputRows += rowCount - numOutputBatches += 1 - val startNative = System.currentTimeMillis() - try { - val handle = jniWrapper - .nativeConvertRowToColumnar(r2cHandle, rowLength.toArray, arrowBuf.memoryAddress()) - val cb = ColumnarBatches.create(handle) - convertTime += System.currentTimeMillis() - startNative - cb - } finally { - arrowBuf.close() - arrowBuf = null - } - } - } - Iterators - .wrap(res) - .protectInvocationFlow() - .recycleIterator { - jniWrapper.close(r2cHandle) - } - .recyclePayload(_.close()) - .create() - } + columnBatchBytes: Long): Iterator[ColumnarBatch] = + SharedRowToColumnarExec.toColumnarBatchIterator( + it, + schema, + numInputRows, + numOutputBatches, + convertTime, + columnBatchSize, + columnBatchBytes) } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
