Support resize on PyBytes when initialised via with_new

[AudriusButkevicius]

I'm building a rust library that I expose to python via pyo3. This library is generating quite a lot of data in rust and handing it over to python (100GB for example).

One problem with how I generate the data is that I don't know exactly how much data there will be until it's generated (decompression), so I cannot allocate a correctly sized PyBytes object ahead of time, I need to first decompress the data in a rust allocated buffer, check its size, then allocate a PyBytes of the right size, and copy the data over. Given the amount of data, the copy from rust storage to PyBytes becomes a large overall cost of the api.

There also seems to be no C python api to allocate a bytes object from an existing block of memory, without copying the data, this is because the object definition has to be immediately before the memory that is the content of the object.

However, C api supports _Py_Resize There are some constraints of when it can and cannot be used, but I think if used from within (or shortly after) the PyBytes::with_new closure, passes all the checks.

My proposal is to either:

1. Allow returning "used" size from within the PyBytes::with_new init closure, which should be equal or less than the initial requested allocation, which if less than initial allocation, calls _Py_Resize after the with_new closure.
2. Implement a new new_with equivalent constructor, that gets a Vec like interface which allows both growing and shrinking the PyBytes object from within the closure.

Currently I'm working around this by allocating a PyObject myself, doing object pointer lifetime juggling myself. Seems to work however I'm not sure I understand who owns the object at what time, or how to free it properly in case ownership never gets handed over to the interpreter. Tons of unsafe code I don't feel comfortable with.

I also cannot use bytearray because I genuinely do not want the data to be modified from python.

I'd be happy to implement 1, however I feel I might be too new to rust to implement 2 (could do with hand holding I guess).

[neachdainn]

If the first proposal is the chosen one, then it would probably be good to have a version that doesn't zero out the bytes first. Something like:

fn with_new_uninit<F>(py: Python<'_>, len: usize, init: F) -> PyResult<&Self>
where
    F: FnMut(&[MaybeUninit<u8>]) -> PyResult<()>

[AudriusButkevicius]

I think that is already the case? pyo3 calls PyBytes_FromStringAndSize with a null pointer, and according to the docs:

If v is NULL, the contents of the bytes object are uninitialized.

[neachdainn]

Yes, but PyO3 writes zeros to the entire slice.

[AudriusButkevicius]

Ah, yes, missed that.

[davidhewitt]

There's some previous discussion to this in #1074

We could explore making it work but _Py_Resize is private, so I'd need to ask the CPython core devs how they would like this to work.

[AudriusButkevicius]

Whats the process to get that going?

I assumed that the function had underscore just because its usable in a very specific context.

[damelLP]

Could we do this by adding an alternative utility function sth like:

   pub fn resizeable_new_with<F>(py: Python<'_>, len: usize, init: F) -> PyResult<&PyBytes>
    where
    ¦   F: FnOnce(&mut [u8]) -> PyResult<usize>,
    {
    ¦   unsafe {
    ¦   ¦   let mut pyptr =
    ¦   ¦   ¦   ffi::PyBytes_FromStringAndSize(std::ptr::null(), len as ffi::Py_ssize_t);

    ¦   ¦   if pyptr.is_null() {
    ¦   ¦   ¦   return Err(PyRuntimeError::new_err(format!(
    ¦   ¦   ¦   ¦   "failed to allocate python bytes object of size {}",
    ¦   ¦   ¦   ¦   len
    ¦   ¦   ¦   )));
    ¦   ¦   }

    ¦   ¦   // Check for an allocation error and return it
    ¦   ¦   let buffer = ffi::PyBytes_AsString(pyptr) as *mut u8;
    ¦   ¦   debug_assert!(!buffer.is_null());

    ¦   ¦   // If init returns an Err, pypybytearray will automatically deallocate the buffer
    ¦   ¦   let new_len = init(std::slice::from_raw_parts_mut(buffer, len))?;

    ¦   ¦   if _PyBytes_Resize(ptr::addr_of_mut!(pyptr), new_len as ffi::Py_ssize_t) != 0 {
    ¦   ¦   ¦   return Err(PyRuntimeError::new_err("failed to resize bytes object"));
    ¦   ¦   }

    ¦   ¦   let pypybytes: Py<PyBytes> = Py::from_owned_ptr_or_err(py, pyptr)?;
    ¦   ¦   Ok(pypybytes.into_ref(py))
    ¦   }
    }

[neachdainn]

@damelLP: The _PyBtyes_Resize that the code uses is the potential problem here, not necessarily the method name.

@AudriusButkevicius and @davidhewitt: I'm talking about things I'm not very familiar with, but wouldn't it be possible to construct the data with PyByteArrays and then simply make a PyBytes object from the array? I'm not fully certain how much overhead that would add but it would be a stable / public way to do this.

[AudriusButkevicius]

That would perform a copy (as far as I understand) which is what I want to avoid. As it stands, I dont think there is a zero copy way to allocate PyBytes

[neachdainn]

It uses the buffer protocol which, if I understand correctly, doesn't copy the underlying data.

While each of these types have their own semantics, they share the common characteristic of being backed by a possibly large memory buffer. It is then desirable, in some situations, to access that buffer directly and without intermediate copying.

[neachdainn]

It should also be possible to create a new type that has the buffer protocol and pass that to PyBytes without copying.

[AudriusButkevicius]

I assumed it was copying. Do you have docs that suggests its zero copy?

[neachdainn]

Scattered throughout the buffer protocol documentation

While each of these types have their own semantics, they share the common characteristic of being backed by a possibly large memory buffer. It is then desirable, in some situations, to access that buffer directly and without intermediate copying.

Buffer structures (or simply “buffers”) are useful as a way to expose the binary data from another object to the Python programmer. They can also be used as a zero-copy slicing mechanism. Using their ability to reference a block of memory, it is possible to expose any data to the Python programmer quite easily. The memory could be a large, constant array in a C extension, it could be a raw block of memory for manipulation before passing to an operating system library, or it could be used to pass around structured data in its native, in-memory format.

These are not things I've used before, but the documentation heavily suggests, if not says outright, that it should be zero copy.

[adamreichold]

It should also be possible to create a new type that has the buffer protocol and pass that to PyBytes without copying.

Considering the resizeable PyBytes might be XY problem, I would also suggest looking into creating a Python object implementing the buffer protocol. For example, in rust-numpy we have PySliceContainer which we use a the base object for NumPy arrays which are backed by Rust-allocated Vec or rather ndarray::Array instances, i.e. writing the data into a Python object is not really required, just a Python object which will properly deallocate the Rust allocation to act as the base object for the buffer protocol implementor.

[davidhewitt]

Sorry that I fell off this thread a little.

Whats the process to get that going?

I assumed that the function had underscore just because its usable in a very specific context.

There's probably good reasons why a general _Py_Resize function is private. The right process would be to create an issue upstream in CPython with a clear proposal of what's needed, and possibly even help implement.

Considering the resizeable PyBytes might be XY problem, I would also suggest looking into creating a Python object implementing the buffer protocol.

👍 on this, and in particular we plan to make this easier with #3148.