vortex_array/arrays/shared/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::fmt::Formatter;
6use std::future::Future;
7use std::sync::Arc;
8use std::sync::OnceLock;
9
10use async_lock::Mutex as AsyncMutex;
11use vortex_error::SharedVortexResult;
12use vortex_error::VortexResult;
13
14use crate::ArrayRef;
15use crate::Canonical;
16use crate::IntoArray;
17use crate::array::Array;
18use crate::array::ArrayParts;
19use crate::array::TypedArrayRef;
20use crate::array_slots;
21use crate::arrays::Shared;
22
23#[array_slots(Shared)]
24pub struct SharedSlots {
25    /// The source array that is shared and lazily computed.
26    #[slot(0)]
27    pub source: ArrayRef,
28}
29
30/// A lazily-executing array wrapper with a one-way transition from source to cached form.
31///
32/// Before materialization, operations delegate to the source array.
33/// After materialization (via `get_or_compute`), operations delegate to the cached result.
34#[derive(Debug, Clone)]
35pub struct SharedData {
36    cached: Arc<OnceLock<SharedVortexResult<ArrayRef>>>,
37    async_compute_lock: Arc<AsyncMutex<()>>,
38}
39
40impl Display for SharedData {
41    fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result {
42        Ok(())
43    }
44}
45
46#[expect(async_fn_in_trait)]
47pub trait SharedArrayExt: TypedArrayRef<Shared> + SharedArraySlotsExt {
48    fn current_array_ref(&self) -> &ArrayRef {
49        match self.cached.get() {
50            Some(Ok(arr)) => arr,
51            _ => self.source(),
52        }
53    }
54
55    fn get_or_compute(
56        &self,
57        f: impl FnOnce(&ArrayRef) -> VortexResult<Canonical>,
58    ) -> VortexResult<ArrayRef> {
59        let result = self
60            .cached
61            .get_or_init(|| f(self.source()).map(|c| c.into_array()).map_err(Arc::new));
62        result.clone().map_err(Into::into)
63    }
64
65    async fn get_or_compute_async<F, Fut>(&self, f: F) -> VortexResult<ArrayRef>
66    where
67        F: FnOnce(ArrayRef) -> Fut,
68        Fut: Future<Output = VortexResult<Canonical>>,
69    {
70        if let Some(result) = self.cached.get() {
71            return result.clone().map_err(Into::into);
72        }
73
74        let _guard = self.async_compute_lock.lock().await;
75
76        if let Some(result) = self.cached.get() {
77            return result.clone().map_err(Into::into);
78        }
79
80        let computed = f(self.source().clone())
81            .await
82            .map(|c| c.into_array())
83            .map_err(Arc::new);
84
85        let result = self.cached.get_or_init(|| computed);
86        result.clone().map_err(Into::into)
87    }
88}
89impl<T: TypedArrayRef<Shared>> SharedArrayExt for T {}
90
91impl SharedData {
92    pub fn new() -> Self {
93        Self {
94            cached: Arc::new(OnceLock::new()),
95            async_compute_lock: Arc::new(AsyncMutex::new(())),
96        }
97    }
98}
99
100impl Default for SharedData {
101    fn default() -> Self {
102        Self::new()
103    }
104}
105
106impl Array<Shared> {
107    /// Creates a new `SharedArray`.
108    pub fn new(source: ArrayRef) -> Self {
109        let dtype = source.dtype().clone();
110        let len = source.len();
111        unsafe {
112            Array::from_parts_unchecked(
113                ArrayParts::new(Shared, dtype, len, SharedData::new())
114                    .with_slots(SharedSlots { source }.into_slots()),
115            )
116        }
117    }
118}
vortex_array/arrays/shared/array.rs

vortex_array/arrays/shared/
array.rs
