vortex-array 0.86.1

Vortex in memory columnar data format
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

//! Session-scoped buffer allocation.

use std::any::Any;

pub use vortex_buffer::BufferAllocator;
pub use vortex_buffer::BufferAllocatorRef;
pub use vortex_buffer::StaticBufferAllocator;
use vortex_session::SessionExt;
use vortex_session::SessionGuard;
use vortex_session::SessionVar;
use vortex_session::VortexSession;

/// Session-scoped memory configuration for Vortex arrays.
#[derive(Clone, Debug)]
pub struct MemorySession {
    allocator: BufferAllocatorRef,
}

impl MemorySession {
    /// Creates a new memory configuration using the provided allocator.
    pub fn new(allocator: BufferAllocatorRef) -> Self {
        Self { allocator }
    }

    /// Returns the configured allocator.
    pub fn allocator(&self) -> BufferAllocatorRef {
        self.allocator.clone()
    }

    /// Updates the configured allocator.
    pub fn set_allocator(&mut self, allocator: BufferAllocatorRef) {
        self.allocator = allocator;
    }
}

impl Default for MemorySession {
    fn default() -> Self {
        Self::new(BufferAllocatorRef::statically_allocated())
    }
}

impl SessionVar for MemorySession {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }
}

/// Extension methods for session-scoped buffer allocation.
pub trait MemorySessionExt: SessionExt {
    /// Returns the memory configuration.
    fn memory(&self) -> SessionGuard<'_, MemorySession> {
        self.get::<MemorySession>()
    }

    /// Returns the configured buffer allocator.
    fn allocator(&self) -> BufferAllocatorRef {
        self.memory().allocator()
    }

    /// Configures the session allocator and returns the session.
    fn with_allocator(self, allocator: BufferAllocatorRef) -> VortexSession {
        let session = self.session();
        session.get_mut::<MemorySession>().set_allocator(allocator);
        session
    }
}

impl<S: SessionExt> MemorySessionExt for S {}

#[cfg(test)]
mod tests {
    use vortex_buffer::BufferAllocatorRef;

    use super::MemorySession;

    #[test]
    fn memory_session_replaces_allocator() {
        let allocator = BufferAllocatorRef::statically_allocated();
        let mut session = MemorySession::default();
        session.set_allocator(allocator);
        let buffer = session.allocator().copy_from([1u32, 2, 3]);
        assert_eq!(buffer.as_slice(), [1, 2, 3]);
    }
}