Skip to content

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

Closed
Aram164 wants to merge 1 commit into
deem-data:mainfrom
Aram164:feature/literal-inlining
Closed

Aram164 wants to merge 1 commit into
deem-data:mainfrom
Aram164:feature/literal-inlining

Conversation

@Aram164

@Aram164 Aram164 commented Oct 20, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Adds a literal-inlining pass to the logical optimizer. Immutable built-in scalar values (None, bool, int, float, str, bytes) are substituted directly into eligible consuming operators. Redundant input edges are removed, and remaining OperandRef indices are remapped accordingly.

Closes #251

Changes

  • Introduces stratum/optimizer/logical/_literal_inlining.py, executed immediately after conversion to the logical operator DAG.
  • Adds OptConfig(literal_inlining=True) (enabled by default), allowing the pass to be disabled for debugging, comparison, or isolated optimizer tests.
  • Updates GetItemOp.update_name() to reflect inlined literal keys in graph and explain output while preserving existing names for unresolved graph-fed keys.

Implementation

The pass operates on the logical DAG in three steps: constant discovery, recursive literal substitution, and graph rewiring.

Step 1 — Constant discovery and consumer eligibility

After convert_to_ops() constructs the logical DAG, the pass traverses its nodes in topological order, looking for ValueOp nodes containing eligible immutable scalar constants.

For each eligible ValueOp, the pass examines its outgoing edges to identify consuming operators and the positions at which the constant appears in their input lists.

Each occurrence is then checked using _inlinable(), which determines whether the consumer can safely accept the literal at that position while respecting operator-specific restrictions and downstream optimization requirements.

Eligible occurrences proceed to recursive field substitution.

Step 2 — Recursive field substitution

Once a consumer input passes the eligibility check, _check_for_constant() iterates over the consumer's fields, retrieving the parameter stored in each field.

For each parameter, it invokes _inline_literal(), passing the parameter, the constant producer's input position, and its literal value.

_inline_literal() recursively traverses supported container types (tuples, lists, dictionaries, and slices), looking for OperandRef instances referencing the specified input position. When a matching reference is found, it is replaced with the literal value. The function propagates both the updated parameter and a boolean indicating whether a substitution occurred back through the recursive calls.

If a substitution occurred, _check_for_constant() updates the corresponding consumer field with the reconstructed parameter. It then returns a boolean indicating whether any of the consumer's fields were modified.

Finally, literal_inlining() uses this result to record the successfully inlined input position in records[consumer], allowing the redundant dependency to be removed during the subsequent graph-rewiring step.

At this stage, the consumer's fields have been modified, but its input and output edges remain unchanged.

Step 3 — Edge removal and operand remapping

After all eligible substitutions have been recorded, _rewire_edges() processes each affected consumer alongside the set of input positions whose constants were inlined.

For each consumer, a new input list is constructed by excluding the recorded positions, effectively removing the redundant dependencies.

The original input list is then traversed to construct a mapping from old input indices to new ones. A running offset tracks how many inputs have been removed:

  • For removed inputs, the offset is incremented. The consumer is also removed from the producer's output list if that producer no longer appears among the consumer's remaining inputs.
  • For retained inputs, the original index is mapped to its new position by subtracting the accumulated offset.

Once the mapping is complete, the consumer's fields are traversed again using remap_operand_refs(), which updates any remaining OperandRef indices to match the compacted input list.

Finally, the consumer's original input list is replaced with the compacted version, completing the DAG rewiring while preserving input/output edge consistency and valid operand references.

By collecting substitutions per consumer before modifying graph edges, the pass avoids repeated input-index shifts and can remap all remaining references consistently.

Inlining restrictions

The pass conservatively preserves dependencies required by downstream optimizations:

  • Split markers: ValueOps marked is_X or is_y are excluded because add_splitting_op relies on these markers.
  • Binary operators: Inlining cannot eliminate every graph input of a BinOp. Shared scalar producers are preserved for binary consumers to retain algebraic rewrite opportunities such as identity elimination and x * x → square(x).
  • Unary operators: UnaryOp is excluded to avoid removing its only graph input.
  • Implicit receivers: Input position 0 of MethodCallOp and GetItemOp remains unchanged.
  • NumPy calls: The primary input is preserved to satisfy numeric extraction assumptions.
  • Concatenation: ConcatOp only permits inlining integer or string axis values.

Performance

The pass replaces eligible graph-backed scalar constants with literal values stored directly in consumer fields.

This has two benefits:

  • Reduced DAG dependencies: Redundant input edges are removed, allowing scalar ValueOp nodes with no remaining consumers to disappear from the reachable optimized plan. This reduces the number of values that may need to be scheduled, materialized as operator outputs, and managed by the runtime.
  • Improved downstream optimization opportunities: Known constants become directly accessible to existing optimization passes, enabling specialized operators and transformations that would otherwise be blocked by graph-fed OperandRef references.

The implementation separates literal substitution from graph rewiring, allowing multiple input positions of the same consumer to be compacted together while preserving operand-index consistency.

The pass does not evaluate expressions or materialize input data.

Tests

Adds tests/optimizer/logical/test_literal_inlining.py, covering:

  • Direct transformations: Scalar eligibility, input compaction, OperandRef remapping, graph validity, nested containers, keyword arguments, slices, dictionary subclasses, and shared producers.
  • Safety restrictions: Protected X/y markers, constant-only BinOps, shared scalar producers, excluded UnaryOps, implicit receivers, NumPy primary inputs, and GetItemOp naming.
  • Cross-pass integration: Demonstrates that inlining enables aggregation fusion (including specifications and ddof), column-projection extraction, constant-aware numeric extraction, and selection folding. Tests validate the logical DAG and compare physical-plan execution with pandas results.
  • Configuration: Confirms that OptConfig(literal_inlining=False) preserves the previous plan structure.

Existing tests in test_numeric_ops.py, test_aggregation_ops.py, test_projection_ops.py, test_selection_ops.py, and test_map_ops.py retain their original inputs and assertions. Affected unit tests explicitly disable literal inlining to isolate the optimization under test, while the new integration tests verify the combined behavior with inlining enabled.

Tests run:

.venv/bin/python -m pytest tests/optimizer/logical/test_literal_inlining.py -q
.venv/bin/python -m pytest tests/optimizer/logical -q

Result: 45 passed (focused file), 999 passed, 1 skipped (logical optimizer suite).

@Aram164
Aram164 force-pushed the feature/literal-inlining branch from 6286638 to fda6f02 Compare October 20, 2026 01:41
@codecov

codecov Bot commented Oct 20, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 91.00000% with 9 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
stratum/optimizer/logical/_ops.py 38.46% 4 Missing and 4 partials ⚠️
stratum/optimizer/logical/_literal_inlining.py 98.79% 0 Missing and 1 partial ⚠️
Flag Coverage Δ
unittests 92.69% <91.00%> (+<0.01%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
stratum/optimizer/_optimize.py 95.98% <100.00%> (+0.04%) ⬆️
stratum/optimizer/logical/_literal_inlining.py 98.79% <98.79%> (ø)
stratum/optimizer/logical/_ops.py 96.10% <38.46%> (-1.11%) ⬇️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@e-strauss

Copy link
Copy Markdown
Collaborator

@Aram164 thanks, during the merge I simplified two things by reusing existing helper function from stratums:

Reuse _resolve_operand for the substitution: a list mapping inlined slots to
their value and kept slots to their new OperandRef inlines and renumbers in one
call, replacing the hand-rolled walker and the separate remapping step. Drive
the pass from consumers instead of ValueOps.

Stop inlining into BinOp. Numeric extraction and selection folding already read
a ValueOp operand as a constant, and the shared-constant and not-all-literal
rules only existed to keep algebraic-rewrite tests on scalar fixtures green.

@e-strauss

Copy link
Copy Markdown
Collaborator

closed in c94554a

@e-strauss e-strauss closed this Oct 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Optimizer] Inline constant scalar ValueOps into consuming operators

2 participants