Skip to content

Refactor decision schema for DatasetBatch lineage and error handling - #353

Open
CeliaBenquet wants to merge 1 commit into
mainfrom
celia/decision-lineage
Open

Refactor decision schema for DatasetBatch lineage and error handling#353
CeliaBenquet wants to merge 1 commit into
mainfrom
celia/decision-lineage

Conversation

@CeliaBenquet

Copy link
Copy Markdown
Collaborator

This pull request improves the robustness and correctness of the data pipeline by refining dependency handling, error reporting, and lineage enforcement in the schema population logic. It also adds unit tests to verify the correct ordering and gating of table population steps.

Dependency handling and error reporting improvements:

  • Enhanced error handling in import_np_sync_schema() by distinguishing between ModuleNotFoundError (missing package) and other exceptions (e.g., schema/DB compatibility issues), providing more specific warnings for each case.
  • Ensured that decision schema population steps are gated on successful completion of core schemas, and added a warning if the core schema stage fails to prevent improper dependency ordering.

Schema population and lineage corrections:

  • Updated ExperimentMember in decision.py to explicitly lineage its key_source to DatasetBatch, ensuring population only occurs after batch resolution. Removed redundant check for DatasetBatch existence in make().
  • Refined the construction of included_combos in both decision.py tables to use .proj() for selecting key fields, improving clarity and correctness.
  • Adjusted the TeensyTTL table's key_source logic for improved batch membership gating.

Testing enhancements:

  • Added new unit tests to verify the correct ordering and gating of schema population steps in both run.py and cron_scenario.py, as well as to ensure the correct lineage of ExperimentMember.

@CeliaBenquet
CeliaBenquet requested a review from lecriste August 14, 2026 11:18

@lecriste lecriste left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Three of the four changes are correct improvements.
But the included_combos.proj() refactor seems a high-severity regression that stops PredictionModel / PredictionModel10Windows from populating at all:

High — included_combos .proj(...) is not equivalent to dj.U(...) & and breaks population

dj_pipeline/vr4mice/schema/decision.py, both PredictionModel.key_source and PredictionModel10Windows.key_source:

included_combos = (InclusionStatus * ExperimentMember & {"included": 1}).proj(
    "set_name", "stage_name", "batch_name"
)
return LabelSet.proj() * ModelParams.proj() * included_combos

InclusionStatus * ExperimentMember is keyed by dataset (both tables' PK is dataset); set_name/stage_name/batch_name are secondary attributes. .proj(...) keeps the source primary key and only includes those three — so included_combos is keyed by dataset (one row per included session), not the deduplicated combo.

The original dj.U("set_name","stage_name","batch_name") & (...) is SELECT DISTINCT set_name, stage_name, batch_name — keyed by the 3-tuple, which is exactly the combo this key_source needs. The refactor changes the resulting key_source PK from (label_set_id, param_id, set_name, stage_name, batch_name) (= the table's own PK) to (label_set_id, param_id, dataset).

Effect: populate then calls make(key) with a key that has no stage_name / set_name / batch_name. In make:

stage_name = (ExperimentStage & key).fetch1("stage_name")

key shares no attribute with ExperimentStage, so the restriction is a no-op
(DataJoint returns the whole table when there are no common attributes) → fetch1() raises because there are multiple stages → caught by make's try/except → a FailedSession entry per key → both model tables populate nothing, silently. Even without the fetch1 failure, InclusionStatus * ExperimentMember & key would then filter by dataset instead of the combo, and the populate "already-done" antijoin weakens to (label_set_id, param_id) only.

Fix — revert both to the original idiom:

included_combos = dj.U("set_name", "stage_name", "batch_name") & (
    InclusionStatus * ExperimentMember & {"included": 1}
)

Low — new decision tests assert source text, not behavior

tests/unit/test_np_sync.py: the four added tests read the source files and assert on string content / ordering (step order in run.py and cron_scenario.py, and that ExperimentMember has def key_source(self): return vr4mice.DatasetBatch). They execute no DataJoint, so they pass regardless of whether the key_source expressions are semantically correct — in particular they would not catch the PredictionModel(.10Windows).key_source regression above, which is the riskiest change in the PR and currently has no behavioral coverage.

Consider at least asserting the resolved primary key, e.g.
set(PredictionModel.key_source.primary_key) == {"label_set_id", "param_id", "set_name", "stage_name", "batch_name"}, or a small integration test that populates a fixture, so a key-source shape regression fails the suite.

Nit — anchor ExperimentMember.key_source to its own PK parent and gate by existence

dj_pipeline/vr4mice/schema/decision.py, ExperimentMember.key_source:

@property
def key_source(self):
    # per-dataset, but only once its batch is resolved
    return Dataset & vr4mice.DatasetBatch

instead of return vr4mice.DatasetBatch. Both schedule population per dataset today (since DatasetBatch's PK is dataset), but returning the whole DatasetBatch table silently assumes DatasetBatch.primary_key == ExperimentMember.primary_key. If DatasetBatch ever gains a primary-key attribute, that override would quietly change this table's population granularity and feed make() keys with an extra field — the same failure mode as the PredictionModel key_source issue.

Dataset & vr4mice.DatasetBatch reads as "per Dataset, restricted to those with a materialized DatasetBatch," keeps the key anchored to ExperimentMember's actual PK parent, and the & restriction always preserves Dataset's key regardless of DatasetBatch's shape.

(Alternatively, declare the dependency in the definition — -> vr4mice.DatasetBatch — and drop the override entirely, since DataJoint's default key_source already gates on primary parents; that enforces a real FK too, at the cost of a schema migration.)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants