From 0db2dfd213bfd7cabdb07dc7b5cbefdce4263d0e Mon Sep 17 00:00:00 2001 From: Liudeng Zhang Date: Sun, 8 Mar 2026 22:29:47 -0500 Subject: [PATCH] Preserve attrs on Dask DataFrame sample() Add "sample" to the list of wrapped methods in _accessor.py so that calling .sample() on a Dask DataFrame retains its .attrs metadata, consistent with the other wrapped methods. --- src/spatialdata/models/_accessor.py | 1 + tests/models/test_accessor.py | 9 +++++++++ 2 files changed, 10 insertions(+) diff --git a/src/spatialdata/models/_accessor.py b/src/spatialdata/models/_accessor.py index 1b635427c..3174ee80d 100644 --- a/src/spatialdata/models/_accessor.py +++ b/src/spatialdata/models/_accessor.py @@ -124,6 +124,7 @@ def __repr__(self) -> str: "copy", "drop", "map_partitions", + "sample", "set_index", ]: wrap_method_with_attrs(method_name=method_name, dask_class=DaskDataFrame) diff --git a/tests/models/test_accessor.py b/tests/models/test_accessor.py index 6f21a1905..6a8fd47a0 100644 --- a/tests/models/test_accessor.py +++ b/tests/models/test_accessor.py @@ -130,6 +130,15 @@ def test_dataframe_set_index_preserves_attrs(): assert isinstance(result.attrs, AttrsAccessor) +def test_dataframe_sample_preserves_attrs(): + """Test that DataFrame.sample preserves attrs.""" + df = dd.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}), npartitions=2) + df.attrs["key"] = "value" + result = df.sample(frac=0.5) + assert result.attrs["key"] == "value" + assert isinstance(result.attrs, AttrsAccessor) + + # ============================================================================ # Series wrapped methods tests # ============================================================================