tianxuan/analysis/entity_detector.py

"""Auto-detect entity column from CSV schema using heuristic scoring.

Provides :func:`detect_entity_column` which scores all string/categorical
columns by keyword-name match, unique-value ratio, and dtype, returning
a ranked list of candidates and a recommended column.

Usage::

    from analysis.entity_detector import detect_entity_column

    result = detect_entity_column(dataset_id)
    result['recommended']  # e.g. 'src_ip'
    result['candidates']   # sorted list of scored columns
"""
from __future__ import annotations

from typing import Optional

import polars as pl

from .session_store import SessionStore

# ---------------------------------------------------------------------------
#  Entity-related keywords (lowercased)
# ---------------------------------------------------------------------------

ENTITY_KEYWORDS: tuple[str, ...] = (
    # User exact column names only — no generic/fuzzy keywords
    ':ips', ':ipd', ':prs', ':prd',
    'server_ip', 'client_ip', 'scnt', 'dcnt',
    'cnam', 'snam', '0ver', '4dur', '8ses', '2tmo', '4ksz',
    'cnrs', 'isrs', '8ack', '8ppk', '8dbd',
    '1ipp', '4dbn', 'tabl', 'name', 'source_node',
    'cipher_suite', 'ecdhe_named_curve',
    '0cph', '0crv', '0rnd', '0rnt', 'row', 'time', 'timestamp',
    'latd', 'lond', 'ispn', 'orgn', 'city',
)


def _score_name(column: str) -> tuple[float, str]:
    """Score a column name against *ENTITY_KEYWORDS*.

    Returns ``(score, matched_keyword_or_empty)`` where *score* is 3.0
    for every matched keyword.  Partial matches (keyword is a substring
    of the lowercased column name) also count.
    """
    col_lower = column.lower().replace('-', '_').replace('.', '_').replace(' ', '_')
    for kw in ENTITY_KEYWORDS:
        if kw == col_lower or col_lower.endswith(f'_{kw}') or col_lower.startswith(f'{kw}_'):
            return 3.0, kw
    return 0.0, ''


def detect_entity_column(dataset_id: str, entity_columns: Optional[list[str]] = None) -> dict:
    """Score all string/categorical columns and return candidates.

    The scoring algorithm:

    *   **Name match** (``+3``): column name (normalised) is a keyword, or
        ends/starts with a keyword via ``_`` separator.
    *   **Unique ratio** (0–10): columns whose ratio of unique values to
        non-null values falls in the range ``[0.01, 0.50]`` receive up to
        10 points (``min(ratio * 100, 10)``).
    *   **String dtype** (``+2``): ``Utf8`` / ``String`` / ``Categorical``
        columns get a bonus.
    *   **Null penalty** (``*0.1``): if null ratio > 50%, score severely reduced.
    *   **Penalty** (``-5``): unique ratio > 0.50 (likely random IDs or
        transaction-level identifiers).

    Parameters
    ----------
    dataset_id:
        Dataset identifier in the session store.
    entity_columns:
        Optional list of column names to restrict detection to. When given,
        only these columns are scored (must exist in the schema).

    Returns
    -------
    dict
        ``candidates``: list of ``{column, score, unique_ratio, matched_keyword}``
        sorted descending by score.

        ``recommended``: highest-scoring column name (or ``''`` when none found).

        ``recommended_multi``: list of the top 3 candidate column names (may be
        shorter than 3 when fewer candidates exist).

        ``dataset_id``: the input identifier (echoed).
    """
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return {
            'dataset_id': dataset_id,
            'candidates': [],
            'recommended': '',
            'error': f'Dataset not found: {dataset_id}',
        }

    lf: pl.LazyFrame = entry['lazyframe']
    schema: dict[str, str] = entry.get('schema', {})

    # Identify string / categorical columns
    string_types = frozenset({'Utf8', 'String', 'Categorical', 'str', 'cat'})
    candidate_scores: list[dict] = []

    try:
        # Collect a sample to compute unique ratios efficiently
        df_sample = lf.collect(streaming=True)
    except Exception as exc:
        return {
            'dataset_id': dataset_id,
            'candidates': [],
            'recommended': '',
            'error': f'Failed to collect dataset: {exc}',
        }

    total_rows = len(df_sample)

    for col_name, dtype_str in schema.items():
        # If entity_columns specified, skip columns not in the list
        if entity_columns is not None and col_name not in entity_columns:
            continue

        # Only score string-like columns
        dtype_base = dtype_str.split('(')[0].strip()
        if dtype_base not in string_types:
            continue

        series_all = df_sample[col_name]
        non_null_series = series_all.drop_nulls()
        non_null_count = len(non_null_series)
        null_count = total_rows - non_null_count

        if non_null_count == 0:
            continue

        unique_count = int(non_null_series.n_unique())
        unique_ratio = unique_count / non_null_count

        # --- scoring ---
        name_score, matched_kw = _score_name(col_name)
        dtype_score = 2.0  # String/categorical bonus
        ratio_score = min(unique_ratio * 100.0, 10.0)
        penalty = -5.0 if unique_ratio > 0.50 else 0.0

        total_score = name_score + dtype_score + ratio_score + penalty

        # Null ratio penalty: if more than 50% null, severely reduce score
        if null_count / max(1, non_null_count) > 0.5:
            total_score *= 0.1

        candidate_scores.append({
            'column': col_name,
            'score': round(total_score, 2),
            'unique_ratio': round(unique_ratio, 4),
            'matched_keyword': matched_kw,
            'non_null_count': non_null_count,
            'unique_count': unique_count,
            'null_count': null_count,
        })

    # Sort descending by score, then by unique_count desc (prefer richer columns)
    candidate_scores.sort(key=lambda c: (c['score'], c['unique_count']), reverse=True)

    recommended = candidate_scores[0]['column'] if candidate_scores else ''
    recommended_multi = [c['column'] for c in candidate_scores[:3]]

    return {
        'dataset_id': dataset_id,
        'candidates': candidate_scores,
        'recommended': recommended,
        'recommended_multi': recommended_multi,  # NEW
        'total_rows': total_rows,
        'scored_columns': len(candidate_scores),
    }