How to attach __record__ parameters to arrays from lazy-loaded parquet files without materializing?

[nikoladze]

I'd like to read a nested array lazily from parquet files and attach behavior to some of the record fields. For example if i would have

array = ak.zip(
    {
        "A" : ak.zip(
            {
                "x" : [[1, 2], [], [3, 4, 5]],
                "y" : [[10, 20], [], [30, 40, 50]]
            }
        ),
        "B" : ak.zip(
            {
                "x": [[1, 2, 3, 4], [5], [6]],
                "y": [[10, 20, 30, 40], [50], [60]]
            }
        )
    }, depth_limit=1
)

ak.to_parquet(array, "array.parquet")

i want to load this with

lazy = ak.from_parquet("array.parquet", lazy=True)

Now, for example if i want to set the __record__ parameter of the "A" record i could do something like

lazy["A"] = ak.with_name(lazy.A, "AClass")

but that will materialize the virtual array under "A". Skimming over the code for ak.from_parquet it seems manually reconstructing that VirtualArray will be quite involved.

Is there a straightforward solution for this?

[agoose77]

Setting the top-level __record__

The reason that this happens is that after setting the __record__ parameter, with_name re-visits the entire layout and simplifies any union types. This is useful, because it removes redundant unions, but has the side-effect of materialising the virtual array.

I wonder @jpivarski if it is sensible to implement a flag for recursively_apply to prevent it from recursing into virtual layouts. In this particular instance, if a virtual layout has not been materialised by the first getfunction then we shouldn't either. Of course, we don't need a flag for this; the second getfunction can add a case to handle it.

In the mean-time, unless Jim has a better solution, you can implement a non-simplifying variant of with_name. You can also just modify the layout directly, but this is cleaner!

def with_name_lazy(array, name, highlevel=True, behavior=None):
    def getfunction(layout):
        if isinstance(layout, ak.layout.RecordArray):
            parameters = dict(layout.parameters)
            parameters["__record__"] = name
            return lambda: ak.layout.RecordArray(
                layout.contents,
                layout.recordlookup,
                len(layout),
                layout.identities,
                parameters,
            )
        else:
            return None

    out = ak._util.recursively_apply(
        ak.operations.convert.to_layout(array), getfunction, pass_depth=False
    )
    if highlevel:
        return ak._util.wrap(out, ak._util.behaviorof(array, behavior=behavior))
    else:
        return out

Now, we should be able to detect whether a virtual array has been generated with layout.peek_cache. However, this doesn't seem to return anything for the arrays lazily loaded by from_parquet. I think this is because of how the caching is setup - I'm not an expert here yet! If you wrap the result with your own cache, e.g.

cache = {}
cached = ak.with_cache(lazy, cache)

Then you can create a simplifying variant of the with_name function that does not materialise virtual arrays:

def with_name_lazy(array, name, highlevel=True, behavior=None):
    def getfunction(layout):
        if isinstance(layout, ak.layout.RecordArray):
            parameters = dict(layout.parameters)
            parameters["__record__"] = name
            return lambda: ak.layout.RecordArray(
                layout.contents,
                layout.recordlookup,
                len(layout),
                layout.identities,
                parameters,
            )
        else:
            return None

    out = ak._util.recursively_apply(
        ak.operations.convert.to_layout(array), getfunction, pass_depth=False
    )

    def getfunction2(layout):
        if isinstance(layout, ak.layout.VirtualArray) and layout.peek_cache is None:
            return None

        if isinstance(layout, ak._util.uniontypes):
            return lambda: layout.simplify(merge=True, mergebool=False)
        else:
            return None

    out2 = ak._util.recursively_apply(out, getfunction2, pass_depth=False)

    if highlevel:
        return ak._util.wrap(out2, ak._util.behaviorof(array, behavior=behavior))
    else:
        return out2

[nikoladze]

Thanks for the quick answer! I should have mentioned that i already tried the "non-simplifying" variant with only the first call of recursively_apply. That's already better, but the part that descends down until the RecordArray still materializes the offsets. I'm checking the cache with lazy._caches:

lazy = ak.from_parquet("array.parquet", lazy=True)
named_lazy_A = with_name_lazy(lazy.A, "AClass") # using the function that does not have the simplify call
lazy._caches

Out[6]:
({'ak.from_parquet:0:off:A.list.item.x:A[0]': <ListOffsetArray64>
     <offsets><Index64 i="[0 2 2 5]" offset="0" length="4" at="0x7f1494000100"/></offsets>
     <content><RecordArray length="5">
         <field index="0" key="x">
             <NumpyArray format="l" shape="5" data="1 2 3 4 5" at="0x7f1494000080"/>
         </field>
     </RecordArray></content>
 </ListOffsetArray64>},)

[nikoladze]

I also attempted to rebuild the VirtualArray layout and modify the form, but that did not seem to work - the generator does not put the record name there:

import awkward as ak
import json

lazy = ak.from_parquet("array.parquet", lazy=True)

def rebuild_virtual_with_name(layout, name):
    generator = layout.generator
    args_list = list(generator.args)
    form = args_list[3]
    form_dict = json.loads(form.tojson())
    form_dict["content"]["parameters"]["__record__"] = "AClass"
    new_form = ak.forms.Form.fromjson(json.dumps(form_dict))
    args_list[3] = new_form
    return ak.Array(
        ak.layout.VirtualArray(
            ak.layout.ArrayGenerator(
                generator.callable,
                args=tuple(args_list),
                form=new_form,
            )
        )
    )

lazy_A = rebuild_virtual_with_name(lazy.A.layout, "AClass")

if i now try to access fields of lazy_A - i get ValueError: generated array does not conform to expected form: and one can see the generated form does not have the parameters.

[agoose77]

Ah, I originally (mis)read that you wanted to set the __record__ of the non-virtual RecordArray. Whoops!

I think that the answer here depends upon how generic you want to be. Ultimately, I don't think that there's any way to do this besides deferring the call to with_name. If you don't want to write the logic to traverse a form dict, you can just replace the virtual layout with one that hides the form information, and just assumes the virtual layout has a length. If you wanted to do this "properly", you could predict what this function would do to the form, and update that information.

Here's a simple implementation of the former:

from collections.abc import Iterable


def _deferred_with_name_lazy(layout, name, where):
    inner = ak.materialized(layout, highlevel=False)
    return with_name_lazy(inner, name, where, highlevel=False)


def with_name_lazy(base, name, where, highlevel=True, behavior=None):
    # Recursively move toward target
    if isinstance(where, str):
        where = [where]

    def getfunction(layout, depth, where):
        # If we reach a record array
        if isinstance(layout, ak.layout.RecordArray):
            # If we have no paths left to traverse
            if len(where) == 0:
                # Then add the parameter
                parameters = dict(layout.parameters)
                parameters["__record__"] = name
                return lambda: ak.layout.RecordArray(
                    layout.contents,
                    layout.recordlookup,
                    len(layout),
                    layout.identities,
                    parameters,
                )
            # Otherwise, continue recursion
            else:
                return where[1:]

        # If we hit a virtual array, we don't want to materialise it,
        # so we return another virtual array that hides
        elif isinstance(layout, ak.layout.VirtualArray):
            return lambda: ak.virtual(
                _deferred_with_name_lazy,
                (layout, name, where),
                length=len(layout),
                highlevel=False,
            )
        else:
            return where

    out = ak._util.recursively_apply(
        ak.operations.convert.to_layout(array), getfunction, pass_user=True, user=where
    )
    if highlevel:
        return ak._util.wrap(out, ak._util.behaviorof(array, behavior=behavior))
    else:
        return out

Then, to use this we have

annotated = with_name_lazy(lazy, "AClass", "A")

or for deeper RecordArrays,

annotated = with_name_lazy(lazy, "AClass", ("A", "child_of_A"))

[jpivarski]

This project has a long history of reorganizing methods to stop them from materializing VirtualArrays. The ultimate solution will be to integrate Awkward Array and Dask (see this talk) because Dask has a .compute() method for the user to be able to say "now is the time to compute."

But until that day comes, we can continue to fix specific operations that we feel should not cause a materialization. If this is a problem with simplify_union triggering a materialization, that can be fixed. The mergeability is checked by each Content class, and VirtualArray::mergeable simply materializes and then uses generated array to check mergeability. If the mergeability check were moved to the Form classes, rather than the Content classes, then VirtualArray::mergeable would be able to use its Form to do this check instead of materializing. (That's how other tests have been made non-materializing, by moving tests that do not depend on a full array into Forms.)

[agoose77]

This is effectively two problems:

1. the existing with_name implementation eagerly materialises VirtualArrays because it visits the entire layout
2. with_name needs to defer its operation if it encounters a VirtualArray during its visit.

I don't think (2) is something we necessarily want to be getting involved with because it opens a whole can of worms w.r.t deferring operations & modifying the form to reflect the pending change etc.

(1) definitely seems like something we could look at, because all we really need to do is make sure that we don't travel any farther than the depth of the layout that was renamed.

[jpivarski]

I'll suggest an alternative.

[jpivarski]

@nikoladze, maybe you can avoid with_name by constructing a new array? There's new documentation on that:

https://awkward-array.org/how-to-create-constructors.html#content-recordarray

Since your array contains VirtualArrays, you need to keep the original ak.Array in scope until you've made a new ak.Array, because the VirtualArrays only keto weak reference to any attached caches:

https://awkward-array.org/how-to-create-constructors.html#content-virtualarray

If your use of with_name attached new arrays more deeply than the top level, it would have had to materialize anyway: in order to zip nested lists, it would have needed the offset arrays of the other lists, to be sure it's broadcasted correctly.

[nikoladze]

@jpivarski, @agoose77, thanks a lot for your detailed answers! I did not really think about using the existing virtual array to generate a new virtual array, but maybe that is the way to go here. But i think i haven't quite understood how the nesting of virtual arrays works. I tried the following:

lazy = ak.from_parquet("array.parquet", lazy=True)

leaf_form = {
    "class": "NumpyArray",
    "itemsize": 8,
    "format": "l",
    "primitive": "int64",
}

def get_virtual_content(virtual_col):

    def generate():
        return ak.materialized(virtual_col).layout.content

    return ak.layout.VirtualArray(
        ak.layout.ArrayGenerator(
            generate,
            form=ak.forms.Form.fromjson(json.dumps(leaf_form))
        )
    )

def get_virtual_record(lazy_col, with_name):
    form = ak.forms.Form.fromjson(json.dumps({
        "class": "ListOffsetArray64",
        "offsets": "i64",
        "content": {
            "class": "RecordArray",
            "contents": {
                "x": leaf_form,
                "y": leaf_form,
            }
        },
        "parameters": {
            "__record__": with_name
        }
    }))

    def generate():
        return ak.layout.ListOffsetArray64(
            ak.materialized(lazy_col.x).layout.offsets,
            ak.layout.RecordArray(
                [
                    get_virtual_content(lazy_col.x.layout),
                    get_virtual_content(lazy_col.y.layout)
                ],
                ["x", "y"]
            ),
            parameters={"__record__": with_name}
        )

    return ak.Array(
        ak.layout.VirtualArray(
            ak.layout.ArrayGenerator(generate, length=len(lazy_col), form=form)
        )
    )

lazy_A = get_virtual_record(lazy.A, "AClass")

However if i now access e.g lazy_A.x it seems also the virtual arrays for the x and y leaf nodes get materialized. I think it happens in the constructor of ak.layout.RecordArray ...

In general i was hoping to get a more straightforward solution to read a coffea.nanoevents - like array. Parquet already allows storing the nested data, awkward array has ak.from_parquet with lazy=True, so in principle that already brings us halfway there. The next step would then to attach names to the collections.

Tagging @nsmith- maybe he has some ideas as well.

[jpivarski]

If the tree structure is VirtualArray → RecordArray → {x: NumpyArray, y: NumpyArray}, then you can't materialize x without materializing y. A VirtualArray is all-or-nothing; it is either materialized or not.

Therefore, the following is a common pattern: RecordArray → {x: VirtualArray → NumpyArray, y: VirtualArray → NumpyArray}. Accessing x won't necessarily access y. NanoEvents uses this pattern, as does ak.from_buffers and ak.from_parquet.

Here, for example, is a utility class (that @nsmith- wrote) to insert VirtualForms into the field of every RecordForm (under the record, not over it). https://github.com/scikit-hep/awkward-1.0/blob/346662d894df1ecab7318ba4c2bf8225ac7b94c6/src/awkward/operations/convert.py#L4792-L4813

[nikoladze]

mhh - i thought i already did RecordArray → {x: VirtualArray → NumpyArray, y: VirtualArray → NumpyArray} since my get_virtual_content function returns a virtual array?

In any case - since ak.from_buffers does it right i could also use that (with a mapping that makes use of the existing lazy array) - the following seems to work:

import awkward as ak

lazy = ak.from_parquet("array.parquet", lazy=True)

def get_leaf_form(form_key):
    return {
        "class": "NumpyArray",
        "itemsize": 8,
        "format": "l",
        "primitive": "int64",
        "form_key": form_key,
    }

def get_collection_form(col_name, with_name):
    return {
        "class": "ListOffsetArray64",
        "offsets": "i64",
        "content": {
            "class": "RecordArray",
            "contents": {
                "x": get_leaf_form(f"{col_name}.x"),
                "y": get_leaf_form(f"{col_name}.y"),
            },
            "parameters": {
                "__record__": with_name,
            }
        },
        "form_key": f"{col_name}.x",
    }

form = {
    "class": "RecordArray",
    "contents": {
        "A": get_collection_form("A", "AClass"),
        "B": get_collection_form("B", "BClass"),
    }
}

class Mapping:

    def __getitem__(self, key):
        part, array_key, what = key.split("-")
        array = ak.materialized(lazy[tuple(array_key.split("."))])
        if what == "offsets":
            return array.layout.offsets
        elif what == "data":
            return array.layout.content

lazy_with_names = ak.from_buffers(form, len(lazy), Mapping(), lazy=True)

[jpivarski]

I was basing my assessment on the fact that your get_virtual_record describes a RecordForm that contains the NumpyForms.

Using ak.from_buffers directly works.

Here is a complete example, in case it helps:

import cachetools

import awkward as ak
import numpy as np

def generate_x():
    print("generating x")
    return ak.Array([1.1, 2.2, 3.3, 4.4, 5.5])

def generate_y():
    print("generating y")
    return ak.Array([10, 20, 30, 40, 50])

form_x = ak.forms.Form.from_numpy(np.dtype(np.float64))
form_y = ak.forms.Form.from_numpy(np.dtype(np.int64))

cache = cachetools.LRUCache(np.inf)

recordarray = ak.layout.RecordArray(
    [
        ak.layout.VirtualArray(
            ak.layout.ArrayGenerator(generate_x, length=5, form=form_x),
            ak.layout.ArrayCache(cache),
        ),
        ak.layout.VirtualArray(
            ak.layout.ArrayGenerator(generate_y, length=5, form=form_y),
            ak.layout.ArrayCache(cache),
        ),
    ],
    ["x", "y"]
)

array = ak.Array(recordarray)

(such a thing could also be set up using ak.virtual, but I'm going the low-level route because I think this is meant to be integrated into something else low-level).

[nsmith-]

Sorry to be so late to the party. It seems to me the magic in NanoEvents is that the form can be manipulated to add __record__ and any other parameter at any level before passing to ak.from_buffers, which will dutifully copy them into the materialized layout: https://github.com/scikit-hep/awkward-1.0/blob/346662d894df1ecab7318ba4c2bf8225ac7b94c6/src/awkward/operations/convert.py#L4356

You start from ak.from_parquet where you have no opportunity to manipulate the form. Ideally the parquet schema would carry this information and https://github.com/scikit-hep/awkward-1.0/blob/346662d894df1ecab7318ba4c2bf8225ac7b94c6/src/awkward/operations/convert.py#L3146 would read it and act as appropriate. I don't know if arrow schemas have such a parameters concept.

[jpivarski]

Arrow and Parquet have string-based metadata, which is where the parameters should eventually go (but haven't yet).