tianxuan/analysis/column_processor.py

"""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