kevy_store/
zset_algebra.rs

1//! zset algebra (Redis 6.2 `ZINTERSTORE` / `ZUNIONSTORE` /
2//! `ZDIFFSTORE` / `ZINTERCARD` semantics): pure combination helpers
3//! over gathered `(member, score)` lists + the store-side
4//! materialization write.
5//!
6//! The pure functions take inputs ALREADY extracted from their keys
7//! (each input = one source key's scored members, sets contributing
8//! score 1.0 per Redis) so both consumers share one semantics:
9//! the embedded facade (reads under shard locks) and the server's
10//! cross-shard gather reducer.
11
12#[cfg(not(feature = "std"))]
13use crate::nostd_prelude::*;
14use crate::{Store, StoreError};
15
16/// Scratch tables for the set-algebra passes: std's hash tables on
17/// `std`, alloc's B-trees without it (`no_std` correctness form —
18/// these are per-call temporaries, not hot state).
19#[cfg(feature = "std")]
20type ScratchMap<K, V> = std::collections::HashMap<K, V>;
21#[cfg(feature = "std")]
22type ScratchSet<T> = std::collections::HashSet<T>;
23#[cfg(not(feature = "std"))]
24type ScratchMap<K, V> = alloc::collections::BTreeMap<K, V>;
25#[cfg(not(feature = "std"))]
26type ScratchSet<T> = alloc::collections::BTreeSet<T>;
27
28/// `AGGREGATE` mode for inter/union (Redis 6.2; default `Sum`).
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub enum ZAggregate {
31    /// Weighted sum of scores (the default).
32    Sum,
33    /// Minimum weighted score.
34    Min,
35    /// Maximum weighted score.
36    Max,
37}
38
39fn agg(a: f64, b: f64, mode: ZAggregate) -> f64 {
40    match mode {
41        ZAggregate::Sum => a + b,
42        ZAggregate::Min => a.min(b),
43        ZAggregate::Max => a.max(b),
44    }
45}
46
47fn weight_of(weights: Option<&[f64]>, i: usize) -> f64 {
48    weights.map_or(1.0, |w| w.get(i).copied().unwrap_or(1.0))
49}
50
51/// `ZUNIONSTORE` combination: every member of any input, scores
52/// aggregated across the inputs it appears in (weighted).
53pub fn zunion(
54    inputs: &[Vec<(Vec<u8>, f64)>],
55    weights: Option<&[f64]>,
56    mode: ZAggregate,
57) -> Vec<(Vec<u8>, f64)> {
58    let mut acc: Vec<(Vec<u8>, f64)> = Vec::new();
59    let mut idx: ScratchMap<Vec<u8>, usize> = ScratchMap::new();
60    for (i, input) in inputs.iter().enumerate() {
61        let w = weight_of(weights, i);
62        for (m, s) in input {
63            let ws = s * w;
64            match idx.get(m) {
65                Some(&slot) => acc[slot].1 = agg(acc[slot].1, ws, mode),
66                None => {
67                    idx.insert(m.clone(), acc.len());
68                    acc.push((m.clone(), ws));
69                }
70            }
71        }
72    }
73    acc
74}
75
76/// `ZINTERSTORE` combination: members present in EVERY input.
77pub fn zinter(
78    inputs: &[Vec<(Vec<u8>, f64)>],
79    weights: Option<&[f64]>,
80    mode: ZAggregate,
81) -> Vec<(Vec<u8>, f64)> {
82    let Some((first, rest)) = inputs.split_first() else {
83        return Vec::new();
84    };
85    // Membership maps for the non-first inputs.
86    let maps: Vec<ScratchMap<&[u8], f64>> = rest
87        .iter()
88        .map(|inp| inp.iter().map(|(m, s)| (m.as_slice(), *s)).collect())
89        .collect();
90    let w0 = weight_of(weights, 0);
91    let mut out = Vec::new();
92    'member: for (m, s) in first {
93        let mut score = s * w0;
94        for (j, map) in maps.iter().enumerate() {
95            match map.get(m.as_slice()) {
96                Some(&sj) => score = agg(score, sj * weight_of(weights, j + 1), mode),
97                None => continue 'member,
98            }
99        }
100        out.push((m.clone(), score));
101    }
102    out
103}
104
105/// `ZDIFFSTORE` combination: members of the first input absent from
106/// every other input; scores from the first input (no weights /
107/// aggregate — Redis 6.2 defines none for ZDIFF).
108pub fn zdiff(inputs: &[Vec<(Vec<u8>, f64)>]) -> Vec<(Vec<u8>, f64)> {
109    let Some((first, rest)) = inputs.split_first() else {
110        return Vec::new();
111    };
112    let excluded: ScratchSet<&[u8]> = rest
113        .iter()
114        .flat_map(|inp| inp.iter().map(|(m, _)| m.as_slice()))
115        .collect();
116    first
117        .iter()
118        .filter(|(m, _)| !excluded.contains(m.as_slice()))
119        .cloned()
120        .collect()
121}
122
123/// `ZINTERCARD` (with optional `LIMIT`, 0 = unlimited): cardinality of
124/// the intersection, short-circuiting at the limit.
125pub fn zintercard(inputs: &[Vec<(Vec<u8>, f64)>], limit: usize) -> usize {
126    let Some((first, rest)) = inputs.split_first() else {
127        return 0;
128    };
129    let maps: Vec<ScratchSet<&[u8]>> = rest
130        .iter()
131        .map(|inp| inp.iter().map(|(m, _)| m.as_slice()).collect())
132        .collect();
133    let mut n = 0;
134    'member: for (m, _) in first {
135        for map in &maps {
136            if !map.contains(m.as_slice()) {
137                continue 'member;
138            }
139        }
140        n += 1;
141        if limit != 0 && n >= limit {
142            return n;
143        }
144    }
145    n
146}
147
148impl Store {
149    /// Extract one source key's scored members for the algebra ops:
150    /// zsets as-is, sets with score 1.0 (Redis semantics), absent key
151    /// = empty, any other type = `WrongType`.
152    pub fn zset_or_set_members(&mut self, key: &[u8]) -> Result<Vec<(Vec<u8>, f64)>, StoreError> {
153        use crate::Value;
154        match self.live_entry(key) {
155            None => Ok(Vec::new()),
156            Some(e) => match &e.value {
157                Value::ZSet(z) => Ok(z
158                    .by_member
159                    .iter()
160                    .map(|(m, s)| (m.to_vec(), *s))
161                    .collect()),
162                Value::SegZSet(z) => {
163                    Ok(z.ordered().map(|(m, s)| (m.to_vec(), s)).collect())
164                }
165                Value::SmallZSetInline(z) => {
166                    Ok(z.iter().map(|(m, s)| (m.to_vec(), s)).collect())
167                }
168                Value::Set(s) => Ok(s.iter().map(|m| (m.to_vec(), 1.0)).collect()),
169                Value::SegSet(s) => Ok(s.keys().map(|m| (m.to_vec(), 1.0)).collect()),
170                Value::SmallSetInline(s) => Ok(s.iter().map(|m| (m.to_vec(), 1.0)).collect()),
171                _ => Err(StoreError::WrongType),
172            },
173        }
174    }
175
176    /// Materialize an algebra result at `dst`: existing value (any
177    /// type) is dropped, result written as a zset — Redis `*STORE`
178    /// overwrite semantics. Empty result deletes `dst` (Redis drops
179    /// the destination rather than storing an empty zset). Returns
180    /// the stored cardinality.
181    pub fn zstore_result(&mut self, dst: &[u8], pairs: &[(Vec<u8>, f64)]) -> usize {
182        self.del(&[dst]);
183        if pairs.is_empty() {
184            return 0;
185        }
186        let scored: Vec<(f64, &[u8])> = pairs.iter().map(|(m, s)| (*s, m.as_slice())).collect();
187        let _ = self.zadd(dst, &scored);
188        pairs.len()
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195
196    fn zs(pairs: &[(&str, f64)]) -> Vec<(Vec<u8>, f64)> {
197        pairs.iter().map(|(m, s)| (m.as_bytes().to_vec(), *s)).collect()
198    }
199
200    #[test]
201    fn union_weights_and_aggregates() {
202        let a = zs(&[("x", 1.0), ("y", 2.0)]);
203        let b = zs(&[("y", 3.0), ("z", 4.0)]);
204        let mut u = zunion(&[a.clone(), b.clone()], None, ZAggregate::Sum);
205        u.sort_by(|x, y| x.0.cmp(&y.0));
206        assert_eq!(u, zs(&[("x", 1.0), ("y", 5.0), ("z", 4.0)]));
207        // WEIGHTS 2 3, MIN
208        let mut u = zunion(&[a, b], Some(&[2.0, 3.0]), ZAggregate::Min);
209        u.sort_by(|x, y| x.0.cmp(&y.0));
210        assert_eq!(u, zs(&[("x", 2.0), ("y", 4.0), ("z", 12.0)]));
211    }
212
213    #[test]
214    fn inter_semantics() {
215        let a = zs(&[("x", 1.0), ("y", 2.0)]);
216        let b = zs(&[("y", 3.0), ("z", 4.0)]);
217        assert_eq!(zinter(&[a.clone(), b.clone()], None, ZAggregate::Sum), zs(&[("y", 5.0)]));
218        assert_eq!(
219            zinter(&[a, b], Some(&[10.0, 1.0]), ZAggregate::Max),
220            zs(&[("y", 20.0)])
221        );
222        assert!(zinter(&[], None, ZAggregate::Sum).is_empty());
223    }
224
225    #[test]
226    fn diff_and_intercard() {
227        let a = zs(&[("x", 1.0), ("y", 2.0), ("z", 3.0)]);
228        let b = zs(&[("y", 9.0)]);
229        let mut d = zdiff(&[a.clone(), b.clone()]);
230        d.sort_by(|x, y| x.0.cmp(&y.0));
231        assert_eq!(d, zs(&[("x", 1.0), ("z", 3.0)]));
232        assert_eq!(zintercard(&[a.clone(), b.clone()], 0), 1);
233        let c = zs(&[("x", 0.0), ("y", 0.0), ("z", 0.0)]);
234        assert_eq!(zintercard(&[a.clone(), c.clone()], 0), 3);
235        assert_eq!(zintercard(&[a, c], 2), 2); // LIMIT short-circuit
236    }
237
238    #[test]
239    fn store_materialization_and_source_extraction() {
240        let mut s = Store::new();
241        s.zadd(b"z", &[(1.0, b"a".as_slice()), (2.0, b"b".as_slice())]).unwrap();
242        s.sadd(b"s", &[b"a".as_slice(), b"c".as_slice()]).unwrap();
243        let mut zm = s.zset_or_set_members(b"z").unwrap();
244        zm.sort_by(|x, y| x.0.cmp(&y.0));
245        assert_eq!(zm, zs(&[("a", 1.0), ("b", 2.0)]));
246        let mut sm = s.zset_or_set_members(b"s").unwrap();
247        sm.sort_by(|x, y| x.0.cmp(&y.0));
248        assert_eq!(sm, zs(&[("a", 1.0), ("c", 1.0)]));
249        assert!(s.zset_or_set_members(b"missing").unwrap().is_empty());
250        s.set(b"str", b"v".to_vec(), None, false, false);
251        assert!(s.zset_or_set_members(b"str").is_err());
252
253        // *STORE overwrite + empty-result-deletes semantics.
254        s.set(b"dst", b"old".to_vec(), None, false, false);
255        assert_eq!(s.zstore_result(b"dst", &zs(&[("m", 7.0)])), 1);
256        assert_eq!(s.zscore(b"dst", b"m").unwrap(), Some(7.0));
257        assert_eq!(s.zstore_result(b"dst", &[]), 0);
258        assert_eq!(s.zcard(b"dst").unwrap(), 0);
259        assert_eq!(s.exists(&[b"dst".as_slice()]), 0);
260    }
261}
kevy_store/zset_algebra.rs

kevy_store/
zset_algebra.rs
