From 2d6f37ecdb07798e62dfa3787f73435ca546c284 Mon Sep 17 00:00:00 2001
From: zhli1142015 <zhli@microsoft.com>
Date: Wed, 11 Dec 2024 11:31:46 +0800
Subject: [PATCH] address comments

---
 velox/docs/functions/spark/json.rst           |  4 +-
 .../sparksql/specialforms/FromJson.cpp        | 46 ++++++-------------
 .../sparksql/specialforms/FromJson.h          |  2 +
 3 files changed, 18 insertions(+), 34 deletions(-)

diff --git a/velox/docs/functions/spark/json.rst b/velox/docs/functions/spark/json.rst
index 5ae258e6668e4..48409c3af0c33 100644
--- a/velox/docs/functions/spark/json.rst
+++ b/velox/docs/functions/spark/json.rst
@@ -23,8 +23,10 @@ JSON Functions
     supported types. For maps, the key type must be VARCHAR. When casting to
     ROW, only JSON objects are supported, and the keys in the JSON object must 
     match the field names of the ROW exactly (case-sensitive).
-    Behaviors of the casts are shown with the examples below:::
+    Behaviors of the casts are shown with the examples below. ::
 
+        SELECT from_json('{"a": true}'); -- {'a'=true} // Output type: ROW(a BOOLEAN)
         SELECT from_json('{"a": 1}'); -- {'a'=1} // Output type: ROW(a INTEGER)
+        SELECT from_json('{"a": 1.0}'); -- {'a'=1.0} // Output type: ROW(a DOUBLE)
         SELECT from_json('["name", "age", "id"]'); -- ['name', 'age', 'id'] // Output type: ARRAY(VARCHAR)
         SELECT from_json('{"a": 1, "b": 2}'); -- {'a'=1, 'b'=2} // Output type: MAP(VARCHAR,INTEGER)
diff --git a/velox/functions/sparksql/specialforms/FromJson.cpp b/velox/functions/sparksql/specialforms/FromJson.cpp
index ce92dfe94891c..f424c3e028136 100644
--- a/velox/functions/sparksql/specialforms/FromJson.cpp
+++ b/velox/functions/sparksql/specialforms/FromJson.cpp
@@ -300,15 +300,12 @@ struct ParseJsonTypedImpl {
       case simdjson::ondemand::json_type::number: {
         SIMDJSON_ASSIGN_OR_RAISE(auto num, value.get_number());
         switch (num.get_number_type()) {
-          case simdjson::ondemand::number_type::floating_point_number:
-            return simdjson::INCORRECT_TYPE;
           case simdjson::ondemand::number_type::signed_integer:
             return convertIfInRange<T>(num.get_int64(), writer);
           case simdjson::ondemand::number_type::unsigned_integer:
             return simdjson::NUMBER_OUT_OF_RANGE;
-          case simdjson::ondemand::number_type::big_integer:
-            VELOX_UNREACHABLE(); // value.get_number() would have failed
-                                 // already.
+          default:
+            return simdjson::INCORRECT_TYPE;
         }
         break;
       }
@@ -386,39 +383,22 @@ class FromJsonFunction final : public exec::VectorFunction {
       const TypePtr& outputType,
       exec::EvalCtx& context,
       VectorPtr& result) const final {
-    VELOX_USER_CHECK_EQ(args.size(), 1, "from_json expects one argument.");
     VELOX_USER_CHECK(
         args[0]->isConstantEncoding() || args[0]->isFlatEncoding(),
         "Single-arg deterministic functions receive their only argument as flat or constant vector.");
     context.ensureWritable(rows, outputType, result);
     result->clearNulls(rows);
-    switch (result->typeKind()) {
-      case TypeKind::ARRAY: {
-        if (args[0]->isConstantEncoding()) {
-          parseJsonConstant<TypeKind::ARRAY>(args[0], context, rows, *result);
-        } else {
-          parseJsonFlat<TypeKind::ARRAY>(args[0], context, rows, *result);
-        }
-        break;
-      }
-      case TypeKind::MAP: {
-        if (args[0]->isConstantEncoding()) {
-          parseJsonConstant<TypeKind::MAP>(args[0], context, rows, *result);
-        } else {
-          parseJsonFlat<TypeKind::MAP>(args[0], context, rows, *result);
-        }
-        break;
-      }
-      case TypeKind::ROW: {
-        if (args[0]->isConstantEncoding()) {
-          parseJsonConstant<TypeKind::ROW>(args[0], context, rows, *result);
-        } else {
-          parseJsonFlat<TypeKind::ROW>(args[0], context, rows, *result);
-        }
-        break;
-      }
-      default:
-        VELOX_UNSUPPORTED("Unsupported type {}.", result->type()->toString());
+    if (args[0]->isConstantEncoding()) {
+      VELOX_DYNAMIC_TYPE_DISPATCH(
+          parseJsonConstant,
+          result->typeKind(),
+          args[0],
+          context,
+          rows,
+          *result);
+    } else {
+      VELOX_DYNAMIC_TYPE_DISPATCH(
+          parseJsonFlat, result->typeKind(), args[0], context, rows, *result);
     }
   }
 
diff --git a/velox/functions/sparksql/specialforms/FromJson.h b/velox/functions/sparksql/specialforms/FromJson.h
index 073db7c0898a1..9a8c6e5efcdd8 100644
--- a/velox/functions/sparksql/specialforms/FromJson.h
+++ b/velox/functions/sparksql/specialforms/FromJson.h
@@ -18,6 +18,7 @@
 #include "velox/expression/FunctionCallToSpecialForm.h"
 
 namespace facebook::velox::functions::sparksql {
+
 class FromJsonCallToSpecialForm : public exec::FunctionCallToSpecialForm {
  public:
   // Throws not supported exception.
@@ -33,4 +34,5 @@ class FromJsonCallToSpecialForm : public exec::FunctionCallToSpecialForm {
 
   static constexpr const char* kFromJson = "from_json";
 };
+
 } // namespace facebook::velox::functions::sparksql