diff --git a/src/Interpreters/RewriteFunctionToSubcolumnVisitor.cpp b/src/Interpreters/RewriteFunctionToSubcolumnVisitor.cpp index b0647baaa7d..debbd9ae8fe 100644 --- a/src/Interpreters/RewriteFunctionToSubcolumnVisitor.cpp +++ b/src/Interpreters/RewriteFunctionToSubcolumnVisitor.cpp @@ -38,7 +38,13 @@ ASTPtr transformCountNullableToSubcolumn(const String & name_in_storage, const S return makeASTFunction("sum", makeASTFunction("not", ast)); } -const std::unordered_map> function_to_subcolumn = +ASTPtr transformMapContainsToSubcolumn(const String & name_in_storage, const String & subcolumn_name, const ASTPtr & arg) +{ + auto ast = transformToSubcolumn(name_in_storage, subcolumn_name); + return makeASTFunction("has", ast, arg); +} + +const std::unordered_map> unary_function_to_subcolumn = { {"length", {TypeIndex::Array, "size0", transformToSubcolumn}}, {"empty", {TypeIndex::Array, "size0", transformEmptyToSubcolumn}}, @@ -50,30 +56,50 @@ const std::unordered_map> binary_function_to_subcolumn +{ + {"mapContains", {TypeIndex::Map, "keys", transformMapContainsToSubcolumn}}, +}; + } void RewriteFunctionToSubcolumnData::visit(ASTFunction & function, ASTPtr & ast) const { const auto & arguments = function.arguments->children; - if (arguments.size() != 1) + if (arguments.empty() || arguments.size() > 2) return; const auto * identifier = arguments[0]->as(); if (!identifier) return; - auto it = function_to_subcolumn.find(function.name); - if (it == function_to_subcolumn.end()) - return; - - const auto & [type_id, subcolumn_name, transformer] = it->second; const auto & columns = metadata_snapshot->getColumns(); const auto & name_in_storage = identifier->name(); - if (columns.has(name_in_storage) - && columns.get(name_in_storage).type->getTypeId() == type_id) + if (!columns.has(name_in_storage)) + return; + + TypeIndex column_type_id = columns.get(name_in_storage).type->getTypeId(); + + if (arguments.size() == 1) { - ast = transformer(name_in_storage, subcolumn_name); + auto it = unary_function_to_subcolumn.find(function.name); + if (it != unary_function_to_subcolumn.end()) + { + const auto & [type_id, subcolumn_name, transformer] = it->second; + if (column_type_id == type_id) + ast = transformer(name_in_storage, subcolumn_name); + } + } + else + { + auto it = binary_function_to_subcolumn.find(function.name); + if (it != binary_function_to_subcolumn.end()) + { + const auto & [type_id, subcolumn_name, transformer] = it->second; + if (column_type_id == type_id) + ast = transformer(name_in_storage, subcolumn_name, arguments[1]); + } } } diff --git a/tests/queries/0_stateless/02115_map_contains.reference b/tests/queries/0_stateless/02115_map_contains.reference new file mode 100644 index 00000000000..975e9876237 --- /dev/null +++ b/tests/queries/0_stateless/02115_map_contains.reference @@ -0,0 +1,4 @@ +SELECT has(`m.keys`, \'a\') +FROM t_map_contains +1 +0 diff --git a/tests/queries/0_stateless/02115_map_contains.sql b/tests/queries/0_stateless/02115_map_contains.sql new file mode 100644 index 00000000000..3c7f21cb4f1 --- /dev/null +++ b/tests/queries/0_stateless/02115_map_contains.sql @@ -0,0 +1,12 @@ +DROP TABLE IF EXISTS t_map_contains; + +CREATE TABLE t_map_contains (m Map(String, UInt32)) ENGINE = Memory; + +INSERT INTO t_map_contains VALUES (map('a', 1, 'b', 2)), (map('c', 3, 'd', 4)); + +SET optimize_functions_to_subcolumns = 1; + +EXPLAIN SYNTAX SELECT mapContains(m, 'a') FROM t_map_contains; +SELECT mapContains(m, 'a') FROM t_map_contains; + +DROP TABLE t_map_contains;