nearest_ranges: honour preserve_input_order for direction="upstream"/"downstream"

pyranges1/core/pyranges_main.py

@@ -29,6 +29,7 @@
     START_COL,
     STRAND_BEHAVIOR_OPPOSITE,
     STRAND_COL,
+    TEMP_INDEX_COL,
     TEMP_TRANSCRIPT_ID_COL,
     USE_STRAND_DEFAULT,
     VALID_BY_TYPES,
@@ -2336,8 +2337,8 @@ def nearest_ranges(  # type: ignore[override]
           int64  |    category        int64    int64  category    str                 int64    int64  str              int64
         -------  ---  ------------  -------  -------  ----------  --------------  ---------  -------  ----------  ----------
               0  |    chr1                3        6  +           chr1                   20       22  +                   15
-              2  |    chr1                8        9  +           chr1                   20       22  +                   12
               1  |    chr1                5        7  -           chr1                    6        7  -                    0
+              2  |    chr1                8        9  +           chr1                   20       22  +                   12
         PyRanges with 3 rows, 9 columns, and 1 index columns.
         Contains 1 chromosomes and 2 strands.
 
@@ -2397,7 +2398,18 @@ def nearest_ranges(  # type: ignore[override]
         # Unpacked rather than zipped against VALID_GENOMIC_STRAND_INFO: pairing the
         # halves with their strand by position would silently mismap if either the
         # constant or split_on_strand's return order changed.
-        forward_self, reverse_self = split_on_strand(self)
+        # Each strand is searched and ordered on its own, so the plain concat below
+        # groups the result by strand (every forward row, then every reverse row) and
+        # drops the interleaving between strands; preserve_input_order would otherwise
+        # have no effect in this branch (issue #169). Carry each row's position in self
+        # through the split so the halves can be merged back into input order, matching
+        # the direction="any" branch. The index cannot stand in for the position, since
+        # k > 1 duplicates it and the caller's index may be non-unique.
+        ordered_self = self
+        if preserve_input_order:
+            ordered_self = self.copy()
+            ordered_self[TEMP_INDEX_COL] = np.arange(len(self))
+        forward_self, reverse_self = split_on_strand(ordered_self)
         per_strand = [
             RangeFrame(strand_self).nearest_ranges(
                 other=_other,
@@ -2416,7 +2428,10 @@ def nearest_ranges(  # type: ignore[override]
             )
         ]
 
-        return ensure_pyranges(pd.concat(per_strand))
+        combined = pd.concat(per_strand)
+        if preserve_input_order:
+            combined = combined.sort_values(TEMP_INDEX_COL, kind="stable").drop(columns=TEMP_INDEX_COL)
+        return ensure_pyranges(combined)
 
     def overlap(  # type: ignore[override]
         self,

tests/unit/test_nearest_ranges.py

@@ -196,6 +196,62 @@ def test_ties_first_applies_per_strand_for_a_directional_query() -> None:
     assert dict(zip(upstream["Strand"], upstream["Start_b"], strict=True)) == {"+": 50, "-": 200}
 
 
+def test_directional_nearest_preserves_input_order_across_strands() -> None:
+    """preserve_input_order must interleave the strand halves back into input order.
+
+    A directional query (upstream/downstream) splits self by strand, searches each
+    half on its own, and concatenates the results. That concat groups the rows by
+    strand -- every + row, then every - row -- so on an interleaved-strand frame the
+    output no longer matches the input and preserve_input_order had no effect at all
+    (issue #169). direction="any" never splits and always honoured the option, so it
+    is the oracle here.
+    """
+    query = pr.PyRanges(
+        {
+            "Chromosome": ["chr1"] * 4,
+            "Start": [100, 300, 500, 700],
+            "End": [120, 320, 520, 720],
+            "Strand": ["+", "-", "+", "-"],
+            "Id": ["a", "b", "c", "d"],
+        }
+    )
+    other = pr.PyRanges(
+        {
+            "Chromosome": ["chr1", "chr1"],
+            "Start": [0, 900],
+            "End": [10, 910],
+            "Strand": ["+", "+"],
+        }
+    )
+
+    # direction="any" is unaffected: it returns rows in input order (the oracle).
+    for preserve in (True, False):
+        result = query.nearest_ranges(
+            other, direction="any", strand_behavior="ignore", preserve_input_order=preserve
+        )
+        assert result["Id"].tolist() == ["a", "b", "c", "d"]
+        assert list(result.index) == [0, 1, 2, 3]
+
+    # Directional queries must honour it too: input order, original index.
+    for direction in ("upstream", "downstream"):
+        result = query.nearest_ranges(
+            other, direction=direction, strand_behavior="ignore", preserve_input_order=True
+        )
+        assert result["Id"].tolist() == ["a", "b", "c", "d"]
+        assert list(result.index) == [0, 1, 2, 3]
+
+    # preserve_input_order=False is unchanged: rows stay grouped by strand, every
+    # forward row before every reverse row, never the interleaved input order.
+    for direction in ("upstream", "downstream"):
+        grouped = query.nearest_ranges(
+            other, direction=direction, strand_behavior="ignore", preserve_input_order=False
+        )
+        ids = grouped["Id"].tolist()
+        assert set(ids) == {"a", "b", "c", "d"}
+        assert set(ids[:2]) == {"a", "c"}
+        assert set(ids[2:]) == {"b", "d"}
+
+
 def test_nearest_ranges_rejects_an_unknown_ties() -> None:
     """An unknown value must raise, not reach the kernel: a Rust panic is not
     an Exception, so `except Exception` cannot catch it."""