Upload data_loader.py with huggingface_hub

data_loader.py

@@ -0,0 +1,65 @@
+"""Load data from HF Dataset with Streamlit caching."""
+
+from __future__ import annotations
+
+import json
+
+import pandas as pd
+import streamlit as st
+from huggingface_hub import hf_hub_download
+
+DATASET_REPO = "buckeyeguy/osc-usage-data"
+
+
+@st.cache_data(ttl=300)
+def load_data() -> tuple[pd.DataFrame, pd.DataFrame, dict]:
+    """Download Parquet + metadata from HF Dataset. Cached for 5 min."""
+    jobs_path = hf_hub_download(repo_id=DATASET_REPO, filename="jobs.parquet", repo_type="dataset")
+    snapshots_path = hf_hub_download(
+        repo_id=DATASET_REPO, filename="snapshots.parquet", repo_type="dataset"
+    )
+    metadata_path = hf_hub_download(
+        repo_id=DATASET_REPO, filename="metadata.json", repo_type="dataset"
+    )
+
+    jobs = pd.read_parquet(jobs_path)
+    snapshots = pd.read_parquet(snapshots_path)
+    with open(metadata_path) as f:
+        metadata = json.load(f)
+
+    # Ensure datetime columns
+    for col in ["submit_time", "start_time", "end_time"]:
+        if col in jobs.columns:
+            jobs[col] = pd.to_datetime(jobs[col])
+
+    # Add derived columns
+    if "end_time" in jobs.columns:
+        jobs["end_date"] = jobs["end_time"].dt.date
+        jobs["end_month"] = jobs["end_time"].dt.to_period("M").astype(str)
+        jobs["end_dow"] = jobs["end_time"].dt.dayofweek  # 0=Mon
+        jobs["end_hour"] = jobs["end_time"].dt.hour
+
+    if "walltime_used" in jobs.columns:
+        jobs["walltime_hours"] = jobs["walltime_used"] / 3600.0
+
+    return jobs, snapshots, metadata
+
+
+def filter_jobs(
+    jobs: pd.DataFrame,
+    date_range: tuple | None = None,
+    projects: list[str] | None = None,
+    users: list[str] | None = None,
+    systems: list[str] | None = None,
+) -> pd.DataFrame:
+    """Apply sidebar filters to jobs DataFrame."""
+    df = jobs.copy()
+    if date_range and "end_date" in df.columns:
+        df = df[df["end_date"].between(date_range[0], date_range[1])]
+    if projects:
+        df = df[df["project_code"].isin(projects)]
+    if users:
+        df = df[df["username"].isin(users)]
+    if systems:
+        df = df[df["system_code"].isin(systems)]
+    return df