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
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
//! 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 = build_store_op(combine, zset_inputs, weights, aggregate, dst);
self.ship_store_op(conn_id, seq, dst_shard, op);
}
/// Step 2 of a `*STORE` orchestrator: re-arm the pending slot so one
/// `Part::Int` (the stored cardinality) folds into the `:n` reply, then
/// run the store Op on the DESTINATION's owning shard. Shared with the
/// geo `*STORE` orchestrator ([`crate::exec_geostore`]) — both compute
/// elsewhere and materialize here.
pub(crate) fn ship_store_op(
&mut self,
conn_id: u64,
seq: u64,
dst_shard: usize,
op: Op,
) {
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,
},
);
}
}
/// 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);
}
/// 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).
pub(crate) 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)
}
/// Combine the gathered inputs per `combine` and build the step-2 store
/// Op. Extracted verbatim from [`Shard::finalize_zstore_agg`] (single
/// call site, `inline(always)`) purely for the 50-LOC fn rule.
#[inline(always)]
fn build_store_op(
combine: ZCombine,
zset_inputs: Vec<Vec<(Vec<u8>, f64)>>,
weights: Option<Vec<f64>>,
aggregate: kevy_store::ZAggregate,
dst: Vec<u8>,
) -> 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 }
}
}
}