Feat: to_json Infinity/-Infinity Nan values support

native/spark-expr/src/json_funcs/to_json.rs

@@ -235,7 +235,7 @@ fn struct_to_json(array: &StructArray, timezone: &str) -> Result<ArrayRef> {
                     json.push_str("\":");
                     // value
                     let string_value = string_arrays[col_index].value(row_index);
-                    if is_string[col_index] {
+                    if is_string[col_index] || is_infinity(string_value) || is_nan(string_value) {
                         json.push('"');
                         json.push_str(&escape_string(string_value));
                         json.push('"');
@@ -252,6 +252,14 @@ fn struct_to_json(array: &StructArray, timezone: &str) -> Result<ArrayRef> {
     Ok(Arc::new(builder.finish()))
 }
 
+fn is_infinity(input: &str) -> bool {
+    input == "Infinity" || input == "-Infinity"
+}
+
+fn is_nan(input: &str) -> bool {
+    input == "NaN"
+}
+
 #[cfg(test)]
 mod test {
     use crate::json_funcs::to_json::struct_to_json;

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

@@ -105,57 +105,37 @@ object CometGetArrayStructFields extends CometExpressionSerde[GetArrayStructFiel
 
 object CometStructsToJson extends CometExpressionSerde[StructsToJson] {
 
-  override def getIncompatibleReasons(): Seq[String] = Seq(
-    "Does not support `+Infinity` and `-Infinity` for numeric types (float, double)." +
-      " (https://github.com/apache/datafusion-comet/issues/3016)")
-
-  override def getSupportLevel(expr: StructsToJson): SupportLevel =
-    Incompatible(
-      Some(
-        "Does not support Infinity/-Infinity for numeric types" +
-          " (https://github.com/apache/datafusion-comet/issues/3016)"))
+  override def getSupportLevel(expr: StructsToJson): SupportLevel = {
+    if (expr.options.nonEmpty) {
+      return Unsupported(Some("StructsToJson with options is not supported"))
+    }
+    val dataType = expr.child.dataType
+    if (!isSupportedType(dataType)) {
+      return Unsupported(Some(s"Struct type: $dataType contains unsupported types"))
+    }
+    Compatible()
+  }
 
   override def convert(
       expr: StructsToJson,
       inputs: Seq[Attribute],
       binding: Boolean): Option[ExprOuterClass.Expr] = {
-    if (expr.options.nonEmpty) {
-      withInfo(expr, "StructsToJson with options is not supported")
-      None
-    } else {
-      val isSupported = expr.child.dataType match {
-        case s: StructType =>
-          s.fields.forall(f => isSupportedType(f.dataType))
-        case _: MapType | _: ArrayType =>
-          // Spark supports map and array in StructsToJson but this is not yet
-          // implemented in Comet
-          false
-        case _ =>
-          false
-      }
-
-      if (isSupported) {
-        exprToProtoInternal(expr.child, inputs, binding) match {
-          case Some(p) =>
-            val toJson = ExprOuterClass.ToJson
-              .newBuilder()
-              .setChild(p)
-              .setTimezone(expr.timeZoneId.getOrElse("UTC"))
-              .setIgnoreNullFields(true)
-              .build()
-            Some(
-              ExprOuterClass.Expr
-                .newBuilder()
-                .setToJson(toJson)
-                .build())
-          case _ =>
-            withInfo(expr, expr.child)
-            None
-        }
-      } else {
-        withInfo(expr, "Unsupported data type", expr.child)
+    exprToProtoInternal(expr.child, inputs, binding) match {
+      case Some(p) =>
+        val toJson = ExprOuterClass.ToJson
+          .newBuilder()
+          .setChild(p)
+          .setTimezone(expr.timeZoneId.getOrElse("UTC"))
+          .setIgnoreNullFields(true)
+          .build()
+        Some(
+          ExprOuterClass.Expr
+            .newBuilder()
+            .setToJson(toJson)
+            .build())
+      case _ =>
+        withInfo(expr, expr.child)
         None
-      }
     }
   }
 

spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala

@@ -27,7 +27,7 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.internal.types.StringTypeWithCollation
 import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, DataTypes, MapType, StringType}
 
-import org.apache.comet.CometConf
+import org.apache.comet.{CometConf, CometExplainInfo}
 import org.apache.comet.CometSparkSessionExtensions.withInfo
 import org.apache.comet.expressions.{CometCast, CometEvalMode}
 import org.apache.comet.serde.{CommonStringExprs, Compatible, ExprOuterClass, Incompatible, SupportLevel}
@@ -139,10 +139,14 @@ trait CometExprShim extends CommonStringExprs {
       case i: Invoke =>
         (i.targetObject, i.functionName, i.arguments) match {
           case (Literal(evaluator: StructsToJsonEvaluator, _), "evaluate", Seq(child)) =>
-            exprToProtoInternal(
-              StructsToJson(evaluator.options, child, evaluator.timeZoneId),
-              inputs,
-              binding)
+            val toJson = StructsToJson(evaluator.options, child, evaluator.timeZoneId)
+            val exprProto = exprToProtoInternal(toJson, inputs, binding)
+            if (exprProto.isEmpty) {
+              toJson
+                .getTagValue(CometExplainInfo.EXTENSION_INFO)
+                .foreach(reasons => i.setTagValue(CometExplainInfo.EXTENSION_INFO, reasons))
+            }
+            exprProto
           case _ => None
         }
 

spark/src/main/spark-4.1/org/apache/comet/shims/CometExprShim.scala

@@ -27,7 +27,7 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.internal.types.StringTypeWithCollation
 import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, DataTypes, MapType, StringType}
 
-import org.apache.comet.CometConf
+import org.apache.comet.{CometConf, CometExplainInfo}
 import org.apache.comet.CometSparkSessionExtensions.withInfo
 import org.apache.comet.expressions.{CometCast, CometEvalMode}
 import org.apache.comet.serde.{CommonStringExprs, Compatible, ExprOuterClass, Incompatible, SupportLevel}
@@ -138,10 +138,14 @@ trait CometExprShim extends CommonStringExprs {
       case i: Invoke =>
         (i.targetObject, i.functionName, i.arguments) match {
           case (Literal(evaluator: StructsToJsonEvaluator, _), "evaluate", Seq(child)) =>
-            exprToProtoInternal(
-              StructsToJson(evaluator.options, child, evaluator.timeZoneId),
-              inputs,
-              binding)
+            val toJson = StructsToJson(evaluator.options, child, evaluator.timeZoneId)
+            val exprProto = exprToProtoInternal(toJson, inputs, binding)
+            if (exprProto.isEmpty) {
+              toJson
+                .getTagValue(CometExplainInfo.EXTENSION_INFO)
+                .foreach(reasons => i.setTagValue(CometExplainInfo.EXTENSION_INFO, reasons))
+            }
+            exprProto
           case _ => None
         }
 

spark/src/main/spark-4.2/org/apache/comet/shims/CometExprShim.scala

@@ -27,7 +27,7 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.internal.types.StringTypeWithCollation
 import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, DataTypes, MapType, StringType}
 
-import org.apache.comet.CometConf
+import org.apache.comet.{CometConf, CometExplainInfo}
 import org.apache.comet.CometSparkSessionExtensions.withInfo
 import org.apache.comet.expressions.{CometCast, CometEvalMode}
 import org.apache.comet.serde.{CommonStringExprs, Compatible, ExprOuterClass, Incompatible, SupportLevel}
@@ -138,10 +138,14 @@ trait CometExprShim extends CommonStringExprs {
       case i: Invoke =>
         (i.targetObject, i.functionName, i.arguments) match {
           case (Literal(evaluator: StructsToJsonEvaluator, _), "evaluate", Seq(child)) =>
-            exprToProtoInternal(
-              StructsToJson(evaluator.options, child, evaluator.timeZoneId),
-              inputs,
-              binding)
+            val toJson = StructsToJson(evaluator.options, child, evaluator.timeZoneId)
+            val exprProto = exprToProtoInternal(toJson, inputs, binding)
+            if (exprProto.isEmpty) {
+              toJson
+                .getTagValue(CometExplainInfo.EXTENSION_INFO)
+                .foreach(reasons => i.setTagValue(CometExplainInfo.EXTENSION_INFO, reasons))
+            }
+            exprProto
           case _ => None
         }
 

spark/src/test/resources/sql-tests/expressions/struct/structs_to_json.sql

@@ -16,14 +16,20 @@
 -- under the License.
 
 statement
-CREATE TABLE test_to_json(a int, b string) USING parquet
+CREATE TABLE test_to_json(a int, b string, f float, d double) USING parquet
 
 statement
-INSERT INTO test_to_json VALUES (1, 'hello'), (NULL, NULL), (0, '')
+INSERT INTO test_to_json VALUES (1, 'hello', cast('NaN' as float), cast('Infinity' as double)), (NULL, NULL, NULL, NULL), (0, '', 0.0, 0.0)
 
-query spark_answer_only
-SELECT to_json(named_struct('a', a, 'b', b)) FROM test_to_json
+query
+SELECT to_json(named_struct('a', a, 'b', b, 'f', f, 'd', d)) FROM test_to_json
 
 -- literal arguments
-query spark_answer_only
+query
 SELECT to_json(named_struct('a', 1, 'b', 'hello'))
+
+query expect_fallback(StructsToJson with options is not supported)
+SELECT to_json(named_struct('a', a, 'b', b), map('timestampFormat', 'dd/MM/yyyy')) FROM test_to_json
+
+query expect_fallback(Struct type: StructType(StructField(a,IntegerType,true),StructField(b,ArrayType(StringType,true),false)) contains unsupported types)
+SELECT to_json(named_struct('a', a, 'b', array(b))) FROM test_to_json

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

@@ -58,7 +58,7 @@ class CometJsonExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelpe
           filename,
           100,
           SchemaGenOptions(generateArray = false, generateStruct = false, generateMap = false),
-          DataGenOptions(generateNaN = false, generateInfinity = false))
+          DataGenOptions(generateNaN = true, generateInfinity = true))
       }
       val table = spark.read.parquet(filename)
       val fieldsNames = table.schema.fields
@@ -70,6 +70,19 @@ class CometJsonExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelpe
     }
   }
 
+  test("to_json - fallback reasons") {
+    withTable("t") {
+      sql("CREATE TABLE t(a INT, b STRING) USING parquet")
+      sql("INSERT INTO t VALUES (1, 'hello')")
+      checkSparkAnswerAndFallbackReason(
+        "SELECT to_json(named_struct('a', a, 'b', b), map('timestampFormat', 'dd/MM/yyyy')) FROM t",
+        "StructsToJson with options is not supported")
+      checkSparkAnswerAndFallbackReason(
+        "SELECT to_json(named_struct('b', array(b))) FROM t",
+        "Struct type: StructType(StructField(b,ArrayType(StringType,true),false)) contains unsupported types")
+    }
+  }
+
   test("from_json - basic primitives") {
     Seq(true, false).foreach { dictionaryEnabled =>
       withParquetTable(