diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/literal.py b/python/cudf_polars/cudf_polars/dsl/expressions/literal.py index cecbbb44a536..46e668c4b086 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.""" diff --git a/python/cudf_polars/cudf_polars/dsl/nodebase.py b/python/cudf_polars/cudf_polars/dsl/nodebase.py index 64375b4a2f84..90074ee68416 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,26 @@ 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: Hashable) -> None: + """ + Feed ``hasher`` a bottom-up structural digest of ``obj``. + + 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): - return _expand_hashable(obj.get_hashable()) + hasher.update(b"N") + hasher.update(obj._get_stable_digest()) 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]): @@ -49,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 @@ -102,17 +117,23 @@ 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. - 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``). Returns ------- @@ -122,8 +143,10 @@ 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) + # 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 7e3985461390..eb7c27208aac 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 @@ -93,9 +93,12 @@ 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(): +def test_get_stable_plan_id() -> None: node = expr.BinOp( DataType(pl.Int64()), plc.binaryop.BinaryOperator.ADD, @@ -115,6 +118,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] diff --git a/python/cudf_polars/tests/expressions/test_literal.py b/python/cudf_polars/tests/expressions/test_literal.py index 8665112f5c9c..adc49d9137c4 100644 --- a/python/cudf_polars/tests/expressions/test_literal.py +++ b/python/cudf_polars/tests/expressions/test_literal.py @@ -172,6 +172,19 @@ 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_struct_literal_not_supported(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [1, 2, 3]}) q = df.select(pl.lit({"x": 1, "y": "foo"}))