Skip to content

Commit 54dd3db

Browse files
timtreisclaude
andcommitted
Eliminate DataFrame copies and double data loading in datashader path
Two additional performance fixes on top of the datashader speedups: 1. Replace .assign() + .rename() with direct column assignment when attaching the color column to the transformed element. Avoids two full DataFrame copies (~320MB saved for 10M points). 2. Add preloaded_color_data parameter to _set_color_source_vec so _render_points can pass already-loaded color data from get_values() instead of triggering a redundant second load from the table. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 06b4859 commit 54dd3db

2 files changed

Lines changed: 18 additions & 10 deletions

File tree

src/spatialdata_plot/pl/render.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -784,6 +784,10 @@ def _render_points(
784784
# from the registered points (see above) avoids duplicate-origin ambiguities.
785785
color_table_name = table_name
786786

787+
# When color was already loaded from a table (line 690), pass it directly
788+
# to avoid a redundant get_values() call inside _set_color_source_vec.
789+
_preloaded = points_pd_with_color[col_for_color] if added_color_from_table and col_for_color else None
790+
787791
color_source_vector, color_vector, _ = _set_color_source_vec(
788792
sdata=sdata_filt,
789793
element=color_element,
@@ -797,6 +801,7 @@ def _render_points(
797801
table_name=color_table_name,
798802
render_type="points",
799803
coordinate_system=coordinate_system,
804+
preloaded_color_data=_preloaded,
800805
)
801806

802807
if added_color_from_table and col_for_color is not None:
@@ -874,7 +879,7 @@ def _render_points(
874879
if isinstance(color_source_vector, pd.Series)
875880
else pd.Series(color_source_vector, index=series_index)
876881
)
877-
transformed_element = transformed_element.assign(col_for_color=source_series)
882+
transformed_element[col_for_color] = source_series
878883
else:
879884
if isinstance(color_vector, dd.Series):
880885
color_vector = color_vector.compute()
@@ -883,8 +888,7 @@ def _render_points(
883888
if isinstance(color_vector, pd.Series)
884889
else pd.Series(color_vector, index=series_index)
885890
)
886-
transformed_element = transformed_element.assign(col_for_color=color_series)
887-
transformed_element = transformed_element.rename(columns={"col_for_color": col_for_color})
891+
transformed_element[col_for_color] = color_series
888892

889893
color_dtype = transformed_element[col_for_color].dtype if col_for_color is not None else None
890894
color_by_categorical = col_for_color is not None and (

src/spatialdata_plot/pl/utils.py

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1019,6 +1019,7 @@ def _set_color_source_vec(
10191019
table_layer: str | None = None,
10201020
render_type: Literal["points", "labels"] | None = None,
10211021
coordinate_system: str | None = None,
1022+
preloaded_color_data: pd.Series | None = None,
10221023
) -> tuple[ArrayLike | pd.Series | None, ArrayLike, bool]:
10231024
if value_to_plot is None and element is not None:
10241025
color = np.full(len(element), na_color.get_hex_with_alpha())
@@ -1046,13 +1047,16 @@ def _set_color_source_vec(
10461047
element_name=element_name,
10471048
table_name=table_name,
10481049
)
1049-
color_source_vector = get_values(
1050-
value_key=value_to_plot,
1051-
sdata=sdata,
1052-
element_name=element_name,
1053-
table_name=table_name,
1054-
table_layer=table_layer,
1055-
)[value_to_plot]
1050+
if preloaded_color_data is not None:
1051+
color_source_vector = preloaded_color_data
1052+
else:
1053+
color_source_vector = get_values(
1054+
value_key=value_to_plot,
1055+
sdata=sdata,
1056+
element_name=element_name,
1057+
table_name=table_name,
1058+
table_layer=table_layer,
1059+
)[value_to_plot]
10561060

10571061
color_series = (
10581062
color_source_vector if isinstance(color_source_vector, pd.Series) else pd.Series(color_source_vector)

0 commit comments

Comments
 (0)