From d684a03138fc351f38b37edc89587e8832478eeb Mon Sep 17 00:00:00 2001 From: anon Date: Thu, 7 May 2026 22:36:06 +0200 Subject: [PATCH 1/4] Fix double .compute() when render_points color= is a native column (#633) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The _preloaded optimisation previously only fired when color data came from an AnnData table (added_color_from_table=True). For element-native color columns the flag stayed False, so _set_color_source_vec fell back to a fresh get_values() call — a second .compute() on the dask DataFrame. Fix: widen the condition to col_for_color in points_pd_with_color.columns, which is True in both the native-column and table-sourced cases. Co-Authored-By: Claude Sonnet 4.6 --- src/spatialdata_plot/pl/render.py | 9 +++++++-- tests/pl/test_render_points.py | 33 +++++++++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/src/spatialdata_plot/pl/render.py b/src/spatialdata_plot/pl/render.py index eb0fc300..2d3d8f46 100644 --- a/src/spatialdata_plot/pl/render.py +++ b/src/spatialdata_plot/pl/render.py @@ -832,8 +832,13 @@ def _render_points( # from the registered points (see above) avoids duplicate-origin ambiguities. color_table_name = table_name - # Reuse color data already loaded from the table to avoid a redundant get_values() call. - _preloaded = points_pd_with_color[col_for_color] if added_color_from_table and col_for_color is not None else None + # Reuse already-loaded color data to avoid a redundant get_values() call — applies to both + # native-column and table-sourced cases, both of which are present in points_pd_with_color. + _preloaded = ( + points_pd_with_color[col_for_color] + if col_for_color is not None and col_for_color in points_pd_with_color.columns + else None + ) color_source_vector, color_vector, _ = _set_color_source_vec( sdata=sdata_filt, diff --git a/tests/pl/test_render_points.py b/tests/pl/test_render_points.py index bdb3218a..c63fc7bc 100644 --- a/tests/pl/test_render_points.py +++ b/tests/pl/test_render_points.py @@ -1,5 +1,6 @@ import math +import dask import dask.dataframe import datashader as ds import matplotlib @@ -1004,3 +1005,35 @@ def test_no_table_fallback_warning_for_element_column(caplog): with logger_no_warns(caplog, logger, match="fallback for color mapping"): sdata.pl.render_points("pts", color="cell_type").pl.show() plt.close("all") + + +def test_render_points_native_color_column_single_compute(): + # Regression test for #633: color column native to the points element triggered + # two .compute() calls instead of one. + dask.config.set({"dataframe.query-planning": False}) + compute_calls = [] + original_compute = dask.dataframe.DataFrame.compute + + def counting_compute(self, **kwargs): + compute_calls.append(True) + return original_compute(self, **kwargs) + + dask.dataframe.DataFrame.compute = counting_compute + try: + rng = np.random.default_rng(42) + n = 100 + df = pd.DataFrame( + { + "x": rng.uniform(0, 100, n), + "y": rng.uniform(0, 100, n), + "cell_type": pd.Categorical(rng.choice(["A", "B", "C"], n)), + } + ) + points = PointsModel.parse(df) + sdata = SpatialData(points={"pts": points}) + sdata.pl.render_points("pts", color="cell_type").pl.show() + plt.close("all") + finally: + dask.dataframe.DataFrame.compute = original_compute + + assert len(compute_calls) == 1, f"Expected 1 .compute() call, got {len(compute_calls)}" From 972c1f5b7b1ae74727235ae5de7ea8c8c01e9884 Mon Sep 17 00:00:00 2001 From: anon Date: Thu, 7 May 2026 23:17:28 +0200 Subject: [PATCH 2/4] simplify: trim comment; scope dask config in regression test Co-Authored-By: Claude Sonnet 4.6 --- src/spatialdata_plot/pl/render.py | 3 +-- tests/pl/test_render_points.py | 33 ++++++++++++++++--------------- 2 files changed, 18 insertions(+), 18 deletions(-) diff --git a/src/spatialdata_plot/pl/render.py b/src/spatialdata_plot/pl/render.py index 2d3d8f46..c8b9a09e 100644 --- a/src/spatialdata_plot/pl/render.py +++ b/src/spatialdata_plot/pl/render.py @@ -832,8 +832,7 @@ def _render_points( # from the registered points (see above) avoids duplicate-origin ambiguities. color_table_name = table_name - # Reuse already-loaded color data to avoid a redundant get_values() call — applies to both - # native-column and table-sourced cases, both of which are present in points_pd_with_color. + # Reuse already-loaded color data to avoid a redundant get_values() call. _preloaded = ( points_pd_with_color[col_for_color] if col_for_color is not None and col_for_color in points_pd_with_color.columns diff --git a/tests/pl/test_render_points.py b/tests/pl/test_render_points.py index c63fc7bc..b5408185 100644 --- a/tests/pl/test_render_points.py +++ b/tests/pl/test_render_points.py @@ -1010,7 +1010,6 @@ def test_no_table_fallback_warning_for_element_column(caplog): def test_render_points_native_color_column_single_compute(): # Regression test for #633: color column native to the points element triggered # two .compute() calls instead of one. - dask.config.set({"dataframe.query-planning": False}) compute_calls = [] original_compute = dask.dataframe.DataFrame.compute @@ -1018,22 +1017,24 @@ def counting_compute(self, **kwargs): compute_calls.append(True) return original_compute(self, **kwargs) - dask.dataframe.DataFrame.compute = counting_compute - try: - rng = np.random.default_rng(42) - n = 100 - df = pd.DataFrame( - { - "x": rng.uniform(0, 100, n), - "y": rng.uniform(0, 100, n), - "cell_type": pd.Categorical(rng.choice(["A", "B", "C"], n)), - } - ) - points = PointsModel.parse(df) - sdata = SpatialData(points={"pts": points}) + rng = np.random.default_rng(42) + n = 100 + df = pd.DataFrame( + { + "x": rng.uniform(0, 100, n), + "y": rng.uniform(0, 100, n), + "cell_type": pd.Categorical(rng.choice(["A", "B", "C"], n)), + } + ) + points = PointsModel.parse(df) + sdata = SpatialData(points={"pts": points}) + + with ( + dask.config.set({"dataframe.query-planning": False}), + pytest.MonkeyPatch().context() as mp, + ): + mp.setattr(dask.dataframe.DataFrame, "compute", counting_compute) sdata.pl.render_points("pts", color="cell_type").pl.show() plt.close("all") - finally: - dask.dataframe.DataFrame.compute = original_compute assert len(compute_calls) == 1, f"Expected 1 .compute() call, got {len(compute_calls)}" From c74f4c2c3507aa3bfd2c8f2dfb4a19082d56ed8a Mon Sep 17 00:00:00 2001 From: anon Date: Fri, 8 May 2026 12:13:34 +0200 Subject: [PATCH 3/4] simplify: use monkeypatch fixture instead of pytest.MonkeyPatch().context() Co-Authored-By: Claude Sonnet 4.6 --- tests/pl/test_render_points.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/tests/pl/test_render_points.py b/tests/pl/test_render_points.py index b5408185..fe57d102 100644 --- a/tests/pl/test_render_points.py +++ b/tests/pl/test_render_points.py @@ -1007,7 +1007,7 @@ def test_no_table_fallback_warning_for_element_column(caplog): plt.close("all") -def test_render_points_native_color_column_single_compute(): +def test_render_points_native_color_column_single_compute(monkeypatch): # Regression test for #633: color column native to the points element triggered # two .compute() calls instead of one. compute_calls = [] @@ -1029,11 +1029,8 @@ def counting_compute(self, **kwargs): points = PointsModel.parse(df) sdata = SpatialData(points={"pts": points}) - with ( - dask.config.set({"dataframe.query-planning": False}), - pytest.MonkeyPatch().context() as mp, - ): - mp.setattr(dask.dataframe.DataFrame, "compute", counting_compute) + with dask.config.set({"dataframe.query-planning": False}): + monkeypatch.setattr(dask.dataframe.DataFrame, "compute", counting_compute) sdata.pl.render_points("pts", color="cell_type").pl.show() plt.close("all") From 75434d6589aec1eedb09c7602010ac967fb216e4 Mon Sep 17 00:00:00 2001 From: anon Date: Fri, 8 May 2026 12:16:35 +0200 Subject: [PATCH 4/4] Remove compute-count regression test from test_render_points Co-Authored-By: Claude Sonnet 4.6 --- tests/pl/test_render_points.py | 31 ------------------------------- 1 file changed, 31 deletions(-) diff --git a/tests/pl/test_render_points.py b/tests/pl/test_render_points.py index fe57d102..bdb3218a 100644 --- a/tests/pl/test_render_points.py +++ b/tests/pl/test_render_points.py @@ -1,6 +1,5 @@ import math -import dask import dask.dataframe import datashader as ds import matplotlib @@ -1005,33 +1004,3 @@ def test_no_table_fallback_warning_for_element_column(caplog): with logger_no_warns(caplog, logger, match="fallback for color mapping"): sdata.pl.render_points("pts", color="cell_type").pl.show() plt.close("all") - - -def test_render_points_native_color_column_single_compute(monkeypatch): - # Regression test for #633: color column native to the points element triggered - # two .compute() calls instead of one. - compute_calls = [] - original_compute = dask.dataframe.DataFrame.compute - - def counting_compute(self, **kwargs): - compute_calls.append(True) - return original_compute(self, **kwargs) - - rng = np.random.default_rng(42) - n = 100 - df = pd.DataFrame( - { - "x": rng.uniform(0, 100, n), - "y": rng.uniform(0, 100, n), - "cell_type": pd.Categorical(rng.choice(["A", "B", "C"], n)), - } - ) - points = PointsModel.parse(df) - sdata = SpatialData(points={"pts": points}) - - with dask.config.set({"dataframe.query-planning": False}): - monkeypatch.setattr(dask.dataframe.DataFrame, "compute", counting_compute) - sdata.pl.render_points("pts", color="cell_type").pl.show() - plt.close("all") - - assert len(compute_calls) == 1, f"Expected 1 .compute() call, got {len(compute_calls)}"