|
| 1 | +#include <Analyzer/Passes/OptimizeGroupByInjectiveFunctionsPass.h> |
| 2 | +#include <Analyzer/ConstantNode.h> |
| 3 | +#include <Analyzer/FunctionNode.h> |
| 4 | +#include <Analyzer/InDepthQueryTreeVisitor.h> |
| 5 | +#include <Analyzer/IQueryTreeNode.h> |
| 6 | +#include <DataTypes/IDataType.h> |
| 7 | +#include <Interpreters/ExternalDictionariesLoader.h> |
| 8 | + |
| 9 | +namespace DB |
| 10 | +{ |
| 11 | + |
| 12 | +namespace |
| 13 | +{ |
| 14 | + |
| 15 | +const std::unordered_set<String> possibly_injective_function_names |
| 16 | +{ |
| 17 | + "dictGet", |
| 18 | + "dictGetString", |
| 19 | + "dictGetUInt8", |
| 20 | + "dictGetUInt16", |
| 21 | + "dictGetUInt32", |
| 22 | + "dictGetUInt64", |
| 23 | + "dictGetInt8", |
| 24 | + "dictGetInt16", |
| 25 | + "dictGetInt32", |
| 26 | + "dictGetInt64", |
| 27 | + "dictGetFloat32", |
| 28 | + "dictGetFloat64", |
| 29 | + "dictGetDate", |
| 30 | + "dictGetDateTime" |
| 31 | +}; |
| 32 | + |
| 33 | +class OptimizeGroupByInjectiveFunctionsVisitor : public InDepthQueryTreeVisitorWithContext<OptimizeGroupByInjectiveFunctionsVisitor> |
| 34 | +{ |
| 35 | + using Base = InDepthQueryTreeVisitorWithContext<OptimizeGroupByInjectiveFunctionsVisitor>; |
| 36 | +public: |
| 37 | + explicit OptimizeGroupByInjectiveFunctionsVisitor(ContextPtr context) |
| 38 | + : Base(std::move(context)) |
| 39 | + {} |
| 40 | + |
| 41 | + void enterImpl(QueryTreeNodePtr & node) |
| 42 | + { |
| 43 | + if (!getSettings().optimize_injective_functions_in_group_by) |
| 44 | + return; |
| 45 | + |
| 46 | + auto * query = node->as<QueryNode>(); |
| 47 | + if (!query) |
| 48 | + return; |
| 49 | + |
| 50 | + if (!query->hasGroupBy()) |
| 51 | + return; |
| 52 | + |
| 53 | + if (query->isGroupByWithCube() || query->isGroupByWithRollup()) |
| 54 | + return; |
| 55 | + |
| 56 | + auto & group_by = query->getGroupBy().getNodes(); |
| 57 | + if (query->isGroupByWithGroupingSets()) |
| 58 | + { |
| 59 | + for (auto & set : group_by) |
| 60 | + { |
| 61 | + auto & grouping_set = set->as<ListNode>()->getNodes(); |
| 62 | + optimizeGroupingSet(grouping_set); |
| 63 | + } |
| 64 | + } |
| 65 | + else |
| 66 | + optimizeGroupingSet(group_by); |
| 67 | + } |
| 68 | + |
| 69 | +private: |
| 70 | + void optimizeGroupingSet(QueryTreeNodes & grouping_set) |
| 71 | + { |
| 72 | + auto context = getContext(); |
| 73 | + |
| 74 | + QueryTreeNodes new_group_by_keys; |
| 75 | + new_group_by_keys.reserve(grouping_set.size()); |
| 76 | + for (auto & group_by_elem : grouping_set) |
| 77 | + { |
| 78 | + std::queue<QueryTreeNodePtr> nodes_to_process; |
| 79 | + nodes_to_process.push(group_by_elem); |
| 80 | + |
| 81 | + while (!nodes_to_process.empty()) |
| 82 | + { |
| 83 | + auto node_to_process = nodes_to_process.front(); |
| 84 | + nodes_to_process.pop(); |
| 85 | + |
| 86 | + auto const * function_node = node_to_process->as<FunctionNode>(); |
| 87 | + if (!function_node) |
| 88 | + { |
| 89 | + // Constant aggregation keys are removed in PlannerExpressionAnalysis.cpp |
| 90 | + new_group_by_keys.push_back(node_to_process); |
| 91 | + continue; |
| 92 | + } |
| 93 | + |
| 94 | + // Aggregate functions are not allowed in GROUP BY clause |
| 95 | + auto function = function_node->getFunctionOrThrow(); |
| 96 | + bool can_be_eliminated = function->isInjective(function_node->getArgumentColumns()); |
| 97 | + |
| 98 | + if (can_be_eliminated) |
| 99 | + { |
| 100 | + for (auto const & argument : function_node->getArguments()) |
| 101 | + { |
| 102 | + // We can skip constants here because aggregation key is already not a constant. |
| 103 | + if (argument->getNodeType() != QueryTreeNodeType::CONSTANT) |
| 104 | + nodes_to_process.push(argument); |
| 105 | + } |
| 106 | + } |
| 107 | + else |
| 108 | + new_group_by_keys.push_back(node_to_process); |
| 109 | + } |
| 110 | + } |
| 111 | + |
| 112 | + grouping_set = std::move(new_group_by_keys); |
| 113 | + } |
| 114 | +}; |
| 115 | + |
| 116 | +} |
| 117 | + |
| 118 | +void OptimizeGroupByInjectiveFunctionsPass::run(QueryTreeNodePtr query_tree_node, ContextPtr context) |
| 119 | +{ |
| 120 | + OptimizeGroupByInjectiveFunctionsVisitor visitor(std::move(context)); |
| 121 | + visitor.visit(query_tree_node); |
| 122 | +} |
| 123 | + |
| 124 | +} |
0 commit comments