fix join non matching table (#813)

* fix join non matching table

* add test non matching element

* removed unused comment

---------

Co-authored-by: LucaMarconato <2664412+LucaMarconato@users.noreply.github.com>

Author: melonora

src/spatialdata/_core/query/relational_query.py

@@ -253,12 +253,12 @@ def _get_joined_table_indices(
     mask = table_instance_key_column.isin(element_indices)
     if joined_indices is None:
         if match_rows == "left":
-            joined_indices = _match_rows(table_instance_key_column, mask, element_indices, match_rows)
+            _, joined_indices = _match_rows(table_instance_key_column, mask, element_indices, match_rows)
         else:
             joined_indices = table_instance_key_column[mask].index
     else:
         if match_rows == "left":
-            add_indices = _match_rows(table_instance_key_column, mask, element_indices, match_rows)
+            _, add_indices = _match_rows(table_instance_key_column, mask, element_indices, match_rows)
             joined_indices = joined_indices.append(add_indices)
         # in place append does not work with pd.Index
         else:
@@ -294,8 +294,14 @@ def _get_masked_element(
     mask = table_instance_key_column.isin(element_indices)
     masked_table_instance_key_column = table_instance_key_column[mask]
     mask_values = mask_values if len(mask_values := masked_table_instance_key_column.values) != 0 else None
-    if match_rows == "right":
-        mask_values = _match_rows(table_instance_key_column, mask, element_indices, match_rows)
+    if match_rows in ["left", "right"]:
+        left_index, _ = _match_rows(table_instance_key_column, mask, element_indices, match_rows)
+
+        if mask_values is not None and len(left_index) != len(mask_values):
+            mask = left_index.isin(mask_values)
+            mask_values = left_index[mask]
+        else:
+            mask_values = left_index
 
     if isinstance(element, DaskDataFrame):
         return element.map_partitions(lambda df: df.loc[mask_values], meta=element)
@@ -404,6 +410,9 @@ def _inner_join_spatialelement_table(
                 element_dict[element_type][name] = None
                 continue
 
+    if joined_indices is not None:
+        joined_indices = joined_indices.dropna() if any(joined_indices.isna()) else joined_indices
+
     joined_table = table[joined_indices, :].copy() if joined_indices is not None else None
     _inplace_fix_subset_categorical_obs(subset_adata=joined_table, original_adata=table)
     return element_dict, joined_table
@@ -483,22 +492,24 @@ def _match_rows(
     mask: pd.Series,
     element_indices: pd.RangeIndex,
     match_rows: str,
-) -> pd.Index:
+) -> tuple[pd.Index, pd.Index]:
     instance_id_df = pd.DataFrame(
         {"instance_id": table_instance_key_column[mask].values, "index_right": table_instance_key_column[mask].index}
     )
     element_index_df = pd.DataFrame({"index_left": element_indices})
-    index_col = "index_left" if match_rows == "right" else "index_right"
 
-    merged_df = pd.merge(
-        element_index_df, instance_id_df, left_on="index_left", right_on="instance_id", how=match_rows
-    )[index_col]
+    merged_df = pd.merge(element_index_df, instance_id_df, left_on="index_left", right_on="instance_id", how=match_rows)
+    index_left = merged_df["index_left"]
+    index_right = merged_df["index_right"]
 
     # With labels it can be that index 0 is NaN
-    if isinstance(merged_df.iloc[0], float) and math.isnan(merged_df.iloc[0]):
-        merged_df = merged_df.iloc[1:]
+    if isinstance(index_left.iloc[0], float) and math.isnan(index_left.iloc[0]):
+        index_left = index_left.iloc[1:]
+
+    if isinstance(index_right.iloc[0], float) and math.isnan(index_right.iloc[0]):
+        index_right = index_right.iloc[1:]
 
-    return pd.Index(merged_df)
+    return pd.Index(index_left), pd.Index(index_right)
 
 
 class JoinTypes(Enum):

tests/core/query/test_relational_query.py

@@ -318,6 +318,64 @@ def test_left_exclusive_and_right_join(sdata_query_aggregation):
     )
 
 
+def test_match_rows_inner_join_non_matching_element(sdata_query_aggregation):
+    sdata = sdata_query_aggregation
+    sdata["values_circles"] = sdata["values_circles"][4:]
+    original_index = sdata["values_circles"].index
+    reversed_instance_id = [3, 5, 8, 7, 6, 4, 1, 2, 0] + list(reversed(range(12)))
+    sdata["table"].obs["instance_id"] = reversed_instance_id
+
+    element_dict, table = join_spatialelement_table(
+        sdata=sdata,
+        spatial_element_names="values_circles",
+        table_name="table",
+        how="inner",
+        match_rows="left",
+    )
+    assert all(table.obs["instance_id"].values == original_index)
+
+    element_dict, table = join_spatialelement_table(
+        sdata=sdata,
+        spatial_element_names="values_circles",
+        table_name="table",
+        how="inner",
+        match_rows="right",
+    )
+
+    assert all(element_dict["values_circles"].index == [5, 8, 7, 6, 4])
+
+
+def test_match_rows_inner_join_non_matching_table(sdata_query_aggregation):
+    sdata = sdata_query_aggregation
+    table = sdata["table"][3:]
+    original_instance_id = table.obs["instance_id"]
+    reversed_instance_id = [6, 7, 8, 3, 4, 5] + list(reversed(range(12)))
+    table.obs["instance_id"] = reversed_instance_id
+    sdata["table"] = table
+
+    element_dict, table = join_spatialelement_table(
+        sdata=sdata,
+        spatial_element_names=["values_circles", "values_polygons"],
+        table_name="table",
+        how="inner",
+        match_rows="left",
+    )
+
+    assert all(table.obs["instance_id"].values == original_instance_id.values)
+
+    element_dict, table = join_spatialelement_table(
+        sdata=sdata,
+        spatial_element_names=["values_circles", "values_polygons"],
+        table_name="table",
+        how="inner",
+        match_rows="right",
+    )
+
+    indices = element_dict["values_circles"].index.append(element_dict["values_polygons"].index)
+
+    assert all(indices == reversed_instance_id)
+
+
 # TODO: there is a lot of dublicate code, simplify with a function that tests both the case sdata=None and sdata=sdata
 def test_match_rows_join(sdata_query_aggregation):
     sdata = sdata_query_aggregation