Zero-copy view of offsets and flattened data Numpy arrays as 2D Awkward arrays

[gipert]

I'm trying to write a robust and zero-copy de-serializer of 2D jagged numerical vectors, written to disk according to this specification, into Awkward 2.x arrays. This is what I have written so far:

def _ak_from_buffers(flattened_data: NDArray, cumulative_length: NDArray):
    # should also run ak.is_valid() before returning?
    return ak.from_buffers(
        {
            "class": "ListOffsetArray",
            "offsets": "i64",  # how to get this from cumulative_length.dtype?
            "content": {
                "class": "NumpyArray",
                "primitive": ak.types.numpytype.dtype_to_primitive(flattened_data.dtype),
                "form_key": "node1",
            },
            "form_key": "node0",
        },
        len(cumulative_length) - 1,
        {
            "node0-offsets": cumulative_length,  # how to prepend the zero but avoid copy?
            "node1-data": flattened_data,
        },
    )

and I have a couple of questions (also inlined in the code):

1. The value of "offsets" in the Form is hardcoded, is there a way to get it from cumulative_length.dtype?
2. The offsets in the container should start from zero, but the input cumulative_length does not. Is there a way to fix this without making a copy?
3. Is there anything else that can be improved?

[agoose77]

Hi @gipert, interesting discussion!

1. We don't publicly expose a direct mapping from dtype to offsets. I'd suggest just writing a mapping yourself (it won't be very big). Also, see (3). Note that our kernels only support a small set of dtypes, so anything that's not int64 will often be copied once the array is manipulated.
2. We don't provide any layout types with an implicit zero start index. I think you're going to have to copy here, unfortunately.
3. You don't need to use ak.from_buffers here. You can also directly construct the layout nodes using our layout system. This would mean that you didn't need to figure out the index-type associated with a dtype.
   e.g.

offsets = np.empty(len(cumulative_length) + 1, dtype=cumulative_length.dtype)
offsets[1:] = cumulative_length
offsets[0] = 0

layout = ak.contents.ListOffsetArray(
    offsets=ak.index.Index(offsets),
    content=ak.contents.NumpyArray(flattened_data)
)

[gipert]

Thanks a lot! I guess I just initialize the array as ak.Array(layout) then?

[gipert]

Bonus question: how could your example be extended to support vectors of the same type but arbitrary dimensionality? In case of ak.from_buffers(), I could just add extra offset arrays to the container dictionary.

[agoose77]

If you have more than one level of nesting, then you'd implement this as an iteration or recursion, building from the leaf node (NumpyArray) to the root list.

import numpy as np

def list_from_cumulative_lengths(cumulative_length, content):
  offsets = np.empty(len(cumulative_length) + 1, dtype=cumulative_length.dtype)
  offsets[1:] = cumulative_length
  offsets[0] = 0
  
  return  ak.contents.ListOffsetArray(
      offsets=ak.index.Index(offsets),
      content=content
  )

leaf = ak.contents.NumpyArray(flattened_data)
list_1 = list_from_cumulative_lengths(cumulative_length_1, leaf)
list_2 = list_from_cumulative_lengths(cumulative_length_2, list_1)
...
list_n = list_from_cumulative_lengths(cumulative_length_n, list_n_sub_1)

Does that help at all? :)

[gipert]

Yes thanks! Neat!