diff --git a/vortex-array/src/mask_future.rs b/vortex-array/src/mask_future.rs index a46107623c1..1fa8e3ee970 100644 --- a/vortex-array/src/mask_future.rs +++ b/vortex-array/src/mask_future.rs @@ -16,12 +16,23 @@ use vortex_error::vortex_panic; use vortex_mask::Mask; /// A future that resolves to a mask. +/// +/// Masks that are known up-front (e.g. scans without a filter) are stored resolved, so awaiting, +/// cloning, and slicing them never allocates a boxed future. #[derive(Clone)] pub struct MaskFuture { - inner: Shared>>, + inner: Inner, len: usize, } +#[derive(Clone)] +enum Inner { + /// The mask is already resolved. + Ready(Mask), + /// The mask is still being computed. + Pending(Shared>>), +} + impl MaskFuture { /// Create a new MaskFuture from a future that returns a mask. pub fn new(len: usize, fut: F) -> Self @@ -29,8 +40,8 @@ impl MaskFuture { F: Future> + Send + 'static, { Self { - inner: fut - .inspect(move |r| { + inner: Inner::Pending( + fut.inspect(move |r| { if let Ok(mask) = r && mask.len() != len { vortex_panic!("MaskFuture created with future that returned mask of incorrect length (expected {}, got {})", len, mask.len()); @@ -39,6 +50,7 @@ impl MaskFuture { .map_err(Arc::new) .boxed() .shared(), + ), len, } } @@ -55,7 +67,10 @@ impl MaskFuture { /// Create a MaskFuture from a ready mask. pub fn ready(mask: Mask) -> Self { - Self::new(mask.len(), async move { Ok(mask) }) + Self { + len: mask.len(), + inner: Inner::Ready(mask), + } } /// Create a MaskFuture that resolves to a mask with all values set to true. @@ -65,25 +80,42 @@ impl MaskFuture { /// Create a MaskFuture that resolves to a slice of the original mask. pub fn slice(&self, range: Range) -> Self { - // Slicing the whole mask is the identity. Cloning shares the existing future instead of - // allocating another boxed, shared one that would await it only to hand the mask back. + // Slicing the whole mask is the identity. Cloning shares the existing state instead of + // allocating another boxed, shared future that would await it only to hand the mask back. if range.start == 0 && range.end == self.len { return self.clone(); } - let inner = self.inner.clone(); - Self::new(range.len(), async move { Ok(inner.await?.slice(range)) }) + match &self.inner { + Inner::Ready(mask) => Self::ready(mask.slice(range)), + Inner::Pending(inner) => { + let inner = inner.clone(); + Self::new(range.len(), async move { Ok(inner.await?.slice(range)) }) + } + } } + /// Observe the resolved mask. + /// + /// An already-resolved mask calls `f` immediately rather than deferring it to the first poll. pub fn inspect( self, f: impl FnOnce(&SharedVortexResult) + 'static + Send + Sync, ) -> Self { let len = self.len; - Self { - inner: self.inner.inspect(f).boxed().shared(), - len, + match self.inner { + Inner::Ready(mask) => { + f(&Ok(mask.clone())); + Self { + inner: Inner::Ready(mask), + len, + } + } + Inner::Pending(inner) => Self { + inner: Inner::Pending(inner.inspect(f).boxed().shared()), + len, + }, } } } @@ -95,12 +127,18 @@ impl Future for MaskFuture { mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>, ) -> std::task::Poll { - self.inner.poll_unpin(cx).map_err(VortexError::from) + match &mut self.inner { + Inner::Ready(mask) => std::task::Poll::Ready(Ok(mask.clone())), + Inner::Pending(inner) => inner.poll_unpin(cx).map_err(VortexError::from), + } } } #[cfg(test)] mod tests { + use std::sync::atomic::AtomicBool; + use std::sync::atomic::Ordering; + use vortex_buffer::BitBuffer; use super::*; @@ -123,4 +161,41 @@ mod tests { Ok(()) }) } + + /// Ready and pending variants must behave identically through slice and await. + #[test] + fn pending_slice_matches_ready_slice() -> VortexResult<()> { + futures::executor::block_on(async { + let mask = Mask::from_buffer(BitBuffer::from_iter([true, false, true, true, false])); + let ready = MaskFuture::ready(mask.clone()); + let pending = { + let mask = mask.clone(); + MaskFuture::new(mask.len(), async move { Ok(mask) }) + }; + + assert_eq!(ready.slice(1..4).await?, pending.slice(1..4).await?); + assert_eq!(ready.clone().await?, pending.clone().await?); + Ok(()) + }) + } + + /// Inspect fires for an already-resolved mask and preserves the resolved value. + #[test] + fn inspect_fires_on_ready_mask() -> VortexResult<()> { + futures::executor::block_on(async { + let mask = Mask::from_buffer(BitBuffer::from_iter([true, false])); + let fired = Arc::new(AtomicBool::new(false)); + let fut = MaskFuture::ready(mask.clone()).inspect({ + let fired = Arc::clone(&fired); + move |r| { + assert!(r.is_ok()); + fired.store(true, Ordering::Relaxed); + } + }); + + assert!(fired.load(Ordering::Relaxed)); + assert_eq!(fut.await?, mask); + Ok(()) + }) + } }