address comments

facebookincubator · Dec 11, 2024 · 2d6f37e · 2d6f37e
1 parent 2762885
commit 2d6f37e
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 34 deletions.
diff --git 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
@@ -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
@@ -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