From 2ab39968b43e1d943568e1593148197844710894 Mon Sep 17 00:00:00 2001
From: darrenhuai <60621295+darrenhuai@users.noreply.github.com>
Date: Thu, 6 Aug 2026 16:36:30 -0700
Subject: [PATCH 1/2] Cover the folding branches in Optimizer, drop dead _debug
 method

optimizer.py was sitting at 86% because fold_case had no tests at
all (the whole method, not just the branches called out in #121),
and fold_zeros only had subtract covered - multiply and add zero
folding were completely untested. Added tests for each folding
method's disabled-optimizer passthrough, the actual folding paths,
and the various "nothing to fold" fallbacks, including the
non-literal IF/ELSE case that's deliberately left unfolded because
of the postgresql FIXME.

Also removed _debug: nothing calls it except itself, and now that
it's gone the unreachable branch after `return expr` in fold_case's
boolean-cast case doesn't need a decision either - it just wasn't
part of the count anymore once I checked, since the compiler drops
statements that follow an unconditional return in the same block.

optimizer.py coverage: 86% -> 100%.
---
 src/orbital/translation/optimizer.py |  11 --
 tests/test_optimizer.py              | 175 +++++++++++++++++++++++++++
 2 files changed, 175 insertions(+), 11 deletions(-)

diff --git a/src/orbital/translation/optimizer.py b/src/orbital/translation/optimizer.py
index ed85521..f72aa5c 100644
--- a/src/orbital/translation/optimizer.py
+++ b/src/orbital/translation/optimizer.py
@@ -327,14 +327,3 @@ def fold_operation(self, expr: ibis.Expr) -> ibis.Expr:
         else:
             # No possible folding
             return expr
-
-    def _debug(self, expr: ibis.Expr, show_args: bool = True) -> str:
-        """Given an expression, return a string representation for debugging."""
-        if isinstance(expr, Literal):
-            return repr(expr.value)
-        elif show_args is False:
-            return type(expr).__name__
-        elif not hasattr(expr, "args"):
-            return f"{type(expr).__name__}(<unknown>)"
-        else:
-            return f"{type(expr).__name__}({', '.join([self._debug(a, show_args=False) for a in expr.args])})"
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 1dd786a..17177cf 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -147,3 +147,178 @@ def test_fold_operation_unary_not(self):
         # not return the unreduced operation tree.
         assert isinstance(result.op(), Literal)
         assert result.op().value is False
+
+    def test_ensure_expr_passthrough_for_expr(self):
+        column = ibis.memtable({"a": [1.0]})["a"]
+        result = self.optimizer._ensure_expr(column)
+        assert result is column
+
+    def test_ensure_expr_wraps_literal_op_node(self):
+        op_node = ibis.literal(5).op()
+        result = self.optimizer._ensure_expr(op_node)
+        assert isinstance(result, ibis.Expr)
+        assert result.execute() == 5
+
+    def test_ensure_expr_raises_for_unsupported_type(self):
+        with pytest.raises(TypeError):
+            self.optimizer._ensure_expr(object())
+
+    def test_fold_contiguous_sum_disabled_returns_unchanged(self):
+        disabled = Optimizer(enabled=False)
+        a, b = ibis.literal(1), ibis.literal(2)
+        result = disabled.fold_contiguous_sum([a, b])
+        assert len(result) == 2
+        assert result[0] is a
+        assert result[1] is b
+
+    def test_fold_case_deferred_raises(self):
+        with pytest.raises(NotImplementedError):
+            self.optimizer.fold_case(ibis._.x)
+
+    def test_fold_case_disabled_returns_unchanged(self):
+        disabled = Optimizer(enabled=False)
+        expr = ibis.cases((ibis.literal(True), ibis.literal(1)), else_=ibis.literal(0))
+        result = disabled.fold_case(expr)
+        assert result is expr
+
+    def test_fold_case_all_results_same_value_folds_to_value(self):
+        # Every result and the default are literal 7, so the case can never
+        # produce anything else - it should collapse to a plain literal.
+        expr = ibis.cases(
+            (ibis.literal(True), ibis.literal(7)),
+            (ibis.literal(False), ibis.literal(7)),
+            else_=ibis.literal(7),
+        )
+        result = self.optimizer.fold_case(expr)
+        assert result.op().value == 7
+
+    def test_fold_case_literal_true_condition_returns_result(self):
+        expr = ibis.cases((ibis.literal(True), ibis.literal(1)), else_=ibis.literal(0))
+        result = self.optimizer.fold_case(expr)
+        assert result.op().value == 1
+
+    def test_fold_case_literal_false_condition_returns_default(self):
+        expr = ibis.cases((ibis.literal(False), ibis.literal(1)), else_=ibis.literal(0))
+        result = self.optimizer.fold_case(expr)
+        assert result.op().value == 0
+
+    def test_fold_case_boolean_ifelse_on_non_literal_condition_not_folded(self):
+        # Folding a single IF/ELSE returning 1 or 0 into a boolean cast is
+        # deliberately disabled (see the FIXME in fold_case: it doesn't work
+        # on postgresql), so a non-literal condition must come back
+        # unchanged rather than being rewritten.
+        table = ibis.memtable({"x": [1.0]})
+        condition = table["x"] > 0
+        expr = ibis.cases((condition, ibis.literal(1)), else_=ibis.literal(0))
+        result = self.optimizer.fold_case(expr)
+        assert result is expr
+
+    def test_fold_case_multiple_cases_returns_unchanged(self):
+        expr = ibis.cases(
+            (ibis.literal(True), ibis.literal("a")),
+            (ibis.literal(False), ibis.literal("b")),
+            else_=ibis.literal("c"),
+        )
+        result = self.optimizer.fold_case(expr)
+        assert result is expr
+
+    def test_fold_cast_disabled_returns_unchanged(self):
+        disabled = Optimizer(enabled=False)
+        expr = ibis.literal(5).cast("float64")
+        result = disabled.fold_cast(expr)
+        assert result is expr
+
+    def test_fold_cast_non_cast_expr_returns_unchanged(self):
+        # Ibis itself drops a cast to a column's own type, so fold_cast
+        # never even sees a Cast node here - it should just hand the
+        # expression back.
+        table = ibis.memtable({"a": [1.0]})
+        column = table["a"]
+        not_a_cast = column.cast("float64")
+        assert not isinstance(not_a_cast.op(), ibis.expr.operations.Cast)
+        result = self.optimizer.fold_cast(not_a_cast)
+        assert result is not_a_cast
+
+    def test_fold_cast_unsupported_literal_type_raises(self):
+        expr = ibis.literal(5).cast("date")
+        with pytest.raises(NotImplementedError):
+            self.optimizer.fold_cast(expr)
+
+    def test_fold_cast_already_target_type_unwraps_manual_cast(self):
+        # Build a Cast node directly (bypassing ibis's own .cast(), which
+        # already drops a same-type cast on its own) so fold_cast has to be
+        # the one to notice the arg is already the target type.
+        table = ibis.memtable({"a": [1.0]})
+        column = table["a"]
+        manual_cast = ibis.expr.operations.Cast(column.op(), to=column.type()).to_expr()
+        result = self.optimizer.fold_cast(manual_cast)
+        assert result.equals(column)
+
+    def test_fold_zeros_disabled_returns_unchanged(self):
+        disabled = Optimizer(enabled=False)
+        expr = ibis.literal(0) * ibis.literal(7)
+        result = disabled.fold_zeros(expr)
+        assert result is expr
+
+    def test_fold_zeros_multiply_left_zero(self):
+        expr = ibis.literal(0) * ibis.literal(7)
+        result = self.optimizer.fold_zeros(expr)
+        assert result.op().value == 0
+
+    def test_fold_zeros_multiply_right_zero(self):
+        expr = ibis.literal(7) * ibis.literal(0)
+        result = self.optimizer.fold_zeros(expr)
+        assert result.op().value == 0
+
+    def test_fold_zeros_add_left_zero(self):
+        table = ibis.memtable({"value": [1.0, 2.0]})
+        column = table["value"]
+        expr = ibis.literal(0) + column
+        result = self.optimizer.fold_zeros(expr)
+        assert result.equals(column)
+
+    def test_fold_zeros_add_right_zero(self):
+        table = ibis.memtable({"value": [1.0, 2.0]})
+        column = table["value"]
+        expr = column + ibis.literal(0)
+        result = self.optimizer.fold_zeros(expr)
+        assert result.equals(column)
+
+    def test_fold_zeros_no_zero_operand_returns_unchanged(self):
+        expr = ibis.literal(3) - ibis.literal(4)
+        result = self.optimizer.fold_zeros(expr)
+        assert result is expr
+
+    def test_fold_operation_disabled_returns_unchanged(self):
+        disabled = Optimizer(enabled=False)
+        expr = ibis.literal(2) + ibis.literal(3)
+        result = disabled.fold_operation(expr)
+        assert result is expr
+
+    def test_fold_operation_python_scalar_wrapped_as_literal(self):
+        # fold_operation can be handed a plain Python value when a previous
+        # fold already reduced a subtree to a scalar; it still has to come
+        # back out as an ibis expression.
+        result = self.optimizer.fold_operation(5)
+        assert isinstance(result, ibis.Expr)
+        assert result.op().value == 5
+
+    def test_fold_operation_mixed_literal_and_column_delegates_to_fold_zeros(self):
+        table = ibis.memtable({"value": [1.0, 2.0]})
+        column = table["value"]
+        expr = column * ibis.literal(0)
+        result = self.optimizer.fold_operation(expr)
+        assert result.op().value == 0
+
+    def test_fold_operation_binary_literal_folding(self):
+        expr = ibis.literal(2) + ibis.literal(3)
+        result = self.optimizer.fold_operation(expr)
+        assert isinstance(result.op(), Literal)
+        assert result.op().value == 5
+
+    def test_fold_operation_unfoldable_op_returns_unchanged(self):
+        # IsNull isn't in BINARY_OPS or UNARY_OPS, so even with an
+        # all-literal input there's nothing fold_operation can precompute.
+        expr = ibis.literal(5).isnull()
+        result = self.optimizer.fold_operation(expr)
+        assert result is expr

From cd064be94678f9620c6c3bd7b0531b523bb250df Mon Sep 17 00:00:00 2001
From: darrenhuai <60621295+darrenhuai@users.noreply.github.com>
Date: Mon, 10 Aug 2026 15:07:39 -0700
Subject: [PATCH 2/2] Keep _debug, exclude it from coverage instead of deleting
 it

amol- uses it on demand while working on the optimizer, so mark it
with a no cover pragma (matching types.py) and say in the docstring
why it isn't tested.
---
 src/orbital/translation/optimizer.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/src/orbital/translation/optimizer.py b/src/orbital/translation/optimizer.py
index f72aa5c..446e384 100644
--- a/src/orbital/translation/optimizer.py
+++ b/src/orbital/translation/optimizer.py
@@ -327,3 +327,20 @@ def fold_operation(self, expr: ibis.Expr) -> ibis.Expr:
         else:
             # No possible folding
             return expr
+
+    def _debug(  # pragma: no cover
+        self, expr: ibis.Expr, show_args: bool = True
+    ) -> str:
+        """Given an expression, return a string representation for debugging.
+
+        Only used on demand while developing the optimizer, so it is
+        intentionally left uncovered by the test suite.
+        """
+        if isinstance(expr, Literal):
+            return repr(expr.value)
+        elif show_args is False:
+            return type(expr).__name__
+        elif not hasattr(expr, "args"):
+            return f"{type(expr).__name__}(<unknown>)"
+        else:
+            return f"{type(expr).__name__}({', '.join([self._debug(a, show_args=False) for a in expr.args])})"