[Optimizer] Add literal inlining pass for scalar ValueOps (#251)

stratum/optimizer/_optimize.py

@@ -5,8 +5,10 @@
 from collections import deque, defaultdict
 from dataclasses import dataclass
 from typing import Any
+
 from .logical._op_cse import apply_op_cse
 from .logical._dataframe_ops import extract_dataframe_op, add_splitting_op
+from .logical._literal_inlining import literal_inlining
 from .logical._numeric_ops import extract_numeric_op
 from .logical._candidate_ops import CollectCandidatesOp, DeclaredScoringOp, ScoreCandidatesOp
 from .logical._ops import BaseEstimatorOp, ChoiceOp, Op, OperandRef, as_op, check_choices_not_shared, remap_operand_refs
@@ -85,6 +87,7 @@ def __init__(
         projection_pushup: bool = True,
         selector: ImplementationSelector | None = None,
         pandas_query: bool = False,
+        literal_inlining: bool = True,
     ):
         """``selector`` overrides the implementation-selection policy for this
         plan. Left ``None``, the policy named by ``FLAGS.implementation_selector``
@@ -112,6 +115,7 @@ def __init__(
         self.propagate_schema = propagate_schema
         self.selector = selector
         self.pandas_query = pandas_query
+        self.literal_inlining = literal_inlining
 
 def _debug_show_graph(root: Op, name: str):
     if diagnostics_enabled():
@@ -197,6 +201,8 @@ def logical_optimize(dag_root: DataOp, config: OptConfig, env: dict = None,
 
     # Convert to logical operator DAG
     root = convert_to_ops(dag_root, env)
+    if config.literal_inlining:
+        root = literal_inlining(root)
 
     # Add splitting op
     root = add_splitting_op(root)

stratum/optimizer/logical/_literal_inlining.py

@@ -0,0 +1,101 @@
+"""Inline scalar ``ValueOp``s into consumers that can hold the literal directly."""
+from collections import defaultdict
+from typing import Any
+
+from stratum.optimizer._op_utils import topological_iterator
+from stratum.optimizer.logical._base import IRNode, rebuild_dict
+from stratum.optimizer.logical._dataframe_ops import ConcatOp
+from stratum.optimizer.logical._ops import BinOp, CallOp, GetItemOp, MethodCallOp, Op, OperandRef, ValueOp, remap_operand_refs
+
+
+def _is_literal(value: Any) -> bool:
+    return type(value) in (type(None), bool, int, float, str, bytes)
+
+
+def _inlinable(out: Op, pos: int) -> bool:
+    if isinstance(out, ConcatOp):
+        if isinstance(out.axis, OperandRef) and out.axis.k == pos:
+            param_type = type(out.inputs[pos].value) 
+            return param_type in (str, int)
+    if not isinstance(out, (CallOp, MethodCallOp, GetItemOp, BinOp)):
+        return False
+    if isinstance(out, (MethodCallOp, GetItemOp)):
+        return pos != 0
+    if isinstance(out, CallOp) and getattr(out.func, "__module__", None) == "numpy":
+        return pos != 0
+    if isinstance(out, BinOp):
+        # Keep a graph input for numeric extraction / algebraic rewrites: never
+        # inline away every input, and never hide a shared producer's identity
+        # (e.g. ``x * x``) behind a literal at one of its BinOp consumers.
+        if len(out.inputs[pos].outputs) != 1:
+            return False
+        return not all(isinstance(i, ValueOp) and _is_literal(i.value) for i in out.inputs)
+    return True
+
+
+def literal_inlining(root: Op) -> Op:
+    records: dict[Op, set[int]] = defaultdict(set)
+    for op in list(topological_iterator(root)):
+        if not isinstance(op, ValueOp) or not _is_literal(op.value) or op.is_X or op.is_y:
+            continue
+        for out in op.outputs:
+            for pos, input_op in enumerate(out.inputs):
+                if input_op is op and _inlinable(out, pos) and _check_for_constant(op, out, pos):
+                    records[out].add(pos)
+                    if isinstance(out, GetItemOp):
+                        out.update_name()
+    _rewire_edges(records)
+    return root
+
+
+def _rewire_edges(records: dict[Op, set[int]]) -> None:
+    for op, removed in records.items():
+        removed_set = set(removed)
+        new_inputs = [input_op for i, input_op in enumerate(op.inputs) if i not in removed_set]
+        mapping, offset = {}, 0
+        for i, input_op in enumerate(op.inputs):
+            if i in removed_set:
+                if not any(inp is input_op for inp in new_inputs):
+                    input_op.outputs = [out for out in input_op.outputs if out is not op]
+                offset += 1
+            else:
+                mapping[i] = i - offset
+        for field in op.fields:
+            param = getattr(op, field)
+            new_param = remap_operand_refs(param, mapping)
+            if param is not new_param:
+                setattr(op, field, new_param)
+        op.inputs = new_inputs
+
+
+def _check_for_constant(value_node: ValueOp, out: IRNode, pos: int) -> bool:
+    changed_any = False
+    for field in out.fields:
+        param = getattr(out, field)
+        changed, new_param = _inline_literal(param, pos, value_node.value)
+        if changed:
+            setattr(out, field, new_param)
+            changed_any = True
+    return changed_any
+
+
+def _inline_literal(param: Any, pos: int, literal: Any) -> tuple[bool, Any]:
+    if isinstance(param, OperandRef):
+        return (True, literal) if param.k == pos else (False, param)
+    if isinstance(param, tuple):
+        results = [_inline_literal(v, pos, literal) for v in param]
+        changed = any(flag for flag, _ in results)
+        return changed, tuple(value for _, value in results) if changed else param
+    if isinstance(param, list):
+        results = [_inline_literal(v, pos, literal) for v in param]
+        changed = any(flag for flag, _ in results)
+        return changed, [value for _, value in results] if changed else param
+    if isinstance(param, dict):
+        results = {k: _inline_literal(v, pos, literal) for k, v in param.items()}
+        changed = any(flag for flag, _ in results.values())
+        return changed, rebuild_dict(param, ((k, result[1]) for k, result in results.items())) if changed else param
+    if isinstance(param, slice):
+        results = [_inline_literal(v, pos, literal) for v in (param.start, param.stop, param.step)]
+        changed = any(flag for flag, _ in results)
+        return changed, slice(*(value for _, value in results)) if changed else param
+    return False, param

stratum/optimizer/logical/_ops.py

@@ -688,6 +688,21 @@ def __init__(self, key=None, name=None, is_filter=False):
             name = str(key)
         super().__init__(name=name)
 
+    def update_name(self):
+        # A graph-fed key has no literal label; keep the construction-time name.
+        stack = [self.key]
+        while stack:
+            value = stack.pop()
+            if isinstance(value, OperandRef):
+                return
+            if isinstance(value, (tuple, list)):
+                stack.extend(value)
+            elif isinstance(value, dict):
+                stack.extend(value.values())
+            elif isinstance(value, slice):
+                stack.extend((value.start, value.stop, value.step))
+        self.name = _getitem_key_name(self.key)
+
     def propagate_output_schema(self):
         """A literal key projects a sub-schema; a row slice or a boolean row mask
         keeps the input schema."""

tests/optimizer/logical/test_aggregation_ops.py

@@ -425,8 +425,8 @@ def setUp(self):
             "w": [4, 5, 6],
         })
 
-    def _agg_ops(self, data, **kwargs):
-        ops = optimize(data, OptConfig(dataframe_ops=True), **kwargs)
+    def _agg_ops(self, data, literal_inlining=False, **kwargs):
+        ops = optimize(data, OptConfig(dataframe_ops=True, literal_inlining=literal_inlining), **kwargs)
         return ops, [o for o in ops if isinstance(o, AggregateOp)]
 
     def test_agg_with_spec_fuses_and_executes(self):