Skip to content

Commit

Permalink
rust constructor for PyArrowBuffer (#242)
Browse files Browse the repository at this point in the history
* rust constructor for PyArrowBuffer

* improved doc
  • Loading branch information
kylebarron authored Oct 16, 2024
1 parent 64a6591 commit 45be4a1
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 3 deletions.
2 changes: 1 addition & 1 deletion pyo3-arrow/src/array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -232,7 +232,7 @@ impl PyArray {
let arr = self.array.as_primitive::<Int64Type>();
let values = arr.values();
let buffer = values.inner().clone();
crate::buffer::PyArrowBuffer::from_arrow(buffer)
crate::buffer::PyArrowBuffer::new(buffer)
}
_ => todo!(),
}
Expand Down
25 changes: 23 additions & 2 deletions pyo3-arrow/src/buffer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,20 @@ use crate::PyArray;

/// A wrapper around an Arrow [Buffer].
///
/// This implements both import and export via the Python buffer protocol.
///
/// ### Buffer import
///
/// This can be very useful as a general way to support ingest of a Python buffer protocol object.
/// The underlying Arrow [Buffer] manages the external memory, automatically calling the Python
/// buffer's release callback when the Arrow [Buffer] reference count reaches 0.
///
/// This does not need to be used with Arrow at all! This can be used with any API where you want
/// to handle both Python-provided and Rust-provided buffers. [`PyArrowBuffer`] implements
/// `AsRef<[u8]>`.
///
/// ### Buffer export
///
/// The Python buffer protocol is implemented on this buffer to enable zero-copy data transfer of
/// the core buffer into Python. This allows for zero-copy data sharing with numpy via
/// `numpy.frombuffer`.
Expand All @@ -35,8 +49,15 @@ impl AsRef<Buffer> for PyArrowBuffer {
}
}

impl AsRef<[u8]> for PyArrowBuffer {
fn as_ref(&self) -> &[u8] {
self.0.as_ref()
}
}

impl PyArrowBuffer {
pub(crate) fn from_arrow(buffer: Buffer) -> Self {
/// Construct a new [PyArrowBuffer]
pub fn new(buffer: Buffer) -> Self {
Self(buffer)
}

Expand All @@ -50,7 +71,7 @@ impl PyArrowBuffer {
impl PyArrowBuffer {
/// new
#[new]
pub fn new(buf: PyArrowBuffer) -> Self {
fn py_new(buf: PyArrowBuffer) -> Self {
buf
}

Expand Down

0 comments on commit 45be4a1

Please sign in to comment.