diff --git a/src/Dictionaries/FlatDictionary.cpp b/src/Dictionaries/FlatDictionary.cpp index c67e9686e10..380da1cf2e0 100644 --- a/src/Dictionaries/FlatDictionary.cpp +++ b/src/Dictionaries/FlatDictionary.cpp @@ -221,6 +221,49 @@ ColumnUInt8::Ptr FlatDictionary::isInHierarchy( return result; } +ColumnPtr FlatDictionary::getDescendands( + ColumnPtr key_column, + const DataTypePtr &, + size_t level) const +{ + PaddedPODArray keys_backup; + const auto & keys = getColumnVectorData(this, key_column, keys_backup); + + size_t hierarchical_attribute_index = *dict_struct.hierarchical_attribute_index; + const auto & hierarchical_attribute = attributes[hierarchical_attribute_index]; + + const UInt64 null_value = std::get(hierarchical_attribute.null_values); + const ContainerType & parent_keys = std::get>(hierarchical_attribute.arrays); + + HashMap parent_to_child; + + for (size_t i = 0; i < parent_keys.size(); ++i) + { + auto parent_key = parent_keys[i]; + parent_to_child[parent_key] = static_cast(i); + } + + auto is_key_valid_func = [&](auto & key) + { + return parent_to_child.find(key) != nullptr; + }; + + auto get_child_key_func = [&](auto & key) + { + std::optional result; + + auto it = parent_to_child.find(key); + + if (it) + result = it->getMapped(); + + return result; + }; + + auto result = getDescendandsArray(keys, null_value, level, is_key_valid_func, get_child_key_func); + return result; +} + void FlatDictionary::createAttributes() { const auto size = dict_struct.attributes.size(); diff --git a/src/Dictionaries/FlatDictionary.h b/src/Dictionaries/FlatDictionary.h index a47ac8c34d8..f4e9ef8a559 100644 --- a/src/Dictionaries/FlatDictionary.h +++ b/src/Dictionaries/FlatDictionary.h @@ -82,6 +82,11 @@ public: ColumnPtr in_key_column, const DataTypePtr & key_type) const override; + ColumnPtr getDescendands( + ColumnPtr key_column, + const DataTypePtr & key_type, + size_t level) const override; + BlockInputStreamPtr getBlockInputStream(const Names & column_names, size_t max_block_size) const override; private: diff --git a/src/Dictionaries/HierarchyDictionariesUtils.h b/src/Dictionaries/HierarchyDictionariesUtils.h index ba2a74db298..7dd0fd311bc 100644 --- a/src/Dictionaries/HierarchyDictionariesUtils.h +++ b/src/Dictionaries/HierarchyDictionariesUtils.h @@ -32,17 +32,17 @@ struct IsKeyValidFuncInterface }; template -struct GetParentKeyFuncInterface +struct GetKeyFuncInterface { std::optional operator()(T key [[maybe_unused]]) { return {}; } }; -template +template ElementsAndOffsets getKeysHierarchy( const PaddedPODArray & hierarchy_keys, const KeyType & hierarchy_null_value, IsKeyValidFunc && is_key_valid_func, - GetParentKeyFunc && get_parent_func) + GetKeyFunc && get_key_func) { size_t hierarchy_keys_size = hierarchy_keys.size(); @@ -114,7 +114,7 @@ ElementsAndOffsets getKeysHierarchy( elements.emplace_back(hierarchy_key); ++current_hierarchy_depth; - std::optional parent_key = std::forward(get_parent_func)(hierarchy_key); + std::optional parent_key = std::forward(get_key_func)(hierarchy_key); if (!parent_key.has_value()) break; @@ -162,7 +162,11 @@ PaddedPODArray isInKeysHierarchy( PaddedPODArray result; result.resize_fill(hierarchy_keys.size()); - ElementsAndOffsets hierarchy = getKeysHierarchy(hierarchy_keys, hierarchy_null_value, std::forward(is_key_valid_func), std::forward(get_parent_func)); + ElementsAndOffsets hierarchy = getKeysHierarchy( + hierarchy_keys, + hierarchy_null_value, + std::forward(is_key_valid_func), + std::forward(get_parent_func)); auto & offsets = hierarchy.offsets; auto & elements = hierarchy.elements; @@ -186,6 +190,53 @@ PaddedPODArray isInKeysHierarchy( return result; } +template +ColumnPtr getDescendandsArray( + const PaddedPODArray & hierarchy_keys, + const KeyType & hierarchy_null_value, + size_t level, + IsKeyValidFunc && is_key_valid_func, + GetDescendantKeyFunc && get_descendant_func) +{ + auto elements_and_offsets = getKeysHierarchy( + hierarchy_keys, + hierarchy_null_value, + std::forward(is_key_valid_func), + std::forward(get_descendant_func)); + + auto & elements = elements_and_offsets.elements; + auto & offsets = elements_and_offsets.offsets; + + PaddedPODArray descendants; + descendants.reserve(elements.size()); + + PaddedPODArray descendants_offsets; + descendants_offsets.reserve(elements.size()); + + for (size_t i = 0; i < offsets.size(); ++i) + { + size_t offset_start_index = i > 0 ? offsets[i - 1] : 0; + size_t offset_end_index = offsets[i]; + size_t size = offset_end_index - offset_start_index; + + if (level == 0) + descendants.insert(elements.begin() + offset_start_index + 1, elements.begin() + offset_end_index); + else if (level < size) + descendants.emplace_back(elements[offset_start_index + level]); + + descendants_offsets.emplace_back(descendants.size()); + } + + auto elements_column = ColumnVector::create(); + elements_column->getData() = std::move(elements_and_offsets.elements); + + auto offsets_column = ColumnVector::create(); + offsets_column->getData() = std::move(offsets); + + auto column_array = ColumnArray::create(std::move(elements_column), std::move(offsets_column)); + return column_array; +} + ColumnPtr getHierarchyDefaultImplementation(const IDictionary * dictionary, ColumnPtr key_column, const DataTypePtr & key_type); ColumnUInt8::Ptr isInHierarchyDefaultImplementation( diff --git a/src/Dictionaries/IDictionary.h b/src/Dictionaries/IDictionary.h index 8b10cd7a819..8bc4e45ffa1 100644 --- a/src/Dictionaries/IDictionary.h +++ b/src/Dictionaries/IDictionary.h @@ -165,7 +165,7 @@ struct IDictionary : public IExternalLoadable const DataTypePtr & key_type [[maybe_unused]]) const { throw Exception(ErrorCodes::NOT_IMPLEMENTED, - "Hierarchy is not supported for {} dictionary.", + "Method getHierarchy is not supported for {} dictionary.", getDictionaryID().getNameForLogs()); } @@ -175,7 +175,17 @@ struct IDictionary : public IExternalLoadable const DataTypePtr & key_type [[maybe_unused]]) const { throw Exception(ErrorCodes::NOT_IMPLEMENTED, - "Hierarchy is not supported for {} dictionary.", + "Method isInHierarchy is not supported for {} dictionary.", + getDictionaryID().getNameForLogs()); + } + + virtual ColumnPtr getDescendands( + ColumnPtr key_column [[maybe_unused]], + const DataTypePtr & key_type [[maybe_unused]], + size_t level [[maybe_unused]]) const + { + throw Exception(ErrorCodes::NOT_IMPLEMENTED, + "Method getDescendands is not supported for {} dictionary.", getDictionaryID().getNameForLogs()); } diff --git a/src/Functions/FunctionsExternalDictionaries.h b/src/Functions/FunctionsExternalDictionaries.h index 9d190644b3c..d1f8ff5a1f8 100644 --- a/src/Functions/FunctionsExternalDictionaries.h +++ b/src/Functions/FunctionsExternalDictionaries.h @@ -806,4 +806,123 @@ private: mutable FunctionDictHelper helper; }; +class FunctionDictGetChildren final : public IFunction +{ +public: + static constexpr auto name = "dictGetChildren"; + + static FunctionPtr create(const Context & context) + { + return std::make_shared(context); + } + + explicit FunctionDictGetChildren(const Context & context_) + : helper(context_) {} + + String getName() const override { return name; } + +private: + size_t getNumberOfArguments() const override { return 2; } + + bool useDefaultImplementationForConstants() const final { return true; } + ColumnNumbers getArgumentsThatAreAlwaysConstant() const final { return {0}; } + + DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override + { + if (!isString(arguments[0])) + throw Exception{"Illegal type " + arguments[0]->getName() + " of first argument of function " + getName() + + ", expected a string.", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT}; + + if (!WhichDataType(arguments[1]).isUInt64()) + throw Exception{"Illegal type " + arguments[1]->getName() + " of second argument of function " + getName() + + ", must be UInt64.", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT}; + + return std::make_shared(std::make_shared()); + } + + bool isDeterministic() const override { return false; } + + ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t input_rows_count) const override + { + if (input_rows_count == 0) + return result_type->createColumn(); + + auto dict = helper.getDictionary(arguments[0]); + + if (!dict->hasHierarchy()) + throw Exception(ErrorCodes::UNSUPPORTED_METHOD, "Dictionary ({}) does not support hierarchy", dict->getFullName()); + + ColumnPtr res = dict->getDescendands(arguments[1].column, std::make_shared(), 0); + + return res; + } + + mutable FunctionDictHelper helper; +}; + +class FunctionDictGetDescendands final : public IFunction +{ +public: + static constexpr auto name = "dictGetDescendands"; + + static FunctionPtr create(const Context & context) + { + return std::make_shared(context); + } + + explicit FunctionDictGetDescendands(const Context & context_) + : helper(context_) {} + + String getName() const override { return name; } + +private: + size_t getNumberOfArguments() const override { return 3; } + + bool useDefaultImplementationForConstants() const final { return true; } + ColumnNumbers getArgumentsThatAreAlwaysConstant() const final { return {0, 2}; } + + DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override + { + if (!isString(arguments[0])) + throw Exception{"Illegal type " + arguments[0]->getName() + " of first argument of function " + getName() + + ", expected a string.", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT}; + + if (!WhichDataType(arguments[1]).isUInt64()) + throw Exception{"Illegal type " + arguments[1]->getName() + " of second argument of function " + getName() + + ", must be UInt64.", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT}; + + if (!WhichDataType(arguments[2]).isUInt64()) + throw Exception{"Illegal type " + arguments[1]->getName() + " of third argument of function " + getName() + + ", must be const UInt64.", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT}; + + return std::make_shared(); + } + + bool isDeterministic() const override { return false; } + + ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t input_rows_count) const override + { + if (input_rows_count == 0) + return result_type->createColumn(); + + const auto * level_const_column = checkAndGetColumnConst>(arguments[2].column.get()); + + if (!level_const_column) + throw Exception{"Illegal type " + arguments[1].type->getName() + " of third argument of function " + getName() + + ", must be const UInt64.", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT}; + + auto dict = helper.getDictionary(arguments[0]); + + if (!dict->hasHierarchy()) + throw Exception(ErrorCodes::UNSUPPORTED_METHOD, "Dictionary ({}) does not support hierarchy", dict->getFullName()); + + size_t level = static_cast(level_const_column->getValue()); + ColumnPtr res = dict->getDescendands(arguments[1].column, std::make_shared(), level); + + return res; + } + + mutable FunctionDictHelper helper; +}; + }