|
9 | 9 | #include "duckdb/optimizer/optimizer.hpp" |
10 | 10 | #include "duckdb/planner/expression/bound_aggregate_expression.hpp" |
11 | 11 | #include "duckdb/planner/expression/bound_case_expression.hpp" |
| 12 | +#include "duckdb/planner/expression/bound_cast_expression.hpp" |
12 | 13 | #include "duckdb/planner/expression/bound_columnref_expression.hpp" |
13 | 14 | #include "duckdb/planner/expression/bound_comparison_expression.hpp" |
14 | 15 | #include "duckdb/planner/expression/bound_constant_expression.hpp" |
@@ -134,6 +135,137 @@ LogicalOperator *FindProjectionAggregateInput(unique_ptr<LogicalOperator> &plan, |
134 | 135 | } |
135 | 136 | return nullptr; |
136 | 137 | } |
| 138 | + |
| 139 | +void InjectSumNonNullCounts(ClientContext &context, unique_ptr<LogicalOperator> &plan) { |
| 140 | + if (!plan) { |
| 141 | + return; |
| 142 | + } |
| 143 | + |
| 144 | + // Find the linear output path down to the aggregate. ORDER BY/LIMIT/TOP_N |
| 145 | + // and MATERIALIZED_CTE are pass-through wrappers; projections may rename or |
| 146 | + // cast aggregate outputs and therefore need to be followed explicitly. |
| 147 | + vector<LogicalProjection *> projections; |
| 148 | + LogicalOperator *node = plan.get(); |
| 149 | + LogicalOperator *agg_search = nullptr; |
| 150 | + while (node) { |
| 151 | + switch (node->type) { |
| 152 | + case LogicalOperatorType::LOGICAL_ORDER_BY: |
| 153 | + case LogicalOperatorType::LOGICAL_LIMIT: |
| 154 | + case LogicalOperatorType::LOGICAL_TOP_N: |
| 155 | + case LogicalOperatorType::LOGICAL_DISTINCT: |
| 156 | + node = node->children.empty() ? nullptr : node->children[0].get(); |
| 157 | + continue; |
| 158 | + case LogicalOperatorType::LOGICAL_MATERIALIZED_CTE: |
| 159 | + node = node->children.size() < 2 ? nullptr : node->children[1].get(); |
| 160 | + continue; |
| 161 | + case LogicalOperatorType::LOGICAL_PROJECTION: |
| 162 | + projections.push_back(&node->Cast<LogicalProjection>()); |
| 163 | + node = node->children.empty() ? nullptr : node->children[0].get(); |
| 164 | + continue; |
| 165 | + case LogicalOperatorType::LOGICAL_FILTER: |
| 166 | + if (!node->children.empty() && |
| 167 | + node->children[0]->type == LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY) { |
| 168 | + agg_search = node->children[0].get(); |
| 169 | + } |
| 170 | + node = nullptr; |
| 171 | + continue; |
| 172 | + case LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY: |
| 173 | + agg_search = node; |
| 174 | + node = nullptr; |
| 175 | + continue; |
| 176 | + default: |
| 177 | + node = nullptr; |
| 178 | + continue; |
| 179 | + } |
| 180 | + } |
| 181 | + if (!agg_search || projections.empty()) { |
| 182 | + return; |
| 183 | + } |
| 184 | + |
| 185 | + auto &agg = agg_search->Cast<LogicalAggregate>(); |
| 186 | + auto &aggregate_projection = *projections.back(); |
| 187 | + auto &output_projection = *projections.front(); |
| 188 | + const idx_t original_aggregate_count = agg.expressions.size(); |
| 189 | + |
| 190 | + struct SumCountOutput { |
| 191 | + idx_t projection_index; |
| 192 | + idx_t aggregate_index; |
| 193 | + }; |
| 194 | + vector<SumCountOutput> outputs; |
| 195 | + for (idx_t projection_index = 0; projection_index < output_projection.expressions.size(); projection_index++) { |
| 196 | + Expression *resolved = output_projection.expressions[projection_index].get(); |
| 197 | + while (resolved) { |
| 198 | + while (resolved->expression_class == ExpressionClass::BOUND_CAST) { |
| 199 | + resolved = resolved->Cast<BoundCastExpression>().child.get(); |
| 200 | + } |
| 201 | + if (resolved->type != ExpressionType::BOUND_COLUMN_REF) { |
| 202 | + break; |
| 203 | + } |
| 204 | + auto &ref = resolved->Cast<BoundColumnRefExpression>(); |
| 205 | + if (ref.binding.table_index == agg.aggregate_index) { |
| 206 | + break; |
| 207 | + } |
| 208 | + LogicalProjection *referenced_projection = nullptr; |
| 209 | + for (idx_t projection_depth = 1; projection_depth < projections.size(); projection_depth++) { |
| 210 | + if (projections[projection_depth]->table_index == ref.binding.table_index) { |
| 211 | + referenced_projection = projections[projection_depth]; |
| 212 | + break; |
| 213 | + } |
| 214 | + } |
| 215 | + if (!referenced_projection || ref.binding.column_index >= referenced_projection->expressions.size()) { |
| 216 | + resolved = nullptr; |
| 217 | + break; |
| 218 | + } |
| 219 | + resolved = referenced_projection->expressions[ref.binding.column_index].get(); |
| 220 | + } |
| 221 | + if (!resolved || resolved->type != ExpressionType::BOUND_COLUMN_REF) { |
| 222 | + continue; |
| 223 | + } |
| 224 | + auto &ref = resolved->Cast<BoundColumnRefExpression>(); |
| 225 | + if (ref.binding.table_index != agg.aggregate_index || ref.binding.column_index >= original_aggregate_count) { |
| 226 | + continue; |
| 227 | + } |
| 228 | + auto &aggregate_expr = agg.expressions[ref.binding.column_index]; |
| 229 | + if (aggregate_expr->expression_class != ExpressionClass::BOUND_AGGREGATE) { |
| 230 | + continue; |
| 231 | + } |
| 232 | + auto &sum = aggregate_expr->Cast<BoundAggregateExpression>(); |
| 233 | + if (sum.function.name != "sum" || sum.IsDistinct() || sum.children.size() != 1) { |
| 234 | + continue; |
| 235 | + } |
| 236 | + if (sum.alias.find(string(openivm::SUM_COL_PREFIX)) == 0 || |
| 237 | + sum.alias.find(string(openivm::VAR_SQ_COL_PREFIX)) == 0 || |
| 238 | + sum.alias.find(string(openivm::VAR_SQP_COL_PREFIX)) == 0) { |
| 239 | + continue; |
| 240 | + } |
| 241 | + |
| 242 | + auto count_func = BindAggregateByName(context, "count", {sum.children[0]->return_type}); |
| 243 | + vector<unique_ptr<Expression>> count_args; |
| 244 | + count_args.push_back(sum.children[0]->Copy()); |
| 245 | + auto count_expr = make_uniq<BoundAggregateExpression>(std::move(count_func), std::move(count_args), nullptr, |
| 246 | + nullptr, AggregateType::NON_DISTINCT); |
| 247 | + count_expr->alias = string(openivm::SUM_COUNT_COL_PREFIX) + to_string(projection_index); |
| 248 | + outputs.push_back({projection_index, agg.expressions.size()}); |
| 249 | + agg.expressions.push_back(std::move(count_expr)); |
| 250 | + } |
| 251 | + if (outputs.empty()) { |
| 252 | + return; |
| 253 | + } |
| 254 | + |
| 255 | + agg.ResolveOperatorTypes(); |
| 256 | + auto aggregate_bindings = agg_search->GetColumnBindings(); |
| 257 | + auto aggregate_types = agg_search->types; |
| 258 | + for (auto &output : outputs) { |
| 259 | + idx_t output_index = agg.groups.size() + output.aggregate_index; |
| 260 | + auto count_ref = |
| 261 | + make_uniq<BoundColumnRefExpression>(aggregate_types[output_index], aggregate_bindings[output_index]); |
| 262 | + count_ref->alias = string(openivm::SUM_COUNT_COL_PREFIX) + to_string(output.projection_index); |
| 263 | + aggregate_projection.expressions.push_back(std::move(count_ref)); |
| 264 | + } |
| 265 | + aggregate_projection.ResolveOperatorTypes(); |
| 266 | + OPENIVM_DEBUG_PRINT("[PlanRewrite] Injected %zu SUM non-NULL counts\n", outputs.size()); |
| 267 | +} |
| 268 | + |
137 | 269 | void RewriteDerivedAggregates(ClientContext &context, unique_ptr<LogicalOperator> &plan, Optimizer &opt, bool is_top) { |
138 | 270 | for (auto &child : plan->children) { |
139 | 271 | RewriteDerivedAggregates(context, child, opt, false); |
|
0 commit comments