refactor v3 data types

[d-v-b]

As per #2750, we need a new model of data types if we want to support more data types. Accordingly, this PR will refactor data types for the zarr v3 side of the codebase and make them extensible. I would also like to handle v2 as well with the same data structures, and confine the v2 / v3 differences to the places where they vary.

In main,all the v3 data types are encoded as variants of an enum (i.e., strings). Enumerating each dtype as a string is cumbersome for datetimes, that are parametrized by a time unit, and plain unworkable for parametric dtypes like fixed-length strings, which are parametrized by their length. This means we need a model of data types that can be parametrized, and I think separate classes is probably the way to go here. Separating the different data types into different objects also gives us a natural way to capture some of the per-data type variability baked into the spec: each data type class can define its own default value, and also define methods for how its scalars should be converted to / from JSON.

This is a very rough draft right now -- I'm mostly posting this for visibility as I iterate on it.

[d-v-b]

copying a comment @nenb made in this zulip discussion:

The first thing that caught my eye was that you are using numpy character codes. What was the motivation for this? numpy character codes are not extensible in their current format, and lead to issues like: jax-ml/ml_dtypes#41.

A feature of the character code is that it provides a way to distinguish parametric types like U* from parametrized instances of those types (like U3). Defining a class with the character code U means instances of the class can be initialized with a "length" parameter, and then we can make U2, U3, etc, as instances of the same class. If instead we bind a concrete numpy dtype as class attributes, we need a separate class for U2, U3, etc, which is undesirable. I do think I can work around this, but I figured the explanation might be helpful.

[nenb]

Summarising from a zulip discussion:

@nenb: How is the endianness of a dtype handled?

@d-v-b: In v3, endianness is specified by the codecs. The exact same data type could be decoded to big or little endian representations based on the state of a codec. This is very much not how v2 does it -- v2 puts the endianness in the dtype.

Proposed solution: Make endianness an attribute on in the dtype instance. This will be an implementation detail used by zarr-python to handle endianness, but won't be part of the dtype on disk (as requuired by the spec).

[d-v-b]

Summarising from a zulip discussion:

@nenb: How is the endianness of a dtype handled?

@d-v-b: In v3, endianness is specified by the codecs. The exact same data type could be decoded to big or little endian representations based on the state of a codec. This is very much not how v2 does it -- v2 puts the endianness in the dtype.

Proposed solution: Make endianness an attribute on in the dtype instance. This will be an implementation detail used by zarr-python to handle endianness, but won't be part of the dtype on disk (as requuired by the spec).

Thanks for the summary! I have implemented the proposed solution.

[d-v-b]

I recently pushed a lot of internal changes. A quick summary:

• I added explicit type annotations to the json-producing functions, so wherever possible these functions are annotated with the actual string / typeddict representation of the metadata form of the dtype. This means the type-checking routines that perform type-narrowing on JSON input now have a much more explicit typeguard annotation. But, it also means there is some impedance mismatch between the JSON type we are using in the rest of the codebase, and the typeddict universe of the dtypes. Ultimately I think we need to fix our JSON type to be compatible with typeddicts, but that's a problem for a later day.

• I separated the zarr v2 JSON data type deserialization from the zarr v3 deserialization. The key difference is that zarr v2 deserialization now takes an optional object_codec_id parameter, which is the name of an "object codec" present in the array metadata document. This parameter is necessary to resolve the data type for zarr v2 arrays saved with dtype "|O", which can only be disambiguated by checking for special codecs like vlen-utf8, vlen-bytes, pickle, etc in filters or compressor. The process of creating zarr v2 metadata documents from JSON has been updated to include this logic. This means we can add zarr-python 2.18-compatible vlen-bytes and vlen-array dtypes. That's an effort for another PR.

• I added a consistency check to the data type registry. When resolving a native data type, if two or more zdtype classes match that native data type, an exception is raised and the user is encouraged to unregister one of the two ambiguous zdtype instances from the registry.

  To be clear, I do not like that dtype resolution currently involves calling a bunch of class methods, succeeding only when exactly one of those class method invocations succeeds. Things would be much simpler with a mapping from dict from numpy data type classes to zarr data type classes, as @nenb suggested in here. But this is complicated by the numpy void dtype, which supports two separate zarr data types: fixed-length raw bytes and structured dtypes. One solution would be to add numpy void to the list of data types we don't dynamically resolve, and then I think the static {numpy dtype: zarr dtype} mapping can work. But I consider this a refinement of the content in this PR, and suitable for a separate effort.

I have not written enough documentation. I would request that we make this happen in another PR, and focus on getting this PR finally reviewed, merged, and included in a pre-release so people can test it more easily.

[nenb]

Expressing my personal opinion here: I would be very much in favour of getting this into a pre-release so that I can test it with a variety of new data types and see what happens. Like you have pointed out, testing in the pre-release will likely be the best way to motivate what parts of the documentation and the (internal) API need to be improved before an actual release.

Great work!

[d-v-b]

@ianhi @dstansby if you have the time could you check and see if the changes you requested have all been addressed?

[dstansby]

Sorry, I'm not going to have time to re-review this in the near future. If you've looked at all my previous comments feel free to dimsiss my review.

[d-v-b]

Sorry, I'm not going to have time to re-review this in the near future. If you've looked at all my previous comments feel free to dimsiss my review.

no worries, and thanks for your patience with this big PR. I'll review your feedback; some things might get spun out into issues (i'm thinking about changes in the config)