adds a helper method for extacting data

[tomjemmett]

• checks for any nulls in columns
• prints a message about how many rows are going to be extracted
• centralises the extraction logic

fixes #145

[Copilot]

Copilot review overview

🟡 Changes recommended

The new shared null-check/extract helper has correctness and reliability gaps (assert-based enforcement, NA/NaN handling, and avoidable repeated Spark recomputation) that should be addressed before approval.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 4 · 1

Open (5)

• Handle empty column lists before calling stack() · New
• Detect NaN values in null-check validation · New
• Replace assert with a real validation exception · New
• Avoid recomputing the DataFrame three times · New
• Rename beds-specific extractor for clarity · New

What changed in this PR

This PR refactors the model-data extract scripts to centralize the common “write parquet + partitioning” logic into a shared decorator, and adds extract-time validation to catch missing values (related to #145).

Changes:

• Introduces a shared @extract(...) decorator and check_extract_for_nulls(...) helper to standardize extraction and (optionally) null-check outputs.
• Updates multiple extract scripts to return DataFrames and delegate writing/partitioning to the decorator.
• Adds a row-count log message during extracts and per-extract column-exclusion lists for null checks.

File	Description
src/​nhp/​data/​model_data/​helpers.py	Adds shared extraction decorator and null-check helper.
src/​nhp/​data/​model_data/​op.py	Migrates OP extract to decorator-based pattern and returns DF.
src/​nhp/​data/​model_data/​ip.py	Migrates IP extract to decorator-based pattern and returns DF.
src/​nhp/​data/​model_data/​aae.py	Migrates A&E extract to decorator-based pattern and returns DF.
src/​nhp/​data/​model_data/​ip_tpmas.py	Splits TPMA extracts into two decorated extract functions.
src/​nhp/​data/​model_data/​ip_functional_areas.py	Splits functional-area extracts into two decorated functions.
src/​nhp/​data/​model_data/​inequalities.py	Migrates inequalities extract to decorator-based pattern (null checks off).
src/​nhp/​data/​model_data/​demographic_factors.py	Migrates demographics extract to decorator-based pattern and returns DF.
src/​nhp/​data/​model_data/​birth_factors.py	Migrates births extract to decorator-based pattern and returns DF.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Copilot review overview

🟡 Changes recommended

The new extract decorator can leak persisted DataFrames on exceptions and includes a misleading null-check error message, which should be corrected before merging.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 1 · 1

Open (2)

• Persisted DataFrame may leak if extract fails · New
• Unordered columns make null-check output nondeterministic · New

Resolved since last review (5)

• Avoid recomputing the DataFrame three times
• Replace assert with a real validation exception
• Detect NaN values in null-check validation
• Handle empty column lists before calling stack()
• Rename beds-specific extractor for clarity

Previously missed (1)

In code that hasn't changed since last review

Inconsistent PySpark import style

src/​nhp/​data/​model_data/​inequalities.py:7

DataFrame is imported from pyspark.sql.dataframe while SparkSession is imported from pyspark.sql; this is inconsistent with the other model_data extract modules and can be simplified to a single import statement.

[Copilot]

Copilot review overview

🔵 Needs a closer look

A few issues in the new shared helper (misleading error/doc text and missing wrapper metadata preservation) should be addressed to avoid confusion and improve maintainability.

Review effort: Lite
Findings: None

Resolved since last review (2)

• Persisted DataFrame may leak if extract fails
• Unordered columns make null-check output nondeterministic

Previously missed (2)

In code that hasn't changed since last review

Decorator drops function metadata without functools.wraps

src/​nhp/​data/​model_data/​helpers.py:132

The extract decorator wraps the original function but does not preserve its metadata (name/docstring/qualname), making stack traces and CLI help output less informative. Consider using functools.wraps (without copying annotations) so the decorated extract functions remain identifiable.

Docstring incorrectly describes null-value checks as null columns

src/​nhp/​data/​model_data/​helpers.py:93

Docstring says "null columns" but the function checks for null values within columns, which is misleading for anyone relying on this helper. Update the wording (and the exclude_cols type in the doc) to match the actual behavior.

[Copilot]

Copilot review overview

🟢 Approval recommended

The refactor is cohesive and consistent across extract scripts, with only minor clarity issues (messaging/annotation readability) remaining.

Review effort: Lite
Findings: 3

Open (3)

• Docstring says null columns instead of null values · New
• Use keyword arguments in decorator for clarity · New
• Use keyword arguments in decorator for clarity · New

[yiwen-h]

Thank you! the extract decorator is neat