[winstxnhdw]

Summary

Currently, the overhead of PyBuffer takes a couple of microseconds to allocate on the heap, which may be too much overhead for some workloads.

This PR implements a pinned stack-allocated PyUntypedBuffer variant, PyUntypedBufferView.

Closes #5836

[davidhewitt]

Thanks very much for this. Various thoughts around the accessor methods.

Also, out of scope for this PR, but I keep wondering if we should provide iterators for these structures. Especially with the strides / suboffsets etc, it's not necessarily trivial to get this right.

[winstxnhdw]

Also, out of scope for this PR, but I keep wondering if we should provide iterators for these structures. Especially with the strides / suboffsets etc, it's not necessarily trivial to get this right.

Yes, I think it would definitely be useful as well. I am not yet sure how this would look like though.

[davidhewitt]

Thanks for the continued work here, looking great!

Implementing Drop directly on PyUntypedBufferView is functionally equivalent to what I had in mind from the "drop guard", so gets a 👍 from me.

Afraid I had quite a few more thoughts around some of the edge cases (plus some ideas we can probably ignore). Hopefully helps us get to the right eventual abstraction!

[codspeed]

Merging this PR will not alter performance

✅ 143 untouched benchmarks
⏩ 6 skipped benchmarks¹

---

_{Comparing winstxnhdw:feat/stacked-pybuffer (c18b6a5) with main (424149e)}

Footnotes

1. 6 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

[davidhewitt]

Very cool!

After reading through this implementation a couple of times, I really like the expressiveness this gives in the type system. I have some ideas which might refine it further (see BufferFlags struct idea).

Main concern is about generic code bloat from explosion of generic parameters. I am unsure if there's a way that we can mitigate that at all.

[davidhewitt]

Thanks for all of this - would be interested to see what you think of these suggestions (none are mandatory, I'm just pushing out ideas of what I think I like, which may not be to others' taste!)

[davidhewitt]

#5870 has now merged, we'll want to add a similar API here.

[winstxnhdw]

Sorry for the delay. I've been a little burnt out from school + work. I've adapted your suggestions and modified them to be what I think is appropriate.

Also, if you do PyBufferFlags::simple().full(), you won't be able to append let's say format() on it anymore because full() already implies format(). Just a nice DX win that I haven't really seen in other libraries like reqwest or tokio.

This is ok.

let bytes = PyBytes::new(py, b"abcde");
let flags = PyBufferFlags::simple()
    .writable()
    .format()
    .c_contiguous()
    .strides();

PyUntypedBufferView::with_flags(&bytes, flags, |view| {
 ...
})

This is not ok.

let bytes = PyBytes::new(py, b"abcde");
let flags = PyBufferFlags::simple()
    .writable()
    .format()
    .full() // errors here
    .c_contiguous()
    .strides();

PyUntypedBufferView::with_flags(&bytes, flags, |view| {
 ...
})

[davidhewitt]

Thanks, this is looking really cool. We're obviously venturing into new territory here so please bear with me if we have a few rounds of review while we experiment with what feels good (both to implement and to use).

A bunch of comments applied. I suspect we'll learn from these.

[davidhewitt]

Thanks, I think this is looking in pretty good shape. Let's see if we get an answer on python/cpython#148431 regarding format / itemsize interactions just in case that encourages us to make some other tweaks.

I would love feedback from some buffer API consumers about how this looks.

Ping @alex @kylebarron, I believe you both use PyO3's buffer API, please forgive the ping. I would love to ask:

• Would this view type be useful to you for performance benefit over the PyBuffer<T> case which incurs an unconditional heap allocation? (I think it seems worth it for some consumers even if it doesn't affect your use cases, so this won't block merge if you say no.)
• As consumers of the API, what do you think of the PyBufferFlags type which encodes the request and what fields from the exporter are meaningful in the type system. See e.g. #5894 (comment) - if the buffer request didn't ask for shape information, then .shape() accessor is not available, for example. I think we've cooked up something pretty nice here, though it would be nice to know whether others agree before we merge it 😅

See also docs for this PyUntypedBufferView on this branch at https://69dbb09a02f842688b3dcc9f--pyo3.netlify.app/main/doc/pyo3/buffer/struct.pyuntypedbufferview

[davidhewitt]

A thought: if we wanted to change CONTIGUITY to be an enum later (when possible) or add SUBOFFSETS, we would need users to be unable to directly name this type.

I think we might be able to achieve that by making this struct something like PyBufferFlagsImpl and having the real PyBufferFlags type be a thin re-export around it which uses the PyBufferFlagsImpl as a type parameter, similar to the way the view types do this.

I'm unsure of the merits of doing so, what do you think? It would be nice to have control over the implementation later if it doesn't add tons of complexity for both us and readers.

[winstxnhdw]

I agree that we shouldn't expose PyBufferFlags to minimise breaking changes. I've made the change, but I think it looks a bit ugly. I've made it a separate commit so that you can easily view the diff, and for me to revert if needed. Let me know what you think.

[davidhewitt]

It doesn't seem too bad to me, what bit do you dislike?

If you don't like the name FlagsImpl, one thought I have is that we could rename PyBufferFlags -> PyBufferRequest and PyBufferFlagsImpl -> PyBufferFlags.

"Compound Requests" is the term given in the Python docs for the flag combinations, after all.

See also my other comment on .format() method which might hide many of the uses of FlagsImpl if applied to all possible patterns.

[kylebarron]

I'm not sure I have enough context to answer these questions. Do you have a full usage example perhaps?

• Would this view type be useful to you for performance benefit over the PyBuffer<T> case which incurs an unconditional heap allocation?

What is PyBuffer<T> allocating? It's not copying the entire buffer, right? Do you mean it's essentially Boxing the reference to the Python data? If so, that seems like a small price to pay to be amortized over the entire input buffer (I tend to use the buffer protocol with large inputs, so for a small overhead that's amortized over the entire input I don't mind much).

• As consumers of the API, what do you think of the PyBufferFlags type which encodes the request and what fields from the exporter are meaningful in the type system. See e.g. #5894 (comment) - if the buffer request didn't ask for shape information, then .shape() accessor is not available, for example. I think we've cooked up something pretty nice here, though it would be nice to know whether others agree before we merge it 😅

That does look cool.

For my usage I tend to want to expose an input Python buffer as either Bytes or an Arrow Buffer. See for example FromPyObject for my PyBytes, which then I (unsafely) implement AsRef<[u8]> so that I can interpret the Python memory region for Bytes::from_owner.

I think this is sound if the Python user doesn't mutate the buffer while the Rust code is running? My use cases usually expect immutable buffers, even though there's no invariant to force that.

I don't think I'd have much benefit from a FnOnce closure that can access the Python region just once, even though that does look safer than my usage.

[winstxnhdw]

I've also been thinking if we could implement the extractor for the view buffers. It would definitely be nicer to use. I'll investigate.

[davidhewitt]

Thanks, some responses to the PyBufferFlagsImpl changes.

[davidhewitt]

It doesn't seem too bad to me, what bit do you dislike?

If you don't like the name FlagsImpl, one thought I have is that we could rename PyBufferFlags -> PyBufferRequest and PyBufferFlagsImpl -> PyBufferFlags.

"Compound Requests" is the term given in the Python docs for the flag combinations, after all.

See also my other comment on .format() method which might hide many of the uses of FlagsImpl if applied to all possible patterns.

[davidhewitt]

What is PyBuffer<T> allocating? It's not copying the entire buffer, right? Do you mean it's essentially Boxing the reference to the Python data? If so, that seems like a small price to pay to be amortized over the entire input buffer (I tend to use the buffer protocol with large inputs, so for a small overhead that's amortized over the entire input I don't mind much).

You're absolutely correct - for large buffers uses the allocation overhead can be irrelevant, and keeping the buffer information on the stack has the downside that you can't keep the buffer export alive for later consumption. OP noticed the performance advantage was relevant to their use case, however.

I can imagine this stack-based form is generally better for short-lived temporary buffers used to read (or write) a small bytes export.

I've also been thinking if we could implement the extractor for the view buffers. It would definitely be nicer to use. I'll investigate.

Can you clarify what you're thinking here? Maybe I can offer ideas :)

[davidhewitt]

I think this is sound if the Python user doesn't mutate the buffer while the Rust code is running? My use cases usually expect immutable buffers, even though there's no invariant to force that.

Yes, that sounds about right. @alex has previously written about how the current buffer API is deficient of guarantees against data races, for now I guess we live with that. I think @ngoldbaum is also thinking about this problem from time to time (with all three of free-threading, numpy, and PyO3 hats on, even).

[winstxnhdw]

I realised that we never explicitly prevent scalar buffers even though the existing API behavior doesn't support scalar buffers as the constructor will reject null shape and stride pointers.

[winstxnhdw]

Hi @davidhewitt, just bumping this in case you forgot about it!

[davidhewitt]

Thanks for all the work here and sorry for the many delays on my end! We recently made some changes in #6381 which adjusted what APIs we expose, so I pushed a commit here to mirror that rather than ask you for yet another revision. Let's merge 👍

[winstxnhdw]

Thank you so much for all your patience. So excited to use this new API!

[winstxnhdw]

Hey @davidhewitt, I am not familiar with WASM but it seems like the ctypes test is failing on WASI. Should we exclude the test WASI builds?

[bschoenmaeckers]

Hey @davidhewitt, I am not familiar with WASM but it seems like the ctypes test is failing on WASI. Should we exclude the test WASI builds?

I think ignoring this test on wasi is fine.