Skip to content

Commit

Permalink
Add a function to directly load file from strax folder (#801)
Browse files Browse the repository at this point in the history
* Add a function to directly load file from strax folder
Using only information inside the folder

* Add test
  • Loading branch information
dachengx authored Feb 15, 2024
1 parent 1a93526 commit 882b545
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 0 deletions.
37 changes: 37 additions & 0 deletions strax/io.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,13 @@

import os
import bz2
import json

import numpy as np
import blosc
import zstd
import lz4.frame as lz4
from ast import literal_eval

import strax

Expand Down Expand Up @@ -95,3 +97,38 @@ def _compress_blosc(data):


COMPRESSORS["blosc"]["compress"] = _compress_blosc


@export
def dry_load_files(dirname, chunk_number=None):
prefix = strax.storage.files.dirname_to_prefix(dirname)
metadata_json = f"{prefix}-metadata.json"
md_path = os.path.join(dirname, metadata_json)

with open(md_path, mode="r") as f:
metadata = json.loads(f.read())

dtype = literal_eval(metadata["dtype"])

results = []
if chunk_number is None:
for chunk_info in metadata["chunks"]:
if chunk_info["n"] != 0:
results.append(
load_file(
os.path.join(dirname, f"{prefix}-{chunk_info['chunk_i']:06d}"),
metadata["compressor"],
dtype,
)
)
results = np.hstack(results)
else:
if chunk_number >= len(metadata["chunks"]):
raise ValueError(f"Chunk {chunk_number:06d} does not exist in {dirname}.")
if metadata["chunks"][chunk_number]["n"] != 0:
results = load_file(
os.path.join(dirname, f"{prefix}-{chunk_number:06d}"),
metadata["compressor"],
dtype,
)
return results if len(results) else np.empty(0, dtype)
11 changes: 11 additions & 0 deletions tests/test_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,17 @@ def get_st_and_fill_frontends(self) -> ty.Tuple[strax.Context, dict]:

return (strax.Context(storage=frontends, **self.context_kwargs), frontend_setup)

def test_dry_load_files(self):
"""Test that dry_load_files can load the data."""
st, frontend_setup = self.get_st_and_fill_frontends()
for sf in st.storage:
key = st.key_for(self.run_id, self.target)
dirname = os.path.join(sf.path, str(key))
strax.io.dry_load_files(dirname)
strax.io.dry_load_files(dirname, 0)
with self.assertRaises(ValueError):
strax.io.dry_load_files(dirname, 99)

def test_close_goes_first_md(self):
"""Let's see that if we get the meta-data, it's from the one with the lowest remoteness.
Expand Down

0 comments on commit 882b545

Please sign in to comment.