polars_core/chunked_array/object/
registry.rs

1//! This is a heap allocated utility that can be used to register an object type.
2//!
3//! That object type will know its own generic type parameter `T` and callers can simply
4//! send `&Any` values and don't have to know the generic type themselves.
5use std::any::Any;
6use std::fmt::{Debug, Formatter};
7use std::ops::Deref;
8use std::sync::{Arc, LazyLock, RwLock};
9
10use polars_arrow::array::builder::ArrayBuilder;
11use polars_arrow::array::{Array, ArrayRef};
12use polars_arrow::datatypes::ArrowDataType;
13use polars_utils::pl_str::PlSmallStr;
14
15use crate::chunked_array::object::builder::ObjectChunkedBuilder;
16use crate::datatypes::AnyValue;
17use crate::prelude::{ListBuilderTrait, ObjectChunked, PolarsObject};
18use crate::series::{IntoSeries, Series};
19
20/// Takes a `name` and `capacity` and constructs a new builder.
21pub type BuilderConstructor =
22    Box<dyn Fn(PlSmallStr, usize) -> Box<dyn AnonymousObjectBuilder> + Send + Sync>;
23pub type ObjectConverter = Arc<dyn Fn(AnyValue) -> Box<dyn Any> + Send + Sync>;
24pub type ObjectArrayGetter = Arc<dyn Fn(&dyn Array, usize) -> Option<AnyValue<'_>> + Send + Sync>;
25pub type WithGIL = Arc<dyn Fn(&mut dyn FnMut()) + Send + Sync>;
26
27pub struct ObjectRegistry {
28    /// A function that creates an object builder
29    pub builder_constructor: BuilderConstructor,
30    // A function that converts AnyValue to Box<dyn Any> of the object type
31    object_converter: Option<ObjectConverter>,
32    pub physical_dtype: ArrowDataType,
33    // A function that gets an AnyValue from a Box<dyn Array>.
34    array_getter: ObjectArrayGetter,
35    // A function which grabs the Python GIL.
36    with_gil: WithGIL,
37}
38
39impl Debug for ObjectRegistry {
40    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
41        write!(f, "object-registry")
42    }
43}
44
45static GLOBAL_OBJECT_REGISTRY: LazyLock<RwLock<Option<ObjectRegistry>>> =
46    LazyLock::new(Default::default);
47
48/// This trait can be registered, after which that global registration
49/// can be used to materialize object types
50pub trait AnonymousObjectBuilder: ArrayBuilder {
51    fn as_array_builder(self: Box<Self>) -> Box<dyn ArrayBuilder>;
52
53    /// # Safety
54    /// Expect `ObjectArray<T>` arrays.
55    unsafe fn from_chunks(self: Box<Self>, chunks: Vec<ArrayRef>) -> Series;
56
57    /// Append a `null` value.
58    fn append_null(&mut self);
59
60    /// Append a `T` of [`ObjectChunked<T>`][ObjectChunked<T>] made generic via the [`Any`] trait.
61    ///
62    /// [ObjectChunked<T>]: crate::chunked_array::object::ObjectChunked
63    fn append_value(&mut self, value: &dyn Any);
64
65    #[inline]
66    fn append_option(&mut self, value: Option<&dyn Any>) {
67        match value {
68            None => self.append_null(),
69            Some(v) => self.append_value(v),
70        }
71    }
72
73    /// Take the current state and materialize as a [`Series`]
74    /// the builder should not be used after that.
75    fn to_series(&mut self) -> Series;
76
77    fn get_list_builder(
78        &self,
79        name: PlSmallStr,
80        values_capacity: usize,
81        list_capacity: usize,
82    ) -> Box<dyn ListBuilderTrait>;
83}
84
85impl<T: PolarsObject> AnonymousObjectBuilder for ObjectChunkedBuilder<T> {
86    /// # Safety
87    /// Expects `ObjectArray<T>` arrays.
88    unsafe fn from_chunks(self: Box<Self>, chunks: Vec<ArrayRef>) -> Series {
89        ObjectChunked::<T>::new_with_compute_len(Arc::new(self.field().clone()), chunks)
90            .into_series()
91    }
92
93    fn as_array_builder(self: Box<Self>) -> Box<dyn ArrayBuilder> {
94        self
95    }
96
97    fn append_null(&mut self) {
98        self.append_null()
99    }
100
101    fn append_value(&mut self, value: &dyn Any) {
102        let value = value.downcast_ref::<T>().unwrap();
103        self.append_value(value.clone())
104    }
105
106    fn to_series(&mut self) -> Series {
107        let builder = std::mem::take(self);
108        builder.finish().into_series()
109    }
110    fn get_list_builder(
111        &self,
112        name: PlSmallStr,
113        values_capacity: usize,
114        list_capacity: usize,
115    ) -> Box<dyn ListBuilderTrait> {
116        Box::new(super::extension::list::ExtensionListBuilder::<T>::new(
117            name,
118            values_capacity,
119            list_capacity,
120        ))
121    }
122}
123
124pub fn register_object_builder(
125    builder_constructor: BuilderConstructor,
126    object_converter: ObjectConverter,
127    physical_dtype: ArrowDataType,
128    array_getter: ObjectArrayGetter,
129    with_gil: WithGIL,
130) {
131    let reg = GLOBAL_OBJECT_REGISTRY.deref();
132    let mut reg = reg.write().unwrap();
133
134    *reg = Some(ObjectRegistry {
135        builder_constructor,
136        object_converter: Some(object_converter),
137        physical_dtype,
138        array_getter,
139        with_gil,
140    })
141}
142
143#[cold]
144pub fn get_object_physical_type() -> ArrowDataType {
145    let reg = GLOBAL_OBJECT_REGISTRY.read().unwrap();
146    let reg = reg.as_ref().unwrap();
147    reg.physical_dtype.clone()
148}
149
150pub fn get_object_builder(name: PlSmallStr, capacity: usize) -> Box<dyn AnonymousObjectBuilder> {
151    let reg = GLOBAL_OBJECT_REGISTRY.read().unwrap();
152    let reg = reg.as_ref().unwrap();
153    (reg.builder_constructor)(name, capacity)
154}
155
156pub fn get_object_converter() -> ObjectConverter {
157    let reg = GLOBAL_OBJECT_REGISTRY.read().unwrap();
158    let reg = reg.as_ref().unwrap();
159    reg.object_converter.as_ref().unwrap().clone()
160}
161
162pub fn get_object_array_getter() -> ObjectArrayGetter {
163    let reg = GLOBAL_OBJECT_REGISTRY.read().unwrap();
164    reg.as_ref().unwrap().array_getter.clone()
165}
166
167/// Run the given function while holding the GIL.
168///
169/// This is sometimes used to avoid the overhead of repeatedly
170/// releasing and acquiring the GIL.
171pub fn run_with_gil<R, F: FnOnce() -> R>(f: F) -> R {
172    let reg = GLOBAL_OBJECT_REGISTRY.read().unwrap();
173    let with_gil = reg.as_ref().unwrap().with_gil.clone();
174    let r = &mut None;
175    let f = &mut Some(f);
176    (with_gil)(&mut || {
177        *r = Some((f.take().unwrap())());
178    });
179    r.take().unwrap()
180}
polars_core/chunked_array/object/registry.rs

polars_core/chunked_array/object/
registry.rs
