1use std::collections::HashMap;
24
25use serde::{Deserialize, Serialize};
26use ts_rs::TS;
27
28use crate::proto;
29
30#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize, TS)]
31#[serde(rename_all = "snake_case")]
32pub enum WindowAggregate {
33 Sum,
34 Avg,
35 Count,
36 Min,
37 Max,
38 Stddev,
39 Var,
40 First,
41 Last,
42 Lag,
43 Lead,
44 Diff,
45 Rate,
46 Ema,
47}
48
49#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Serialize)]
50#[serde(rename_all = "snake_case")]
51pub enum WindowFrame {
52 Rows(u32),
53 Range(f64),
54 Cumulative,
55}
56
57#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize, TS)]
61#[serde(rename_all = "snake_case")]
62pub enum WindowSortDir {
63 #[default]
64 Asc,
65 Desc,
66}
67
68impl std::fmt::Display for WindowSortDir {
69 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
70 f.write_str(match self {
71 Self::Asc => "asc",
72 Self::Desc => "desc",
73 })
74 }
75}
76
77#[derive(Clone, Debug, Deserialize, PartialEq, Serialize, TS)]
81pub struct WindowSort(pub String, pub WindowSortDir);
82
83#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize, TS)]
88pub struct Windows(#[ts(as = "HashMap<String, RawWindowSpec>")] pub HashMap<String, WindowSpec>);
89
90impl std::ops::Deref for Windows {
91 type Target = HashMap<String, WindowSpec>;
92
93 fn deref(&self) -> &Self::Target {
94 &self.0
95 }
96}
97
98impl std::ops::DerefMut for Windows {
99 fn deref_mut(&mut self) -> &mut Self::Target {
100 &mut self.0
101 }
102}
103
104#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
107#[serde(try_from = "RawWindowSpec", into = "RawWindowSpec")]
108pub struct WindowSpec {
109 pub column: String,
110 pub aggregate: WindowAggregate,
111 pub partition_by: Vec<String>,
112 pub order_by: Option<WindowSort>,
113 pub frame: Option<WindowFrame>,
114 pub offset: Option<u32>,
115 pub alpha: Option<f64>,
116}
117
118#[derive(Clone, Debug, Deserialize, Serialize, TS)]
120#[serde(deny_unknown_fields)]
121#[ts(rename = "WindowSpec")]
122struct RawWindowSpec {
123 column: String,
126
127 aggregate: WindowAggregate,
128
129 #[serde(default)]
132 #[serde(skip_serializing_if = "Vec::is_empty")]
133 partition_by: Vec<String>,
134
135 #[serde(default)]
137 #[serde(skip_serializing_if = "Option::is_none")]
138 order_by: Option<WindowSort>,
139
140 #[serde(default)]
143 #[serde(skip_serializing_if = "Option::is_none")]
144 rows: Option<u32>,
145
146 #[serde(default)]
150 #[serde(skip_serializing_if = "Option::is_none")]
151 range: Option<f64>,
152
153 #[serde(default)]
156 #[serde(skip_serializing_if = "Option::is_none")]
157 cumulative: Option<bool>,
158
159 #[serde(default)]
161 #[serde(skip_serializing_if = "Option::is_none")]
162 offset: Option<u32>,
163
164 #[serde(default)]
166 #[serde(skip_serializing_if = "Option::is_none")]
167 alpha: Option<f64>,
168}
169
170impl From<WindowSpec> for RawWindowSpec {
171 fn from(value: WindowSpec) -> Self {
172 let (rows, range, cumulative) = match value.frame {
173 Some(WindowFrame::Rows(n)) => (Some(n), None, None),
174 Some(WindowFrame::Range(x)) => (None, Some(x), None),
175 Some(WindowFrame::Cumulative) => (None, None, Some(true)),
176 None => (None, None, None),
177 };
178
179 RawWindowSpec {
180 column: value.column,
181 aggregate: value.aggregate,
182 partition_by: value.partition_by,
183 order_by: value.order_by,
184 rows,
185 range,
186 cumulative,
187 offset: value.offset,
188 alpha: value.alpha,
189 }
190 }
191}
192
193impl TryFrom<RawWindowSpec> for WindowSpec {
194 type Error = String;
195
196 fn try_from(value: RawWindowSpec) -> Result<Self, Self::Error> {
197 let frame = match (value.rows, value.range, value.cumulative) {
198 (None, None, None) => None,
199 (Some(n), None, None) => Some(WindowFrame::Rows(n)),
200 (None, Some(x), None) => Some(WindowFrame::Range(x)),
201 (None, None, Some(true)) => Some(WindowFrame::Cumulative),
202 (None, None, Some(false)) => {
203 return Err("`cumulative` must be `true` when present".to_string());
204 },
205 _ => {
206 return Err("`rows`, `range` and `cumulative` are mutually exclusive".to_string());
207 },
208 };
209
210 Ok(WindowSpec {
211 column: value.column,
212 aggregate: value.aggregate,
213 partition_by: value.partition_by,
214 order_by: value.order_by,
215 frame,
216 offset: value.offset,
217 alpha: value.alpha,
218 })
219 }
220}
221
222impl From<WindowAggregate> for proto::WindowAggregate {
223 fn from(value: WindowAggregate) -> Self {
224 match value {
225 WindowAggregate::Sum => Self::Sum,
226 WindowAggregate::Avg => Self::Avg,
227 WindowAggregate::Count => Self::Count,
228 WindowAggregate::Min => Self::Min,
229 WindowAggregate::Max => Self::Max,
230 WindowAggregate::Stddev => Self::Stddev,
231 WindowAggregate::Var => Self::Var,
232 WindowAggregate::First => Self::First,
233 WindowAggregate::Last => Self::Last,
234 WindowAggregate::Lag => Self::Lag,
235 WindowAggregate::Lead => Self::Lead,
236 WindowAggregate::Diff => Self::Diff,
237 WindowAggregate::Rate => Self::Rate,
238 WindowAggregate::Ema => Self::Ema,
239 }
240 }
241}
242
243impl From<proto::WindowAggregate> for WindowAggregate {
244 fn from(value: proto::WindowAggregate) -> Self {
245 match value {
246 proto::WindowAggregate::Sum => Self::Sum,
247 proto::WindowAggregate::Avg => Self::Avg,
248 proto::WindowAggregate::Count => Self::Count,
249 proto::WindowAggregate::Min => Self::Min,
250 proto::WindowAggregate::Max => Self::Max,
251 proto::WindowAggregate::Stddev => Self::Stddev,
252 proto::WindowAggregate::Var => Self::Var,
253 proto::WindowAggregate::First => Self::First,
254 proto::WindowAggregate::Last => Self::Last,
255 proto::WindowAggregate::Lag => Self::Lag,
256 proto::WindowAggregate::Lead => Self::Lead,
257 proto::WindowAggregate::Diff => Self::Diff,
258 proto::WindowAggregate::Rate => Self::Rate,
259 proto::WindowAggregate::Ema => Self::Ema,
260 }
261 }
262}
263
264impl From<WindowFrame> for proto::window_spec::Frame {
265 fn from(value: WindowFrame) -> Self {
266 match value {
267 WindowFrame::Rows(n) => Self::Rows(n),
268 WindowFrame::Range(x) => Self::Range(x),
269 WindowFrame::Cumulative => Self::Cumulative(0),
270 }
271 }
272}
273
274impl From<proto::window_spec::Frame> for WindowFrame {
275 fn from(value: proto::window_spec::Frame) -> Self {
276 match value {
277 proto::window_spec::Frame::Rows(n) => Self::Rows(n),
278 proto::window_spec::Frame::Range(x) => Self::Range(x),
279 proto::window_spec::Frame::Cumulative(_) => Self::Cumulative,
280 }
281 }
282}
283
284impl From<WindowSort> for proto::window_spec::Order {
285 fn from(value: WindowSort) -> Self {
286 proto::window_spec::Order {
287 column: value.0,
288 desc: value.1 == WindowSortDir::Desc,
289 }
290 }
291}
292
293impl From<proto::window_spec::Order> for WindowSort {
294 fn from(value: proto::window_spec::Order) -> Self {
295 WindowSort(
296 value.column,
297 if value.desc {
298 WindowSortDir::Desc
299 } else {
300 WindowSortDir::Asc
301 },
302 )
303 }
304}
305
306impl From<WindowSpec> for proto::WindowSpec {
307 fn from(value: WindowSpec) -> Self {
308 proto::WindowSpec {
309 source: value.column,
310 op: proto::WindowAggregate::from(value.aggregate) as i32,
311 partition_by: value.partition_by,
312 order_by: value.order_by.map(|x| x.into()),
313 frame: value.frame.map(|x| x.into()),
314 offset: value.offset,
315 alpha: value.alpha,
316 }
317 }
318}
319
320impl From<proto::WindowSpec> for WindowSpec {
321 fn from(value: proto::WindowSpec) -> Self {
322 WindowSpec {
323 column: value.source,
324 aggregate: proto::WindowAggregate::try_from(value.op)
325 .unwrap_or(proto::WindowAggregate::Sum)
326 .into(),
327 partition_by: value.partition_by,
328 order_by: value.order_by.map(WindowSort::from),
329 frame: value.frame.map(|x| x.into()),
330 offset: value.offset,
331 alpha: value.alpha,
332 }
333 }
334}
335
336#[cfg(test)]
337mod tests {
338 use super::*;
339
340 fn spec(frame: Option<WindowFrame>) -> WindowSpec {
341 WindowSpec {
342 column: "price".to_string(),
343 aggregate: WindowAggregate::Sum,
344 partition_by: vec![],
345 order_by: None,
346 frame,
347 offset: None,
348 alpha: None,
349 }
350 }
351
352 #[test]
353 fn test_frame_roundtrips_flattened() {
354 for (frame, json) in [
355 (None, r#"{"column":"price","aggregate":"sum"}"#),
356 (
357 Some(WindowFrame::Rows(19)),
358 r#"{"column":"price","aggregate":"sum","rows":19}"#,
359 ),
360 (
361 Some(WindowFrame::Range(5000.0)),
362 r#"{"column":"price","aggregate":"sum","range":5000.0}"#,
363 ),
364 (
365 Some(WindowFrame::Cumulative),
366 r#"{"column":"price","aggregate":"sum","cumulative":true}"#,
367 ),
368 ] {
369 assert_eq!(serde_json::to_string(&spec(frame)).unwrap(), json);
370 assert_eq!(
371 serde_json::from_str::<WindowSpec>(json).unwrap(),
372 spec(frame)
373 );
374 }
375 }
376
377 #[test]
378 fn test_frame_rejects_invalid_combinations() {
379 for json in [
380 r#"{"column":"price","aggregate":"sum","rows":19,"range":1.0}"#,
381 r#"{"column":"price","aggregate":"sum","rows":19,"cumulative":true}"#,
382 r#"{"column":"price","aggregate":"sum","cumulative":false}"#,
383 r#"{"column":"price","aggregate":"sum","frame":"cumulative"}"#,
384 r#"{"column":"price","aggregate":"sum","rowz":19}"#,
385 ] {
386 assert!(serde_json::from_str::<WindowSpec>(json).is_err(), "{json}");
387 }
388 }
389}