(datafusion-comet) branch main updated: feat: Add config option to enable converting CSV to columnar (#871)

[ag...@apache.org]

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 3d1dbb03 feat: Add config option to enable converting CSV to columnar (#871)
3d1dbb03 is described below

commit 3d1dbb03eff2ead2f863371b6d93c9073f3619b4
Author: Andy Grove <ag...@apache.org>
AuthorDate: Tue Aug 27 08:38:29 2024 -0600

    feat: Add config option to enable converting CSV to columnar (#871)
    
    * CSV
    
    * save progress
    
    * test passes
    
    * Revert additional json test
    
    * fix path
    
    * revert a change
    
    * update docs
    
    * scalastyle
---
 common/src/main/scala/org/apache/comet/CometConf.scala   |  9 +++++++++
 docs/source/user-guide/configs.md                        |  1 +
 docs/source/user-guide/datasources.md                    |  5 +++++
 .../org/apache/comet/CometSparkSessionExtensions.scala   |  4 ++++
 spark/src/test/resources/test-data/csv-test-1.csv        |  3 +++
 .../scala/org/apache/comet/exec/CometExecSuite.scala     | 16 ++++++++++++++++
 6 files changed, 38 insertions(+)

diff --git a/common/src/main/scala/org/apache/comet/CometConf.scala b/common/src/main/scala/org/apache/comet/CometConf.scala
index ee429b80..b02f526e 100644
--- a/common/src/main/scala/org/apache/comet/CometConf.scala
+++ b/common/src/main/scala/org/apache/comet/CometConf.scala
@@ -92,6 +92,15 @@ object CometConf extends ShimCometConf {
       .booleanConf
       .createWithDefault(false)
 
+  val COMET_CONVERT_FROM_CSV_ENABLED: ConfigEntry[Boolean] =
+    conf("spark.comet.convert.csv.enabled")
+      .doc(
+        "When enabled, data from Spark (non-native) CSV v1 and v2 scans will be converted to " +
+          "Arrow format. Note that to enable native vectorized execution, both this config and " +
+          "'spark.comet.exec.enabled' need to be enabled.")
+      .booleanConf
+      .createWithDefault(false)
+
   val COMET_EXEC_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.enabled")
     .doc(
       "Whether to enable Comet native vectorized execution for Spark. This controls whether " +
diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md
index d23e4c71..0679bacb 100644
--- a/docs/source/user-guide/configs.md
+++ b/docs/source/user-guide/configs.md
@@ -30,6 +30,7 @@ Comet provides the following configuration settings.
 | spark.comet.columnar.shuffle.async.max.thread.num | Maximum number of threads on an executor used for Comet async columnar shuffle. By default, this config is 100. This is the upper bound of total number of shuffle threads per executor. In other words, if the number of cores * the number of shuffle threads per task `spark.comet.columnar.shuffle.async.thread.num` is larger than this config. Comet will use this config as the number of shuffle threads per executor instead. | 100 |
 | spark.comet.columnar.shuffle.async.thread.num | Number of threads used for Comet async columnar shuffle per shuffle task. By default, this config is 3. Note that more threads means more memory requirement to buffer shuffle data before flushing to disk. Also, more threads may not always improve performance, and should be set based on the number of cores available. | 3 |
 | spark.comet.columnar.shuffle.memory.factor | Fraction of Comet memory to be allocated per executor process for Comet shuffle. Comet memory size is specified by `spark.comet.memoryOverhead` or calculated by `spark.comet.memory.overhead.factor` * `spark.executor.memory`. By default, this config is 1.0. | 1.0 |
+| spark.comet.convert.csv.enabled | When enabled, data from Spark (non-native) CSV v1 and v2 scans will be converted to Arrow format. Note that to enable native vectorized execution, both this config and 'spark.comet.exec.enabled' need to be enabled. | false |
 | spark.comet.convert.json.enabled | When enabled, data from Spark (non-native) JSON v1 and v2 scans will be converted to Arrow format. Note that to enable native vectorized execution, both this config and 'spark.comet.exec.enabled' need to be enabled. | false |
 | spark.comet.convert.parquet.enabled | When enabled, data from Spark (non-native) Parquet v1 and v2 scans will be converted to Arrow format. Note that to enable native vectorized execution, both this config and 'spark.comet.exec.enabled' need to be enabled. | false |
 | spark.comet.debug.enabled | Whether to enable debug mode for Comet. By default, this config is false. When enabled, Comet will do additional checks for debugging purpose. For example, validating array when importing arrays from JVM at native side. Note that these checks may be expensive in performance and should only be enabled for debugging purpose. | false |
diff --git a/docs/source/user-guide/datasources.md b/docs/source/user-guide/datasources.md
index 92c13f67..9607ba60 100644
--- a/docs/source/user-guide/datasources.md
+++ b/docs/source/user-guide/datasources.md
@@ -26,6 +26,11 @@ in the schema are supported. When this option is not enabled, the scan will fall
 enabling `spark.comet.convert.parquet.enabled` will immediately convert the data into Arrow format, allowing native 
 execution to happen after that, but the process may not be efficient.
 
+## CSV
+
+Comet does not provide native CSV scan, but when `spark.comet.convert.csv.enabled` is enabled, data is immediately
+converted into Arrow format, allowing native execution to happen after that.
+
 ## JSON
 
 Comet does not provide native JSON scan, but when `spark.comet.convert.json.enabled` is enabled, data is immediately
diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
index 2b2c0ebe..268264a5 100644
--- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
+++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
@@ -35,9 +35,11 @@ import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, BroadcastQueryStageExec, ShuffleQueryStageExec}
 import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec}
 import org.apache.spark.sql.execution.datasources._
+import org.apache.spark.sql.execution.datasources.csv.CSVFileFormat
 import org.apache.spark.sql.execution.datasources.json.JsonFileFormat
 import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
 import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
+import org.apache.spark.sql.execution.datasources.v2.csv.CSVScan
 import org.apache.spark.sql.execution.datasources.v2.json.JsonScan
 import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan
 import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec, ShuffleExchangeExec}
@@ -1111,6 +1113,7 @@ object CometSparkSessionExtensions extends Logging {
         // Convert Spark DS v1 scan to Arrow format
         case scan: FileSourceScanExec =>
           scan.relation.fileFormat match {
+            case _: CSVFileFormat => CometConf.COMET_CONVERT_FROM_CSV_ENABLED.get(conf)
             case _: JsonFileFormat => CometConf.COMET_CONVERT_FROM_JSON_ENABLED.get(conf)
             case _: ParquetFileFormat => CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.get(conf)
             case _ => isSparkToArrowEnabled(conf, op)
@@ -1118,6 +1121,7 @@ object CometSparkSessionExtensions extends Logging {
         // Convert Spark DS v2 scan to Arrow format
         case scan: BatchScanExec =>
           scan.scan match {
+            case _: CSVScan => CometConf.COMET_CONVERT_FROM_CSV_ENABLED.get(conf)
             case _: JsonScan => CometConf.COMET_CONVERT_FROM_JSON_ENABLED.get(conf)
             case _: ParquetScan => CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.get(conf)
             case _ => isSparkToArrowEnabled(conf, op)
diff --git a/spark/src/test/resources/test-data/csv-test-1.csv b/spark/src/test/resources/test-data/csv-test-1.csv
new file mode 100644
index 00000000..c919cbd5
--- /dev/null
+++ b/spark/src/test/resources/test-data/csv-test-1.csv
@@ -0,0 +1,3 @@
+1,2,3
+4,5,6
+7,0,8
\ No newline at end of file
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 b7839abd..bfddd74d 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
@@ -1658,6 +1658,22 @@ class CometExecSuite extends CometTestBase {
       })
   }
 
+  test("read CSV file") {
+    Seq("", "csv").foreach { v1List =>
+      withSQLConf(
+        SQLConf.USE_V1_SOURCE_LIST.key -> v1List,
+        CometConf.COMET_EXPLAIN_FALLBACK_ENABLED.key -> "true",
+        CometConf.COMET_CONVERT_FROM_CSV_ENABLED.key -> "true") {
+        spark.read
+          .csv("src/test/resources/test-data/csv-test-1.csv")
+          .createOrReplaceTempView("tbl")
+        // use a projection with an expression otherwise we end up with
+        // just the file scan
+        checkSparkAnswerAndOperator("SELECT cast(_c0 as int), _c1, _c2 FROM tbl")
+      }
+    }
+  }
+
   test("read JSON file") {
     Seq("", "json").foreach { v1List =>
       withSQLConf(


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org