Context.export_dataset method to export to another file format

[matbryan52]

NOTE The implementation has changed, leaving the below for the record...

Supports export to npy and raw, could be later extended. In particular in response to #1366 and likely other Issues.

User-facing is fairly minimal:

class DataSet:
    ...

    def save_as(
        self,
        path: os.PathLike,
        format: Optional[Literal['raw', 'npy']] = None,
        progress: bool = False,
        save_dtype: Optional['nt.DTypeLike'] = None,
    ):
        ...

Works via partition.get_macrotile() so should be quite well-supported. Writes data serially partition-by-partition but this could of course be improved.

I did see some old code for writing data notably io.writers.base.WriteHandle, but I couldn't quite make sense of it. Instead I wrote the writer support into the file common.writers.

A further quality-of-life extension would be a method to actually load the data direct to an array in memory.

I'm expecting to discuss / re-work this before merging, of course, but I'm putting up the PR for discussion!

Contributor Checklist:

• I have added or updated my entry in the creators.json file
• I have added a changelog entry for my contribution
• I have added/updated documentation for all user-facing changes
• I have added/updated test cases
• I have included the rebuilt production build of the client (only if changes were made to the GUI)

Reviewer Checklist:

• /azp run libertem.libertem-data passed
• No import of GPL code from MIT code

[uellue]

Thx for addressing this! :-)

Maybe one could also use the Dask integration? That already allows loading to memory, storing into anything that allows slice assignment, HDF5, TileDB, Zarr, NPY stacks: https://docs.dask.org/en/stable/array-creation.html#store-dask-arrays. Probably it just has to be documented properly?

Also, in LiberTEM-live we have https://github.com/LiberTEM/LiberTEM-live/blob/master/src/libertem_live/udf/record.py that performs parallel writing to an NPY file on the workers using an UDF. That approach could be the way to go for best performance. Instead of NPY one could also look into ZARR -- it just has to support parallel assignment to non-overlapping slices.

[matbryan52]

Maybe one could also use the Dask integration? That already allows loading to memory, storing into anything that allows slice assignment, HDF5, TileDB, Zarr, NPY stacks: https://docs.dask.org/en/stable/array-creation.html#store-dask-arrays. Probably it just has to be documented properly?

Dask integration is a good way to do this, I agree. The reason that I didn't go with that for this first implementation (even in a hidden way) is that right now we can't ensure that partitions are axis-aligned so saving via Dask could incur rechunking (or even loading all the data into memory). It is a good reason to revisit #1264 and do it properly.

It's also possible that the delayed functions will be executed on different workers / machines leading IPC for the array chunks, though this is probably a minor issue.

Also, in LiberTEM-live we have https://github.com/LiberTEM/LiberTEM-live/blob/master/src/libertem_live/udf/record.py that performs parallel writing to an NPY file on the workers using an UDF. That approach could be the way to go for best performance. Instead of NPY one could also look into ZARR -- it just has to support parallel assignment to non-overlapping slices.

Great, I can look at implementing this approach, though we would have to upstream RecordUDF or similar rather than importing from libertem-live I think.

[sk1p]

Some comments:

The API: if we want to extend to a parallel version later, there will probably be some interaction with a Context - we should think about how that would look like, and if we can easily implement that in a backwards-compatible way. Also, running the export in the main process / node breaks with one "tradition" of the DataSet API, which is to only perform I/O "remotely" on an executor.

Also, in LiberTEM-live we have https://github.com/LiberTEM/LiberTEM-live/blob/master/src/libertem_live/udf/record.py that performs parallel writing to an NPY file on the workers using an UDF.

As far as I can see, with this we (will) have multiple ways of converting/saving:

1. As a side-effect of running RecordUDF
2. "Imperatively" using DataSet.save_as
3. live: controlling the data acquisition software of the detector and telling it to acquire into a specified path (saved in the detector-specific format, like .mib or HDF5 with bslz4 or similar)
4. live: specialized data acquisition / writing routines that are built into our receiving libraries (example: the upcoming k2is support in LiberTEM-live can write to disk very efficiently and can keep up with the data rates without any issues; data writing can be controlled via the Python API)

Both 3) and 4) can be summarized as: built for the live case; "declarative" writing; limited format support and full integration into libertem-live is still TODO.

Should we attempt to unify these in some way, at least in the interface? Having different interfaces for offline conversion vs. live acquisition would also be a good option IMHO.

That approach could be the way to go for best performance.

Agreed for offline conversion of data; I'm not convinced that it will work in the general case for live acquisition. I haven't tried it yet in a multi-process execution, but having a single-threaded mmap based writer in the pipeline of the k2is receiver kills/severely limits performance (IIRC: page faults and memory pressure -> jitter and less efficient memory re-use).

though we would have to upstream RecordUDF or similar rather than importing from libertem-live I think.

Agreed.

[uellue]

Good point about the K2 IS and dedicated high-performance low-overhead writers! In that sense, it would be a dataset option in an UDF run and just saving or converting would be running an empty UDF set with the dataset saving option enabled?

[matbryan52]

After discussion with @sk1p updated this to instead use a method on the Context:

    def export_dataset(
        self,
        dataset: DataSet,
        *
        path: os.PathLike,
        progress: bool = False,
        overwrite: bool = False,
    ):
        """
        Export the dataset to another format on disk

with an at-this-time minimal interface. Also this PR now upstreams RecordUDF to take advantage of its performance.

To be discussed if this is something we want to add, but it is undeniably a very convenient way to work with some odd datasets (offline) with other tools !

[codecov]

Codecov Report

Patch coverage: 94.44% and project coverage change: +0.23 🎉

Comparison is base (4f9de10) 68.47% compared to head (36f3695) 68.70%.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1379      +/-   ##
==========================================
+ Coverage   68.47%   68.70%   +0.23%     
==========================================
  Files         305      158     -147     
  Lines       17874    15541    -2333     
  Branches     3201     2771     -430     
==========================================
- Hits        12239    10678    -1561     
+ Misses       5115     4436     -679     
+ Partials      520      427      -93     

Impacted Files	Coverage Δ
src/libertem/udf/record.py	91.30% <91.30%> (ø)
src/libertem/api.py	91.09% <100.00%> (+1.52%)	⬆️

... and 182 files with indirect coverage changes

☔ View full report in Codecov by Sentry.
📢 Do you have feedback about the report comment? Let us know in this issue.

[sk1p]

Looks good overall, I've added some comments in-line, mostly details.

[matbryan52]

Thanks for the review again, sorry to have missed all those doc changes!

[matbryan52]

/azp run libertem.libertem-data

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[sk1p]

LGTM!