This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new f9f32a1d perf: Optimize CometSparkToColumnar for columnar input (#892)
f9f32a1d is described below
commit f9f32a1d5e0efe5d57e26184d67a645486da61b9
Author: Matt Butrovich <[email protected]>
AuthorDate: Tue Sep 3 17:02:15 2024 -0400
perf: Optimize CometSparkToColumnar for columnar input (#892)
* Modify `CometSparkToColumnarExec` to use `ColumnarArray` rather than row
iteration when the child is producing columnar data.
* Modify `CometSparkToColumnarExec` to use `ColumnarArray` rather than row
iteration when the child is producing columnar data.
* Run `mvn spotless:apply` to clean up formatting, changed to camelCase for
variables, do not call `ColumnarArray.numElements()` in a tight loop, added
`writeColNoNull` as an optimization, reserve the initial vector capacity in
`writeCol` as an optimization.
* Fix one more variable name.
* Clean up duplicated code.
* Add test for basic types.
* Add more types to test.
* Revert testing changes to CometExecBenchmark.
* Add benchmark.
* Link to issue in ArrowWriters.
* Address review feedback.
* make format
* Address review feedback, mostly naming suggestions. Ran `make format`
again.
---
.../sql/comet/execution/arrow/ArrowWriters.scala | 35 ++++++
.../execution/arrow/CometArrowConverters.scala | 133 ++++++++++++++++-----
.../spark/sql/comet/CometSparkToColumnarExec.scala | 51 ++++----
.../org/apache/comet/exec/CometExecSuite.scala | 52 ++++++++
.../spark/sql/benchmark/CometExecBenchmark.scala | 10 ++
5 files changed, 228 insertions(+), 53 deletions(-)
diff --git
a/common/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala
b/common/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala
index 8d9f373f..f12ab7c4 100644
---
a/common/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala
+++
b/common/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala
@@ -28,12 +28,15 @@ import
org.apache.spark.sql.catalyst.expressions.SpecializedGetters
import org.apache.spark.sql.comet.util.Utils
import org.apache.spark.sql.errors.QueryExecutionErrors
import org.apache.spark.sql.types._
+import org.apache.spark.sql.vectorized.ColumnarArray
/**
* This file is mostly copied from Spark SQL's
* org.apache.spark.sql.execution.arrow.ArrowWriter.scala. Comet shadows Arrow
classes to avoid
* potential conflicts with Spark's Arrow dependencies, hence we cannot reuse
Spark's ArrowWriter
* directly.
+ *
+ * Performance enhancement:
https://github.com/apache/datafusion-comet/issues/888
*/
private[arrow] object ArrowWriter {
def create(root: VectorSchemaRoot): ArrowWriter = {
@@ -103,6 +106,16 @@ class ArrowWriter(val root: VectorSchemaRoot, fields:
Array[ArrowFieldWriter]) {
count += 1
}
+ def writeCol(input: ColumnarArray, columnIndex: Int): Unit = {
+ fields(columnIndex).writeCol(input)
+ count = input.numElements()
+ }
+
+ def writeColNoNull(input: ColumnarArray, columnIndex: Int): Unit = {
+ fields(columnIndex).writeColNoNull(input)
+ count = input.numElements()
+ }
+
def finish(): Unit = {
root.setRowCount(count)
fields.foreach(_.finish())
@@ -137,6 +150,28 @@ private[arrow] abstract class ArrowFieldWriter {
count += 1
}
+ def writeCol(input: ColumnarArray): Unit = {
+ val inputNumElements = input.numElements()
+ valueVector.setInitialCapacity(inputNumElements)
+ while (count < inputNumElements) {
+ if (input.isNullAt(count)) {
+ setNull()
+ } else {
+ setValue(input, count)
+ }
+ count += 1
+ }
+ }
+
+ def writeColNoNull(input: ColumnarArray): Unit = {
+ val inputNumElements = input.numElements()
+ valueVector.setInitialCapacity(inputNumElements)
+ while (count < inputNumElements) {
+ setValue(input, count)
+ count += 1
+ }
+ }
+
def finish(): Unit = {
valueVector.setValueCount(count)
}
diff --git
a/common/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala
b/common/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala
index 9dbd8dcd..e0489ae7 100644
---
a/common/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala
+++
b/common/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala
@@ -21,18 +21,19 @@ package org.apache.spark.sql.comet.execution.arrow
import org.apache.arrow.memory.{BufferAllocator, RootAllocator}
import org.apache.arrow.vector.VectorSchemaRoot
+import org.apache.arrow.vector.types.pojo.Schema
import org.apache.spark.TaskContext
import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.comet.util.Utils
import org.apache.spark.sql.types.StructType
-import org.apache.spark.sql.vectorized.ColumnarBatch
+import org.apache.spark.sql.vectorized.{ColumnarArray, ColumnarBatch}
import org.apache.comet.vector.NativeUtil
object CometArrowConverters extends Logging {
// TODO: we should reuse the same root allocator in the comet code base?
- val rootAllocator: BufferAllocator = new RootAllocator(Long.MaxValue)
+ private val rootAllocator: BufferAllocator = new RootAllocator(Long.MaxValue)
// This is similar how Spark converts internal row to Arrow format except
that it is transforming
// the result batch to Comet's ColumnarBatch instead of serialized bytes.
@@ -44,24 +45,23 @@ object CometArrowConverters extends Logging {
// exported process increases the reference count of the Arrow vectors. The
reference count is
// only decreased when the native plan is done with the vectors, which is
usually longer than
// all the ColumnarBatches are consumed.
- private[sql] class ArrowBatchIterator(
- rowIter: Iterator[InternalRow],
+
+ abstract private[sql] class ArrowBatchIterBase(
schema: StructType,
- maxRecordsPerBatch: Long,
timeZoneId: String,
context: TaskContext)
extends Iterator[ColumnarBatch]
with AutoCloseable {
- private val arrowSchema = Utils.toArrowSchema(schema, timeZoneId)
+ protected val arrowSchema: Schema = Utils.toArrowSchema(schema, timeZoneId)
// Reuse the same root allocator here.
- private val allocator =
+ protected val allocator: BufferAllocator =
rootAllocator.newChildAllocator(s"to${this.getClass.getSimpleName}", 0,
Long.MaxValue)
- private val root = VectorSchemaRoot.create(arrowSchema, allocator)
- private val arrowWriter = ArrowWriter.create(root)
+ protected val root: VectorSchemaRoot =
VectorSchemaRoot.create(arrowSchema, allocator)
+ protected val arrowWriter: ArrowWriter = ArrowWriter.create(root)
- private var currentBatch: ColumnarBatch = null
- private var closed: Boolean = false
+ protected var currentBatch: ColumnarBatch = null
+ protected var closed: Boolean = false
Option(context).foreach {
_.addTaskCompletionListener[Unit] { _ =>
@@ -69,9 +69,27 @@ object CometArrowConverters extends Logging {
}
}
- override def hasNext: Boolean = rowIter.hasNext || {
+ override def close(): Unit = {
close(false)
- false
+ }
+
+ protected def close(closeAllocator: Boolean): Unit = {
+ try {
+ if (!closed) {
+ if (currentBatch != null) {
+ arrowWriter.reset()
+ currentBatch.close()
+ currentBatch = null
+ }
+ root.close()
+ closed = true
+ }
+ } finally {
+ // the allocator shall be closed when the task is finished
+ if (closeAllocator) {
+ allocator.close()
+ }
+ }
}
override def next(): ColumnarBatch = {
@@ -79,11 +97,25 @@ object CometArrowConverters extends Logging {
currentBatch
}
- override def close(): Unit = {
+ protected def nextBatch(): ColumnarBatch
+
+ }
+
+ private[sql] class RowToArrowBatchIter(
+ rowIter: Iterator[InternalRow],
+ schema: StructType,
+ maxRecordsPerBatch: Long,
+ timeZoneId: String,
+ context: TaskContext)
+ extends ArrowBatchIterBase(schema, timeZoneId, context)
+ with AutoCloseable {
+
+ override def hasNext: Boolean = rowIter.hasNext || {
close(false)
+ false
}
- private def nextBatch(): ColumnarBatch = {
+ override protected def nextBatch(): ColumnarBatch = {
if (rowIter.hasNext) {
// the arrow writer shall be reset before writing the next batch
arrowWriter.reset()
@@ -99,33 +131,68 @@ object CometArrowConverters extends Logging {
null
}
}
+ }
- private def close(closeAllocator: Boolean): Unit = {
- try {
- if (!closed) {
- if (currentBatch != null) {
- arrowWriter.reset()
- currentBatch.close()
- currentBatch = null
+ def rowToArrowBatchIter(
+ rowIter: Iterator[InternalRow],
+ schema: StructType,
+ maxRecordsPerBatch: Long,
+ timeZoneId: String,
+ context: TaskContext): Iterator[ColumnarBatch] = {
+ new RowToArrowBatchIter(rowIter, schema, maxRecordsPerBatch, timeZoneId,
context)
+ }
+
+ private[sql] class ColumnBatchToArrowBatchIter(
+ colBatch: ColumnarBatch,
+ schema: StructType,
+ maxRecordsPerBatch: Int,
+ timeZoneId: String,
+ context: TaskContext)
+ extends ArrowBatchIterBase(schema, timeZoneId, context)
+ with AutoCloseable {
+
+ private var rowsProduced: Int = 0
+
+ override def hasNext: Boolean = rowsProduced < colBatch.numRows() || {
+ close(false)
+ false
+ }
+
+ override protected def nextBatch(): ColumnarBatch = {
+ val rowsInBatch = colBatch.numRows()
+ if (rowsProduced < rowsInBatch) {
+ // the arrow writer shall be reset before writing the next batch
+ arrowWriter.reset()
+ val rowsToProduce =
+ if (maxRecordsPerBatch <= 0) rowsInBatch - rowsProduced
+ else Math.min(maxRecordsPerBatch, rowsInBatch - rowsProduced)
+
+ for (columnIndex <- 0 until colBatch.numCols()) {
+ val column = colBatch.column(columnIndex)
+ val columnArray = new ColumnarArray(column, rowsProduced,
rowsToProduce)
+ if (column.hasNull) {
+ arrowWriter.writeCol(columnArray, columnIndex)
+ } else {
+ arrowWriter.writeColNoNull(columnArray, columnIndex)
}
- root.close()
- closed = true
- }
- } finally {
- // the allocator shall be closed when the task is finished
- if (closeAllocator) {
- allocator.close()
}
+
+ rowsProduced += rowsToProduce
+
+ arrowWriter.finish()
+ NativeUtil.rootAsBatch(root)
+ } else {
+ null
}
}
}
- def toArrowBatchIterator(
- rowIter: Iterator[InternalRow],
+ def columnarBatchToArrowBatchIter(
+ colBatch: ColumnarBatch,
schema: StructType,
- maxRecordsPerBatch: Long,
+ maxRecordsPerBatch: Int,
timeZoneId: String,
context: TaskContext): Iterator[ColumnarBatch] = {
- new ArrowBatchIterator(rowIter, schema, maxRecordsPerBatch, timeZoneId,
context)
+ new ColumnBatchToArrowBatchIter(colBatch, schema, maxRecordsPerBatch,
timeZoneId, context)
}
}
diff --git
a/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala
b/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala
index d615846c..640bb877 100644
---
a/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala
+++
b/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala
@@ -19,8 +19,6 @@
package org.apache.spark.sql.comet
-import scala.collection.JavaConverters._
-
import org.apache.spark.TaskContext
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.rdd.RDD
@@ -65,31 +63,44 @@ case class CometSparkToColumnarExec(child: SparkPlan)
val timeZoneId = conf.sessionLocalTimeZone
val schema = child.schema
- val rdd: RDD[InternalRow] = if (child.supportsColumnar) {
+ if (child.supportsColumnar) {
child
.executeColumnar()
.mapPartitionsInternal { iter =>
- iter.flatMap(_.rowIterator().asScala)
+ iter.flatMap { columnBatch =>
+ val context = TaskContext.get()
+ CometArrowConverters.columnarBatchToArrowBatchIter(
+ columnBatch,
+ schema,
+ maxRecordsPerBatch,
+ timeZoneId,
+ context)
+ }
+ }
+ .map { batch =>
+ numInputRows += batch.numRows()
+ numOutputBatches += 1
+ batch
}
} else {
- child.execute()
+ child
+ .execute()
+ .mapPartitionsInternal { iter =>
+ val context = TaskContext.get()
+ CometArrowConverters.rowToArrowBatchIter(
+ iter,
+ schema,
+ maxRecordsPerBatch,
+ timeZoneId,
+ context)
+ }
+ .map { batch =>
+ numInputRows += batch.numRows()
+ numOutputBatches += 1
+ batch
+ }
}
- rdd
- .mapPartitionsInternal { iter =>
- val context = TaskContext.get()
- CometArrowConverters.toArrowBatchIterator(
- iter,
- schema,
- maxRecordsPerBatch,
- timeZoneId,
- context)
- }
- .map { batch =>
- numInputRows += batch.numRows()
- numOutputBatches += 1
- batch
- }
}
override protected def withNewChildInternal(newChild: SparkPlan):
CometSparkToColumnarExec =
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
index ad963fb8..ad2242d2 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
@@ -1547,6 +1547,58 @@ class CometExecSuite extends CometTestBase {
})
}
+ test("SparkToColumnar over BatchScan (Spark Parquet reader)") {
+ Seq("", "parquet").foreach { v1List =>
+ Seq(true, false).foreach { parquetVectorized =>
+ Seq(
+ "cast(id as tinyint)",
+ "cast(id as smallint)",
+ "cast(id as integer)",
+ "cast(id as bigint)",
+ "cast(id as float)",
+ "cast(id as double)",
+ "cast(id as decimal)",
+ "cast(id as timestamp)",
+ "cast(id as string)",
+ "cast(id as binary)",
+ "struct(id)").foreach { valueType =>
+ {
+ withSQLConf(
+ SQLConf.USE_V1_SOURCE_LIST.key -> v1List,
+ CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false",
+ CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.key -> "true",
+ SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key ->
parquetVectorized.toString) {
+ withTempPath { dir =>
+ var df = spark
+ .range(10000)
+ .selectExpr("id as key", s"$valueType as value")
+ .toDF("key", "value")
+
+ df.write.parquet(dir.toString)
+
+ df = spark.read.parquet(dir.toString)
+ checkSparkAnswerAndOperator(
+ df.select("*").groupBy("key", "value").count(),
+ includeClasses = Seq(classOf[CometSparkToColumnarExec]))
+
+ // Verify that the BatchScanExec nodes supported columnar
output when requested for Spark 3.4+.
+ // Earlier versions support columnar output for fewer type.
+ if (isSpark34Plus) {
+ val leaves = df.queryExecution.executedPlan.collectLeaves()
+ if (parquetVectorized && isSpark34Plus) {
+ assert(leaves.forall(_.supportsColumnar))
+ } else {
+ assert(!leaves.forall(_.supportsColumnar))
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+
test("SparkToColumnar over InMemoryTableScanExec") {
Seq("true", "false").foreach(aqe => {
Seq("true", "false").foreach(cacheVectorized => {
diff --git
a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExecBenchmark.scala
b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExecBenchmark.scala
index 400d9b82..7205484e 100644
---
a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExecBenchmark.scala
+++
b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExecBenchmark.scala
@@ -96,6 +96,16 @@ object CometExecBenchmark extends CometBenchmarkBase {
}
}
+ benchmark.addCase("SQL Parquet - Spark (Scan), Comet (Exec)") { _ =>
+ withSQLConf(
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.key -> "true") {
+ spark.sql("select c2 + 1, c1 + 2 from parquetV1Table where c1 + 1
> 0").noop()
+ }
+ }
+
benchmark.run()
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
