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

kevy_store/
zset_algebra.rs
