mz_timely_util/
temporal.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Utilities to efficiently store future updates.
17
18use std::collections::BTreeMap;
19
20use mz_ore::cast::CastFrom;
21use timely::progress::Timestamp;
22use timely::progress::frontier::AntichainRef;
23
24/// Timestamp extension for timestamps that can advance by `2^exponent`.
25///
26/// Most likely, this is only relevant for totally ordered timestamps.
27pub trait BucketTimestamp: Timestamp {
28    /// The number of bits in the timestamp.
29    const DOMAIN: usize = size_of::<Self>() * 8;
30    /// Advance this timestamp by `2^exponent`. Returns `None` if the
31    /// timestamp would overflow.
32    fn advance_by_power_of_two(&self, exponent: u32) -> Option<Self>;
33}
34
35/// A type that can be split into two parts based on a timestamp.
36pub trait Bucket: Sized {
37    /// The timestamp type associated with this storage.
38    type Timestamp: BucketTimestamp;
39    /// Split self in two, based on the timestamp. The result is a pair of self, where the first
40    /// element contains all data with a timestamp strictly less than `timestamp`, and the second
41    /// all other data.
42    fn split(self, timestamp: &Self::Timestamp, fuel: &mut i64) -> (Self, Self);
43}
44
45/// A sorted list of buckets, representing data bucketed by timestamp.
46///
47/// Bucket chains support three main APIs: finding buckets for a given timestamp, peeling
48/// off buckets up to a frontier, and restoring the chain property. All operations aim to be
49/// amortized logarithmic in the number of outstanding timestamps.
50///
51/// We achieve this by storing buckets of increasing size for timestamps that are further out
52/// in the future. At the same time, in a well-formed chain, adjacent buckets span a time range at
53/// most factor 4 different. Factor 4 means that we can skip at most one bucket size for adjacent
54/// buckets, limiting the amount of work we need to do when peeling off buckets or restoring the
55/// chain property.
56///
57/// A bucket chain is well-formed if all buckets are within two bits of each other, with an imaginary
58/// bucket of -2 bits at the start. A chain does not need to be well-formed at all times, and supports
59/// peeling and finding even if not well-formed. However, `peel` might need to split more buckets to
60/// extract the desired data.
61///
62/// The `restore` method can be used to restore the chain property. It needs to be called repeatedly
63/// with a positive amount of fuel while the remaining fuel after the call is non-positive. This
64/// allows the caller to control the amount of work done in a single call and interleave it with
65/// other work.
66#[derive(Debug)]
67pub struct BucketChain<S: Bucket> {
68    content: BTreeMap<S::Timestamp, (u32, S)>,
69}
70
71impl<S: Bucket> BucketChain<S> {
72    /// Construct a new bucket chain. Spans the whole time domain.
73    #[inline]
74    pub fn new(storage: S) -> Self {
75        let bits = S::Timestamp::DOMAIN.try_into().expect("Must fit");
76        Self {
77            // The initial bucket starts at the minimum timestamp and spans the whole domain.
78            content: BTreeMap::from([(Timestamp::minimum(), (bits, storage))]),
79        }
80    }
81
82    /// Find the time range for the bucket that contains data for time `timestamp`.
83    /// Returns a range of the lower (inclusive) and upper (exclusive) time bound of the bucket,
84    /// or `None` if there is no bucket for the requested time. Only times that haven't
85    /// been peeled can still be found.
86    ///
87    /// The bounds are only valid until the next call to `peel` or `restore`.
88    #[inline]
89    pub fn range_of(&self, timestamp: &S::Timestamp) -> Option<std::ops::Range<S::Timestamp>> {
90        let (time, (bits, _)) = self.content.range(..=timestamp).next_back()?;
91        let top = time
92            .advance_by_power_of_two(bits.saturating_sub(1))
93            .expect("must exist");
94        Some(time.clone()..top)
95    }
96
97    /// Find the bucket that contains data for time `timestamp`. Returns a reference to the bucket,
98    /// or `None` if there is no bucket for the requested time.
99    ///
100    /// Only times that haven't been peeled can still be found.
101    #[inline]
102    pub fn find(&self, timestamp: &S::Timestamp) -> Option<&S> {
103        self.content
104            .range(..=timestamp)
105            .next_back()
106            .map(|(_, (_, storage))| storage)
107    }
108
109    /// Find the bucket that contains data for time `timestamp`. Returns a mutable reference to
110    /// the bucket, or `None` if there is no bucket for the requested time.
111    ///
112    /// Only times that haven't been peeled can still be found.
113    #[inline]
114    pub fn find_mut(&mut self, timestamp: &S::Timestamp) -> Option<&mut S> {
115        self.content
116            .range_mut(..=timestamp)
117            .next_back()
118            .map(|(_, (_, storage))| storage)
119    }
120
121    /// Peel off all data up to `frontier`, where the returned buckets contain all
122    /// data strictly less than the frontier.
123    #[inline]
124    pub fn peel(&mut self, frontier: AntichainRef<S::Timestamp>) -> Vec<S> {
125        let mut peeled = vec![];
126        // While there are buckets, and the frontier is not less than the lowest offset, peel off
127        while let Some(min_entry) = self.content.first_entry()
128            && !frontier.less_equal(min_entry.key())
129        {
130            let (offset, (bits, storage)) = self.content.pop_first().expect("must exist");
131            let upper = offset.advance_by_power_of_two(bits);
132
133            // Split the bucket if it spans the frontier.
134            if upper.is_none() && !frontier.is_empty()
135                || upper.is_some() && frontier.less_than(&upper.unwrap())
136            {
137                // We need to split the bucket, no matter how much fuel we have.
138                self.split_and_insert(&mut 0, bits, offset, storage);
139            } else {
140                // Bucket is ready to go.
141                peeled.push(storage);
142            }
143        }
144        peeled
145    }
146
147    /// Restore the chain property by splitting buckets as necessary.
148    ///
149    /// The chain is well-formed if all buckets are within two bits of each other, with an imaginary
150    /// bucket of -2 bits just before the smallest bucket.
151    #[inline]
152    pub fn restore(&mut self, fuel: &mut i64) {
153        // We could write this in terms of a cursor API, but it's not stable yet. Instead, we
154        // allocate a new map and move elements over.
155        let mut new = BTreeMap::default();
156        let mut last_bits = -2;
157        while *fuel > 0
158            && let Some((time, (bits, storage))) = self.content.pop_first()
159        {
160            // Insert bucket if the size is correct.
161            if isize::cast_from(bits) <= last_bits + 2 {
162                new.insert(time, (bits, storage));
163                last_bits = isize::cast_from(bits);
164            } else {
165                // Otherwise, we need to split it.
166                self.split_and_insert(fuel, bits, time, storage);
167            }
168        }
169        // Move remaining elements if we ran out of fuel.
170        new.append(&mut self.content);
171        self.content = new;
172    }
173
174    /// Returns `true` if the chain is empty. This means there are no outstanding times left.
175    #[inline(always)]
176    pub fn is_empty(&self) -> bool {
177        self.content.is_empty()
178    }
179
180    /// The number of buckets in the chain.
181    #[inline(always)]
182    pub fn len(&self) -> usize {
183        self.content.len()
184    }
185
186    /// Split the bucket specified by `(bits, offset, storage)` and insert the new buckets.
187    /// Updates `fuel`.
188    ///
189    /// Panics if the bucket cannot be split, i.e, it covers 0 bits.
190    #[inline(always)]
191    fn split_and_insert(&mut self, fuel: &mut i64, bits: u32, offset: S::Timestamp, storage: S) {
192        let bits = bits - 1;
193        let midpoint = offset.advance_by_power_of_two(bits).expect("must exist");
194        let (bot, top) = storage.split(&midpoint, fuel);
195        self.content.insert(offset, (bits, bot));
196        self.content.insert(midpoint, (bits, top));
197    }
198}
199
200#[cfg(test)]
201mod tests {
202    use super::*;
203
204    impl BucketTimestamp for u8 {
205        fn advance_by_power_of_two(&self, bits: u32) -> Option<Self> {
206            self.checked_add(1_u8.checked_shl(bits)?)
207        }
208    }
209
210    impl BucketTimestamp for u64 {
211        fn advance_by_power_of_two(&self, bits: u32) -> Option<Self> {
212            self.checked_add(1_u64.checked_shl(bits)?)
213        }
214    }
215
216    struct TestStorage<T> {
217        inner: Vec<T>,
218    }
219
220    impl<T: std::fmt::Debug> std::fmt::Debug for TestStorage<T> {
221        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
222            self.inner.fmt(f)
223        }
224    }
225
226    impl<T: BucketTimestamp> Bucket for TestStorage<T> {
227        type Timestamp = T;
228        fn split(self, timestamp: &T, fuel: &mut i64) -> (Self, Self) {
229            *fuel = fuel.saturating_sub(self.inner.len().try_into().expect("must fit"));
230            let (left, right) = self.inner.into_iter().partition(|d| *d < *timestamp);
231            (Self { inner: left }, Self { inner: right })
232        }
233    }
234
235    fn collect_and_sort<T: BucketTimestamp>(peeled: Vec<TestStorage<T>>) -> Vec<T> {
236        let mut collected: Vec<_> = peeled
237            .iter()
238            .flat_map(|b| b.inner.iter().cloned())
239            .collect();
240        collected.sort();
241        collected
242    }
243
244    #[mz_ore::test]
245    fn test_bucket_chain_empty_peel_all() {
246        let mut chain = BucketChain::new(TestStorage::<u8> { inner: vec![] });
247        let mut fuel = 1000;
248        chain.restore(&mut fuel);
249        assert!(fuel > 0);
250        let peeled = chain.peel(AntichainRef::new(&[]));
251        assert!(collect_and_sort(peeled).is_empty());
252        assert!(chain.is_empty());
253    }
254
255    #[mz_ore::test]
256    fn test_bucket_chain_u8() {
257        let mut chain = BucketChain::new(TestStorage::<u8> {
258            inner: (0..=255).collect(),
259        });
260        let mut fuel = -1;
261        while fuel <= 0 {
262            fuel = 100;
263            chain.restore(&mut fuel);
264        }
265        let peeled = chain.peel(AntichainRef::new(&[1]));
266        assert_eq!(peeled.len(), 1);
267        assert_eq!(peeled[0].inner[0], 0);
268        assert!(collect_and_sort(peeled).into_iter().eq(0..1));
269        let mut fuel = 1000;
270        chain.restore(&mut fuel);
271        assert!(fuel > 0);
272        let peeled = chain.peel(AntichainRef::new(&[63]));
273        let mut fuel = 1000;
274        chain.restore(&mut fuel);
275        assert!(fuel > 0);
276        assert!(collect_and_sort(peeled).into_iter().eq(1..63));
277        let peeled = chain.peel(AntichainRef::new(&[65]));
278        let mut fuel = 1000;
279        chain.restore(&mut fuel);
280        assert!(fuel > 0);
281        assert!(collect_and_sort(peeled).into_iter().eq(63..65));
282        let peeled = chain.peel(AntichainRef::new(&[]));
283        let mut fuel = 1000;
284        chain.restore(&mut fuel);
285        assert!(fuel > 0);
286        assert!(collect_and_sort(peeled).into_iter().eq(65..=255));
287    }
288
289    /// Test a chain with 10M disjoint elements.
290    #[mz_ore::test]
291    #[cfg_attr(miri, ignore)] // slow
292    fn test_bucket_10m() {
293        let limit = 10_000_000;
294
295        let mut chain = BucketChain::new(TestStorage::<u64> { inner: Vec::new() });
296        let mut fuel = 1000;
297        chain.restore(&mut fuel);
298        assert!(fuel > 0);
299
300        let now = 1739276664_u64;
301
302        let peeled = chain.peel(AntichainRef::new(&[now]));
303        let mut fuel = 1000;
304        chain.restore(&mut fuel);
305        assert!(fuel > 0);
306        let peeled = collect_and_sort(peeled);
307        assert!(peeled.is_empty());
308
309        for i in now..now + limit {
310            chain.find_mut(&i).expect("must exist").inner.push(i);
311        }
312
313        let mut offset = now;
314        let step = 1000;
315        while offset < now + limit {
316            let peeled = chain.peel(AntichainRef::new(&[offset + step]));
317            assert!(
318                collect_and_sort(peeled)
319                    .into_iter()
320                    .eq(offset..offset + step)
321            );
322            offset += step;
323            let mut fuel = 1000;
324            chain.restore(&mut fuel);
325        }
326    }
327}
mz_timely_util/temporal.rs

mz_timely_util/
temporal.rs
