[GLUTEN-3927][CH] Improve the performance of element_at

backends-clickhouse/src/test/scala/io/glutenproject/execution/GlutenFunctionValidateSuite.scala

@@ -485,4 +485,32 @@ class GlutenFunctionValidateSuite extends GlutenClickHouseWholeStageTransformerS
       "select * from date_table where to_date(from_unixtime(ts)) = '2019-01-01'",
       noFallBack = true) { _ => }
   }
+
+  test("test element_at function") {
+    withSQLConf(
+      SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
+        (ConstantFolding.ruleName + "," + NullPropagation.ruleName)) {
+      // input type is array<array<int>>
+      runQueryAndCompare(
+        "SELECT array(array(1,2,3), array(4,5,6))[1], " +
+          "array(array(id,id+1,id+2), array(id+3,id+4,id+5)) from range(100)",
+        noFallBack = true
+      )(checkOperatorMatch[ProjectExecTransformer])
+
+      // input type is array<array<string>>
+      runQueryAndCompare(
+        "SELECT array(array('1','2','3'), array('4','5','6'))[1], " +
+          "array(array('1','2',cast(id as string)), array('4','5',cast(id as string)))[1] " +
+          "from range(100)",
+        noFallBack = true
+      )(checkOperatorMatch[ProjectExecTransformer])
+
+      // input type is array<map<string, int>>
+      runQueryAndCompare(
+        "SELECT array(map(cast(id as string), id), map(cast(id+1 as string), id+1))[1] " +
+          "from range(100)",
+        noFallBack = true
+      )(checkOperatorMatch[ProjectExecTransformer])
+    }
+  }
 }

cpp-ch/clickhouse.version

@@ -1,3 +1,3 @@
 CH_ORG=Kyligence
 CH_BRANCH=rebase_ch/20231206
-CH_COMMIT=5a611598662
+CH_COMMIT=6bd969f2c79

cpp-ch/local-engine/Parser/scalar_function_parser/elementAt.cpp

@@ -21,27 +21,25 @@
 
 namespace local_engine
 {
-    class FunctionParserElementAt : public FunctionParserArrayElement
-    {
-    public:
-        explicit FunctionParserElementAt(SerializedPlanParser * plan_parser_) : FunctionParserArrayElement(plan_parser_) { }
-        ~FunctionParserElementAt() override = default;
-        static constexpr auto name = "element_at";
-        String getName() const override { return name; }
+class FunctionParserElementAt : public FunctionParserArrayElement
+{
+public:
+    explicit FunctionParserElementAt(SerializedPlanParser * plan_parser_) : FunctionParserArrayElement(plan_parser_) { }
+    ~FunctionParserElementAt() override = default;
+    static constexpr auto name = "element_at";
+    String getName() const override { return name; }
 
-        const ActionsDAG::Node * parse(
-        const substrait::Expression_ScalarFunction & substrait_func,
-        ActionsDAGPtr & actions_dag) const override
-        {
-            auto parsed_args = parseFunctionArguments(substrait_func, "", actions_dag);
-            if (parsed_args.size() != 2)
-                throw Exception(DB::ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Function {} requires exactly two arguments", getName());
-            if (isMap(removeNullable(parsed_args[0]->result_type)))
-                return toFunctionNode(actions_dag, "arrayElement", parsed_args);
-            else
-                return FunctionParserArrayElement::parse(substrait_func, actions_dag);
-        }
-    };
+    const ActionsDAG::Node * parse(const substrait::Expression_ScalarFunction & substrait_func, ActionsDAGPtr & actions_dag) const override
+    {
+        auto parsed_args = parseFunctionArguments(substrait_func, "", actions_dag);
+        if (parsed_args.size() != 2)
+            throw Exception(DB::ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Function {} requires exactly two arguments", getName());
+        if (isMap(removeNullable(parsed_args[0]->result_type)))
+            return toFunctionNode(actions_dag, "arrayElement", parsed_args);
+        else
+            return FunctionParserArrayElement::parse(substrait_func, actions_dag);
+    }
+};
 
-    static FunctionParserRegister<FunctionParserElementAt> register_element_at;
+static FunctionParserRegister<FunctionParserElementAt> register_element_at;
 }