Release v0.0.113 with Polars/Arrow speedups

adpena · adpena · commit 0d6a9a157ad9 · 2026-02-20T12:06:31.000-06:00
diff --git a/.cache/entities_Current_Districts_2025_Schools_2024_to_2025.sqlite b/.cache/entities_Current_Districts_2025_Schools_2024_to_2025.sqlite
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f8cbbcd83c58ab1c7a711338003553599919e1f8a1c9f1576406904493b1012f
-size 259493888
+oid sha256:c26d4dc73740fb798315f7447365c57e004757d126481b2371ca88ee565f0e8b
+size 259403776
diff --git a/.cache/map_payloads_Current_Districts_2025_Schools_2024_to_2025.sqlite b/.cache/map_payloads_Current_Districts_2025_Schools_2024_to_2025.sqlite
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6a3dd6d364d372719d6208dd33b594ba09ed02a13552daa333f1f51613d5294c
+oid sha256:026e07a67363a0d163e196ec98edfcf9e1cc567531941e9464f15b16512ce241
 size 317960192
diff --git a/.cache/repo_Current_Districts_2025_Schools_2024_to_2025.pkl.gz b/.cache/repo_Current_Districts_2025_Schools_2024_to_2025.pkl.gz
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:61670ea71e1d14a245420a884c406efcce72188bfd6dd3b8faaceca05e3b0706
-size 30717546
+oid sha256:fc5cd1f3118b6d3a38d33c02a3630f7418a462e2960169305c9d1241e559ba52
+size 30896931
diff --git a/.cache/repo_districts_campuses.pkl.gz b/.cache/repo_districts_campuses.pkl.gz
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dc605919048190e61ad82d2451086a71d184301448d7e094f60474db3e8260e6
-size 525
+oid sha256:d2a06b4b61f75bd97fd886f2eec8836b7cc1679562f33cbbfebed848647bba43
+size 524
diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "teadata"
-version = "0.0.112"
+version = "0.0.113"
 description = "Pythonic data engine for Texas education datasets"
 authors = [{name="Alejandro Peña", email="adpena@gmail.com"}]
 readme = "README.md"
@@ -42,6 +42,7 @@ classifiers = [
 # Faster I/O / performance add-ons
 speedups = [
     "pyogrio>=0.7",
+    "polars>=1.0",
     "pygeos; python_version<'3.11'",  # no-op on modern, left for legacy envs
 ]
 
diff --git a/teadata/.cache/entities_Current_Districts_2025_Schools_2024_to_2025.sqlite b/teadata/.cache/entities_Current_Districts_2025_Schools_2024_to_2025.sqlite
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f8cbbcd83c58ab1c7a711338003553599919e1f8a1c9f1576406904493b1012f
-size 259493888
+oid sha256:c26d4dc73740fb798315f7447365c57e004757d126481b2371ca88ee565f0e8b
+size 259403776
diff --git a/teadata/.cache/map_payloads_Current_Districts_2025_Schools_2024_to_2025.sqlite b/teadata/.cache/map_payloads_Current_Districts_2025_Schools_2024_to_2025.sqlite
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6a3dd6d364d372719d6208dd33b594ba09ed02a13552daa333f1f51613d5294c
+oid sha256:026e07a67363a0d163e196ec98edfcf9e1cc567531941e9464f15b16512ce241
 size 317960192
diff --git a/teadata/.cache/repo_Current_Districts_2025_Schools_2024_to_2025.pkl.gz b/teadata/.cache/repo_Current_Districts_2025_Schools_2024_to_2025.pkl.gz
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:61670ea71e1d14a245420a884c406efcce72188bfd6dd3b8faaceca05e3b0706
-size 30717546
+oid sha256:fc5cd1f3118b6d3a38d33c02a3630f7418a462e2960169305c9d1241e559ba52
+size 30896931
diff --git a/teadata/.cache/repo_districts_campuses.pkl.gz b/teadata/.cache/repo_districts_campuses.pkl.gz
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dc605919048190e61ad82d2451086a71d184301448d7e094f60474db3e8260e6
-size 525
+oid sha256:d2a06b4b61f75bd97fd886f2eec8836b7cc1679562f33cbbfebed848647bba43
+size 524
diff --git a/teadata/engine.py b/teadata/engine.py
@@ -3014,30 +3014,25 @@ def norm(x):
         dst_series = df[dst_col].map(norm)
         counts = df[count_col].to_numpy()
 
-        # Local cache to avoid repeated lookups against _campus_by_number
+        # Precompute campus-id lookups for all observed normalized keys
         id_cache: Dict[str, Optional[uuid.UUID]] = {}
-
-        def _lookup(normed: str | None) -> Optional[uuid.UUID]:
-            if normed is None:
-                return None
-            hit = id_cache.get(normed)
-            if hit is not None or normed in id_cache:
-                return hit
+        observed_keys = set(src_series.dropna())
+        observed_keys.update(dst_series.dropna())
+        for normed in observed_keys:
             digits = normed[1:] if normed.startswith("'") else normed
-            cid = self._campus_by_number.get(normed) or self._campus_by_number.get(
-                digits
-            )
-            id_cache[normed] = cid
-            return cid
+            id_cache[normed] = self._campus_by_number.get(
+                normed
+            ) or self._campus_by_number.get(digits)
 
-        for src_key, dst_key, raw in zip(
-            src_series.to_numpy(), dst_series.to_numpy(), counts
+        src_ids = src_series.map(id_cache).to_numpy()
+        dst_ids = dst_series.map(id_cache).to_numpy()
+
+        for src_key, dst_key, src_id, dst_id, raw in zip(
+            src_series.to_numpy(), dst_series.to_numpy(), src_ids, dst_ids, counts
         ):
             if not src_key or not dst_key:
                 continue
 
-            src_id = _lookup(src_key)
-            dst_id = _lookup(dst_key)
             if src_id is None or dst_id is None:
                 if src_id is None and dst_id is None:
                     self._xfers_missing["either"] = (
@@ -3072,9 +3067,6 @@ def _sort_key(edge):
 
         for k in list(self._xfers_in.keys()):
             # for transfers_in sort by the incoming count, same logic
-            from_id, cnt, _ = (
-                self._xfers_in[k][0] if self._xfers_in[k] else (None, None, None)
-            )
             self._xfers_in[k].sort(
                 key=lambda ed: (
                     -(ed[1] or -1),
diff --git a/teadata/enrichment/campuses.py b/teadata/enrichment/campuses.py
@@ -12,6 +12,11 @@
     normalize_campus_number_column,
 )
 
+try:  # optional performance dependency
+    import polars as pl
+except Exception:  # pragma: no cover - optional dependency
+    pl = None  # type: ignore[assignment]
+
 DEFAULT_PEIMS_FINANCIAL_COLUMNS: list[str] = [
     "instruction_af_perc",
     "transportation_af_per_student",
@@ -44,6 +49,77 @@ def _canon_series(series: pd.Series) -> pd.Series:
     return series.map(canonical_campus_number)
 
 
+def _is_text_series(series: pd.Series, *, sample_size: int = 64) -> bool:
+    if pd.api.types.is_string_dtype(series):
+        return True
+    if series.dtype != object:
+        return False
+    non_null = series[series.notna()]
+    if non_null.empty:
+        return False
+    sample = non_null.head(sample_size)
+    return bool(len(sample)) and all(isinstance(v, str) for v in sample)
+
+
+def _clean_value_columns(df: pd.DataFrame, columns: list[str]) -> None:
+    text_cols = [col for col in columns if col in df.columns and _is_text_series(df[col])]
+    if text_cols:
+        if pl is not None:
+            try:
+                text_frame = pl.from_pandas(df[text_cols], include_index=False)
+                cleaned = text_frame.with_columns(
+                    [
+                        pl.when(pl.col(col).is_null())
+                        .then(None)
+                        .otherwise(pl.col(col).str.strip_chars())
+                        .alias(col)
+                        for col in text_cols
+                    ]
+                ).with_columns(
+                    [
+                        pl.when(pl.col(col) == "").then(None).otherwise(pl.col(col)).alias(
+                            col
+                        )
+                        for col in text_cols
+                    ]
+                )
+                cleaned_pdf = cleaned.to_pandas(use_pyarrow_extension_array=True)
+                for col in text_cols:
+                    df[col] = cleaned_pdf[col]
+            except Exception:
+                for col in text_cols:
+                    cleaned = df[col].astype("string").str.strip()
+                    cleaned = cleaned.replace("", pd.NA)
+                    df[col] = cleaned.where(cleaned.notna(), None)
+        else:
+            for col in text_cols:
+                cleaned = df[col].astype("string").str.strip()
+                cleaned = cleaned.replace("", pd.NA)
+                df[col] = cleaned.where(cleaned.notna(), None)
+
+    for col in columns:
+        if col in df.columns:
+            df[col] = df[col].where(pd.notna(df[col]), None)
+
+
+def _build_mapping(sub: pd.DataFrame, use_cols: list[str]) -> dict[str, dict[str, Any]]:
+    sub = sub.where(pd.notnull(sub), None)
+    base_mapping = sub.set_index("campus_number", drop=True)[use_cols].to_dict(
+        orient="index"
+    )
+
+    mapping: dict[str, dict[str, Any]] = {}
+    for key, record in base_mapping.items():
+        if not key:
+            continue
+        mapping[key] = record
+        digits = key[1:] if isinstance(key, str) and key.startswith("'") else key
+        mapping.setdefault(digits, record)
+        if isinstance(digits, str) and digits.isdigit():
+            mapping.setdefault(str(int(digits)), record)
+    return mapping
+
+
 def _build_campus_multi_index(repo) -> dict[str, Any]:
     """Return a dict mapping multiple key shapes to Campus.id for robust matching.
     Keys created per campus_number:
@@ -159,29 +235,11 @@ def _apply_campus_enrichment_table(
                 "campus enrichment requires a non-empty `select` collection of column names"
             )
 
-    for c in use_cols:
-        if c in df.columns:
-            df[c] = df[c].apply(
-                lambda x: (
-                    None
-                    if (pd.isna(x) or (isinstance(x, str) and x.strip() == ""))
-                    else (x.strip() if isinstance(x, str) else x)
-                )
-            )
+    _clean_value_columns(df, use_cols)
 
     # Build mapping from canonical campus_number -> selected attrs
-    mapping: dict[str, dict[str, Any]] = {}
     sub = df[["campus_number"] + use_cols].drop_duplicates("campus_number")
-    for r in sub.itertuples(index=False):
-        key = getattr(r, "campus_number")
-        if not key:
-            continue
-        record = {k: getattr(r, k) for k in use_cols}
-        mapping[key] = record
-        digits = key[1:]
-        mapping.setdefault(digits, record)
-        if digits.isdigit():
-            mapping.setdefault(str(int(digits)), record)
+    mapping = _build_mapping(sub, use_cols)
 
     if _profile_enabled():
         _debug(
@@ -315,21 +373,9 @@ def _apply_campus_peims_financials(
             + ", ".join(sorted(missing))
         )
 
+    _clean_value_columns(df, use_cols)
     sub = df[["campus_number"] + use_cols].drop_duplicates("campus_number")
-
-    mapping: dict[str, dict[str, Any]] = {}
-    for r in sub.itertuples(index=False):
-        key = getattr(r, "campus_number")
-        if not key:
-            continue
-
-        record = {k: getattr(r, k) for k in use_cols}
-
-        mapping[key] = record
-        digits = key[1:]
-        mapping.setdefault(digits, record)
-        if digits.isdigit():
-            mapping.setdefault(str(int(digits)), record)
+    mapping = _build_mapping(sub, use_cols)
 
     if _profile_enabled():
         _debug(
diff --git a/teadata/enrichment/districts.py b/teadata/enrichment/districts.py
@@ -10,6 +10,11 @@
 from teadata.teadata_config import canonical_district_number
 from teadata.teadata_config import normalize_district_number_column
 
+try:  # optional performance dependency
+    import polars as pl
+except Exception:  # pragma: no cover - optional dependency
+    pl = None  # type: ignore[assignment]
+
 
 def _profile_enabled() -> bool:
     return bool(getattr(classes_mod, "ENABLE_PROFILING", False))
@@ -40,6 +45,59 @@ def _canon_district_number(x: Any) -> Optional[str]:
         return None
 
 
+def _is_text_series(series: pd.Series, *, sample_size: int = 64) -> bool:
+    if pd.api.types.is_string_dtype(series):
+        return True
+    if series.dtype != object:
+        return False
+    non_null = series[series.notna()]
+    if non_null.empty:
+        return False
+    sample = non_null.head(sample_size)
+    return bool(len(sample)) and all(isinstance(v, str) for v in sample)
+
+
+def _clean_value_columns(df: pd.DataFrame, columns: list[str]) -> None:
+    text_cols = [col for col in columns if col in df.columns and _is_text_series(df[col])]
+    if text_cols:
+        if pl is not None:
+            try:
+                text_frame = pl.from_pandas(df[text_cols], include_index=False)
+                cleaned = text_frame.with_columns(
+                    [
+                        pl.when(pl.col(col).is_null())
+                        .then(None)
+                        .otherwise(pl.col(col).str.strip_chars())
+                        .alias(col)
+                        for col in text_cols
+                    ]
+                ).with_columns(
+                    [
+                        pl.when(pl.col(col) == "").then(None).otherwise(pl.col(col)).alias(
+                            col
+                        )
+                        for col in text_cols
+                    ]
+                )
+                cleaned_pdf = cleaned.to_pandas(use_pyarrow_extension_array=True)
+                for col in text_cols:
+                    df[col] = cleaned_pdf[col]
+            except Exception:
+                for col in text_cols:
+                    cleaned = df[col].astype("string").str.strip()
+                    cleaned = cleaned.replace("", pd.NA)
+                    df[col] = cleaned.where(cleaned.notna(), None)
+        else:
+            for col in text_cols:
+                cleaned = df[col].astype("string").str.strip()
+                cleaned = cleaned.replace("", pd.NA)
+                df[col] = cleaned.where(cleaned.notna(), None)
+
+    for col in columns:
+        if col in df.columns:
+            df[col] = df[col].where(pd.notna(df[col]), None)
+
+
 # -----------------------------
 # Core enrichment routine used by both the @enricher class and legacy wrapper
 # -----------------------------
@@ -123,23 +181,18 @@ def _apply_district_enrichment_table(
             )
 
     # Basic whitespace/NA cleanup on the value columns
-    for c in use_cols:
-        if c in df.columns:
-            df[c] = df[c].apply(
-                lambda x: (
-                    None
-                    if (pd.isna(x) or (isinstance(x, str) and x.strip() == ""))
-                    else (x.strip() if isinstance(x, str) else x)
-                )
-            )
+    _clean_value_columns(df, use_cols)
 
     # Build mapping keyed by canonical district_number → {attr: value}
     sub = df[["district_number"] + use_cols].drop_duplicates("district_number")
-    mapping: Dict[str, Dict[str, Any]] = {}
-    for r in sub.itertuples(index=False):
-        key = getattr(r, "district_number")
-        if key:
-            mapping[key] = {k: getattr(r, k) for k in use_cols}
+    sub = sub.where(pd.notnull(sub), None)
+    mapping = {
+        key: attrs
+        for key, attrs in sub.set_index("district_number", drop=True)[use_cols]
+        .to_dict(orient="index")
+        .items()
+        if key
+    }
 
     if _profile_enabled():
         _debug(
diff --git a/teadata/load_data.py b/teadata/load_data.py
diff --git a/teadata/scripts/check_env.py b/teadata/scripts/check_env.py
diff --git a/tests/test_transfer_speedups.py b/tests/test_transfer_speedups.py
diff --git a/uv.lock b/uv.lock
