wehub-resource-sync
568 lines
17 KiB
Python
568 lines
17 KiB
Python
"""Data quality and filtering utilities for data loading.
|
|
|
|
This module provides centralized functions for:
|
|
- Coverage summaries (rows, symbols, date range)
|
|
- OHLC invariant checks
|
|
- Null rate analysis
|
|
- Gap detection in time series
|
|
- Symbol subsetting for test-mode execution
|
|
|
|
Usage:
|
|
>>> from utils.data_quality import describe_coverage, check_ohlc_invariants
|
|
>>> coverage = describe_coverage(df, time_col="timestamp", asset_col="symbol")
|
|
>>> invariants = check_ohlc_invariants(df)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import random
|
|
from datetime import timedelta
|
|
from typing import TYPE_CHECKING
|
|
|
|
import polars as pl
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Sequence
|
|
|
|
|
|
def apply_max_symbols(
|
|
data: pl.DataFrame | pl.LazyFrame,
|
|
max_symbols: int,
|
|
symbol_col: str = "symbol",
|
|
seed: int = 42,
|
|
) -> pl.DataFrame | pl.LazyFrame:
|
|
"""Limit data to a random subset of symbols for fast-path testing.
|
|
|
|
Selects a reproducible random sample of symbols using a fixed seed.
|
|
Returns data unchanged if max_symbols <= 0 or >= total symbols.
|
|
"""
|
|
if max_symbols <= 0:
|
|
return data
|
|
|
|
if isinstance(data, pl.LazyFrame):
|
|
all_symbols = data.select(pl.col(symbol_col).unique()).collect()[symbol_col].to_list()
|
|
else:
|
|
all_symbols = data[symbol_col].unique().to_list()
|
|
|
|
if max_symbols >= len(all_symbols):
|
|
return data
|
|
|
|
rng = random.Random(seed)
|
|
selected = rng.sample(sorted(all_symbols), max_symbols)
|
|
return data.filter(pl.col(symbol_col).is_in(selected))
|
|
|
|
|
|
def describe_coverage(
|
|
df: pl.DataFrame,
|
|
time_col: str = "timestamp",
|
|
asset_col: str = "symbol",
|
|
) -> dict:
|
|
"""Return coverage summary for a dataset.
|
|
|
|
Args:
|
|
df: DataFrame with time and asset columns
|
|
time_col: Name of the timestamp/date column
|
|
asset_col: Name of the asset identifier column
|
|
|
|
Returns:
|
|
Dictionary with rows, assets, time_min, time_max, unique_times
|
|
"""
|
|
return {
|
|
"rows": df.height,
|
|
"assets": df[asset_col].n_unique() if asset_col in df.columns else 0,
|
|
"time_min": df[time_col].min(),
|
|
"time_max": df[time_col].max(),
|
|
"unique_times": df[time_col].n_unique(),
|
|
}
|
|
|
|
|
|
def print_coverage(
|
|
df: pl.DataFrame,
|
|
time_col: str = "timestamp",
|
|
asset_col: str = "symbol",
|
|
dataset_name: str = "Dataset",
|
|
) -> None:
|
|
"""Print formatted coverage summary."""
|
|
cov = describe_coverage(df, time_col, asset_col)
|
|
print(f"=== {dataset_name} Coverage ===")
|
|
print(f" Rows: {cov['rows']:,}")
|
|
print(f" Assets: {cov['assets']:,}")
|
|
print(f" Time range: {cov['time_min']} to {cov['time_max']}")
|
|
print(f" Unique times: {cov['unique_times']:,}")
|
|
|
|
|
|
def check_ohlc_invariants(
|
|
df: pl.DataFrame,
|
|
open_col: str = "open",
|
|
high_col: str = "high",
|
|
low_col: str = "low",
|
|
close_col: str = "close",
|
|
volume_col: str = "volume",
|
|
) -> pl.DataFrame:
|
|
"""Check OHLC data quality invariants.
|
|
|
|
Validates:
|
|
- high >= low
|
|
- high >= open
|
|
- high >= close
|
|
- low <= open
|
|
- low <= close
|
|
- volume >= 0 (if volume column exists)
|
|
|
|
For each check, only rows where all relevant columns are non-null are
|
|
considered. This prevents null comparisons from distorting percentages
|
|
(important for TAQ data where trade columns may be null for no-trade bars).
|
|
|
|
Args:
|
|
df: DataFrame with OHLC columns
|
|
open_col, high_col, low_col, close_col: Column names for OHLC
|
|
volume_col: Column name for volume (optional)
|
|
|
|
Returns:
|
|
DataFrame with check names and valid_pct columns
|
|
"""
|
|
results = []
|
|
total_rows = df.height
|
|
cols = set(df.columns)
|
|
|
|
def _check_invariant(name: str, condition: pl.Expr, required_cols: list[str]) -> None:
|
|
"""Check an invariant on rows where all required columns are non-null."""
|
|
# Filter to rows where all required columns are non-null
|
|
not_null_filter = pl.all_horizontal([pl.col(c).is_not_null() for c in required_cols])
|
|
applicable = df.filter(not_null_filter)
|
|
n_applicable = applicable.height
|
|
|
|
if n_applicable == 0:
|
|
return # Skip if no applicable rows
|
|
|
|
valid_pct = applicable.select(condition.mean()).item() * 100
|
|
results.append(
|
|
{
|
|
"check": name,
|
|
"valid_pct": valid_pct,
|
|
"applicable_rows": n_applicable,
|
|
"total_rows": total_rows,
|
|
}
|
|
)
|
|
|
|
# Define checks with their required columns
|
|
if {high_col, low_col}.issubset(cols):
|
|
_check_invariant(
|
|
"high_gte_low",
|
|
pl.col(high_col) >= pl.col(low_col),
|
|
[high_col, low_col],
|
|
)
|
|
|
|
if {high_col, open_col}.issubset(cols):
|
|
_check_invariant(
|
|
"high_gte_open",
|
|
pl.col(high_col) >= pl.col(open_col),
|
|
[high_col, open_col],
|
|
)
|
|
|
|
if {high_col, close_col}.issubset(cols):
|
|
_check_invariant(
|
|
"high_gte_close",
|
|
pl.col(high_col) >= pl.col(close_col),
|
|
[high_col, close_col],
|
|
)
|
|
|
|
if {low_col, open_col}.issubset(cols):
|
|
_check_invariant(
|
|
"low_lte_open",
|
|
pl.col(low_col) <= pl.col(open_col),
|
|
[low_col, open_col],
|
|
)
|
|
|
|
if {low_col, close_col}.issubset(cols):
|
|
_check_invariant(
|
|
"low_lte_close",
|
|
pl.col(low_col) <= pl.col(close_col),
|
|
[low_col, close_col],
|
|
)
|
|
|
|
if volume_col in cols:
|
|
_check_invariant(
|
|
"volume_non_negative",
|
|
pl.col(volume_col) >= 0,
|
|
[volume_col],
|
|
)
|
|
|
|
if not results:
|
|
return pl.DataFrame({"check": [], "valid_pct": [], "applicable_rows": [], "total_rows": []})
|
|
|
|
return pl.DataFrame(results)
|
|
|
|
|
|
def print_ohlc_invariants(
|
|
df: pl.DataFrame,
|
|
open_col: str = "open",
|
|
high_col: str = "high",
|
|
low_col: str = "low",
|
|
close_col: str = "close",
|
|
volume_col: str = "volume",
|
|
show_coverage: bool = False,
|
|
) -> None:
|
|
"""Print OHLC invariant check results.
|
|
|
|
Args:
|
|
show_coverage: If True, show how many rows each check applies to
|
|
"""
|
|
result = check_ohlc_invariants(df, open_col, high_col, low_col, close_col, volume_col)
|
|
print("=== OHLC Invariants ===")
|
|
for row in result.iter_rows(named=True):
|
|
status = "[OK]" if row["valid_pct"] >= 99.99 else "[WARN]"
|
|
coverage = ""
|
|
if show_coverage and row["applicable_rows"] < row["total_rows"]:
|
|
coverage = f" ({row['applicable_rows']:,}/{row['total_rows']:,} rows)"
|
|
print(f" {status} {row['check']}: {row['valid_pct']:.2f}%{coverage}")
|
|
|
|
|
|
def null_rate(
|
|
df: pl.DataFrame,
|
|
cols: Sequence[str] | None = None,
|
|
) -> pl.DataFrame:
|
|
"""Calculate null rates for specified columns.
|
|
|
|
Args:
|
|
df: DataFrame to analyze
|
|
cols: Columns to check (default: all columns)
|
|
|
|
Returns:
|
|
DataFrame with column names and null_pct
|
|
"""
|
|
if cols is None:
|
|
cols = df.columns
|
|
else:
|
|
cols = [c for c in cols if c in df.columns]
|
|
|
|
if not cols:
|
|
return pl.DataFrame({"column": [], "null_pct": []})
|
|
|
|
rates = df.select([pl.col(c).is_null().mean().alias(c) for c in cols])
|
|
|
|
return pl.DataFrame(
|
|
{
|
|
"column": list(rates.columns),
|
|
"null_pct": [rates[col].item() * 100 for col in rates.columns],
|
|
}
|
|
)
|
|
|
|
|
|
def print_null_rates(
|
|
df: pl.DataFrame,
|
|
cols: Sequence[str] | None = None,
|
|
threshold: float = 0.0,
|
|
) -> None:
|
|
"""Print null rates for columns exceeding threshold.
|
|
|
|
Args:
|
|
df: DataFrame to analyze
|
|
cols: Columns to check (default: all columns)
|
|
threshold: Only print columns with null_pct > threshold
|
|
"""
|
|
result = null_rate(df, cols)
|
|
result = result.filter(pl.col("null_pct") > threshold)
|
|
print("=== Null Rates ===")
|
|
if result.height == 0:
|
|
print(" No nulls detected")
|
|
else:
|
|
for row in result.iter_rows(named=True):
|
|
print(f" {row['column']}: {row['null_pct']:.2f}%")
|
|
|
|
|
|
def gap_summary(
|
|
df: pl.DataFrame,
|
|
time_col: str = "timestamp",
|
|
group_col: str | None = "symbol",
|
|
expected_delta: timedelta | None = None,
|
|
) -> pl.DataFrame:
|
|
"""Identify gaps in time series data.
|
|
|
|
Args:
|
|
df: DataFrame with time series data
|
|
time_col: Name of timestamp column
|
|
group_col: Column to group by (e.g., symbol). None for ungrouped.
|
|
expected_delta: Expected time between rows (e.g., timedelta(hours=1))
|
|
|
|
Returns:
|
|
DataFrame with gap statistics per group (if grouped) or overall
|
|
"""
|
|
df_sorted = df.sort([group_col, time_col] if group_col else [time_col])
|
|
|
|
# Calculate time differences
|
|
if group_col:
|
|
df_gaps = df_sorted.with_columns(pl.col(time_col).diff().over(group_col).alias("time_diff"))
|
|
else:
|
|
df_gaps = df_sorted.with_columns(pl.col(time_col).diff().alias("time_diff"))
|
|
|
|
# If expected_delta provided, filter to gaps exceeding it
|
|
if expected_delta is not None:
|
|
df_gaps = df_gaps.filter(
|
|
(pl.col("time_diff") > expected_delta) | pl.col("time_diff").is_null()
|
|
)
|
|
|
|
# Aggregate
|
|
if group_col:
|
|
return (
|
|
df_gaps.filter(pl.col("time_diff").is_not_null())
|
|
.group_by(group_col)
|
|
.agg(
|
|
pl.len().alias("gap_count"),
|
|
pl.col("time_diff").max().alias("max_gap"),
|
|
)
|
|
.sort(group_col)
|
|
)
|
|
else:
|
|
gaps = df_gaps.filter(pl.col("time_diff").is_not_null())
|
|
if gaps.height == 0:
|
|
return pl.DataFrame({"gap_count": [0], "max_gap": [None]})
|
|
return pl.DataFrame(
|
|
{
|
|
"gap_count": [gaps.height],
|
|
"max_gap": [gaps["time_diff"].max()],
|
|
}
|
|
)
|
|
|
|
|
|
def per_asset_stats(
|
|
df: pl.DataFrame,
|
|
time_col: str = "timestamp",
|
|
asset_col: str = "symbol",
|
|
price_col: str = "close",
|
|
volume_col: str | None = "volume",
|
|
) -> pl.DataFrame:
|
|
"""Calculate per-asset summary statistics.
|
|
|
|
Args:
|
|
df: DataFrame with time series data
|
|
time_col: Timestamp column name
|
|
asset_col: Asset identifier column name
|
|
price_col: Price column for mean calculation
|
|
volume_col: Volume column (optional)
|
|
|
|
Returns:
|
|
DataFrame with rows, start, end, avg_price per asset
|
|
"""
|
|
aggs = [
|
|
pl.len().alias("rows"),
|
|
pl.col(time_col).min().alias("start"),
|
|
pl.col(time_col).max().alias("end"),
|
|
pl.col(price_col).mean().alias("avg_price"),
|
|
]
|
|
|
|
if volume_col and volume_col in df.columns:
|
|
aggs.append(pl.col(volume_col).mean().alias("avg_volume"))
|
|
|
|
return df.group_by(asset_col).agg(aggs).sort(asset_col)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Modeling pipeline quality gates
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def validate_prices(
|
|
df: pl.DataFrame,
|
|
price_cols: Sequence[str] = ("open", "high", "low", "close"),
|
|
asset_col: str = "symbol",
|
|
time_col: str = "timestamp",
|
|
) -> list[str]:
|
|
"""Check price columns for negative values, infinities, and NaN.
|
|
|
|
Returns a list of warning/error strings. Empty list = all clean.
|
|
"""
|
|
issues: list[str] = []
|
|
cols_present = [c for c in price_cols if c in df.columns]
|
|
|
|
for col in cols_present:
|
|
n_neg = df.filter(pl.col(col) < 0).height
|
|
n_inf = df.filter(pl.col(col).is_infinite()).height
|
|
n_nan = df.filter(pl.col(col).is_nan()).height
|
|
|
|
if n_neg > 0:
|
|
# Show which assets have negative prices
|
|
neg_assets = df.filter(pl.col(col) < 0).select(asset_col).unique().to_series().to_list()
|
|
issues.append(
|
|
f"CRITICAL: {col} has {n_neg} negative values "
|
|
f"(assets: {neg_assets[:5]}{'...' if len(neg_assets) > 5 else ''})"
|
|
)
|
|
if n_inf > 0:
|
|
issues.append(f"CRITICAL: {col} has {n_inf} infinite values")
|
|
if n_nan > 0:
|
|
issues.append(f"WARNING: {col} has {n_nan} NaN values")
|
|
|
|
return issues
|
|
|
|
|
|
def validate_labels(
|
|
df: pl.DataFrame,
|
|
label_col: str,
|
|
max_abs_return: float = 0.5,
|
|
) -> list[str]:
|
|
"""Check forward return labels for data quality issues.
|
|
|
|
Args:
|
|
df: DataFrame containing the label column
|
|
label_col: Name of the forward return column
|
|
max_abs_return: Maximum plausible absolute return (e.g., 0.5 = 50%)
|
|
|
|
Returns list of warning/error strings.
|
|
"""
|
|
issues: list[str] = []
|
|
vals = df[label_col].drop_nulls()
|
|
|
|
n_inf = vals.filter(vals.is_infinite()).len()
|
|
n_nan = vals.filter(vals.is_nan()).len()
|
|
n_extreme = vals.filter(vals.abs() > max_abs_return).len()
|
|
n_total = vals.len()
|
|
|
|
if n_inf > 0:
|
|
issues.append(f"CRITICAL: {label_col} has {n_inf} infinite values")
|
|
if n_nan > 0:
|
|
issues.append(f"CRITICAL: {label_col} has {n_nan} NaN values")
|
|
if n_extreme > 0:
|
|
pct = n_extreme / n_total * 100
|
|
issues.append(
|
|
f"WARNING: {label_col} has {n_extreme} values with |ret| > {max_abs_return:.0%} "
|
|
f"({pct:.2f}% of {n_total:,} rows)"
|
|
)
|
|
|
|
return issues
|
|
|
|
|
|
def validate_features(
|
|
df: pl.DataFrame,
|
|
feature_cols: Sequence[str],
|
|
max_abs_value: float = 1e6,
|
|
) -> list[str]:
|
|
"""Check feature columns for infinities, all-null, and extreme values.
|
|
|
|
Args:
|
|
df: DataFrame containing feature columns
|
|
feature_cols: List of feature column names to validate
|
|
max_abs_value: Threshold for flagging extreme values
|
|
|
|
Returns list of warning/error strings.
|
|
"""
|
|
issues: list[str] = []
|
|
n_rows = df.height
|
|
|
|
inf_cols = []
|
|
null_cols = []
|
|
extreme_cols = []
|
|
|
|
for col in feature_cols:
|
|
if col not in df.columns:
|
|
continue
|
|
|
|
series = df[col]
|
|
n_null = series.null_count()
|
|
non_null = series.drop_nulls()
|
|
|
|
if n_null == n_rows:
|
|
null_cols.append(col)
|
|
continue
|
|
|
|
if non_null.len() > 0:
|
|
n_inf = non_null.filter(non_null.is_infinite()).len()
|
|
if n_inf > 0:
|
|
inf_cols.append((col, n_inf))
|
|
|
|
n_extreme = non_null.filter(non_null.abs() > max_abs_value).len()
|
|
if n_extreme > 0:
|
|
extreme_cols.append((col, n_extreme))
|
|
|
|
if inf_cols:
|
|
details = ", ".join(f"{c}({n})" for c, n in inf_cols[:10])
|
|
issues.append(f"CRITICAL: {len(inf_cols)} features have infinite values: {details}")
|
|
|
|
if null_cols:
|
|
issues.append(
|
|
f"WARNING: {len(null_cols)} features are entirely null: "
|
|
f"{null_cols[:10]}{'...' if len(null_cols) > 10 else ''}"
|
|
)
|
|
|
|
if extreme_cols:
|
|
details = ", ".join(f"{c}({n})" for c, n in extreme_cols[:10])
|
|
issues.append(
|
|
f"WARNING: {len(extreme_cols)} features have values |x| > {max_abs_value:.0e}: {details}"
|
|
)
|
|
|
|
return issues
|
|
|
|
|
|
def validate_modeling_inputs(
|
|
features_df: pl.DataFrame,
|
|
label_df: pl.DataFrame,
|
|
feature_cols: Sequence[str],
|
|
label_col: str,
|
|
join_cols: Sequence[str] = ("timestamp", "symbol"),
|
|
price_cols: Sequence[str] = (),
|
|
asset_col: str = "symbol",
|
|
max_abs_return: float = 0.5,
|
|
max_abs_feature: float = 1e6,
|
|
fail_on_critical: bool = True,
|
|
) -> dict:
|
|
"""Run all data quality checks before modeling.
|
|
|
|
This is the gate between data preparation (labels + features) and
|
|
model training. Call this at the start of evaluation notebooks.
|
|
|
|
Args:
|
|
features_df: Feature DataFrame
|
|
label_df: Label DataFrame with forward returns
|
|
feature_cols: Feature column names to validate
|
|
label_col: Forward return column name
|
|
join_cols: Columns used to join features and labels
|
|
price_cols: Price columns to check (if present in features_df)
|
|
asset_col: Asset identifier column name
|
|
max_abs_return: Max plausible absolute return for labels
|
|
max_abs_feature: Max plausible absolute feature value
|
|
fail_on_critical: If True, raise ValueError on CRITICAL issues
|
|
|
|
Returns:
|
|
Dict with 'issues' (list of strings), 'n_critical', 'n_warning'
|
|
|
|
Raises:
|
|
ValueError: If fail_on_critical=True and any CRITICAL issues found
|
|
"""
|
|
all_issues: list[str] = []
|
|
|
|
# 1. Price checks (if price columns present)
|
|
if price_cols:
|
|
all_issues.extend(validate_prices(features_df, price_cols, asset_col=asset_col))
|
|
|
|
# 2. Label checks
|
|
all_issues.extend(validate_labels(label_df, label_col, max_abs_return))
|
|
|
|
# 3. Feature checks
|
|
all_issues.extend(validate_features(features_df, feature_cols, max_abs_feature))
|
|
|
|
# Summarize
|
|
n_critical = sum(1 for i in all_issues if i.startswith("CRITICAL"))
|
|
n_warning = sum(1 for i in all_issues if i.startswith("WARNING"))
|
|
|
|
# Print results
|
|
if all_issues:
|
|
print(f"Data Quality Gate: {n_critical} CRITICAL, {n_warning} WARNING")
|
|
for issue in all_issues:
|
|
marker = "[X]" if issue.startswith("CRITICAL") else "[!]"
|
|
print(f" {marker} {issue}")
|
|
else:
|
|
print("Data Quality Gate: ALL CLEAR")
|
|
|
|
result = {
|
|
"issues": all_issues,
|
|
"n_critical": n_critical,
|
|
"n_warning": n_warning,
|
|
}
|
|
|
|
if fail_on_critical and n_critical > 0:
|
|
raise ValueError(
|
|
f"Data quality gate FAILED: {n_critical} critical issues. "
|
|
f"Fix upstream data before modeling."
|
|
)
|
|
|
|
return result
|