Skip to content

Commit f5ebd73

Browse files
mwiewiorclaude
mwiewior
and
claude
authored
Add nearest UDTF with k support and bioframe parity tests (#13)
* Add nearest UDTF with k support and bioframe parity tests * Address PR review feedback for nearest behavior * Optimize nearest k=1 hot path: remove HashMap, specialize stream loop - Remove `by_position` FnvHashMap from NearestIntervalIndex, replacing HashMap lookups with direct field extraction via extract_coitree_meta() and interval_meta_from() (eliminates O(n) construction + per-query lookups) - Specialize k=1 path in get_nearest_stream to call nearest_one() directly, bypassing Vec clear/push/iterate overhead - Pre-allocate output vectors with_capacity(num_rows * k) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Optimize nearest: cache contig lookups and intern strings in index build Cache the last contig hash-map lookup in the stream loop so sorted genomic data (~10M rows, ~24 contigs) drops from ~10M FNV hashes to ~24. Intern contig strings in build_nearest_indexes() by grouping on borrowed &str from Arrow, reducing String allocations from ~1M to ~25. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Optimize nearest: remove RepartitionExec, resolve arrays per-batch, lazy by_end - Remove forced RepartitionExec on right plan; use natural partitioning (matches interval_join CollectLeft behavior, avoids cross-thread shuffle) - Add PosArray::resolve() returning Cow<[i32]> — resolves enum variant once per batch; inner loop uses direct slice indexing with zero dispatch and no Result overhead per row - Make NearestIntervalIndex::by_end lazily initialized via OnceLock — for k=1 include_overlaps=true hot path, avoids ~24MB/contig allocation and cache pressure from the sorted copy Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Address PR review feedback for optimization commit - Add null guard in PosArray::resolve(): reject arrays with null coordinate values early instead of silently reading raw buffer garbage - Add comment explaining intentional k=1 vs k>1 loop duplication (hot-path avoids Vec alloc and clear per row) - Replace FnvHashMap<Position, ()> with FnvHashSet<Position> in nearest_k seen-set for idiomatic usage - Consolidate extract_coitree_position/extract_coitree_meta cfg blocks into per-platform modules so both functions share a single predicate; adding a new target only requires updating one block Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Replace fnv with ahash for all hash maps/sets ahash uses hardware AES-NI instructions on modern CPUs, providing better throughput than FNV for all key sizes. It was already a direct dependency (and transitive via hashbrown/DataFusion). Removes the fnv dependency entirely. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Fix import ordering after fnv→ahash migration cargo fmt reorders ahash imports alphabetically with other external crate imports. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 2b954c1 commit f5ebd73

12 files changed

Lines changed: 1702 additions & 157 deletions

File tree

‎README.md‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ This workspace provides a collection of Rust crates that implement DataFusion UD
3333
- **Interval Join**: SQL joins with range overlap conditions, automatically optimized via a custom physical planner
3434
- **Coverage**: Base-pair overlap depth between two interval sets via `SELECT * FROM coverage('reads', 'targets')`
3535
- **Count Overlaps**: Count overlapping intervals per region via `SELECT * FROM count_overlaps('reads', 'targets')`
36-
- **Nearest**: Nearest-neighbor interval matching (one match per query interval)
36+
- **Nearest**: Nearest-neighbor interval matching via SQL join (`CoitreesNearest`) or `nearest(...)` table function
3737
- **Multiple Algorithms**: Coitrees (default), IntervalTree, ArrayIntervalTree, Lapper, SuperIntervals — selectable via `SET bio.interval_join_algorithm`
3838
- **Transparent Optimization**: Hash/nested-loop joins with range conditions are automatically replaced with interval joins
3939

@@ -163,15 +163,35 @@ SELECT * FROM coverage('reads', 'targets', 'chrom', 'start', 'end', 'contig', 'p
163163

164164
-- 0-based half-open coordinates
165165
SELECT * FROM count_overlaps('reads', 'targets', 'contig', 'pos_start', 'pos_end', 'strict')
166+
167+
-- Nearest neighbors (left_* and right_* output columns)
168+
SELECT * FROM nearest('targets', 'reads')
169+
170+
-- Top-3 nearest neighbors per right interval, ignoring overlaps
171+
SELECT * FROM nearest('targets', 'reads', 3, false)
172+
173+
-- Default k=1, ignore overlaps
174+
SELECT * FROM nearest('targets', 'reads', false)
175+
176+
-- 0-based half-open coordinates
177+
SELECT * FROM nearest('targets', 'reads', 1, true, 'contig', 'pos_start', 'pos_end', 'strict')
166178
```
167179

180+
`nearest()` accepted forms:
181+
- `nearest('left', 'right')`
182+
- `nearest('left', 'right', k)`
183+
- `nearest('left', 'right', overlap_bool)`
184+
- `nearest('left', 'right', k, overlap_bool)`
185+
- optional shared or separate column names, with optional trailing `'strict'|'weak'`
186+
168187
### Nearest Interval Matching
169188

170189
```sql
171190
SET bio.interval_join_algorithm = CoitreesNearest;
172191

173192
-- Returns exactly one match per right-side row:
174193
-- the overlapping interval if one exists, otherwise the nearest by distance
194+
-- if no keyed candidate exists, left columns are NULL
175195
SELECT *
176196
FROM targets
177197
JOIN reads

‎datafusion/bio-function-ranges/Cargo.toml‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,15 +14,13 @@ tokio = { workspace = true }
1414
futures = { workspace = true }
1515
log = { workspace = true }
1616
async-trait = "0.1.88"
17-
ahash = "0.8.11"
17+
ahash = "0.8.6"
1818
coitrees = "0.4.0"
19-
fnv = "1.0.7"
20-
bio = "2.0.1"
19+
bio = "3.0.0"
2120
rust-lapper = "1.1.0"
2221
superintervals = { path = "superintervals" }
2322
parking_lot = "0.12.3"
24-
hashbrown = "0.14.5"
2523

2624
[dev-dependencies]
2725
tokio = { workspace = true, features = ["rt-multi-thread", "macros"] }
28-
rstest = "0.22.0"
26+
rstest = "0.26.1"

‎datafusion/bio-function-ranges/README.md‎

Lines changed: 52 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
# datafusion-bio-function-ranges
22

3-
Interval join, coverage, and count-overlaps for Apache DataFusion.
3+
Interval join, coverage, count-overlaps, and nearest-neighbor operations for Apache DataFusion.
44

55
This crate provides optimized genomic interval operations as DataFusion extensions:
66

@@ -34,7 +34,7 @@ register_ranges_functions(&ctx);
3434

3535
### `register_ranges_functions(ctx)`
3636

37-
Registers the `coverage` and `count_overlaps` SQL table functions on an existing `SessionContext`. This is analogous to `register_pileup_functions` in the pileup crate.
37+
Registers the `coverage`, `count_overlaps`, and `nearest` SQL table functions on an existing `SessionContext`. This is analogous to `register_pileup_functions` in the pileup crate.
3838

3939
```rust
4040
use datafusion_bio_function_ranges::register_ranges_functions;
@@ -48,7 +48,7 @@ Convenience function that creates a `SessionContext` with:
4848
- Custom query planner for automatic interval join detection
4949
- Physical optimizer rule that converts hash/nested-loop joins to interval joins
5050
- `BioConfig` extension for algorithm selection via `SET bio.*` statements
51-
- `coverage()` and `count_overlaps()` SQL table functions
51+
- `coverage()`, `count_overlaps()`, and `nearest()` SQL table functions
5252

5353
```rust
5454
use datafusion_bio_function_ranges::create_bio_session;
@@ -84,6 +84,52 @@ Counts overlapping intervals. Same interface as `coverage`, but returns the coun
8484
SELECT * FROM count_overlaps('reads', 'targets')
8585
```
8686

87+
### `nearest(left_table, right_table [, k] [, overlap] [, columns...] [, filter_op])`
88+
89+
Returns up to `k` nearest left intervals for each right interval.
90+
91+
- `k` default: `1` (must be `>= 1`)
92+
- `overlap` default: `true`
93+
- `true`: overlapping intervals are returned first, then nearest non-overlaps if needed
94+
- `false`: overlaps are ignored, only nearest non-overlaps are returned
95+
- If no keyed candidate exists for a right row, a row is still emitted with `NULL` in `left_*` columns.
96+
- Deterministic tie-break order is by `(start, end, row_position)` on the left side.
97+
98+
Output columns are prefixed with `left_` and `right_` to avoid ambiguity.
99+
100+
Accepted call forms:
101+
- `nearest('left', 'right')`
102+
- `nearest('left', 'right', 3)`
103+
- `nearest('left', 'right', false)` (use default `k=1`, disable overlap candidates)
104+
- `nearest('left', 'right', 3, false)`
105+
- `nearest('left', 'right', 3, false, 'contig', 'pos_start', 'pos_end')`
106+
- `nearest('left', 'right', 3, false, 'l_contig', 'l_start', 'l_end', 'r_contig', 'r_start', 'r_end')`
107+
- append `'strict'` or `'weak'` to any columns form above
108+
109+
```sql
110+
-- Default k=1, include overlaps
111+
SELECT * FROM nearest('targets', 'reads')
112+
113+
-- Top-3 nearest per right interval, ignoring overlaps
114+
SELECT * FROM nearest('targets', 'reads', 3, false)
115+
116+
-- Default k=1, ignore overlaps
117+
SELECT * FROM nearest('targets', 'reads', false)
118+
119+
-- 0-based half-open coordinates
120+
SELECT * FROM nearest('targets', 'reads', 1, true, 'contig', 'pos_start', 'pos_end', 'strict')
121+
122+
-- Custom column names for left and right
123+
SELECT * FROM nearest(
124+
'left_tbl',
125+
'right_tbl',
126+
5,
127+
true,
128+
'chrom', 'start', 'end',
129+
'chr', 'from', 'to'
130+
)
131+
```
132+
87133
### Filter Operations
88134

89135
| Value | Description | Use when |
@@ -130,6 +176,7 @@ JOIN reads
130176
```
131177

132178
Returns exactly one match per right-side row: the overlapping interval if one exists, otherwise the nearest interval by distance.
179+
If there is no matching key group on the left side, the row is emitted with nulls on left columns.
133180

134181
## Programmatic API
135182

@@ -207,11 +254,12 @@ register_ranges_functions(&ctx); // registers coverage() and count_overlaps() U
207254

208255
### New SQL Table Functions
209256

210-
The `coverage` and `count_overlaps` operations are now available as SQL table functions (previously only accessible via the Rust `CountOverlapsProvider` API):
257+
The `coverage`, `count_overlaps`, and `nearest` operations are now available as SQL table functions:
211258

212259
```sql
213260
SELECT * FROM coverage('reads', 'targets')
214261
SELECT * FROM count_overlaps('reads', 'targets')
262+
SELECT * FROM nearest('targets', 'reads')
215263
```
216264

217265
### Dependency Update

‎datafusion/bio-function-ranges/src/array_utils.rs‎

Lines changed: 69 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
1+
use std::borrow::Cow;
2+
13
use datafusion::arrow::array::{
2-
GenericStringArray, Int32Array, Int64Array, RecordBatch, StringViewArray, UInt32Array,
4+
Array, GenericStringArray, Int32Array, Int64Array, RecordBatch, StringViewArray, UInt32Array,
35
UInt64Array,
46
};
57
use datafusion::arrow::datatypes::DataType;
@@ -61,6 +63,72 @@ impl PosArray<'_> {
6163
}
6264
}
6365
}
66+
67+
/// Resolve the entire array to a contiguous `&[i32]` slice, converting
68+
/// non-Int32 types once per batch. The inner loop can then use direct
69+
/// indexing with zero dispatch and no `Result` overhead per row.
70+
///
71+
/// Returns an error if the array contains null values, since
72+
/// `values()` exposes raw buffer contents at null positions.
73+
pub fn resolve(&self) -> Result<Cow<'_, [i32]>> {
74+
let null_count = match self {
75+
PosArray::Int32(arr) => arr.null_count(),
76+
PosArray::Int64(arr) => arr.null_count(),
77+
PosArray::UInt32(arr) => arr.null_count(),
78+
PosArray::UInt64(arr) => arr.null_count(),
79+
};
80+
if null_count > 0 {
81+
return Err(DataFusionError::Execution(
82+
"coordinate column contains null values; nearest requires non-null coordinates"
83+
.to_string(),
84+
));
85+
}
86+
match self {
87+
PosArray::Int32(arr) => Ok(Cow::Borrowed(arr.values())),
88+
PosArray::Int64(arr) => arr
89+
.values()
90+
.iter()
91+
.enumerate()
92+
.map(|(i, &v)| {
93+
i32::try_from(v).map_err(|_| {
94+
DataFusionError::Execution(format!(
95+
"coordinate value {v} at row {i} overflows i32 (max {})",
96+
i32::MAX
97+
))
98+
})
99+
})
100+
.collect::<Result<Vec<i32>>>()
101+
.map(Cow::Owned),
102+
PosArray::UInt32(arr) => arr
103+
.values()
104+
.iter()
105+
.enumerate()
106+
.map(|(i, &v)| {
107+
i32::try_from(v).map_err(|_| {
108+
DataFusionError::Execution(format!(
109+
"coordinate value {v} at row {i} overflows i32 (max {})",
110+
i32::MAX
111+
))
112+
})
113+
})
114+
.collect::<Result<Vec<i32>>>()
115+
.map(Cow::Owned),
116+
PosArray::UInt64(arr) => arr
117+
.values()
118+
.iter()
119+
.enumerate()
120+
.map(|(i, &v)| {
121+
i32::try_from(v).map_err(|_| {
122+
DataFusionError::Execution(format!(
123+
"coordinate value {v} at row {i} overflows i32 (max {})",
124+
i32::MAX
125+
))
126+
})
127+
})
128+
.collect::<Result<Vec<i32>>>()
129+
.map(Cow::Owned),
130+
}
131+
}
64132
}
65133

66134
/// Extract contig, start, and end column arrays from a [`RecordBatch`].

‎datafusion/bio-function-ranges/src/count_overlaps.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use std::any::Any;
22
use std::fmt::{Debug, Formatter};
33
use std::sync::Arc;
44

5+
use ahash::AHashMap;
56
use async_trait::async_trait;
67
use coitrees::COITree;
78
use datafusion::arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef};
@@ -16,7 +17,6 @@ use datafusion::physical_plan::{
1617
DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
1718
};
1819
use datafusion::prelude::{Expr, SessionContext};
19-
use fnv::FnvHashMap;
2020

2121
use crate::filter_op::FilterOp;
2222
use crate::interval_tree::{build_coitree_from_batches, get_stream};
@@ -156,7 +156,7 @@ impl TableProvider for CountOverlapsProvider {
156156

157157
struct CountOverlapsExec {
158158
schema: SchemaRef,
159-
trees: Arc<FnvHashMap<String, COITree<(), u32>>>,
159+
trees: Arc<AHashMap<String, COITree<(), u32>>>,
160160
right: Arc<dyn ExecutionPlan>,
161161
columns_2: Arc<(String, String, String)>,
162162
filter_op: FilterOp,

‎datafusion/bio-function-ranges/src/interval_tree.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,21 @@
11
use std::cmp::{max, min};
22
use std::sync::Arc;
33

4+
use ahash::AHashMap;
45
use coitrees::{COITree, Interval, IntervalTree};
56
use datafusion::arrow::array::{Int64Array, RecordBatch};
67
use datafusion::arrow::datatypes::SchemaRef;
78
use datafusion::common::Result;
89
use datafusion::execution::{SendableRecordBatchStream, TaskContext};
910
use datafusion::physical_plan::ExecutionPlan;
1011
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
11-
use fnv::FnvHashMap;
1212
use futures::StreamExt;
1313
use futures::stream::BoxStream;
1414

1515
use crate::array_utils::get_join_col_arrays;
1616
use crate::filter_op::FilterOp;
1717

18-
type IntervalHashMap = FnvHashMap<String, Vec<Interval<()>>>;
18+
type IntervalHashMap = AHashMap<String, Vec<Interval<()>>>;
1919

2020
pub fn merge_intervals(mut intervals: Vec<Interval<()>>) -> Vec<Interval<()>> {
2121
if intervals.is_empty() {
@@ -44,7 +44,7 @@ pub fn build_coitree_from_batches(
4444
batches: Vec<RecordBatch>,
4545
columns: (&str, &str, &str),
4646
coverage: bool,
47-
) -> Result<FnvHashMap<String, COITree<(), u32>>> {
47+
) -> Result<AHashMap<String, COITree<(), u32>>> {
4848
let mut nodes = IntervalHashMap::default();
4949

5050
for batch in batches {
@@ -65,7 +65,7 @@ pub fn build_coitree_from_batches(
6565
}
6666
}
6767

68-
let mut trees = FnvHashMap::<String, COITree<(), u32>>::default();
68+
let mut trees = AHashMap::<String, COITree<(), u32>>::default();
6969
for (seqname, seqname_nodes) in nodes {
7070
if coverage {
7171
trees.insert(seqname, COITree::new(&merge_intervals(seqname_nodes)));
@@ -88,7 +88,7 @@ pub fn get_coverage(tree: &COITree<(), u32>, start: i32, end: i32) -> i32 {
8888
#[allow(clippy::too_many_arguments)]
8989
pub fn get_stream(
9090
right_plan: Arc<dyn ExecutionPlan>,
91-
trees: Arc<FnvHashMap<String, COITree<(), u32>>>,
91+
trees: Arc<AHashMap<String, COITree<(), u32>>>,
9292
new_schema: SchemaRef,
9393
columns_2: Arc<(String, String, String)>,
9494
filter_op: FilterOp,

‎datafusion/bio-function-ranges/src/lib.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@ pub mod array_utils;
22
pub mod count_overlaps;
33
pub mod filter_op;
44
pub mod interval_tree;
5+
pub mod nearest;
6+
pub mod nearest_index;
57
pub mod physical_planner;
68
pub mod session_context;
79
pub mod table_function;

0 commit comments

Comments
 (0)