Wrapping ak.Array

[fleble]

Dear experts,

For various reasons, I wrote a class inheriting from ak.Array implementing a few methods. For instance:

class Wrapper(ak.Array):
    def __init__(self, data):
        super().__init__(data)

    def test(self):
        print("OK")

ak_array = Wrapper([[1,2], [], [3]])
ak_array.test()

This worked with version 1.5.1 of awkward array, but does not work anymore with version 1.7.0. The error is the following:

  File "/home/fleble/miniconda3/envs/svj/lib/python3.8/site-packages/awkward/highlevel.py", line 1129, in __getattr__
    raise AttributeError(
AttributeError: no field named 'test'

Indeed, the type of ak_array is <class 'awkward.highlevel.Array'>.

Is there a way to write a class deriving from ak.Array or should this be prohibited?

[agoose77]

This change was caused by ccd7537.

I don't know what @jpivarski expects to happen with custom array classes in this context. A solution that is guaranteed to work in future is the __array__ parameter:

ak.behavior['wrapper_cls'] = Wrapper

ak_array = ak.with_parameter(
  [[1,2], [], [3]],
   "__array__", 
  "wrapper_cls"
)

Awkward uses the behaviours mechanism to allow an association between methods and data, without binding them together explicitly. This mechanism extends to records, i.e. collections of arrays, which can provide OO interfaces by defining a particular behaviour class for the record. For more information, see here: https://awkward-array.readthedocs.io/en/latest/ak.behavior.html

[agoose77]

@grst in general, you would be encouraged not to do this. Awkward makes new versions of arrays all over the place, and it does this by assuming that arrays are layouts + behaviour dictionary. This is intentional; the behaviour mechanism is like a lightweight view over an existing array, so we don't want any additional mechanisms for configuration besides the array itself.

You might use the existing parameters method to attach metadata to an array. I don't believe we have any strong policy about what can be put in there yet (ping @jpivarski). Additionally, I believe that most operations that would cause these parameters to be lost would also lead to the __array__ parameter that associates the wrapper with the particular array to be lost. The reverse is perhaps not true though; we might somewhere remove only the __array__ parameter. This likely wouldn't be a problem in this case, though.

Could you perhaps expand on your use case here? :)

[grst]

Our use case is related to scverse/anndata#647: AnnData is a container with several matrices that are aligned to each other and can be numpy arrays, sparse arrays, pandas data frames, or now also awkward arrays.

We are implementing "views" with a copy-on-write mechanism for AnnData. When slicing an AnnData object, it becomes an AnnDataView, when a user tries to modify the view, it is implicitly copied.

# X = main matrix, obsm = matrices aligned to axis 0 ("observations") of X
>>> adata= anndata.AnnData(X = np.ones((4,4)), obsm={"test": np.ones((4, 2))})
>>> adata
AnnData object with n_obs × n_vars = 4 × 4
    obsm: 'test'

# Let's create a view by slicing
>>> view = adata[:3, ]
>>> view 
View of AnnData object with n_obs × n_vars = 3 × 4
    obsm: 'test'

# Let's update the "test" matrix in `.obsm`
>>> view.obsm["test"]
ArrayView([[1., 1.],
           [1., 1.],
           [1., 1.]])
>>> view.obsm["test"][1, :] = 42
/home/sturm/tmp/ipykernel_57822/3480696444.py:1: ImplicitModificationWarning: Trying to modify attribute `.obsm` of view, initializing view as actual.
  view.obsm["test"][1, :] = 42

# Now `view` is a regular AnnData object and not a view anymore. 
>>> view
AnnData object with n_obs × n_vars = 3 × 4
    obsm: 'test

With the changes introduced in the PR above, .obsm could contain an AwkwardArray:

>>> adata.obsm["awk"] = ak.Array([{'a': 0, 'b': 1}] * 4)

When we change that array, we need to "notify" the AnnDataView that holds the view of the awkward array, to turn itself into a copy.

# Now we need to trigger the `ImplicitModificationWarning` and make an actual copy
>>> adata.obsm["awk"]["b"] = np.array([42, 42, 42, 42])

The way we currently solve this is to have a View Class for each data type (e.g. SparseCSRView) with a _SetItemMixin that overrides the __setitem__ function. This overridden __setitem__ function makes use of a reference to the AnnData view to trigger the copy. This reference is what I would want to pass as parameter to the AwkwardArrayView which I tried implementing as a "behavior".

I hope I managed to give enough context without going too much into anndata-specific details. LMK if anything is unclear.

[agoose77]

I'm sure I could find it in the PR, but it would be easier for me to ask here - does adata.obsm["awk"] = ak.Array([{'a': 0, 'b': 1}] * 4) store the ak.Array or the underlying layout in your AnnData container?

Assuming it stores the ak.Array, then on the face of it, parameters would seem useful here. We have a restriction that parameters needs to be JSON serialisable. So, I suggest a weak registry:

import awkward._v2 as ak
import weakref

_registry = weakref.WeakValueDictionary()


ANNDATA_KEY_PARAMETER = "_anndata_listener_key"

def find_listener(key: str):
    return _registry[key]


def compute_key(array: ak.Array, listener) -> str:
    return f"target-{id(listener)}"


def register_listener(array: ak.Array, listener):
    key = compute_key(array, listener)
    _registry[key] = listener
    return ak.with_parameter(array, ANNDATA_KEY_PARAMETER, key)


class ArrayView(ak.Array):
    def _handle_missing_key(self, name, value):
        pass

    def _handle_invalid_key(self, name, value, key):
        pass

    def __setitem__(self, name, value):
        key = self.layout.parameter(ANNDATA_KEY_PARAMETER)
        if key is None:
            self._handle_missing_key(name, value)
        else:
            try:
                ref = find_listener(key)
            except KeyError:
                self._handle_invalid_key(name, value, key)
            else:
                ref.notify(name, value, self)
        super().__setitem__(name, value)


ak.behavior['*', 'ArrayView'] = ArrayView


class Listener:
    def notify(self, name, value, notifier):
        print(f"Notified by {notifier} for {name}={value!r}")

# Associate target with array
listener = Listener()

array = register_listener(
    ak.zip({'x': [[1, 2, 3], [4, 5, ]]}, with_name="ArrayView"),
    listener
)

array['x'] = array.x * 2

Would something like this work?

[grst]

'm sure I could find it in the PR, but it would be easier for me to ask here - does adata.obsm["awk"] = ak.Array([{'a': 0, 'b': 1}] * 4) store the ak.Array or the underlying layout in your AnnData container?

Correct, it stores the ak.Array!

@ivirshup, what do you think of this approach? Looks reasonable to me!

[jpivarski]

There's no limitation on the scope of what goes into parameters except

• dunder keys (__array__, __record__, __doc__, ...) are reserved, since they have some active meaning (interpreted by Awkward internals)
• parameters have to be JSON (so that they can be serialized and saved forever or transmitted to workers on a network).

Virtually every operation takes an array and makes a new array from it: a lot of attention is given in the code to pass on parameters through all the operations in which the output is "the same array" in some sense. Since you want metadata that follows the array around, this sounds like the right place to put it.

But also, you're using this to make weak references, which are very transient, non-JSON-like things. That argues against making them parameters. @agoose77's solution of using parameters to store keys of a weak reference registry is a good one, though if you ever serialize these arrays, they'll retain keys that point to a registry that no longer exists. If you're ever reading and writing or transmitting arrays with these keys, be sure to distinguish between keys relevant for this Python process with its global registry of transient weak references and some other one in another Python process. (Maybe always strip the keys out before serializing?)

On the other other hand, maybe this is a job for the new ak.Array-level metadata proposed in #1391. (I just mentioned this issue a few minutes ago, and it's high in our priority list.) This metadata would follow "the same array" around as ak.Array instances beget new ak.Array instances, like parameters, but it wouldn't be part of the serialized array, wouldn't have the JSON restriction, and therefore could hold transients like weak references. As in the discussion a few minutes ago, we want this mechanism to be general enough that it supports a variety of use-cases. (The original use-case is just to not lose xarray metadata, if we ever convert to or from xarrays.)

[jpivarski]

Thanks, @agoose77, for answering this! You made all the points I wanted to except the "why".

You can make subclasses of ak.Array and ak.Record, but you have to assign them through ak.behavior. There's a big write-up of that here:

https://awkward-array.readthedocs.io/en/latest/ak.behavior.html

When you do array manipulations, you end up creating new array objects all the time—every slice is a new array object from Python's point of view (though the data buffers themselves share a lot of memory for efficiency). The ak.behavior is a set of rules for applying your subclasses to new array objects depending on "parameters" in the array, so even if your data type is buried deeply in a data structure (and there are no Python classes associated with it), when it comes to the surface through a slice or something, it will be wrapped by the right Python class (which must be a subclass of ak.Array or ak.Record, depending on what you're doing).

This change was caused by ccd7537.

Oh, yeah, the old behavior worked by accident. Before that fix, it wasn't checking ak.behavior, which let the class be stale (and therefore wrong) in some cases and allowed your Wrapper without going through the ak.behavior mechanism, @fleble. The new version revealed unintended behavior that you were using. Switching to ak.behavior would make it more stable anyway: if you have records inside the Array that need to be interpreted a particular way, this would wrap them any time they come to the surface through a slice, not just when you explicitly call Wrapper.

[fleble]

@agoose77 and @jpivarski thank you very much for your answers. I had completely missed the point of ak.behavior. I see it now. Thank you very much for the very clear, thorough and very fast reply!

If my class also defines some attributes, and a class inherits from another one, I figured out I can proceed like this:

class Wrapper0(ak.Array):
    def __init__(self, data):
        ak.behavior["Wrapper0"] = Wrapper0
        if "__array__" not in ak.parameters(data):
            data = ak.with_parameter(
                data,
                "__array__",
                "Wrapper0",
            )
        super().__init__(data)
        self.attr0 = "attr0"

    def test0(self):
        print("0")


class Wrapper1(Wrapper0):
    def __init__(self, data):
        ak.behavior["Wrapper1"] = Wrapper1
        if "__array__" not in ak.parameters(data):
            data = ak.with_parameter(
                data,
                "__array__",
                "Wrapper1",
            )
        super().__init__(data)
        self.attr1 = "attr1"

    def test1(self):
        print(self.fields)


data = {
    "Jet_pt": [[1,2], [], [3]],
    "MET": [1, 2, 3],
}
ak_array = Wrapper1(ak.Array(data))

ak_array.test0()
ak_array.test1()
print(ak_array.attr0)
print(ak_array.attr1)

(posting in case it is of interest for someone or in case you have a comment about it).

Thank you very much again.

[agoose77]

I would advise against registering the behaviour on the initialiser of your classes. It is better to do this outside of the class, e.g.

class Wrapper0(ak.Array):
    def test0(self):
        print("0")

	@property
    def attr0(self):
        return "attr0"

ak.behavior['Wrapper0'] = Wrapper0


class Wrapper1(Wrapper0):

    def test1(self):
        print(self.fields)

	@property
    def attr1(self):
        return "attr1"

ak.behavior['Wrapper1'] = Wrapper1


data = {
    "Jet_pt": [[1,2], [], [3]],
    "MET": [1, 2, 3],
}
ak_array = ak.with_parameter(ak.zip(data, depth_limit=1), "__array__", "Wrapper1")

ak_array.test0()
ak_array.test1()
print(ak_array.attr0)
print(ak_array.attr1)

However, there's more to ak.behavior than just __array__. If your behaviour corresponds to operations on fields, i.e. if you want to define a class for a particular collection of fields, there's a very similar parameter called __record__. This allows you to specify behavior for singular records as well as arrays of records (RecordArrays).

class MyRecord(ak.Record):
    def what_am_i(self):
        print("I am a record")

ak.behavior['MyRecord'] = MyRecord


class MyRecordArray(ak.Array):
    def what_am_i(self):
        print("I am an array of record")

ak.behavior['*', 'MyRecord'] = MyRecordArray



data = {
    "Jet_pt": [[1,2], [], [3]],
    "MET": [1, 2, 3],
}
ak_array = ak.zip(data, with_name="MyRecord", depth_limit=1)

Then we have

>>> ak_array.what_am_i()
I am an array of record
>>> ak_array[0].what_am_i()
I am a record

It is useful to have separate classes for records and arrays when you can't use the same logic e.g., if you are subscripting one of your fields:
This works for arrays

>>> ak_array.MET[0] # OK
1

But fails for records

>>> ak_array[0].MET[0] # NOT OK
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-27-75468926cd2d> in <module>
      1 ak_array.MET[0] # OK
----> 2 ak_array[0].MET[0] # NOT OK

TypeError: 'int' object is not subscriptable

However, often your methods / properties work equally well upon records and arrays. In these cases, there's a useful mixin_class decorator to reduce the amount of boilerplate:

@ak.mixin_class(ak.behavior) # Register with ak.behavior for arrays and records!
class MyRecord:
    def what_am_i(self):
        print("I am an array / record")


data = {
    "Jet_pt": [[1,2], [], [3]],
    "MET": [1, 2, 3],
}
ak_array = ak.zip(data, with_name="MyRecord", depth_limit=1)

Then we have

>>> ak_array.what_am_i()
I am an array / record
>>> ak_array[0].what_am_i()
I am an array / record

In terms of naming, you don't really want Record or Wrapper in your names. Instead, something meaningful e.g. Point3D that represents the kind of object you have.