Repository navigation
Conversation
6286638 to
fda6f02
Compare
Codecov Report❌ Patch coverage is
Flags with carried forward coverage won't be shown. Click here to find out more.
... and 1 file with indirect coverage changes 🚀 New features to boost your workflow:
|
|
@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 Stop inlining into BinOp. Numeric extraction and selection folding already read |
|
closed in c94554a |
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 remainingOperandRefindices are remapped accordingly.Closes #251
Changes
stratum/optimizer/logical/_literal_inlining.py, executed immediately after conversion to the logical operator DAG.OptConfig(literal_inlining=True)(enabled by default), allowing the pass to be disabled for debugging, comparison, or isolated optimizer tests.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 forValueOpnodes 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 forOperandRefinstances 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 inrecords[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:
Once the mapping is complete, the consumer's fields are traversed again using
remap_operand_refs(), which updates any remainingOperandRefindices 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:
ValueOps markedis_Xoris_yare excluded becauseadd_splitting_oprelies on these markers.BinOp. Shared scalar producers are preserved for binary consumers to retain algebraic rewrite opportunities such as identity elimination andx * x → square(x).UnaryOpis excluded to avoid removing its only graph input.MethodCallOpandGetItemOpremains unchanged.ConcatOponly 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:
ValueOpnodes 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.OperandRefreferences.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:OperandRefremapping, graph validity, nested containers, keyword arguments, slices, dictionary subclasses, and shared producers.BinOps, shared scalar producers, excludedUnaryOps, implicit receivers, NumPy primary inputs, andGetItemOpnaming.ddof), column-projection extraction, constant-aware numeric extraction, and selection folding. Tests validate the logical DAG and compare physical-plan execution with pandas results.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, andtest_map_ops.pyretain 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:
Result: 45 passed (focused file), 999 passed, 1 skipped (logical optimizer suite).