Use pyo3-bytes from obstore (#128)

* Use pyo3-bytes from obstore

* Use pyo3 bytes in put

Author: kylebarron

Cargo.lock

@@ -28,6 +28,7 @@ object_store = { workspace = true }
 pyo3 = { workspace = true, features = ["chrono"] }
 pyo3-arrow = "0.6"
 pyo3-async-runtimes = { workspace = true, features = ["tokio-runtime"] }
+pyo3-bytes = "0.1.2"
 pyo3-file = { workspace = true }
 pyo3-object_store = { path = "../pyo3-object_store" }
 tokio = { workspace = true, features = [

obstore/Cargo.toml

@@ -1,12 +1,12 @@
 use std::io::SeekFrom;
 use std::sync::Arc;
 
-use arrow::buffer::Buffer;
+use bytes::Bytes;
 use object_store::buffered::BufReader;
 use pyo3::exceptions::{PyIOError, PyStopAsyncIteration, PyStopIteration};
 use pyo3::prelude::*;
-use pyo3_arrow::buffer::PyArrowBuffer;
 use pyo3_async_runtimes::tokio::future_into_py;
+use pyo3_bytes::PyBytes;
 use pyo3_object_store::{PyObjectStore, PyObjectStoreError, PyObjectStoreResult};
 use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncSeekExt, Lines};
 use tokio::sync::Mutex;
@@ -154,34 +154,34 @@ impl PyReadableFile {
     }
 }
 
-async fn read(reader: Arc<Mutex<BufReader>>, size: Option<usize>) -> PyResult<PyArrowBuffer> {
+async fn read(reader: Arc<Mutex<BufReader>>, size: Option<usize>) -> PyResult<PyBytes> {
     let mut reader = reader.lock().await;
     if let Some(size) = size {
         let mut buf = vec![0; size as _];
         reader.read_exact(&mut buf).await?;
-        Ok(PyArrowBuffer::new(Buffer::from_vec(buf)))
+        Ok(Bytes::from(buf).into())
     } else {
         let mut buf = Vec::new();
         reader.read_to_end(&mut buf).await?;
-        Ok(PyArrowBuffer::new(Buffer::from_vec(buf)))
+        Ok(Bytes::from(buf).into())
     }
 }
 
-async fn readline(reader: Arc<Mutex<BufReader>>) -> PyResult<PyArrowBuffer> {
+async fn readline(reader: Arc<Mutex<BufReader>>) -> PyResult<PyBytes> {
     let mut reader = reader.lock().await;
     let mut buf = String::new();
     reader.read_line(&mut buf).await?;
-    Ok(PyArrowBuffer::new(Buffer::from_vec(buf.into_bytes())))
+    Ok(Bytes::from(buf.into_bytes()).into())
 }
 
-async fn readlines(reader: Arc<Mutex<BufReader>>, hint: i64) -> PyResult<Vec<PyArrowBuffer>> {
+async fn readlines(reader: Arc<Mutex<BufReader>>, hint: i64) -> PyResult<Vec<PyBytes>> {
     let mut reader = reader.lock().await;
     if hint <= 0 {
         let mut lines = Vec::new();
         loop {
             let mut buf = String::new();
             let n = reader.read_line(&mut buf).await?;
-            lines.push(PyArrowBuffer::new(Buffer::from_vec(buf.into_bytes())));
+            lines.push(Bytes::from(buf.into_bytes()).into());
             // Ok(0) signifies EOF
             if n == 0 {
                 return Ok(lines);
@@ -198,7 +198,7 @@ async fn readlines(reader: Arc<Mutex<BufReader>>, hint: i64) -> PyResult<Vec<PyA
             let mut buf = String::new();
             let n = reader.read_line(&mut buf).await?;
             byte_count += n;
-            lines.push(PyArrowBuffer::new(Buffer::from_vec(buf.into_bytes())));
+            lines.push(Bytes::from(buf.into_bytes()).into());
             // Ok(0) signifies EOF
             if n == 0 {
                 return Ok(lines);

obstore/src/buffered.rs

@@ -1,16 +1,13 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
-use arrow::buffer::Buffer;
 use bytes::Bytes;
 use chrono::{DateTime, Utc};
 use futures::stream::{BoxStream, Fuse};
 use futures::StreamExt;
 use object_store::{GetOptions, GetRange, GetResult, ObjectStore};
 use pyo3::exceptions::{PyStopAsyncIteration, PyStopIteration, PyValueError};
 use pyo3::prelude::*;
-use pyo3::types::PyBytes;
-use pyo3_arrow::buffer::PyArrowBuffer;
 use pyo3_object_store::{PyObjectStore, PyObjectStoreError, PyObjectStoreResult};
 use tokio::sync::Mutex;
 
@@ -302,7 +299,7 @@ impl PyBytesWrapper {
 // support the buffer protocol in the future (e.g. for get_range) you may need to have a separate
 // wrapper of Bytes
 impl<'py> IntoPyObject<'py> for PyBytesWrapper {
-    type Target = PyBytes;
+    type Target = pyo3::types::PyBytes;
     type Output = Bound<'py, Self::Target>;
     type Error = PyErr;
 
@@ -311,7 +308,7 @@ impl<'py> IntoPyObject<'py> for PyBytesWrapper {
 
         // Copy all internal Bytes objects into a single PyBytes
         // Since our inner callback is infallible, this will only panic on out of memory
-        PyBytes::new_with(py, total_len, |target| {
+        pyo3::types::PyBytes::new_with(py, total_len, |target| {
             let mut offset = 0;
             for buf in self.0.iter() {
                 target[offset..offset + buf.len()].copy_from_slice(buf);
@@ -370,11 +367,11 @@ pub(crate) fn get_range(
     path: String,
     start: usize,
     end: usize,
-) -> PyObjectStoreResult<PyArrowBuffer> {
+) -> PyObjectStoreResult<pyo3_bytes::PyBytes> {
     let runtime = get_runtime(py)?;
     py.allow_threads(|| {
         let out = runtime.block_on(store.as_ref().get_range(&path.into(), start..end))?;
-        Ok::<_, PyObjectStoreError>(PyArrowBuffer::new(Buffer::from_bytes(out.into())))
+        Ok::<_, PyObjectStoreError>(pyo3_bytes::PyBytes::new(out))
     })
 }
 
@@ -392,7 +389,7 @@ pub(crate) fn get_range_async(
             .get_range(&path.into(), start..end)
             .await
             .map_err(PyObjectStoreError::ObjectStoreError)?;
-        Ok(PyArrowBuffer::new(Buffer::from_bytes(out.into())))
+        Ok(pyo3_bytes::PyBytes::new(out))
     })
 }
 
@@ -403,7 +400,7 @@ pub(crate) fn get_ranges(
     path: String,
     starts: Vec<usize>,
     ends: Vec<usize>,
-) -> PyObjectStoreResult<Vec<PyArrowBuffer>> {
+) -> PyObjectStoreResult<Vec<pyo3_bytes::PyBytes>> {
     let runtime = get_runtime(py)?;
     let ranges = starts
         .into_iter()
@@ -412,11 +409,7 @@ pub(crate) fn get_ranges(
         .collect::<Vec<_>>();
     py.allow_threads(|| {
         let out = runtime.block_on(store.as_ref().get_ranges(&path.into(), &ranges))?;
-        Ok::<_, PyObjectStoreError>(
-            out.into_iter()
-                .map(|buf| PyArrowBuffer::new(Buffer::from_bytes(buf.into())))
-                .collect(),
-        )
+        Ok::<_, PyObjectStoreError>(out.into_iter().map(|buf| buf.into()).collect())
     })
 }
 
@@ -441,7 +434,7 @@ pub(crate) fn get_ranges_async(
             .map_err(PyObjectStoreError::ObjectStoreError)?;
         Ok(out
             .into_iter()
-            .map(|buf| PyArrowBuffer::new(Buffer::from_bytes(buf.into())))
+            .map(pyo3_bytes::PyBytes::new)
             .collect::<Vec<_>>())
     })
 }

obstore/src/get.rs

@@ -12,8 +12,9 @@ use object_store::{
 };
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
-use pyo3::pybacked::{PyBackedBytes, PyBackedStr};
+use pyo3::pybacked::PyBackedStr;
 use pyo3::types::PyDict;
+use pyo3_bytes::PyBytes;
 use pyo3_file::PyFileLikeObject;
 use pyo3_object_store::{PyObjectStore, PyObjectStoreResult};
 
@@ -61,7 +62,7 @@ impl<'py> FromPyObject<'py> for PyUpdateVersion {
 pub(crate) enum MultipartPutInput {
     File(BufReader<File>),
     FileLike(PyFileLikeObject),
-    Buffer(Cursor<PyBackedBytes>),
+    Buffer(Cursor<PyBytes>),
 }
 
 impl MultipartPutInput {
@@ -83,7 +84,7 @@ impl<'py> FromPyObject<'py> for MultipartPutInput {
     fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
         if let Ok(path) = ob.extract::<PathBuf>() {
             Ok(Self::File(BufReader::new(File::open(path)?)))
-        } else if let Ok(buffer) = ob.extract::<PyBackedBytes>() {
+        } else if let Ok(buffer) = ob.extract::<PyBytes>() {
             Ok(Self::Buffer(Cursor::new(buffer)))
         } else {
             Ok(Self::FileLike(PyFileLikeObject::with_requirements(