diff --git a/pyranges1/core/pyranges_main.py b/pyranges1/core/pyranges_main.py index 70f9bcd..58f5441 100644 --- a/pyranges1/core/pyranges_main.py +++ b/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, diff --git a/tests/unit/test_nearest_ranges.py b/tests/unit/test_nearest_ranges.py index 3e94fcb..1b122cd 100644 --- a/tests/unit/test_nearest_ranges.py +++ b/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."""