✨ rewrite CASE WHEN x IS NULL THEN y ELSE x as IFNULL(x, y) #2293

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
 

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