address comments

facebookincubator · Jan 8, 2025 · dd458b9 · dd458b9
1 parent 48024c2
commit dd458b9
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 26 deletions.
diff --git a/velox/docs/functions/spark/json.rst b/velox/docs/functions/spark/json.rst
@@ -45,9 +45,9 @@ JSON Functions
         SELECT json_object_keys(1); -- NULL
         SELECT json_object_keys('"hello"'); -- NULL
 
-.. spark:function:: from_json(jsonString) -> [json object]
+.. spark:function:: from_json(jsonString) -> array / map / row
 
-    Casts a JSON string to an ARRAY, MAP, or ROW type, with the output type 
+    Casts ``jsonString`` to an ARRAY, MAP, or ROW type, with the output type 
     determined by the expression. Returns NULL, if the input string is unparsable.
     Supported element types include BOOLEAN, TINYINT, SMALLINT, INTEGER, BIGINT, 
     REAL, DOUBLE, VARCHAR, ARRAY, MAP and ROW. When casting to ARRAY or MAP, 
@@ -58,7 +58,7 @@ JSON Functions
 
     * Does not support user provided options.
 
-    * Does not support enablePartialResults = false.
+    * Only supports partial result mode, which requires spark configuration spark.sql.json.enablePartialResults = true.
 
     * Does not support single quotes as delimiters.  
 
@@ -71,3 +71,4 @@ JSON Functions
         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())
+        SELECT from_json('{"a": {"b": 1}}'); -- {'a'={b=1}} // Output type: ROW({"a"}, {ROW({"b"}, {INTEGER()})})
diff --git a/velox/functions/sparksql/specialforms/FromJson.cpp b/velox/functions/sparksql/specialforms/FromJson.cpp
@@ -56,13 +56,10 @@ struct ExtractJsonTypeImpl {
     apply(Input value, exec::GenericWriter& writer, bool /*isRoot*/) {
       SIMDJSON_ASSIGN_OR_RAISE(auto type, value.type());
       std::string_view s;
-      switch (type) {
-        case simdjson::ondemand::json_type::string: {
-          SIMDJSON_ASSIGN_OR_RAISE(s, value.get_string());
-          break;
-        }
-        default:
-          s = value.raw_json();
+      if (type == simdjson::ondemand::json_type::string) {
+        SIMDJSON_ASSIGN_OR_RAISE(s, value.get_string());
+      } else {
+        s = value.raw_json();
       }
       writer.castTo<Varchar>().append(s);
       return simdjson::SUCCESS;
@@ -74,16 +71,12 @@ struct ExtractJsonTypeImpl {
     static simdjson::error_code
     apply(Input value, exec::GenericWriter& writer, bool /*isRoot*/) {
       SIMDJSON_ASSIGN_OR_RAISE(auto type, value.type());
-      switch (type) {
-        case simdjson::ondemand::json_type::boolean: {
-          auto& w = writer.castTo<bool>();
-          SIMDJSON_ASSIGN_OR_RAISE(w, value.get_bool());
-          break;
-        }
-        default:
-          return simdjson::INCORRECT_TYPE;
+      if (type == simdjson::ondemand::json_type::boolean) {
+        auto& w = writer.castTo<bool>();
+        SIMDJSON_ASSIGN_OR_RAISE(w, value.get_bool());
+        return simdjson::SUCCESS;
       }
-      return simdjson::SUCCESS;
+      return simdjson::INCORRECT_TYPE;
     }
   };
 
@@ -216,13 +209,11 @@ struct ExtractJsonTypeImpl {
       if (type == simdjson::ondemand::json_type::object) {
         SIMDJSON_ASSIGN_OR_RAISE(auto object, value.get_object());
 
-        bool allFieldsAreAscii = true;
-        const auto size = rowType.size();
-        for (auto i = 0; i < size; ++i) {
-          const auto& name = rowType.nameOf(i);
-          allFieldsAreAscii &=
-              functions::stringCore::isAscii(name.data(), name.size());
-        }
+        const auto& names = rowType.names();
+        bool allFieldsAreAscii =
+            std::all_of(names.begin(), names.end(), [](const auto& name) {
+              return functions::stringCore::isAscii(name.data(), name.size());
+            });
 
         auto fieldIndices = makeFieldIndicesMap(rowType, allFieldsAreAscii);
 
@@ -349,6 +340,7 @@ struct ExtractJsonTypeImpl {
     return simdjson::SUCCESS;
   }
 
+  // Creates a map of lower case field names to their indices in the row type.
   static folly::F14FastMap<std::string, int32_t> makeFieldIndicesMap(
       const RowType& rowType,
       bool allFieldsAreAscii) {