Filtering/querying a multi-key btree index in Rust

user_14471985 0
Viewed 106

I'm trying to use a BTreeMap as an index in an in-memory database, with multiple keys. Something like this:

let mut map = BTreeMap::new();
map.insert(("a".to_string(), "x".to_string()), "ax".to_string());
map.insert(("a".to_string(), "y".to_string()), "ay".to_string());

Now my question is, what would be the best way to actually query this? Say for instance I'd like to get ("a", *), i.e. all entries with "a" as the first key and anything as second.

I tried something like this:

use std::{collections::{BTreeMap}, cmp::Ordering};
use std::ops::Bound::{Included, Unbounded};

#[derive(Clone, Debug, Hash)]
pub enum StringKey {
    Exact(String),
    Any,
}
impl PartialOrd for StringKey {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        match (self, other) {
            (StringKey::Exact(a), StringKey::Exact(b)) => Some(a.cmp(b)),
            (StringKey::Exact(_), StringKey::Any) => Some(Ordering::Equal),
            (StringKey::Any, StringKey::Exact(_)) => Some(Ordering::Equal),
            (StringKey::Any, StringKey::Any) => Some(Ordering::Equal),
        }
    }
}
impl Ord for StringKey {
    fn cmp(&self, other: &Self) -> Ordering {
        self.partial_cmp(other).unwrap()
    }
}
impl PartialEq for StringKey {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Self::Exact(a), Self::Exact(b)) => a == b,
            (Self::Exact(_), Self::Any) => true,
            (Self::Any, Self::Exact(_)) => true,
            (Self::Any, Self::Any) => true,
        }
    }
}
impl Eq for StringKey {
}

fn main() {
    let mut map = BTreeMap::new();
    map.insert((StringKey::Exact("a".to_string()), StringKey::Exact("x".to_string())), "ax".to_string());
    map.insert((StringKey::Exact("a".to_string()), StringKey::Exact("y".to_string())), "ay".to_string());
    map.insert((StringKey::Exact("b".to_string()), StringKey::Exact("x".to_string())), "bx".to_string());
    let query = (StringKey::Exact("a".to_string()), StringKey::Any);
    // Would be easier to do (Included(query), Included(query)), but that only returns one item
    let iter = map.range((Included(query.clone()), Unbounded)); 
    for (key, value) in iter {
        if key > &query {
            return;
        }
        println!("{:?} {:?}", key, value);
    }
}

Which prints

(Exact("a"), Exact("x")) "ax"
(Exact("a"), Exact("y")) "ay"

So that works, but the problem is just that I'm violating the Ord/Eq rules here (things like StringKey::Exact("a") == StringKey::Any == StringKey::Exact("b")), which doesn't seem like the right way to go about it. It's also a problem since I'd like to store the StringKey in HashMaps in other parts of the code but that doesn't really work when Eq is implemented like this.

So, again; is there a better way to do this?

Thankful for any help!

rust
2 Answers

Your current implementation breaks the contract of Ord. This means that the only thing you reliably say about calling map.range() is that it will not break memory safety. The values that are returned by the function, or whether it returns at all instead of panicking, could change in different versions of rust and cannot be relied on. In short, you really want your Ord implementation to be a total order.

The best way to solve your problem (which also plays nicely into the interface of range) is to have two special values - Min and Max. Then a query for ("a", *) would effectively translate to ("a", Min) .. ("a", Max).

use std::{collections::BTreeMap, cmp::Ordering};
use std::ops::Bound::{Included, Unbounded};

#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum StringKey {
    Min,
    Exact(String),
    Max,
}

fn main() {
    let mut map = BTreeMap::new();
    map.insert((StringKey::Exact("a".to_string()), StringKey::Exact("x".to_string())), "ax".to_string());
    map.insert((StringKey::Exact("a".to_string()), StringKey::Exact("y".to_string())), "ay".to_string());
    map.insert((StringKey::Exact("b".to_string()), StringKey::Exact("x".to_string())), "bx".to_string());
    
    let query_min = (StringKey::Exact("a".to_string()), StringKey::Min);
    let query_max = (StringKey::Exact("a".to_string()), StringKey::Max);
    let iter = map.range(query_min..query_max); 
    for (key, value) in iter {
        println!("{:?} {:?}", key, value);
    }
}
user_53041240

For filtering it may be enough to match the tuple pattern.

As a simple example:

use std::collections::BTreeMap;

fn main() {
    let mut map = BTreeMap::new();
    map.insert(("a", "x"), "ax");
    map.insert(("a", "y"), "ay");
    map.insert(("b", "y"), "by");
    for (k, v) in &map {
        if matches!(k, ("a", _)) {
            println!("{:?} -> {}", k, v);
        }
    }
}

Playground

The solution with range is more elegant (probably more performant too). But it is also more convoluted. So it will depend on your needs.

user_16951720
Related