Files
smartdb/rust/crates/rustdb-query/src/sort.rs

138 lines
4.1 KiB
Rust

use bson::{Bson, Document};
use crate::field_path::get_nested_value;
/// Sort documents according to a sort specification.
/// Sort spec: `{ field1: 1, field2: -1 }` where 1 = ascending, -1 = descending.
pub fn sort_documents(docs: &mut [Document], sort_spec: &Document) {
if sort_spec.is_empty() {
return;
}
docs.sort_by(|a, b| {
for (field, direction) in sort_spec {
let ascending = match direction {
Bson::Int32(n) => *n > 0,
Bson::Int64(n) => *n > 0,
Bson::String(s) => !s.eq_ignore_ascii_case("desc") && !s.eq_ignore_ascii_case("descending"),
_ => true,
};
let a_val = get_value(a, field);
let b_val = get_value(b, field);
let ord = compare_bson_values(&a_val, &b_val);
let ord = if ascending { ord } else { ord.reverse() };
if ord != std::cmp::Ordering::Equal {
return ord;
}
}
std::cmp::Ordering::Equal
});
}
fn get_value(doc: &Document, field: &str) -> Option<Bson> {
if field.contains('.') {
get_nested_value(doc, field)
} else {
doc.get(field).cloned()
}
}
/// Compare two BSON values for sorting purposes.
/// BSON type sort order: null < numbers < strings < objects < arrays < binData < ObjectId < bool < date
fn compare_bson_values(a: &Option<Bson>, b: &Option<Bson>) -> std::cmp::Ordering {
use std::cmp::Ordering;
match (a, b) {
(None, None) => Ordering::Equal,
(None, Some(Bson::Null)) => Ordering::Equal,
(Some(Bson::Null), None) => Ordering::Equal,
(None, Some(_)) => Ordering::Less,
(Some(_), None) => Ordering::Greater,
(Some(Bson::Null), Some(Bson::Null)) => Ordering::Equal,
(Some(Bson::Null), Some(_)) => Ordering::Less,
(Some(_), Some(Bson::Null)) => Ordering::Greater,
(Some(av), Some(bv)) => compare_typed(av, bv),
}
}
fn compare_typed(a: &Bson, b: &Bson) -> std::cmp::Ordering {
use std::cmp::Ordering;
// Cross-type numeric comparison
let a_num = to_f64(a);
let b_num = to_f64(b);
if let (Some(an), Some(bn)) = (a_num, b_num) {
return an.partial_cmp(&bn).unwrap_or(Ordering::Equal);
}
match (a, b) {
(Bson::String(x), Bson::String(y)) => x.cmp(y),
(Bson::Boolean(x), Bson::Boolean(y)) => x.cmp(y),
(Bson::DateTime(x), Bson::DateTime(y)) => x.cmp(y),
(Bson::ObjectId(x), Bson::ObjectId(y)) => x.cmp(y),
_ => {
let ta = type_order(a);
let tb = type_order(b);
ta.cmp(&tb)
}
}
}
fn to_f64(v: &Bson) -> Option<f64> {
match v {
Bson::Int32(n) => Some(*n as f64),
Bson::Int64(n) => Some(*n as f64),
Bson::Double(n) => Some(*n),
_ => None,
}
}
fn type_order(v: &Bson) -> u8 {
match v {
Bson::Null => 0,
Bson::Int32(_) | Bson::Int64(_) | Bson::Double(_) | Bson::Decimal128(_) => 1,
Bson::String(_) => 2,
Bson::Document(_) => 3,
Bson::Array(_) => 4,
Bson::Binary(_) => 5,
Bson::ObjectId(_) => 7,
Bson::Boolean(_) => 8,
Bson::DateTime(_) => 9,
_ => 10,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_sort_ascending() {
let mut docs = vec![
bson::doc! { "x": 3 },
bson::doc! { "x": 1 },
bson::doc! { "x": 2 },
];
sort_documents(&mut docs, &bson::doc! { "x": 1 });
assert_eq!(docs[0].get_i32("x").unwrap(), 1);
assert_eq!(docs[1].get_i32("x").unwrap(), 2);
assert_eq!(docs[2].get_i32("x").unwrap(), 3);
}
#[test]
fn test_sort_descending() {
let mut docs = vec![
bson::doc! { "x": 1 },
bson::doc! { "x": 3 },
bson::doc! { "x": 2 },
];
sort_documents(&mut docs, &bson::doc! { "x": -1 });
assert_eq!(docs[0].get_i32("x").unwrap(), 3);
assert_eq!(docs[1].get_i32("x").unwrap(), 2);
assert_eq!(docs[2].get_i32("x").unwrap(), 1);
}
}