1use std::collections::HashSet;
2
3use polars_core::prelude::{AnyValue, Column, DataFrame, DataType, NamedFrom, Series};
4use ruranges_core::{
5 boundary, cluster, complement_single, extend, group_cumsum, max_disjoint, merge, nearest,
6 outside_bounds, overlaps, sorts, split, subtract, tile,
7};
8use rustc_hash::FxHashMap;
9
10use crate::error::{RangeFrameError, Result};
11use crate::factorize::{factorize, factorize_pair, factorize_pair_by};
12use crate::range_frame::{extract_coordinates, take_rows, RangeView};
13
14#[derive(Clone, Debug)]
15pub struct MergeOptions {
16 pub count_column: Option<String>,
17 pub match_by: Vec<String>,
18 pub slack: i64,
19}
20
21impl Default for MergeOptions {
22 fn default() -> Self {
23 Self {
24 count_column: None,
25 match_by: Vec::new(),
26 slack: 0,
27 }
28 }
29}
30
31impl MergeOptions {
32 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
33 where
34 I: IntoIterator<Item = S>,
35 S: Into<String>,
36 {
37 self.match_by = columns.into_iter().map(Into::into).collect();
38 self
39 }
40}
41
42#[derive(Clone, Debug)]
43pub struct ClusterOptions {
44 pub match_by: Vec<String>,
45 pub cluster_column: String,
46 pub slack: i64,
47}
48
49impl Default for ClusterOptions {
50 fn default() -> Self {
51 Self {
52 match_by: Vec::new(),
53 cluster_column: "Cluster".to_owned(),
54 slack: 0,
55 }
56 }
57}
58
59impl ClusterOptions {
60 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
61 where
62 I: IntoIterator<Item = S>,
63 S: Into<String>,
64 {
65 self.match_by = columns.into_iter().map(Into::into).collect();
66 self
67 }
68}
69
70#[derive(Clone, Debug)]
71pub struct CountOverlapsOptions {
72 pub match_by: Vec<String>,
73 pub slack: i64,
74}
75
76impl Default for CountOverlapsOptions {
77 fn default() -> Self {
78 Self {
79 match_by: Vec::new(),
80 slack: 0,
81 }
82 }
83}
84
85impl CountOverlapsOptions {
86 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
87 where
88 I: IntoIterator<Item = S>,
89 S: Into<String>,
90 {
91 self.match_by = columns.into_iter().map(Into::into).collect();
92 self
93 }
94}
95
96#[derive(Clone, Copy, Debug, Eq, PartialEq)]
97pub enum JoinType {
98 Inner,
99 Left,
100 Right,
101 Outer,
102}
103
104#[derive(Clone, Debug)]
105pub struct JoinOverlapsOptions {
106 pub match_by: Vec<String>,
107 pub multiple: crate::OverlapMode,
108 pub slack: i64,
109 pub suffix: String,
110 pub contained_intervals_only: bool,
111 pub join_type: JoinType,
112 pub report_overlap_column: Option<String>,
113 pub preserve_input_order: bool,
114}
115
116impl Default for JoinOverlapsOptions {
117 fn default() -> Self {
118 Self {
119 match_by: Vec::new(),
120 multiple: crate::OverlapMode::All,
121 slack: 0,
122 suffix: "_b".to_owned(),
123 contained_intervals_only: false,
124 join_type: JoinType::Inner,
125 report_overlap_column: None,
126 preserve_input_order: true,
127 }
128 }
129}
130
131impl JoinOverlapsOptions {
132 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
133 where
134 I: IntoIterator<Item = S>,
135 S: Into<String>,
136 {
137 self.match_by = columns.into_iter().map(Into::into).collect();
138 self
139 }
140}
141
142#[derive(Clone, Debug)]
143pub struct IntersectOverlapsOptions {
144 pub match_by: Vec<String>,
145 pub multiple: crate::OverlapMode,
146 pub slack: i64,
147 pub contained_intervals_only: bool,
148 pub preserve_input_order: bool,
149}
150
151impl Default for IntersectOverlapsOptions {
152 fn default() -> Self {
153 Self {
154 match_by: Vec::new(),
155 multiple: crate::OverlapMode::All,
156 slack: 0,
157 contained_intervals_only: false,
158 preserve_input_order: true,
159 }
160 }
161}
162
163impl IntersectOverlapsOptions {
164 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
165 where
166 I: IntoIterator<Item = S>,
167 S: Into<String>,
168 {
169 self.match_by = columns.into_iter().map(Into::into).collect();
170 self
171 }
172}
173
174#[derive(Clone, Debug)]
175pub struct SetIntersectOverlapsOptions {
176 pub match_by: Vec<String>,
177 pub multiple: crate::OverlapMode,
178 pub preserve_input_order: bool,
179}
180
181impl Default for SetIntersectOverlapsOptions {
182 fn default() -> Self {
183 Self {
184 match_by: Vec::new(),
185 multiple: crate::OverlapMode::All,
186 preserve_input_order: true,
187 }
188 }
189}
190
191impl SetIntersectOverlapsOptions {
192 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
193 where
194 I: IntoIterator<Item = S>,
195 S: Into<String>,
196 {
197 self.match_by = columns.into_iter().map(Into::into).collect();
198 self
199 }
200}
201
202#[derive(Clone, Debug, Default)]
203pub struct SetUnionOverlapsOptions {
204 pub match_by: Vec<String>,
205}
206
207impl SetUnionOverlapsOptions {
208 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
209 where
210 I: IntoIterator<Item = S>,
211 S: Into<String>,
212 {
213 self.match_by = columns.into_iter().map(Into::into).collect();
214 self
215 }
216}
217
218#[derive(Clone, Copy, Debug, Eq, PartialEq)]
219pub enum NearestDirection {
220 Any,
221 Forward,
222 Backward,
223}
224
225impl NearestDirection {
226 fn as_kernel_str(self) -> &'static str {
227 match self {
228 Self::Any => "any",
229 Self::Forward => "forward",
230 Self::Backward => "backward",
231 }
232 }
233}
234
235#[derive(Clone, Debug)]
236pub struct NearestOptions {
237 pub direction: NearestDirection,
238 pub k: usize,
239 pub exclude_overlaps: bool,
240 pub preserve_input_order: bool,
241 pub suffix: String,
242 pub distance_column: Option<String>,
243 pub match_by: Vec<String>,
244 pub left_start_col: String,
245 pub left_end_col: String,
246 pub right_start_col: String,
247 pub right_end_col: String,
248}
249
250impl Default for NearestOptions {
251 fn default() -> Self {
252 Self {
253 direction: NearestDirection::Any,
254 k: 1,
255 exclude_overlaps: false,
256 preserve_input_order: true,
257 suffix: "_b".to_owned(),
258 distance_column: Some("Distance".to_owned()),
259 match_by: Vec::new(),
260 left_start_col: "Start".to_owned(),
261 left_end_col: "End".to_owned(),
262 right_start_col: "Start".to_owned(),
263 right_end_col: "End".to_owned(),
264 }
265 }
266}
267
268impl NearestOptions {
269 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
270 where
271 I: IntoIterator<Item = S>,
272 S: Into<String>,
273 {
274 self.match_by = columns.into_iter().map(Into::into).collect();
275 self
276 }
277
278 pub fn with_interval_columns(
279 mut self,
280 start_col: impl Into<String>,
281 end_col: impl Into<String>,
282 ) -> Self {
283 let start = start_col.into();
284 let end = end_col.into();
285 self.left_start_col = start.clone();
286 self.left_end_col = end.clone();
287 self.right_start_col = start;
288 self.right_end_col = end;
289 self
290 }
291
292 pub fn with_left_interval_columns(
293 mut self,
294 start_col: impl Into<String>,
295 end_col: impl Into<String>,
296 ) -> Self {
297 self.left_start_col = start_col.into();
298 self.left_end_col = end_col.into();
299 self
300 }
301
302 pub fn with_right_interval_columns(
303 mut self,
304 start_col: impl Into<String>,
305 end_col: impl Into<String>,
306 ) -> Self {
307 self.right_start_col = start_col.into();
308 self.right_end_col = end_col.into();
309 self
310 }
311}
312
313#[derive(Clone, Debug)]
314pub struct SubtractOptions {
315 pub match_by: Vec<String>,
316 pub preserve_input_order: bool,
317}
318
319#[derive(Clone, Debug)]
320pub struct SortRangesOptions {
321 pub match_by: Vec<String>,
322 pub negative_strand_column: Option<String>,
323}
324
325impl Default for SortRangesOptions {
326 fn default() -> Self {
327 Self {
328 match_by: Vec::new(),
329 negative_strand_column: None,
330 }
331 }
332}
333
334impl SortRangesOptions {
335 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
336 where
337 I: IntoIterator<Item = S>,
338 S: Into<String>,
339 {
340 self.match_by = columns.into_iter().map(Into::into).collect();
341 self
342 }
343}
344
345#[derive(Clone, Debug, Default)]
346pub struct ExtendRangesOptions {
347 pub match_by: Vec<String>,
348 pub ext: Option<i64>,
349 pub ext_3: Option<i64>,
350 pub ext_5: Option<i64>,
351 pub negative_strand_column: Option<String>,
352}
353
354impl ExtendRangesOptions {
355 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
356 where
357 I: IntoIterator<Item = S>,
358 S: Into<String>,
359 {
360 self.match_by = columns.into_iter().map(Into::into).collect();
361 self
362 }
363}
364
365#[derive(Clone, Debug)]
366pub struct TileRangesOptions {
367 pub match_by: Vec<String>,
368 pub tile_size: i64,
369 pub negative_strand_column: Option<String>,
370 pub overlap_column: Option<String>,
371}
372
373impl TileRangesOptions {
374 pub fn new(tile_size: i64) -> Self {
375 Self {
376 match_by: Vec::new(),
377 tile_size,
378 negative_strand_column: None,
379 overlap_column: None,
380 }
381 }
382
383 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
384 where
385 I: IntoIterator<Item = S>,
386 S: Into<String>,
387 {
388 self.match_by = columns.into_iter().map(Into::into).collect();
389 self
390 }
391}
392
393#[derive(Clone, Debug)]
394pub struct ClipRangesOptions {
395 pub chromsizes: Option<FxHashMap<String, i64>>,
396 pub remove: bool,
397 pub only_right: bool,
398 pub group_sizes_col: String,
399}
400
401impl Default for ClipRangesOptions {
402 fn default() -> Self {
403 Self {
404 chromsizes: None,
405 remove: false,
406 only_right: false,
407 group_sizes_col: "Chromosome".to_owned(),
408 }
409 }
410}
411
412#[derive(Clone, Debug)]
413pub struct GroupCumsumOptions {
414 pub match_by: Vec<String>,
415 pub forward_strand_column: Option<String>,
416 pub cumsum_start_column: Option<String>,
417 pub cumsum_end_column: Option<String>,
418 pub keep_order: bool,
419}
420
421impl Default for GroupCumsumOptions {
422 fn default() -> Self {
423 Self {
424 match_by: Vec::new(),
425 forward_strand_column: None,
426 cumsum_start_column: None,
427 cumsum_end_column: None,
428 keep_order: true,
429 }
430 }
431}
432
433impl GroupCumsumOptions {
434 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
435 where
436 I: IntoIterator<Item = S>,
437 S: Into<String>,
438 {
439 self.match_by = columns.into_iter().map(Into::into).collect();
440 self
441 }
442}
443
444#[derive(Clone, Debug)]
445pub struct MaxDisjointOptions {
446 pub match_by: Vec<String>,
447 pub slack: i64,
448 pub preserve_input_order: bool,
449}
450
451impl Default for MaxDisjointOptions {
452 fn default() -> Self {
453 Self {
454 match_by: Vec::new(),
455 slack: 0,
456 preserve_input_order: true,
457 }
458 }
459}
460
461impl MaxDisjointOptions {
462 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
463 where
464 I: IntoIterator<Item = S>,
465 S: Into<String>,
466 {
467 self.match_by = columns.into_iter().map(Into::into).collect();
468 self
469 }
470}
471
472#[derive(Clone, Debug)]
473pub struct SplitOverlapsOptions {
474 pub match_by: Vec<String>,
475 pub between: bool,
476}
477
478impl Default for SplitOverlapsOptions {
479 fn default() -> Self {
480 Self {
481 match_by: Vec::new(),
482 between: false,
483 }
484 }
485}
486
487impl SplitOverlapsOptions {
488 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
489 where
490 I: IntoIterator<Item = S>,
491 S: Into<String>,
492 {
493 self.match_by = columns.into_iter().map(Into::into).collect();
494 self
495 }
496}
497
498#[derive(Clone, Debug, Default)]
499pub struct OuterRangesOptions {
500 pub match_by: Vec<String>,
501}
502
503impl OuterRangesOptions {
504 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
505 where
506 I: IntoIterator<Item = S>,
507 S: Into<String>,
508 {
509 self.match_by = columns.into_iter().map(Into::into).collect();
510 self
511 }
512}
513
514#[derive(Clone, Debug)]
515pub struct ComplementRangesOptions {
516 pub match_by: Vec<String>,
517 pub include_first_interval: bool,
518 pub chromsizes: Option<FxHashMap<String, i64>>,
519 pub group_sizes_col: String,
520}
521
522impl Default for ComplementRangesOptions {
523 fn default() -> Self {
524 Self {
525 match_by: Vec::new(),
526 include_first_interval: false,
527 chromsizes: None,
528 group_sizes_col: "Chromosome".to_owned(),
529 }
530 }
531}
532
533impl ComplementRangesOptions {
534 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
535 where
536 I: IntoIterator<Item = S>,
537 S: Into<String>,
538 {
539 self.match_by = columns.into_iter().map(Into::into).collect();
540 self
541 }
542}
543
544impl Default for SubtractOptions {
545 fn default() -> Self {
546 Self {
547 match_by: Vec::new(),
548 preserve_input_order: true,
549 }
550 }
551}
552
553impl SubtractOptions {
554 pub fn with_match_by<I, S>(mut self, columns: I) -> Self
555 where
556 I: IntoIterator<Item = S>,
557 S: Into<String>,
558 {
559 self.match_by = columns.into_iter().map(Into::into).collect();
560 self
561 }
562}
563
564impl<'a> RangeView<'a> {
565 pub(crate) fn merge_overlaps(&self, options: &MergeOptions) -> Result<DataFrame> {
566 let groups = factorize(self.df, &options.match_by)?;
567 let coords = extract_coordinates(self.df, self.start_col, self.end_col)?;
568
569 if let Some((starts, ends)) = coords.as_i32() {
570 if let Ok(slack) = i32::try_from(options.slack) {
571 let (idx, merged_starts, merged_ends, counts) =
572 merge::sweep_line_merge(&groups, starts, ends, slack);
573 return build_merge_result_i32(
574 self,
575 &idx,
576 merged_starts,
577 merged_ends,
578 counts,
579 options,
580 );
581 }
582 }
583
584 if let Some((starts, ends)) = coords.as_i64() {
585 let (idx, merged_starts, merged_ends, counts) =
586 merge::sweep_line_merge(&groups, starts, ends, options.slack);
587 return build_merge_result_i64(self, &idx, merged_starts, merged_ends, counts, options);
588 }
589
590 let (starts, ends) = coords.into_i64_buffers();
591 let (idx, merged_starts, merged_ends, counts) =
592 merge::sweep_line_merge(&groups, &starts, &ends, options.slack);
593 build_merge_result_i64(self, &idx, merged_starts, merged_ends, counts, options)
594 }
595
596 pub(crate) fn cluster_overlaps(&self, options: &ClusterOptions) -> Result<DataFrame> {
597 let groups = factorize(self.df, &options.match_by)?;
598 let coords = extract_coordinates(self.df, self.start_col, self.end_col)?;
599
600 if let Some((starts, ends)) = coords.as_i32() {
601 if let Ok(slack) = i32::try_from(options.slack) {
602 let (cluster_ids, idx) = cluster::sweep_line_cluster(&groups, starts, ends, slack);
603 return build_cluster_result(self.df, &idx, cluster_ids, &options.cluster_column);
604 }
605 }
606
607 if let Some((starts, ends)) = coords.as_i64() {
608 let (cluster_ids, idx) =
609 cluster::sweep_line_cluster(&groups, starts, ends, options.slack);
610 return build_cluster_result(self.df, &idx, cluster_ids, &options.cluster_column);
611 }
612
613 let (starts, ends) = coords.into_i64_buffers();
614 let (cluster_ids, idx) =
615 cluster::sweep_line_cluster(&groups, &starts, &ends, options.slack);
616 build_cluster_result(self.df, &idx, cluster_ids, &options.cluster_column)
617 }
618
619 pub(crate) fn count_overlaps(
620 &self,
621 other: &RangeView<'_>,
622 options: &CountOverlapsOptions,
623 ) -> Result<Series> {
624 let (left_groups, right_groups) = factorize_pair(self.df, other.df, &options.match_by)?;
625 let left_coords = extract_coordinates(self.df, self.start_col, self.end_col)?;
626 let right_coords = extract_coordinates(other.df, other.start_col, other.end_col)?;
627
628 if let (Some((left_starts, left_ends)), Some((right_starts, right_ends))) =
629 (left_coords.as_i32(), right_coords.as_i32())
630 {
631 if let Ok(slack) = i32::try_from(options.slack) {
632 let counts = overlaps::count_overlaps(
633 &left_groups,
634 left_starts,
635 left_ends,
636 &right_groups,
637 right_starts,
638 right_ends,
639 slack,
640 );
641 return Ok(Series::new("Count".into(), counts));
642 }
643 }
644
645 if let (Some((left_starts, left_ends)), Some((right_starts, right_ends))) =
646 (left_coords.as_i64(), right_coords.as_i64())
647 {
648 let counts = overlaps::count_overlaps(
649 &left_groups,
650 left_starts,
651 left_ends,
652 &right_groups,
653 right_starts,
654 right_ends,
655 options.slack,
656 );
657 return Ok(Series::new("Count".into(), counts));
658 }
659
660 let (left_starts, left_ends) = left_coords.into_i64_buffers();
661 let (right_starts, right_ends) = right_coords.into_i64_buffers();
662 let counts = overlaps::count_overlaps(
663 &left_groups,
664 &left_starts,
665 &left_ends,
666 &right_groups,
667 &right_starts,
668 &right_ends,
669 options.slack,
670 );
671 Ok(Series::new("Count".into(), counts))
672 }
673
674 pub(crate) fn join_overlaps(
675 &self,
676 other: &RangeView<'_>,
677 options: &JoinOverlapsOptions,
678 ) -> Result<DataFrame> {
679 let overlap_options = crate::OverlapOptions {
680 multiple: options.multiple,
681 slack: options.slack,
682 contained_intervals_only: options.contained_intervals_only,
683 match_by: options.match_by.clone(),
684 preserve_input_order: options.preserve_input_order,
685 ..crate::OverlapOptions::default()
686 };
687 let pairs = self.overlap_pairs(other, &overlap_options)?;
688 build_join_result(self.df, other.df, &pairs.left, &pairs.right, options)
689 }
690
691 pub(crate) fn intersect_overlaps(
692 &self,
693 other: &RangeView<'_>,
694 options: &IntersectOverlapsOptions,
695 ) -> Result<DataFrame> {
696 let overlap_options = crate::OverlapOptions {
697 multiple: options.multiple,
698 slack: options.slack,
699 contained_intervals_only: options.contained_intervals_only,
700 match_by: options.match_by.clone(),
701 preserve_input_order: options.preserve_input_order,
702 ..crate::OverlapOptions::default()
703 };
704 let pairs = self.overlap_pairs(other, &overlap_options)?;
705 build_intersect_result(self, other, &pairs.left, &pairs.right)
706 }
707
708 pub(crate) fn set_intersect_overlaps(
709 &self,
710 other: &RangeView<'_>,
711 options: &SetIntersectOverlapsOptions,
712 ) -> Result<DataFrame> {
713 let left_merged = self.merge_overlaps(&MergeOptions {
714 count_column: None,
715 match_by: options.match_by.clone(),
716 slack: 0,
717 })?;
718 let right_merged = other.merge_overlaps(&MergeOptions {
719 count_column: None,
720 match_by: options.match_by.clone(),
721 slack: 0,
722 })?;
723 let left = RangeView::with_columns(&left_merged, self.start_col, self.end_col)?;
724 let right = RangeView::with_columns(&right_merged, other.start_col, other.end_col)?;
725 left.intersect_overlaps(
726 &right,
727 &IntersectOverlapsOptions {
728 match_by: options.match_by.clone(),
729 multiple: options.multiple,
730 slack: 0,
731 contained_intervals_only: false,
732 preserve_input_order: options.preserve_input_order,
733 },
734 )
735 }
736
737 pub(crate) fn set_union_overlaps(
738 &self,
739 other: &RangeView<'_>,
740 options: &SetUnionOverlapsOptions,
741 ) -> Result<DataFrame> {
742 let _ = factorize_pair(self.df, other.df, &options.match_by)?;
743 let left_idx = all_indices(self.df.height())?;
744 let right_idx = all_indices(other.df.height())?;
745 let mut combined = take_merge_projection(
746 self.df,
747 &left_idx,
748 self.start_col,
749 self.end_col,
750 &options.match_by,
751 )?;
752 let right = take_merge_projection(
753 other.df,
754 &right_idx,
755 other.start_col,
756 other.end_col,
757 &options.match_by,
758 )?;
759 combined.vstack_mut(&right)?;
760 let combined = RangeView::with_columns(&combined, self.start_col, self.end_col)?;
761 combined.merge_overlaps(&MergeOptions {
762 count_column: None,
763 match_by: options.match_by.clone(),
764 slack: 0,
765 })
766 }
767
768 pub(crate) fn sort_ranges(&self, options: &SortRangesOptions) -> Result<DataFrame> {
769 let groups = factorize(self.df, &options.match_by)?;
770 let reverse = match options.negative_strand_column.as_deref() {
771 Some(column) => Some(negative_strand_flags(self.df, column)?),
772 None => None,
773 };
774 let coords = extract_coordinates(self.df, self.start_col, self.end_col)?;
775
776 if let Some((starts, ends)) = coords.as_i32() {
777 let idx = sorts::sort_order_idx(&groups, starts, ends, reverse.as_deref());
778 return take_rows(self.df, &idx);
779 }
780
781 if let Some((starts, ends)) = coords.as_i64() {
782 let idx = sorts::sort_order_idx(&groups, starts, ends, reverse.as_deref());
783 return take_rows(self.df, &idx);
784 }
785
786 let (starts, ends) = coords.into_i64_buffers();
787 let idx = sorts::sort_order_idx(&groups, &starts, &ends, reverse.as_deref());
788 take_rows(self.df, &idx)
789 }
790
791 pub(crate) fn extend_ranges(&self, options: &ExtendRangesOptions) -> Result<DataFrame> {
792 let (ext_3, ext_5) = resolve_extend_options(options)?;
793 let groups = if options.match_by.is_empty() {
794 all_indices(self.df.height())?
795 } else {
796 factorize(self.df, &options.match_by)?
797 };
798 let negative_strand = match options.negative_strand_column.as_deref() {
799 Some(column) => negative_strand_flags(self.df, column)?,
800 None => vec![false; self.df.height()],
801 };
802 let coords = extract_coordinates(self.df, self.start_col, self.end_col)?;
803
804 let mut out = self.df.clone();
805 if let Some((starts, ends)) = coords.as_i32() {
806 if let (Ok(ext_3), Ok(ext_5)) = (i32::try_from(ext_3), i32::try_from(ext_5)) {
807 let (starts, ends) =
808 extend::extend_grp(&groups, starts, ends, &negative_strand, ext_3, ext_5);
809 replace_interval_columns_i32(&mut out, self.start_col, self.end_col, starts, ends)?;
810 return Ok(out);
811 }
812 }
813
814 if let Some((starts, ends)) = coords.as_i64() {
815 let (starts, ends) =
816 extend::extend_grp(&groups, starts, ends, &negative_strand, ext_3, ext_5);
817 replace_interval_columns_i64(&mut out, self.start_col, self.end_col, starts, ends)?;
818 return Ok(out);
819 }
820
821 let (starts, ends) = coords.into_i64_buffers();
822 let (starts, ends) =
823 extend::extend_grp(&groups, &starts, &ends, &negative_strand, ext_3, ext_5);
824 replace_interval_columns_i64(&mut out, self.start_col, self.end_col, starts, ends)?;
825 Ok(out)
826 }
827
828 pub(crate) fn tile_ranges(&self, options: &TileRangesOptions) -> Result<DataFrame> {
829 if options.tile_size <= 0 {
830 return Err(compute_error("tile_size must be greater than 0"));
831 }
832
833 let _ = factorize(self.df, &options.match_by)?;
834 let negative_strand = match options.negative_strand_column.as_deref() {
835 Some(column) => negative_strand_flags(self.df, column)?,
836 None => vec![false; self.df.height()],
837 };
838 let coords = extract_coordinates(self.df, self.start_col, self.end_col)?;
839
840 if let Some((starts, ends)) = coords.as_i32() {
841 if let Ok(tile_size) = i32::try_from(options.tile_size) {
842 let (starts, ends, idx, overlaps) =
843 tile::tile(starts, ends, &negative_strand, tile_size);
844 return build_tile_result_i32(self, &idx, starts, ends, overlaps, options);
845 }
846 }
847
848 if let Some((starts, ends)) = coords.as_i64() {
849 let (starts, ends, idx, overlaps) =
850 tile::tile(starts, ends, &negative_strand, options.tile_size);
851 return build_tile_result_i64(self, &idx, starts, ends, overlaps, options);
852 }
853
854 let (starts, ends) = coords.into_i64_buffers();
855 let (starts, ends, idx, overlaps) =
856 tile::tile(&starts, &ends, &negative_strand, options.tile_size);
857 build_tile_result_i64(self, &idx, starts, ends, overlaps, options)
858 }
859
860 pub(crate) fn clip_ranges(&self, options: &ClipRangesOptions) -> Result<DataFrame> {
861 let groups = factorize(self.df, std::slice::from_ref(&options.group_sizes_col))?;
862 let coords = extract_coordinates(self.df, self.start_col, self.end_col)?;
863
864 if let Some((starts, ends)) = coords.as_i32() {
865 let chrom_lens =
866 row_aligned_bounds_i32(self.df, &options.group_sizes_col, ends, options)?;
867 let (idx, starts, ends) = outside_bounds::outside_bounds(
868 &groups,
869 starts,
870 ends,
871 &chrom_lens,
872 !options.remove,
873 options.only_right,
874 )
875 .map_err(compute_error)?;
876 return build_clip_result_i32(self, &idx, starts, ends, options);
877 }
878
879 if let Some((starts, ends)) = coords.as_i64() {
880 let chrom_lens =
881 row_aligned_bounds_i64(self.df, &options.group_sizes_col, ends, options)?;
882 let (idx, starts, ends) = outside_bounds::outside_bounds(
883 &groups,
884 starts,
885 ends,
886 &chrom_lens,
887 !options.remove,
888 options.only_right,
889 )
890 .map_err(compute_error)?;
891 return build_clip_result_i64(self, &idx, starts, ends, options);
892 }
893
894 let (starts, ends) = coords.into_i64_buffers();
895 let chrom_lens = row_aligned_bounds_i64(self.df, &options.group_sizes_col, &ends, options)?;
896 let (idx, starts, ends) = outside_bounds::outside_bounds(
897 &groups,
898 &starts,
899 &ends,
900 &chrom_lens,
901 !options.remove,
902 options.only_right,
903 )
904 .map_err(compute_error)?;
905 build_clip_result_i64(self, &idx, starts, ends, options)
906 }
907
908 pub(crate) fn group_cumsum(&self, options: &GroupCumsumOptions) -> Result<DataFrame> {
909 let groups = factorize(self.df, &options.match_by)?;
910 let forward_strand = match options.forward_strand_column.as_deref() {
911 Some(column) => forward_strand_flags(self.df, column)?,
912 None => vec![true; self.df.height()],
913 };
914 let coords = extract_coordinates(self.df, self.start_col, self.end_col)?;
915
916 if let Some((starts, ends)) = coords.as_i32() {
917 let (idx, cumsum_starts, cumsum_ends) = group_cumsum::sweep_line_cumsum(
918 &groups,
919 starts,
920 ends,
921 &forward_strand,
922 options.keep_order,
923 );
924 return build_group_cumsum_result_i32(self, &idx, cumsum_starts, cumsum_ends, options);
925 }
926
927 if let Some((starts, ends)) = coords.as_i64() {
928 let (idx, cumsum_starts, cumsum_ends) = group_cumsum::sweep_line_cumsum(
929 &groups,
930 starts,
931 ends,
932 &forward_strand,
933 options.keep_order,
934 );
935 return build_group_cumsum_result_i64(self, &idx, cumsum_starts, cumsum_ends, options);
936 }
937
938 let (starts, ends) = coords.into_i64_buffers();
939 let (idx, cumsum_starts, cumsum_ends) = group_cumsum::sweep_line_cumsum(
940 &groups,
941 &starts,
942 &ends,
943 &forward_strand,
944 options.keep_order,
945 );
946 build_group_cumsum_result_i64(self, &idx, cumsum_starts, cumsum_ends, options)
947 }
948
949 pub(crate) fn max_disjoint_overlaps(&self, options: &MaxDisjointOptions) -> Result<DataFrame> {
950 let groups = factorize(self.df, &options.match_by)?;
951 let coords = extract_coordinates(self.df, self.start_col, self.end_col)?;
952
953 if let Some((starts, ends)) = coords.as_i32() {
954 if let Ok(slack) = i32::try_from(options.slack) {
955 let idx = max_disjoint::max_disjoint(
956 &groups,
957 starts,
958 ends,
959 slack,
960 options.preserve_input_order,
961 );
962 return take_rows(self.df, &idx);
963 }
964 }
965
966 if let Some((starts, ends)) = coords.as_i64() {
967 let idx = max_disjoint::max_disjoint(
968 &groups,
969 starts,
970 ends,
971 options.slack,
972 options.preserve_input_order,
973 );
974 return take_rows(self.df, &idx);
975 }
976
977 let (starts, ends) = coords.into_i64_buffers();
978 let idx = max_disjoint::max_disjoint(
979 &groups,
980 &starts,
981 &ends,
982 options.slack,
983 options.preserve_input_order,
984 );
985 take_rows(self.df, &idx)
986 }
987
988 pub(crate) fn split_overlaps(&self, options: &SplitOverlapsOptions) -> Result<DataFrame> {
989 let groups = factorize(self.df, &options.match_by)?;
990 let coords = extract_coordinates(self.df, self.start_col, self.end_col)?;
991
992 if let Some((starts, ends)) = coords.as_i32() {
993 let (idx, starts, ends) =
994 split::sweep_line_split(&groups, starts, ends, 0_i32, options.between);
995 return build_projected_interval_result_i32(
996 self,
997 &idx,
998 starts,
999 ends,
1000 &options.match_by,
1001 );
1002 }
1003
1004 if let Some((starts, ends)) = coords.as_i64() {
1005 let (idx, starts, ends) =
1006 split::sweep_line_split(&groups, starts, ends, 0_i64, options.between);
1007 return build_projected_interval_result_i64(
1008 self,
1009 &idx,
1010 starts,
1011 ends,
1012 &options.match_by,
1013 );
1014 }
1015
1016 let (starts, ends) = coords.into_i64_buffers();
1017 let (idx, starts, ends) =
1018 split::sweep_line_split(&groups, &starts, &ends, 0_i64, options.between);
1019 build_projected_interval_result_i64(self, &idx, starts, ends, &options.match_by)
1020 }
1021
1022 pub(crate) fn outer_ranges(&self, options: &OuterRangesOptions) -> Result<DataFrame> {
1023 let groups = factorize(self.df, &options.match_by)?;
1024 let coords = extract_coordinates(self.df, self.start_col, self.end_col)?;
1025
1026 if let Some((starts, ends)) = coords.as_i32() {
1027 let (idx, starts, ends, _counts) = boundary::sweep_line_boundary(&groups, starts, ends);
1028 return build_projected_interval_result_i32(
1029 self,
1030 &idx,
1031 starts,
1032 ends,
1033 &options.match_by,
1034 );
1035 }
1036
1037 if let Some((starts, ends)) = coords.as_i64() {
1038 let (idx, starts, ends, _counts) = boundary::sweep_line_boundary(&groups, starts, ends);
1039 return build_projected_interval_result_i64(
1040 self,
1041 &idx,
1042 starts,
1043 ends,
1044 &options.match_by,
1045 );
1046 }
1047
1048 let (starts, ends) = coords.into_i64_buffers();
1049 let (idx, starts, ends, _counts) = boundary::sweep_line_boundary(&groups, &starts, &ends);
1050 build_projected_interval_result_i64(self, &idx, starts, ends, &options.match_by)
1051 }
1052
1053 pub(crate) fn complement_ranges(&self, options: &ComplementRangesOptions) -> Result<DataFrame> {
1054 let merged = self.merge_overlaps(&MergeOptions {
1055 count_column: None,
1056 match_by: options.match_by.clone(),
1057 slack: 0,
1058 })?;
1059 let view = RangeView::with_columns(&merged, self.start_col, self.end_col)?;
1060 view.complement_ranges_merged(options)
1061 }
1062
1063 fn complement_ranges_merged(&self, options: &ComplementRangesOptions) -> Result<DataFrame> {
1064 let groups = factorize(self.df, &options.match_by)?;
1065 let coords = extract_coordinates(self.df, self.start_col, self.end_col)?;
1066
1067 if let Some((starts, ends)) = coords.as_i32() {
1068 let chrom_lens = complement_lengths_i32(self.df, &groups, options)?;
1069 let (_groups, starts, ends, idx) = complement_single::sweep_line_complement(
1070 &groups,
1071 starts,
1072 ends,
1073 0_i32,
1074 &chrom_lens,
1075 options.include_first_interval,
1076 );
1077 return build_projected_interval_result_i32(
1078 self,
1079 &idx,
1080 starts,
1081 ends,
1082 &options.match_by,
1083 );
1084 }
1085
1086 if let Some((starts, ends)) = coords.as_i64() {
1087 let chrom_lens = complement_lengths_i64(self.df, &groups, options)?;
1088 let (_groups, starts, ends, idx) = complement_single::sweep_line_complement(
1089 &groups,
1090 starts,
1091 ends,
1092 0_i64,
1093 &chrom_lens,
1094 options.include_first_interval,
1095 );
1096 return build_projected_interval_result_i64(
1097 self,
1098 &idx,
1099 starts,
1100 ends,
1101 &options.match_by,
1102 );
1103 }
1104
1105 let (starts, ends) = coords.into_i64_buffers();
1106 let chrom_lens = complement_lengths_i64(self.df, &groups, options)?;
1107 let (_groups, starts, ends, idx) = complement_single::sweep_line_complement(
1108 &groups,
1109 &starts,
1110 &ends,
1111 0_i64,
1112 &chrom_lens,
1113 options.include_first_interval,
1114 );
1115 build_projected_interval_result_i64(self, &idx, starts, ends, &options.match_by)
1116 }
1117
1118 pub(crate) fn nearest_ranges_by_match(
1119 &self,
1120 other: &RangeView<'_>,
1121 left_match_by: &[String],
1122 right_match_by: &[String],
1123 options: &NearestOptions,
1124 ) -> Result<DataFrame> {
1125 if options.k == 0 {
1126 return Err(RangeFrameError::InvalidNearestK { k: options.k });
1127 }
1128
1129 let (left_groups, right_groups) =
1130 factorize_pair_by(self.df, other.df, left_match_by, right_match_by)?;
1131 let left_coords = extract_coordinates(self.df, self.start_col, self.end_col)?;
1132 let right_coords = extract_coordinates(other.df, other.start_col, other.end_col)?;
1133
1134 if let (Some((left_starts, left_ends)), Some((right_starts, right_ends))) =
1135 (left_coords.as_i32(), right_coords.as_i32())
1136 {
1137 let (left_idx, right_idx, distances) = nearest::nearest(
1138 &left_groups,
1139 left_starts,
1140 left_ends,
1141 &right_groups,
1142 right_starts,
1143 right_ends,
1144 0_i32,
1145 options.k,
1146 !options.exclude_overlaps,
1147 options.direction.as_kernel_str(),
1148 options.preserve_input_order,
1149 );
1150 return build_nearest_result_i32(
1151 self.df, other.df, &left_idx, &right_idx, distances, options,
1152 );
1153 }
1154
1155 if let (Some((left_starts, left_ends)), Some((right_starts, right_ends))) =
1156 (left_coords.as_i64(), right_coords.as_i64())
1157 {
1158 let (left_idx, right_idx, distances) = nearest::nearest(
1159 &left_groups,
1160 left_starts,
1161 left_ends,
1162 &right_groups,
1163 right_starts,
1164 right_ends,
1165 0_i64,
1166 options.k,
1167 !options.exclude_overlaps,
1168 options.direction.as_kernel_str(),
1169 options.preserve_input_order,
1170 );
1171 return build_nearest_result_i64(
1172 self.df, other.df, &left_idx, &right_idx, distances, options,
1173 );
1174 }
1175
1176 let (left_starts, left_ends) = left_coords.into_i64_buffers();
1177 let (right_starts, right_ends) = right_coords.into_i64_buffers();
1178 let (left_idx, right_idx, distances) = nearest::nearest(
1179 &left_groups,
1180 &left_starts,
1181 &left_ends,
1182 &right_groups,
1183 &right_starts,
1184 &right_ends,
1185 0_i64,
1186 options.k,
1187 !options.exclude_overlaps,
1188 options.direction.as_kernel_str(),
1189 options.preserve_input_order,
1190 );
1191 build_nearest_result_i64(self.df, other.df, &left_idx, &right_idx, distances, options)
1192 }
1193
1194 pub(crate) fn subtract_overlaps(
1195 &self,
1196 other: &RangeView<'_>,
1197 options: &SubtractOptions,
1198 ) -> Result<DataFrame> {
1199 let (left_groups, right_groups) = factorize_pair(self.df, other.df, &options.match_by)?;
1200 let left_coords = extract_coordinates(self.df, self.start_col, self.end_col)?;
1201 let right_coords = extract_coordinates(other.df, other.start_col, other.end_col)?;
1202
1203 if let (Some((left_starts, left_ends)), Some((right_starts, right_ends))) =
1204 (left_coords.as_i32(), right_coords.as_i32())
1205 {
1206 let (idx, starts, ends) = subtract::sweep_line_subtract(
1207 &left_groups,
1208 left_starts,
1209 left_ends,
1210 &right_groups,
1211 right_starts,
1212 right_ends,
1213 options.preserve_input_order,
1214 );
1215 return build_subtract_result_i32(self, &idx, starts, ends);
1216 }
1217
1218 if let (Some((left_starts, left_ends)), Some((right_starts, right_ends))) =
1219 (left_coords.as_i64(), right_coords.as_i64())
1220 {
1221 let (idx, starts, ends) = subtract::sweep_line_subtract(
1222 &left_groups,
1223 left_starts,
1224 left_ends,
1225 &right_groups,
1226 right_starts,
1227 right_ends,
1228 options.preserve_input_order,
1229 );
1230 return build_subtract_result_i64(self, &idx, starts, ends);
1231 }
1232
1233 let (left_starts, left_ends) = left_coords.into_i64_buffers();
1234 let (right_starts, right_ends) = right_coords.into_i64_buffers();
1235 let (idx, starts, ends) = subtract::sweep_line_subtract(
1236 &left_groups,
1237 &left_starts,
1238 &left_ends,
1239 &right_groups,
1240 &right_starts,
1241 &right_ends,
1242 options.preserve_input_order,
1243 );
1244 build_subtract_result_i64(self, &idx, starts, ends)
1245 }
1246}
1247
1248fn build_merge_result_i32(
1249 view: &RangeView<'_>,
1250 indices: &[u32],
1251 starts: Vec<i32>,
1252 ends: Vec<i32>,
1253 counts: Vec<u32>,
1254 options: &MergeOptions,
1255) -> Result<DataFrame> {
1256 let mut out = take_merge_projection(
1257 view.df,
1258 indices,
1259 view.start_col,
1260 view.end_col,
1261 &options.match_by,
1262 )?;
1263 replace_interval_columns_i32(&mut out, view.start_col, view.end_col, starts, ends)?;
1264
1265 if let Some(count_column) = options.count_column.as_deref() {
1266 out.with_column(Column::from(Series::new(count_column.into(), counts)))?;
1267 }
1268
1269 Ok(out)
1270}
1271
1272fn build_merge_result_i64(
1273 view: &RangeView<'_>,
1274 indices: &[u32],
1275 starts: Vec<i64>,
1276 ends: Vec<i64>,
1277 counts: Vec<u32>,
1278 options: &MergeOptions,
1279) -> Result<DataFrame> {
1280 let mut out = take_merge_projection(
1281 view.df,
1282 indices,
1283 view.start_col,
1284 view.end_col,
1285 &options.match_by,
1286 )?;
1287 replace_interval_columns_i64(&mut out, view.start_col, view.end_col, starts, ends)?;
1288
1289 if let Some(count_column) = options.count_column.as_deref() {
1290 out.with_column(Column::from(Series::new(count_column.into(), counts)))?;
1291 }
1292
1293 Ok(out)
1294}
1295
1296fn build_cluster_result(
1297 df: &DataFrame,
1298 indices: &[u32],
1299 cluster_ids: Vec<u32>,
1300 cluster_column: &str,
1301) -> Result<DataFrame> {
1302 let mut out = take_rows(df, indices)?;
1303 out.with_column(Column::from(Series::new(
1304 cluster_column.into(),
1305 cluster_ids,
1306 )))?;
1307 Ok(out)
1308}
1309
1310fn build_nearest_result_i32(
1311 left_df: &DataFrame,
1312 right_df: &DataFrame,
1313 left_idx: &[u32],
1314 right_idx: &[u32],
1315 distances: Vec<i32>,
1316 options: &NearestOptions,
1317) -> Result<DataFrame> {
1318 let mut out = take_rows(left_df, left_idx)?;
1319 append_suffixed_columns(&mut out, right_df, right_idx, &options.suffix)?;
1320
1321 if let Some(distance_column) = options.distance_column.as_deref() {
1322 out.with_column(Column::from(Series::new(distance_column.into(), distances)))?;
1323 }
1324
1325 Ok(out)
1326}
1327
1328fn build_nearest_result_i64(
1329 left_df: &DataFrame,
1330 right_df: &DataFrame,
1331 left_idx: &[u32],
1332 right_idx: &[u32],
1333 distances: Vec<i64>,
1334 options: &NearestOptions,
1335) -> Result<DataFrame> {
1336 let mut out = take_rows(left_df, left_idx)?;
1337 append_suffixed_columns(&mut out, right_df, right_idx, &options.suffix)?;
1338
1339 if let Some(distance_column) = options.distance_column.as_deref() {
1340 out.with_column(Column::from(Series::new(distance_column.into(), distances)))?;
1341 }
1342
1343 Ok(out)
1344}
1345
1346fn build_subtract_result_i32(
1347 view: &RangeView<'_>,
1348 indices: &[u32],
1349 starts: Vec<i32>,
1350 ends: Vec<i32>,
1351) -> Result<DataFrame> {
1352 let mut out = take_rows(view.df, indices)?;
1353 replace_interval_columns_i32(&mut out, view.start_col, view.end_col, starts, ends)?;
1354 Ok(out)
1355}
1356
1357fn build_subtract_result_i64(
1358 view: &RangeView<'_>,
1359 indices: &[u32],
1360 starts: Vec<i64>,
1361 ends: Vec<i64>,
1362) -> Result<DataFrame> {
1363 let mut out = take_rows(view.df, indices)?;
1364 replace_interval_columns_i64(&mut out, view.start_col, view.end_col, starts, ends)?;
1365 Ok(out)
1366}
1367
1368fn build_tile_result_i32(
1369 view: &RangeView<'_>,
1370 indices: &[usize],
1371 starts: Vec<i32>,
1372 ends: Vec<i32>,
1373 overlaps: Vec<f64>,
1374 options: &TileRangesOptions,
1375) -> Result<DataFrame> {
1376 let indices = usize_indices_to_u32(indices)?;
1377 let mut out = take_rows(view.df, &indices)?;
1378 replace_interval_columns_i32(&mut out, view.start_col, view.end_col, starts, ends)?;
1379 if let Some(column) = options.overlap_column.as_deref() {
1380 out.with_column(Column::from(Series::new(column.into(), overlaps)))?;
1381 }
1382 Ok(out)
1383}
1384
1385fn build_tile_result_i64(
1386 view: &RangeView<'_>,
1387 indices: &[usize],
1388 starts: Vec<i64>,
1389 ends: Vec<i64>,
1390 overlaps: Vec<f64>,
1391 options: &TileRangesOptions,
1392) -> Result<DataFrame> {
1393 let indices = usize_indices_to_u32(indices)?;
1394 let mut out = take_rows(view.df, &indices)?;
1395 replace_interval_columns_i64(&mut out, view.start_col, view.end_col, starts, ends)?;
1396 if let Some(column) = options.overlap_column.as_deref() {
1397 out.with_column(Column::from(Series::new(column.into(), overlaps)))?;
1398 }
1399 Ok(out)
1400}
1401
1402fn build_clip_result_i32(
1403 view: &RangeView<'_>,
1404 indices: &[u32],
1405 starts: Vec<i32>,
1406 ends: Vec<i32>,
1407 options: &ClipRangesOptions,
1408) -> Result<DataFrame> {
1409 let mut out = take_rows(view.df, indices)?;
1410 if !options.remove {
1411 replace_interval_columns_i32(&mut out, view.start_col, view.end_col, starts, ends)?;
1412 }
1413 Ok(out)
1414}
1415
1416fn build_clip_result_i64(
1417 view: &RangeView<'_>,
1418 indices: &[u32],
1419 starts: Vec<i64>,
1420 ends: Vec<i64>,
1421 options: &ClipRangesOptions,
1422) -> Result<DataFrame> {
1423 let mut out = take_rows(view.df, indices)?;
1424 if !options.remove {
1425 replace_interval_columns_i64(&mut out, view.start_col, view.end_col, starts, ends)?;
1426 }
1427 Ok(out)
1428}
1429
1430fn build_group_cumsum_result_i32(
1431 view: &RangeView<'_>,
1432 indices: &[u32],
1433 cumsum_starts: Vec<i32>,
1434 cumsum_ends: Vec<i32>,
1435 options: &GroupCumsumOptions,
1436) -> Result<DataFrame> {
1437 let mut out = take_rows(view.df, indices)?;
1438 append_or_replace_cumsum_i32(
1439 &mut out,
1440 view.start_col,
1441 view.end_col,
1442 cumsum_starts,
1443 cumsum_ends,
1444 options,
1445 )?;
1446 Ok(out)
1447}
1448
1449fn build_group_cumsum_result_i64(
1450 view: &RangeView<'_>,
1451 indices: &[u32],
1452 cumsum_starts: Vec<i64>,
1453 cumsum_ends: Vec<i64>,
1454 options: &GroupCumsumOptions,
1455) -> Result<DataFrame> {
1456 let mut out = take_rows(view.df, indices)?;
1457 append_or_replace_cumsum_i64(
1458 &mut out,
1459 view.start_col,
1460 view.end_col,
1461 cumsum_starts,
1462 cumsum_ends,
1463 options,
1464 )?;
1465 Ok(out)
1466}
1467
1468fn build_join_result(
1469 left_df: &DataFrame,
1470 right_df: &DataFrame,
1471 left_idx: &[u32],
1472 right_idx: &[u32],
1473 options: &JoinOverlapsOptions,
1474) -> Result<DataFrame> {
1475 let right_names = renamed_right_column_names(left_df, right_df, &options.suffix);
1476 let mut frames = Vec::new();
1477
1478 let mut inner = take_rows(left_df, left_idx)?;
1479 append_renamed_columns(&mut inner, right_df, right_idx, &right_names)?;
1480 if let Some(column) = options.report_overlap_column.as_deref() {
1481 append_overlap_length_column(&mut inner, column, "Start", "End", &options.suffix)?;
1482 }
1483 frames.push(inner);
1484
1485 if matches!(options.join_type, JoinType::Left | JoinType::Outer) {
1486 let missing = missing_indices(left_df.height(), left_idx);
1487 if !missing.is_empty() {
1488 let mut left_missing = take_rows(left_df, &missing)?;
1489 let nulls = null_columns_for(right_df, &right_names, missing.len());
1490 left_missing.hstack_mut(&nulls)?;
1491 append_null_overlap_length_column(&mut left_missing, options, missing.len())?;
1492 frames.push(left_missing);
1493 }
1494 }
1495
1496 if matches!(options.join_type, JoinType::Right | JoinType::Outer) {
1497 let missing = missing_indices(right_df.height(), right_idx);
1498 if !missing.is_empty() {
1499 let mut right_missing = null_frame_for(left_df, missing.len())?;
1500 append_renamed_columns(&mut right_missing, right_df, &missing, &right_names)?;
1501 append_null_overlap_length_column(&mut right_missing, options, missing.len())?;
1502 frames.push(right_missing);
1503 }
1504 }
1505
1506 concat_same_schema(frames)
1507}
1508
1509fn build_intersect_result(
1510 left: &RangeView<'_>,
1511 right: &RangeView<'_>,
1512 left_idx: &[u32],
1513 right_idx: &[u32],
1514) -> Result<DataFrame> {
1515 let mut out = take_rows(left.df, left_idx)?;
1516 let left_taken = take_rows(left.df, left_idx)?;
1517 let right_taken = take_rows(right.df, right_idx)?;
1518 let left_coords = extract_coordinates(&left_taken, left.start_col, left.end_col)?;
1519 let right_coords = extract_coordinates(&right_taken, right.start_col, right.end_col)?;
1520 let (left_starts, left_ends) = left_coords.into_i64_buffers();
1521 let (right_starts, right_ends) = right_coords.into_i64_buffers();
1522
1523 let starts = left_starts
1524 .into_iter()
1525 .zip(right_starts)
1526 .map(|(left, right)| left.max(right))
1527 .collect::<Vec<_>>();
1528 let ends = left_ends
1529 .into_iter()
1530 .zip(right_ends)
1531 .map(|(left, right)| left.min(right))
1532 .collect::<Vec<_>>();
1533
1534 replace_interval_columns_i64(&mut out, left.start_col, left.end_col, starts, ends)?;
1535 Ok(out)
1536}
1537
1538fn build_projected_interval_result_i32(
1539 view: &RangeView<'_>,
1540 indices: &[u32],
1541 starts: Vec<i32>,
1542 ends: Vec<i32>,
1543 match_by: &[String],
1544) -> Result<DataFrame> {
1545 let mut out =
1546 take_interval_projection(view.df, indices, view.start_col, view.end_col, match_by)?;
1547 replace_interval_columns_i32(&mut out, view.start_col, view.end_col, starts, ends)?;
1548 Ok(out)
1549}
1550
1551fn build_projected_interval_result_i64(
1552 view: &RangeView<'_>,
1553 indices: &[u32],
1554 starts: Vec<i64>,
1555 ends: Vec<i64>,
1556 match_by: &[String],
1557) -> Result<DataFrame> {
1558 let mut out =
1559 take_interval_projection(view.df, indices, view.start_col, view.end_col, match_by)?;
1560 replace_interval_columns_i64(&mut out, view.start_col, view.end_col, starts, ends)?;
1561 Ok(out)
1562}
1563
1564fn take_merge_projection(
1565 df: &DataFrame,
1566 indices: &[u32],
1567 start_col: &str,
1568 end_col: &str,
1569 match_by: &[String],
1570) -> Result<DataFrame> {
1571 let taken = take_rows(df, indices)?;
1572 let columns = taken
1573 .columns()
1574 .iter()
1575 .filter(|column| {
1576 let name = column.name().as_str();
1577 name == start_col || name == end_col || match_by.iter().any(|value| value == name)
1578 })
1579 .cloned()
1580 .collect::<Vec<_>>();
1581 Ok(DataFrame::new_infer_height(columns)?)
1582}
1583
1584fn take_interval_projection(
1585 df: &DataFrame,
1586 indices: &[u32],
1587 start_col: &str,
1588 end_col: &str,
1589 match_by: &[String],
1590) -> Result<DataFrame> {
1591 take_merge_projection(df, indices, start_col, end_col, match_by)
1592}
1593
1594fn append_suffixed_columns(
1595 left: &mut DataFrame,
1596 right_df: &DataFrame,
1597 right_idx: &[u32],
1598 suffix: &str,
1599) -> Result<()> {
1600 let right_taken = take_rows(right_df, right_idx)?;
1601 let mut columns = Vec::with_capacity(right_taken.width());
1602
1603 for column in right_taken.columns() {
1604 let mut column = column.clone();
1605 column.rename(format!("{}{}", column.name(), suffix).into());
1606 columns.push(column);
1607 }
1608
1609 left.hstack_mut(&columns)?;
1610 Ok(())
1611}
1612
1613fn renamed_right_column_names(left: &DataFrame, right: &DataFrame, suffix: &str) -> Vec<String> {
1614 let left_names = left
1615 .get_column_names()
1616 .into_iter()
1617 .map(|name| name.as_str().to_owned())
1618 .collect::<HashSet<_>>();
1619
1620 right
1621 .get_column_names()
1622 .into_iter()
1623 .map(|name| {
1624 if left_names.contains(name.as_str()) {
1625 format!("{}{}", name.as_str(), suffix)
1626 } else {
1627 name.as_str().to_owned()
1628 }
1629 })
1630 .collect()
1631}
1632
1633fn append_renamed_columns(
1634 left: &mut DataFrame,
1635 right_df: &DataFrame,
1636 right_idx: &[u32],
1637 names: &[String],
1638) -> Result<()> {
1639 let right_taken = take_rows(right_df, right_idx)?;
1640 let mut columns = Vec::with_capacity(right_taken.width());
1641
1642 for (column, name) in right_taken.columns().iter().zip(names) {
1643 let mut column = column.clone();
1644 column.rename(name.as_str().into());
1645 columns.push(column);
1646 }
1647
1648 left.hstack_mut(&columns)?;
1649 Ok(())
1650}
1651
1652fn null_columns_for(df: &DataFrame, names: &[String], len: usize) -> Vec<Column> {
1653 df.columns()
1654 .iter()
1655 .zip(names)
1656 .map(|(column, name)| Column::full_null(name.as_str().into(), len, column.dtype()))
1657 .collect()
1658}
1659
1660fn null_frame_for(df: &DataFrame, len: usize) -> Result<DataFrame> {
1661 let columns = df
1662 .columns()
1663 .iter()
1664 .map(|column| Column::full_null(column.name().clone(), len, column.dtype()))
1665 .collect::<Vec<_>>();
1666 Ok(DataFrame::new_infer_height(columns)?)
1667}
1668
1669fn missing_indices(len: usize, present: &[u32]) -> Vec<u32> {
1670 let mut seen = vec![false; len];
1671 for &idx in present {
1672 if let Some(slot) = seen.get_mut(idx as usize) {
1673 *slot = true;
1674 }
1675 }
1676 seen.into_iter()
1677 .enumerate()
1678 .filter_map(|(idx, seen)| (!seen).then_some(idx as u32))
1679 .collect()
1680}
1681
1682fn concat_same_schema(mut frames: Vec<DataFrame>) -> Result<DataFrame> {
1683 let mut iter = frames.drain(..);
1684 let mut out = iter.next().unwrap_or_else(DataFrame::empty);
1685 for frame in iter {
1686 out.vstack_mut(&frame)?;
1687 }
1688 Ok(out)
1689}
1690
1691fn append_overlap_length_column(
1692 df: &mut DataFrame,
1693 column_name: &str,
1694 start_col: &str,
1695 end_col: &str,
1696 suffix: &str,
1697) -> Result<()> {
1698 let right_start = format!("{start_col}{suffix}");
1699 let right_end = format!("{end_col}{suffix}");
1700 let left_starts = integer_column_to_i64(df, start_col)?;
1701 let left_ends = integer_column_to_i64(df, end_col)?;
1702 let right_starts = integer_column_to_i64(df, &right_start)?;
1703 let right_ends = integer_column_to_i64(df, &right_end)?;
1704 let overlap = left_starts
1705 .into_iter()
1706 .zip(left_ends)
1707 .zip(right_starts)
1708 .zip(right_ends)
1709 .map(|(((left_start, left_end), right_start), right_end)| {
1710 left_end.min(right_end) - left_start.max(right_start)
1711 })
1712 .collect::<Vec<_>>();
1713 df.with_column(Column::from(Series::new(column_name.into(), overlap)))?;
1714 Ok(())
1715}
1716
1717fn append_null_overlap_length_column(
1718 df: &mut DataFrame,
1719 options: &JoinOverlapsOptions,
1720 len: usize,
1721) -> Result<()> {
1722 if let Some(column_name) = options.report_overlap_column.as_deref() {
1723 df.with_column(Column::full_null(column_name.into(), len, &DataType::Int64))?;
1724 }
1725 Ok(())
1726}
1727
1728fn integer_column_to_i64(df: &DataFrame, column_name: &str) -> Result<Vec<i64>> {
1729 let column = df.column(column_name)?;
1730 if column.null_count() > 0 {
1731 return Err(RangeFrameError::NullValues {
1732 column: column_name.to_owned(),
1733 });
1734 }
1735
1736 let series = column.as_materialized_series();
1737 let casted = match series.dtype() {
1738 DataType::Int8
1739 | DataType::Int16
1740 | DataType::Int32
1741 | DataType::Int64
1742 | DataType::UInt8
1743 | DataType::UInt16
1744 | DataType::UInt32
1745 | DataType::UInt64 => series.cast(&DataType::Int64)?,
1746 dtype => {
1747 return Err(RangeFrameError::InvalidCoordinateDtype {
1748 column: column_name.to_owned(),
1749 dtype: dtype.to_string(),
1750 });
1751 }
1752 };
1753 Ok(casted.i64()?.into_no_null_iter().collect())
1754}
1755
1756fn complement_lengths_i32(
1757 df: &DataFrame,
1758 groups: &[u32],
1759 options: &ComplementRangesOptions,
1760) -> Result<FxHashMap<u32, i32>> {
1761 let mut out = FxHashMap::default();
1762 let Some(chromsizes) = options.chromsizes.as_ref() else {
1763 return Ok(out);
1764 };
1765
1766 let col = df.column(&options.group_sizes_col)?;
1767 for (row_idx, group) in groups.iter().copied().enumerate() {
1768 if out.contains_key(&group) {
1769 continue;
1770 }
1771 let key = any_value_key(col.get(row_idx)?);
1772 if let Some(length) = chromsizes.get(&key) {
1773 if let Ok(length) = i32::try_from(*length) {
1774 out.insert(group, length);
1775 }
1776 }
1777 }
1778 Ok(out)
1779}
1780
1781fn complement_lengths_i64(
1782 df: &DataFrame,
1783 groups: &[u32],
1784 options: &ComplementRangesOptions,
1785) -> Result<FxHashMap<u32, i64>> {
1786 let mut out = FxHashMap::default();
1787 let Some(chromsizes) = options.chromsizes.as_ref() else {
1788 return Ok(out);
1789 };
1790
1791 let col = df.column(&options.group_sizes_col)?;
1792 for (row_idx, group) in groups.iter().copied().enumerate() {
1793 if out.contains_key(&group) {
1794 continue;
1795 }
1796 let key = any_value_key(col.get(row_idx)?);
1797 if let Some(length) = chromsizes.get(&key) {
1798 out.insert(group, *length);
1799 }
1800 }
1801 Ok(out)
1802}
1803
1804fn any_value_key(value: AnyValue<'_>) -> String {
1805 match value {
1806 AnyValue::String(value) => value.to_owned(),
1807 AnyValue::StringOwned(value) => value.as_str().to_owned(),
1808 _ => value.to_string(),
1809 }
1810}
1811
1812fn all_indices(len: usize) -> Result<Vec<u32>> {
1813 if len > u32::MAX as usize {
1814 return Err(RangeFrameError::TooManyRows { len });
1815 }
1816 Ok((0..len as u32).collect())
1817}
1818
1819fn usize_indices_to_u32(indices: &[usize]) -> Result<Vec<u32>> {
1820 indices
1821 .iter()
1822 .map(|&idx| u32::try_from(idx).map_err(|_| RangeFrameError::TooManyRows { len: idx }))
1823 .collect()
1824}
1825
1826fn resolve_extend_options(options: &ExtendRangesOptions) -> Result<(i64, i64)> {
1827 let has_ext = options.ext.is_some();
1828 let has_ends = options.ext_3.is_some() || options.ext_5.is_some();
1829 if has_ext == has_ends {
1830 return Err(compute_error(
1831 "must use at least one and not both of ext and ext_3/ext_5",
1832 ));
1833 }
1834
1835 if let Some(ext) = options.ext {
1836 Ok((ext, ext))
1837 } else {
1838 Ok((options.ext_3.unwrap_or(0), options.ext_5.unwrap_or(0)))
1839 }
1840}
1841
1842fn negative_strand_flags(df: &DataFrame, column_name: &str) -> Result<Vec<bool>> {
1843 let column = df
1844 .column(column_name)
1845 .map_err(|_| RangeFrameError::MissingStrandColumn {
1846 column: column_name.to_owned(),
1847 })?;
1848 if column.null_count() > 0 {
1849 return Err(RangeFrameError::NullValues {
1850 column: column_name.to_owned(),
1851 });
1852 }
1853
1854 let values = column.as_materialized_series().str().map_err(|_| {
1855 RangeFrameError::UnsupportedKeyDtype {
1856 column: column_name.to_owned(),
1857 dtype: column.dtype().to_string(),
1858 }
1859 })?;
1860 Ok(values
1861 .into_no_null_iter()
1862 .map(|value| value == "-")
1863 .collect())
1864}
1865
1866fn forward_strand_flags(df: &DataFrame, column_name: &str) -> Result<Vec<bool>> {
1867 Ok(negative_strand_flags(df, column_name)?
1868 .into_iter()
1869 .map(|is_negative| !is_negative)
1870 .collect())
1871}
1872
1873fn row_aligned_bounds_i32(
1874 df: &DataFrame,
1875 group_col: &str,
1876 ends: &[i32],
1877 options: &ClipRangesOptions,
1878) -> Result<Vec<i32>> {
1879 let bounds = row_aligned_bounds_i64(
1880 df,
1881 group_col,
1882 &ends.iter().map(|&v| v as i64).collect::<Vec<_>>(),
1883 options,
1884 )?;
1885 bounds
1886 .into_iter()
1887 .map(|value| i32::try_from(value).map_err(|_| compute_error("chromsize exceeds i32 range")))
1888 .collect()
1889}
1890
1891fn row_aligned_bounds_i64(
1892 df: &DataFrame,
1893 group_col: &str,
1894 ends: &[i64],
1895 options: &ClipRangesOptions,
1896) -> Result<Vec<i64>> {
1897 if options.chromsizes.is_none() {
1898 let max_end = ends.iter().copied().max().unwrap_or(0);
1899 return Ok(vec![max_end; df.height()]);
1900 }
1901
1902 let chromsizes = options.chromsizes.as_ref().expect("checked above");
1903 let col = df.column(group_col)?;
1904 let mut out = Vec::with_capacity(df.height());
1905 let mut missing = HashSet::new();
1906 for row_idx in 0..df.height() {
1907 let key = any_value_key(col.get(row_idx)?);
1908 if let Some(length) = chromsizes.get(&key) {
1909 out.push(*length);
1910 } else {
1911 missing.insert(key);
1912 out.push(0);
1913 }
1914 }
1915
1916 if !missing.is_empty() {
1917 return Err(compute_error(format!(
1918 "not all groups were in the chromsize map; missing keys: {missing:?}"
1919 )));
1920 }
1921
1922 Ok(out)
1923}
1924
1925fn append_or_replace_cumsum_i32(
1926 df: &mut DataFrame,
1927 start_col: &str,
1928 end_col: &str,
1929 starts: Vec<i32>,
1930 ends: Vec<i32>,
1931 options: &GroupCumsumOptions,
1932) -> Result<()> {
1933 match (
1934 options.cumsum_start_column.as_deref(),
1935 options.cumsum_end_column.as_deref(),
1936 ) {
1937 (Some(start_name), Some(end_name)) => {
1938 df.with_column(Column::from(Series::new(start_name.into(), starts)))?;
1939 df.with_column(Column::from(Series::new(end_name.into(), ends)))?;
1940 }
1941 (None, None) => replace_interval_columns_i32(df, start_col, end_col, starts, ends)?,
1942 _ => {
1943 return Err(compute_error(
1944 "both cumsum_start_column and cumsum_end_column must be set",
1945 ))
1946 }
1947 }
1948 Ok(())
1949}
1950
1951fn append_or_replace_cumsum_i64(
1952 df: &mut DataFrame,
1953 start_col: &str,
1954 end_col: &str,
1955 starts: Vec<i64>,
1956 ends: Vec<i64>,
1957 options: &GroupCumsumOptions,
1958) -> Result<()> {
1959 match (
1960 options.cumsum_start_column.as_deref(),
1961 options.cumsum_end_column.as_deref(),
1962 ) {
1963 (Some(start_name), Some(end_name)) => {
1964 df.with_column(Column::from(Series::new(start_name.into(), starts)))?;
1965 df.with_column(Column::from(Series::new(end_name.into(), ends)))?;
1966 }
1967 (None, None) => replace_interval_columns_i64(df, start_col, end_col, starts, ends)?,
1968 _ => {
1969 return Err(compute_error(
1970 "both cumsum_start_column and cumsum_end_column must be set",
1971 ))
1972 }
1973 }
1974 Ok(())
1975}
1976
1977fn compute_error(message: impl Into<String>) -> RangeFrameError {
1978 RangeFrameError::Polars(polars_core::prelude::PolarsError::ComputeError(
1979 message.into().into(),
1980 ))
1981}
1982
1983fn replace_interval_columns_i32(
1984 df: &mut DataFrame,
1985 start_col: &str,
1986 end_col: &str,
1987 starts: Vec<i32>,
1988 ends: Vec<i32>,
1989) -> Result<()> {
1990 df.with_column(Column::from(Series::new(start_col.into(), starts)))?;
1991 df.with_column(Column::from(Series::new(end_col.into(), ends)))?;
1992 Ok(())
1993}
1994
1995fn replace_interval_columns_i64(
1996 df: &mut DataFrame,
1997 start_col: &str,
1998 end_col: &str,
1999 starts: Vec<i64>,
2000 ends: Vec<i64>,
2001) -> Result<()> {
2002 df.with_column(Column::from(Series::new(start_col.into(), starts)))?;
2003 df.with_column(Column::from(Series::new(end_col.into(), ends)))?;
2004 Ok(())
2005}
2006
2007#[cfg(test)]
2008mod tests {
2009 use polars_core::prelude::{Column, DataFrame, NamedFrom, Series};
2010 use rustc_hash::FxHashMap;
2011
2012 use super::{
2013 ClipRangesOptions, ClusterOptions, ComplementRangesOptions, CountOverlapsOptions,
2014 ExtendRangesOptions, GroupCumsumOptions, IntersectOverlapsOptions, JoinOverlapsOptions,
2015 JoinType, MaxDisjointOptions, MergeOptions, NearestDirection, NearestOptions,
2016 OuterRangesOptions, SetIntersectOverlapsOptions, SetUnionOverlapsOptions,
2017 SortRangesOptions, SplitOverlapsOptions, SubtractOptions, TileRangesOptions,
2018 };
2019 use crate::dataframe_ext::DataFrameRanges;
2020
2021 fn make_df(columns: Vec<Series>) -> DataFrame {
2022 let columns = columns.into_iter().map(Column::from).collect::<Vec<_>>();
2023 DataFrame::new_infer_height(columns).unwrap()
2024 }
2025
2026 #[test]
2027 fn merge_overlaps_matches_pyranges_style_projection_and_counts() {
2028 let df = make_df(vec![
2029 Series::new("Chrom".into(), &["chr1", "chr1", "chr2"]),
2030 Series::new("Start".into(), &[1_i64, 4, 10]),
2031 Series::new("End".into(), &[5_i64, 8, 12]),
2032 Series::new("Score".into(), &[1_i32, 2, 3]),
2033 ]);
2034
2035 let result = df
2036 .merge_overlaps(MergeOptions {
2037 count_column: Some("Count".to_owned()),
2038 ..MergeOptions::default().with_match_by(["Chrom"])
2039 })
2040 .unwrap();
2041
2042 assert_eq!(
2043 result.get_column_names(),
2044 &["Chrom", "Start", "End", "Count"]
2045 );
2046 assert_eq!(result.height(), 2);
2047 assert_eq!(
2048 result
2049 .column("Count")
2050 .unwrap()
2051 .u32()
2052 .unwrap()
2053 .into_no_null_iter()
2054 .collect::<Vec<_>>(),
2055 vec![2, 1]
2056 );
2057 }
2058
2059 #[test]
2060 fn cluster_overlaps_adds_cluster_column() {
2061 let df = make_df(vec![
2062 Series::new("Chrom".into(), &["chr1", "chr1", "chr1", "chr2"]),
2063 Series::new("Start".into(), &[1_i64, 3, 10, 1]),
2064 Series::new("End".into(), &[4_i64, 5, 12, 2]),
2065 ]);
2066
2067 let result = df
2068 .cluster_overlaps(ClusterOptions::default().with_match_by(["Chrom"]))
2069 .unwrap();
2070
2071 assert_eq!(result.height(), 4);
2072 assert_eq!(
2073 result
2074 .column("Cluster")
2075 .unwrap()
2076 .u32()
2077 .unwrap()
2078 .into_no_null_iter()
2079 .collect::<Vec<_>>(),
2080 vec![0, 0, 1, 3]
2081 );
2082 }
2083
2084 #[test]
2085 fn nearest_ranges_appends_suffixed_columns_and_distance() {
2086 let left = make_df(vec![
2087 Series::new("Chrom".into(), &["chr1", "chr1"]),
2088 Series::new("Start".into(), &[1_i64, 20]),
2089 Series::new("End".into(), &[5_i64, 25]),
2090 Series::new("Name".into(), &["a", "b"]),
2091 ]);
2092 let right = make_df(vec![
2093 Series::new("Chrom".into(), &["chr1", "chr1"]),
2094 Series::new("Start".into(), &[8_i64, 30]),
2095 Series::new("End".into(), &[10_i64, 35]),
2096 Series::new("Name".into(), &["x", "y"]),
2097 ]);
2098
2099 let result = left
2100 .nearest_ranges(
2101 &right,
2102 NearestOptions {
2103 direction: NearestDirection::Forward,
2104 ..NearestOptions::default().with_match_by(["Chrom"])
2105 },
2106 )
2107 .unwrap();
2108
2109 assert_eq!(result.height(), 2);
2110 assert!(result.column("Chrom_b").is_ok());
2111 assert!(result.column("Distance").is_ok());
2112 }
2113
2114 #[test]
2115 fn subtract_overlaps_splits_intervals_and_preserves_other_columns() {
2116 let left = make_df(vec![
2117 Series::new("Chrom".into(), &["chr1"]),
2118 Series::new("Start".into(), &[1_i64]),
2119 Series::new("End".into(), &[10_i64]),
2120 Series::new("Name".into(), &["a"]),
2121 ]);
2122 let right = make_df(vec![
2123 Series::new("Chrom".into(), &["chr1"]),
2124 Series::new("Start".into(), &[4_i64]),
2125 Series::new("End".into(), &[6_i64]),
2126 ]);
2127
2128 let result = left
2129 .subtract_overlaps(&right, SubtractOptions::default().with_match_by(["Chrom"]))
2130 .unwrap();
2131
2132 assert_eq!(result.height(), 2);
2133 assert_eq!(
2134 result
2135 .column("Start")
2136 .unwrap()
2137 .i64()
2138 .unwrap()
2139 .into_no_null_iter()
2140 .collect::<Vec<_>>(),
2141 vec![1, 6]
2142 );
2143 assert_eq!(
2144 result
2145 .column("End")
2146 .unwrap()
2147 .i64()
2148 .unwrap()
2149 .into_no_null_iter()
2150 .collect::<Vec<_>>(),
2151 vec![4, 10]
2152 );
2153 }
2154
2155 #[test]
2156 fn count_overlaps_returns_one_count_per_left_interval() {
2157 let left = make_df(vec![
2158 Series::new("Chrom".into(), &["chr1", "chr1", "chr1"]),
2159 Series::new("Start".into(), &[1_i64, 10, 30]),
2160 Series::new("End".into(), &[5_i64, 20, 40]),
2161 ]);
2162 let right = make_df(vec![
2163 Series::new("Chrom".into(), &["chr1", "chr1", "chr1", "chr1"]),
2164 Series::new("Start".into(), &[3_i64, 11, 18, 35]),
2165 Series::new("End".into(), &[4_i64, 12, 19, 36]),
2166 ]);
2167
2168 let counts = left
2169 .count_overlaps(
2170 &right,
2171 CountOverlapsOptions::default().with_match_by(["Chrom"]),
2172 )
2173 .unwrap();
2174
2175 assert_eq!(
2176 counts
2177 .u32()
2178 .unwrap()
2179 .into_no_null_iter()
2180 .collect::<Vec<_>>(),
2181 vec![1, 2, 1]
2182 );
2183 }
2184
2185 #[test]
2186 fn join_overlaps_suffixes_right_collisions_and_supports_left_join() {
2187 let left = make_df(vec![
2188 Series::new("Chrom".into(), &["chr1", "chr1"]),
2189 Series::new("Start".into(), &[1_i64, 10]),
2190 Series::new("End".into(), &[5_i64, 20]),
2191 Series::new("Name".into(), &["a", "b"]),
2192 ]);
2193 let right = make_df(vec![
2194 Series::new("Chrom".into(), &["chr1"]),
2195 Series::new("Start".into(), &[3_i64]),
2196 Series::new("End".into(), &[4_i64]),
2197 Series::new("Score".into(), &[7_i32]),
2198 ]);
2199
2200 let result = left
2201 .join_overlaps(
2202 &right,
2203 JoinOverlapsOptions {
2204 join_type: JoinType::Left,
2205 ..JoinOverlapsOptions::default().with_match_by(["Chrom"])
2206 },
2207 )
2208 .unwrap();
2209
2210 assert_eq!(result.height(), 2);
2211 assert!(result.column("Start_b").is_ok());
2212 assert!(result.column("Score").is_ok());
2213 assert_eq!(result.column("Score").unwrap().null_count(), 1);
2214 }
2215
2216 #[test]
2217 fn join_overlaps_left_join_keeps_report_column_for_unmatched_rows() {
2218 let left = make_df(vec![
2219 Series::new("Chrom".into(), &["chr1", "chr1"]),
2220 Series::new("Start".into(), &[1_i64, 20]),
2221 Series::new("End".into(), &[5_i64, 25]),
2222 ]);
2223 let right = make_df(vec![
2224 Series::new("Chrom".into(), &["chr1"]),
2225 Series::new("Start".into(), &[3_i64]),
2226 Series::new("End".into(), &[7_i64]),
2227 ]);
2228
2229 let result = left
2230 .join_overlaps(
2231 &right,
2232 JoinOverlapsOptions {
2233 join_type: JoinType::Left,
2234 report_overlap_column: Some("Overlap".to_owned()),
2235 ..JoinOverlapsOptions::default().with_match_by(["Chrom"])
2236 },
2237 )
2238 .unwrap();
2239
2240 assert_eq!(result.height(), 2);
2241 let overlap = result.column("Overlap").unwrap().i64().unwrap();
2242 assert_eq!(overlap.get(0), Some(2));
2243 assert_eq!(overlap.get(1), None);
2244 }
2245
2246 #[test]
2247 fn intersect_overlaps_replaces_coordinates_with_intersection() {
2248 let left = make_df(vec![
2249 Series::new("Chrom".into(), &["chr1"]),
2250 Series::new("Start".into(), &[1_i64]),
2251 Series::new("End".into(), &[10_i64]),
2252 Series::new("Name".into(), &["a"]),
2253 ]);
2254 let right = make_df(vec![
2255 Series::new("Chrom".into(), &["chr1"]),
2256 Series::new("Start".into(), &[4_i64]),
2257 Series::new("End".into(), &[6_i64]),
2258 ]);
2259
2260 let result = left
2261 .intersect_overlaps(
2262 &right,
2263 IntersectOverlapsOptions::default().with_match_by(["Chrom"]),
2264 )
2265 .unwrap();
2266
2267 assert_eq!(
2268 result
2269 .column("Start")
2270 .unwrap()
2271 .i64()
2272 .unwrap()
2273 .into_no_null_iter()
2274 .collect::<Vec<_>>(),
2275 vec![4]
2276 );
2277 assert_eq!(
2278 result
2279 .column("End")
2280 .unwrap()
2281 .i64()
2282 .unwrap()
2283 .into_no_null_iter()
2284 .collect::<Vec<_>>(),
2285 vec![6]
2286 );
2287 }
2288
2289 #[test]
2290 fn set_intersect_overlaps_merges_inputs_first() {
2291 let left = make_df(vec![
2292 Series::new("Chrom".into(), &["chr1", "chr1", "chr1"]),
2293 Series::new("Start".into(), &[5_i64, 20, 40]),
2294 Series::new("End".into(), &[10_i64, 30, 50]),
2295 ]);
2296 let right = make_df(vec![
2297 Series::new("Chrom".into(), &["chr1", "chr1", "chr1", "chr1"]),
2298 Series::new("Start".into(), &[7_i64, 18, 25, 28]),
2299 Series::new("End".into(), &[9_i64, 22, 33, 32]),
2300 ]);
2301
2302 let result = left
2303 .set_intersect_overlaps(
2304 &right,
2305 SetIntersectOverlapsOptions::default().with_match_by(["Chrom"]),
2306 )
2307 .unwrap();
2308
2309 assert_eq!(
2310 result
2311 .column("Start")
2312 .unwrap()
2313 .i64()
2314 .unwrap()
2315 .into_no_null_iter()
2316 .collect::<Vec<_>>(),
2317 vec![7, 20, 25]
2318 );
2319 assert_eq!(
2320 result
2321 .column("End")
2322 .unwrap()
2323 .i64()
2324 .unwrap()
2325 .into_no_null_iter()
2326 .collect::<Vec<_>>(),
2327 vec![9, 22, 30]
2328 );
2329 }
2330
2331 #[test]
2332 fn set_union_overlaps_projects_and_merges_both_inputs() {
2333 let left = make_df(vec![
2334 Series::new("Chrom".into(), &["chr1", "chr1", "chr1"]),
2335 Series::new("Start".into(), &[1_i64, 4, 10]),
2336 Series::new("End".into(), &[3_i64, 9, 11]),
2337 Series::new("Name".into(), &["a", "b", "c"]),
2338 ]);
2339 let right = make_df(vec![
2340 Series::new("Chrom".into(), &["chr1", "chr1", "chr1"]),
2341 Series::new("Start".into(), &[2_i64, 2, 9]),
2342 Series::new("End".into(), &[3_i64, 9, 10]),
2343 ]);
2344
2345 let result = left
2346 .set_union_overlaps(
2347 &right,
2348 SetUnionOverlapsOptions::default().with_match_by(["Chrom"]),
2349 )
2350 .unwrap();
2351
2352 assert_eq!(result.get_column_names(), &["Chrom", "Start", "End"]);
2353 assert_eq!(
2354 result
2355 .column("Start")
2356 .unwrap()
2357 .i64()
2358 .unwrap()
2359 .into_no_null_iter()
2360 .collect::<Vec<_>>(),
2361 vec![1, 9, 10]
2362 );
2363 assert_eq!(
2364 result
2365 .column("End")
2366 .unwrap()
2367 .i64()
2368 .unwrap()
2369 .into_no_null_iter()
2370 .collect::<Vec<_>>(),
2371 vec![9, 10, 11]
2372 );
2373 }
2374
2375 #[test]
2376 fn sort_ranges_uses_factorized_groups_and_reverse_strand_order() {
2377 let df = make_df(vec![
2378 Series::new("Chrom".into(), &["chr1", "chr1", "chr1", "chr1"]),
2379 Series::new("Strand".into(), &["+", "+", "-", "-"]),
2380 Series::new("Start".into(), &[40_i64, 1, 10, 70]),
2381 Series::new("End".into(), &[60_i64, 11, 25, 80]),
2382 ]);
2383
2384 let result = df
2385 .sort_ranges(SortRangesOptions {
2386 match_by: vec!["Chrom".to_owned(), "Strand".to_owned()],
2387 negative_strand_column: Some("Strand".to_owned()),
2388 })
2389 .unwrap();
2390
2391 assert_eq!(
2392 result
2393 .column("Start")
2394 .unwrap()
2395 .i64()
2396 .unwrap()
2397 .into_no_null_iter()
2398 .collect::<Vec<_>>(),
2399 vec![1, 40, 70, 10]
2400 );
2401 }
2402
2403 #[test]
2404 fn extend_ranges_supports_per_row_and_grouped_extensions() {
2405 let df = make_df(vec![
2406 Series::new("Chrom".into(), &["chr1", "chr1", "chr1"]),
2407 Series::new("Start".into(), &[3_i64, 8, 5]),
2408 Series::new("End".into(), &[6_i64, 9, 7]),
2409 Series::new("Strand".into(), &["+", "+", "-"]),
2410 Series::new("Tx".into(), &["a", "a", "b"]),
2411 ]);
2412
2413 let per_row = df
2414 .extend_ranges(ExtendRangesOptions {
2415 ext: Some(3),
2416 ..ExtendRangesOptions::default()
2417 })
2418 .unwrap();
2419 assert_eq!(
2420 per_row
2421 .column("Start")
2422 .unwrap()
2423 .i64()
2424 .unwrap()
2425 .into_no_null_iter()
2426 .collect::<Vec<_>>(),
2427 vec![0, 5, 2]
2428 );
2429
2430 let grouped = df
2431 .extend_ranges(ExtendRangesOptions {
2432 match_by: vec!["Tx".to_owned()],
2433 ext_3: Some(3),
2434 ext_5: Some(0),
2435 negative_strand_column: Some("Strand".to_owned()),
2436 ..ExtendRangesOptions::default()
2437 })
2438 .unwrap();
2439 assert_eq!(
2440 grouped
2441 .column("End")
2442 .unwrap()
2443 .i64()
2444 .unwrap()
2445 .into_no_null_iter()
2446 .collect::<Vec<_>>(),
2447 vec![6, 12, 7]
2448 );
2449 assert_eq!(
2450 grouped
2451 .column("Start")
2452 .unwrap()
2453 .i64()
2454 .unwrap()
2455 .into_no_null_iter()
2456 .collect::<Vec<_>>(),
2457 vec![3, 8, 2]
2458 );
2459 }
2460
2461 #[test]
2462 fn tile_ranges_emits_tiles_and_overlap_fraction() {
2463 let df = make_df(vec![
2464 Series::new("Chrom".into(), &["chr1", "chr1"]),
2465 Series::new("Start".into(), &[99_i64, 100]),
2466 Series::new("End".into(), &[100_i64, 250]),
2467 ]);
2468
2469 let result = df
2470 .tile_ranges(TileRangesOptions {
2471 overlap_column: Some("TileOverlap".to_owned()),
2472 ..TileRangesOptions::new(100).with_match_by(["Chrom"])
2473 })
2474 .unwrap();
2475
2476 assert_eq!(
2477 result
2478 .column("Start")
2479 .unwrap()
2480 .i64()
2481 .unwrap()
2482 .into_no_null_iter()
2483 .collect::<Vec<_>>(),
2484 vec![0, 100, 200]
2485 );
2486 assert_eq!(
2487 result
2488 .column("TileOverlap")
2489 .unwrap()
2490 .f64()
2491 .unwrap()
2492 .into_no_null_iter()
2493 .collect::<Vec<_>>(),
2494 vec![0.01, 1.0, 0.5]
2495 );
2496 }
2497
2498 #[test]
2499 fn clip_ranges_clips_or_removes_out_of_bounds_intervals() {
2500 let df = make_df(vec![
2501 Series::new("Chromosome".into(), &["1", "1", "3"]),
2502 Series::new("Start".into(), &[1_i64, 249250600, 5]),
2503 Series::new("End".into(), &[2_i64, 249250640, 7]),
2504 ]);
2505 let mut chromsizes = FxHashMap::default();
2506 chromsizes.insert("1".to_owned(), 249250621);
2507 chromsizes.insert("3".to_owned(), 500);
2508
2509 let clipped = df
2510 .clip_ranges(ClipRangesOptions {
2511 chromsizes: Some(chromsizes.clone()),
2512 ..ClipRangesOptions::default()
2513 })
2514 .unwrap();
2515 assert_eq!(
2516 clipped
2517 .column("End")
2518 .unwrap()
2519 .i64()
2520 .unwrap()
2521 .into_no_null_iter()
2522 .collect::<Vec<_>>(),
2523 vec![2, 249250621, 7]
2524 );
2525
2526 let removed = df
2527 .clip_ranges(ClipRangesOptions {
2528 chromsizes: Some(chromsizes),
2529 remove: true,
2530 ..ClipRangesOptions::default()
2531 })
2532 .unwrap();
2533 assert_eq!(removed.height(), 2);
2534 }
2535
2536 #[test]
2537 fn group_cumsum_adds_running_coordinates_per_group() {
2538 let df = make_df(vec![
2539 Series::new("Chrom".into(), &["chr1", "chr1", "chr1"]),
2540 Series::new("Start".into(), &[1_i64, 40, 70]),
2541 Series::new("End".into(), &[11_i64, 60, 80]),
2542 Series::new("Tx".into(), &["t1", "t1", "t2"]),
2543 ]);
2544
2545 let result = df
2546 .group_cumsum(GroupCumsumOptions {
2547 match_by: vec!["Chrom".to_owned(), "Tx".to_owned()],
2548 cumsum_start_column: Some("CStart".to_owned()),
2549 cumsum_end_column: Some("CEnd".to_owned()),
2550 ..GroupCumsumOptions::default()
2551 })
2552 .unwrap();
2553
2554 assert_eq!(
2555 result
2556 .column("CStart")
2557 .unwrap()
2558 .i64()
2559 .unwrap()
2560 .into_no_null_iter()
2561 .collect::<Vec<_>>(),
2562 vec![0, 10, 0]
2563 );
2564 assert_eq!(
2565 result
2566 .column("CEnd")
2567 .unwrap()
2568 .i64()
2569 .unwrap()
2570 .into_no_null_iter()
2571 .collect::<Vec<_>>(),
2572 vec![10, 30, 10]
2573 );
2574 }
2575
2576 #[test]
2577 fn max_disjoint_overlaps_selects_non_overlapping_subset() {
2578 let df = make_df(vec![
2579 Series::new("Chrom".into(), &["chr1", "chr1", "chr1"]),
2580 Series::new("Start".into(), &[1_i64, 4, 10]),
2581 Series::new("End".into(), &[5_i64, 7, 14]),
2582 ]);
2583
2584 let result = df
2585 .max_disjoint_overlaps(MaxDisjointOptions::default().with_match_by(["Chrom"]))
2586 .unwrap();
2587
2588 assert_eq!(
2589 result
2590 .column("Start")
2591 .unwrap()
2592 .i64()
2593 .unwrap()
2594 .into_no_null_iter()
2595 .collect::<Vec<_>>(),
2596 vec![1, 10]
2597 );
2598 }
2599
2600 #[test]
2601 fn split_overlaps_emits_atomic_segments() {
2602 let df = make_df(vec![
2603 Series::new("Chrom".into(), &["chr1", "chr1"]),
2604 Series::new("Start".into(), &[3_i64, 5]),
2605 Series::new("End".into(), &[6_i64, 9]),
2606 ]);
2607
2608 let result = df
2609 .split_overlaps(SplitOverlapsOptions::default().with_match_by(["Chrom"]))
2610 .unwrap();
2611
2612 assert_eq!(
2613 result
2614 .column("Start")
2615 .unwrap()
2616 .i64()
2617 .unwrap()
2618 .into_no_null_iter()
2619 .collect::<Vec<_>>(),
2620 vec![3, 5, 6]
2621 );
2622 assert_eq!(
2623 result
2624 .column("End")
2625 .unwrap()
2626 .i64()
2627 .unwrap()
2628 .into_no_null_iter()
2629 .collect::<Vec<_>>(),
2630 vec![5, 6, 9]
2631 );
2632 }
2633
2634 #[test]
2635 fn outer_ranges_returns_group_boundaries() {
2636 let df = make_df(vec![
2637 Series::new("Transcript".into(), &["tr1", "tr1", "tr2"]),
2638 Series::new("Start".into(), &[1_i64, 60, 110]),
2639 Series::new("End".into(), &[40_i64, 68, 130]),
2640 ]);
2641
2642 let result = df
2643 .outer_ranges(OuterRangesOptions::default().with_match_by(["Transcript"]))
2644 .unwrap();
2645
2646 assert_eq!(
2647 result
2648 .column("Start")
2649 .unwrap()
2650 .i64()
2651 .unwrap()
2652 .into_no_null_iter()
2653 .collect::<Vec<_>>(),
2654 vec![1, 110]
2655 );
2656 assert_eq!(
2657 result
2658 .column("End")
2659 .unwrap()
2660 .i64()
2661 .unwrap()
2662 .into_no_null_iter()
2663 .collect::<Vec<_>>(),
2664 vec![68, 130]
2665 );
2666 }
2667
2668 #[test]
2669 fn complement_ranges_returns_internal_gaps_after_merge() {
2670 let df = make_df(vec![
2671 Series::new("Chrom".into(), &["chr1", "chr1", "chr1"]),
2672 Series::new("Start".into(), &[2_i64, 10, 12]),
2673 Series::new("End".into(), &[5_i64, 11, 18]),
2674 ]);
2675
2676 let result = df
2677 .complement_ranges(ComplementRangesOptions::default().with_match_by(["Chrom"]))
2678 .unwrap();
2679
2680 assert_eq!(
2681 result
2682 .column("Start")
2683 .unwrap()
2684 .i64()
2685 .unwrap()
2686 .into_no_null_iter()
2687 .collect::<Vec<_>>(),
2688 vec![5, 11]
2689 );
2690 assert_eq!(
2691 result
2692 .column("End")
2693 .unwrap()
2694 .i64()
2695 .unwrap()
2696 .into_no_null_iter()
2697 .collect::<Vec<_>>(),
2698 vec![10, 12]
2699 );
2700 }
2701}