diff --git a/analysis/column_processor.py b/analysis/column_processor.py new file mode 100644 index 0000000..53e6342 --- /dev/null +++ b/analysis/column_processor.py @@ -0,0 +1,162 @@ +"""Unified column processor — dispatches per-type distance computation. + +Replaces the if-elif chain in :func:`distance.run_preprocessing_pipeline` +with :data:`~analysis.types.TYPE_REGISTRY` dispatch. + +Usage:: + + from analysis.types import classify_schema + from analysis.column_processor import ColumnProcessor + + column_types = classify_schema(lf) + proc = ColumnProcessor(lf, column_types) + lf = proc.compute_distances() # Applies each type's .distance() + lf = proc.normalize() # StandardScaler on numeric dist cols + lf, noise = proc.svd_denoise() # SVD with noise profile +""" + +from __future__ import annotations + +import numpy as np +import polars as pl + +from sklearn.preprocessing import StandardScaler +from sklearn.decomposition import TruncatedSVD + + +# ── Prefixes produced by each type's .distance() ── +_DIST_PREFIXES: tuple[str, ...] = ( + '_ip_dist_', '_str_dist_', '_bool_dist_', '_bytes_dist_', + '_ts_dist_', '_enum_dist_', '_raw_', +) + +# ── Types that produce numeric distance/raw output (all 8 do) ── +_ALL_TYPES: set[str] = { + "枚举", "字节", "浮点", "整数", "布尔", "字符串", "IPv4", "时间戳", +} + + +class ColumnProcessor: + """Processes all columns in a dataset by their registered types. + + Groups columns by type, dispatches each group to the corresponding + :class:`~analysis.types.ColumnType` instance in + :data:`~analysis.types.TYPE_REGISTRY`, then chains normalisation + and SVD denoising. + + Parameters + ---------- + lf: + Input LazyFrame with raw data columns. + column_types: + ``{col_name: ColumnType}`` mapping from + :func:`~analysis.types.classify_schema`. + """ + + def __init__( + self, + lf: pl.LazyFrame, + column_types: dict, + ) -> None: + self._lf = lf + self._column_types = column_types + self._feature_cols: list[str] = [] + self._noise_profile: dict = {} + + # ── public API ────────────────────────────────────────────────────── + + @property + def feature_cols(self) -> list[str]: + """Numeric feature columns available after :meth:`compute_distances`.""" + return list(self._feature_cols) + + @property + def noise_profile(self) -> dict: + """SVD noise profile populated by :meth:`svd_denoise`.""" + return dict(self._noise_profile) + + @property + def lf(self) -> pl.LazyFrame: + """Current LazyFrame state.""" + return self._lf + + def compute_distances(self) -> pl.LazyFrame: + """Group columns by type and dispatch to TYPE_REGISTRY. + + Each :class:`~analysis.types.ColumnType` subclass handles all + columns of that type in one call. IP columns with fewer than + 2 entries are skipped (they need pairs). + + Returns the LazyFrame with distance/raw columns appended. + """ + from analysis.types import TYPE_REGISTRY + + # Group columns by type name + groups: dict[str, list[str]] = {} + for col_name, ct in self._column_types.items(): + groups.setdefault(ct.name, []).append(col_name) + + result = self._lf + + # Dispatch each group to its type's distance method + for type_name, cols in groups.items(): + type_obj = TYPE_REGISTRY.get(type_name) + if type_obj is None: + continue + # IP distance needs at least 2 columns (pairs) + if type_name == "IPv4" and len(cols) < 2: + continue + result = type_obj.distance(result, cols) + + # ── Collect feature columns produced by distance computation ── + schema = result.collect_schema() + self._feature_cols = [ + c for c in schema.names() + if any(c.startswith(p) for p in _DIST_PREFIXES) + ] + + self._lf = result + return result + + def normalize(self) -> pl.LazyFrame: + """Apply StandardScaler normalisation to feature columns. + + Uses :func:`analysis.distance.normalize_features` internally. + Requires :meth:`compute_distances` to have been called first. + """ + if not self._feature_cols: + return self._lf + + from analysis.distance import normalize_features + self._lf = normalize_features(self._lf, self._feature_cols) + return self._lf + + def svd_denoise(self, variance_threshold: float = 0.95) -> tuple[pl.LazyFrame, dict]: + """SVD noise removal on normalised feature columns. + + Uses :func:`analysis.distance.svd_denoise` internally. + Requires :meth:`normalize` to have been called first. + + Parameters + ---------- + variance_threshold: + Fraction of total variance to preserve (0.0–1.0). + + Returns + ------- + tuple[pl.LazyFrame, dict] + ``(denoised_lf, noise_profile)``. + """ + if not self._feature_cols: + self._noise_profile = { + 'kept_components': 0, 'total_components': 0, + 'kept_variance': 0.0, 'noise_components': [], + } + return self._lf, self._noise_profile + + from analysis.distance import svd_denoise as _svd_denoise + norm_cols = [f'_norm_{c}' for c in self._feature_cols] + self._lf, self._noise_profile = _svd_denoise( + self._lf, norm_cols, variance_threshold, + ) + return self._lf, self._noise_profile