Discussion: dtype system and integrating record types

[aldanor]

I've been looking at how record types can be integrated in rust-numpy and here's an unsorted collection of thoughts for discussion.

Let's look at Element:

pub unsafe trait Element: Clone + Send {
    const DATA_TYPE: DataType;
    fn is_same_type(dtype: &PyArrayDescr) -> bool;
    fn npy_type() -> NPY_TYPES { ... }
    fn get_dtype(py: Python) -> &PyArrayDescr { ... }
}

• npy_type() is used in PyArray::new() and the like. Instead, one should use PyArray_NewFromDescr() to make use of the custom descriptor. Should all places where npy_type() is used split between "simple type, use New" and "user type, use NewFromDescr"? Or, alternatively, should arrays always be constructed from descriptor? (in which case, npy_type() becomes redundant and should be removed)
• Why is same_type() needed at all? It is only used in FromPyObject::extract where one could simply use PyArray_EquivTypes (like it's done in pybind11). Isn't it largely redundant? (or does it exist for optimization purposes? In which case, is it even noticeable performance-wise?)
• DATA_TYPE constant is really only used to check if it's an object or not in 2 places, like this:

  if T::DATA_TYPE != DataType::Object

  Isn't this redundant as well? Given that one can always do

  T::get_dtype().get_datatype() != Some(DataType::Object)
  // or, can add something like: T::get_dtype().is_object()

• With all the notes above, Element essentially is just

   pub unsafe trait Element: Clone + Send {
       fn get_dtype(py: Python) -> &PyArrayDescr;
   }

• For structured types, do we want to stick the type descriptor into DataType? E.g.:

  enum DataType { ..., Record(RecordType) }

  Or, alternatively, just keep it as DataType::Void? In which case, how does one recover record type descriptor? (it can always be done through numpy C API of course, via PyArrayDescr).
• In order to enable user-defined record dtypes, having to return &PyArrayDescr would probably require:
  • Maintaining a global static thread-safe registry of registered dtypes (kind of like it's done in pybind11)
  • Initializing this registry somewhere
  • Any other options?
• Element should probably be implemented for tuples and fixed-size arrays.
• In order to implement structured dtypes, we'll inevitably have to resort to proc-macros. A few random thoughts and examples of how it can be done (any suggestions?):

  • #[numpy(record)]
    #[derive(Clone, Copy)]
    #[repr(packed)]
    struct Foo { x: i32, u: Bar } // where Bar is a registered numpy dtype as well
    // dtype = [('x', '<i4'), ('u', ...)]

  • We probably have to require either of #[repr(C)], #[repr(packed)] or #[repr(transparent)]
  • If repr is required, it can be an argument of the macro, e.g. #[numpy(record, repr = "C")]. (or not)
  • We also have to require Copy? (or not? technically, you could have object-type fields inside)
  • For wrapper types, we can allow something like this:

  • #[numpy(transparent)]
    #[repr(transparent)]
    struct Wrapper(pub i32);
    // dtype = '<i4'

  • For object types, the current suggestion in the docs is to implement a wrapper type and then impl Element for it manually. This seems largely redundant, given that the DATA_TYPE will always be Object. It would be nice if any #[pyclass]-wrapped types could automatically implement Element, but it would be impossible due to orphan rule. An alternative would be something like this:

    #[pyclass]
    #[numpy] // i.e., #[numpy(object)]
    struct Foo {}

  • How does one register dtypes for foreign (remote) types? I.e., OrderedFloat<f32> or Wrapping<u64> or some PyClassFromOtherCrate? We can try doing something like what serde does for remote types.

[aldanor]

Related: #234, #186.

[adamreichold]

W.r.t. the question on why Element looks it does, I can only answer

DATA_TYPE constant is really only used to check if it's an object or not in 2 places, like this:

for lack of involvement in the other design decisions. Here, we discussed that a simple const TRIVIALLY_COPYABLE: bool; could serve the same purpose but having the enum does not really hurt either. In any case, this being a compile-time constant is certainly nice as it will completely remove one of the two branches from the generated code where this is checked

Maintaining a global static thread-safe registry of registered dtypes (kind of like it's done in pybind11)

I think having a GIL-protected lazily initialized static like in

rust-numpy/src/slice_container.rs

Lines 109 to 112 in eb1068b

	fn type_object_raw(py: pyo3::Python) -> *mut ffi::PyTypeObject {
	static TYPE_OBJECT: LazyStaticType = LazyStaticType::new();
	TYPE_OBJECT.get_or_init::<Self>(py)
	}

might be preferable to a single global data structure?

Should all places where npy_type() is used split between "simple type, use New" and "user type, use NewFromDescr"?

While the overhead of producing a PyArrayDesc from a suitably synchronized data structure is probably small, I suspect that not having to produce it is still faster. Or does NumPy resort to PyArray_DescrFromType internally in any case?

[adamreichold]

Or does NumPy resort to PyArray_DescrFromType internally in any case?

That seems to be the case indeed:

https://github.com/numpy/numpy/blob/1798a7d463590de6dc0dfdad69735d92288313f3/numpy/core/src/multiarray/ctors.c#L1105

So we can probably replace our calls to PyArray_New. The only remaining difference I see is that we have to go through PY_ARRAY_API which implies an additional acquire load.

In order to implement structured dtypes, we'll inevitably have to resort to proc-macros.

I think it would actually be helpful to ignore this part completely for now and design the necessary types without any macro support first. The proc-macros should essentially only relieve programmers from writing repetitive and error prone code, shouldn't they?

[aldanor]

we discussed that a simple const TRIVIALLY_COPYABLE: bool; could serve the same purpose but having the enum does not really hurt either. In any case, this being a compile-time constant is certainly nice as it will completely remove one of the two branches from the generated code where this is checked

Then my initial guess was correct. I wonder though whether it's actually noticeable, i.e. has anyone tried to measure it? With all the other Python cruft on top, I would strongly suspect there'd be no visible difference at all. I understand where it comes from, it's just there's another implicit logical requirement in the Element trait - and it's that Self::get_dtype(py).get_datatype() == Some(Self::DATA_TYPE) - and this requirement is imposed only (?) to have DATA_TYPE as a compile time thing instead of runtime (where at runtime it's a single load dtype_ptr->type_num).

I think it would actually be helpful to ignore this part completely for now and design the necessary types without any macro support first. The proc-macros should essentially only relieve programmers from writing repetitive and error prone code, shouldn't they?

With record types, it's pretty much a requirement; anything else would be extremely boilerplate-ish and unsafe (with the user having to provide all offsets and descriptors for all fields manually). It might be nice to leave an option to construct descriptors manually at runtime, but I can't see a single real use case for it off top of my head. So, the necessary types and the infrastructure have to be designed in a way that would make it trivial to integrate them into proc macros, from the get go.

Or does NumPy resort to PyArray_DescrFromType internally in any case?

Yes, I think it does (and so does pybind11).

I think having a GIL-protected lazily initialized static like in ... might be preferable to a single global data structure?

Yep - thanks for pointing that out, something like that should be sufficient.

[aldanor]

Just checked it out: if one replaces

rust-numpy/src/array.rs

Line 850 in eb1068b

if T::DATA_TYPE != DataType::Object {

with something runtime like

        if !T::get_dtype(py).is_object() {

the difference in from_slice runtime for an array of 5 elements is 1ns (20ns vs 21ns).

For comparison, creating arange(1, 5, 1) is 31ns.

So yea... you could say it's noticeable (but only if you create hundreds millions of arrays).

[adamreichold]

and this requirement is imposed only (?) to have DATA_TYPE as a compile time thing instead of runtime

ATM, the DATA_TYPE member is used to generate default implementations for the other methods, i.e. it is a utility to simplify the implementation. If same_type and npy_type are dropped in favour of using get_dtype we could certainly remove remove this member. But I also do not really see a significant cost to having at least a const TRIVIAL: bool; member.

(where at runtime it's a single load dtype_ptr->type_num).

I think if anything has measurable overhead, then it is acquiring the reference to a PyArrayDesc itself, either by accessing PY_ARRAY_API.PyArray_DescrFromType or by accessing a lazily initialized static.

It might be nice to leave an option to construct descriptors manually at runtime, but I can't see a single real use case for it off top of my head. So, the necessary types and the infrastructure have to be designed in a way that would make it trivial to integrate them into proc macros, from the get go.

proc-macros generate code that is injected into the dependent crate, i.e. they have to use this crate's public API and all the code they produce must have been possible to be written by hand. The point is not that I expect user code to manually produce those trait implementations (but avoiding the build time overhead of proc-macros would be one reason to do so), but that we can focus the discussion on the mechanisms and leave the proc-macro UI for later.

[adamreichold]

So yea... you could say it's noticeable (but only if you create hundreds millions of arrays).

One possible caveat with a targeted benchmark here is that deciding this at compile time is about removing code and thereby reducing instruction cache pressure which could only be noticeable when other code shares the address space.

[adamreichold]

As a first actionable result of your findings, why not start with a PR to simplify the Element trait to avoid having same_type and npy_type and directly producing a descriptor via PyArrayDescr::from_npy_type? I think that would certainly be a nice simplification. Personally, I would like this combined with a const TRIVIAL: bool flag, but I understand that I am most likely not the only one with an opinion on this.

[adamreichold]

For object types, the current suggestion in the docs is to implement a wrapper type and then impl Element for it manually. This seems largely redundant, given that the DATA_TYPE will always be Object.

I fear that it might actually be unsound: During extraction, we only check the data type to be object, but not whether all the elements are of type Py<T> for the givenT: PyClass. So e.g. .readonly().as_array() would enable safe code to call methods on T even though completely different - possibly heterogeneously typed - objects are stored in the array.

I fear we need to remove this suggestion and limit ourselves to Py<PyAny> as the only reasonable Element implementation for now. I also wonder how checking for homogeneity during extraction is best handled if we do want to properly support PyArray<Py<T>, D>.

[aldanor]

One possible caveat with a targeted benchmark here is that deciding this at compile time is about removing code and thereby reducing instruction cache pressure which could only be noticeable when other code shares the address space.

While I'm often guilty in micro-optimizing Rust code down to every single branch hint myself, I just don't think it matters here, the last thing you're going to want to be doing in numpy bindings is instantiate millions of array objects in a loop - that was kind of my point.

Here's a few more thoughts on what would need to happen if const DATA_TYPE remains in place:

• NB: DATA_TYPE != DataType::Object would no longer be sufficient - because record dtypes can be both trivially copyable (requires all fields to be trivially copyable) and not.
• If we want to keep the data type const, this would imply DataType::Record(<...>) would need to be const-constructible as well. This implies that e.g. all fields names need to be &'static, all offsets are const and so on. It's probably doable but will be tough (see below).
• Note also that MSRV of 1.48.0 doesn't allow us access to min const generics, so messing with consts is going to be even harder. Also, std::ptr::addr_of!() has been only stabilized in the same release (1.51.0), if I remember correctly.

Oh yea, and also, record types can be nested - that would be fun to implement as const as well (at runtime we can just box it all on the heap).

[aldanor]

The note above only applies if it remains as DATA_TYPE, of course.

If it's changed to IS_POD: bool, then DataType can become a complex non-const type (to host the record dtype descriptor), etc.

[aldanor]

all the elements are of type Py<T> for the givenT: PyClass. So e.g. .readonly().as_array() would enable safe code to call methods on T even though completely different - possibly heterogeneously typed - objects are stored in the array.

I fear we need to remove this suggest and limit ourselves to Py<PyAny> as the only reasonable Element implementation for now. I also wonder how checking for homogeneity during extraction is best handled if we do want to properly support PyArray<Py<T>, D>.

Yea, I was asking myself whether I was reading that code correctly - it only checks for Object dtype but not the type of objects themselves.

There's three options:

• Always safe: Py<PyAny>, in which case it's always good and safe.
• Py<T>, in which case, unfortunately, we need to isinstance-check every single element in order to keep it safe.
• Unchecked and unsafe version of Py<T> which would be nice to keep available.

[aldanor]

As a first actionable result of your findings, why not start with a PR to simplify the Element trait to avoid having same_type and npy_type and directly producing a descriptor via PyArrayDescr::from_npy_type? I think that would certainly be a nice simplification. Personally, I would like this combined with a const TRIVIAL: bool flag, but I understand that I am most likely not the only one with an opinion on this.

Yes, I've already started on this, as a first step, will try to push something shortly. There's actually a few more subtle dtype-related bugs I've uncovered in the process that I'll fix as well (the tests are very sparse so I guess noone noticed anything so far...).

[aldanor]

(This is a bit of a wall of text, so thanks in advance to whoever reads through it in its entirety. I tried to keep it as organized as possible.)

tagging @adamreichold @kngwyu @davidhewitt

Ok, now that step 1 (#256) is ready to be merged, let's think about next steps. I'll use pybind11 as reference here since that's what I'm most familiar with having implemented a big chunk of that myself in the past.

Descriptor (format) strings and PEP 3118

In pybind11, dtypes can be constructed from buffer descriptor strings (see PEP 3118 for details). It actually calls back to NumPy's Python API to do that (numpy.core._internal._dtype_from_pep3118), but it's possible to reimplement it manually if we want to - I would actually learn towards the latter, provided we borrow and verify all the related tests from numpy itself.

The benefit of using descriptor strings - they're easy to hash if a registry is used, there's no nesting, etc - it's just a simple string that can cover any dtype possible, including scalar types, object types, multi-dimensionals subarrays, recarrays, and any combination of the above. We can also generate them with proc-macros at compile time and make them &'static if need be.

Now, if we go this route, we'll have to delve into "format chars", "type kinds", "buffer descriptor strings" etc. There's obviously a big overlap with pyo3::buffer module as there's ElementType::from_format and all the stuff around it, but it won't be sufficient without some reworking. Which raises the question: do we want to copy/paste stuff in pyo3::buffer and reimplement it to suit our needs? Or do we want to make it pybind11-like so that numpy arrays and buffers are tightly bound together, sharing the same infrastructure? See the next section for more details.

Buffers, pyo3::buffer and rust-numpy

In pybind11, buffers and arrays are bound together (like they should be), i.e. you can instantiate one from another (both ways) while controlling who owns the data. There's no such thing in rust-numpy and I think it's a pretty big gap.

There's pyo3::buffer that we can try to integrate with, but, off top of my head, there's some issues with it. But before we continue, consider an example:

>>> dt = np.dtype([('x', '<i8'), ('y', '<U1'), ('z', 'O', (1, 2))])

>>> arr = np.rec.fromarrays(
...     [[1, 2], ['a', 'b'], [[[int, str]], [[float, list]]]], 
...     dtype=dt,
... )

>>> arr
rec.array([(1, 'a', [[<class 'int'>, <class 'str'>]]),
           (2, 'b', [[<class 'float'>, <class 'list'>]])],
          dtype=[('x', '<i8'), ('y', '<U1'), ('z', 'O', (1, 2))])

>>> arr.data.format
'T{=q:x:@1w:y:(1,2)O:z:}'

Here, arr.data is the matching Python buffer (which you can also get via memoryview(arr)). Normally, you can do np.array(arr.data) and get the identical array back, but funnily enough, this exact example triggers the padding bug I have already reported in numpy upstream 6 years ago (numpy/numpy#7797) and which was then patched manually in pybind11 by hand.

Back to pyo3::buffer:

pub enum ElementType {
    SignedInteger { bytes: usize },
    UnsignedInteger { bytes: usize },
    Bool,
    Float { bytes: usize },
    Unknown,
}

Here's what's wrong with this enum:

• It allows unsupported types, like Float(3)
• It doesn't support complex types
• It doesn't support struct types
• It doesn't support subarrays
• It doesn't support object types

It's basically a distant relative of DataType that we have just removed from rust-numpy in favor of a single dtype pointer. With these limitations, there's no way we can make proper use of pyo3::buffer and we'll end up duplicating lots of its logic in rust-numpy.

And here's the trait:

pub unsafe trait Element: Copy {
    fn is_compatible_format(format: &CStr) -> bool;
}

Note that this is similar to numpy::dtype::Element and is_compatible_format is equivalent to same_type() which we have just cut out. And again, there's some issues with it:

• Copy is unnecessary as it blocks object types
  • ... and if Copy bound is removed, then some methods like PyBuffer::copy_to_slice() are no longer valid and should be rewritten same way as in rust-numpy (taking IS_COPY into account)
• It's impossible to recover the format descriptor from a type (e.g. u32) - kind of like we can do with numpy::dtype::<u32>(py).

One thing that could be done to improve it would be this:

pub unsafe trait Element: Clone + Send {
    const IS_COPY: bool;
    fn format() -> &'static CStr;
}

fn can_convert(from: &CStr, to: &CStr) -> bool { ... } // or whichever other way it's done

And then ElementType can be removed. Alternatively, it can be replaced with FormatString(&CStr) which would then be returned by Element::format().

Note how close it is now to the newly reworked numpy::Element - the only difference being in nul-terminated string pointer being returned instead of the dtype pointer.

Note also that any valid T: pyo3::buffer::Element is (logically speaking) also a valid numpy::dtype::Element since dtype ptr can always be built from a valid format string (the inverse is not always true, e.g. due to M and m types; but in most cases it can be worked around if need be).

What to do next

Should the buffers be touched up in pyo3 first, so we can then use that in rust-numpy to add support for record types, subarrays, etc? In which case, what exactly should be done?

One interesting point for discussion may be - could we share some of the functionality between buffers and numpy arrays? (in pybind11, array is derived from buffer and array_t is derived from array - which is quite nice and avoid lots of duplication) Just some random thinking, e.g. maybe buffer-like types like the buffer itself or numpy array could deref to some common base type which would provide shared functionality. Or maybe we could somehow try and connect the two Element traits together.