Skip to main content

radiate_expr/
builder.rs

1use crate::{
2    Expr, ExprNode,
3    ops::{RollupOp, ScheduleOp},
4};
5use crate::{
6    When,
7    ops::{BinaryOp, SelectOp, TrinaryOp, UnaryOp},
8};
9use radiate_utils::{AnyValue, DataType, SmallStr, WindowBuffer};
10use std::ops::{Add, Div, Mul, Neg, Not, Sub};
11
12impl Expr {
13    pub fn identity() -> Expr {
14        Expr::from(SelectOp::Identity)
15    }
16
17    pub fn lit(value: impl Into<AnyValue<'static>>) -> Expr {
18        Expr::from(value.into())
19    }
20
21    pub fn range(sel: impl Into<std::ops::Range<usize>>) -> Expr {
22        let range = sel.into();
23        Expr::from(SelectOp::Range(range.start, range.end))
24    }
25
26    pub fn select(name: impl Into<SmallStr>) -> Expr {
27        Expr::from(SelectOp::Field(name.into()))
28    }
29
30    pub fn warmup(period: usize) -> When {
31        When::new(Expr::new(ExprNode::Schedule(ScheduleOp::Warmup {
32            period,
33            current: 0,
34        })))
35    }
36
37    pub fn when(cond: impl Into<Expr>) -> When {
38        When::new(cond.into())
39    }
40
41    pub fn every(interval: usize) -> When {
42        When::new(Expr::new(ExprNode::Schedule(ScheduleOp::Interval {
43            count: 0,
44            limit: interval,
45        })))
46    }
47
48    pub fn throttle(duration: std::time::Duration) -> When {
49        When::new(Expr::new(ExprNode::Schedule(ScheduleOp::Duration {
50            last: None,
51            interval: duration,
52        })))
53    }
54
55    pub fn time(self) -> Expr {
56        self.cast(DataType::Duration)
57    }
58
59    pub fn value(self) -> Expr {
60        self.cast(DataType::Float32)
61    }
62
63    pub fn debug(self) -> Expr {
64        self.unary(UnaryOp::Debug)
65    }
66
67    pub fn attr(self, attr: impl Into<SmallStr>) -> Expr {
68        match self.node {
69            ExprNode::Selector(selector) => Expr::from(SelectOp::Nested {
70                parent: Box::new(selector),
71                child: Box::new(SelectOp::Field(attr.into())),
72            }),
73            _ => self,
74        }
75    }
76
77    pub fn rolling(self, window_size: usize) -> Expr {
78        Expr::new(ExprNode::Rolling {
79            child: Box::new(self),
80            buffer: WindowBuffer::with_capacity(window_size),
81        })
82    }
83
84    pub fn coalesce(self, rhs: impl Into<Expr>) -> Expr {
85        self.binary(rhs.into(), BinaryOp::Coalesce)
86    }
87
88    pub fn first(self) -> Expr {
89        self.reducer(RollupOp::First)
90    }
91
92    pub fn last(self) -> Expr {
93        self.reducer(RollupOp::Last)
94    }
95
96    pub fn sum(self) -> Expr {
97        self.reducer(RollupOp::Sum)
98    }
99
100    pub fn mean(self) -> Expr {
101        self.reducer(RollupOp::Mean)
102    }
103
104    pub fn stddev(self) -> Expr {
105        self.reducer(RollupOp::StdDev)
106    }
107
108    pub fn min(self) -> Expr {
109        self.reducer(RollupOp::Min)
110    }
111
112    pub fn max(self) -> Expr {
113        self.reducer(RollupOp::Max)
114    }
115
116    pub fn var(self) -> Expr {
117        self.reducer(RollupOp::Var)
118    }
119
120    pub fn skew(self) -> Expr {
121        self.reducer(RollupOp::Skew)
122    }
123
124    pub fn count(self) -> Expr {
125        self.reducer(RollupOp::Count)
126    }
127
128    pub fn slope(self) -> Expr {
129        self.reducer(RollupOp::Slope)
130    }
131
132    pub fn unique(self) -> Expr {
133        self.reducer(RollupOp::Unique)
134    }
135
136    pub fn pow(self, exp: impl Into<Expr>) -> Expr {
137        self.binary(exp.into(), BinaryOp::Pow)
138    }
139
140    pub fn lt(self, rhs: impl Into<Expr>) -> Expr {
141        self.binary(rhs.into(), BinaryOp::Lt)
142    }
143
144    pub fn lte(self, rhs: impl Into<Expr>) -> Expr {
145        self.binary(rhs.into(), BinaryOp::Lte)
146    }
147
148    pub fn gt(self, rhs: impl Into<Expr>) -> Expr {
149        self.binary(rhs.into(), BinaryOp::Gt)
150    }
151
152    pub fn gte(self, rhs: impl Into<Expr>) -> Expr {
153        self.binary(rhs.into(), BinaryOp::Gte)
154    }
155
156    pub fn eq(self, rhs: impl Into<Expr>) -> Expr {
157        self.binary(rhs.into(), BinaryOp::Eq)
158    }
159
160    pub fn ne(self, rhs: impl Into<Expr>) -> Expr {
161        self.binary(rhs.into(), BinaryOp::Ne)
162    }
163
164    pub fn between(self, low: impl Into<Expr>, high: impl Into<Expr>) -> Expr {
165        let low = low.into();
166        let high = high.into();
167        self.clone().gte(low).and(self.lte(high))
168    }
169
170    pub fn and(self, rhs: impl Into<Expr>) -> Expr {
171        self.binary(rhs.into(), BinaryOp::And)
172    }
173
174    pub fn or(self, rhs: impl Into<Expr>) -> Expr {
175        self.binary(rhs.into(), BinaryOp::Or)
176    }
177
178    #[allow(clippy::should_implement_trait)]
179    pub fn not(self) -> Expr {
180        self.unary(UnaryOp::Not)
181    }
182
183    #[allow(clippy::should_implement_trait)]
184    pub fn neg(self) -> Expr {
185        self.unary(UnaryOp::Neg)
186    }
187
188    pub fn abs(self) -> Expr {
189        self.unary(UnaryOp::Abs)
190    }
191
192    #[allow(clippy::should_implement_trait)]
193    pub fn add(self, rhs: impl Into<Expr>) -> Expr {
194        self.binary(rhs.into(), BinaryOp::Add)
195    }
196
197    #[allow(clippy::should_implement_trait)]
198    pub fn sub(self, rhs: impl Into<Expr>) -> Expr {
199        self.binary(rhs.into(), BinaryOp::Sub)
200    }
201
202    #[allow(clippy::should_implement_trait)]
203    pub fn mul(self, rhs: impl Into<Expr>) -> Expr {
204        self.binary(rhs.into(), BinaryOp::Mul)
205    }
206
207    #[allow(clippy::should_implement_trait)]
208    pub fn div(self, rhs: impl Into<Expr>) -> Expr {
209        self.binary(rhs.into(), BinaryOp::Div)
210    }
211
212    pub fn clamp(self, min: impl Into<Expr>, max: impl Into<Expr>) -> Expr {
213        self.trinary(min.into(), max.into(), TrinaryOp::Clamp)
214    }
215
216    pub fn or_else(self, rhs: impl Into<Expr>) -> Expr {
217        self.binary(rhs.into(), BinaryOp::Coalesce)
218    }
219
220    pub fn min_with(self, rhs: impl Into<Expr>) -> Expr {
221        self.binary(rhs.into(), BinaryOp::Min)
222    }
223    pub fn max_with(self, rhs: impl Into<Expr>) -> Expr {
224        self.binary(rhs.into(), BinaryOp::Max)
225    }
226
227    pub fn quantile(self, q: f32) -> Expr {
228        self.reducer(RollupOp::Quantile(q))
229    }
230
231    pub fn stagnation(self, epsilon: f32) -> Expr {
232        self.unary(UnaryOp::Stagnation {
233            epsilon,
234            last_value: None,
235            count: 0,
236        })
237    }
238
239    pub fn cast(self, to: DataType) -> Expr {
240        self.unary(UnaryOp::Cast(to))
241    }
242
243    /// Relative error from a target: `(self - target) / target`. Fuses into
244    /// a single Affine node. `target == 0` produces a degenerate expression
245    /// (division by zero shows up as a NaN/Inf at eval time, then propagates
246    /// to the outer Clamp).
247    pub fn error(self, target: f32) -> Expr {
248        // (x - target) / target == x * (1/target) + (-1)
249        self.binary(Expr::from(1.0 / target), BinaryOp::Mul)
250            .add(Expr::from(-1.0))
251            .compile()
252    }
253
254    fn unary(self, op: UnaryOp) -> Expr {
255        Expr::new(ExprNode::Unary {
256            child: Box::new(self),
257            op,
258        })
259    }
260
261    fn binary(self, rhs: Expr, op: BinaryOp) -> Expr {
262        Expr::new(ExprNode::Binary {
263            lhs: Box::new(self),
264            rhs: Box::new(rhs),
265            op,
266        })
267    }
268
269    fn trinary(self, second: Expr, third: Expr, op: TrinaryOp) -> Expr {
270        Expr::new(ExprNode::Trinary {
271            first: Box::new(self),
272            second: Box::new(second),
273            third: Box::new(third),
274            op,
275        })
276    }
277
278    fn reducer(self, rollup: RollupOp) -> Expr {
279        Expr::new(ExprNode::Reduce {
280            child: Box::new(self),
281            rollup,
282        })
283    }
284}
285
286macro_rules! impl_from_literal {
287    ($($ty:ty => $variant:ident),*) => {
288        $(
289            impl From<$ty> for Expr {
290                fn from(value: $ty) -> Self {
291                    use crate::ExprNode;
292                    Expr::new(ExprNode::Literal(value.into()))
293                }
294            }
295        )*
296    };
297}
298
299impl_from_literal!(
300    u8 => UInt8,
301    u16 => UInt16,
302    u32 => UInt32,
303    u64 => UInt64,
304    u128 => UInt128,
305
306    i8 => Int8,
307    i16 => Int16,
308    i32 => Int32,
309    i64 => Int64,
310    i128 => Int128,
311
312    f32 => Float32,
313    f64 => Float64,
314
315    bool => Bool,
316    char => Char,
317    String => Str,
318
319    usize => Usize
320);
321
322impl<T> Add<T> for Expr
323where
324    T: Into<Expr>,
325{
326    type Output = Expr;
327    fn add(self, rhs: T) -> Expr {
328        Expr::new(ExprNode::Binary {
329            lhs: Box::new(self),
330            rhs: Box::new(rhs.into()),
331            op: BinaryOp::Add,
332        })
333    }
334}
335
336impl<T> Sub<T> for Expr
337where
338    T: Into<Expr>,
339{
340    type Output = Expr;
341    fn sub(self, rhs: T) -> Expr {
342        Expr::new(ExprNode::Binary {
343            lhs: Box::new(self),
344            rhs: Box::new(rhs.into()),
345            op: BinaryOp::Sub,
346        })
347    }
348}
349
350impl<T> Mul<T> for Expr
351where
352    T: Into<Expr>,
353{
354    type Output = Expr;
355    fn mul(self, rhs: T) -> Expr {
356        Expr::new(ExprNode::Binary {
357            lhs: Box::new(self),
358            rhs: Box::new(rhs.into()),
359            op: BinaryOp::Mul,
360        })
361    }
362}
363
364impl<T> Div<T> for Expr
365where
366    T: Into<Expr>,
367{
368    type Output = Expr;
369    fn div(self, rhs: T) -> Expr {
370        Expr::new(ExprNode::Binary {
371            lhs: Box::new(self),
372            rhs: Box::new(rhs.into()),
373            op: BinaryOp::Div,
374        })
375    }
376}
377
378impl Neg for Expr {
379    type Output = Expr;
380    fn neg(self) -> Expr {
381        Expr::new(ExprNode::Unary {
382            child: Box::new(self),
383            op: UnaryOp::Neg,
384        })
385    }
386}
387
388impl Not for Expr {
389    type Output = Expr;
390    fn not(self) -> Expr {
391        Expr::new(ExprNode::Unary {
392            child: Box::new(self),
393            op: UnaryOp::Not,
394        })
395    }
396}