Merge pull request #297 from cjtitus/rechunk_fix

Fix minor bug in event_model's rechunk_event_pages
bluesky · Mar 28, 2024 · 707a108 · 707a108
2 parents bb2fa03 + 330566a
commit 707a108
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/event_model/__init__.py b/event_model/__init__.py
@@ -2772,7 +2772,8 @@ def page_chunks(page: dict, chunk_size: int, remainder: int) -> Generator:
                     for key in page["timestamps"].keys()
                 },
                 "filled": {
-                    key: page["filled"][key][start:stop] for key in page["data"].keys()
+                    key: page["filled"][key][start:stop]
+                    for key in page["filled"].keys()
                 },
             }
 
@@ -2831,7 +2832,7 @@ def merge_event_pages(event_pages: Iterable[EventPage]) -> EventPage:
             key: list(
                 itertools.chain.from_iterable([page["filled"][key] for page in pages])
             )
-            for key in pages[0]["data"].keys()
+            for key in pages[0]["filled"].keys()
         },
     )
     return cast(EventPage, doc)

diff --git a/event_model/tests/test_em.py b/event_model/tests/test_em.py
@@ -838,7 +838,8 @@ def test_single_run_document_router():
         sr("descriptor", desc_bundle.descriptor_doc)
 
 
-def test_rechunk_event_pages():
+@pytest.mark.parametrize("filled", [True, False])
+def test_rechunk_event_pages(filled):
     def event_page_gen(page_size, num_pages):
         """
         Generator event_pages for testing.
@@ -851,7 +852,7 @@ def event_page_gen(page_size, num_pages):
                 **{key: list(range(page_size)) for key in array_keys},
                 "data": {key: list(range(page_size)) for key in data_keys},
                 "timestamps": {key: list(range(page_size)) for key in data_keys},
-                "filled": {key: list(range(page_size)) for key in data_keys},
+                "filled": {key: list(range(page_size)) for key in data_keys if filled},
             }
 
     # Get a list of event pages of size 13.