Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Replace OnceCell::Lazy with std::sync::LazyLock ? #124

Open
bionicles opened this issue Jan 2, 2025 · 2 comments
Open

Replace OnceCell::Lazy with std::sync::LazyLock ? #124

bionicles opened this issue Jan 2, 2025 · 2 comments

Comments

@bionicles
Copy link

bionicles commented Jan 2, 2025

in lib.rs:

use once_cell::sync::Lazy;

pub(crate) static POLARS: Lazy<PyObject> = Lazy::new(|| {
    Python::with_gil(|py| PyModule::import_bound(py, "polars").unwrap().to_object(py))
});

pub(crate) static SERIES: Lazy<PyObject> =
    Lazy::new(|| Python::with_gil(|py| POLARS.getattr(py, "Series").unwrap()));

we could remove this dependency on once_cell by using LazyLock in std::sync

use std::sync::LazyLock;

pub(crate) static POLARS: LazyLock<Py<PyModule>> = LazyLock::new(|| {
    Python::with_gil(|py| {
        let x = PyModule::import(py, "polars").unwrap().unbind();
        x
    })
});

pub(crate) static SERIES: LazyLock<PyObject> =
    LazyLock::new(|| Python::with_gil(|py| POLARS.getattr(py, "Series").unwrap()));

benefit: drops an external dependency on once_cell
downside: may increase MSRV (minimum supported rust version)

@bionicles
Copy link
Author

Ah, found another use of once_cell in the alloc module:

/// A memory allocator that relays allocations to the allocator used by Polars.
///
/// You can use it as the global memory allocator:
///
/// ```rust
/// use pyo3_polars::PolarsAllocator;
///
/// #[global_allocator]
/// static ALLOC: PolarsAllocator = PolarsAllocator::new();
/// ```
///
/// If the allocator capsule (`polars.polars._allocator`) is not available,
/// this allocator fallbacks to [`std::alloc::System`].
pub struct PolarsAllocator(OnceRef<'static, AllocatorCapsule>);

impl PolarsAllocator {
    fn get_allocator(&self) -> &'static AllocatorCapsule {
        // Do not allocate in this function,
        // otherwise it will cause infinite recursion.
        self.0.get_or_init(|| {
            let r = (unsafe { Py_IsInitialized() } != 0)
                .then(|| {
                    Python::with_gil(|_| unsafe {
                        (PyCapsule_Import(ALLOCATOR_CAPSULE_NAME.as_ptr() as *const c_char, 0)
                            as *const AllocatorCapsule)
                            .as_ref()
                    })
                })
                .flatten();
            #[cfg(debug_assertions)]
            if r.is_none() {
                // Do not use eprintln; it may alloc.
                let msg = b"failed to get allocator capsule\n";
                // Message length type is platform-dependent.
                let msg_len = msg.len().try_into().unwrap();
                unsafe { libc::write(2, msg.as_ptr() as *const libc::c_void, msg_len) };
            }
            r.unwrap_or(&FALLBACK_ALLOCATOR_CAPSULE)
        })
    }

    /// Create a `PolarsAllocator`.
    pub const fn new() -> Self {
        PolarsAllocator(OnceRef::new())
    }
}

Checking if I can replace this too ... might make a PR to fix for latest PyO3

@bionicles
Copy link
Author

Might be OK with std::sync::OnceLock as it has a similar get_or_init method:

use std::sync::OnceLock;

/// A memory allocator that relays allocations to the allocator used by Polars.
///
/// You can use it as the global memory allocator:
///
/// ```rust
/// use pyo3_polars::PolarsAllocator;
///
/// #[global_allocator]
/// static ALLOC: PolarsAllocator = PolarsAllocator::new();
/// ```
///
/// If the allocator capsule (`polars.polars._allocator`) is not available,
/// this allocator fallbacks to [`std::alloc::System`].
pub struct PolarsAllocator(OnceLock<&'static AllocatorCapsule>);

impl PolarsAllocator {
    fn get_allocator(&self) -> &'static AllocatorCapsule {
        // Do not allocate in this function,
        // otherwise it will cause infinite recursion.
        self.0.get_or_init(|| {
            let r = (unsafe { Py_IsInitialized() } != 0)
                .then(|| {
                    Python::with_gil(|_| unsafe {
                        (PyCapsule_Import(ALLOCATOR_CAPSULE_NAME.as_ptr() as *const c_char, 0)
                            as *const AllocatorCapsule)
                            .as_ref()
                    })
                })
                .flatten();
            #[cfg(debug_assertions)]
            if r.is_none() {
                // Do not use eprintln; it may alloc.
                let msg = b"failed to get allocator capsule\n";
                // Message length type is platform-dependent.
                // allow "useless" conversion to disable warning about platform-dependent try-into of msg_len
                #[allow(clippy::useless_conversion)]
                let msg_len = msg.len().try_into().unwrap();
                unsafe { libc::write(2, msg.as_ptr() as *const libc::c_void, msg_len) };
            }
            r.unwrap_or(&FALLBACK_ALLOCATOR_CAPSULE)
        })
    }

    /// Create a `PolarsAllocator`.
    pub const fn new() -> Self {
        PolarsAllocator(OnceLock::new())
    }
}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant