Going over all events once without a cache - speed issue

[Teddy-Curtis]

Hey!

So I am having a slight problem with the speed of reading events from a parquet dataset.
Essentially, I have a large dataset in parquet files that I need to loop over to save individually for Pytorch.
Example data point:
ak.Record of :

{'node_data' : [1,4,2,5,7,8,....],  'edge_data' : [[1,2], [0,1], .....], .....}

Now the data is too large to be loaded in at once so I read it lazily:

dataset = ak.from_parquet("Documents/Data", lazy=True)
for data in dataset:
    do.....

However, with this, the cache ends up getting bigger and bigger until it crashes, so instead I try:

dataset = ak.from_parquet("Documents/Data", lazy=True, lazy_cache=None)
for data in dataset:
    do.....

But now, even though I'm still only reading each datapoint once, it runs far slower than without setting lazy_cache=None. Is there any reason for this to be so much slower even though I am only reading each data point once?

Does anyone have any advice on how I should proceed?
I'm guessing there might be a way to limit the amount that is cached so it still runs quickly without using up more and more ram, however I'm still a bit confused as to why the speed should change between the two scenarios.

Any help would be extremely appreciated!
Thanks!

[agoose77]

Without seeing your file, I can't give a detailed explanation as to why this happens. But, roughly:

When you pass lazy=True, Awkward reads the structure of the array, but not the contents. For partitioned arrays (Parquet files with row groups), this means that Awkward can read only a single row group if you slice it appropriately. The caching mechanism means that if you try and read two parts of the array that refer to the same portion of the underlying file, it will re-use the data from the cache, saving you a read() call.

Setting lazy_cache=None means that Awkward must request the data from the file whenever the contents need to be resolved, because it does not use a cache. This will be slower in cases where you're reading overlapping parts of the same file. In your case, if a row-group spans more than one row of your file, it will re-read the entire row group for each iteration.

What you probably want to do here is read a full row-group, and handle the laziness manually, e.g.

def iter_rows_lazy(file, **kwargs):
    array = ak.from_parquet(file, lazy=True, **kwargs)
    n_row_groups = len(ak.partitions(array))
   
    for i in range(n_row_groups):
        yield from ak.from_parquet(file, row_groups=i, **kwargs)
 
 
 for row in iter_rows_lazy("Documents/Data"):
     do_something(row)

I've not tested this yet, so it might need some tweaks.

[Teddy-Curtis]

Ahhhh OK that makes a lot of sense now!
However, even with the function you gave above:

def iter_rows_lazy(file, **kwargs):
    array = ak.from_parquet(file, lazy=True)
    n_row_groups = len(ak.partitions(array))
   
    for i in range(n_row_groups):
        yield from ak.from_parquet(file, row_groups=i, **kwargs)

I'm still getting a memory error after about 30,000 iterations: (skipping a couple of the first lines

Traceback (most recent call last):
...
...
...
  File "HGCal/GNN/Data/directed2LayersK5/datasetClass.py", line 82, in process
    for i, graph in enumerate(iter_rows_lazy(file)):
  File "HGCal/GNN/Data/directed2LayersK5/datasetClass.py", line 74, in iter_rows_lazy
    yield from ak.from_parquet(file, row_groups=i, **kwargs)
  File "/home/user1/miniconda3/envs/tGPU/lib/python3.8/site-packages/awkward/operations/convert.py", line 3960, in from_parquet
    batches = dataset.read_row_group_batches()
  File "/home/user1/miniconda3/envs/tGPU/lib/python3.8/site-packages/awkward/operations/convert.py", line 3469, in read_row_group_batches
    batches.extend(self.read_row_group(i, columns).to_batches())
  File "/home/user1/miniconda3/envs/tGPU/lib/python3.8/site-packages/awkward/operations/convert.py", line 3678, in read_row_group
    return single_file.read_row_group(
  File "/home/user1/miniconda3/envs/tGPU/lib/python3.8/site-packages/pyarrow/parquet/__init__.py", line 421, in read_row_group
    return self.reader.read_row_group(i, column_indices=column_indices,
  File "pyarrow/_parquet.pyx", line 1334, in pyarrow._parquet.ParquetReader.read_row_group
  File "pyarrow/_parquet.pyx", line 1353, in pyarrow._parquet.ParquetReader.read_row_groups
  File "pyarrow/error.pxi", line 117, in pyarrow.lib.check_status
pyarrow.lib.ArrowMemoryError: realloc of size 1073741824 failed

So I get the idea behind it, but it seems like it isn't completely refreshing the cache every time this is called:

for i in range(n_row_groups):
    yield from ak.from_parquet(file, row_groups=i, **kwargs)

So maybe it's easier doing the other way with my own cache that I reset say every time a new row group is opened?
Would this be done but just passing cache = {} into lazy_cache = cache then clearing that?
For example I've quickly written this:

def iter_lazy(file, **kwargs):
    cache = {}
    dataset = ak.from_parquet(file, lazy=True, lazy_cache=cache, **kwargs)
    partitions = ak.partitions(dataset)
    part_idx = 0
    part = partitions[part_idx]
    for i, data in enumerate(dataset):
        if i > part:
            cache = {}
            part_idx += 1
            part += partitions[part_idx]
        yield data

Sorry for the mountain of questions!

[agoose77]

That looks like quite a big re-allocation. Do you have many columns in your dataset? You can

pip install parquet-tools
parquet-tools inspect /path/to/parquet-file.parquet

to inspect your file and e.g. determine how many row groups it has. It seems like you have a very large row group, so I'm wondering if the files are themselves not properly partitioned.

[Teddy-Curtis]

Using parquet-tools on my file gives me:

############ file meta data ############
created_by: parquet-cpp-arrow version 9.0.0
num_columns: 7
num_rows: 11141
num_row_groups: 1
format_version: 2.6
serialized_size: 2168


############ Columns ############
item
item
item
gen_matched
pt
bdteg
quality

############ Column(item) ############
name: item
path: node_pos.list.item.list.item
max_definition_level: 2
max_repetition_level: 2
physical_type: DOUBLE
logical_type: None
converted_type (legacy): NONE
compression: SNAPPY (space_saved: 12%)

############ Column(item) ############
name: item
path: node_data.list.item.list.item
max_definition_level: 2
max_repetition_level: 2
physical_type: DOUBLE
logical_type: None
converted_type (legacy): NONE
compression: SNAPPY (space_saved: 3%)

############ Column(item) ############
name: item
path: edge_data.list.item.list.item
max_definition_level: 2
max_repetition_level: 2
physical_type: INT64
logical_type: None
converted_type (legacy): NONE
compression: SNAPPY (space_saved: 19%)

############ Column(gen_matched) ############
name: gen_matched
path: gen_matched
max_definition_level: 0
max_repetition_level: 0
physical_type: INT64
logical_type: None
converted_type (legacy): NONE
compression: SNAPPY (space_saved: -4%)

############ Column(pt) ############
name: pt
path: pt
max_definition_level: 0
max_repetition_level: 0
physical_type: DOUBLE
logical_type: None
converted_type (legacy): NONE
compression: SNAPPY (space_saved: 24%)

############ Column(bdteg) ############
name: bdteg
path: bdteg
max_definition_level: 0
max_repetition_level: 0
physical_type: DOUBLE
logical_type: None
converted_type (legacy): NONE
compression: SNAPPY (space_saved: 18%)

############ Column(quality) ############
name: quality
path: quality
max_definition_level: 0
max_repetition_level: 0
physical_type: INT64
logical_type: None
converted_type (legacy): NONE
compression: SNAPPY (space_saved: 7%)

Seems a bit weird that a couple of the columns are called 'item' instead they should be called node_pos, node_data, edge_data. Maybe that's what's causing me the issue? The *_data columns can be quite large but I only have 11141 rows and 7 columns. The file sizes aren't huge either with the biggest ones being ~50mb.
Also, if I output ak.partitions(dataset) for my entire dataset I get:

[11141, 13500, 13500, 13500, 13500, 11141, 13500, 13500, 13500, 13500, 11141, 13500, 13500, 13500, 13500, 13500, 13500, 3622, 11141, 11141, 13500, 222, 11141, 11141, 11141, 13500, 11141, 11141]

So none are particularly huge compared to any others

[agoose77]

Sorry, I should have clarified - can you do this for all of your files? The file you uploaded is fairly small, but I'm wondering if one of the files is much larger / unpartitioned vs the others.

RE the item name, it's the way that the Parquet schema describes a list. I assume you have a list-of-lists in node_pos, e.g. a list of 3-coordinate vectors.

[Teddy-Curtis]

RE the item name, it's the way that the Parquet schema describes a list. I assume you have a list-of-lists in node_pos, e.g. a list of 3-coordinate vectors.

Ok that's exactly what I have so that should be fine.
As for the other files, if I find num_row_groups, they all equal 1, the most amount of rows is 13500 and the largest file size is 50.1MiB.
Maybe I could try and split all the files into smaller ones first and then try and run the same loop and see if that does anything?

[jpivarski]

Thanks for answering this, @agoose77! The one point I want to add is that Awkward 1.x's lazy arrays are being removed: this is the biggest change from 1.x to 2.x. The v2 from_parquet function doesn't have a lazy option, but it does let you specify which columns and row groups you want to select from the file.

Also, dask-awkward runs on v2, so you'd be replacing any workflows like, "Load lazy arrays and try to keep them from accidentally reading the whole file while performing computations on a subset of the file," with workflows like, "Load Dask arrays, write computation, then call .compute() to have it load what it needs to, given that it can see your whole computation graph before it starts." This is a heads-up in case you want to avoid your current problems by trying out the new version.