vortex_array/
mask_future.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::future::Future;
5use std::ops::Range;
6use std::sync::Arc;
7
8use futures::FutureExt;
9use futures::TryFutureExt;
10use futures::future::BoxFuture;
11use futures::future::Shared;
12use vortex_error::SharedVortexResult;
13use vortex_error::VortexError;
14use vortex_error::VortexResult;
15use vortex_error::vortex_panic;
16use vortex_mask::Mask;
17
18/// A future that resolves to a mask.
19#[derive(Clone)]
20pub struct MaskFuture {
21    inner: Shared<BoxFuture<'static, SharedVortexResult<Mask>>>,
22    len: usize,
23}
24
25impl MaskFuture {
26    /// Create a new MaskFuture from a future that returns a mask.
27    pub fn new<F>(len: usize, fut: F) -> Self
28    where
29        F: Future<Output = VortexResult<Mask>> + Send + 'static,
30    {
31        Self {
32            inner: fut
33                .inspect(move |r| {
34                    if let Ok(mask) = r
35                        && mask.len() != len {
36                            vortex_panic!("MaskFuture created with future that returned mask of incorrect length (expected {}, got {})", len, mask.len());
37                        }
38                })
39                .map_err(Arc::new)
40                .boxed()
41                .shared(),
42            len,
43        }
44    }
45
46    /// Returns the length of the mask.
47    pub fn len(&self) -> usize {
48        self.len
49    }
50
51    /// Returns true if the mask is empty.
52    pub fn is_empty(&self) -> bool {
53        self.len == 0
54    }
55
56    /// Create a MaskFuture from a ready mask.
57    pub fn ready(mask: Mask) -> Self {
58        Self::new(mask.len(), async move { Ok(mask) })
59    }
60
61    /// Create a MaskFuture that resolves to a mask with all values set to true.
62    pub fn new_true(row_count: usize) -> Self {
63        Self::ready(Mask::new_true(row_count))
64    }
65
66    /// Create a MaskFuture that resolves to a slice of the original mask.
67    pub fn slice(&self, range: Range<usize>) -> Self {
68        // Slicing the whole mask is the identity. Cloning shares the existing future instead of
69        // allocating another boxed, shared one that would await it only to hand the mask back.
70        if range.start == 0 && range.end == self.len {
71            return self.clone();
72        }
73
74        let inner = self.inner.clone();
75        Self::new(range.len(), async move { Ok(inner.await?.slice(range)) })
76    }
77
78    pub fn inspect(
79        self,
80        f: impl FnOnce(&SharedVortexResult<Mask>) + 'static + Send + Sync,
81    ) -> Self {
82        let len = self.len;
83
84        Self {
85            inner: self.inner.inspect(f).boxed().shared(),
86            len,
87        }
88    }
89}
90
91impl Future for MaskFuture {
92    type Output = VortexResult<Mask>;
93
94    fn poll(
95        mut self: std::pin::Pin<&mut Self>,
96        cx: &mut std::task::Context<'_>,
97    ) -> std::task::Poll<Self::Output> {
98        self.inner.poll_unpin(cx).map_err(VortexError::from)
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use vortex_buffer::BitBuffer;
105
106    use super::*;
107
108    /// Slicing resolves to the same mask the equivalent [`Mask::slice`] would produce, for both
109    /// the full range (which takes the identity fast path) and a sub-range.
110    #[test]
111    fn slice_resolves_to_sliced_mask() -> VortexResult<()> {
112        futures::executor::block_on(async {
113            let mask = Mask::from_buffer(BitBuffer::from_iter([true, false, true, true, false]));
114            let fut = MaskFuture::ready(mask.clone());
115
116            let full = fut.slice(0..mask.len());
117            assert_eq!(full.len(), mask.len());
118            assert_eq!(full.await?, mask);
119
120            let partial = fut.slice(0..mask.len() - 1);
121            assert_eq!(partial.len(), mask.len() - 1);
122            assert_eq!(partial.await?, mask.slice(0..mask.len() - 1));
123            Ok(())
124        })
125    }
126}
vortex_array/mask_future.rs

vortex_array/
mask_future.rs
