awkward-kaitai-lazy.py

# https://gist.github.com/agoose77/e1567e9c7ee328c2da9b2a81dc9c1c7d
from __future__ import annotations

from collections.abc import Iterator

import awkward_katai as ak_ka

import awkward as ak


def ak_ka_iterate(array: ak.Array, step_size: int) -> Iterator[ak.Array]:
    layout = ak.to_layout(array, ...)

    virtual_key = layout.purelist_parameter("_virtual_key")
    readers = array.attrs['@readers']
    reader = readers[virtual_key]

    while True:
        yield reader(step_size)


array = ak_ka.load("...", lazy={'@geometry.coordinates'})

# for geometry_coordinates in ak_ka.materialize_in_steps(array.geometry.coordinates, step_size=1_000):
#     geometry_coordinates.show()
    
for geometry_coordinates in array.attrs['@geometry.coordinates'].iterate(step_size=1_000):
    geometry_coordinates.show()

# array_without_large_fields, geometry_coordinates_view = ak_ka.load_partial(..., step_size=1000)
# for geometry_coordinates in geometry_coordinates_view:
#     ...


load_context = {"geom": ak_ka.LazyLoader("geometry.coordinates")}
array = ak_ka.load("...", context=load_context)
    
for geometry_coordinates in load_context["geom"].iterate(step_size=1_000):
    geometry_coordinates.show()