WIP

Author: Kazantsev Maksim

docs/source/user-guide/latest/configs.md

@@ -324,7 +324,7 @@ These settings can be used to determine which parts of the plan are accelerated
 | `spark.comet.expression.StringTrimBoth.enabled` | Enable Comet acceleration for `StringTrimBoth` | true |
 | `spark.comet.expression.StringTrimLeft.enabled` | Enable Comet acceleration for `StringTrimLeft` | true |
 | `spark.comet.expression.StringTrimRight.enabled` | Enable Comet acceleration for `StringTrimRight` | true |
-Add| `spark.comet.expression.StructsToCsv.enabled` | Enable Comet acceleration for `StructsToCsv` | true |
+| `spark.comet.expression.StructsToCsv.enabled` | Enable Comet acceleration for `StructsToCsv` | true |
 | `spark.comet.expression.StructsToJson.enabled` | Enable Comet acceleration for `StructsToJson` | true |
 | `spark.comet.expression.Substring.enabled` | Enable Comet acceleration for `Substring` | true |
 | `spark.comet.expression.Subtract.enabled` | Enable Comet acceleration for `Subtract` | true |

native/spark-expr/src/csv_funcs/to_csv.rs

@@ -16,11 +16,11 @@
 // under the License.
 
 use arrow::array::{
-    Array, ArrayRef, BooleanArray, Int16Array, Int32Array, Int64Array, Int8Array, LargeStringArray,
-    StringArray, StringBuilder,
+    as_boolean_array, as_largestring_array, as_string_array, Array, ArrayRef, StringBuilder,
 };
 use arrow::array::{RecordBatch, StructArray};
 use arrow::datatypes::{DataType, Schema};
+use datafusion::common::cast::{as_int16_array, as_int32_array, as_int64_array, as_int8_array};
 use datafusion::common::{exec_err, Result};
 use datafusion::logical_expr::ColumnarValue;
 use datafusion::physical_expr::PhysicalExpr;
@@ -96,6 +96,10 @@ impl PhysicalExpr for ToCsv {
         Ok(DataType::Utf8)
     }
 
+    fn nullable(&self, input_schema: &Schema) -> Result<bool> {
+        self.expr.nullable(input_schema)
+    }
+
     fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
         let input_value = self.expr.evaluate(batch)?.into_array(batch.num_rows())?;
 
@@ -134,6 +138,7 @@ impl PhysicalExpr for ToCsv {
 fn struct_to_csv(array: &StructArray, delimiter: &str, null_value: &str) -> Result<ArrayRef> {
     let mut builder = StringBuilder::with_capacity(array.len(), array.len() * 16);
     let mut csv_string = String::with_capacity(array.len() * 16);
+
     for row_idx in 0..array.len() {
         if array.is_null(row_idx) {
             builder.append_null();
@@ -146,8 +151,7 @@ fn struct_to_csv(array: &StructArray, delimiter: &str, null_value: &str) -> Resu
                 if column.is_null(row_idx) {
                     csv_string.push_str(null_value);
                 } else {
-                    let value = convert_to_string(column, row_idx)?;
-                    csv_string.push_str(&value);
+                    convert_to_string(column, &mut csv_string, row_idx)?;
                 }
             }
         }
@@ -156,38 +160,40 @@ fn struct_to_csv(array: &StructArray, delimiter: &str, null_value: &str) -> Resu
     Ok(Arc::new(builder.finish()))
 }
 
-fn convert_to_string(array: &ArrayRef, row_idx: usize) -> Result<String> {
+#[inline]
+fn convert_to_string(array: &ArrayRef, csv_string: &mut String, row_idx: usize) -> Result<()> {
     match array.data_type() {
         DataType::Boolean => {
-            let array = array.as_any().downcast_ref::<BooleanArray>().unwrap();
-            Ok(array.value(row_idx).to_string())
+            let array = as_boolean_array(array);
+            csv_string.push_str(&array.value(row_idx).to_string())
         }
         DataType::Int8 => {
-            let array = array.as_any().downcast_ref::<Int8Array>().unwrap();
-            Ok(array.value(row_idx).to_string())
+            let array = as_int8_array(array)?;
+            csv_string.push_str(&array.value(row_idx).to_string())
         }
         DataType::Int16 => {
-            let array = array.as_any().downcast_ref::<Int16Array>().unwrap();
-            Ok(array.value(row_idx).to_string())
+            let array = as_int16_array(array)?;
+            csv_string.push_str(&array.value(row_idx).to_string())
         }
         DataType::Int32 => {
-            let array = array.as_any().downcast_ref::<Int32Array>().unwrap();
-            Ok(array.value(row_idx).to_string())
+            let array = as_int32_array(array)?;
+            csv_string.push_str(&array.value(row_idx).to_string())
         }
         DataType::Int64 => {
-            let array = array.as_any().downcast_ref::<Int64Array>().unwrap();
-            Ok(array.value(row_idx).to_string())
+            let array = as_int64_array(array)?;
+            csv_string.push_str(&array.value(row_idx).to_string())
         }
         DataType::Utf8 => {
-            let array = array.as_any().downcast_ref::<StringArray>().unwrap();
-            Ok(array.value(row_idx).to_string())
+            let array = as_string_array(array);
+            csv_string.push_str(&array.value(row_idx).to_string())
         }
         DataType::LargeUtf8 => {
-            let array = array.as_any().downcast_ref::<LargeStringArray>().unwrap();
-            Ok(array.value(row_idx).to_string())
+            let array = as_largestring_array(array);
+            csv_string.push_str(&array.value(row_idx).to_string())
         }
-        _ => exec_err!("to_csv not implemented for type: {:?}", array.data_type()),
+        _ => return exec_err!("to_csv not implemented for type: {:?}", array.data_type()),
     }
+    Ok(())
 }
 
 #[cfg(test)]

spark/src/test/scala/org/apache/comet/CometCsvExpressionSuite.scala

@@ -49,8 +49,9 @@ class CometCsvExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper
         val df = spark.read
           .parquet(filename)
           .select(to_csv(struct(col("c0"), col("c1"), col("c2"))))
-        checkSparkAnswerAndOperator(df)
-
+        df.explain(true)
+        df.printSchema()
+        checkSparkAnswer(df)
       }
     }
   }

spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala

@@ -22,6 +22,7 @@ package org.apache.spark.sql
 import java.util.concurrent.atomic.AtomicInteger
 
 import scala.concurrent.duration._
+import scala.jdk.CollectionConverters._
 import scala.reflect.ClassTag
 import scala.reflect.runtime.universe.TypeTag
 import scala.util.{Success, Try}
@@ -43,7 +44,7 @@ import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 import org.apache.spark.sql.internal._
 import org.apache.spark.sql.test._
-import org.apache.spark.sql.types.{DecimalType, StructType}
+import org.apache.spark.sql.types.{DecimalType, StringType, StructType}
 
 import org.apache.comet._
 import org.apache.comet.shims.ShimCometSparkSessionExtensions
@@ -119,6 +120,10 @@ abstract class CometTestBase
     if (withTol.isDefined) {
       checkAnswerWithTolerance(dfComet, expected, withTol.get)
     } else {
+      val df =
+        spark.createDataFrame(expected.toList.asJava, new StructType().add("value", StringType))
+      df.show(false)
+      df.printSchema()
       checkAnswer(dfComet, expected)
     }
 

spark/src/test/scala/org/apache/spark/sql/benchmark/CometCsvExpressionBenchmark.scala

@@ -43,9 +43,8 @@ case class CsvExprConfig(
 // spotless:off
 /**
  * Benchmark to measure performance of Comet CSV expressions. To run this benchmark:
- * `SPARK_GENERATE_BENCHMARK_FILES=1 make
- * benchmark-org.apache.spark.sql.benchmark.CometCsvExpressionBenchmark` Results will be written
- * to "spark/benchmarks/CometCsvExpressionBenchmark-**results.txt".
+ * `SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometCsvExpressionBenchmark`
+ * Results will be written to "spark/benchmarks/CometCsvExpressionBenchmark-**results.txt".
  */
 // spotless:on
 object CometCsvExpressionBenchmark extends CometBenchmarkBase {