PM-pinou
334 lines
12 KiB
Python
334 lines
12 KiB
Python
"""Data profiling for LLM decision-making.
|
||
|
||
Produces column-level statistics and a correlation matrix, with built-in
|
||
context‑window truncation so the output stays under a configurable budget
|
||
(default 16 KB).
|
||
"""
|
||
import json
|
||
from typing import Any, Optional
|
||
|
||
import polars as pl
|
||
import numpy as np
|
||
|
||
|
||
MAX_PROFILE_SIZE_KB = 16
|
||
"""Maximum serialised profile size before truncation kicks in."""
|
||
|
||
MAX_SAMPLE_ROWS = 5000
|
||
"""Maximum rows to collect for profiling."""
|
||
|
||
MAX_COLUMNS_FULL = 100
|
||
"""Beyond this column count the profile is always truncated."""
|
||
|
||
MAX_CORRELATION_COLS = 20
|
||
"""Maximum numeric columns to include in the correlation matrix."""
|
||
|
||
_ENTITY_KEYWORDS = frozenset({
|
||
'ip', 'host', 'domain', 'user', 'client', 'src_ip', 'dst_ip',
|
||
'source', 'destination', 'address', 'entity', 'name', 'id',
|
||
'username', 'hostname', 'mac', 'email', 'phone', 'session',
|
||
})
|
||
|
||
|
||
def _is_float_dtype(dtype: pl.DataType) -> bool:
|
||
return dtype in (pl.Float32, pl.Float64)
|
||
|
||
|
||
def _is_int_dtype(dtype: pl.DataType) -> bool:
|
||
return dtype in (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
|
||
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64)
|
||
|
||
|
||
def _is_numeric_dtype(dtype: pl.DataType) -> bool:
|
||
return _is_float_dtype(dtype) or _is_int_dtype(dtype)
|
||
|
||
|
||
def _is_string_dtype(dtype: pl.DataType) -> bool:
|
||
return dtype in (pl.Utf8, pl.String, pl.Categorical)
|
||
|
||
|
||
def _is_temporal_dtype(dtype: pl.DataType) -> bool:
|
||
return dtype in (pl.Date, pl.Datetime, pl.Time, pl.Duration)
|
||
|
||
|
||
# ── column statistics ───────────────────────────────────────────────────
|
||
|
||
def _compute_column_stats(series: pl.Series) -> dict:
|
||
"""Return a dict of statistics for a single column, based on its dtype.
|
||
|
||
Numeric columns get min/max/mean/std/median/percentiles.
|
||
String columns get unique count, missing count, and top-5 frequencies.
|
||
Boolean columns get true/false/null counts.
|
||
Temporal columns get min/max/null_count.
|
||
"""
|
||
dtype = series.dtype
|
||
stats: dict = {
|
||
'dtype': str(dtype),
|
||
'null_count': int(series.null_count()),
|
||
'non_null_count': int(len(series) - series.null_count()),
|
||
}
|
||
|
||
if _is_numeric_dtype(dtype):
|
||
non_null = series.drop_nulls()
|
||
if len(non_null) > 0:
|
||
stats.update({
|
||
'min': _safe_float(non_null.min()),
|
||
'max': _safe_float(non_null.max()),
|
||
'mean': _safe_float(non_null.mean()),
|
||
'std': _safe_float(non_null.std()),
|
||
'median': _safe_float(non_null.median()),
|
||
'q25': _safe_float(non_null.quantile(0.25)),
|
||
'q75': _safe_float(non_null.quantile(0.75)),
|
||
'unique_count': int(non_null.n_unique()),
|
||
})
|
||
else:
|
||
stats.update({
|
||
'min': None, 'max': None, 'mean': None, 'std': None,
|
||
'median': None, 'q25': None, 'q75': None, 'unique_count': 0,
|
||
})
|
||
stats['type_category'] = 'numeric'
|
||
|
||
elif _is_string_dtype(dtype):
|
||
non_null = series.drop_nulls()
|
||
stats['unique_count'] = int(non_null.n_unique()) if len(non_null) > 0 else 0
|
||
stats['type_category'] = 'categorical'
|
||
|
||
# Top-5 value counts
|
||
if len(non_null) > 0:
|
||
vc = non_null.value_counts().sort('count', descending=True).head(5)
|
||
# Polars >=1.0 value_counts() uses original col name, not 'index'
|
||
val_col = vc.columns[0]
|
||
stats['top_values'] = [
|
||
{'value': str(row[val_col]), 'count': int(row['count'])}
|
||
for row in vc.iter_rows(named=True)
|
||
]
|
||
else:
|
||
stats['top_values'] = []
|
||
|
||
# Detect entity column
|
||
stats['is_entity_candidate'] = str(series.name).lower() in _ENTITY_KEYWORDS
|
||
|
||
elif dtype == pl.Boolean:
|
||
non_null = series.drop_nulls()
|
||
if len(non_null) > 0:
|
||
true_count = int(non_null.sum())
|
||
stats.update({
|
||
'true_count': true_count,
|
||
'false_count': int(len(non_null)) - true_count,
|
||
})
|
||
else:
|
||
stats.update({'true_count': 0, 'false_count': 0})
|
||
stats['type_category'] = 'boolean'
|
||
|
||
else:
|
||
stats['type_category'] = 'other'
|
||
stats['unique_count'] = int(series.drop_nulls().n_unique())
|
||
|
||
return stats
|
||
|
||
|
||
def _safe_float(val: Any) -> Optional[float]:
|
||
"""Convert *val* to float or return None on failure."""
|
||
if val is None:
|
||
return None
|
||
try:
|
||
f = float(val)
|
||
return None if np.isnan(f) or np.isinf(f) else round(f, 6)
|
||
except (ValueError, TypeError):
|
||
return None
|
||
|
||
|
||
# ── correlation matrix ──────────────────────────────────────────────────
|
||
|
||
def _compute_correlation_matrix(
|
||
df: pl.DataFrame,
|
||
top_n: int = MAX_CORRELATION_COLS,
|
||
) -> Optional[list[dict]]:
|
||
"""Pearson correlation for numeric columns.
|
||
|
||
Returns a list of ``{column_i, column_j, correlation}`` dicts for the
|
||
top *top_n* columns by mean absolute correlation, sorted descending by
|
||
absolute value. Returns *None* when fewer than 2 numeric columns exist.
|
||
"""
|
||
numeric_cols = [
|
||
col for col in df.columns
|
||
if _is_numeric_dtype(df[col].dtype)
|
||
]
|
||
if len(numeric_cols) < 2:
|
||
return None
|
||
|
||
# Limit to top_n numeric columns (pick by variance — highest first)
|
||
if len(numeric_cols) > top_n:
|
||
variances = []
|
||
for col in numeric_cols:
|
||
v = df[col].drop_nulls().std()
|
||
variances.append((col, 0.0 if v is None else float(v)))
|
||
variances.sort(key=lambda x: x[1], reverse=True)
|
||
numeric_cols = [c for c, _ in variances[:top_n]]
|
||
|
||
# Build correlation matrix
|
||
try:
|
||
import pandas as pd
|
||
pdf = df.select(numeric_cols).to_pandas()
|
||
corr = pdf.corr(method='pearson')
|
||
pairs: list[dict] = []
|
||
for i, col_i in enumerate(numeric_cols):
|
||
for col_j in numeric_cols[i + 1:]:
|
||
val = corr.loc[col_i, col_j]
|
||
if np.isnan(val):
|
||
continue
|
||
pairs.append({
|
||
'column_i': col_i,
|
||
'column_j': col_j,
|
||
'correlation': round(float(val), 4),
|
||
})
|
||
except ImportError:
|
||
# Pandas not available — compute simple correlation matrix via numpy
|
||
data = df.select(numeric_cols).to_numpy()
|
||
data = np.nan_to_num(data, nan=0.0)
|
||
corr_np = np.corrcoef(data.T)
|
||
pairs = []
|
||
for i in range(len(numeric_cols)):
|
||
for j in range(i + 1, len(numeric_cols)):
|
||
val = corr_np[i, j]
|
||
if np.isnan(val):
|
||
continue
|
||
pairs.append({
|
||
'column_i': numeric_cols[i],
|
||
'column_j': numeric_cols[j],
|
||
'correlation': round(float(val), 4),
|
||
})
|
||
|
||
# Sort by absolute correlation descending
|
||
pairs.sort(key=lambda x: abs(x['correlation']), reverse=True)
|
||
return pairs[:top_n * 2] # cap output
|
||
|
||
|
||
# ── truncation ──────────────────────────────────────────────────────────
|
||
|
||
def _estimate_json_size(obj: Any) -> int:
|
||
"""Return the length of *obj* when serialised to JSON."""
|
||
return len(json.dumps(obj, default=str, ensure_ascii=False))
|
||
|
||
|
||
def truncate_profile(
|
||
profile: dict,
|
||
max_kb: int = MAX_PROFILE_SIZE_KB,
|
||
) -> dict:
|
||
"""Remove low-value columns from *profile* until it fits in *max_kb*.
|
||
|
||
Sets ``truncated=True`` and adds ``omitted_columns`` if truncation
|
||
occurred. Columns with the least variance (or highest null rate) are
|
||
removed first.
|
||
"""
|
||
max_bytes = max_kb * 1024
|
||
if _estimate_json_size(profile) <= max_bytes:
|
||
profile['truncated'] = False
|
||
profile.setdefault('omitted_columns', [])
|
||
return profile
|
||
|
||
column_stats = profile.get('column_stats', {})
|
||
if not column_stats:
|
||
profile['truncated'] = True
|
||
return profile
|
||
|
||
# Score columns: low variance / high null rate → remove first
|
||
scored: list[tuple[str, float]] = []
|
||
for col_name, stats in column_stats.items():
|
||
null_rate = stats.get('null_count', 0) / max(stats.get('non_null_count', 1), 1)
|
||
# Numeric columns get a variance estimate
|
||
std = stats.get('std')
|
||
variance_score = float(std) if std is not None else 0.0
|
||
# Remove low-variance OR high-null columns
|
||
removal_score = -variance_score + null_rate * 100
|
||
scored.append((col_name, removal_score))
|
||
|
||
scored.sort(key=lambda x: x[1]) # most valuable first
|
||
|
||
kept_columns = []
|
||
omitted_columns = []
|
||
for col_name, _ in scored:
|
||
test_profile = dict(profile)
|
||
test_profile['column_stats'] = {
|
||
c: column_stats[c] for c in kept_columns + [col_name]
|
||
}
|
||
if _estimate_json_size(test_profile) <= max_bytes:
|
||
kept_columns.append(col_name)
|
||
else:
|
||
omitted_columns.append(col_name)
|
||
|
||
profile['column_stats'] = {
|
||
c: column_stats[c] for c in kept_columns
|
||
}
|
||
profile['truncated'] = True
|
||
profile['omitted_columns'] = omitted_columns
|
||
profile['truncation_note'] = (
|
||
f"{len(omitted_columns)} of {len(omitted_columns) + len(kept_columns)} "
|
||
f"columns omitted to stay under {max_kb} KB budget"
|
||
)
|
||
return profile
|
||
|
||
|
||
# ── main profiling entry point ──────────────────────────────────────────
|
||
|
||
def profile_dataset(
|
||
lf: pl.LazyFrame,
|
||
sample_size: int = 1000,
|
||
columns: Optional[list[str]] = None,
|
||
max_kb: int = MAX_PROFILE_SIZE_KB,
|
||
max_correlation_cols: int = MAX_CORRELATION_COLS,
|
||
) -> dict:
|
||
"""Compute column-level statistics and a correlation matrix.
|
||
|
||
Args:
|
||
lf: LazyFrame to profile.
|
||
sample_size: Maximum rows to collect (sampled if larger).
|
||
columns: Subset of columns to profile (``None`` = all).
|
||
max_kb: Truncation budget in KB.
|
||
max_correlation_cols: Max numeric columns for correlation.
|
||
|
||
Returns:
|
||
A dict with keys ``column_stats``, ``correlation_matrix``,
|
||
``truncated``, ``omitted_columns``, ``total_columns``,
|
||
``profiled_rows``, and ``sample_ratio``.
|
||
"""
|
||
# Collect (possibly sampled) dataframe
|
||
df = lf.collect(streaming=True)
|
||
total_rows = len(df)
|
||
sample_ratio = 1.0
|
||
|
||
if total_rows > sample_size:
|
||
df = df.sample(n=sample_size, seed=42)
|
||
sample_ratio = sample_size / total_rows
|
||
|
||
# Determine columns to profile
|
||
if columns:
|
||
profile_columns = [c for c in columns if c in df.columns]
|
||
else:
|
||
profile_columns = list(df.columns)
|
||
|
||
if len(profile_columns) > MAX_COLUMNS_FULL:
|
||
profile_columns = profile_columns[:MAX_COLUMNS_FULL]
|
||
|
||
# Per-column stats
|
||
column_stats: dict = {}
|
||
for col_name in profile_columns:
|
||
column_stats[col_name] = _compute_column_stats(df[col_name])
|
||
|
||
# Correlation matrix
|
||
correlation_matrix = _compute_correlation_matrix(df, top_n=max_correlation_cols)
|
||
|
||
result = {
|
||
'column_stats': column_stats,
|
||
'correlation_matrix': correlation_matrix,
|
||
'truncated': False,
|
||
'omitted_columns': [],
|
||
'total_columns': len(df.columns),
|
||
'profiled_columns': len(profile_columns),
|
||
'profiled_rows': len(df),
|
||
'total_rows': total_rows,
|
||
'sample_ratio': round(sample_ratio, 4),
|
||
}
|
||
|
||
return truncate_profile(result, max_kb=max_kb)
|