feat:Nearest interval function (#59)

* Nearest interval function

* Disabling JoinSelection rule for nearest function

* CoitresNearest fix

* Fixing macos, windows and linux x86 builds

* add nearest test

* add more intervals to nearest test

* remove join_selection optimization rule

---------

Co-authored-by: Pavel Khamutou <[email protected]>

Author: mwiewior

sandbox/closest.md

@@ -6,6 +6,7 @@ use crate::physical_planner::joins::utils::{
 use crate::session_context::Algorithm;
 use ahash::RandomState;
 use bio::data_structures::interval_tree as rust_bio;
+use coitrees::{COITree, Interval};
 use datafusion::arrow::array::{Array, AsArray, PrimitiveArray, PrimitiveBuilder, RecordBatch};
 use datafusion::arrow::compute;
 use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef, UInt32Type};
@@ -702,6 +703,7 @@ enum IntervalJoinAlgorithm {
     ArrayIntervalTree(FnvHashMap<u64, rust_bio::ArrayBackedIntervalTree<i32, Position>>),
     AIList(FnvHashMap<u64, scailist::ScAIList<Position>>),
     Lapper(FnvHashMap<u64, rust_lapper::Lapper<u32, Position>>),
+    CoitresNearest(FnvHashMap<u64, (COITree<Position, u32>, Vec<Interval<Position>>)>),
 }
 
 impl Debug for IntervalJoinAlgorithm {
@@ -715,6 +717,8 @@ impl Debug for IntervalJoinAlgorithm {
                     .collect::<HashMap<_, _>>();
                 f.debug_struct("Coitrees").field("0", &q).finish()
             }
+            &IntervalJoinAlgorithm::CoitresNearest(_) => todo!(),
+
             IntervalJoinAlgorithm::IntervalTree(m) => {
                 f.debug_struct("IntervalTree").field("0", m).finish()
             }
@@ -749,6 +753,28 @@ impl IntervalJoinAlgorithm {
 
                 IntervalJoinAlgorithm::Coitrees(hashmap)
             }
+            Algorithm::CoitreesNearest => {
+                use coitrees::{COITree, Interval, IntervalTree};
+
+                let hashmap = hash_map
+                    .into_iter()
+                    .map(|(k, v)| {
+                        let mut intervals = v
+                            .into_iter()
+                            .map(SequilaInterval::into_coitrees)
+                            .collect::<Vec<Interval<Position>>>();
+
+                        // can hold up to u32::MAX intervals
+                        let tree: COITree<Position, u32> = COITree::new(intervals.iter());
+                        intervals.sort_by(|a, b| {
+                            a.first.cmp(&b.first).then_with(|| a.last.cmp(&b.last))
+                        });
+                        (k, (tree, intervals))
+                    })
+                    .collect::<FnvHashMap<u64, (COITree<Position, u32>, Vec<Interval<Position>>)>>(
+                    );
+                IntervalJoinAlgorithm::CoitresNearest(hashmap)
+            }
             Algorithm::IntervalTree => {
                 let hashmap = hash_map
                     .into_iter()
@@ -810,10 +836,22 @@ impl IntervalJoinAlgorithm {
     }
 
     /// unoptimized on Linux x64 (without target-cpu=native)
-    #[cfg(all(
-        target_os = "linux",
-        target_arch = "x86_64",
-        not(target_feature = "avx")
+    #[cfg(any(
+        all(
+            target_os = "linux",
+            target_arch = "x86_64",
+            not(target_feature = "avx")
+        ),
+        all(
+            target_os = "macos",
+            target_arch = "x86_64",
+            not(target_feature = "avx")
+        ),
+        all(
+            target_os = "windows",
+            target_arch = "x86_64",
+            not(target_feature = "avx")
+        ),
     ))]
     fn extract_position(&self, node: &coitrees::IntervalNode<Position, u32>) -> Position {
         node.metadata
@@ -823,12 +861,56 @@ impl IntervalJoinAlgorithm {
     #[cfg(any(
         all(target_os = "macos", target_arch = "aarch64"),
         all(target_os = "macos", target_arch = "x86_64", target_feature = "avx"),
-        all(target_os = "linux", target_arch = "x86_64", target_feature = "avx")
+        all(target_os = "linux", target_arch = "x86_64", target_feature = "avx"),
+        all(target_os = "windows", target_arch = "x86_64", target_feature = "avx")
     ))]
     fn extract_position(&self, node: &coitrees::Interval<&Position>) -> Position {
         *node.metadata
     }
 
+    fn nearest(&self, start: i32, end: i32, ranges2: &[Interval<Position>]) -> Option<Position> {
+        if ranges2.is_empty() {
+            return None;
+        }
+
+        let sorted_ranges2 = ranges2;
+
+        let mut closest_idx = None;
+        let mut min_distance = i32::MAX;
+
+        let mut left = 0;
+        let mut right = sorted_ranges2.len();
+
+        // Binary search to narrow down candidates in ranges2
+        while left < right {
+            let mid = (left + right) / 2;
+            if sorted_ranges2[mid].first < end {
+                left = mid + 1;
+            } else {
+                right = mid;
+            }
+        }
+
+        // Check ranges around the binary search result for nearest distance
+        for &i in [left.saturating_sub(1), left].iter() {
+            if let Some(r2) = sorted_ranges2.get(i) {
+                let distance = if end < r2.first {
+                    r2.first - end
+                } else if r2.last < start {
+                    start - r2.last
+                } else {
+                    0
+                };
+
+                if distance < min_distance {
+                    min_distance = distance;
+                    closest_idx = Some(r2.metadata);
+                }
+            }
+        }
+
+        closest_idx
+    }
     fn get<F>(&self, k: u64, start: i32, end: i32, mut f: F)
     where
         F: FnMut(Position),
@@ -843,6 +925,25 @@ impl IntervalJoinAlgorithm {
                     });
                 }
             }
+            IntervalJoinAlgorithm::CoitresNearest(hashmap) => {
+                use coitrees::IntervalTree;
+                if let Some(tree) = hashmap.get(&k) {
+                    let mut i = 0;
+                    // first look for overlaps and return an arbitrary one (see: https://web.mit.edu/~r/current/arch/i386_linux26/lib/R/library/IRanges/html/nearest-methods.html)
+                    tree.0.query(start, end, |node| {
+                        let position: Position = self.extract_position(node);
+                        if i == 0 {
+                            f(position);
+                            i += 1;
+                        }
+                    });
+                    // found no overlaps in the tree - try to look for nearest intervals
+                    if i == 0 {
+                        let position = self.nearest(start, end, &tree.1);
+                        f(position.unwrap());
+                    }
+                }
+            }
             IntervalJoinAlgorithm::IntervalTree(hashmap) => {
                 if let Some(tree) = hashmap.get(&k) {
                     for entry in tree.find(start..end + 1) {
@@ -1086,8 +1187,23 @@ impl IntervalJoinStream {
                 .get(*hash_val, start.value(i), end.value(i), |pos| {
                     pos_vect.push(pos as u32);
                 });
-            rle_right.push(pos_vect.len() as u32);
-            builder_left.append_slice(&pos_vect);
+            match &build_side.hash_map {
+                IntervalJoinAlgorithm::CoitresNearest(_t) => {
+                    // even if there is no hit we need to preserve the right side
+                    rle_right.push(1);
+                    if pos_vect.len() == 0 {
+                        builder_left.append_null();
+                    } else {
+                        builder_left.append_slice(&pos_vect);
+                    }
+                }
+                _ => {
+                    rle_right.push(pos_vect.len() as u32);
+                    builder_left.append_slice(&pos_vect);
+                }
+            }
+
+            // builder_left.append_slice(&pos_vect);
             pos_vect.clear();
         }
         let left_indexes = builder_left.finish();

sandbox/complement.md

@@ -5,6 +5,7 @@ use datafusion::common::extensions_options;
 use datafusion::config::ConfigExtension;
 use datafusion::execution::runtime_env::RuntimeEnv;
 use datafusion::execution::SessionStateBuilder;
+use datafusion::physical_optimizer::optimizer::PhysicalOptimizer;
 use datafusion::prelude::{SessionConfig, SessionContext};
 use log::info;
 use std::str::FromStr;
@@ -26,12 +27,16 @@ impl SeQuiLaSessionExt for SessionContext {
     }
 
     fn with_config_rt_sequila(config: SessionConfig, runtime: Arc<RuntimeEnv>) -> SessionContext {
+        let mut rules = PhysicalOptimizer::new().rules;
+        rules.retain(|rule| rule.name() != "join_selection");
+        rules.push(Arc::new(IntervalJoinPhysicalOptimizationRule));
+
         let ctx: SessionContext = SessionStateBuilder::new()
             .with_config(config)
             .with_runtime_env(runtime)
             .with_default_features()
             .with_query_planner(Arc::new(SeQuiLaQueryPlanner))
-            .with_physical_optimizer_rule(Arc::new(IntervalJoinPhysicalOptimizationRule))
+            .with_physical_optimizer_rules(rules)
             .build()
             .into();
 
@@ -61,6 +66,7 @@ pub enum Algorithm {
     ArrayIntervalTree,
     AIList,
     Lapper,
+    CoitreesNearest,
 }
 
 #[derive(Debug)]
@@ -85,6 +91,7 @@ impl FromStr for Algorithm {
             "arrayintervaltree" => Ok(Algorithm::ArrayIntervalTree),
             "ailist" => Ok(Algorithm::AIList),
             "lapper" => Ok(Algorithm::Lapper),
+            "coitreesnearest" => Ok(Algorithm::CoitreesNearest),
             _ => Err(ParseAlgorithmError(format!(
                 "Can't parse '{}' as Algorithm",
                 s
@@ -101,6 +108,7 @@ impl std::fmt::Display for Algorithm {
             Algorithm::ArrayIntervalTree => "ArrayIntervalTree",
             Algorithm::AIList => "AIList",
             Algorithm::Lapper => "Lapper",
+            Algorithm::CoitreesNearest => "CoitreesNearest",
         };
         write!(f, "{}", val)
     }

sequila/sequila-core/src/physical_planner/joins/interval_join.rs

@@ -348,3 +348,52 @@ async fn test_all_gt_lt_conditions(ctx: SessionContext) -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test(flavor = "multi_thread")]
+#[rstest::rstest]
+async fn test_nearest(ctx: SessionContext) -> Result<()> {
+    let a = r#"
+        CREATE TABLE a (contig TEXT, strand TEXT, start INTEGER, end INTEGER) AS VALUES
+        ('a', 's', 5, 10)
+    "#;
+
+    let b = r#"
+        CREATE TABLE b (contig TEXT, strand TEXT, start INTEGER, end INTEGER) AS VALUES
+        ('a', 's', 11, 13),
+        ('a', 's', 20, 21),
+        ('a', 'x', 0, 1),
+        ('b', 's', 1, 2)
+    "#;
+
+    ctx.sql("SET sequila.interval_join_algorithm TO CoitreesNearest")
+        .await?;
+
+    ctx.sql(a).await?;
+    ctx.sql(b).await?;
+
+    let q = r#"
+        SELECT * FROM a JOIN b
+        ON a.contig = b.contig AND a.strand = b.strand
+            AND a.start < b.end AND a.end > b.start
+    "#;
+
+    let result = ctx.sql(q).await?;
+    result.clone().show().await?;
+
+    let results = result.collect().await?;
+
+    let expected = [
+        "+--------+--------+-------+-----+--------+--------+-------+-----+",
+        "| contig | strand | start | end | contig | strand | start | end |",
+        "+--------+--------+-------+-----+--------+--------+-------+-----+",
+        "|        |        |       |     | a      | x      | 0     | 1   |",
+        "|        |        |       |     | b      | s      | 1     | 2   |",
+        "| a      | s      | 5     | 10  | a      | s      | 11    | 13  |",
+        "| a      | s      | 5     | 10  | a      | s      | 20    | 21  |",
+        "+--------+--------+-------+-----+--------+--------+-------+-----+",
+    ];
+
+    assert_batches_sorted_eq!(expected, &results);
+
+    Ok(())
+}