Merge pull request #12179 from azat/GROUP-BY-injective-elimination-dictGet-fixes

alexey-milovidov · web-flow · commit 0c37fe9c7573 · 2020-07-09T04:29:26.000+03:00
Fix dictGet arguments check during GROUP BY injective functions elimination
diff --git a/src/Functions/FunctionsExternalDictionaries.h b/src/Functions/FunctionsExternalDictionaries.h
@@ -99,6 +99,10 @@ class FunctionDictHelper
 
     bool isDictGetFunctionInjective(const Block & sample_block)
     {
+        /// Assume non-injective by default
+        if (!sample_block)
+            return false;
+
         if (sample_block.columns() != 3 && sample_block.columns() != 4)
             throw Exception{"Function dictGet... takes 3 or 4 arguments", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH};
 
diff --git a/src/Functions/IFunction.h b/src/Functions/IFunction.h
@@ -133,6 +133,10 @@ class IFunctionBase
       * But we assume, that it is injective. This could be documented as implementation-specific behaviour.
       *
       * sample_block should contain data types of arguments and values of constants, if relevant.
+      * NOTE: to check is function injective with any arguments, you can pass
+      *       empty block as sample_block (since most of the time function will
+      *       ignore it anyway, and creating arguments just for checking is
+      *       function injective or not is overkill).
       */
     virtual bool isInjective(const Block & /*sample_block*/) const { return false; }
 
diff --git a/src/Interpreters/SyntaxAnalyzer.cpp b/src/Interpreters/SyntaxAnalyzer.cpp
@@ -248,6 +248,7 @@ void executeScalarSubqueries(ASTPtr & query, const Context & context, size_t sub
 
 const std::unordered_set<String> possibly_injective_function_names
 {
+        "dictGet",
         "dictGetString",
         "dictGetUInt8",
         "dictGetUInt16",
@@ -327,10 +328,18 @@ void optimizeGroupBy(ASTSelectQuery * select_query, const NameSet & source_colum
                     continue;
                 }
 
-                const auto & dict_name = function->arguments->children[0]->as<ASTLiteral &>().value.safeGet<String>();
-                const auto & dict_ptr = context.getExternalDictionariesLoader().getDictionary(dict_name);
-                const auto & attr_name = function->arguments->children[1]->as<ASTLiteral &>().value.safeGet<String>();
+                const auto * dict_name_ast = function->arguments->children[0]->as<ASTLiteral>();
+                const auto * attr_name_ast = function->arguments->children[1]->as<ASTLiteral>();
+                if (!dict_name_ast || !attr_name_ast)
+                {
+                    ++i;
+                    continue;
+                }
+
+                const auto & dict_name = dict_name_ast->value.safeGet<String>();
+                const auto & attr_name = attr_name_ast->value.safeGet<String>();
 
+                const auto & dict_ptr = context.getExternalDictionariesLoader().getDictionary(dict_name);
                 if (!dict_ptr->isInjective(attr_name))
                 {
                     ++i;
diff --git a/tests/queries/0_stateless/01375_GROUP_BY_injective_elimination_dictGet_BAD_ARGUMENTS.reference b/tests/queries/0_stateless/01375_GROUP_BY_injective_elimination_dictGet_BAD_ARGUMENTS.reference
diff --git a/tests/queries/0_stateless/01375_GROUP_BY_injective_elimination_dictGet_BAD_ARGUMENTS.sql b/tests/queries/0_stateless/01375_GROUP_BY_injective_elimination_dictGet_BAD_ARGUMENTS.sql
@@ -0,0 +1 @@
+SELECT dictGetString(concat('default', '.countryId'), 'country', toUInt64(number)) AS country FROM numbers(2) GROUP BY country; -- { serverError 36; }
diff --git a/tests/queries/0_stateless/01376_GROUP_BY_injective_elimination_dictGet.reference b/tests/queries/0_stateless/01376_GROUP_BY_injective_elimination_dictGet.reference
@@ -0,0 +1 @@
+1.1
diff --git a/tests/queries/0_stateless/01376_GROUP_BY_injective_elimination_dictGet.sql b/tests/queries/0_stateless/01376_GROUP_BY_injective_elimination_dictGet.sql
@@ -0,0 +1,31 @@
+-- https://github.com/ClickHouse/ClickHouse/issues/11469
+SELECT dictGet('default.countryId', 'country', toUInt64(number)) AS country FROM numbers(2) GROUP BY country; -- { serverError 36; }
+
+
+-- with real dictionary
+DROP TABLE IF EXISTS dictdb_01376.table_for_dict;
+DROP DICTIONARY IF EXISTS dictdb_01376.dict_exists;
+DROP DATABASE IF EXISTS dictdb_01376;
+
+CREATE DATABASE dictdb_01376 ENGINE = Ordinary;
+
+CREATE TABLE dictdb_01376.table_for_dict
+(
+  key_column UInt64,
+  value Float64
+)
+ENGINE = Memory();
+
+INSERT INTO dictdb_01376.table_for_dict VALUES (1, 1.1);
+
+CREATE DICTIONARY IF NOT EXISTS dictdb_01376.dict_exists
+(
+  key_column UInt64,
+  value Float64 DEFAULT 77.77
+)
+PRIMARY KEY key_column
+SOURCE(CLICKHOUSE(HOST 'localhost' PORT 9000 USER 'default' TABLE 'table_for_dict' DB 'dictdb_01376'))
+LIFETIME(1)
+LAYOUT(FLAT());
+
+SELECT dictGet('dictdb_01376.dict_exists', 'value', toUInt64(1)) as val FROM numbers(2) GROUP BY val;

Original file line number	Diff line number	Diff line change
`@@ -99,6 +99,10 @@ class FunctionDictHelper`
`99`	`99`
`100`	`100`	`bool isDictGetFunctionInjective(const Block & sample_block)`
`101`	`101`	`{`
	`102`	`+ /// Assume non-injective by default`
	`103`	`+ if (!sample_block)`
	`104`	`+ return false;`
	`105`	`+`
`102`	`106`	`if (sample_block.columns() != 3 && sample_block.columns() != 4)`
`103`	`107`	`throw Exception{"Function dictGet... takes 3 or 4 arguments", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH};`
`104`	`108`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+SELECT dictGetString(concat('default', '.countryId'), 'country', toUInt64(number)) AS country FROM numbers(2) GROUP BY country; -- { serverError 36; }`