From fda6f02c5d93e7a130c9a62ea14428e8b9f1eca9 Mon Sep 17 00:00:00 2001 From: Aram164 Date: Tue, 20 Oct 2026 03:17:48 +0200 Subject: [PATCH] [Optimizer] Add literal inlining pass for scalar ValueOps (#251) --- stratum/optimizer/_optimize.py | 6 + .../optimizer/logical/_literal_inlining.py | 101 ++++++ stratum/optimizer/logical/_ops.py | 15 + .../optimizer/logical/test_aggregation_ops.py | 4 +- .../logical/test_literal_inlining.py | 311 ++++++++++++++++++ tests/optimizer/logical/test_map_ops.py | 2 +- tests/optimizer/logical/test_numeric_ops.py | 10 +- .../optimizer/logical/test_projection_ops.py | 5 +- tests/optimizer/logical/test_selection_ops.py | 2 +- 9 files changed, 445 insertions(+), 11 deletions(-) create mode 100644 stratum/optimizer/logical/_literal_inlining.py create mode 100644 tests/optimizer/logical/test_literal_inlining.py diff --git a/stratum/optimizer/_optimize.py b/stratum/optimizer/_optimize.py index 3e3b5531..b0487efb 100644 --- a/stratum/optimizer/_optimize.py +++ b/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) diff --git a/stratum/optimizer/logical/_literal_inlining.py b/stratum/optimizer/logical/_literal_inlining.py new file mode 100644 index 00000000..35e604a8 --- /dev/null +++ b/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 diff --git a/stratum/optimizer/logical/_ops.py b/stratum/optimizer/logical/_ops.py index ca53fa3b..4cfb76de 100644 --- a/stratum/optimizer/logical/_ops.py +++ b/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.""" diff --git a/tests/optimizer/logical/test_aggregation_ops.py b/tests/optimizer/logical/test_aggregation_ops.py index 0c7ea03d..5f25f2cc 100644 --- a/tests/optimizer/logical/test_aggregation_ops.py +++ b/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): diff --git a/tests/optimizer/logical/test_literal_inlining.py b/tests/optimizer/logical/test_literal_inlining.py new file mode 100644 index 00000000..c1cca4d1 --- /dev/null +++ b/tests/optimizer/logical/test_literal_inlining.py @@ -0,0 +1,311 @@ +"""Literal inlining on the logical Op DAG.""" + +from collections import defaultdict +import operator + +import numpy as np +import pandas as pd +import pytest +from sklearn.linear_model import Ridge + +import stratum as st +from stratum._api import evaluate +from stratum.optimizer._optimize import OptConfig, extract_numeric_operators, logical_optimize, optimize +from stratum.optimizer._op_utils import topological_iterator, validate_dag +from stratum.optimizer.logical._aggregation_ops import AggregateOp +from stratum.optimizer.logical._column_expr import BinOpExpr, Col, Const +from stratum.optimizer.logical._dataframe_ops import ConcatOp, MaskSelectionOp +from stratum.optimizer.logical._literal_inlining import literal_inlining +from stratum.optimizer.logical._numeric_ops import NumericOp, NumericOpType +from stratum.optimizer.logical._projection_ops import ColumnProjectionOp +from stratum.runtime._buffer_pool import BufferPool +from stratum.optimizer.logical._ops import BinOp, CallOp, GetItemOp, MethodCallOp, OperandRef, PredictorOp, UnaryOp, ValueOp +from stratum.optimizer.physical._impl_selection import bind_op, DefaultImplementationSelector +from stratum.optimizer.physical._lowering import lower_to_physical +from stratum.optimizer.physical._plan_context import PlanContext + +R0, R1, R2 = OperandRef(0), OperandRef(1), OperandRef(2) + + +def wire(op, *inputs): + op.inputs = list(inputs) + for source in inputs: + source.add_output(op) + return op + + +def call(func, *args, **kwargs): + return CallOp(func=func, args=args, kwargs=kwargs) + + +def inline(make, values): + """Wire ``make()`` over one ValueOp per value and run the pass on it.""" + sources = [ValueOp(value) for value in values] + root = wire(make(), *sources) + assert literal_inlining(root) is root + return root, sources + + +def assert_kept(root, sources, kept): + """Exactly ``sources[kept]`` stay inputs of ``root``, with matching output edges.""" + assert root.inputs == [sources[i] for i in kept] + for i, source in enumerate(sources): + assert (root in source.outputs) == (i in kept) + validate_dag(root) + + +def execute(root): + root = lower_to_physical(root, PlanContext.from_flags()) + values = {} + for op in topological_iterator(root): + bind_op(op, PlanContext.from_flags(), selector=DefaultImplementationSelector()) + values[op] = op.process("fit_transform", [values[p] for p in op.inputs]) + return values[root] + + +@pytest.mark.parametrize("value, inlined", [ + *[(v, True) for v in (None, False, 7, 1.25, "BRASS", b"BRASS")], + # Exact built-in types only: numpy scalars and containers stay graph inputs. + *[(v, False) for v in ([1], {"a": 1}, np.array([1]), np.float64(1.0), pd.DataFrame({"a": [1]}), object())], +]) +def test_only_builtin_scalars_are_inlined(value, inlined): + root, sources = inline(lambda: call(lambda x: x, R0), [value]) + + assert_kept(root, sources, kept=[] if inlined else [0]) + assert root.args == ((value,) if inlined else (R0,)) + assert execute(root) is value + + +@pytest.mark.parametrize("pattern", ["LLL", "DDD", "LDL", "DLD"]) +def test_inputs_compact_and_refs_remap(pattern): + # L = literal (inlined), D = list (stays an input). + values = [i if kind == "L" else [i] for i, kind in enumerate(pattern)] + root, sources = inline(lambda: call(lambda *a: a, R0, R1, R2), values) + + assert_kept(root, sources, kept=[i for i, kind in enumerate(pattern) if kind == "D"]) + assert execute(root) == tuple(values) + + +def test_refs_in_nested_containers_and_kwargs(): + nested = [R1, slice(None, R0, None), defaultdict(lambda: "missing", {"a": R0})] + root, sources = inline(lambda: call(lambda *a, **k: (a, k), R0, nested, k=R0, d=R1), [3, [4]]) + + assert_kept(root, sources, kept=[1]) + assert root.args[:2] == (3, [R0, slice(None, 3, None), root.args[1][2]]) + assert type(root.args[1][2]) is defaultdict and root.args[1][2]["absent"] == "missing" + assert root.kwargs == {"k": 3, "d": R0} + args, kwargs = execute(root) + assert args[1][0] == [4] and args[1][2]["a"] == 3 and kwargs == {"k": 3, "d": [4]} + + +def test_shared_literal_and_idempotence(): + source = ValueOp("a") + method = wire(MethodCallOp("replace", args=(R0, "x"), kwargs={}), source) + other = wire(call(lambda x: x, R0), source) + root = wire(call(lambda a, b: (a, b), R0, R1), method, other) + + for _ in range(2): + literal_inlining(root) + # The method's receiver is implicit input 0, so that edge must stay. + assert source.outputs == [method] and method.inputs == [source] + assert method.args == (R0, "x") + assert other.inputs == [] and other.args == ("a",) + assert execute(root) == ("x", "a") + + +_UNCHECKED = object() + + +@pytest.mark.parametrize("make, values, kept, expected", [ + pytest.param(lambda: GetItemOp(key=R1), ["BRASS", 3], [0], "S", id="getitem-keeps-container"), + pytest.param(lambda: MethodCallOp("endswith", args=(R1,), kwargs={}), ["BRASS", "SS"], [0], True, + id="method-keeps-receiver"), + pytest.param(lambda: call(np.add, R0, R1), [2.0, 3.0], [0], 5.0, id="numpy-keeps-primary"), + pytest.param(lambda: BinOp(operator.sub, R0, R1), [np.array([10, 20]), 3], [0], np.array([7, 17]), + id="binop-right"), + pytest.param(lambda: BinOp(operator.sub, R0, R1), [3, np.array([10, 20])], [1], np.array([-7, -17]), + id="binop-left"), + pytest.param(lambda: PredictorOp(estimator=Ridge(), param_refs={"alpha": R1}), + [pd.DataFrame({"x": [1, 2]}), 0.5], [0, 1], _UNCHECKED, id="unsupported-consumer"), +]) +def test_consumer_rules(make, values, kept, expected): + root, sources = inline(make, values) + + assert_kept(root, sources, kept) + if expected is not _UNCHECKED: + np.testing.assert_equal(execute(root), expected) + + +@pytest.mark.parametrize("axis, inlined", [(1, True), ("columns", True), (True, False), (1.0, False), (None, False)]) +def test_concat_inlines_only_int_or_str_axis(axis, inlined): + # `first` is a literal too, but only the axis position is eligible. + frame = pd.DataFrame({"a": [1, 2]}) + root, sources = inline(lambda: ConcatOp(first=R0, others=[R1], axis=R2), [None, frame, axis]) + + assert_kept(root, sources, kept=[0, 1] if inlined else [0, 1, 2]) + if inlined: + pd.testing.assert_frame_equal(execute(root), pd.concat([None, frame], axis=axis)) + + +@pytest.mark.parametrize("make, values", [ + pytest.param(lambda: call(np.add, R0, R1), [2.0, 3.0], id="numpy-literal-primary"), + pytest.param(lambda: call(np.arctan2, R0, R1), [np.array([2.0, 4.0]), 1.0], id="numpy-literal-second"), + pytest.param(lambda: BinOp(operator.sub, R0, R1), [np.array([2, 4]), 10], id="binop-literal-right"), + pytest.param(lambda: BinOp(operator.sub, R0, R1), [10, np.array([2, 4])], id="binop-literal-left"), +]) +def test_inlined_constants_feed_numeric_extraction(make, values): + reference = make() + expected = reference.func(*values) if isinstance(reference, CallOp) else reference.op(*values) + root, _ = inline(make, values) + + root = extract_numeric_operators(root) + + assert isinstance(root, NumericOp) + np.testing.assert_allclose(execute(root), expected) + + +@pytest.mark.parametrize("flag", ["is_X", "is_y"]) +def test_marked_x_y_values_are_not_inlined(flag): + source = ValueOp(3) + setattr(source, flag, True) + root = wire(call(lambda x: x, R0), source) + + literal_inlining(root) + + assert_kept(root, [source], kept=[0]) + assert root.args == (R0,) + + +def test_constant_only_binop_and_unary_keep_inputs(): + root, sources = inline(lambda: BinOp(operator.add, R0, R1), [2, 3]) + assert_kept(root, sources, kept=[0, 1]) + assert (root.left, root.right) == (R0, R1) + + root, sources = inline(lambda: UnaryOp(operator.neg, R0), [3]) + assert_kept(root, sources, kept=[0]) + assert root.operand == R0 + assert execute(root) == -3 + + +def test_shared_scalar_is_not_inlined_into_binops(): + shared = ValueOp(5) + left = wire(BinOp(operator.add, R0, R1), ValueOp(np.array([1, 2])), shared) + right = wire(BinOp(operator.mul, R0, R1), ValueOp(np.array([3, 4])), shared) + root = wire(call(lambda a, b: (a, b), R0, R1), left, right) + + literal_inlining(root) + + assert shared.outputs == [left, right] + assert left.inputs[1] is shared and right.inputs[1] is shared + validate_dag(root) + + +def test_getitem_name_reflects_inlined_key(): + root, sources = inline(lambda: GetItemOp(key=R1, name="Value"), ["BRASS", 3]) + + assert_kept(root, sources, kept=[0]) + assert root.key == 3 and root.name == "3" + + +def test_optconfig_disables_literal_inlining(): + frame = pd.DataFrame({"part": ["BRASS", "STEEL"]}) + dag = st.as_data_op(frame)[st.as_data_op("part")] + + ops, *_ = optimize(dag, OptConfig(literal_inlining=False)) + + assert any(isinstance(op, ValueOp) and op.value == "part" for op in ops) + assert any(isinstance(op, GetItemOp) and op.key == OperandRef(1) for op in ops) + + +def run_plan(ops): + pool = BufferPool() + for op in ops: + pool.put(op, op.process("fit_transform", [pool.pin(key) for key in op.inputs])) + return pool.pin(ops[-1]) + + +def test_inlining_enables_aggregation_fusion(): + frame = pd.DataFrame({"g": ["a", "a", "b"], "v": [1, 2, 3], "w": [4, 5, 6]}) + dag = st.as_data_op(frame).groupby("g").agg(st.var("spec")) + config = OptConfig(dataframe_ops=True) + + root = logical_optimize(dag, config, env={"spec": "sum"}) + validate_dag(root) + assert len([o for o in topological_iterator(root) if isinstance(o, AggregateOp)]) == 1 + + plan, *_ = optimize(dag, config, env={"spec": "sum"}) + pd.testing.assert_frame_equal(run_plan(plan), frame.groupby("g").agg("sum")) + + +def test_inlining_enables_aggregation_parameter_fusion(): + frame = pd.DataFrame({"g": ["a", "a", "b"], "v": [1.0, 2.0, 3.0], "w": [4.0, 5.0, 6.0]}) + dag = st.as_data_op(frame).groupby("g").std(ddof=st.var("d")) + config = OptConfig(dataframe_ops=True) + + root = logical_optimize(dag, config, env={"d": 0}) + validate_dag(root) + agg_ops = [o for o in topological_iterator(root) if isinstance(o, AggregateOp)] + assert len(agg_ops) == 1 and agg_ops[0].aggregations[0][1].params == {"ddof": 0} + + plan, *_ = optimize(dag, config, env={"d": 0}) + pd.testing.assert_frame_equal(run_plan(plan), frame.groupby("g").std(ddof=0)) + + +def test_inlining_enables_projection_extraction(): + frame = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}) + dag = st.as_data_op(frame)[st.as_data_op("x")] + config = OptConfig(dataframe_ops=True) + + root = logical_optimize(dag, config) + validate_dag(root) + ops = list(topological_iterator(root)) + proj = [o for o in ops if isinstance(o, ColumnProjectionOp)] + assert len(proj) == 1 and proj[0].key == "x" + assert not any(isinstance(o, GetItemOp) for o in ops) + + plan, *_ = optimize(dag, config) + assert run_plan(plan).tolist() == [1, 2, 3] + + +def test_inlining_enables_numeric_extraction_with_constant(): + dag = st.as_data_op(2).skb.apply_func(np.add, st.as_data_op(3)) + config = OptConfig() + + root = logical_optimize(dag, config) + validate_dag(root) + numeric = [o for o in topological_iterator(root) if isinstance(o, NumericOp)] + assert len(numeric) == 1 and numeric[0].type is NumericOpType.ADD and numeric[0].constant == 3 + + plan, *_ = optimize(dag, config) + assert run_plan(plan) == 5 + + +def test_inlining_enables_selection_folding(): + frame = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}) + data = st.as_data_op(frame) + thr = st.as_data_op(1) + dag = data[data["x"] > thr].assign(base=thr) + config = OptConfig(dataframe_ops=True) + + root = logical_optimize(dag, config) + validate_dag(root) + ops = list(topological_iterator(root)) + sels = [o for o in ops if isinstance(o, MaskSelectionOp)] + assert len(sels) == 1 + assert sels[0].predicate == BinOpExpr(operator.gt, Col("x"), Const(1)) + assert len(sels[0].inputs) == 1 + assert not any(isinstance(o, ValueOp) for o in ops) + + plan, *_ = optimize(dag, config) + pd.testing.assert_frame_equal(run_plan(plan), frame[frame["x"] > 1].assign(base=1)) + + +def test_optimize_inlines_constants_and_env_variables(): + frame = pd.DataFrame({"part": ["BRASS", "STEEL", None, "RED BRASS"]}) + # `evaluate` resolves `st.var` through its env, so both literals reach the pass as ValueOps. + dag = st.as_data_op(frame)[st.as_data_op("part")].str.endswith(st.var("suffix", "BRASS"), na=False) + + ops, *_ = optimize(dag, env={"suffix": "BRASS"}) + + assert not any(isinstance(op, ValueOp) and isinstance(op.value, str) for op in ops) + pd.testing.assert_series_equal(evaluate(dag), frame["part"].str.endswith("BRASS", na=False)) diff --git a/tests/optimizer/logical/test_map_ops.py b/tests/optimizer/logical/test_map_ops.py index a32c7e96..c47e9555 100644 --- a/tests/optimizer/logical/test_map_ops.py +++ b/tests/optimizer/logical/test_map_ops.py @@ -788,7 +788,7 @@ def test_projection_fused_get_item_with_choice(self): sub_dag1 = data.apply(lambda x, a: x + a, a=st.as_data_op(1)) sub_dag2 = data root = st.choose_from([sub_dag1, sub_dag2]).as_data_op() - ops = optimize(root) + ops = optimize(root, OptConfig(literal_inlining=False)) self.assertEqual(5, len(ops)) self.assertIsInstance(ops[1], ColumnProjectionOp) self.assertIsInstance(ops[3], ApplyUDFOp) diff --git a/tests/optimizer/logical/test_numeric_ops.py b/tests/optimizer/logical/test_numeric_ops.py index 797444d9..76bc58cf 100644 --- a/tests/optimizer/logical/test_numeric_ops.py +++ b/tests/optimizer/logical/test_numeric_ops.py @@ -6,7 +6,7 @@ from sklearn.dummy import DummyRegressor from stratum.optimizer.logical._numeric_ops import NumericOp, NumericOpType, make_binary_numeric_op from stratum.optimizer.logical._ops import CallOp, OperandRef, ValueOp -from stratum.optimizer._optimize import optimize +from stratum.optimizer._optimize import OptConfig, optimize class TestNumericOps(unittest.TestCase): def setUp(self): @@ -274,28 +274,28 @@ def test_extract_np_multiply_callop(self): def test_extract_callop_add_var_var(self): df1 = st.as_data_op(2) df2 = st.as_data_op(3) - out, *_ = optimize(df1.skb.apply_func(np.add, df2)) + out, *_ = optimize(df1.skb.apply_func(np.add, df2), OptConfig(literal_inlining=False)) op = self._assert_var_var_extracted(out, NumericOpType.ADD) self.assertEqual(op.process("fit", [2, 3]), 5) def test_extract_callop_subtract_var_var(self): df1 = st.as_data_op(5) df2 = st.as_data_op(3) - out, *_ = optimize(df1.skb.apply_func(np.subtract, df2)) + out, *_ = optimize(df1.skb.apply_func(np.subtract, df2), OptConfig(literal_inlining=False)) op = self._assert_var_var_extracted(out, NumericOpType.SUBTRACT) self.assertEqual(op.process("fit", [5, 3]), 2) def test_extract_callop_multiply_var_var(self): df1 = st.as_data_op(2) df2 = st.as_data_op(3) - out, *_ = optimize(df1.skb.apply_func(np.multiply, df2)) + out, *_ = optimize(df1.skb.apply_func(np.multiply, df2), OptConfig(literal_inlining=False)) op = self._assert_var_var_extracted(out, NumericOpType.MULTIPLY) self.assertEqual(op.process("fit", [2, 3]), 6) def test_extract_callop_divide_var_var(self): df1 = st.as_data_op(6) df2 = st.as_data_op(2) - out, *_ = optimize(df1.skb.apply_func(np.divide, df2)) + out, *_ = optimize(df1.skb.apply_func(np.divide, df2), OptConfig(literal_inlining=False)) op = self._assert_var_var_extracted(out, NumericOpType.DIVIDE) self.assertEqual(op.process("fit", [6, 2]), 3.0) diff --git a/tests/optimizer/logical/test_projection_ops.py b/tests/optimizer/logical/test_projection_ops.py index 6270ac21..823db299 100644 --- a/tests/optimizer/logical/test_projection_ops.py +++ b/tests/optimizer/logical/test_projection_ops.py @@ -38,8 +38,9 @@ class TestColumnProjectionOp(unittest.TestCase): def setUp(self): self.df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}) - def _ops(self, data_op): - return optimize(data_op, OptConfig(dataframe_ops=True)) + def _ops(self, data_op, **config_kwargs): + config_kwargs.setdefault("literal_inlining", False) + return optimize(data_op, OptConfig(dataframe_ops=True, **config_kwargs)) def test_single_column_parses_to_series_projection(self): ops = self._ops(st.as_data_op(self.df)["x"]) diff --git a/tests/optimizer/logical/test_selection_ops.py b/tests/optimizer/logical/test_selection_ops.py index ef33ac03..383a5c4b 100644 --- a/tests/optimizer/logical/test_selection_ops.py +++ b/tests/optimizer/logical/test_selection_ops.py @@ -279,7 +279,7 @@ def test_shared_graph_fed_constant_stays_leaf(self): data = st.as_data_op(self.df) thr = st.as_data_op(1) ops = optimize(data[data["x"] > thr].assign(base=thr), - OptConfig(dataframe_ops=True)) + OptConfig(dataframe_ops=True, literal_inlining=False)) sel = self._mask(ops) self.assertEqual( BinOpExpr(operator.gt, Col("x"), OperandLeaf(OperandRef(1))),