polars-ooc 0.54.2

Out-of-core processing support for Polars
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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

use std::fmt::Debug;
use std::ops::{Deref, DerefMut};

use polars_core::frame::DataFrame;

use crate::spill_context::ParameterFreeSpillContext;
use crate::{PinnedMut, PinnedRef, SpillToken, Spillable, memory_manager};

impl Spillable for DataFrame {
    // TODO: just a dummy spill for now. Boxed to reduce size.
    type Spilled = Box<DataFrame>;

    fn estimate_byte_size(&self) -> usize {
        self.estimated_size()
    }

    async fn spill(&self) -> Self::Spilled {
        Box::new(self.clone())
    }

    async fn unspill(location: &Self::Spilled) -> Self {
        (**location).clone()
    }
}

pub struct SpillFrame {
    token: SpillToken<DataFrame>,
    height: usize,
}

impl AsRef<SpillToken<DataFrame>> for SpillFrame {
    fn as_ref(&self) -> &SpillToken<DataFrame> {
        &self.token
    }
}

impl SpillFrame {
    pub fn new_unregistered(df: DataFrame) -> Self {
        let height = df.height();
        let token = SpillToken::new(df);
        Self { token, height }
    }

    pub async fn new<C: ParameterFreeSpillContext>(df: DataFrame, ctx: &C) -> Self {
        let slf = Self::new_unregistered(df);
        ctx.register(&slf);
        memory_manager().spill().await;
        slf
    }

    pub fn new_blocking<C: ParameterFreeSpillContext>(df: DataFrame, ctx: &C) -> Self {
        let slf = Self::new_unregistered(df);
        ctx.register(&slf);
        memory_manager().spill_blocking();
        slf
    }

    /// The height of the contained DataFrame. Does not need to unspill DataFrame.
    pub fn height(&self) -> usize {
        self.height
    }

    /// Get a reference to the underlying DataFrame, returning None if it was spilled.
    pub fn try_get(&self) -> Option<PinnedRef<'_, DataFrame>> {
        self.token.try_get()
    }

    /// Get a reference to the underlying DataFrame, unspilling it if it
    /// was spilled.
    pub async fn get(&self) -> PinnedRef<'_, DataFrame> {
        self.token.get().await
    }

    /// Blocking version of get.
    pub fn get_blocking(&self) -> PinnedRef<'_, DataFrame> {
        self.token.get_blocking()
    }

    /// Get a mutable reference to the underlying DataFrame, unspilling it if it
    /// was spilled.
    pub async fn get_mut(&mut self) -> PinnedFrameMut<'_> {
        PinnedFrameMut {
            inner: self.token.get_mut().await,
            height: &mut self.height,
        }
    }

    /// Blocking version of get_mut.
    pub fn get_mut_blocking(&mut self) -> PinnedFrameMut<'_> {
        PinnedFrameMut {
            inner: self.token.get_mut_blocking(),
            height: &mut self.height,
        }
    }

    /// Consumes this SpillFrame, unspilling it if it were spilled.
    pub async fn into_df(self) -> DataFrame {
        self.token.into_inner().await
    }

    /// Blocking version of into_df.
    pub fn into_df_blocking(self) -> DataFrame {
        self.token.into_inner_blocking()
    }
}

impl Debug for SpillFrame {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut s = f.debug_struct("SpillFrame");
        match self.token.try_get() {
            Some(df) => s.field("df", &*df),
            None => s.field("df", &"spilled"),
        };
        s.finish()
    }
}

pub struct PinnedFrameMut<'a> {
    height: &'a mut usize,
    inner: PinnedMut<'a, DataFrame>,
}

impl<'a> Deref for PinnedFrameMut<'a> {
    type Target = DataFrame;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<'a> DerefMut for PinnedFrameMut<'a> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

impl<'a> Drop for PinnedFrameMut<'a> {
    fn drop(&mut self) {
        *self.height = self.inner.height();
    }
}