PERF: Optimize cudf-polars stable hashing - #23582
Conversation
This PR improves the performance of `Node.get_stable_id()` and its callers, most notably `lower_ir_graph_with_node_map`. The old `_expand_hashable` didn't populate / cache the stable hash values of its children, so `lower_ir_graph_with_node_map` had to repeatedly rewalk subtrees and recompute hash values. Now, we go child first. This lets us reuse the cached value from node.get_stable_id().
TomAugspurger
added
improvement
📝 WalkthroughSummary by CodeRabbit
WalkthroughStable node IDs now use cached structural digests. Literal IDs now depend on frozen content rather than object identity. Tests cover digest derivation, equivalent plans, and differing literal values. ChangesStable hashing
Estimated code review effort: 3 (Moderate) | ~15 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudf_polars/cudf_polars/dsl/nodebase.py`:
- Around line 40-42: Update the fallback hashing logic used by
Literal.get_hashable() and LiteralColumn.get_hashable() so it does not encode
process-dependent object identities such as id(self.value). Canonically
serialize supported values into stable bytes, preserving identical
get_stable_id() results for equivalent plans across processes while retaining
distinct encodings for different value types.
- Around line 32-34: Update the Node hashing logic around Node.get_stable_id()
so parent hashing incorporates each child’s full internal digest rather than
only its truncated 32-bit ID. Preserve the 32-bit conversion only when producing
the external stable ID, ensuring distinct child subtrees remain distinguishable
in SerializablePlan.nodes and lower_ir_graph_with_node_map.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a602f4fc-81af-430f-87d7-28fb96df2a81
📒 Files selected for processing (2)
python/cudf_polars/cudf_polars/dsl/nodebase.pypython/cudf_polars/tests/dsl/test_nodebase.py
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
python/cudf_polars/cudf_polars/dsl/expressions/literal.py (1)
110-115: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueComplete the
is_equaldocstring.Document the
otherparameter and the boolean return value. The current docstring only describes behavior.As per coding guidelines, “all public API methods have complete docstrings documenting parameters, return values, and behavior.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/cudf_polars/dsl/expressions/literal.py` around lines 110 - 115, Complete the public LiteralColumn.is_equal docstring by documenting the other parameter, the boolean return value, and its existing content-based equality behavior; do not change the implementation.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudf_polars/cudf_polars/dsl/expressions/literal.py`:
- Around line 33-34: Update the dict-handling branch in _freeze_for_hash to
canonicalize keys with a type-aware, comparable encoding before sorting, so
heterogeneous keys such as integers and strings cannot trigger TypeError while
preserving distinct key identities for Literal.get_stable_id().
In `@python/cudf_polars/tests/expressions/test_literal.py`:
- Around line 175-195: Extend test_literal_stable_id_is_content_based and
test_literal_column_stable_id_is_content_based to cover nested list/dictionary
scalar values, empty Series, all-null Series, and single-element Series,
asserting equal IDs for equivalent values and distinct IDs for differing values.
Add mixed-type Series coverage where the supported Polars dtype permits it,
including nulls or heterogeneous elements as appropriate, to exercise
_freeze_for_hash boundary and recursive handling.
---
Nitpick comments:
In `@python/cudf_polars/cudf_polars/dsl/expressions/literal.py`:
- Around line 110-115: Complete the public LiteralColumn.is_equal docstring by
documenting the other parameter, the boolean return value, and its existing
content-based equality behavior; do not change the implementation.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5967bfb6-f7b6-4b3e-a39e-bb6285e61f6d
📒 Files selected for processing (4)
python/cudf_polars/cudf_polars/dsl/expressions/literal.pypython/cudf_polars/cudf_polars/dsl/nodebase.pypython/cudf_polars/tests/dsl/test_nodebase.pypython/cudf_polars/tests/expressions/test_literal.py
🚧 Files skipped from review as they are similar to previous changes (1)
- python/cudf_polars/cudf_polars/dsl/nodebase.py
| if isinstance(value, dict): | ||
| return tuple(sorted((k, _freeze_for_hash(v)) for k, v in value.items())) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Handle heterogeneous dictionary keys without ordering them directly.
A dictionary can contain keys such as 1 and "1". sorted then compares these keys and raises TypeError. This prevents Literal.get_stable_id() from completing.
Use a type-aware canonical key encoding before sorting, or reject unsupported dictionary key types during construction.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudf_polars/cudf_polars/dsl/expressions/literal.py` around lines 33 -
34, Update the dict-handling branch in _freeze_for_hash to canonicalize keys
with a type-aware, comparable encoding before sorting, so heterogeneous keys
such as integers and strings cannot trigger TypeError while preserving distinct
key identities for Literal.get_stable_id().
| def test_literal_stable_id_is_content_based(): | ||
| a = Literal(DataType(pl.Int64()), 42) | ||
| b = Literal(DataType(pl.Int64()), 42) | ||
| c = Literal(DataType(pl.Int64()), 43) | ||
| assert a.get_stable_id() == b.get_stable_id() | ||
| assert a.get_stable_id() != c.get_stable_id() | ||
| # Distinct Python types with the same dtype stay distinct via value. | ||
| assert ( | ||
| Literal(DataType(pl.String()), "42").get_stable_id() | ||
| != Literal(DataType(pl.Int64()), 42).get_stable_id() | ||
| ) | ||
|
|
||
|
|
||
| def test_literal_column_stable_id_is_content_based(): | ||
| from cudf_polars.dsl.expressions.literal import LiteralColumn | ||
|
|
||
| a = LiteralColumn(DataType(pl.Int64()), pl.Series([1, 2, 3], dtype=pl.Int64())) | ||
| b = LiteralColumn(DataType(pl.Int64()), pl.Series([1, 2, 3], dtype=pl.Int64())) | ||
| c = LiteralColumn(DataType(pl.Int64()), pl.Series([1, 2, 4], dtype=pl.Int64())) | ||
| assert a.get_stable_id() == b.get_stable_id() | ||
| assert a.get_stable_id() != c.get_stable_id() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Add stable-ID coverage for recursive and boundary values.
These tests cover scalar literals and a non-empty, non-null integer Series only. Add cases for nested list or dictionary values, empty Series, all-null Series, and single-element Series. Add mixed-type coverage when the Polars dtype supports it.
This validates the new _freeze_for_hash path and null or size boundaries.
As per coding guidelines, “Ensure test files provide comprehensive edge case coverage (empty, all-null, single-element, mixed types).”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudf_polars/tests/expressions/test_literal.py` around lines 175 - 195,
Extend test_literal_stable_id_is_content_based and
test_literal_column_stable_id_is_content_based to cover nested list/dictionary
scalar values, empty Series, all-null Series, and single-element Series,
asserting equal IDs for equivalent values and distinct IDs for differing values.
Add mixed-type Series coverage where the supported Polars dtype permits it,
including nulls or heterogeneous elements as appropriate, to exercise
_freeze_for_hash boundary and recursive handling.
Source: Coding guidelines
|
|
||
| def _expand_hashable(obj: Any) -> Any: | ||
| """Expand nested Node instances to their hashable form.""" | ||
| def _update_stable_hasher(hasher: Any, obj: Hashable) -> None: |
There was a problem hiding this comment.
Perhaps a nicer way of doing this is something like the following.
The pre-order traversal of an expression uniquely defines it. A non-leaf object is therefore uniquely defined by the structure induced by:
[(type(n), n._non_child) for n in traversal([root])]
I think?
There was a problem hiding this comment.
Ohh I like the idea of using its position in the traversal. Let me see how that works.
There was a problem hiding this comment.
I wasn't immediately able to get this to work :/
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/cudf_polars/cudf_polars/dsl/expressions/literal.py (1)
103-108: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftSeparate local identity from stable hashing.
id(self.value)makesLiteralColumnIDs change after pickle reconstruction and across processes. Provide a process-independent stable representation while retaining identity-based hashing for local replacement. Add a cross-process serialization test.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/cudf_polars/dsl/expressions/literal.py` around lines 103 - 108, Separate LiteralColumn.get_hashable()’s local replacement identity from the representation used for serialization or stable hashing, replacing id(self.value) in the cross-process path with a deterministic process-independent value while preserving identity-based behavior locally. Add a serialization test that pickles and reconstructs a LiteralColumn, including across process boundaries, and verifies the stable representation remains unchanged.
🧹 Nitpick comments (1)
python/cudf_polars/cudf_polars/dsl/expressions/literal.py (1)
103-104: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winComplete the
get_hashable()docstring.The method returns a tuple containing the expression type,
self.dtype.plc_type, andid(self.value). It does not compute a content hash. Add aReturnssection and document the process-local identity behavior.Proposed documentation
- """Compute a hash of the column.""" + """Return the hashable identity key for this literal column. + + Returns + ------- + Hashable + Tuple containing the expression type, PLC dtype, and + process-local identity of ``self.value``. + """As per coding guidelines, all public API methods must document parameters, return values, and behavior.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/cudf_polars/dsl/expressions/literal.py` around lines 103 - 104, Update the public Literal.get_hashable() docstring to describe that it returns a tuple of the expression type, self.dtype.plc_type, and id(self.value), rather than computing a content hash. Add a Returns section and document that the value identity is process-local; no implementation changes are needed.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@python/cudf_polars/cudf_polars/dsl/expressions/literal.py`:
- Around line 103-108: Separate LiteralColumn.get_hashable()’s local replacement
identity from the representation used for serialization or stable hashing,
replacing id(self.value) in the cross-process path with a deterministic
process-independent value while preserving identity-based behavior locally. Add
a serialization test that pickles and reconstructs a LiteralColumn, including
across process boundaries, and verifies the stable representation remains
unchanged.
---
Nitpick comments:
In `@python/cudf_polars/cudf_polars/dsl/expressions/literal.py`:
- Around line 103-104: Update the public Literal.get_hashable() docstring to
describe that it returns a tuple of the expression type, self.dtype.plc_type,
and id(self.value), rather than computing a content hash. Add a Returns section
and document that the value identity is process-local; no implementation changes
are needed.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d0934a44-918d-4b99-96b6-b807620f88a1
📒 Files selected for processing (2)
python/cudf_polars/cudf_polars/dsl/expressions/literal.pypython/cudf_polars/tests/expressions/test_literal.py
💤 Files with no reviewable changes (1)
- python/cudf_polars/tests/expressions/test_literal.py
| # Distinct Python types with the same dtype stay distinct via value. | ||
| assert ( | ||
| Literal(DataType(pl.String()), "42").get_stable_id() | ||
| != Literal(DataType(pl.Int64()), 42).get_stable_id() | ||
| ) |
There was a problem hiding this comment.
Those have different dtypes, so at the very least this comment is wrong.
3 participants
Description
This PR improves the performance of
Node.get_stable_id()and its callers, most notablylower_ir_graph_with_node_map. The old_expand_hashabledidn't populate / cache the stable hash values of its children, solower_ir_graph_with_node_maphad to repeatedly rewalk subtrees and recompute hash values.Here's an impossible to read flamegraph of
lower_ir_graph_with_node_mapon a benchmark that looks like tpc-h query 8. Just notice the shape (very nested):Now, we go child first. This lets us reuse the cached value from node.get_stable_id().
On
main,lower_ir_graph_with_node_maptook 55-75ms for me on tpc-h query 8 at sf-1k. On this branch, it takes ~5-6ms.