use super::super::ast::{NullsPlacement, OrderItem};
use crate::datatypes::values::Value;
use std::cmp::Ordering;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct SortSpec {
pub(crate) ascending: bool,
pub(crate) nulls: NullsPlacement,
}
impl SortSpec {
pub(crate) fn from_order_item(item: &OrderItem) -> Self {
SortSpec {
ascending: item.ascending,
nulls: item.effective_nulls(),
}
}
}
pub(crate) fn compare_sort_keys(a: &[Value], b: &[Value], specs: &[SortSpec]) -> Ordering {
for (i, spec) in specs.iter().enumerate() {
let key_a = a.get(i).unwrap_or(&Value::Null);
let key_b = b.get(i).unwrap_or(&Value::Null);
let a_null = matches!(key_a, Value::Null);
let b_null = matches!(key_b, Value::Null);
match (a_null, b_null) {
(true, true) => continue,
(true, false) => {
return match spec.nulls {
NullsPlacement::First => Ordering::Less,
NullsPlacement::Last => Ordering::Greater,
};
}
(false, true) => {
return match spec.nulls {
NullsPlacement::First => Ordering::Greater,
NullsPlacement::Last => Ordering::Less,
};
}
(false, false) => {}
}
let ordering = compare_one(key_a, key_b);
let oriented = if spec.ascending {
ordering
} else {
ordering.reverse()
};
if oriented != Ordering::Equal {
return oriented;
}
}
Ordering::Equal
}
#[inline]
fn compare_one(a: &Value, b: &Value) -> Ordering {
match (a, b) {
(Value::Float64(x), Value::Float64(y)) => {
crate::graph::core::filtering::cmp_f64_total(*x, *y)
}
(Value::Int64(x), Value::Int64(y)) => x.cmp(y),
(Value::String(x), Value::String(y)) => x.cmp(y),
_ => crate::graph::core::filtering::total_order(a, b),
}
}
#[inline]
fn fast_lane(keys: &[Value], sign: f64) -> f64 {
match keys.first() {
Some(Value::Float64(f)) => sign * f,
Some(Value::Int64(i)) => sign * (*i as f64),
Some(Value::UniqueId(u)) => sign * (*u as f64),
_ => f64::NAN,
}
}
struct Entry<P> {
keys: Vec<Value>,
seq: usize,
lane: f64,
payload: P,
}
pub(crate) struct TopKCollector<P> {
limit: usize,
specs: Vec<SortSpec>,
sign: f64,
heap: Vec<Entry<P>>,
}
impl<P> TopKCollector<P> {
pub(crate) fn new(specs: Vec<SortSpec>, limit: usize) -> Self {
let sign = match specs.first() {
Some(spec) if spec.ascending => 1.0,
Some(_) => -1.0,
None => f64::NAN,
};
TopKCollector {
limit,
specs,
sign,
heap: Vec::with_capacity(limit.min(1024)),
}
}
#[inline]
fn rank(&self, a: &Entry<P>, b: &Entry<P>) -> Ordering {
match a.lane.partial_cmp(&b.lane) {
Some(Ordering::Less) => Ordering::Less,
Some(Ordering::Greater) => Ordering::Greater,
_ => compare_sort_keys(&a.keys, &b.keys, &self.specs).then_with(|| a.seq.cmp(&b.seq)),
}
}
#[inline]
fn rank_candidate(&self, keys: &[Value], lane: f64, seq: usize, other: &Entry<P>) -> Ordering {
match lane.partial_cmp(&other.lane) {
Some(Ordering::Less) => Ordering::Less,
Some(Ordering::Greater) => Ordering::Greater,
_ => {
compare_sort_keys(keys, &other.keys, &self.specs).then_with(|| seq.cmp(&other.seq))
}
}
}
pub(crate) fn accepts(&self, keys: &[Value], seq: usize) -> bool {
if self.limit == 0 {
return false;
}
if self.heap.len() < self.limit {
return true;
}
match self.heap.first() {
Some(root) => {
self.rank_candidate(keys, fast_lane(keys, self.sign), seq, root) == Ordering::Less
}
None => true,
}
}
pub(crate) fn push(&mut self, keys: &[Value], seq: usize, payload: P) {
if self.limit == 0 {
return;
}
let lane = fast_lane(keys, self.sign);
if self.heap.len() < self.limit {
self.heap.push(Entry {
keys: keys.to_vec(),
seq,
lane,
payload,
});
self.sift_up(self.heap.len() - 1);
return;
}
if self.rank_candidate(keys, lane, seq, &self.heap[0]) != Ordering::Less {
return;
}
let root = &mut self.heap[0];
root.keys.clear();
root.keys.extend_from_slice(keys);
root.seq = seq;
root.lane = lane;
root.payload = payload;
self.sift_down(0);
}
fn sift_up(&mut self, mut idx: usize) {
while idx > 0 {
let parent = (idx - 1) / 2;
if self.rank(&self.heap[idx], &self.heap[parent]) != Ordering::Greater {
break;
}
self.heap.swap(idx, parent);
idx = parent;
}
}
fn sift_down(&mut self, mut idx: usize) {
let len = self.heap.len();
loop {
let (left, right) = (2 * idx + 1, 2 * idx + 2);
let mut worst = idx;
if left < len && self.rank(&self.heap[left], &self.heap[worst]) == Ordering::Greater {
worst = left;
}
if right < len && self.rank(&self.heap[right], &self.heap[worst]) == Ordering::Greater {
worst = right;
}
if worst == idx {
return;
}
self.heap.swap(idx, worst);
idx = worst;
}
}
pub(crate) fn into_sorted(mut self) -> Vec<(Vec<Value>, P)> {
let specs = std::mem::take(&mut self.specs);
self.heap.sort_by(|a, b| {
compare_sort_keys(&a.keys, &b.keys, &specs).then_with(|| a.seq.cmp(&b.seq))
});
self.heap
.into_iter()
.map(|entry| (entry.keys, entry.payload))
.collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn asc() -> SortSpec {
SortSpec {
ascending: true,
nulls: NullsPlacement::Last,
}
}
fn desc() -> SortSpec {
SortSpec {
ascending: false,
nulls: NullsPlacement::First,
}
}
fn full_sort(rows: &[Vec<Value>], specs: &[SortSpec], limit: usize) -> Vec<usize> {
let mut idx: Vec<usize> = (0..rows.len()).collect();
idx.sort_by(|&a, &b| compare_sort_keys(&rows[a], &rows[b], specs));
idx.truncate(limit);
idx
}
fn collect_top_k(rows: &[Vec<Value>], specs: &[SortSpec], limit: usize) -> Vec<usize> {
let mut collector: TopKCollector<usize> = TopKCollector::new(specs.to_vec(), limit);
for (seq, keys) in rows.iter().enumerate() {
if collector.accepts(keys, seq) {
collector.push(keys, seq, seq);
}
}
collector
.into_sorted()
.into_iter()
.map(|(_, payload)| payload)
.collect()
}
fn push_no_reuse(c: &mut TopKCollector<usize>, keys: &[Value], seq: usize, payload: usize) {
if c.heap.len() < c.limit {
c.push(keys, seq, payload);
return;
}
let lane = fast_lane(keys, c.sign);
if c.rank_candidate(keys, lane, seq, &c.heap[0]) != Ordering::Less {
return;
}
c.heap[0] = Entry {
keys: keys.to_vec(),
seq,
lane,
payload,
};
c.sift_down(0);
}
fn best<R>(rounds: usize, mut f: impl FnMut() -> R) -> std::time::Duration {
use std::time::Instant;
for _ in 0..3 {
std::hint::black_box(f());
}
let mut best = std::time::Duration::MAX;
for _ in 0..rounds {
let t = Instant::now();
std::hint::black_box(f());
let d = t.elapsed();
if d < best {
best = d;
}
}
best
}
#[test]
#[ignore]
fn top_k_retention_cost() {
let n = 50_000usize;
let numeric: Vec<Value> = (0..n).map(|i| Value::Float64(i as f64)).collect();
let strings: Vec<Value> = (0..n)
.map(|i| Value::String(format!("hc_{}", i % (n / 2))))
.collect();
let cmp = best(30, || {
let mut sink = 0usize;
for i in 0..n {
if compare_sort_keys(
std::slice::from_ref(&numeric[i]),
std::slice::from_ref(&numeric[(i + 7) % n]),
&[desc()],
) == Ordering::Less
{
sink += 1;
}
}
sink
});
println!("compare_sort_keys x{n} (Float64): min={cmp:>10.3?}");
for (label, col) in [("numeric", &numeric), ("string", &strings)] {
let ctrl = best(30, || {
let mut buf: Vec<Value> = Vec::with_capacity(1);
let mut sink = 0usize;
for v in col.iter() {
buf.clear();
buf.push(v.clone());
sink += buf.len();
}
sink
});
println!("{label:8} loop-only min={ctrl:>10.3?}");
for (dir, spec) in [("asc", asc()), ("desc", desc())] {
let mut accepted = 0usize;
let full = best(30, || {
let mut buf: Vec<Value> = Vec::with_capacity(1);
accepted = 0;
let mut c: TopKCollector<usize> = TopKCollector::new(vec![spec], 10);
for (seq, v) in col.iter().enumerate() {
buf.clear();
buf.push(v.clone());
if c.accepts(&buf, seq) {
accepted += 1;
c.push(&buf, seq, seq);
}
}
c.into_sorted().len()
});
let no_reuse = best(30, || {
let mut buf: Vec<Value> = Vec::with_capacity(1);
let mut d: TopKCollector<usize> = TopKCollector::new(vec![spec], 10);
for (seq, v) in col.iter().enumerate() {
buf.clear();
buf.push(v.clone());
if d.accepts(&buf, seq) {
push_no_reuse(&mut d, &buf, seq, seq);
}
}
d.into_sorted().len()
});
println!(
"{label:8} {dir:5} reuse={full:>10.3?} fresh_alloc={no_reuse:>10.3?} \
accepted={accepted:6}"
);
}
}
}
#[test]
fn fast_lane_never_disagrees_with_the_full_sort() {
const BIG: i64 = (1i64 << 53) + 1;
let rows: Vec<Vec<Value>> = (0..400)
.map(|i| {
let key0 = match i % 7 {
0 => Value::Int64(BIG + (i as i64 % 3)),
1 => Value::Float64((i % 5) as f64),
2 => Value::Int64((i % 5) as i64),
3 => Value::UniqueId((i % 5) as u32),
4 => Value::Null,
5 => Value::String(format!("k{}", i % 4)),
_ => Value::Float64(f64::from(-(i % 4))),
};
vec![key0, Value::Int64((i % 11) as i64)]
})
.collect();
for specs in [
vec![asc()],
vec![desc()],
vec![asc(), desc()],
vec![desc(), asc()],
vec![
SortSpec {
ascending: false,
nulls: NullsPlacement::Last,
},
asc(),
],
] {
for limit in [1usize, 3, 17, 400] {
assert_eq!(
collect_top_k(&rows, &specs, limit),
full_sort(&rows, &specs, limit),
"fast lane diverged from full sort at limit {limit} for {specs:?}"
);
}
}
}
#[test]
fn nulls_place_by_spec_not_by_compare_values() {
let specs = [desc()];
assert_eq!(
compare_sort_keys(&[Value::Null], &[Value::Int64(9)], &specs),
Ordering::Less
);
let specs = [asc()];
assert_eq!(
compare_sort_keys(&[Value::Null], &[Value::Int64(9)], &specs),
Ordering::Greater
);
}
#[test]
fn later_keys_break_ties_independently_of_direction() {
let specs = [desc(), asc()];
let a = vec![Value::Int64(1), Value::Int64(5)];
let b = vec![Value::Int64(1), Value::Int64(7)];
assert_eq!(compare_sort_keys(&a, &b, &specs), Ordering::Less);
let specs = [desc(), desc()];
assert_eq!(compare_sort_keys(&a, &b, &specs), Ordering::Greater);
}
#[test]
fn a_cross_type_key_is_decided_by_the_type_rank() {
let specs = [asc(), asc()];
let a = vec![Value::String("x".into()), Value::Int64(9)];
let b = vec![Value::Boolean(true), Value::Int64(1)];
assert_eq!(compare_sort_keys(&a, &b, &specs), Ordering::Less);
let specs = [desc(), asc()];
assert_eq!(compare_sort_keys(&a, &b, &specs), Ordering::Greater);
}
fn one_of_every_rank_class() -> Vec<Value> {
use crate::datatypes::values::{NodeValue, PathValue, RelValue};
let node = NodeValue {
id: 1,
labels: vec!["N".into()],
properties: crate::datatypes::PropMap::new(),
};
let rel = RelValue {
id: 1,
start_id: 1,
end_id: 2,
rel_type: "R".into(),
properties: crate::datatypes::PropMap::new(),
};
vec![
Value::Map(crate::datatypes::PropMap::from_iter([(
"k",
Value::Int64(1),
)])),
Value::Node(Box::new(node.clone())),
Value::NodeRef(3),
Value::Relationship(Box::new(rel.clone())),
Value::List(vec![Value::Int64(1)]),
Value::Path(Box::new(PathValue {
nodes: vec![node],
rels: vec![rel],
})),
Value::DateTime(chrono::NaiveDate::from_ymd_opt(2024, 1, 1).unwrap()),
Value::Timestamp(
chrono::NaiveDate::from_ymd_opt(2024, 1, 1)
.unwrap()
.and_hms_opt(12, 0, 0)
.unwrap(),
),
Value::Duration {
months: 0,
days: 1,
seconds: 0,
},
Value::Point { lat: 1.0, lon: 2.0 },
Value::String("s".into()),
Value::Boolean(false),
Value::Boolean(true),
Value::Float64(-1.5),
Value::Int64(0),
Value::UniqueId(1),
Value::Float64(1.5),
Value::Float64(f64::NAN),
Value::Null,
]
}
#[test]
fn the_type_rank_orders_every_value_class() {
let expected = one_of_every_rank_class();
let rows: Vec<Vec<Value>> = expected.iter().cloned().map(|v| vec![v]).collect();
let mut idx: Vec<usize> = (0..rows.len()).collect();
idx.shuffle_deterministically();
idx.sort_by(|&a, &b| {
compare_sort_keys(
&rows[a],
&rows[b],
&[SortSpec {
ascending: true,
nulls: NullsPlacement::Last,
}],
)
});
let sorted: Vec<Value> = idx.iter().map(|&i| expected[i].clone()).collect();
assert_eq!(
format!("{sorted:?}"),
format!("{expected:?}"),
"ascending total order does not match the documented rank table"
);
let mut idx: Vec<usize> = (0..rows.len()).collect();
idx.shuffle_deterministically();
idx.sort_by(|&a, &b| {
compare_sort_keys(
&rows[a],
&rows[b],
&[SortSpec {
ascending: false,
nulls: NullsPlacement::First,
}],
)
});
let sorted: Vec<Value> = idx.iter().map(|&i| expected[i].clone()).collect();
let mut reversed = expected.clone();
reversed.reverse();
assert_eq!(
format!("{sorted:?}"),
format!("{reversed:?}"),
"descending order is not the reverse of ascending"
);
}
trait ShuffleDeterministically {
fn shuffle_deterministically(&mut self);
}
impl ShuffleDeterministically for Vec<usize> {
fn shuffle_deterministically(&mut self) {
let mut state: u64 = 0x9E37_79B9_7F4A_7C15;
for i in (1..self.len()).rev() {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
self.swap(i, (state % (i as u64 + 1)) as usize);
}
}
}
#[test]
fn the_total_order_is_antisymmetric_and_transitive() {
use crate::graph::core::filtering::total_order;
let mut values = one_of_every_rank_class();
values.extend([
Value::Int64(0),
Value::Float64(0.0),
Value::Float64(-0.0),
Value::UniqueId(0),
Value::Int64((1i64 << 53) + 1),
Value::Int64(1i64 << 53),
Value::Float64((1u64 << 53) as f64),
Value::Float64(f64::INFINITY),
Value::Float64(f64::NEG_INFINITY),
Value::String(String::new()),
Value::List(vec![]),
Value::List(vec![Value::Int64(1), Value::String("a".into())]),
]);
for a in &values {
assert_eq!(total_order(a, a), Ordering::Equal, "not reflexive: {a:?}");
for b in &values {
assert_eq!(
total_order(a, b),
total_order(b, a).reverse(),
"not antisymmetric: {a:?} vs {b:?}"
);
}
}
for a in &values {
for b in &values {
let ab = total_order(a, b);
if ab == Ordering::Greater {
continue;
}
for c in &values {
let bc = total_order(b, c);
if bc == Ordering::Greater {
continue;
}
let ac = total_order(a, c);
assert_ne!(ac, Ordering::Greater, "not transitive: {a:?} {b:?} {c:?}");
if ab == Ordering::Equal && bc == Ordering::Equal {
assert_eq!(
ac,
Ordering::Equal,
"equality not transitive: {a:?} {b:?} {c:?}"
);
}
}
}
}
}
#[test]
fn the_inline_shortcut_agrees_with_the_total_order() {
use crate::graph::core::filtering::total_order;
let mut values = one_of_every_rank_class();
values.extend([
Value::Int64(-7),
Value::Float64(f64::NAN),
Value::String("s".into()),
Value::String("t".into()),
Value::Float64(1.5),
]);
for a in &values {
for b in &values {
assert_eq!(
compare_one(a, b),
total_order(a, b),
"shortcut disagrees for {a:?} vs {b:?}"
);
}
}
}
#[test]
fn integers_past_2_pow_53_compare_exactly_against_floats() {
use crate::graph::core::filtering::total_order;
const BIG: i64 = (1i64 << 53) + 1;
let float = Value::Float64((1u64 << 53) as f64);
assert_eq!(total_order(&Value::Int64(BIG), &float), Ordering::Greater);
assert_eq!(
total_order(&Value::Int64(1i64 << 53), &float),
Ordering::Equal
);
assert_eq!(total_order(&float, &Value::Int64(BIG)), Ordering::Less);
assert_eq!(
total_order(&Value::Int64(2), &Value::Float64(2.5)),
Ordering::Less
);
assert_eq!(
total_order(&Value::Int64(-2), &Value::Float64(-2.5)),
Ordering::Greater
);
assert_eq!(
total_order(&Value::Int64(i64::MAX), &Value::Float64(f64::INFINITY)),
Ordering::Less
);
assert_eq!(
total_order(&Value::Int64(i64::MIN), &Value::Float64(f64::NEG_INFINITY)),
Ordering::Greater
);
assert_eq!(
total_order(&Value::Int64(i64::MAX), &Value::Float64(f64::NAN)),
Ordering::Less
);
assert_eq!(
total_order(&Value::UniqueId(3), &Value::Int64(3)),
Ordering::Equal
);
}
fn mixed_rows(n: usize) -> Vec<Vec<Value>> {
let mut state: u64 = 0x2545_F491_4F6C_DD1D;
let mut next = move || {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
state
};
(0..n)
.map(|_| {
let key = if next() % 2 == 0 {
Value::Int64((next() % 7) as i64)
} else {
Value::String(format!("s{}", next() % 7))
};
vec![key]
})
.collect()
}
#[test]
fn mixed_type_column_sorts_without_panicking() {
for n in [21usize, 24, 32, 64, 400] {
let rows = mixed_rows(n);
for specs in [vec![asc()], vec![desc()]] {
let mut idx: Vec<usize> = (0..rows.len()).collect();
idx.sort_by(|&a, &b| compare_sort_keys(&rows[a], &rows[b], &specs));
assert_eq!(idx.len(), n);
}
}
}
#[test]
fn top_k_matches_a_stable_full_sort() {
let mut state: u64 = 0x2545_F491_4F6C_DD1D;
let mut next = move || {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
state
};
let rows: Vec<Vec<Value>> = (0..500)
.map(|_| {
let a = next() % 7;
let b = next() % 5;
vec![
if a == 0 {
Value::Null
} else {
Value::Int64(a as i64)
},
if b == 0 {
Value::Null
} else {
Value::String(format!("s{b}"))
},
]
})
.collect();
for specs in [
vec![asc(), asc()],
vec![desc(), asc()],
vec![asc(), desc()],
vec![desc(), desc()],
vec![
SortSpec {
ascending: true,
nulls: NullsPlacement::First,
},
SortSpec {
ascending: false,
nulls: NullsPlacement::Last,
},
],
] {
for limit in [0usize, 1, 3, 25, 500, 600] {
assert_eq!(
collect_top_k(&rows, &specs, limit),
full_sort(&rows, &specs, limit),
"top-K diverged from full sort at limit {limit} for {specs:?}"
);
}
}
}
}