From 39322cdbe22c23e62d3456dc5eab678aec5d4700 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=87=8C=E6=B6=9B?= Date: Thu, 28 Oct 2021 17:15:32 +0800 Subject: [PATCH] add a new function mapExtractKeyLike --- .../functions/tuple-map-functions.md | 38 +++++ src/Functions/map.cpp | 136 ++++++++++++++++++ ...02111_function_mapExtractKeyLike.reference | 23 +++ .../02111_function_mapExtractKeyLike.sql | 22 +++ 4 files changed, 219 insertions(+) create mode 100644 tests/queries/0_stateless/02111_function_mapExtractKeyLike.reference create mode 100644 tests/queries/0_stateless/02111_function_mapExtractKeyLike.sql diff --git a/docs/en/sql-reference/functions/tuple-map-functions.md b/docs/en/sql-reference/functions/tuple-map-functions.md index 46ce350377c..45cee5d24bc 100644 --- a/docs/en/sql-reference/functions/tuple-map-functions.md +++ b/docs/en/sql-reference/functions/tuple-map-functions.md @@ -390,5 +390,43 @@ Result: └─────────────────────────────┘ ``` +## mapExtractKeyLike {#mapExtractKeyLike} + +**Syntax** + +```sql +mapExtractKeyLike(map, pattern) +``` + +**Parameters** + +- `map` — Map. [Map](../../sql-reference/data-types/map.md). +- `pattern` - String pattern to match. + +**Returned value** + +- A map contained elements the key of which matchs the specified pattern. If there are no elements matched the pattern, it will return an empty map. + +**Example** + +Query: + +```sql +CREATE TABLE test (a Map(String,String)) ENGINE = Memory; + +INSERT INTO test VALUES ({'abc':'abc','def':'def'}), ({'hij':'hij','klm':'klm'}); + +SELECT mapExtractKeyLike(a, 'a%') FROM test; +``` + +Result: + +```text +┌─mapExtractKeyLike(a, 'a%')─┐ +│ {'abc':'abc'} │ +│ {} │ +└────────────────────────────┘ +``` + [Original article](https://clickhouse.com/docs/en/sql-reference/functions/tuple-map-functions/) diff --git a/src/Functions/map.cpp b/src/Functions/map.cpp index edb0c28c980..029c4f6f80e 100644 --- a/src/Functions/map.cpp +++ b/src/Functions/map.cpp @@ -382,6 +382,141 @@ public: bool useDefaultImplementationForConstants() const override { return true; } }; +class FunctionExtractKeyLike : public IFunction +{ +public: + static constexpr auto name = "mapExtractKeyLike"; + static FunctionPtr create(ContextPtr) { return std::make_shared(); } + + String getName() const override + { + return name; + } + + bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*info*/) const override { return true; } + + size_t getNumberOfArguments() const override { return 2; } + + DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override + { + if (arguments.size() != 2) + throw Exception("Number of arguments for function " + getName() + " doesn't match: passed " + + toString(arguments.size()) + ", should be 2", + ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH); + + + const DataTypeMap * map_type = checkAndGetDataType(arguments[0].type.get()); + + if (!map_type) + throw Exception{"First argument for function " + getName() + " must be a map", + ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT}; + + + auto key_type = map_type->getKeyType(); + + WhichDataType which(key_type); + + if (!which.isStringOrFixedString()) + throw Exception{"Function " + getName() + "only support the map with String or FixedString key", + ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT}; + + if (!isStringOrFixedString(arguments[1].type)) + throw Exception{"Second argument passed to function " + getName() + " must be String or FixedString", + ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT}; + + return std::make_shared(map_type->getKeyType(), map_type->getValueType()); + } + + bool useDefaultImplementationForConstants() const override { return true; } + + ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t input_rows_count) const override + { + bool is_const = isColumnConst(*arguments[0].column); + const ColumnMap * col_map = typeid_cast(arguments[0].column.get()); + + //It may not be necessary to check this condition, cause it will be checked in getReturnTypeImpl function + if (!col_map) + return nullptr; + + const DataTypeMap * map_type = checkAndGetDataType(arguments[0].type.get()); + auto key_type = map_type->getKeyType(); + auto value_type = map_type->getValueType(); + + const auto & nested_column = col_map->getNestedColumn(); + const auto & values_column = col_map->getNestedData().getColumn(1); + const ColumnString * keys_string_column = checkAndGetColumn(col_map->getNestedData().getColumn(0)); + const ColumnFixedString * keys_fixed_string_column = checkAndGetColumn(col_map->getNestedData().getColumn(0)); + + FunctionLike func_like; + + //create result data + MutableColumnPtr keys_data = key_type->createColumn(); + MutableColumnPtr values_data = value_type->createColumn(); + MutableColumnPtr offsets = DataTypeNumber().createColumn(); + + IColumn::Offset current_offset = 0; + + for (size_t row = 0; row < input_rows_count; row++) + { + size_t element_start_row = row != 0 ? nested_column.getOffsets()[row-1] : 0; + size_t element_size = nested_column.getOffsets()[row]- element_start_row; + + ColumnsWithTypeAndName new_arguments; + ColumnPtr sub_map_column; + DataTypePtr data_type; + + if (keys_string_column) + { + sub_map_column = keys_string_column->cut(element_start_row, element_size); + data_type = std::make_shared(); + } + else + { + sub_map_column = keys_fixed_string_column->cut(element_start_row, element_size); + data_type =std::make_shared(checkAndGetColumn(sub_map_column.get())->getN()); + } + + size_t col_key_size = sub_map_column->size(); + auto column = is_const? ColumnConst::create(std::move(sub_map_column), std::move(col_key_size)) : std::move(sub_map_column); + + new_arguments = { + { + column, + data_type, + "" + }, + arguments[1] + }; + + auto res = func_like.executeImpl(new_arguments, result_type, input_rows_count); + const auto & container = checkAndGetColumn(res.get())->getData(); + + for (size_t row_num = 0; row_num < element_size; row_num++) + { + if (container[row_num] == 1) + { + auto key_ref = keys_string_column ? + keys_string_column->getDataAt(element_start_row + row_num) : + keys_fixed_string_column->getDataAt(element_start_row + row_num); + auto value_ref = values_column.getDataAt(element_start_row + row_num); + + keys_data->insertData(key_ref.data, key_ref.size); + values_data->insertData(value_ref.data, value_ref.size); + current_offset += 1; + } + } + + offsets->insert(current_offset); + } + + auto result_nested_column = ColumnArray::create( + ColumnTuple::create(Columns{std::move(keys_data), std::move(values_data)}), + std::move(offsets)); + + return ColumnMap::create(result_nested_column); + } +}; + } void registerFunctionsMap(FunctionFactory & factory) @@ -391,6 +526,7 @@ void registerFunctionsMap(FunctionFactory & factory) factory.registerFunction(); factory.registerFunction(); factory.registerFunction(); + factory.registerFunction(); } } diff --git a/tests/queries/0_stateless/02111_function_mapExtractKeyLike.reference b/tests/queries/0_stateless/02111_function_mapExtractKeyLike.reference new file mode 100644 index 00000000000..45edbc24c75 --- /dev/null +++ b/tests/queries/0_stateless/02111_function_mapExtractKeyLike.reference @@ -0,0 +1,23 @@ +The data of table: +1 {'P1-K1':'1-V1','P2-K2':'1-V2'} +2 {'P1-K1':'2-V1','P2-K2':'2-V2'} +3 {'P1-K1':'3-V1','P2-K2':'3-V2'} +4 {'P1-K1':'4-V1','P2-K2':'4-V2'} +5 {'5-K1':'5-V1','5-K2':'5-V2'} +6 {'P3-K1':'6-V1','P4-K2':'6-V2'} + +The results of query: SELECT id, mapExtractKeyLike(map, \'P1%\') FROM map_extractKeyLike_test ORDER BY id; +1 {'P1-K1':'1-V1'} +2 {'P1-K1':'2-V1'} +3 {'P1-K1':'3-V1'} +4 {'P1-K1':'4-V1'} +5 {} +6 {} + +The results of query: SELECT id, mapExtractKeyLike(map, \'5-K1\') FROM map_extractKeyLike_test ORDER BY id; +1 {} +2 {} +3 {} +4 {} +5 {'5-K1':'5-V1'} +6 {} diff --git a/tests/queries/0_stateless/02111_function_mapExtractKeyLike.sql b/tests/queries/0_stateless/02111_function_mapExtractKeyLike.sql new file mode 100644 index 00000000000..31f53642b74 --- /dev/null +++ b/tests/queries/0_stateless/02111_function_mapExtractKeyLike.sql @@ -0,0 +1,22 @@ +DROP TABLE IF EXISTS map_extractKeyLike_test; + +CREATE TABLE map_extractKeyLike_test (id UInt32, map Map(String, String)) Engine=MergeTree() ORDER BY id settings index_granularity=2; + +INSERT INTO map_extractKeyLike_test VALUES (1, {'P1-K1':'1-V1','P2-K2':'1-V2'}),(2,{'P1-K1':'2-V1','P2-K2':'2-V2'}); +INSERT INTO map_extractKeyLike_test VALUES (3, {'P1-K1':'3-V1','P2-K2':'3-V2'}),(4,{'P1-K1':'4-V1','P2-K2':'4-V2'}); +INSERT INTO map_extractKeyLike_test VALUES (5, {'5-K1':'5-V1','5-K2':'5-V2'}),(6, {'P3-K1':'6-V1','P4-K2':'6-V2'}); + +SELECT 'The data of table:'; +SELECT * FROM map_extractKeyLike_test ORDER BY id; + +SELECT ''; + +SELECT 'The results of query: SELECT id, mapExtractKeyLike(map, \'P1%\') FROM map_extractKeyLike_test ORDER BY id;'; +SELECT id, mapExtractKeyLike(map, 'P1%') FROM map_extractKeyLike_test ORDER BY id; + +SELECT ''; + +SELECT 'The results of query: SELECT id, mapExtractKeyLike(map, \'5-K1\') FROM map_extractKeyLike_test ORDER BY id;'; +SELECT id, mapExtractKeyLike(map, '5-K1') FROM map_extractKeyLike_test ORDER BY id; + +DROP TABLE map_extractKeyLike_test;