Skip table fallback when coloring by element DataFrame column (#571)

Author: timtreis

src/spatialdata_plot/pl/utils.py

@@ -1101,9 +1101,15 @@ def _set_color_source_vec(
             processed = processed.reorder_categories(sorted(processed.categories))
         color_source_vector = processed  # convert, e.g., `pd.Series`
 
+        # When the value lives on the element's own DataFrame (origin="df"),
+        # there is no reason to look up a table for .uns colors.
+        value_from_element = origins[0].origin == "df"
+
         # Use the provided table_name parameter, fall back to only one present
         table_to_use: str | None
-        if table_name is not None and table_name in sdata.tables:
+        if value_from_element:
+            table_to_use = None
+        elif table_name is not None and table_name in sdata.tables:
             table_to_use = table_name
         elif table_name is not None and table_name not in sdata.tables:
             logger.warning(f"Table '{table_name}' not found in `sdata.tables`. Falling back to default behavior.")
@@ -1120,14 +1126,18 @@ def _set_color_source_vec(
 
         adata_for_mapping = sdata[table_to_use] if table_to_use is not None else None
 
-        # Check if custom colors exist in the table's .uns slot
-        if value_to_plot is not None and _has_colors_in_uns(sdata, table_name, value_to_plot):
+        # Check if custom colors exist in the resolved table's .uns slot
+        if (
+            value_to_plot is not None
+            and table_to_use is not None
+            and _has_colors_in_uns(sdata, table_to_use, value_to_plot)
+        ):
             # Extract colors directly from the table's .uns slot
             # Convert Color to ColorLike (str) for the function
             na_color_like: ColorLike = na_color.get_hex() if isinstance(na_color, Color) else na_color
             color_mapping = _extract_colors_from_table_uns(
                 sdata=sdata,
-                table_name=table_name,
+                table_name=table_to_use,
                 col_to_colorby=value_to_plot,
                 color_source_vector=color_source_vector,
                 na_color=na_color_like,

tests/pl/test_render_points.py

@@ -953,3 +953,34 @@ def test_datashader_na_color_none_no_nan_overlay_points(sdata_blobs: SpatialData
         "datashader is still rendering an opaque NaN overlay despite na_color=None"
     )
     plt.close(fig)
+
+
+def test_no_table_fallback_warning_for_element_column(caplog):
+    """No spurious 'using X as fallback for color mapping' warning when coloring by an element column (#562)."""
+    n = 50
+    rng = np.random.default_rng(42)
+    coords = pd.DataFrame({"x": rng.uniform(0, 100, n), "y": rng.uniform(0, 100, n)})
+    coords["cell_type"] = pd.Categorical(rng.choice(["A", "B", "C"], n))
+    points = PointsModel.parse(coords)
+
+    # Two tables so the multi-table fallback path is triggered
+    obs1 = pd.DataFrame(index=pd.RangeIndex(10).astype(str))
+    obs1["region"] = "s1"
+    obs1["instance_id"] = range(10)
+    table1 = AnnData(obs=obs1, X=np.zeros((10, 1)))
+    table1 = TableModel.parse(table1, region="s1", region_key="region", instance_key="instance_id")
+
+    obs2 = pd.DataFrame(index=pd.RangeIndex(10).astype(str))
+    obs2["region"] = "s2"
+    obs2["instance_id"] = range(10)
+    table2 = AnnData(obs=obs2, X=np.zeros((10, 1)))
+    table2 = TableModel.parse(table2, region="s2", region_key="region", instance_key="instance_id")
+
+    sdata = SpatialData(
+        points={"pts": points},
+        tables={"RNA": table1, "protein": table2},
+    )
+
+    with logger_no_warns(caplog, logger, match="fallback for color mapping"):
+        sdata.pl.render_points("pts", color="cell_type").pl.show()
+    plt.close("all")