From 77599ce13c556d486b30a5cc4bd9d9778fae9aee Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 7 Aug 2026 07:34:34 -0700 Subject: [PATCH 1/4] PERF: Optimize cudf-polars stable hashing 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(). --- .../cudf_polars/cudf_polars/dsl/nodebase.py | 40 ++++++++++++------- python/cudf_polars/tests/dsl/test_nodebase.py | 5 ++- 2 files changed, 28 insertions(+), 17 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/nodebase.py b/python/cudf_polars/cudf_polars/dsl/nodebase.py index 64375b4a2f84..d510f5022a3a 100644 --- a/python/cudf_polars/cudf_polars/dsl/nodebase.py +++ b/python/cudf_polars/cudf_polars/dsl/nodebase.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Base class for IR nodes, and utilities.""" @@ -21,13 +21,25 @@ T = TypeVar("T", bound="Node[Any]") -def _expand_hashable(obj: Any) -> Any: - """Expand nested Node instances to their hashable form.""" +def _update_stable_hasher(hasher: Any, obj: Any) -> None: + """ + Feed ``hasher`` a bottom-up structural digest of ``obj``. + + Nested :class:`Node` instances contribute their cached + :meth:`~Node.get_stable_id` rather than a full subtree expansion, so + hashing an entire DAG is linear in total local payload size. + """ if isinstance(obj, Node): - return _expand_hashable(obj.get_hashable()) + hasher.update(b"N") + hasher.update(obj.get_stable_id().to_bytes(4, "big")) elif isinstance(obj, tuple): - return tuple(_expand_hashable(x) for x in obj) - return obj + hasher.update(b"T") + hasher.update(len(obj).to_bytes(4, "big")) + for x in obj: + _update_stable_hasher(hasher, x) + else: + hasher.update(b"L") + hasher.update(repr(obj).encode("utf-8")) class Node(Generic[T]): @@ -106,13 +118,10 @@ def get_stable_id(self) -> int: """ Compute a stable identifier for Node. - Uses MD5 hash of the node's hashable representation for determinism - across process boundaries (Python's hash() uses PYTHONHASHSEED). - - Parameters - ---------- - ir_node - The IR node. + Digests :meth:`get_hashable` bottom-up with MD5 so the result is + deterministic across process boundaries (Python's ``hash()`` uses + ``PYTHONHASHSEED``). Nested nodes contribute their own stable ids + rather than re-expanding their subtrees. Returns ------- @@ -122,8 +131,9 @@ def get_stable_id(self) -> int: try: return self._stable_hash_value except AttributeError: - content = repr(_expand_hashable(self)).encode("utf-8") - self._stable_hash_value = int(hashlib.md5(content).hexdigest()[:8], 16) + h = hashlib.md5(usedforsecurity=False) + _update_stable_hasher(h, self.get_hashable()) + self._stable_hash_value = int(h.hexdigest()[:8], 16) return self._stable_hash_value def get_stable_plan_id(self) -> uuid.UUID: diff --git a/python/cudf_polars/tests/dsl/test_nodebase.py b/python/cudf_polars/tests/dsl/test_nodebase.py index 7e3985461390..a20f29853fd5 100644 --- a/python/cudf_polars/tests/dsl/test_nodebase.py +++ b/python/cudf_polars/tests/dsl/test_nodebase.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -95,7 +95,7 @@ def test_get_stable_id(node: Node): assert node.get_stable_id() == node_id -def test_get_stable_plan_id(): +def test_get_stable_plan_id() -> None: node = expr.BinOp( DataType(pl.Int64()), plc.binaryop.BinaryOperator.ADD, @@ -115,6 +115,7 @@ def test_get_stable_plan_id(): expr.Literal(DataType(pl.Int64()), 1), ) assert node2.get_stable_plan_id() == plan_id + assert node.children[0].get_stable_id() == node2.children[0].get_stable_id() # And uniqueness node3 = node.children[0] From 931f0172b7641b2dcde8db8f454283a1658cb8a8 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 7 Aug 2026 07:59:26 -0700 Subject: [PATCH 2/4] cleanup --- python/cudf_polars/cudf_polars/dsl/nodebase.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/nodebase.py b/python/cudf_polars/cudf_polars/dsl/nodebase.py index d510f5022a3a..eb9c770d8faa 100644 --- a/python/cudf_polars/cudf_polars/dsl/nodebase.py +++ b/python/cudf_polars/cudf_polars/dsl/nodebase.py @@ -120,8 +120,7 @@ def get_stable_id(self) -> int: Digests :meth:`get_hashable` bottom-up with MD5 so the result is deterministic across process boundaries (Python's ``hash()`` uses - ``PYTHONHASHSEED``). Nested nodes contribute their own stable ids - rather than re-expanding their subtrees. + ``PYTHONHASHSEED``). Returns ------- From a3a536d3271c4a77feafe09511553b0cb2d9f527 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 7 Aug 2026 08:23:42 -0700 Subject: [PATCH 3/4] bug fixes --- .../cudf_polars/dsl/expressions/literal.py | 40 ++++++++++++++----- .../cudf_polars/cudf_polars/dsl/nodebase.py | 30 ++++++++++---- python/cudf_polars/tests/dsl/test_nodebase.py | 3 ++ .../tests/expressions/test_literal.py | 23 +++++++++++ 4 files changed, 78 insertions(+), 18 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/literal.py b/python/cudf_polars/cudf_polars/dsl/expressions/literal.py index cecbbb44a536..0887f0bda331 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/literal.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/literal.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # TODO: remove need for this # ruff: noqa: D101 @@ -23,6 +23,20 @@ __all__ = ["Literal", "LiteralColumn"] +def _freeze_for_hash(value: Any) -> Hashable: + """ + Convert ``value`` into a process-independent hashable form. + + Nested ``list`` / ``dict`` values (e.g. list/struct literals) are frozen + into tuples so they can appear in :meth:`Node.get_hashable` results. + """ + if isinstance(value, dict): + return tuple(sorted((k, _freeze_for_hash(v)) for k, v in value.items())) + if isinstance(value, list): + return tuple(_freeze_for_hash(v) for v in value) + return value + + class Literal(Expr): __slots__ = ("value",) _non_child = ("dtype", "value") @@ -57,9 +71,8 @@ def agg_request(self) -> NoReturn: # noqa: D102 "Not expecting to require agg request of literal" ) # pragma: no cover - def get_hashable(self) -> Hashable: - """Get the hash of the literal.""" - return (type(self), self.dtype.plc_type, id(self.value)) + def get_hashable(self) -> Hashable: # noqa: D102 + return (type(self), self.dtype.plc_type, _freeze_for_hash(self.value)) def astype(self, dtype: DataType) -> Literal: """Cast self to dtype.""" @@ -87,12 +100,19 @@ def __init__(self, dtype: DataType, value: pl.Series) -> None: self.children = () self.is_pointwise = True - def get_hashable(self) -> Hashable: - """Compute a hash of the column.""" - # This is stricter than necessary, but we only need this hash - # for identity in groupby replacements so it's OK. And this - # way we avoid doing potentially expensive compute. - return (type(self), self.dtype.plc_type, id(self.value)) + def get_hashable(self) -> Hashable: # noqa: D102 + return ( + type(self), + self.dtype.plc_type, + _freeze_for_hash(self.value.to_list()), + ) + + def is_equal(self, other: LiteralColumn) -> bool: + """Equality that compares Series values by content, not identity.""" + if self is other: + return True + # pl.Series.__eq__ is elementwise and cannot be used as a scalar bool. + return self.dtype == other.dtype and self.value.equals(other.value) def do_evaluate( self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME diff --git a/python/cudf_polars/cudf_polars/dsl/nodebase.py b/python/cudf_polars/cudf_polars/dsl/nodebase.py index eb9c770d8faa..90074ee68416 100644 --- a/python/cudf_polars/cudf_polars/dsl/nodebase.py +++ b/python/cudf_polars/cudf_polars/dsl/nodebase.py @@ -21,17 +21,18 @@ T = TypeVar("T", bound="Node[Any]") -def _update_stable_hasher(hasher: Any, obj: Any) -> None: +def _update_stable_hasher(hasher: Any, obj: Hashable) -> None: """ Feed ``hasher`` a bottom-up structural digest of ``obj``. - Nested :class:`Node` instances contribute their cached - :meth:`~Node.get_stable_id` rather than a full subtree expansion, so - hashing an entire DAG is linear in total local payload size. + Nested :class:`Node` instances contribute their full cached digest + rather than a truncated id or a full subtree expansion, so hashing an + entire DAG is linear in total local payload size while distinct child + subtrees remain distinguishable when composing parent digests. """ if isinstance(obj, Node): hasher.update(b"N") - hasher.update(obj.get_stable_id().to_bytes(4, "big")) + hasher.update(obj._get_stable_digest()) elif isinstance(obj, tuple): hasher.update(b"T") hasher.update(len(obj).to_bytes(4, "big")) @@ -61,11 +62,13 @@ class Node(Generic[T]): __slots__ = ( "_hash_value", "_repr_value", + "_stable_digest", "_stable_hash_value", "_stable_plan_id", "children", ) _hash_value: int + _stable_digest: bytes _stable_hash_value: int _stable_plan_id: uuid.UUID _repr_value: str @@ -114,6 +117,16 @@ def get_hashable(self) -> Hashable: """ return (type(self), self._ctor_arguments(self.children)) + def _get_stable_digest(self) -> bytes: + """Return the full MD5 digest for this node, computing it if needed.""" + try: + return self._stable_digest + except AttributeError: + h = hashlib.md5(usedforsecurity=False) + _update_stable_hasher(h, self.get_hashable()) + self._stable_digest = h.digest() + return self._stable_digest + def get_stable_id(self) -> int: """ Compute a stable identifier for Node. @@ -130,9 +143,10 @@ def get_stable_id(self) -> int: try: return self._stable_hash_value except AttributeError: - h = hashlib.md5(usedforsecurity=False) - _update_stable_hasher(h, self.get_hashable()) - self._stable_hash_value = int(h.hexdigest()[:8], 16) + # First 4 digest bytes == int(hexdigest()[:8], 16). + self._stable_hash_value = int.from_bytes( + self._get_stable_digest()[:4], "big" + ) return self._stable_hash_value def get_stable_plan_id(self) -> uuid.UUID: diff --git a/python/cudf_polars/tests/dsl/test_nodebase.py b/python/cudf_polars/tests/dsl/test_nodebase.py index a20f29853fd5..eb7c27208aac 100644 --- a/python/cudf_polars/tests/dsl/test_nodebase.py +++ b/python/cudf_polars/tests/dsl/test_nodebase.py @@ -93,6 +93,9 @@ def test_get_stable_id(node: Node): assert isinstance(node_id, int) # Second call should return the cached value assert node.get_stable_id() == node_id + # External id is a truncation of the full digest used for composition. + assert node_id == int.from_bytes(node._get_stable_digest()[:4], "big") + assert len(node._get_stable_digest()) == 16 def test_get_stable_plan_id() -> None: diff --git a/python/cudf_polars/tests/expressions/test_literal.py b/python/cudf_polars/tests/expressions/test_literal.py index 8665112f5c9c..6d2736604b38 100644 --- a/python/cudf_polars/tests/expressions/test_literal.py +++ b/python/cudf_polars/tests/expressions/test_literal.py @@ -172,6 +172,29 @@ def test_literal_hash(dtype, val): assert isinstance(hash(Literal(DataType(dtype), val)), int) +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() + + def test_struct_literal_not_supported(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [1, 2, 3]}) q = df.select(pl.lit({"x": 1, "y": "foo"})) From 150aaa056ce83f0d610c62ea4ad22a4e57a42472 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 7 Aug 2026 09:13:38 -0700 Subject: [PATCH 4/4] revert LiteralColumn changes --- .../cudf_polars/dsl/expressions/literal.py | 19 ++++++------------- .../tests/expressions/test_literal.py | 10 ---------- 2 files changed, 6 insertions(+), 23 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/literal.py b/python/cudf_polars/cudf_polars/dsl/expressions/literal.py index 0887f0bda331..46e668c4b086 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/literal.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/literal.py @@ -100,19 +100,12 @@ def __init__(self, dtype: DataType, value: pl.Series) -> None: self.children = () self.is_pointwise = True - def get_hashable(self) -> Hashable: # noqa: D102 - return ( - type(self), - self.dtype.plc_type, - _freeze_for_hash(self.value.to_list()), - ) - - def is_equal(self, other: LiteralColumn) -> bool: - """Equality that compares Series values by content, not identity.""" - if self is other: - return True - # pl.Series.__eq__ is elementwise and cannot be used as a scalar bool. - return self.dtype == other.dtype and self.value.equals(other.value) + def get_hashable(self) -> Hashable: + """Compute a hash of the column.""" + # This is stricter than necessary, but we only need this hash + # for identity in groupby replacements so it's OK. And this + # way we avoid doing potentially expensive compute. + return (type(self), self.dtype.plc_type, id(self.value)) def do_evaluate( self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME diff --git a/python/cudf_polars/tests/expressions/test_literal.py b/python/cudf_polars/tests/expressions/test_literal.py index 6d2736604b38..adc49d9137c4 100644 --- a/python/cudf_polars/tests/expressions/test_literal.py +++ b/python/cudf_polars/tests/expressions/test_literal.py @@ -185,16 +185,6 @@ def test_literal_stable_id_is_content_based(): ) -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() - - def test_struct_literal_not_supported(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [1, 2, 3]}) q = df.select(pl.lit({"x": 1, "y": "foo"}))