diff --git a/opteryx/planner/optimizer/strategies/predicate_rewriter.py b/opteryx/planner/optimizer/strategies/predicate_rewriter.py index 2ec1ac85..f6317e96 100644 --- a/opteryx/planner/optimizer/strategies/predicate_rewriter.py +++ b/opteryx/planner/optimizer/strategies/predicate_rewriter.py @@ -182,6 +182,24 @@ def _rewrite_predicate(predicate, statistics: QueryStatistics): return predicate +def _rewrite_function(function, statistics: QueryStatistics): + """ + Rewrite CASE WHEN x IS NULL THEN y ELSE x END to IFNULL(x, y) + """ + if function.node_type == NodeType.FUNCTION and function.value == "CASE": + if len(function.parameters) == 2 and function.parameters[0].parameters[0].value == "IsNull": + compare_column = function.parameters[0].parameters[0].centre + target_column = function.parameters[1].parameters[1] + value_if_null = function.parameters[1].parameters[0] + + if compare_column.schema_column.identity == target_column.schema_column.identity: + statistics.optimization_predicate_rewriter_case_to_ifnull += 1 + function.value = "IFNULL" + function.parameters = [compare_column, value_if_null] + return function + return function + + class PredicateRewriteStrategy(OptimizationStrategy): def visit(self, node: LogicalPlanNode, context: OptimizerContext) -> OptimizerContext: if not context.optimized_plan: @@ -195,8 +213,10 @@ def visit(self, node: LogicalPlanNode, context: OptimizerContext) -> OptimizerCo new_columns = [] for column in node.columns: new_column = _rewrite_predicate(column, self.statistics) + new_column = _rewrite_function(new_column, self.statistics) new_columns.append(new_column) node.columns = new_columns + context.optimized_plan[context.node_id] = node return context diff --git a/tests/plan_optimization/test_optimizations_invoked.py b/tests/plan_optimization/test_optimizations_invoked.py index 767d4d13..8d7883ec 100644 --- a/tests/plan_optimization/test_optimizations_invoked.py +++ b/tests/plan_optimization/test_optimizations_invoked.py @@ -23,9 +23,10 @@ ("SELECT name FROM $astronauts WHERE name = 'Neil A. Armstrong'", "optimization_predicate_pushdown"), ("SELECT name FROM $planets WHERE name LIKE '%'", "optimization_constant_fold_reduce"), # rewritten to `name is not null` ("SELECT name FROM $planets WHERE name ILIKE '%'", "optimization_constant_fold_reduce"), # rewritten to `name is not null` - ("SELECT name FROM $planets WHERE name ILIKE '%th%'", "optimization_predicate_rewriter_replace_like_with_in_string"), + ("SELECT name FROM $planets WHERE name ILIKE '%th%'", "optimization_predicate_rewriter_replace_like_with_in_string"), ("SELECT name FROM $planets WHERE NOT name NOT ILIKE '%th%'", "optimization_boolean_rewrite_inversion"), ("SELECT * FROM $planets WHERE NOT name != 'Earth'", "optimization_boolean_rewrite_inversion"), + ("SELECT CASE WHEN surface_pressure IS NULL THEN -100.00 ELSE surface_pressure END FROM $planets", "optimization_predicate_rewriter_case_to_ifnull"), ] # fmt:on