Fix exception in compute_mapped for partially reducing pipelines

src/sciline/pipeline.py

@@ -248,12 +248,9 @@ def get_mapped_node_names(
         raise ValueError(
             f"Multiple mapped nodes with name '{base_name}' found: {candidates}"
         )
-    # Drops unrelated indices
-    graph = graph[candidates[0]]  # type: ignore[index]
-    indices = graph.indices
-    if index_names is not None:
-        indices = {name: indices[name] for name in indices if name in index_names}
-    index_names = tuple(indices)
+
+    index_names = tuple(reversed(candidates[0].indices))
+    indices = {name: idx for name, idx in graph.indices.items() if name in index_names}
 
     index = pd.MultiIndex.from_product(indices.values(), names=index_names)
     keys = tuple(NodeName(base_name, IndexValues(index_names, idx)) for idx in index)

tests/pipeline_map_reduce_test.py

@@ -202,7 +202,26 @@ def ab_to_c(a: A, b: B) -> C:
         sl.compute_mapped(pl, C)
 
 
-def test_compute_mapped_indices_selects_between_multiple_candidates() -> None:
+def test_compute_mapped_with_partial_reduction_identifies_correct_index() -> None:
+    def ab_to_c(a: A, b: B) -> C:
+        return C(a + b)
+
+    D = NewType('D', int)
+
+    pl = sl.Pipeline((ab_to_c,))
+    paramsA = pd.DataFrame(
+        {A: [A(10 * i) for i in range(3)]}, index=['a', 'b', 'c']
+    ).rename_axis('x')
+    paramsB = pd.DataFrame(
+        {B: [B(i) for i in range(2)]}, index=['aa', 'bb']
+    ).rename_axis('y')
+    pl = pl.map(paramsA).map(paramsB).reduce(func=max, name=D, index='x')
+    result = sl.compute_mapped(pl, D)
+    assert result['aa'] == C(20)
+    assert result['bb'] == C(21)
+
+
+def test_compute_mapped_index_names_selects_between_multiple_candidates() -> None:
     def ab_to_c(a: A, b: B) -> C:
         return C(a + b)