shared array to avoid duplicate execute

Signed-off-by: Onur Satici <onur@spiraldb.com>

Author: onursatici

Cargo.lock

@@ -99,6 +99,7 @@ arrow-schema = "57.1"
 arrow-select = "57.1"
 arrow-string = "57.1"
 async-compat = "0.2.5"
+async-lock = "3.4"
 async-fs = "2.2.0"
 async-stream = "0.3.6"
 async-trait = "0.1.89"

Cargo.toml

@@ -22,6 +22,7 @@ workspace = true
 [dependencies]
 arbitrary = { workspace = true, optional = true }
 arcref = { workspace = true }
+async-lock = { workspace = true }
 arrow-arith = { workspace = true }
 arrow-array = { workspace = true, features = ["ffi"] }
 arrow-buffer = { workspace = true }

vortex-array/Cargo.toml

@@ -1,22 +1,21 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::future::Future;
 use std::sync::Arc;
 
-use parking_lot::RwLock;
+use async_lock::Mutex;
 use vortex_dtype::DType;
 use vortex_error::VortexResult;
-use vortex_error::vortex_panic;
 
 use crate::ArrayRef;
 use crate::Canonical;
-use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::stats::ArrayStats;
 
 #[derive(Debug, Clone)]
 pub struct SharedArray {
-    pub(super) state: Arc<RwLock<SharedState>>,
+    pub(super) state: Arc<Mutex<SharedState>>,
     pub(super) dtype: DType,
     pub(super) stats: ArrayStats,
 }
@@ -28,60 +27,57 @@ pub(super) enum SharedState {
 }
 
 impl SharedArray {
+    /// Creates a new `SharedArray` wrapping the given source array.
     pub fn new(source: ArrayRef) -> Self {
         Self {
             dtype: source.dtype().clone(),
-            state: Arc::new(RwLock::new(SharedState::Source(source))),
+            state: Arc::new(Mutex::new(SharedState::Source(source))),
             stats: ArrayStats::default(),
         }
     }
 
-    pub fn cached(&self) -> Option<Canonical> {
-        match &*self.state.read() {
-            SharedState::Cached(canonical) => Some(canonical.clone()),
-            SharedState::Source(_) => None,
-        }
-    }
-
-    pub fn cache_or_return(&self, canonical: Canonical) -> Canonical {
-        let mut state = self.state.write();
+    pub fn get_or_compute(
+        &self,
+        f: impl FnOnce(&ArrayRef) -> VortexResult<Canonical>,
+    ) -> VortexResult<Canonical> {
+        let mut state = self.state.lock_blocking();
         match &*state {
-            SharedState::Cached(existing) => existing.clone(),
-            SharedState::Source(_) => {
+            SharedState::Cached(canonical) => Ok(canonical.clone()),
+            SharedState::Source(source) => {
+                let canonical = f(source)?;
                 *state = SharedState::Cached(canonical.clone());
-                canonical
+                Ok(canonical)
             }
         }
     }
 
-    pub fn as_source(&self) -> ArrayRef {
-        let SharedState::Source(source) = &*self.state.read() else {
-            vortex_panic!("already cached");
-        };
-        source.clone()
-    }
-
-    pub(super) fn canonicalize(&self, ctx: &mut ExecutionCtx) -> VortexResult<Canonical> {
-        let source = {
-            let state = self.state.read();
-            match &*state {
-                SharedState::Cached(existing) => return Ok(existing.clone()),
-                SharedState::Source(source) => source.clone(),
+    pub async fn get_or_compute_async<F, Fut>(&self, f: F) -> VortexResult<Canonical>
+    where
+        F: FnOnce(ArrayRef) -> Fut,
+        Fut: Future<Output = VortexResult<Canonical>>,
+    {
+        let mut state = self.state.lock().await;
+        match &*state {
+            SharedState::Cached(canonical) => Ok(canonical.clone()),
+            SharedState::Source(source) => {
+                let source = source.clone();
+                let canonical = f(source).await?;
+                *state = SharedState::Cached(canonical.clone());
+                Ok(canonical)
             }
-        };
-        let canonical = source.execute::<Canonical>(ctx)?;
-        Ok(self.cache_or_return(canonical))
+        }
     }
 
     pub(super) fn current_array_ref(&self) -> ArrayRef {
-        match &*self.state.read() {
+        let state = self.state.lock_blocking();
+        match &*state {
             SharedState::Source(source) => source.clone(),
             SharedState::Cached(canonical) => canonical.clone().into_array(),
         }
     }
 
     pub(super) fn set_source(&mut self, source: ArrayRef) {
         self.dtype = source.dtype().clone();
-        *self.state.write() = SharedState::Source(source);
+        *self.state.lock_blocking() = SharedState::Source(source);
     }
 }

vortex-array/src/arrays/shared/array.rs

@@ -2,8 +2,10 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_buffer::buffer;
+use vortex_error::VortexResult;
 use vortex_session::VortexSession;
 
+use crate::Canonical;
 use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::arrays::PrimitiveArray;
@@ -14,28 +16,22 @@ use crate::session::ArraySession;
 use crate::validity::Validity;
 
 #[test]
-fn shared_array_caches_on_canonicalize() -> vortex_error::VortexResult<()> {
+fn shared_array_caches_on_canonicalize() -> VortexResult<()> {
     let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable).into_array();
     let shared = SharedArray::new(array);
 
-    assert!(shared.cached().is_none());
-
     let session = VortexSession::empty().with::<ArraySession>();
     let mut ctx = ExecutionCtx::new(session);
 
-    let first = shared.canonicalize(&mut ctx)?;
-    let cached = shared.cached().expect("canonicalize should cache result");
-    assert!(
-        cached
-            .as_ref()
-            .array_eq(first.as_ref(), HashPrecision::Value)
-    );
+    let first = shared.get_or_compute(|source| source.clone().execute::<Canonical>(&mut ctx))?;
+
+    // Second call should return cached without invoking the closure.
+    let second = shared.get_or_compute(|_| panic!("should not execute twice"))?;
 
-    let second = shared.canonicalize(&mut ctx)?;
     assert!(
-        second
+        first
             .as_ref()
-            .array_eq(first.as_ref(), HashPrecision::Value)
+            .array_eq(second.as_ref(), HashPrecision::Value)
     );
 
     Ok(())

vortex-array/src/arrays/shared/tests.rs

@@ -12,6 +12,7 @@ use vortex_session::VortexSession;
 use crate::ArrayBufferVisitor;
 use crate::ArrayChildVisitor;
 use crate::ArrayRef;
+use crate::Canonical;
 use crate::EmptyMetadata;
 use crate::ExecutionCtx;
 use crate::IntoArray;
@@ -96,7 +97,9 @@ impl VTable for SharedVTable {
     }
 
     fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
-        Ok(array.canonicalize(ctx)?.into_array())
+        Ok(array
+            .get_or_compute(|source| source.clone().execute::<Canonical>(ctx))?
+            .into_array())
     }
 }
 

vortex-array/src/arrays/shared/vtable.rs

@@ -28,11 +28,8 @@ impl CudaExecute for SharedExecutor {
             .ok()
             .vortex_expect("Array is not a Shared array");
 
-        if let Some(cached) = shared.cached() {
-            return Ok(cached);
-        }
-
-        let canonical = shared.as_source().execute_cuda(ctx).await?;
-        Ok(shared.cache_or_return(canonical))
+        shared
+            .get_or_compute_async(|source| source.execute_cuda(ctx))
+            .await
     }
 }