You may find some helpful information in Efficiently slicing random windows for reduced xarray dataset, which also considered efficiently selecting only “valid” examples.
We’re working on the xbatcher library for this type of use case - our roadmap for the next few months focuses on improving efficiency for the data loaders and we’re considering support for filtering examples.