polars_core/chunked_array/object/extension/
mod.rs

1pub(crate) mod drop;
2pub(super) mod list;
3pub(crate) mod polars_extension;
4
5use std::mem;
6
7use polars_arrow::array::FixedSizeBinaryArray;
8use polars_arrow::bitmap::BitmapBuilder;
9use polars_arrow::datatypes::ExtensionType;
10use polars_buffer::Buffer;
11use polars_extension::PolarsExtension;
12use polars_utils::format_pl_smallstr;
13use polars_utils::relaxed_cell::RelaxedCell;
14
15use crate::PROCESS_ID;
16use crate::prelude::*;
17
18static POLARS_ALLOW_EXTENSION: RelaxedCell<bool> = RelaxedCell::new_bool(false);
19
20/// Control whether extension types may be created.
21///
22/// If the environment variable POLARS_ALLOW_EXTENSION is set, this function has no effect.
23pub fn set_polars_allow_extension(toggle: bool) {
24    POLARS_ALLOW_EXTENSION.store(toggle)
25}
26
27/// Invariants
28/// `ptr` must point to start a `T` allocation
29/// `n_t_vals` must represent the correct number of `T` values in that allocation
30unsafe fn create_drop<T: Sized>(mut ptr: *const u8, n_t_vals: usize) -> Box<dyn FnMut()> {
31    Box::new(move || {
32        let t_size = size_of::<T>() as isize;
33        for _ in 0..n_t_vals {
34            let _ = std::ptr::read_unaligned(ptr as *const T);
35            ptr = ptr.offset(t_size)
36        }
37    })
38}
39
40#[allow(clippy::type_complexity)]
41struct ExtensionSentinel {
42    drop_fn: Option<Box<dyn FnMut()>>,
43    // A function on the heap that take a `array: FixedSizeBinary` and a `name: PlSmallStr`
44    // and returns a `Series` of `ObjectChunked<T>`
45    pub(crate) to_series_fn: Option<Box<dyn Fn(&FixedSizeBinaryArray, &PlSmallStr) -> Series>>,
46}
47
48impl Drop for ExtensionSentinel {
49    fn drop(&mut self) {
50        if let Some(mut drop_fn) = self.drop_fn.take() {
51            (drop_fn)()
52        }
53    }
54}
55
56// https://stackoverflow.com/questions/28127165/how-to-convert-struct-to-u8d
57// not entirely sure if padding bytes in T are initialized or not.
58unsafe fn any_as_u8_slice<T: Sized>(p: &T) -> &[u8] {
59    std::slice::from_raw_parts((p as *const T) as *const u8, size_of::<T>())
60}
61
62/// Create an extension Array that can be sent to arrow and (once wrapped in [`PolarsExtension`] will
63/// also call drop on `T`, when the array is dropped.
64pub(crate) fn create_extension<I: Iterator<Item = Option<T>> + TrustedLen, T: Sized + Default>(
65    iter: I,
66) -> PolarsExtension {
67    let env = "POLARS_ALLOW_EXTENSION";
68    if !(POLARS_ALLOW_EXTENSION.load() || std::env::var(env).is_ok()) {
69        panic!("creating extension types not allowed - try setting the environment variable {env}")
70    }
71    let t_size = size_of::<T>();
72    let t_alignment = align_of::<T>();
73    let n_t_vals = iter.size_hint().1.unwrap();
74
75    let mut buf = Vec::with_capacity(n_t_vals * t_size);
76    let mut validity = BitmapBuilder::with_capacity(n_t_vals);
77
78    // when we transmute from &[u8] to T, T must be aligned correctly,
79    // so we pad with bytes until the alignment matches
80    let n_padding = (buf.as_ptr() as usize) % t_alignment;
81    buf.extend(std::iter::repeat_n(0, n_padding));
82
83    // transmute T as bytes and copy in buffer
84    for opt_t in iter.into_iter() {
85        match opt_t {
86            Some(t) => {
87                unsafe { buf.extend_from_slice(any_as_u8_slice(&t)) };
88                validity.push(true);
89                mem::forget(t);
90            },
91            None => {
92                // The drop function frees this copy, so the original must not be dropped too.
93                let t = T::default();
94                unsafe { buf.extend_from_slice(any_as_u8_slice(&t)) };
95                validity.push(false);
96                mem::forget(t);
97            },
98        }
99    }
100    // The drop function frees exactly `n_t_vals` values.
101    assert_eq!(
102        validity.len(),
103        n_t_vals,
104        "object iterator did not yield its declared length"
105    );
106
107    // We slice the buffer because we want to ignore the padding bytes from here
108    // they can be forgotten.
109    let buf: Buffer<u8> = Buffer::from_vec(buf).sliced(n_padding..);
110    // ptr to start of T, not to start of padding
111    let ptr = buf.as_slice().as_ptr();
112
113    // SAFETY: ptr and t are correct.
114    let drop_fn = unsafe { create_drop::<T>(ptr, n_t_vals) };
115    let et = Box::new(ExtensionSentinel {
116        drop_fn: Some(drop_fn),
117        to_series_fn: None,
118    });
119    let et_ptr = &*et as *const ExtensionSentinel;
120    std::mem::forget(et);
121
122    let metadata = format_pl_smallstr!("{};{}", *PROCESS_ID, et_ptr as usize);
123
124    let physical_type = ArrowDataType::FixedSizeBinary(t_size);
125    let extension_type = ArrowDataType::Extension(Box::new(ExtensionType {
126        name: PlSmallStr::from_static(POLARS_OBJECT_EXTENSION_NAME),
127        inner: physical_type,
128        metadata: Some(metadata),
129    }));
130
131    let array = FixedSizeBinaryArray::new(extension_type, buf, validity.into_opt_validity());
132
133    // SAFETY: we just heap allocated the ExtensionSentinel, so its alive.
134    unsafe { PolarsExtension::new(array) }
135}
136
137#[cfg(test)]
138mod test {
139    use std::fmt::{Display, Formatter};
140    use std::hash::{Hash, Hasher};
141
142    use polars_utils::total_ord::TotalHash;
143
144    use super::*;
145
146    #[derive(Clone, Debug, Default, Eq, Hash, PartialEq)]
147    struct Foo {
148        pub a: i32,
149        pub b: u8,
150        pub other_heap: String,
151    }
152
153    impl TotalEq for Foo {
154        fn tot_eq(&self, other: &Self) -> bool {
155            self == other
156        }
157    }
158
159    impl TotalHash for Foo {
160        fn tot_hash<H>(&self, state: &mut H)
161        where
162            H: Hasher,
163        {
164            self.hash(state);
165        }
166    }
167
168    impl Display for Foo {
169        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
170            write!(f, "{self:?}")
171        }
172    }
173
174    impl PolarsObject for Foo {
175        fn type_name() -> &'static str {
176            "object"
177        }
178    }
179
180    #[test]
181    fn test_create_extension() {
182        set_polars_allow_extension(true);
183        // Run this under MIRI.
184        let foo = Foo {
185            a: 1,
186            b: 1,
187            other_heap: "foo".into(),
188        };
189        let foo2 = Foo {
190            a: 1,
191            b: 1,
192            other_heap: "bar".into(),
193        };
194
195        let vals = vec![Some(foo), Some(foo2)];
196        create_extension(vals.into_iter());
197    }
198}
polars_core/chunked_array/object/extension/mod.rs

polars_core/chunked_array/object/extension/
mod.rs
