vortex_array/
memory.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Session-scoped buffer allocation.
5
6use std::any::Any;
7
8pub use vortex_buffer::BufferAllocator;
9pub use vortex_buffer::BufferAllocatorRef;
10pub use vortex_buffer::StaticBufferAllocator;
11use vortex_session::SessionExt;
12use vortex_session::SessionGuard;
13use vortex_session::SessionVar;
14use vortex_session::VortexSession;
15
16/// Session-scoped memory configuration for Vortex arrays.
17#[derive(Clone, Debug)]
18pub struct MemorySession {
19    allocator: BufferAllocatorRef,
20}
21
22impl MemorySession {
23    /// Creates a new memory configuration using the provided allocator.
24    pub fn new(allocator: BufferAllocatorRef) -> Self {
25        Self { allocator }
26    }
27
28    /// Returns the configured allocator.
29    pub fn allocator(&self) -> BufferAllocatorRef {
30        self.allocator.clone()
31    }
32
33    /// Updates the configured allocator.
34    pub fn set_allocator(&mut self, allocator: BufferAllocatorRef) {
35        self.allocator = allocator;
36    }
37}
38
39impl Default for MemorySession {
40    fn default() -> Self {
41        Self::new(BufferAllocatorRef::statically_allocated())
42    }
43}
44
45impl SessionVar for MemorySession {
46    fn as_any(&self) -> &dyn Any {
47        self
48    }
49
50    fn as_any_mut(&mut self) -> &mut dyn Any {
51        self
52    }
53}
54
55/// Extension methods for session-scoped buffer allocation.
56pub trait MemorySessionExt: SessionExt {
57    /// Returns the memory configuration.
58    fn memory(&self) -> SessionGuard<'_, MemorySession> {
59        self.get::<MemorySession>()
60    }
61
62    /// Returns the configured buffer allocator.
63    fn allocator(&self) -> BufferAllocatorRef {
64        self.memory().allocator()
65    }
66
67    /// Configures the session allocator and returns the session.
68    fn with_allocator(self, allocator: BufferAllocatorRef) -> VortexSession {
69        let session = self.session();
70        session.get_mut::<MemorySession>().set_allocator(allocator);
71        session
72    }
73}
74
75impl<S: SessionExt> MemorySessionExt for S {}
76
77#[cfg(test)]
78mod tests {
79    use vortex_buffer::BufferAllocatorRef;
80
81    use super::MemorySession;
82
83    #[test]
84    fn memory_session_replaces_allocator() {
85        let allocator = BufferAllocatorRef::statically_allocated();
86        let mut session = MemorySession::default();
87        session.set_allocator(allocator);
88        let buffer = session.allocator().copy_from([1u32, 2, 3]);
89        assert_eq!(buffer.as_slice(), [1, 2, 3]);
90    }
91}
vortex_array/memory.rs

vortex_array/
memory.rs
