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 # ============================================================================