[GLUTEN-3412][CH]Bug fix element_at 1st argument is map type (#3413)

* Bug fix element_at 1st element is map type * review fix
apache · Oct 24, 2023 · 277d1c4 · 277d1c4
1 parent edf5c50
commit 277d1c4
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/...ends-clickhouse/src/test/scala/io/glutenproject/execution/GlutenClickHouseTPCHSuite.scala b/...ends-clickhouse/src/test/scala/io/glutenproject/execution/GlutenClickHouseTPCHSuite.scala
@@ -474,10 +474,12 @@ class GlutenClickHouseTPCHSuite extends GlutenClickHouseTPCHAbstractSuite {
     val data_insert_sql = "insert into test_tbl_3271 values(1, 'ab')"
     val select_sql_1 = "select id, split(data, ',')[1] from test_tbl_3271 where id = 1"
     val select_sql_2 = "select id, element_at(split(data, ','), 2) from test_tbl_3271 where id = 1"
+    val select_sql_3 = "select id, element_at(map(id, data), 1) from test_tbl_3271 where id = 1"
     spark.sql(table_create_sql);
     spark.sql(data_insert_sql)
     compareResultsAgainstVanillaSpark(select_sql_1, true, { _ => })
     compareResultsAgainstVanillaSpark(select_sql_2, true, { _ => })
+    compareResultsAgainstVanillaSpark(select_sql_3, true, { _ => })
 
     spark.sql(table_drop_sql)
   }

diff --git a/cpp-ch/local-engine/Parser/scalar_function_parser/elementAt.cpp b/cpp-ch/local-engine/Parser/scalar_function_parser/elementAt.cpp
@@ -16,6 +16,8 @@
  */
 
 #include <Parser/scalar_function_parser/arrayElement.h>
+#include <DataTypes/DataTypeMap.h>
+#include <DataTypes/IDataType.h>
 
 namespace local_engine
 {
@@ -26,6 +28,19 @@ namespace local_engine
         ~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);
+        }
     };
 
     static FunctionParserRegister<FunctionParserElementAt> register_element_at;