facett-grid/src/sort.rs
//! **Multi-column stable sort** — click a header to sort by it; shift-click adds
//! a secondary/tertiary key. The result is a **row-index permutation** (`order`);
//! the data itself is never moved (cheap, reversible, and the original row id is
//! preserved for selection/copy). Stable on ties so prior sort keys survive.
use std::cmp::Ordering;
use crate::data::{GridData, Keys};
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum Dir {
Asc,
Desc,
}
impl Dir {
pub fn flip(self) -> Dir {
match self {
Dir::Asc => Dir::Desc,
Dir::Desc => Dir::Asc,
}
}
pub fn arrow(self) -> &'static str {
match self {
Dir::Asc => "▲",
Dir::Desc => "▼",
}
}
}
/// One sort key: a column index + direction.
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct SortKey {
pub col: usize,
pub dir: Dir,
}
/// The ordered list of active sort keys (primary first).
#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct SortSpec {
pub keys: Vec<SortKey>,
}
impl SortSpec {
pub fn is_empty(&self) -> bool {
self.keys.is_empty()
}
/// The sort priority (1-based) + direction for a column, if it's a key.
pub fn rank_of(&self, col: usize) -> Option<(usize, Dir)> {
self.keys.iter().position(|k| k.col == col).map(|i| (i + 1, self.keys[i].dir))
}
/// Click `col`: if `additive` (shift), toggle/append it as a secondary key;
/// otherwise it becomes the sole key. Re-clicking an existing key flips its
/// direction.
pub fn click(&mut self, col: usize, additive: bool) {
if additive {
if let Some(i) = self.keys.iter().position(|k| k.col == col) {
self.keys[i].dir = self.keys[i].dir.flip();
} else {
self.keys.push(SortKey { col, dir: Dir::Asc });
}
} else {
let new_dir = match self.keys.first() {
Some(k) if k.col == col && self.keys.len() == 1 => k.dir.flip(),
_ => Dir::Asc,
};
self.keys = vec![SortKey { col, dir: new_dir }];
}
}
/// Clear all sort keys (back to source order).
pub fn clear(&mut self) {
self.keys.clear();
}
}
/// `dir` reverses the value order only, never the null placement or the tie-break.
fn directed(o: Ordering, dir: Dir) -> Ordering {
if dir == Dir::Desc { o.reverse() } else { o }
}
/// Nulls stay last in both directions; only non-null values are reversed.
fn cmp_nullable<K>(a: Option<&K>, b: Option<&K>, dir: Dir, cmp: impl Fn(&K, &K) -> Ordering) -> Ordering {
match (a, b) {
(Some(a), Some(b)) => directed(cmp(a, b), dir),
(None, None) => Ordering::Equal,
(None, Some(_)) => Ordering::Greater,
(Some(_), None) => Ordering::Less,
}
}
impl Keys<'_> {
/// Compare source rows `a` and `b` under this key.
fn cmp_rows(&self, a: usize, b: usize, dir: Dir) -> Ordering {
match self {
Keys::Int(k) => cmp_nullable(k[a].as_ref(), k[b].as_ref(), dir, Ord::cmp),
Keys::UInt(k) => cmp_nullable(k[a].as_ref(), k[b].as_ref(), dir, Ord::cmp),
Keys::Float(k) => cmp_nullable(k[a].as_ref(), k[b].as_ref(), dir, f64::total_cmp),
Keys::Bool(k) => cmp_nullable(k[a].as_ref(), k[b].as_ref(), dir, Ord::cmp),
Keys::Text(k) => cmp_nullable(k[a].as_ref(), k[b].as_ref(), dir, Ord::cmp),
}
}
}
/// Build the row-index permutation that sorts `data` under `spec`. Returns the
/// identity (`0..nrows`) when `spec` is empty. Stable (preserves source order on
/// full ties).
///
/// Each key column is read ONCE, edits applied, into contiguous typed memory
/// (`Column::keys`); the primary key is then sorted as `(key, row)` pairs in one
/// `Vec`, so a comparison reads two adjacent values instead of two cells through
/// `Column::value`. Floats compare by `total_cmp`.
/// Secondary keys are looked up by row only on a primary tie, and the source row
/// breaks the last tie, which makes the unstable sort's result unique: stable.
pub fn sorted_order(data: &GridData, spec: &SortSpec) -> Vec<usize> {
let n = data.nrows();
// A key naming a missing column is skipped.
let mut keys = spec.keys.iter().filter_map(|k| Some((data.columns.get(k.col)?.keys(), k.dir)));
let Some((primary, dir)) = keys.next() else {
return (0..n).collect();
};
let rest: Vec<_> = keys.collect();
match &primary {
Keys::Int(k) => sort_by_primary(k, dir, Ord::cmp, &rest),
Keys::UInt(k) => sort_by_primary(k, dir, Ord::cmp, &rest),
Keys::Float(k) => sort_by_primary(k, dir, f64::total_cmp, &rest),
Keys::Bool(k) => sort_by_primary(k, dir, Ord::cmp, &rest),
Keys::Text(k) => sort_by_primary(k, dir, Ord::cmp, &rest),
}
}
fn sort_by_primary<K: Copy>(
keys: &[Option<K>],
dir: Dir,
cmp: impl Fn(&K, &K) -> Ordering,
rest: &[(Keys<'_>, Dir)],
) -> Vec<usize> {
let tie = |a: usize, b: usize| {
rest.iter()
.map(|(k, d)| k.cmp_rows(a, b, *d))
.find(|o| o.is_ne())
.unwrap_or(Ordering::Equal)
.then(a.cmp(&b))
};
let mut pairs = Vec::with_capacity(keys.len());
let mut nulls = Vec::new();
for (row, key) in keys.iter().enumerate() {
match key {
Some(k) => pairs.push((*k, row)),
None => nulls.push(row),
}
}
pairs.sort_unstable_by(|(ka, a), (kb, b)| directed(cmp(ka, kb), dir).then_with(|| tie(*a, *b)));
// Null primaries go last. They tie on the primary, so only the rest can
// reorder them; with no rest they are already in source order.
if !rest.is_empty() {
nulls.sort_unstable_by(|&a, &b| tie(a, b));
}
pairs.into_iter().map(|(_, row)| row).chain(nulls).collect()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::data::Column;
fn data() -> GridData {
GridData::new(vec![
("grp".into(), Column::Text { v: vec!["b".into(), "a".into(), "a".into(), "b".into()], nulls: vec![false; 4] }),
("n".into(), Column::Int { v: vec![2, 5, 1, 9], nulls: vec![false; 4] }),
])
}
#[test]
fn single_key_sorts_and_flips() {
let d = data();
let mut s = SortSpec::default();
s.click(1, false); // sort by n asc
assert_eq!(sorted_order(&d, &s), vec![2, 0, 1, 3]); // 1,2,5,9
s.click(1, false); // re-click → desc
assert_eq!(sorted_order(&d, &s), vec![3, 1, 0, 2]); // 9,5,2,1
}
#[test]
fn multi_key_is_stable_within_groups() {
let d = data();
let mut s = SortSpec::default();
s.click(0, false); // grp asc
s.click(1, true); // then n asc
// grp: a,a,b,b ; within a → n 1(idx2),5(idx1); within b → 2(idx0),9(idx3)
assert_eq!(sorted_order(&d, &s), vec![2, 1, 0, 3]);
assert_eq!(s.rank_of(0), Some((1, Dir::Asc)));
assert_eq!(s.rank_of(1), Some((2, Dir::Asc)));
}
#[test]
fn nulls_sort_last() {
let d = GridData::new(vec![(
"n".into(),
Column::Int { v: vec![3, 0, 1], nulls: vec![false, true, false] },
)]);
let mut s = SortSpec::default();
s.click(0, false);
// 1, 3, then null
assert_eq!(sorted_order(&d, &s), vec![2, 0, 1]);
s.click(0, false);
assert_eq!(sorted_order(&d, &s), vec![0, 2, 1]);
}
#[test]
fn empty_spec_is_identity() {
let d = data();
assert_eq!(sorted_order(&d, &SortSpec::default()), vec![0, 1, 2, 3]);
}
// ── The sort's contract, pinned against PR #2's comparator ───────────────
//
// `reference_order` is PR #2's sort (`a294d344`) kept as the specification:
// every comparison reads both cells through `Column::value` (edits applied),
// nulls last in both directions, `total_cmp` for floats, and the source row as
// the final tie-break. Whatever the sort does to be fast, it must return
// exactly this permutation.
use crate::arrow::{grid_data_from_batch, ArrowColumn};
use crate::data::Value;
use arrow_array::{
ArrayRef, BooleanArray, Float32Array, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray,
UInt64Array,
};
use arrow_schema::{Field, Schema};
use std::cmp::Ordering;
use std::sync::Arc;
fn reference_order(data: &GridData, spec: &SortSpec) -> Vec<usize> {
let mut order: Vec<usize> = (0..data.nrows()).collect();
order.sort_by(|&a, &b| {
for k in &spec.keys {
let Some(col) = data.columns.get(k.col) else { continue };
let o = match (col.value(a), col.value(b)) {
(None, None) => Ordering::Equal,
(None, Some(_)) => Ordering::Greater,
(Some(_), None) => Ordering::Less,
(Some(x), Some(y)) => {
let o = match (x, y) {
(Value::Int(x), Value::Int(y)) => x.cmp(&y),
(Value::UInt(x), Value::UInt(y)) => x.cmp(&y),
(Value::Float(x), Value::Float(y)) => x.total_cmp(&y),
(Value::Bool(x), Value::Bool(y)) => x.cmp(&y),
(Value::Text(x), Value::Text(y)) => x.cmp(y),
_ => unreachable!("a column has one value type"),
};
if k.dir == Dir::Desc { o.reverse() } else { o }
}
};
if o.is_ne() {
return o;
}
}
a.cmp(&b)
});
order
}
fn batch(cols: Vec<(&str, ArrayRef)>) -> GridData {
let fields: Vec<Field> = cols.iter().map(|(n, a)| Field::new(*n, a.data_type().clone(), true)).collect();
let arrays = cols.into_iter().map(|(_, a)| a).collect();
grid_data_from_batch(&RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).unwrap())
}
fn arrow_mut(data: &mut GridData, col: usize) -> &mut ArrowColumn {
match &mut data.columns[col] {
Column::Arrow(c) => c,
_ => panic!("column {col} is not shared Arrow"),
}
}
fn spec(keys: &[(usize, Dir)]) -> SortSpec {
SortSpec { keys: keys.iter().map(|&(col, dir)| SortKey { col, dir }).collect() }
}
/// Sort `col` both ways, check `asc`/`desc` and the reference.
fn assert_both_ways(d: &GridData, col: usize, asc: &[usize], desc: &[usize]) {
for (dir, want) in [(Dir::Asc, asc), (Dir::Desc, desc)] {
let s = spec(&[(col, dir)]);
assert_eq!(sorted_order(d, &s), want, "column {col} {dir:?}");
assert_eq!(sorted_order(d, &s), reference_order(d, &s), "column {col} {dir:?} vs PR #2's comparator");
}
}
#[test]
fn floats_order_by_total_cmp_with_nan_and_signed_zeros() {
let v = vec![
Some(f64::NAN), // 0
Some(-0.0), // 1
Some(0.0), // 2
Some(f64::NEG_INFINITY), // 3
Some(f64::INFINITY), // 4
None, // 5
Some(1.0), // 6
Some(-f64::NAN), // 7: sign bit set, below -inf under total_cmp
Some(-0.0), // 8
None, // 9
Some(0.0), // 10
];
let d = batch(vec![("f", Arc::new(Float64Array::from(v)))]);
// -NaN < -inf < -0.0 < +0.0 < 1 < +inf < NaN, then nulls; ties by source row.
let asc = [7, 3, 1, 8, 2, 10, 6, 4, 0, 5, 9];
// Only the values reverse: ties stay in source-row order, nulls stay last.
let desc = [0, 4, 6, 2, 10, 1, 8, 3, 7, 5, 9];
assert_both_ways(&d, 0, &asc, &desc);
}
#[test]
fn nulls_sort_last_in_both_directions_for_every_type() {
let d = batch(vec![
("i64", Arc::new(Int64Array::from(vec![None, Some(3), Some(-1), None, Some(2)]))),
("i32", Arc::new(Int32Array::from(vec![None, Some(3), Some(-1), None, Some(2)]))),
("u64", Arc::new(UInt64Array::from(vec![None, Some(u64::MAX), Some(1), None, Some(2)]))),
("f32", Arc::new(Float32Array::from(vec![None, Some(3.5), Some(-1.5), None, Some(2.5)]))),
("f64", Arc::new(Float64Array::from(vec![None, Some(3.5), Some(-1.5), None, Some(2.5)]))),
("utf8", Arc::new(StringArray::from(vec![None, Some("c"), Some("a"), None, Some("b")]))),
]);
for col in 0..d.ncols() {
assert_both_ways(&d, col, &[2, 4, 1, 0, 3], &[1, 4, 2, 0, 3]);
}
let b = batch(vec![("bool", Arc::new(BooleanArray::from(vec![None, Some(true), Some(false), None, Some(true)])))]);
assert_both_ways(&b, 0, &[2, 1, 4, 0, 3], &[1, 4, 2, 0, 3]);
// The owned (non-Arrow) storage obeys the same contract.
let owned = GridData::new(vec![
("i".into(), Column::Int { v: vec![0, 3, -1, 0, 2], nulls: vec![true, false, false, true, false] }),
("f".into(), Column::Float { v: vec![0.0, 3.5, -1.5, 0.0, 2.5], nulls: vec![true, false, false, true, false] }),
(
"t".into(),
Column::Text { v: ["", "c", "a", "", "b"].map(String::from).to_vec(), nulls: vec![true, false, false, true, false] },
),
(
"d".into(),
Column::Date { v: ["", "2003", "2001", "", "2002"].map(String::from).to_vec(), nulls: vec![true, false, false, true, false] },
),
]);
for col in 0..owned.ncols() {
assert_both_ways(&owned, col, &[2, 4, 1, 0, 3], &[1, 4, 2, 0, 3]);
}
let owned_bool = GridData::new(vec![(
"b".into(),
Column::Bool { v: vec![false, true, false, false, true], nulls: vec![true, false, false, true, false] },
)]);
assert_both_ways(&owned_bool, 0, &[2, 1, 4, 0, 3], &[1, 4, 2, 0, 3]);
}
#[test]
fn edits_are_overlaid_before_sorting() {
let mut d = batch(vec![
("f", Arc::new(Float64Array::from(vec![Some(3.0), Some(1.0), None, Some(2.0), Some(5.0)]))),
("u", Arc::new(UInt64Array::from(vec![Some(3), Some(1), None, Some(2), Some(5)]))),
("s", Arc::new(StringArray::from(vec![Some("d"), Some("b"), None, Some("c"), Some("e")]))),
("b", Arc::new(BooleanArray::from(vec![Some(true), Some(true), None, Some(false), Some(true)]))),
("i", Arc::new(Int32Array::from(vec![Some(3), Some(1), None, Some(2), Some(5)]))),
]);
// Move row 0 to the front, give the null row 2 a value, and an edit that
// does not parse reads as null (as `Column::value` reads it).
arrow_mut(&mut d, 0).edit(0, "0.5");
arrow_mut(&mut d, 0).edit(2, "2.5");
arrow_mut(&mut d, 0).edit(4, "not a number");
assert_both_ways(&d, 0, &[0, 1, 3, 2, 4], &[2, 3, 1, 0, 4]);
arrow_mut(&mut d, 1).edit(0, "0");
arrow_mut(&mut d, 1).edit(2, &u64::MAX.to_string());
assert_both_ways(&d, 1, &[0, 1, 3, 4, 2], &[2, 4, 3, 1, 0]);
arrow_mut(&mut d, 2).edit(0, "a");
arrow_mut(&mut d, 2).edit(2, "cc");
assert_both_ways(&d, 2, &[0, 1, 3, 2, 4], &[4, 2, 3, 1, 0]);
arrow_mut(&mut d, 3).edit(2, "FALSE");
arrow_mut(&mut d, 3).edit(3, "1");
assert_both_ways(&d, 3, &[2, 0, 1, 3, 4], &[0, 1, 3, 4, 2]);
arrow_mut(&mut d, 4).edit(4, "-7");
assert_both_ways(&d, 4, &[4, 1, 3, 0, 2], &[0, 3, 1, 4, 2]);
}
#[test]
fn ties_break_by_source_row_in_both_directions() {
let d = batch(vec![
("i", Arc::new(Int64Array::from(vec![2, 1, 2, 1, 2]))),
("f", Arc::new(Float64Array::from(vec![2.0, 1.0, 2.0, 1.0, 2.0]))),
("s", Arc::new(StringArray::from(vec!["b", "a", "b", "a", "b"]))),
]);
for col in 0..3 {
assert_both_ways(&d, col, &[1, 3, 0, 2, 4], &[0, 2, 4, 1, 3]);
}
}
#[test]
fn multi_key_sorts_match_the_per_comparison_reference() {
// Small value ranges so every key ties often and the later keys decide;
// nulls, NaN, signed zeros and edits throughout.
let mut seed = 0x2545_F491_4F6C_DD1Du64;
let mut next = move || {
seed ^= seed << 13;
seed ^= seed >> 7;
seed ^= seed << 17;
seed
};
let n = 3_000;
let mut pick = |k: u64| next() % k;
let ints: Vec<Option<i64>> = (0..n).map(|_| (pick(9) != 0).then(|| pick(5) as i64 - 2)).collect();
let uints: Vec<Option<u64>> = (0..n).map(|_| (pick(9) != 0).then(|| u64::MAX - pick(4))).collect();
let floats: Vec<Option<f64>> = (0..n)
.map(|_| (pick(9) != 0).then(|| [f64::NAN, -0.0, 0.0, 1.5, -1.5, f64::INFINITY, -f64::NAN][pick(7) as usize]))
.collect();
let texts: Vec<Option<&str>> = (0..n).map(|_| (pick(9) != 0).then(|| ["", "a", "ab", "b"][pick(4) as usize])).collect();
let bools: Vec<Option<bool>> = (0..n).map(|_| (pick(9) != 0).then(|| pick(2) == 1)).collect();
let mut d = batch(vec![
("i", Arc::new(Int64Array::from(ints))),
("u", Arc::new(UInt64Array::from(uints))),
("f", Arc::new(Float64Array::from(floats))),
("s", Arc::new(StringArray::from(texts))),
("b", Arc::new(BooleanArray::from(bools))),
]);
for _ in 0..200 {
let row = pick(n as u64) as usize;
match pick(5) {
0 => arrow_mut(&mut d, 0).edit(row, &(pick(5) as i64 - 2).to_string()),
1 => arrow_mut(&mut d, 1).edit(row, &(u64::MAX - pick(4)).to_string()),
2 => arrow_mut(&mut d, 2).edit(row, ["NaN", "-0", "0", "1.5", "x"][pick(5) as usize]),
3 => arrow_mut(&mut d, 3).edit(row, ["a", "b", ""][pick(3) as usize]),
_ => arrow_mut(&mut d, 4).edit(row, ["true", "false"][pick(2) as usize]),
}
}
for _ in 0..60 {
let keys: Vec<(usize, Dir)> = (0..1 + pick(3))
// Column 7 does not exist: a stale key is skipped, as before.
.map(|_| ([0, 1, 2, 3, 4, 7][pick(6) as usize], if pick(2) == 0 { Dir::Asc } else { Dir::Desc }))
.collect();
let s = spec(&keys);
assert_eq!(sorted_order(&d, &s), reference_order(&d, &s), "spec {keys:?}");
}
}
}