Merge branch 'opt_nested_funcs_1' into opt_nested_funcs

BIGO · BIGO · commit b86d9e769c29 · 2025-01-21T18:06:06.000+08:00
diff --git a/cpp-ch/local-engine/Functions/SparkFunctionGetJsonObject.cpp b/cpp-ch/local-engine/Functions/SparkFunctionGetJsonObject.cpp
@@ -142,4 +142,8 @@ REGISTER_FUNCTION(FlattenJSONStringOnRequiredFunction)
 {
     factory.registerFunction<FlattenJSONStringOnRequiredFunction>();
 }
+REGISTER_FUNCTION(FlattenJSONStringOnCollapsedFunctions)
+{
+    factory.registerFunction<FlattenJSONStringOnCollapsedFunctions>();
+}
 }
diff --git a/cpp-ch/local-engine/Functions/SparkFunctionGetJsonObject.h b/cpp-ch/local-engine/Functions/SparkFunctionGetJsonObject.h
@@ -637,13 +637,6 @@ class FlattenJSONStringOnRequiredFunction : public DB::IFunction
             GetJsonObjectImpl<DB::DummyJSONParser, DB::DefaultJSONStringSerializer<DB::DummyJSONParser::Element>>>(arguments);
     }
 
-private:
-    DB::ContextPtr context;
-    /// If too many rows cannot be parsed by simdjson directly, we will normalize the json text at first;
-    mutable bool is_most_normal_json_text = true;
-    mutable size_t total_parsed_rows = 0;
-    mutable size_t total_normalized_rows = 0;
-
     template <typename JSONParser>
     bool safeParseJson(std::string_view str, JSONParser & parser, JSONParser::Element & doc) const
     {
@@ -675,6 +668,52 @@ class FlattenJSONStringOnRequiredFunction : public DB::IFunction
         return is_doc_ok;
     }
 
+    virtual void insertResultToColumn(
+        const DB::MutableColumns & res,
+        DB::SimdJSONParser::Element & document,
+        std::vector<std::shared_ptr<DB::GeneratorJSONPath<DB::SimdJSONParser>>> & generator_json_paths,
+        DB::SimdJSONParser &,
+        GetJsonObjectImpl<DB::SimdJSONParser, DB::JSONStringSerializer<DB::SimdJSONParser::Element, DB::SimdJSONElementFormatter>> & impl) const
+    {
+        for (size_t j = 0; j < res.size(); ++j)
+        {
+            generator_json_paths[j]->reinitialize();
+            if (!impl.insertResultToColumn(*res[j], document, *generator_json_paths[j], true))
+            {
+                res[j]->insertDefault();
+            }
+        }
+    }
+
+    virtual void insertResultToColumn(
+        const DB::MutableColumns & res,
+        DB::DummyJSONParser::Element & document,
+        std::vector<std::shared_ptr<DB::GeneratorJSONPath<DB::DummyJSONParser>>> & generator_json_paths,
+        DB::DummyJSONParser &,
+        GetJsonObjectImpl<DB::DummyJSONParser, DB::DefaultJSONStringSerializer<DB::DummyJSONParser::Element>> & impl) const
+    {
+        for (size_t j = 0; j < res.size(); ++j)
+        {
+            generator_json_paths[j]->reinitialize();
+            if (!impl.insertResultToColumn(*res[j], document, *generator_json_paths[j], true))
+            {
+                res[j]->insertDefault();
+            }
+        }
+    }
+
+    virtual DB::ColumnPtr makeResultToReturn(DB::MutableColumns & res) const
+    {
+        return DB::ColumnTuple::create(std::move(res));
+    }
+
+private:
+    DB::ContextPtr context;
+    /// If too many rows cannot be parsed by simdjson directly, we will normalize the json text at first;
+    mutable bool is_most_normal_json_text = true;
+    mutable size_t total_parsed_rows = 0;
+    mutable size_t total_normalized_rows = 0;
+
     template <typename JSONParser, typename Impl>
     DB::ColumnPtr innerExecuteImpl(const DB::ColumnsWithTypeAndName & arguments) const
     {
@@ -773,14 +812,7 @@ class FlattenJSONStringOnRequiredFunction : public DB::IFunction
             }
             if (document_ok)
             {
-                for (size_t j = 0; j < tuple_size; ++j)
-                {
-                    generator_json_paths[j]->reinitialize();
-                    if (!impl.insertResultToColumn(*tuple_columns[j], document, *generator_json_paths[j], true))
-                    {
-                        tuple_columns[j]->insertDefault();
-                    }
-                }
+                insertResultToColumn(tuple_columns, document, generator_json_paths, parser, impl);
             }
             else
             {
@@ -791,7 +823,123 @@ class FlattenJSONStringOnRequiredFunction : public DB::IFunction
             }
         }
 
-        return DB::ColumnTuple::create(std::move(tuple_columns));
+        return makeResultToReturn(tuple_columns);
+    }
+};
+
+class FlattenJSONStringOnCollapsedFunctions : public FlattenJSONStringOnRequiredFunction
+{
+public:
+    static constexpr auto name = "FlattenJSONStringOnCollapsedFunctions";
+    static DB::FunctionPtr create(const DB::ContextPtr & context) { return std::make_shared<FlattenJSONStringOnCollapsedFunctions>(context); }
+    explicit FlattenJSONStringOnCollapsedFunctions(DB::ContextPtr context_) : FlattenJSONStringOnRequiredFunction(context_) { }
+    ~FlattenJSONStringOnCollapsedFunctions() override = default;
+    String getName() const override { return name; }
+
+     void insertResultToColumn(
+        const DB::MutableColumns & res,
+        DB::SimdJSONParser::Element & document,
+        std::vector<std::shared_ptr<DB::GeneratorJSONPath<DB::SimdJSONParser>>> & generator_json_paths,
+        DB::SimdJSONParser & parser,
+        GetJsonObjectImpl<DB::SimdJSONParser, DB::JSONStringSerializer<DB::SimdJSONParser::Element, DB::SimdJSONElementFormatter>> & impl) const override
+    {
+        using JSONStringSerializer = DB::JSONStringSerializer<DB::SimdJSONParser::Element, DB::SimdJSONElementFormatter>;
+        insertResultToColumnImpl<DB::SimdJSONParser, JSONStringSerializer>(res, document, generator_json_paths, parser, impl);
+    }
+
+    void insertResultToColumn(
+        const DB::MutableColumns & res,
+        DB::DummyJSONParser::Element & document,
+        std::vector<std::shared_ptr<DB::GeneratorJSONPath<DB::DummyJSONParser>>> & generator_json_paths,
+        DB::DummyJSONParser & parser,
+        GetJsonObjectImpl<DB::DummyJSONParser, DB::DefaultJSONStringSerializer<DB::DummyJSONParser::Element>> & impl) const override
+    {
+        using JSONStringSerializer = DB::DefaultJSONStringSerializer<DB::DummyJSONParser::Element>;
+        insertResultToColumnImpl<DB::DummyJSONParser, JSONStringSerializer>(res, document, generator_json_paths, parser, impl);
+    }
+
+    DB::ColumnPtr makeResultToReturn(DB::MutableColumns & res) const override
+    {
+        DB::Columns final_cols;
+        final_cols.emplace_back(std::move(res[res.size() - 1]));
+        return DB::ColumnTuple::create(std::move(final_cols));
+    }
+
+private:
+    template<typename JSONParser, typename JSONStringSerializer>
+    void insertResultToColumnImpl(
+        const DB::MutableColumns & res,
+        typename JSONParser::Element & document,
+        std::vector<std::shared_ptr<DB::GeneratorJSONPath<JSONParser>>> & generator_json_paths,
+        JSONParser & parser,
+        GetJsonObjectImpl<JSONParser, JSONStringSerializer> &) const
+    {
+        size_t json_path_pos = 0;
+        size_t last_col_index = res.size() - 1;
+        if (!insertResultToColumnImpl<JSONParser, JSONStringSerializer>(*res[last_col_index], document, generator_json_paths, json_path_pos))
+        {
+            if (document.isString())
+            {
+                typename JSONParser::Element t;
+                bool parsed = safeParseJson(document.getString(), parser, t);
+                if (!parsed)
+                    res[last_col_index]->insertDefault();
+                else if(!insertResultToColumnImpl<JSONParser, JSONStringSerializer>(*res[last_col_index], t, generator_json_paths, json_path_pos))
+                {
+                    res[last_col_index]->insertDefault();
+                }
+            }
+            else
+                res[last_col_index]->insertDefault();
+        }
+    }
+
+    template<typename JSONParser, typename JSONStringSerializer>
+    bool insertResultToColumnImpl(DB::IColumn & dest, typename JSONParser::Element & root, std::vector<std::shared_ptr<DB::GeneratorJSONPath<JSONParser>>> & generator_json_paths, size_t & json_path_pos) const
+    {
+        DB::VisitorStatus status = DB::VisitorStatus::Ok;
+        bool success = false;
+        for (size_t i = json_path_pos; i < generator_json_paths.size(); ++i)
+        {
+            std::shared_ptr<DB::GeneratorJSONPath<JSONParser>> generator_json_path = generator_json_paths[i];
+            generator_json_path->reinitialize();
+            status = DB::VisitorStatus::Ok;
+            while (status != DB::VisitorStatus::Exhausted)
+            {
+                status = generator_json_path->getNextItem(root);
+                if (status == DB::VisitorStatus::Ok)
+                {
+                    success = true;
+                }
+                else if (status == DB::VisitorStatus::Error)
+                {
+                    success = false;
+                }
+            }
+            json_path_pos = i;
+            if (!success)
+            {
+                break;
+            }
+        }
+        if (!success)
+        {
+            return false;
+        }
+        DB::ColumnNullable & nullable_col_str = assert_cast<DB::ColumnNullable &>(dest);
+        DB::ColumnString * col_str = assert_cast<DB::ColumnString *>(&nullable_col_str.getNestedColumn());
+        JSONStringSerializer serializer(*col_str);
+        nullable_col_str.getNullMapData().push_back(0);
+        if (root.isString())
+        {
+            serializer.addRawString(root.getString());
+        }
+        else
+        {
+            serializer.addElement(root);
+        }
+        serializer.commit();
+        return true;
     }
 };
 
diff --git a/cpp-ch/local-engine/Parser/scalar_function_parser/getJSONObject.cpp b/cpp-ch/local-engine/Parser/scalar_function_parser/getJSONObject.cpp
@@ -59,10 +59,26 @@ class GetJSONObjectParser : public FunctionParser
         DB::ActionsDAG & actions_dag) const override
     {
         const auto & args = substrait_func.arguments();
-        if (args.size() != 2)
+        if (args.size() < 2)
         {
-            throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Function {} requires 2 arguments", getCHFunctionName(substrait_func));
+            throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Function {} requires at least 2 arguments", getCHFunctionName(substrait_func));
         }
+        std::string json_paths = "";
+        for (size_t i = 1; i < args.size(); ++i)
+        {
+            if (args[i].value().has_literal())
+            {
+                const auto & literal_expr = args[i].value().literal();
+                if (literal_expr.has_string())
+                {
+                    std::string path = literal_expr.string();
+                    json_paths += path;
+                    if (i != args.size() - 1)
+                        json_paths += "|";
+                }
+            }
+        }
+
         if (args[0].value().has_scalar_function()
             && args[0].value().scalar_function().function_reference() == SelfDefinedFunctionReference::GET_JSON_OBJECT)
         {
@@ -72,8 +88,16 @@ class GetJSONObjectParser : public FunctionParser
             {
                 const auto flatten_function_pb = args[0].value().scalar_function();
                 const auto * flatten_arg0 = parseExpression(actions_dag, flatten_function_pb.arguments(0).value());
-                const auto * flatten_arg1 = parseExpression(actions_dag, flatten_function_pb.arguments(1).value());
-                flatten_json_column_node = toFunctionNode(actions_dag, FlattenJSONStringOnRequiredFunction::name, flatten_json_column_name, {flatten_arg0, flatten_arg1});
+                if (args.size() > 2)
+                {
+                    const auto * flatten_path_arg = addColumnToActionsDAG(actions_dag, std::make_shared<DB::DataTypeString>(), json_paths);
+                    flatten_json_column_node = toFunctionNode(actions_dag, FlattenJSONStringOnCollapsedFunctions::name, flatten_json_column_name, {flatten_arg0, flatten_path_arg});
+                }
+                else
+                {
+                    const auto * flatten_arg1 = parseExpression(actions_dag, flatten_function_pb.arguments(1).value());
+                    flatten_json_column_node = toFunctionNode(actions_dag, FlattenJSONStringOnRequiredFunction::name, flatten_json_column_name, {flatten_arg0, flatten_arg1});
+                }
                 actions_dag.addOrReplaceInOutputs(*flatten_json_column_node);
             }
             return {flatten_json_column_node, parseExpression(actions_dag, args[1].value())};
@@ -85,10 +109,12 @@ class GetJSONObjectParser : public FunctionParser
     }
 
 private:
+    
     static String getFlatterJsonColumnName(const substrait::Expression & arg)
     {
         return arg.ShortDebugString();
     }
+
 };
 
 static FunctionParserRegister<GetJSONObjectParser> register_get_json_object_parser;

Original file line number	Diff line number	Diff line change
`@@ -142,4 +142,8 @@ REGISTER_FUNCTION(FlattenJSONStringOnRequiredFunction)`
`142`	`142`	`{`
`143`	`143`	`factory.registerFunction<FlattenJSONStringOnRequiredFunction>();`
`144`	`144`	`}`
	`145`	`+REGISTER_FUNCTION(FlattenJSONStringOnCollapsedFunctions)`
	`146`	`+{`
	`147`	`+ factory.registerFunction<FlattenJSONStringOnCollapsedFunctions>();`
	`148`	`+}`
`145`	`149`	`}`