Skip to content

Commit

Permalink
Fix and test 'chunks' escape hatch.
Browse files Browse the repository at this point in the history
  • Loading branch information
danielballan committed Oct 4, 2024
1 parent bff9b78 commit 6f5895a
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 1 deletion.
2 changes: 1 addition & 1 deletion databroker/mongo_normalized.py
Original file line number Diff line number Diff line change
Expand Up @@ -160,7 +160,7 @@ def structure_from_descriptor(descriptor, sub_dict, max_seq_num, unicode_columns
numpy_dtype = dtype.to_numpy_dtype()
if "chunks" in field_metadata:
# If the Event Descriptor tells us a preferred chunking, use that.
suggested_chunks = field_metadata["chunks"]
suggested_chunks = tuple(tuple(chunks) for chunks in field_metadata["chunks"])
elif (0 in shape) or (numpy_dtype.itemsize == 0):
# special case to avoid warning from dask
suggested_chunks = shape
Expand Down
47 changes: 47 additions & 0 deletions databroker/tests/test_broker.py
Original file line number Diff line number Diff line change
Expand Up @@ -1308,3 +1308,50 @@ def test_direct_img_read(db, RE, hw):
c = db.v2
uid, = get_uids(RE(count([hw.direct_img], 5)))
c[uid]["primary"]["data"]["img"][:]


def test_img_explicit_chunks(db, RE, hw, tmpdir):
"Test using explicit chunk size"
from ophyd import sim

RE.subscribe(db.insert)
if not hasattr(db, "v2"):
raise pytest.skip("v0 has no v2 accessor")
c = db.v2


class Detector1(sim.SynSignalWithRegistry):
def describe(self):
res = super().describe()
(key,) = res
shape = res[key]["shape"]
assert len(shape) == 2
res[key]["chunks"] = [[5], [1] * shape[0], [shape[1]]]
return res

class Detector2(sim.SynSignalWithRegistry):
def describe(self):
res = super().describe()
(key,) = res
shape = res[key]["shape"]
res[key]["chunks"] = [[5], [2] * (shape[0] // 2), [shape[1]]]
return res

img1 = Detector2(
func=lambda: np.array(np.ones((10, 10))),
name="img",
labels={"detectors"},
save_path=str(tmpdir),
)
img2 = Detector2(
func=lambda: np.array(np.ones((10, 10))),
name="img",
labels={"detectors"},
save_path=str(tmpdir),
)
uid, = get_uids(RE(count([img1], 5)))
c[uid]["primary"]["data"]["img"][:]
assert c[uid]["primary"]["data"]["img"].chunks[1] == tuple([2] * 5)
uid, = get_uids(RE(count([img2], 5)))
c[uid]["primary"]["data"]["img"][:]
assert c[uid]["primary"]["data"]["img"].chunks[1] == tuple([2] * 5)

0 comments on commit 6f5895a

Please sign in to comment.