return to main post

1. (particle) data IO

Full notebook available here.

yt reads particle and grid-based data by iterating across the chunks, with frontend-specific IO functions. For gridded data, each frontend implements a _read_fluid_selection (e.g., yt.frontend.amrvac.AMRVACIOHandler._read_fluid_selection) that iterates over chunks and returns a flat dictionary with numpy arrays concatenated across each chunk. The particle data, frontends must implement a similar function, _read_particle_fields, that typically gets invoked within the BaseIOHanlder._read_particle_selection function. In both cases, the read functions accept the chunks iterator, the fields to read and a selector object:

def _read_particle_fields(self, chunks, ptf, selector):

def _read_fluid_selection(self, chunks, selector, fields, size):

On reading a single chunk, the selector is applied so that we only return selected data from each chunk.

In order to construct a dask.delayed() Task Graph for IO, there are a number of changes required. First, dask uses pickle to serialize the functions and arguments that get distributed to workers, so all arguments to the _read functions above must be pickleable. A recent PR (PR #2934) added __getstate__ and __setstate__ methods to the geometric Cython selector objects, but the chunks can be more challenging. We touch on this further in the following section on the Particle IO prototype.

Second, the _read_particle_fields and _read_fluid_selection typically iterate over chunks internally, but it is more straightforward to construct a chunk-parallel dask process with a single-chunk read function. In the following Particle IO prototype, we avoid a large scale refactor by calling _read_particle_fields repeatedly with a list of a single-chunk, i.e.,_read_particle_fields([chunk],ptf,selector), but separating the chunk iteration and concatenation into a single chunk read function may improve maintainability.

Finally, one question on IO is what to return? The present read routines return a dictionary of numpy arrays. These arrays are typically pre-allocated though the recent PR #2416 removes particle pre-allocation. But in terms of dask, we could return a distritubed dask.array or dask.dataframe, allowing subsequent functions to compute in parallel. The dask.array needs to know chunk sizes a priori, which presents some difficulty. A dask.dataframe does not need to know the chunk sizes, just the expected data type for each column. So at present, in the following prototype, we construct a Task Graph that constructs a dask.dataframe from delayed chunk reads. In order to return the dictionary of flat numpy arrays expected by _read_particle_fields, prototype reduces the distributed dataframes to standard numpy arrays but in principle we could return the dask.dataframes, allowing subsequent calculations to leverage distributed calculations by using the distributed dataframes directly.

The dxl yt-dask-experiment repo includes some initial chunking experiments with dask in the dask_chunking directory. That folder primarily experiments with using dask to do some manual IO on a gadget dataset. In this experiment, the gadget file reading is split into a series of dask-delayed operations that include file IO and selector application. The notebook here gives an overview of that attempt, so I won't go into details here, but will instead focus on an attempt to use dask for file IO natively within yt.

The focus of this attempt falls within BaseIOHandler._read_particle_selection(). In this function, yt currently pre-allocates arrays for all the particle fields that we're reading and then read in the fields for each chunk using the self._read_particle_fields(chunks, ptf, selector) generator.

One conceptually straightforward way to leverage dask is by building a dask.dataframe from delayed objects. Dask dataframes, unlike dask arrays, do not need to know the total size a priori, but we do need to specify a metadata dictionary to declare the column names and datatypes.

So to build a dask.dataframe, we can do the following (from within BaseIOHandler._read_particle_selection()):

        ptypes = list(ptf.keys())
        delayed_dfs = {}
        for ptype in ptypes:
            # build a dataframe from delayed for each particle type
            this_ptf = {ptype: ptf[ptype]}
            delayed_chunks = [               
                dask.delayed(self._read_single_ptype)(
                    ch, this_ptf, selector, ptype_meta[ptype]
                )
                for ch in chunks
            ]
            delayed_dfs[ptype] = ddf.from_delayed(delayed_chunks, meta=ptype_meta[ptype])

In this snippet, we build up a dictionary of dask dataframes organized by particle type (ptypes). First, let's focus on the innermost loop in the snippet, the actual application of the dask.delayed decorator:

dask.delayed(self._read_single_ptype)(
                    ch, this_ptf, selector, ptype_meta[ptype]
                )

This decorator wraps a _read_single_ptype function, which takes a single chunk, a single particle type dictionary (with multiple fields) and a column-datatype metadata dictionary and returns a normal pandas dataframe:

    def _read_single_ptype(self, chunk, ptf, selector, meta_dict):
        # read a single chunk and single particle type into a pandas dataframe so that 
        # we can use dask.dataframe.from_delayed! fields within a particle type should
        # have the same length?

        chunk_results = pd.DataFrame(meta_dict)
        # each particle type could be a different dataframe...
        for field_r, vals in self._read_particle_fields([chunk], ptf, selector):
            chunk_results[field_r[1]] = vals        
            

        return chunk_results

We can see this function just calls our normal self._read_particle_fields and pulls out the field values as usual. We are storing the fields in columns of our single chunk and single particle type in a pandas dataframe, chunk_results.

So for a single particle type, we build our delayed dataframe

            delayed_chunks = [               
                dask.delayed(self._read_single_ptype)(
                    ch, this_ptf, selector, ptype_meta[ptype]
                )
                for ch in chunks
            ]
            delayed_dfs[ptype] = ddf.from_delayed(delayed_chunks, meta=ptype_meta[ptype])

The delayed_dfs object is a dictionary with a delayed dataframe for each particle type. The reason we're organizing by particle type is the issue that different particle types may have a different number of records in a given chunk (otherwise we'd have to store a large number of null values). In order to return the expected in-memory dict that _read_particle_selection() should return, we can very easily pull all the records from across our chunks and particle typeswith:

        rv = {}
        for ptype in ptypes:
            for col in delayed_dfs[ptype].columns:
                rv[(ptype, col)] = delayed_dfs[ptype][col].values.compute()

So this is a fairly promising approach, particularly since the dataframes do not need to know the expected array size. And it does indeed work to read data from our particle front ends, with some modifications to the ParticleContainer class.

The notebook here shows the above approach in action, with notes on the required changes to the ParticleContainer class. To load actually execute in parallel, the only requirement is to spin up a dask Client:

from dask.distributed import Client 
import yt 
c = Client(threads_per_worker=3,n_workers=4)
ds = yt.load_sample("snapshot_033")
ad = ds.all_data()
d = ad[('PartType0','Density')

The dask.delayed and dask.compute will find and connect to the Client. If no client is present, the Task Graph will be executed sequentially. Here's a snapshot of the Dask Dashboard's Task Stream during a parallel particle read:

return to main post