Fix stable ID for DataFrameScan (rapidsai#22091)

cuDF-Polars tracing infrastructure relies on the concept of "stable ID" for each IR node. This ID is **not** currently stable for `DataFrameScan`. This PR adds a simple `_expand_hashable` helper utility to correct this.

Authors:
  - Richard (Rick) Zamora (https://github.com/rjzamora)

Approvers:
  - Tom Augspurger (https://github.com/TomAugspurger)

URL: rapidsai#22091

Author: rjzamora

python/cudf_polars/cudf_polars/dsl/nodebase.py

@@ -18,6 +18,15 @@
 T = TypeVar("T", bound="Node[Any]")
 
 
+def _expand_hashable(obj: Any) -> Any:
+    """Expand nested Node instances to their hashable form."""
+    if isinstance(obj, Node):
+        return _expand_hashable(obj.get_hashable())
+    elif isinstance(obj, tuple):
+        return tuple(_expand_hashable(x) for x in obj)
+    return obj
+
+
 class Node(Generic[T]):
     """
     An abstract node type.
@@ -103,7 +112,7 @@ def get_stable_id(self) -> int:
         try:
             return self._stable_hash_value
         except AttributeError:
-            content = repr(self.get_hashable()).encode("utf-8")
+            content = repr(_expand_hashable(self)).encode("utf-8")
             self._stable_hash_value = int(hashlib.md5(content).hexdigest()[:8], 16)
             return self._stable_hash_value
 

python/cudf_polars/tests/experimental/test_dataframescan.py

@@ -10,11 +10,17 @@
 import polars as pl
 
 from cudf_polars import Translator
+from cudf_polars.dsl.traversal import traversal
 from cudf_polars.experimental.parallel import lower_ir_graph
 from cudf_polars.testing.asserts import assert_gpu_result_equal
 from cudf_polars.utils.config import ConfigOptions
 
 
+def _assert_stable_ids_match(orig, loaded) -> None:
+    for a, b in zip(traversal([orig]), traversal([loaded]), strict=True):
+        assert a.get_stable_id() == b.get_stable_id()
+
+
 @pytest.fixture(scope="module")
 def df():
     return pl.LazyFrame(
@@ -63,6 +69,21 @@ def test_dataframescan_concat(df, engine):
     assert_gpu_result_equal(df2, engine=engine)
 
 
+def test_join_in_memory_lazy_stable_id_pickle():
+    engine = pl.GPUEngine(
+        raise_on_fail=True,
+        executor="streaming",
+        executor_options={"max_rows_per_partition": 1_000},
+    )
+    left = pl.LazyFrame({"k": [1, 2, 3], "x": [10, 20, 30]}).collect().lazy()
+    right = pl.LazyFrame({"k": [2, 3, 4], "y": [1, 2, 3]}).collect().lazy()
+    ir, _, _ = lower_ir_graph(
+        Translator(left.join(right, on="k")._ldf.visit(), engine).translate_ir(),
+        ConfigOptions.from_polars_engine(engine),
+    )
+    _assert_stable_ids_match(ir, pickle.loads(pickle.dumps(ir)))
+
+
 def test_dataframescan_pickle(df):
     _engine = pl.GPUEngine(
         raise_on_fail=True,
@@ -79,3 +100,4 @@ def test_dataframescan_pickle(df):
     # Verify the unpickled IR is equivalent
     assert type(unpickled_ir) is type(ir)
     assert unpickled_ir.schema == ir.schema
+    _assert_stable_ids_match(ir, unpickled_ir)