1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
//! v2.2 — zset/set algebra `*STORE` orchestrator, step-1→step-2
//! transition (same two-hop shape as [`crate::exec_rename`]).
//!
//! Step 1 (built by `build_zalgebra_store`): scored/set gathers fan
//! out to the source keys' shards, folding into
//! [`Agg::ZStoreGather`]. When the last gather lands, `fold` routes
//! the agg here: the origin computes the combination via the shared
//! `kevy_store` pure algebra, re-arms the slot as a plain
//! `Agg::SumInt`, and ships `Op::ZStoreResult` / `Op::SetStoreResult`
//! to `dst`'s owning shard (or executes locally). Step 2's
//! `Part::Int(cardinality)` folds through the SumInt into the `:n`
//! reply — no dedicated finalize needed.
use std::collections::HashMap;
/// One source key's scored members in request order.
type ScoredInput = Vec<(Vec<u8>, f64)>;
use crate::Commands;
use crate::message::{Agg, Gathered, Inbound, Op, SmallReply, ZCombine};
use crate::shard::Shard;
impl<C: Commands> Shard<C> {
pub(crate) fn finalize_zstore_agg(&mut self, conn_id: u64, seq: u64, agg: Agg) {
let Agg::ZStoreGather {
combine,
weights,
aggregate,
dst,
keys,
got,
} = agg
else {
return;
};
// Rebuild inputs in request-key order; WRONGTYPE on any source
// aborts with the standard error (Redis behavior).
let (zset_inputs, wrongtype) = collect_scored(&keys, &got);
if wrongtype {
self.fill_zstore_slot(
conn_id,
seq,
b"-WRONGTYPE Operation against a key holding the wrong kind of value\r\n".to_vec(),
);
return;
}
let dst_shard = self.shard_of(&dst);
let op = match combine {
ZCombine::ZInter | ZCombine::ZUnion | ZCombine::ZDiff => {
let pairs = match combine {
ZCombine::ZInter => {
kevy_store::zinter(&zset_inputs, weights.as_deref(), aggregate)
}
ZCombine::ZUnion => {
kevy_store::zunion(&zset_inputs, weights.as_deref(), aggregate)
}
_ => kevy_store::zdiff(&zset_inputs),
};
Op::ZStoreResult { dst, pairs }
}
ZCombine::SInter | ZCombine::SUnion | ZCombine::SDiff => {
// Set forms ride the Members gather: inputs carry score
// 1.0 placeholders only when sources were sets; strip.
let sets: Vec<Vec<Vec<u8>>> = zset_inputs
.into_iter()
.map(|inp| inp.into_iter().map(|(m, _)| m).collect())
.collect();
let members = match combine {
ZCombine::SInter => crate::reduce::set_intersect(&sets),
ZCombine::SUnion => crate::reduce::set_union(&sets),
_ => crate::reduce::set_diff(&sets),
};
Op::SetStoreResult { dst, members }
}
};
// Re-arm the slot for step 2: one Int part sums into the reply.
if let Some(c) = self.conns.get_mut(&conn_id) {
let idx = (seq - c.next_emit) as usize;
if let Some(slot) = c.pending.get_mut(idx) {
slot.remaining = 1;
slot.agg = Agg::SumInt(0);
}
}
if dst_shard == self.id {
let part = self.exec_op(op);
self.fold(conn_id, seq, part);
} else {
self.send_to(
dst_shard,
Inbound::Request {
origin: self.id,
conn: conn_id,
seq,
op,
},
);
}
}
/// v2.6: re-arm the slot for a continuation phase and fan the new
/// argv to every shard (stateless two-phase — see exec.rs fold).
pub(crate) fn start_extension_phase(&mut self, conn_id: u64, seq: u64, argv: Vec<Vec<u8>>) {
if let Some(c) = self.conns.get_mut(&conn_id) {
let idx = (seq - c.next_emit) as usize;
if let Some(slot) = c.pending.get_mut(idx) {
slot.remaining = self.nshards as u32;
slot.agg = Agg::ExtensionGather { argv: argv.clone(), chunks: Vec::new() };
}
}
let targets: Vec<(usize, Op)> = (0..self.nshards)
.map(|s| (s, Op::Extension { argv: argv.clone() }))
.collect();
self.dispatch_targets(conn_id, seq, targets);
}
/// v2.5: complete an extension fan-out slot with the reduced reply.
pub(crate) fn fill_extension_slot(&mut self, conn_id: u64, seq: u64, reply: Vec<u8>) {
self.fill_zstore_slot(conn_id, seq, reply);
}
/// Complete the pending slot with a pre-encoded reply (parse /
/// WRONGTYPE aborts).
fn fill_zstore_slot(&mut self, conn_id: u64, seq: u64, reply: Vec<u8>) {
if let Some(c) = self.conns.get_mut(&conn_id) {
let idx = (seq - c.next_emit) as usize;
if let Some(slot) = c.pending.get_mut(idx) {
slot.remaining = 1;
slot.agg = Agg::First(None);
}
}
self.fold(conn_id, seq, crate::message::Part::Reply(SmallReply::from_vec(reply)));
}
}
/// Inputs in request order: `Scored` as-is, set `Members` at 1.0,
/// missing = empty. Second return = a WRONGTYPE was gathered.
fn collect_scored(
keys: &[Vec<u8>],
got: &HashMap<Vec<u8>, Gathered>,
) -> (Vec<ScoredInput>, bool) {
let mut inputs = Vec::with_capacity(keys.len());
for k in keys {
match got.get(k) {
Some(Gathered::Scored(p)) => inputs.push(p.clone()),
Some(Gathered::Members(m)) => {
inputs.push(m.iter().map(|v| (v.clone(), 1.0)).collect());
}
Some(Gathered::WrongType) => return (inputs, true),
_ => inputs.push(Vec::new()),
}
}
(inputs, false)
}