1use std::collections::HashMap;
24
25use serde::{Deserialize, Serialize};
26use ts_rs::TS;
27
28use crate::proto;
29
30pub type WindowAggregate = String;
32
33#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Serialize)]
34#[serde(rename_all = "snake_case")]
35pub enum WindowFrame {
36 Rows(u32),
37 Range(f64),
38 Cumulative,
39}
40
41#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize, TS)]
45#[serde(rename_all = "snake_case")]
46pub enum WindowSortDir {
47 #[default]
48 Asc,
49 Desc,
50}
51
52impl std::fmt::Display for WindowSortDir {
53 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54 f.write_str(match self {
55 Self::Asc => "asc",
56 Self::Desc => "desc",
57 })
58 }
59}
60
61#[derive(Clone, Debug, Deserialize, PartialEq, Serialize, TS)]
65pub struct WindowSort(pub String, pub WindowSortDir);
66
67#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize, TS)]
72pub struct Windows(#[ts(as = "HashMap<String, RawWindowSpec>")] pub HashMap<String, WindowSpec>);
73
74impl std::ops::Deref for Windows {
75 type Target = HashMap<String, WindowSpec>;
76
77 fn deref(&self) -> &Self::Target {
78 &self.0
79 }
80}
81
82impl std::ops::DerefMut for Windows {
83 fn deref_mut(&mut self) -> &mut Self::Target {
84 &mut self.0
85 }
86}
87
88#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
91#[serde(try_from = "RawWindowSpec", into = "RawWindowSpec")]
92pub struct WindowSpec {
93 pub column: String,
94 pub aggregate: WindowAggregate,
95 pub partition_by: Vec<String>,
96 pub order_by: Option<WindowSort>,
97 pub frame: Option<WindowFrame>,
98 pub offset: Option<u32>,
99 pub alpha: Option<f64>,
100}
101
102#[derive(Clone, Debug, Deserialize, Serialize, TS)]
104#[serde(deny_unknown_fields)]
105#[ts(rename = "WindowSpec")]
106struct RawWindowSpec {
107 column: String,
110
111 aggregate: WindowAggregate,
112
113 #[serde(default)]
116 #[serde(skip_serializing_if = "Vec::is_empty")]
117 partition_by: Vec<String>,
118
119 #[serde(default)]
123 #[serde(skip_serializing_if = "Option::is_none")]
124 order_by: Option<WindowSort>,
125
126 #[serde(default)]
129 #[serde(skip_serializing_if = "Option::is_none")]
130 rows: Option<u32>,
131
132 #[serde(default)]
136 #[serde(skip_serializing_if = "Option::is_none")]
137 range: Option<f64>,
138
139 #[serde(default)]
142 #[serde(skip_serializing_if = "Option::is_none")]
143 cumulative: Option<bool>,
144
145 #[serde(default)]
147 #[serde(skip_serializing_if = "Option::is_none")]
148 offset: Option<u32>,
149
150 #[serde(default)]
152 #[serde(skip_serializing_if = "Option::is_none")]
153 alpha: Option<f64>,
154}
155
156impl From<WindowSpec> for RawWindowSpec {
157 fn from(value: WindowSpec) -> Self {
158 let (rows, range, cumulative) = match value.frame {
159 Some(WindowFrame::Rows(n)) => (Some(n), None, None),
160 Some(WindowFrame::Range(x)) => (None, Some(x), None),
161 Some(WindowFrame::Cumulative) => (None, None, Some(true)),
162 None => (None, None, None),
163 };
164
165 RawWindowSpec {
166 column: value.column,
167 aggregate: value.aggregate,
168 partition_by: value.partition_by,
169 order_by: value.order_by,
170 rows,
171 range,
172 cumulative,
173 offset: value.offset,
174 alpha: value.alpha,
175 }
176 }
177}
178
179impl TryFrom<RawWindowSpec> for WindowSpec {
180 type Error = String;
181
182 fn try_from(value: RawWindowSpec) -> Result<Self, Self::Error> {
183 let frame = match (value.rows, value.range, value.cumulative) {
184 (None, None, None) => None,
185 (Some(n), None, None) => Some(WindowFrame::Rows(n)),
186 (None, Some(x), None) => Some(WindowFrame::Range(x)),
187 (None, None, Some(true)) => Some(WindowFrame::Cumulative),
188 (None, None, Some(false)) => {
189 return Err("`cumulative` must be `true` when present".to_string());
190 },
191 _ => {
192 return Err("`rows`, `range` and `cumulative` are mutually exclusive".to_string());
193 },
194 };
195
196 Ok(WindowSpec {
197 column: value.column,
198 aggregate: value.aggregate,
199 partition_by: value.partition_by,
200 order_by: value.order_by,
201 frame,
202 offset: value.offset,
203 alpha: value.alpha,
204 })
205 }
206}
207
208impl From<WindowFrame> for proto::window_spec::Frame {
209 fn from(value: WindowFrame) -> Self {
210 match value {
211 WindowFrame::Rows(n) => Self::Rows(n),
212 WindowFrame::Range(x) => Self::Range(x),
213 WindowFrame::Cumulative => Self::Cumulative(0),
214 }
215 }
216}
217
218impl From<proto::window_spec::Frame> for WindowFrame {
219 fn from(value: proto::window_spec::Frame) -> Self {
220 match value {
221 proto::window_spec::Frame::Rows(n) => Self::Rows(n),
222 proto::window_spec::Frame::Range(x) => Self::Range(x),
223 proto::window_spec::Frame::Cumulative(_) => Self::Cumulative,
224 }
225 }
226}
227
228impl From<WindowSort> for proto::window_spec::Order {
229 fn from(value: WindowSort) -> Self {
230 proto::window_spec::Order {
231 column: value.0,
232 desc: value.1 == WindowSortDir::Desc,
233 }
234 }
235}
236
237impl From<proto::window_spec::Order> for WindowSort {
238 fn from(value: proto::window_spec::Order) -> Self {
239 WindowSort(
240 value.column,
241 if value.desc {
242 WindowSortDir::Desc
243 } else {
244 WindowSortDir::Asc
245 },
246 )
247 }
248}
249
250impl From<WindowSpec> for proto::WindowSpec {
251 fn from(value: WindowSpec) -> Self {
252 proto::WindowSpec {
253 source: value.column,
254 op: value.aggregate,
255 partition_by: value.partition_by,
256 order_by: value.order_by.map(|x| x.into()),
257 frame: value.frame.map(|x| x.into()),
258 offset: value.offset,
259 alpha: value.alpha,
260 }
261 }
262}
263
264impl From<proto::WindowSpec> for WindowSpec {
265 fn from(value: proto::WindowSpec) -> Self {
266 WindowSpec {
267 column: value.source,
268 aggregate: value.op,
269 partition_by: value.partition_by,
270 order_by: value.order_by.map(WindowSort::from),
271 frame: value.frame.map(|x| x.into()),
272 offset: value.offset,
273 alpha: value.alpha,
274 }
275 }
276}
277
278#[cfg(test)]
279mod tests {
280 use super::*;
281
282 fn spec(frame: Option<WindowFrame>) -> WindowSpec {
283 WindowSpec {
284 column: "price".to_string(),
285 aggregate: "sum".to_string(),
286 partition_by: vec![],
287 order_by: None,
288 frame,
289 offset: None,
290 alpha: None,
291 }
292 }
293
294 #[test]
295 fn test_frame_roundtrips_flattened() {
296 for (frame, json) in [
297 (None, r#"{"column":"price","aggregate":"sum"}"#),
298 (
299 Some(WindowFrame::Rows(19)),
300 r#"{"column":"price","aggregate":"sum","rows":19}"#,
301 ),
302 (
303 Some(WindowFrame::Range(5000.0)),
304 r#"{"column":"price","aggregate":"sum","range":5000.0}"#,
305 ),
306 (
307 Some(WindowFrame::Cumulative),
308 r#"{"column":"price","aggregate":"sum","cumulative":true}"#,
309 ),
310 ] {
311 assert_eq!(serde_json::to_string(&spec(frame)).unwrap(), json);
312 assert_eq!(
313 serde_json::from_str::<WindowSpec>(json).unwrap(),
314 spec(frame)
315 );
316 }
317 }
318
319 #[test]
320 fn test_frame_rejects_invalid_combinations() {
321 for json in [
322 r#"{"column":"price","aggregate":"sum","rows":19,"range":1.0}"#,
323 r#"{"column":"price","aggregate":"sum","rows":19,"cumulative":true}"#,
324 r#"{"column":"price","aggregate":"sum","cumulative":false}"#,
325 r#"{"column":"price","aggregate":"sum","frame":"cumulative"}"#,
326 r#"{"column":"price","aggregate":"sum","rowz":19}"#,
327 ] {
328 assert!(serde_json::from_str::<WindowSpec>(json).is_err(), "{json}");
329 }
330 }
331}