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

---
 .../sparksql/specialforms/FromJson.cpp        | 30 +++++++++----------
 1 file changed, 15 insertions(+), 15 deletions(-)

diff --git a/velox/functions/sparksql/specialforms/FromJson.cpp b/velox/functions/sparksql/specialforms/FromJson.cpp
index f424c3e02813..0d59ee585816 100644
--- a/velox/functions/sparksql/specialforms/FromJson.cpp
+++ b/velox/functions/sparksql/specialforms/FromJson.cpp
@@ -375,6 +375,7 @@ struct ParseJsonTypedImpl {
 /// - Map: Keys must be `VARCHAR` type.
 /// - Row: Partial parsing is supported, but JSON arrays cannot be parsed into a
 /// ROW type.
+template <TypeKind kind>
 class FromJsonFunction final : public exec::VectorFunction {
  public:
   void apply(
@@ -389,21 +390,13 @@ class FromJsonFunction final : public exec::VectorFunction {
     context.ensureWritable(rows, outputType, result);
     result->clearNulls(rows);
     if (args[0]->isConstantEncoding()) {
-      VELOX_DYNAMIC_TYPE_DISPATCH(
-          parseJsonConstant,
-          result->typeKind(),
-          args[0],
-          context,
-          rows,
-          *result);
+      parseJsonConstant(args[0], context, rows, *result);
     } else {
-      VELOX_DYNAMIC_TYPE_DISPATCH(
-          parseJsonFlat, result->typeKind(), args[0], context, rows, *result);
+      parseJsonFlat(args[0], context, rows, *result);
     }
   }
 
  private:
-  template <TypeKind kind>
   void parseJsonConstant(
       VectorPtr& input,
       exec::EvalCtx& context,
@@ -432,7 +425,7 @@ class FromJsonFunction final : public exec::VectorFunction {
       context.applyToSelectedNoThrow(rows, [&](auto row) {
         writer.setOffset(row);
         if (error != simdjson::SUCCESS ||
-            paseJsonOneRow<kind>(jsonDoc, writer) != simdjson::SUCCESS) {
+            paseJsonOneRow(jsonDoc, writer) != simdjson::SUCCESS) {
           writer.commitNull();
         }
       });
@@ -441,7 +434,6 @@ class FromJsonFunction final : public exec::VectorFunction {
     writer.finish();
   }
 
-  template <TypeKind kind>
   void parseJsonFlat(
       VectorPtr& input,
       exec::EvalCtx& context,
@@ -473,14 +465,13 @@ class FromJsonFunction final : public exec::VectorFunction {
       simdjson::ondemand::document doc;
       auto error = simdjsonParse(paddedInput).get(doc);
       if (error != simdjson::SUCCESS ||
-          paseJsonOneRow<kind>(doc, writer) != simdjson::SUCCESS) {
+          paseJsonOneRow(doc, writer) != simdjson::SUCCESS) {
         writer.commitNull();
       }
     });
     writer.finish();
   }
 
-  template <TypeKind kind>
   static simdjson::error_code paseJsonOneRow(
       simdjson::ondemand::document& doc,
       exec::VectorWriter<Any>& writer) {
@@ -560,10 +551,19 @@ exec::ExprPtr FromJsonCallToSpecialForm::constructSpecialForm(
     VELOX_UNSUPPORTED("Unsupported type {}.", type->toString());
   }
 
+  std::shared_ptr<exec::VectorFunction> func;
+  if (type->kind() == TypeKind::ARRAY) {
+    func = std::make_shared<FromJsonFunction<TypeKind::ARRAY>>();
+  } else if (type->kind() == TypeKind::MAP) {
+    func = std::make_shared<FromJsonFunction<TypeKind::MAP>>();
+  } else {
+    func = std::make_shared<FromJsonFunction<TypeKind::ROW>>();
+  }
+
   return std::make_shared<exec::Expr>(
       type,
       std::move(args),
-      std::make_shared<FromJsonFunction>(),
+      func,
       exec::VectorFunctionMetadata{},
       kFromJson,
       trackCpuUsage);