refactor: move json_extract_scalar back to simple function

Bikramjeet Vig · facebook-github-bot · commit e10c0a35bc94 · 2025-03-11T11:11:29.000-07:00
Summary:
We initially moved json_extract_scalar to a vector function
implementation so that we can reuse the JSON parsing logic and apply
it to VARCHAR inputs. JSON parsing achieves the following for VARCHAR
input: (1) it formats spaces, (2) it sorts the keys inside objects,
(3) it flags invalid JSON, and (4) it unescapes string scalars. For
json_extract_scalar, we only care about (3) and (4) since the output
is a scalar. Therefore, we can avoid the costly operation of
rewriting the large input JSON in a normalized form and only apply
the operations we require. To achieve (3), SIMD JSON can be used to
quickly check for valid JSON, and for (4), SIMD JSON provides an API
to extract a string scalar after removing escaping.

This change moves json_extract_scalar back to simple function
implementation and ensures (3) and (4) are applied.

Differential Revision: D70978059
diff --git a/velox/functions/prestosql/JsonFunctions.cpp b/velox/functions/prestosql/JsonFunctions.cpp
@@ -541,9 +541,6 @@ class JsonParseFunction : public exec::VectorFunction {
 
 class JsonExtractFunction : public exec::VectorFunction {
  public:
-  JsonExtractFunction(bool extractScalarOnly)
-      : extractScalarOnly_(extractScalarOnly) {}
-
   void apply(
       const SelectivityVector& rows,
       std::vector<VectorPtr>& args,
@@ -564,21 +561,7 @@ class JsonExtractFunction : public exec::VectorFunction {
     context.moveOrCopyResult(localResult, rows, result);
   }
 
-  static std::vector<std::shared_ptr<exec::FunctionSignature>> signatures(
-      bool extractScalarOnly) {
-    if (extractScalarOnly) {
-      return {
-          exec::FunctionSignatureBuilder()
-              .returnType("varchar")
-              .argumentType("json")
-              .argumentType("varchar")
-              .build(),
-          exec::FunctionSignatureBuilder()
-              .returnType("varchar")
-              .argumentType("varchar")
-              .argumentType("varchar")
-              .build()};
-    }
+  static std::vector<std::shared_ptr<exec::FunctionSignature>> signatures() {
     return {
         exec::FunctionSignatureBuilder()
             .returnType("json")
@@ -612,13 +595,8 @@ class JsonExtractFunction : public exec::VectorFunction {
         auto jsonValue = json->as<ConstantVector<StringView>>()->valueAt(0);
         auto pathValue = path->as<ConstantVector<StringView>>()->valueAt(0);
         try {
-          if (extractScalarOnly_) {
-            nullResult = processJsonExtractScalar(
-                             jsonValue, pathValue, output) != simdjson::SUCCESS;
-          } else {
-            nullResult = processJsonExtract(jsonValue, pathValue, output) !=
-                simdjson::SUCCESS;
-          }
+          nullResult = processJsonExtract(jsonValue, pathValue, output) !=
+              simdjson::SUCCESS;
         } catch (const VeloxException& e) {
           if (!e.isUserError()) {
             throw;
@@ -645,39 +623,21 @@ class JsonExtractFunction : public exec::VectorFunction {
     exec::VectorWriter<Json> resultWriter;
     resultWriter.init(*flatResult);
 
-    if (extractScalarOnly_) {
-      context.applyToSelectedNoThrow(rows, [&](auto row) {
-        VELOX_DCHECK(!decodedPath->isNullAt(row));
-        resultWriter.setOffset(row);
-        std::string output;
-        if (!decodedJson->isNullAt(row) &&
-            processJsonExtractScalar(
-                decodedJson->valueAt<StringView>(row),
-                decodedPath->valueAt<StringView>(row),
-                output) == simdjson::SUCCESS) {
-          resultWriter.current() = output;
-          resultWriter.commit(true);
-        } else {
-          resultWriter.commit(false);
-        }
-      });
-    } else {
-      context.applyToSelectedNoThrow(rows, [&](auto row) {
-        VELOX_DCHECK(!decodedPath->isNullAt(row));
-        resultWriter.setOffset(row);
-        std::string output;
-        if (!decodedJson->isNullAt(row) &&
-            processJsonExtract(
-                decodedJson->valueAt<StringView>(row),
-                decodedPath->valueAt<StringView>(row),
-                output) == simdjson::SUCCESS) {
-          resultWriter.current() = output;
-          resultWriter.commit(true);
-        } else {
-          resultWriter.commit(false);
-        }
-      });
-    }
+    context.applyToSelectedNoThrow(rows, [&](auto row) {
+      VELOX_DCHECK(!decodedPath->isNullAt(row));
+      resultWriter.setOffset(row);
+      std::string output;
+      if (!decodedJson->isNullAt(row) &&
+          processJsonExtract(
+              decodedJson->valueAt<StringView>(row),
+              decodedPath->valueAt<StringView>(row),
+              output) == simdjson::SUCCESS) {
+        resultWriter.current() = output;
+        resultWriter.commit(true);
+      } else {
+        resultWriter.commit(false);
+      }
+    });
     resultWriter.finish();
   }
 
@@ -750,59 +710,6 @@ class JsonExtractFunction : public exec::VectorFunction {
     return simdjson::SUCCESS;
   }
 
-  FOLLY_ALWAYS_INLINE simdjson::error_code processJsonExtractScalar(
-      const StringView& json,
-      const StringView& jsonPath,
-      std::string& output) const {
-    bool resultPopulated = false;
-    std::optional<std::string> resultStr;
-    auto consumer = [&resultStr, &resultPopulated](auto& v) {
-      if (resultPopulated) {
-        // We should just get a single value, if we see multiple, it's an error
-        // and we should return null.
-        resultStr = std::nullopt;
-        return simdjson::SUCCESS;
-      }
-
-      resultPopulated = true;
-
-      SIMDJSON_ASSIGN_OR_RAISE(auto vtype, v.type());
-      switch (vtype) {
-        case simdjson::ondemand::json_type::boolean: {
-          SIMDJSON_ASSIGN_OR_RAISE(bool vbool, v.get_bool());
-          resultStr = vbool ? "true" : "false";
-          break;
-        }
-        case simdjson::ondemand::json_type::string: {
-          SIMDJSON_ASSIGN_OR_RAISE(resultStr, v.get_string());
-          break;
-        }
-        case simdjson::ondemand::json_type::object:
-        case simdjson::ondemand::json_type::array:
-        case simdjson::ondemand::json_type::null:
-          // Do nothing.
-          break;
-        default: {
-          SIMDJSON_ASSIGN_OR_RAISE(resultStr, simdjson::to_json_string(v));
-        }
-      }
-      return simdjson::SUCCESS;
-    };
-
-    auto& extractor = SIMDJsonExtractor::getInstance(jsonPath);
-    bool isDefinitePath = true;
-    simdjson::padded_string paddedJson(json.data(), json.size());
-    SIMDJSON_TRY(extractor.extract(paddedJson, consumer, isDefinitePath));
-
-    if (resultStr.has_value()) {
-      output = std::move(resultStr.value());
-      return simdjson::SUCCESS;
-    } else {
-      return simdjson::NO_SUCH_FIELD;
-    }
-  }
-
-  bool extractScalarOnly_{false};
   JsonParseImpl parser_;
 };
 
@@ -856,37 +763,13 @@ VELOX_DECLARE_VECTOR_FUNCTION(
     JsonFormatFunction::signatures(),
     std::make_unique<JsonFormatFunction>());
 
-VELOX_DECLARE_STATEFUL_VECTOR_FUNCTION(
-    udf_json_extract_scalar,
-    JsonExtractFunction::signatures(true),
-    [](const std::string& /*name*/,
-       const std::vector<exec::VectorFunctionArg>&,
-       const velox::core::QueryConfig&) {
-      return std::make_shared<JsonExtractFunction>(true);
-    });
-
-// Only used internally at Meta.
-VELOX_DECLARE_STATEFUL_VECTOR_FUNCTION(
-    udf_json_extract_scalar_varchar_only,
-    (std::vector<std::shared_ptr<exec::FunctionSignature>>{
-        facebook::velox::exec::FunctionSignatureBuilder()
-            .returnType("varchar")
-            .argumentType("varchar")
-            .argumentType("varchar")
-            .build()}),
-    [](const std::string& /*name*/,
-       const std::vector<exec::VectorFunctionArg>&,
-       const velox::core::QueryConfig&) {
-      return std::make_shared<JsonExtractFunction>(true);
-    });
-
 VELOX_DECLARE_STATEFUL_VECTOR_FUNCTION(
     udf_json_extract,
-    JsonExtractFunction::signatures(false),
+    JsonExtractFunction::signatures(),
     [](const std::string& /*name*/,
        const std::vector<exec::VectorFunctionArg>&,
        const velox::core::QueryConfig&) {
-      return std::make_shared<JsonExtractFunction>(false);
+      return std::make_shared<JsonExtractFunction>();
     });
 
 VELOX_DECLARE_STATEFUL_VECTOR_FUNCTION(
diff --git a/velox/functions/prestosql/JsonFunctions.h b/velox/functions/prestosql/JsonFunctions.h
@@ -169,6 +169,86 @@ struct JsonArrayGetFunction {
   }
 };
 
+// json_extract_scalar(json, json_path) -> varchar
+// Like json_extract(), but returns the result value as a string (as opposed
+// to being encoded as JSON). The value referenced by json_path must be a scalar
+// (boolean, number or string)
+template <typename T>
+struct JsonExtractScalarFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE bool call(
+      out_type<Varchar>& result,
+      const arg_type<Json>& json,
+      const arg_type<Varchar>& jsonPath) {
+    return callImpl(result, json, jsonPath) == simdjson::SUCCESS;
+  }
+
+ private:
+  FOLLY_ALWAYS_INLINE bool callImpl(
+      out_type<Varchar>& result,
+      const arg_type<Json>& json,
+      const arg_type<Varchar>& jsonPath) {
+    bool resultPopulated = false;
+    std::optional<std::string> resultStr;
+
+    auto consumer = [&resultStr, &resultPopulated](auto& v) {
+      if (resultPopulated) {
+        // We should just get a single value, if we see multiple, it's an error
+        // and we should return null.
+        resultStr = std::nullopt;
+        return simdjson::SUCCESS;
+      }
+
+      resultPopulated = true;
+
+      SIMDJSON_ASSIGN_OR_RAISE(auto vtype, v.type());
+      switch (vtype) {
+        case simdjson::ondemand::json_type::boolean: {
+          SIMDJSON_ASSIGN_OR_RAISE(bool vbool, v.get_bool());
+          resultStr = vbool ? "true" : "false";
+          break;
+        }
+        case simdjson::ondemand::json_type::string: {
+          SIMDJSON_ASSIGN_OR_RAISE(resultStr, v.get_string());
+          break;
+        }
+        case simdjson::ondemand::json_type::object:
+        case simdjson::ondemand::json_type::array:
+        case simdjson::ondemand::json_type::null:
+          // Do nothing.
+          break;
+        default: {
+          SIMDJSON_ASSIGN_OR_RAISE(resultStr, simdjson::to_json_string(v));
+        }
+      }
+      return simdjson::SUCCESS;
+    };
+
+    auto& extractor = SIMDJsonExtractor::getInstance(jsonPath);
+
+    // Check for valid json
+    simdjson::padded_string paddedJson(json.data(), json.size());
+    {
+      SIMDJSON_ASSIGN_OR_RAISE(auto jsonDoc, simdjsonParse(paddedJson));
+      simdjson::ondemand::document parsedDoc;
+      if (simdjsonParse(paddedJson).get(parsedDoc)) {
+        return false;
+      }
+    }
+
+    bool isDefinitePath = true;
+    SIMDJSON_TRY(extractor.extract(paddedJson, consumer, isDefinitePath));
+
+    if (resultStr.has_value()) {
+      result.copy_from(*resultStr);
+      return simdjson::SUCCESS;
+    } else {
+      return simdjson::NO_SUCH_FIELD;
+    }
+  }
+};
+
 template <typename T>
 struct JsonSizeFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);
diff --git a/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp b/velox/functions/prestosql/registration/JsonFunctionsRegistration.cpp
@@ -27,6 +27,11 @@ void registerJsonFunctions(const std::string& prefix) {
   registerFunction<IsJsonScalarFunction, bool, Varchar>(
       {prefix + "is_json_scalar"});
 
+  registerFunction<JsonExtractScalarFunction, Varchar, Json, Varchar>(
+      {prefix + "json_extract_scalar"});
+  registerFunction<JsonExtractScalarFunction, Varchar, Varchar, Varchar>(
+      {prefix + "json_extract_scalar"});
+
   registerFunction<JsonArrayLengthFunction, int64_t, Json>(
       {prefix + "json_array_length"});
   registerFunction<JsonArrayLengthFunction, int64_t, Varchar>(
@@ -61,9 +66,6 @@ void registerJsonFunctions(const std::string& prefix) {
 
   VELOX_REGISTER_VECTOR_FUNCTION(udf_json_extract, prefix + "json_extract");
 
-  VELOX_REGISTER_VECTOR_FUNCTION(
-      udf_json_extract_scalar, prefix + "json_extract_scalar");
-
   VELOX_REGISTER_VECTOR_FUNCTION(udf_json_format, prefix + "json_format");
 
   VELOX_REGISTER_VECTOR_FUNCTION(udf_json_parse, prefix + "json_parse");
diff --git a/velox/functions/prestosql/tests/JsonExtractScalarTest.cpp b/velox/functions/prestosql/tests/JsonExtractScalarTest.cpp
@@ -193,6 +193,19 @@ TEST_F(JsonExtractScalarTest, invalidPath) {
       jsonExtractScalar(R"({"k1":"v1)", "$.k1]"), "Invalid JSON path");
 }
 
+TEST_F(JsonExtractScalarTest, invalidJson) {
+  // Verify that we return null on invalid JSON regardless of whether the
+  // invalid section is after the target path.
+  EXPECT_EQ(jsonExtractScalar(R"({"a": "b", "c": "d})", "$.a"), std::nullopt);
+  EXPECT_EQ(jsonExtractScalar(R"([["a"], ["b]])", "$[0][0]"), std::nullopt);
+}
+
+TEST_F(JsonExtractScalarTest, escapedString) {
+  // Verify the the returned string is unescaped.
+  EXPECT_EQ(
+      jsonExtractScalar(R"({"x": {"a" : 1, "b" : "b\/c"} })", "$.x.b"), "b/c");
+}
+
 // simdjson, like Presto java, returns the large number as-is as a string,
 // without trying to convert it to an integer.
 TEST_F(JsonExtractScalarTest, overflow) {
