pond_rs/
query.rs

1//! Finding rows by position, by key and by a test: `at`, `bisect`,
2//! `includes_key`, `at_or_before` / `at_or_after`, `nearest`, and `find` /
3//! `some` / `every`, as in pond-ts.
4//!
5//! A key to look for is an [`EventKey`]: a point in time (a plain `i64`
6//! converts), a time range or a labelled interval. Keys compare as pond-ts
7//! compares them: by start, then end, then kind (interval, then time, then
8//! time range), then, between two intervals, by label (numbers before text,
9//! text in the same order as JavaScript's `localeCompare`).
10//!
11//! ```
12//! # use pond_rs::{Column, EventKey, TimeSeries};
13//! let s = TimeSeries::new("cpu", vec![0, 1000, 2500, 4000], vec![
14//!     Column::required("v", vec![1.0, 2.0, 3.0, 4.0]),
15//! ])?;
16//! assert_eq!(s.bisect(2000), 2); // where 2000 would go
17//! assert!(s.includes_key(2500));
18//! assert_eq!(s.at_or_before(2000).map(|r| r.begin()), Some(1000));
19//! assert_eq!(s.at_or_after(2000).map(|r| r.begin()), Some(2500));
20//! assert_eq!(s.nearest(3300).map(|r| r.begin()), Some(4000));
21//! assert_eq!(s.at(-1).map(|r| r.begin()), Some(4000));
22//! let busy = s.find(|r| Ok(r.number("v")? > Some(2.5)))?;
23//! assert_eq!(busy.map(|r| r.index()), Some(2));
24//! # Ok::<(), pond_rs::PondError>(())
25//! ```
26
27use std::cmp::Ordering;
28
29use icu_collator::CollatorBorrowed;
30
31use crate::error::{PondError, Result};
32use crate::join::root_collator;
33use crate::rows::Row;
34use crate::series::{Key, Labels, TimeSeries};
35
36/// One interval label: a number or text, as pond-ts allows. (A whole
37/// column of them is [`Labels`](crate::Labels).)
38#[derive(Debug, Clone, PartialEq)]
39pub enum Label {
40    /// A number label. Must be finite.
41    Number(f64),
42    /// A text label.
43    String(String),
44}
45
46impl From<f64> for Label {
47    fn from(v: f64) -> Self {
48        Label::Number(v)
49    }
50}
51
52/// Exact for whole numbers up to 2^53.
53impl From<i64> for Label {
54    fn from(v: i64) -> Self {
55        Label::Number(v as f64)
56    }
57}
58
59impl From<&str> for Label {
60    fn from(v: &str) -> Self {
61        Label::String(v.to_string())
62    }
63}
64
65impl From<String> for Label {
66    fn from(v: String) -> Self {
67        Label::String(v)
68    }
69}
70
71/// One key: a point in time, a time range or a labelled interval, in
72/// epoch milliseconds, as pond-ts `Time` / `TimeRange` / `Interval`. Used to
73/// look rows up ([`TimeSeries::bisect`] and friends) and given by
74/// [`Row::key`]. A plain `i64` converts to a point in time. (A series'
75/// whole key column is a [`Key`](crate::Key).)
76#[derive(Debug, Clone, PartialEq)]
77pub struct EventKey {
78    begin: i64,
79    end: i64,
80    label: Option<Label>,
81    kind: KeyKind,
82}
83
84#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
85enum KeyKind {
86    // In pond-ts's order: kind names compared with `localeCompare`.
87    Interval,
88    Time,
89    TimeRange,
90}
91
92impl EventKey {
93    /// A point in time.
94    pub fn time(at: i64) -> Self {
95        Self {
96            begin: at,
97            end: at,
98            label: None,
99            kind: KeyKind::Time,
100        }
101    }
102
103    /// A time range from `begin` to `end`. An error when `begin > end`, as
104    /// in pond-ts.
105    pub fn time_range(begin: i64, end: i64) -> Result<Self> {
106        check_order("EventKey::time_range", begin, end)?;
107        Ok(Self {
108            begin,
109            end,
110            label: None,
111            kind: KeyKind::TimeRange,
112        })
113    }
114
115    /// An interval from `begin` to `end`, labelled `label`. An error when
116    /// `begin > end`, as in pond-ts, or when a number label is not finite
117    /// (pond-rs labels are finite; pond-ts does not check).
118    pub fn interval(label: impl Into<Label>, begin: i64, end: i64) -> Result<Self> {
119        const CALL: &str = "EventKey::interval";
120        check_order(CALL, begin, end)?;
121        let label = label.into();
122        if let Label::Number(v) = label
123            && !v.is_finite()
124        {
125            return Err(PondError::InvalidData {
126                call: CALL,
127                message: format!("a label must be finite, got {v}"),
128            });
129        }
130        Ok(Self {
131            begin,
132            end,
133            label: Some(label),
134            kind: KeyKind::Interval,
135        })
136    }
137
138    /// The start, in epoch milliseconds.
139    pub fn begin(&self) -> i64 {
140        self.begin
141    }
142
143    /// The end, in epoch milliseconds (the start, for a point in time).
144    pub fn end(&self) -> i64 {
145        self.end
146    }
147
148    /// The label of an interval; `None` for other kinds.
149    pub fn label(&self) -> Option<&Label> {
150        self.label.as_ref()
151    }
152
153    /// `"time"`, `"timeRange"` or `"interval"`, as pond-ts `type()`.
154    pub fn kind(&self) -> &'static str {
155        match self.kind {
156            KeyKind::Time => "time",
157            KeyKind::TimeRange => "timeRange",
158            KeyKind::Interval => "interval",
159        }
160    }
161}
162
163impl From<i64> for EventKey {
164    fn from(at: i64) -> Self {
165        EventKey::time(at)
166    }
167}
168
169fn check_order(call: &'static str, begin: i64, end: i64) -> Result<()> {
170    if begin > end {
171        return Err(PondError::InvalidData {
172            call,
173            message: format!("start must be <= end, got {begin} > {end}"),
174        });
175    }
176    Ok(())
177}
178
179/// Two keys in pond-ts's order (`a.compare(b)`): start, then end, then kind,
180/// then (two intervals) label: numbers before text, text by the root
181/// collator as JavaScript's `localeCompare`.
182#[cfg(feature = "live")]
183#[inline]
184pub(crate) fn compare_keys(a: &EventKey, b: &EventKey) -> Ordering {
185    // Starts differ for almost every pair a live series compares.
186    if a.begin != b.begin {
187        return a.begin.cmp(&b.begin);
188    }
189    compare_same_begin(a, b)
190}
191
192/// [`compare_keys`] for two keys with the same start: the rare case, kept
193/// out of line.
194#[cfg(feature = "live")]
195#[inline(never)]
196fn compare_same_begin(a: &EventKey, b: &EventKey) -> Ordering {
197    let c = a.end.cmp(&b.end).then(a.kind.cmp(&b.kind));
198    if c != Ordering::Equal {
199        return c;
200    }
201    match (&a.label, &b.label) {
202        (Some(Label::Number(x)), Some(Label::Number(y))) => {
203            x.partial_cmp(y).unwrap_or(Ordering::Equal)
204        }
205        (Some(Label::Number(_)), Some(Label::String(_))) => Ordering::Less,
206        (Some(Label::String(_)), Some(Label::Number(_))) => Ordering::Greater,
207        (Some(Label::String(x)), Some(Label::String(y))) if x != y => root_collator().compare(x, y),
208        _ => Ordering::Equal,
209    }
210}
211
212/// pond-ts `a.equals(b)`: same kind, start and end, and for intervals the
213/// same label (`===`).
214#[cfg(feature = "live")]
215pub(crate) fn keys_equal(a: &EventKey, b: &EventKey) -> bool {
216    a.kind == b.kind && a.begin == b.begin && a.end == b.end && a.label == b.label
217}
218
219/// Compares a series' rows with one key, as pond-ts `keyAt(i).compare(key)`.
220struct KeyCompare<'a> {
221    key: &'a Key,
222    query: &'a EventKey,
223    collator: Option<CollatorBorrowed<'static>>,
224}
225
226impl<'a> KeyCompare<'a> {
227    fn new(key: &'a Key, query: &'a EventKey) -> Self {
228        Self {
229            key,
230            query,
231            collator: None,
232        }
233    }
234
235    fn kind(&self) -> KeyKind {
236        match self.key {
237            Key::Time(_) => KeyKind::Time,
238            Key::TimeRange(_) => KeyKind::TimeRange,
239            Key::Interval(_) => KeyKind::Interval,
240        }
241    }
242
243    /// Row `i` against the query.
244    fn cmp(&mut self, i: usize) -> Ordering {
245        let q = self.query;
246        let begin = self.key.begin()[i];
247        let end = self.key.end().map_or(begin, |e| e[i]);
248        let c = begin
249            .cmp(&q.begin)
250            .then(end.cmp(&q.end))
251            .then(self.kind().cmp(&q.kind));
252        if c != Ordering::Equal {
253            return c;
254        }
255        let (Key::Interval(k), Some(ql)) = (self.key, &q.label) else {
256            return c;
257        };
258        match (&k.label, ql) {
259            // Labels are finite, so this is a total order.
260            (Labels::Number(a), Label::Number(b)) => a[i].partial_cmp(b).unwrap_or(Ordering::Equal),
261            // Numbers sort before text.
262            (Labels::Number(_), Label::String(_)) => Ordering::Less,
263            (Labels::String(_), Label::Number(_)) => Ordering::Greater,
264            (Labels::String(a), Label::String(b)) => {
265                let a = a.get(i).unwrap_or("");
266                if a == b {
267                    return Ordering::Equal;
268                }
269                self.collator
270                    .get_or_insert_with(root_collator)
271                    .compare(a, b)
272            }
273        }
274    }
275
276    /// Row `i` is the query exactly, as pond-ts `keyAt(i).equals(key)`:
277    /// same kind, start and end, and for intervals the same label (`===`,
278    /// so text must match byte for byte).
279    fn equals(&self, i: usize) -> bool {
280        let q = self.query;
281        let begin = self.key.begin()[i];
282        let end = self.key.end().map_or(begin, |e| e[i]);
283        if self.kind() != q.kind || begin != q.begin || end != q.end {
284            return false;
285        }
286        match (self.key, &q.label) {
287            (Key::Interval(k), Some(ql)) => match (&k.label, ql) {
288                (Labels::Number(a), Label::Number(b)) => a[i] == *b,
289                (Labels::String(a), Label::String(b)) => a.get(i) == Some(b.as_str()),
290                _ => false,
291            },
292            _ => true,
293        }
294    }
295
296    /// The first row not before the query: pond-ts's binary search, step
297    /// for step, so rows out of label order land where pond-ts lands.
298    fn bisect(&mut self) -> usize {
299        let (mut low, mut high) = (0, self.key.len());
300        while low < high {
301            let mid = (low + high) >> 1;
302            if self.cmp(mid) == Ordering::Less {
303                low = mid + 1;
304            } else {
305                high = mid;
306            }
307        }
308        low
309    }
310}
311
312impl TimeSeries {
313    /// The row at `index`; negative positions count from the end (`-1` is
314    /// the last row), as pond-ts `at` and JavaScript's `Array.prototype.at`.
315    /// `None` past either end.
316    pub fn at(&self, index: i64) -> Option<Row<'_>> {
317        let n = self.len() as i64;
318        let i = if index < 0 { n + index } else { index };
319        if i < 0 || i >= n {
320            return None;
321        }
322        self.row(i as usize)
323    }
324
325    /// Where `key` would go: the first row whose key is not before it (so
326    /// before any rows equal to it), as pond-ts `bisect`. The length of the
327    /// series when every row is before it. O(log N).
328    pub fn bisect(&self, key: impl Into<EventKey>) -> usize {
329        let key = key.into();
330        KeyCompare::new(self.key(), &key).bisect()
331    }
332
333    /// Whether a row has exactly `key`, as pond-ts `includesKey`. Only the
334    /// row [`bisect`](TimeSeries::bisect) lands on is checked, as in
335    /// pond-ts.
336    pub fn includes_key(&self, key: impl Into<EventKey>) -> bool {
337        let key = key.into();
338        let mut c = KeyCompare::new(self.key(), &key);
339        let i = c.bisect();
340        i < self.len() && c.equals(i)
341    }
342
343    /// The row with exactly `key`, else the last row before it, as pond-ts
344    /// `atOrBefore`. `None` when every row is after it.
345    pub fn at_or_before(&self, key: impl Into<EventKey>) -> Option<Row<'_>> {
346        let key = key.into();
347        let mut c = KeyCompare::new(self.key(), &key);
348        let i = c.bisect();
349        if i < self.len() && c.equals(i) {
350            return self.row(i);
351        }
352        i.checked_sub(1).and_then(|i| self.row(i))
353    }
354
355    /// The first row not before `key` (the row
356    /// [`bisect`](TimeSeries::bisect) gives), as pond-ts `atOrAfter`. `None`
357    /// when every row is before it.
358    pub fn at_or_after(&self, key: impl Into<EventKey>) -> Option<Row<'_>> {
359        self.row(self.bisect(key))
360    }
361
362    /// The row whose start is closest to `key`'s start, as pond-ts
363    /// `nearest`; a tie goes to the earlier row. A key outside the series
364    /// gives the first or last row. `None` only for an empty series.
365    pub fn nearest(&self, key: impl Into<EventKey>) -> Option<Row<'_>> {
366        let n = self.len();
367        if n == 0 {
368            return None;
369        }
370        let key = key.into();
371        let i = KeyCompare::new(self.key(), &key).bisect();
372        if i == 0 {
373            return self.row(0);
374        }
375        if i >= n {
376            return self.row(n - 1);
377        }
378        let begin = self.key().begin();
379        let target = key.begin as i128;
380        let (before, after) = (begin[i - 1] as i128, begin[i] as i128);
381        self.row(if target - before <= after - target {
382            i - 1
383        } else {
384            i
385        })
386    }
387
388    /// The first row for which `test` returns `Ok(true)`, as pond-ts
389    /// `find`. An error from `test` stops the walk and is returned.
390    pub fn find<F>(&self, mut test: F) -> Result<Option<Row<'_>>>
391    where
392        F: FnMut(&Row<'_>) -> Result<bool>,
393    {
394        for row in self.rows() {
395            if test(&row)? {
396                return Ok(Some(row));
397            }
398        }
399        Ok(None)
400    }
401
402    /// Whether `test` returns `Ok(true)` for some row, as pond-ts `some`
403    /// (`false` for an empty series). Stops at the first match.
404    pub fn some<F>(&self, test: F) -> Result<bool>
405    where
406        F: FnMut(&Row<'_>) -> Result<bool>,
407    {
408        Ok(self.find(test)?.is_some())
409    }
410
411    /// Whether `test` returns `Ok(true)` for every row, as pond-ts `every`
412    /// (`true` for an empty series). Stops at the first row that fails.
413    pub fn every<F>(&self, mut test: F) -> Result<bool>
414    where
415        F: FnMut(&Row<'_>) -> Result<bool>,
416    {
417        for row in self.rows() {
418            if !test(&row)? {
419                return Ok(false);
420            }
421        }
422        Ok(true)
423    }
424}
425
426impl Row<'_> {
427    /// The row's key, as pond-ts `event.key()`.
428    pub fn key(&self) -> EventKey {
429        let key = self.series().key();
430        let i = self.index();
431        let (begin, end) = (self.begin(), self.end());
432        match key {
433            Key::Time(_) => EventKey::time(begin),
434            Key::TimeRange(_) => EventKey {
435                begin,
436                end,
437                label: None,
438                kind: KeyKind::TimeRange,
439            },
440            Key::Interval(k) => EventKey {
441                begin,
442                end,
443                label: Some(match &k.label {
444                    Labels::Number(v) => Label::Number(v[i]),
445                    Labels::String(s) => Label::String(s.get(i).unwrap_or("").to_string()),
446                }),
447                kind: KeyKind::Interval,
448            },
449        }
450    }
451}
452
453#[cfg(test)]
454mod tests {
455    use super::*;
456    use crate::Column;
457    use crate::series::{IntervalKeys, TimeRangeKeys};
458
459    fn times() -> TimeSeries {
460        TimeSeries::new(
461            "t",
462            vec![0, 1000, 1000, 2500],
463            vec![Column::required("v", vec![1.0, 2.0, 3.0, 4.0])],
464        )
465        .unwrap()
466    }
467
468    #[test]
469    fn bisect_lands_before_equal_rows() {
470        let s = times();
471        assert_eq!(s.bisect(-5), 0);
472        assert_eq!(s.bisect(1000), 1);
473        assert_eq!(s.bisect(1001), 3);
474        assert_eq!(s.bisect(9000), 4);
475    }
476
477    #[test]
478    fn at_counts_back_from_the_end() {
479        let s = times();
480        assert_eq!(s.at(0).map(|r| r.index()), Some(0));
481        assert_eq!(s.at(-1).map(|r| r.index()), Some(3));
482        assert_eq!(s.at(-4).map(|r| r.index()), Some(0));
483        assert!(s.at(-5).is_none());
484        assert!(s.at(4).is_none());
485        assert!(s.at(i64::MIN).is_none());
486    }
487
488    #[test]
489    fn nearest_ties_go_to_the_earlier_row() {
490        let s = times();
491        assert_eq!(s.nearest(500).map(|r| r.index()), Some(0));
492        assert_eq!(s.nearest(501).map(|r| r.index()), Some(1));
493        assert_eq!(s.nearest(-100).map(|r| r.index()), Some(0));
494        assert_eq!(s.nearest(1 << 40).map(|r| r.index()), Some(3));
495        assert!(s.slice(0..0).nearest(0).is_none());
496    }
497
498    #[test]
499    fn a_range_key_is_not_a_time_key() {
500        let s = times();
501        // Same start and end, other kind: never equal; a time sorts before
502        // a time range.
503        let r = EventKey::time_range(1000, 1000).unwrap();
504        assert!(!s.includes_key(r.clone()));
505        assert_eq!(s.bisect(r), 3);
506        assert!(EventKey::time_range(2, 1).is_err());
507        assert!(EventKey::interval(f64::NAN, 0, 1).is_err());
508    }
509
510    #[test]
511    fn interval_labels_order_numbers_before_text() {
512        let s = TimeSeries::from_parts(
513            "i".into(),
514            "interval".into(),
515            Key::Interval(IntervalKeys {
516                label: vec!["a", "B", "b"].into(),
517                begin: vec![0, 0, 0],
518                end: vec![10, 10, 10],
519            }),
520            vec![],
521        );
522        let k = |l: &str| EventKey::interval(l, 0, 10).unwrap();
523        assert_eq!(s.bisect(EventKey::interval(5.0, 0, 10).unwrap()), 0);
524        assert_eq!(s.bisect(k("b")), 1); // root collation: b < B
525        assert!(s.includes_key(k("B")));
526        assert_eq!(s.at(1).unwrap().key(), k("B"));
527        let r = TimeSeries::from_parts(
528            "r".into(),
529            "timeRange".into(),
530            Key::TimeRange(TimeRangeKeys {
531                begin: vec![0],
532                end: vec![5],
533            }),
534            vec![],
535        );
536        assert_eq!(r.at(0).unwrap().key(), EventKey::time_range(0, 5).unwrap());
537        assert_eq!(r.at(0).unwrap().key().kind(), "timeRange");
538    }
539
540    #[test]
541    fn row_key_gives_time_and_number_labelled_keys() {
542        let t = times().at(1).unwrap().key();
543        assert_eq!(
544            (t.clone(), t.kind(), t.label()),
545            (EventKey::time(1000), "time", None)
546        );
547        let n = TimeSeries::from_parts(
548            "n".into(),
549            "interval".into(),
550            Key::Interval(IntervalKeys {
551                label: vec![0.0, 5000.0].into(),
552                begin: vec![0, 5000],
553                end: vec![5000, 10000],
554            }),
555            vec![],
556        );
557        let k = n.at(-1).unwrap().key();
558        assert_eq!(k, EventKey::interval(5000.0, 5000, 10000).unwrap());
559        assert_eq!(k.label(), Some(&Label::Number(5000.0)));
560        assert!(n.includes_key(k));
561    }
562
563    #[test]
564    fn find_some_every_stop_early_and_pass_errors() {
565        let s = times();
566        let mut seen = 0;
567        let hit = s
568            .find(|r| {
569                seen += 1;
570                Ok(r.number("v")? == Some(2.0))
571            })
572            .unwrap();
573        assert_eq!((hit.map(|r| r.index()), seen), (Some(1), 2));
574        assert!(s.some(|r| Ok(r.begin() == 2500)).unwrap());
575        assert!(!s.every(|r| Ok(r.begin() > 0)).unwrap());
576        assert!(s.slice(0..0).every(|_| Ok(false)).unwrap());
577        assert!(!s.slice(0..0).some(|_| Ok(true)).unwrap());
578        assert!(s.find(|r| r.number("nope").map(|_| true)).is_err());
579    }
580}
pond_rs/query.rs

pond_rs/
query.rs
