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

vortex_layout/arrays/shared/
array.rs
