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

vortex_array/arrays/shared/
array.rs
