feat: Implement ToPrettyString (#1921)

Author: andygrove

docs/source/user-guide/expressions.md

@@ -208,9 +208,10 @@ The following Spark expressions are currently available. Any known compatibility
 ## Other
 
 | Expression              | Notes                                                                           |
-| ----------------------- | ------------------------------------------------------------------------------- |
+|-------------------------| ------------------------------------------------------------------------------- |
 | Cast                    | See compatibility guide for list of supported cast expressions and known issues |
 | BloomFilterMightContain |                                                                                 |
 | ScalarSubquery          |                                                                                 |
 | Coalesce                |                                                                                 |
 | NormalizeNaNAndZero     |                                                                                 |
+| ToPrettyString          |                                                                                 |

native/core/src/execution/planner.rs

@@ -739,6 +739,25 @@ impl PhysicalPlanner {
                 let child = self.create_expr(expr.child.as_ref().unwrap(), input_schema)?;
                 Ok(Arc::new(ToJson::new(child, &expr.timezone)))
             }
+            ExprStruct::ToPrettyString(expr) => {
+                let mut spark_cast_options =
+                    SparkCastOptions::new(EvalMode::Try, &expr.timezone, true);
+                let null_string = "NULL";
+                spark_cast_options.null_string = null_string.to_string();
+                let child = self.create_expr(expr.child.as_ref().unwrap(), input_schema)?;
+                let cast = Arc::new(Cast::new(
+                    Arc::clone(&child),
+                    DataType::Utf8,
+                    spark_cast_options,
+                ));
+                Ok(Arc::new(IfExpr::new(
+                    Arc::new(IsNullExpr::new(child)),
+                    Arc::new(Literal::new(ScalarValue::Utf8(Some(
+                        null_string.to_string(),
+                    )))),
+                    cast,
+                )))
+            }
             ExprStruct::ListExtract(expr) => {
                 let child =
                     self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&input_schema))?;

native/proto/src/proto/expr.proto

@@ -82,6 +82,7 @@ message Expr {
     GetArrayStructFields get_array_struct_fields = 57;
     ArrayInsert array_insert = 58;
     MathExpr integral_divide = 59;
+    ToPrettyString to_pretty_string = 60;
   }
 }
 
@@ -282,6 +283,11 @@ message ToJson {
   bool ignore_null_fields = 6;
 }
 
+message ToPrettyString {
+  Expr child = 1;
+  string timezone = 2;
+}
+
 message Hour {
   Expr child = 1;
   string timezone = 2;

native/spark-expr/src/conversion_funcs/cast.rs

@@ -812,6 +812,8 @@ pub struct SparkCastOptions {
     /// We also use the cast logic for adapting Parquet schemas, so this flag is used
     /// for that use case
     pub is_adapting_schema: bool,
+    /// String to use to represent null values
+    pub null_string: String,
 }
 
 impl SparkCastOptions {
@@ -822,6 +824,7 @@ impl SparkCastOptions {
             allow_incompat,
             allow_cast_unsigned_ints: false,
             is_adapting_schema: false,
+            null_string: "null".to_string(),
         }
     }
 
@@ -832,6 +835,7 @@ impl SparkCastOptions {
             allow_incompat,
             allow_cast_unsigned_ints: false,
             is_adapting_schema: false,
+            null_string: "null".to_string(),
         }
     }
 }
@@ -1141,7 +1145,7 @@ fn casts_struct_to_string(
                     str.push_str(", ");
                 }
                 if field.is_null(row_index) {
-                    str.push_str("null");
+                    str.push_str(&spark_cast_options.null_string);
                 } else {
                     str.push_str(field.value(row_index));
                 }

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -850,6 +850,42 @@ object QueryPlanSerde extends Logging with CometExprShim {
           None
         }
 
+      // ToPrettyString is new in Spark 3.5
+      case _
+          if expr.getClass.getSimpleName == "ToPrettyString" && expr
+            .isInstanceOf[UnaryExpression] && expr.isInstanceOf[TimeZoneAwareExpression] =>
+        val child = expr.asInstanceOf[UnaryExpression].child
+        val timezoneId = expr.asInstanceOf[TimeZoneAwareExpression].timeZoneId
+
+        handleCast(
+          expr,
+          child,
+          inputs,
+          binding,
+          DataTypes.StringType,
+          timezoneId,
+          CometEvalMode.TRY) match {
+          case Some(_) =>
+            exprToProtoInternal(child, inputs, binding) match {
+              case Some(p) =>
+                val toPrettyString = ExprOuterClass.ToPrettyString
+                  .newBuilder()
+                  .setChild(p)
+                  .setTimezone(timezoneId.getOrElse("UTC"))
+                  .build()
+                Some(
+                  ExprOuterClass.Expr
+                    .newBuilder()
+                    .setToPrettyString(toPrettyString)
+                    .build())
+              case _ =>
+                withInfo(expr, child)
+                None
+            }
+          case None =>
+            None
+        }
+
       case StructsToJson(options, child, timezoneId) =>
         if (options.nonEmpty) {
           withInfo(expr, "StructsToJson with options is not supported")

spark/src/test/spark-3.5/org/apache/spark/sql/CometToPrettyStringSuite.scala

@@ -0,0 +1,94 @@
+/*
+ * 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.spark.sql
+
+import org.apache.comet.CometConf
+import org.apache.comet.expressions.{CometCast, CometEvalMode, Compatible}
+import org.apache.comet.testing.{DataGenOptions, ParquetGenerator}
+import org.apache.commons.io.FileUtils
+import org.apache.spark.sql.catalyst.TableIdentifier
+import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
+import org.apache.spark.sql.catalyst.expressions.{Alias, ToPrettyString}
+import org.apache.spark.sql.catalyst.plans.logical.Project
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.DataTypes
+
+import java.io.File
+import java.text.SimpleDateFormat
+import scala.util.Random
+
+class CometToPrettyStringSuite extends CometTestBase {
+
+  private var filename: String = null
+
+  /**
+   * We use Asia/Kathmandu because it has a non-zero number of minutes as the offset, so is an
+   * interesting edge case. Also, this timezone tends to be different from the default system
+   * timezone.
+   *
+   * Represents UTC+5:45
+   */
+  private val defaultTimezone = "Asia/Kathmandu"
+
+  override def beforeAll(): Unit = {
+    super.beforeAll()
+    val tempDir = System.getProperty("java.io.tmpdir")
+    filename = s"$tempDir/CometFuzzTestSuite_${System.currentTimeMillis()}.parquet"
+    val random = new Random(42)
+    withSQLConf(
+      CometConf.COMET_ENABLED.key -> "false",
+      SQLConf.SESSION_LOCAL_TIMEZONE.key -> defaultTimezone) {
+      val options =
+        DataGenOptions(
+          generateArray = true,
+          generateStruct = true,
+          generateMap = true,
+          generateNegativeZero = false,
+          // override base date due to known issues with experimental scans
+          baseDate =
+            new SimpleDateFormat("YYYY-MM-DD hh:mm:ss").parse("2024-05-25 12:34:56").getTime)
+      ParquetGenerator.makeParquetFile(random, spark, filename, 1000, options)
+    }
+  }
+
+  protected override def afterAll(): Unit = {
+    super.afterAll()
+    FileUtils.deleteDirectory(new File(filename))
+  }
+
+  test("ToPrettyString") {
+    val df = spark.read.parquet(filename)
+    df.createOrReplaceTempView("t1")
+    val table = spark.sessionState.catalog.lookupRelation(TableIdentifier("t1"))
+
+    for (field <- df.schema.fields) {
+      val col = field.name
+      val prettyExpr = Alias(ToPrettyString(UnresolvedAttribute(col)), s"pretty_$col")()
+      val plan = Project(Seq(prettyExpr), table)
+      val analyzed = spark.sessionState.analyzer.execute(plan)
+      val result: DataFrame = Dataset.ofRows(spark, analyzed)
+      CometCast.isSupported(field.dataType, DataTypes.StringType, Some(spark.sessionState.conf.sessionLocalTimeZone), CometEvalMode.TRY) match {
+        case _: Compatible => checkSparkAnswerAndOperator(result)
+        case _ => checkSparkAnswer(result)
+      }
+    }
+  }
+
+}