polars-core 0.53.0

Core of the Polars DataFrame library
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

use std::mem::ManuallyDrop;

use super::*;
use crate::prelude::*;

pub struct PolarsExtension {
    array: Option<FixedSizeBinaryArray>,
}

impl PolarsExtension {
    pub(crate) unsafe fn new(array: FixedSizeBinaryArray) -> Self {
        Self { array: Some(array) }
    }

    /// Take the Array hold by [`PolarsExtension`] and forget polars extension,
    /// so that drop is not called
    pub(crate) fn take_and_forget(self) -> FixedSizeBinaryArray {
        let mut md = ManuallyDrop::new(self);
        md.array.take().unwrap()
    }

    /// Apply a function with the sentinel, without the sentinels drop being called
    unsafe fn with_sentinel<T, F: FnOnce(&mut ExtensionSentinel) -> T>(&self, fun: F) -> T {
        let mut sentinel = self.get_sentinel();
        let out = fun(&mut sentinel);
        std::mem::forget(sentinel);
        out
    }

    /// Load the sentinel from the heap.
    /// be very careful, this dereferences a raw pointer on the heap,
    unsafe fn get_sentinel(&self) -> Box<ExtensionSentinel> {
        if let ArrowDataType::Extension(ext) = self.array.as_ref().unwrap().dtype() {
            let metadata = ext
                .metadata
                .as_ref()
                .expect("should have metadata in extension type");
            let mut iter = metadata.split(';');

            let pid = iter.next().unwrap().parse::<u128>().unwrap();
            let ptr = iter.next().unwrap().parse::<usize>().unwrap();
            if pid == *PROCESS_ID {
                Box::from_raw(ptr as *const ExtensionSentinel as *mut ExtensionSentinel)
            } else {
                panic!("pid did not mach process id")
            }
        } else {
            panic!("should be extension type")
        }
    }

    /// Calls the heap allocated function in the [`ExtensionSentinel`] that knows
    /// how to convert the [`FixedSizeBinaryArray`] to a `Series` of type [`ObjectChunked<T>`]
    pub(crate) unsafe fn get_series(&self, name: &PlSmallStr) -> Series {
        self.with_sentinel(|sent| {
            (sent.to_series_fn.as_ref().unwrap())(self.array.as_ref().unwrap(), name)
        })
    }

    // heap allocates a function that converts the binary array to a Series of [`ObjectChunked<T>`]
    // the `name` will be the `name` of the output `Series` when this function is called (later).
    pub(crate) unsafe fn set_to_series_fn<T: PolarsObject>(&mut self) {
        let f = Box::new(move |arr: &FixedSizeBinaryArray, name: &PlSmallStr| {
            let iter = arr.iter().map(|opt| {
                opt.map(|bytes| {
                    let t = std::ptr::read_unaligned(bytes.as_ptr() as *const T);

                    let ret = t.clone();
                    std::mem::forget(t);
                    ret
                })
            });

            let ca = ObjectChunked::<T>::from_iter_options(name.clone(), iter);
            ca.into_series()
        });
        self.with_sentinel(move |sent| {
            sent.to_series_fn = Some(f);
        });
    }
}

impl Drop for PolarsExtension {
    fn drop(&mut self) {
        // implicitly drop by taking ownership
        unsafe { self.get_sentinel() };
    }
}