refactor(types+svd): OOP type system + SVD denoise replaces columns instead of _svd_ append
This commit is contained in:
PM-pinou
+17
-44
@@ -15,7 +15,7 @@ from typing import Optional
|
||||
import polars as pl
|
||||
import yaml
|
||||
|
||||
from analysis.type_classifier import classify_schema, DataType
|
||||
from analysis.types import classify_schema, TYPE_REGISTRY
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -624,7 +624,7 @@ def load_csv_directory(
|
||||
merged_schema = {k: v for k, v in merged_schema.items() if k not in existing_drop}
|
||||
logger.info('[CLEAN] dropped columns: %s', existing_drop)
|
||||
|
||||
# ── Cleaning: BOOL_ENUM normalisation ────────────────────────────────
|
||||
# ── Cleaning: BOOL normalisation ─────────────────────────────────────
|
||||
try:
|
||||
type_map = classify_schema(
|
||||
merged_lf.head(1000),
|
||||
@@ -633,21 +633,20 @@ def load_csv_directory(
|
||||
except Exception:
|
||||
type_map = {}
|
||||
|
||||
bool_enum_cols = [
|
||||
bool_cols = [
|
||||
c for c, t in type_map.items()
|
||||
if t == DataType.BOOL_ENUM and c in merged_schema
|
||||
if t.name == "布尔" and c in merged_schema
|
||||
]
|
||||
if bool_enum_cols:
|
||||
# Only clean string-typed BOOL_ENUM columns; native Boolean is already clean
|
||||
if bool_cols:
|
||||
schema_dtypes = dict(zip(
|
||||
merged_lf.collect_schema().names(),
|
||||
merged_lf.collect_schema().dtypes(),
|
||||
))
|
||||
bool_exprs = []
|
||||
cleaned_names: list[str] = []
|
||||
for c in bool_enum_cols:
|
||||
for c in bool_cols:
|
||||
if schema_dtypes.get(c) not in (pl.Utf8, pl.String):
|
||||
continue # native Boolean / numeric — skip
|
||||
continue
|
||||
c_str = pl.col(c).cast(pl.Utf8).str.strip_chars()
|
||||
cleaned = (
|
||||
pl.when(c_str == "")
|
||||
@@ -664,37 +663,12 @@ def load_csv_directory(
|
||||
merged_lf = merged_lf.with_columns(bool_exprs)
|
||||
for c in cleaned_names:
|
||||
merged_schema[c] = 'Utf8'
|
||||
logger.info('[CLEAN] BOOL_ENUM columns normalised: %s', cleaned_names)
|
||||
logger.info('[CLEAN] BOOL columns normalised: %s', cleaned_names)
|
||||
|
||||
# ── Cleaning: LAT_LON normalisation ────────────────────────────────
|
||||
latlon_cols = [
|
||||
c for c, t in type_map.items()
|
||||
if t == DataType.LAT_LON and c in merged_schema
|
||||
]
|
||||
if latlon_cols:
|
||||
schema_dtypes_ll = dict(zip(
|
||||
merged_lf.collect_schema().names(),
|
||||
merged_lf.collect_schema().dtypes(),
|
||||
))
|
||||
latlon_exprs = []
|
||||
for c in latlon_cols:
|
||||
dtype = schema_dtypes_ll.get(c)
|
||||
if dtype in (pl.Utf8, pl.String):
|
||||
latlon_exprs.append(
|
||||
pl.col(c).map_batches(_coerce_to_float, return_dtype=pl.Float64).alias(c)
|
||||
)
|
||||
merged_schema[c] = 'Float64'
|
||||
if latlon_exprs:
|
||||
merged_lf = merged_lf.with_columns(latlon_exprs)
|
||||
logger.info('[CLEAN] LAT_LON columns normalised to Float64: %s', latlon_cols)
|
||||
|
||||
# ── Cleaning: Generic numeric coercion ────────────────────────────
|
||||
# For any remaining string column (not IPv4/URL/HEX/ENUM/etc.),
|
||||
# try to coerce to Float64. Whatever can't parse becomes null.
|
||||
# This prevents "mean on str column" crashes for ANY column.
|
||||
_SKIP_NUM_COERCE = (DataType.IPv4, DataType.URL, DataType.HEX,
|
||||
DataType.ENUM, DataType.LAT_LON, DataType.BOOL_ENUM)
|
||||
# Also protect columns with these name patterns (MAC addresses, etc.)
|
||||
# ── Generic numeric coercion ────────────────────────────────────────
|
||||
# For remaining string columns (not IPv4/BYTES/ENUM/BOOL etc.),
|
||||
# try to coerce to Float64.
|
||||
_SKIP_NUM_COERCE = {"IPv4", "字节", "枚举", "布尔"}
|
||||
_SKIP_NAME_PATTERNS = ('mac', 'mac_', '_mac', 'addr', 'address')
|
||||
schema_dtypes_gn = dict(zip(
|
||||
merged_lf.collect_schema().names(),
|
||||
@@ -702,14 +676,13 @@ def load_csv_directory(
|
||||
))
|
||||
generic_numeric_cols = []
|
||||
for c, t in type_map.items():
|
||||
if t in _SKIP_NUM_COERCE:
|
||||
if t.name in _SKIP_NUM_COERCE:
|
||||
continue
|
||||
if c not in merged_schema:
|
||||
continue
|
||||
dtype = schema_dtypes_gn.get(c)
|
||||
if dtype not in (pl.Utf8, pl.String):
|
||||
continue # already numeric
|
||||
# Skip columns with protected name patterns
|
||||
continue
|
||||
cl = c.lower().replace('-', '_').replace(' ', '_')
|
||||
if any(p in cl for p in _SKIP_NAME_PATTERNS):
|
||||
continue
|
||||
@@ -727,10 +700,10 @@ def load_csv_directory(
|
||||
merged_schema[c] = 'Float64'
|
||||
logger.info('[CLEAN] Generic numeric coercion applied: %s', generic_numeric_cols)
|
||||
|
||||
# ── Cleaning: HEX preprocessing (hex_pairs_count) ────────────────────
|
||||
# ── Cleaning: BYTES (hex) preprocessing ─────────────────────────────
|
||||
hex_cols = [
|
||||
c for c, t in type_map.items()
|
||||
if t == DataType.HEX and c in merged_schema
|
||||
if t.name == "字节" and c in merged_schema
|
||||
]
|
||||
if hex_cols:
|
||||
hex_exprs = []
|
||||
@@ -741,7 +714,7 @@ def load_csv_directory(
|
||||
merged_lf = merged_lf.with_columns(hex_exprs)
|
||||
for c in hex_cols:
|
||||
merged_schema[f"{c}_hex_pairs_count"] = 'UInt32'
|
||||
logger.info('[CLEAN] HEX columns added hex_pairs_count: %s', hex_cols)
|
||||
logger.info('[CLEAN] BYTES columns added hex_pairs_count: %s', hex_cols)
|
||||
|
||||
# Row count estimate (fast path: use first file row count × files)
|
||||
# For accurate count we'd need to collect, but that defeats lazy.
|
||||
|
||||
+74
-45
@@ -679,40 +679,49 @@ def svd_denoise(
|
||||
lf: pl.LazyFrame,
|
||||
feature_cols: list[str],
|
||||
variance_threshold: float = 0.95,
|
||||
) -> pl.LazyFrame:
|
||||
) -> tuple[pl.LazyFrame, dict]:
|
||||
"""Remove noise from numeric features via TruncatedSVD.
|
||||
|
||||
1. Collects *feature_cols* → numpy matrix.
|
||||
2. Fits ``TruncatedSVD(n_components=min(n_rows, n_features − 1))``.
|
||||
3. Keeps *k* components where the cumulative explained variance
|
||||
ratio ``trace(S[:k]) / trace(S) ≥ variance_threshold``
|
||||
(default 0.95).
|
||||
4. Inverse-transforms to reconstruct the denoised matrix.
|
||||
5. Returns *lf* extended with ``_svd_{col}`` columns.
|
||||
1. Fits TruncatedSVD on *feature_cols*.
|
||||
2. Keeps *k* components where cumulative explained variance
|
||||
``trace(S[:k]) / trace(S) ≥ variance_threshold`` (default 0.95).
|
||||
3. Inverse-transforms to reconstruct the denoised matrix.
|
||||
4. **Replaces** the original *feature_cols* with denoised values.
|
||||
|
||||
The noise components (removed by SVD) are returned as a dict for
|
||||
display, NOT as new feature columns.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
lf:
|
||||
LazyFrame containing *feature_cols*.
|
||||
feature_cols:
|
||||
Numeric column names to denoise. Columns not present in
|
||||
*lf* are silently skipped.
|
||||
Numeric column names to denoise. These columns are replaced
|
||||
in-place by their denoised reconstructions.
|
||||
variance_threshold:
|
||||
Fraction of total variance to preserve (``0.0–1.0``).
|
||||
Default ``0.95`` keeps 95 % of information.
|
||||
Fraction of total variance to preserve (0.0–1.0).
|
||||
Default 0.95 keeps 95 % of signal and treats 5 % as noise.
|
||||
|
||||
Returns
|
||||
-------
|
||||
pl.LazyFrame
|
||||
Input frame with additional ``_svd_{col}`` columns.
|
||||
|
||||
Memory
|
||||
------
|
||||
Peak ≈ numpy matrix + SVD workspace; typically <800 MB.
|
||||
tuple[pl.LazyFrame, dict]
|
||||
``(denoised_lf, noise_profile)`` where *noise_profile* contains:
|
||||
- ``kept_components`` : int — number of components kept
|
||||
- ``total_components`` : int — max possible components
|
||||
- ``kept_variance`` : float — fraction of variance kept
|
||||
- ``noise_components`` : list[dict] — per-component noise info
|
||||
(component index, explained variance, top-5 contributing features)
|
||||
"""
|
||||
noise_profile = {
|
||||
'kept_components': 0,
|
||||
'total_components': 0,
|
||||
'kept_variance': 0.0,
|
||||
'noise_components': [],
|
||||
}
|
||||
|
||||
available = [c for c in feature_cols if c in lf.collect_schema().names()]
|
||||
if not available:
|
||||
return lf
|
||||
return lf, noise_profile
|
||||
|
||||
df = lf.select(available).collect(streaming=True)
|
||||
mat = df.to_numpy().astype(np.float64)
|
||||
@@ -721,17 +730,15 @@ def svd_denoise(
|
||||
n, d = mat.shape
|
||||
max_components = min(n, d - 1)
|
||||
if max_components < 1:
|
||||
# Not enough dimensions; return as-is
|
||||
return lf
|
||||
return lf, noise_profile
|
||||
|
||||
from sklearn.decomposition import TruncatedSVD
|
||||
svd = TruncatedSVD(n_components=max_components, random_state=42)
|
||||
_ = svd.fit_transform(mat)
|
||||
|
||||
# Determine k: cumulative variance threshold
|
||||
total_variance = float(np.sum(svd.explained_variance_))
|
||||
if total_variance <= 1e-10:
|
||||
# All-zero variance; denoising would be a no-op
|
||||
return lf
|
||||
return lf, noise_profile
|
||||
|
||||
cumsum = np.cumsum(svd.explained_variance_)
|
||||
for k in range(1, max_components + 1):
|
||||
@@ -740,18 +747,44 @@ def svd_denoise(
|
||||
else:
|
||||
k = max_components
|
||||
|
||||
# Re-fit with chosen k and inverse-transform
|
||||
# Re-fit with chosen k
|
||||
svd_k = TruncatedSVD(n_components=k, random_state=42)
|
||||
reduced = svd_k.fit_transform(mat)
|
||||
reconstructed = svd_k.inverse_transform(reduced)
|
||||
|
||||
# --- Build noise profile (removed components) ---
|
||||
# The "noise" is the variance in components k..max_components
|
||||
kept_var_ratio = cumsum[k - 1] / total_variance if k > 0 else 0.0
|
||||
noise_profile = {
|
||||
'kept_components': k,
|
||||
'total_components': max_components,
|
||||
'kept_variance': round(float(kept_var_ratio), 4),
|
||||
}
|
||||
|
||||
# Per-component noise info for removed components (k and beyond)
|
||||
noise_info = []
|
||||
for comp_idx in range(k, min(k + 10, max_components)):
|
||||
loading = np.abs(svd.components_[comp_idx])
|
||||
top5_idx = np.argsort(loading)[-5:][::-1]
|
||||
top5 = [
|
||||
{'feature': available[i], 'strength': round(float(loading[i]), 4)}
|
||||
for i in top5_idx
|
||||
]
|
||||
noise_info.append({
|
||||
'component': int(comp_idx + 1),
|
||||
'explained_var_ratio': round(float(svd.explained_variance_ratio_[comp_idx]) * 100, 3),
|
||||
'top_features': top5,
|
||||
})
|
||||
noise_profile['noise_components'] = noise_info
|
||||
|
||||
# --- Replace original columns with denoised values ---
|
||||
df_result = lf.collect(streaming=True)
|
||||
for i, col in enumerate(available):
|
||||
df_result = df_result.with_columns(
|
||||
pl.Series(f'_svd_{col}', reconstructed[:, i])
|
||||
pl.Series(col, reconstructed[:, i])
|
||||
)
|
||||
|
||||
return df_result.lazy()
|
||||
return df_result.lazy(), noise_profile
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -795,12 +828,9 @@ def run_preprocessing_pipeline(
|
||||
(``_svd__norm__*_dist_*`` etc.) plus original row anchors.
|
||||
"""
|
||||
if column_types is None:
|
||||
from analysis.type_classifier import classify_schema
|
||||
from analysis.types import classify_schema
|
||||
column_types = classify_schema(lf)
|
||||
|
||||
# Group columns by type
|
||||
from analysis.type_classifier import DataType
|
||||
|
||||
ip_cols: list[str] = []
|
||||
str_cols: list[str] = []
|
||||
bool_cols: list[str] = []
|
||||
@@ -809,24 +839,23 @@ def run_preprocessing_pipeline(
|
||||
float_cols: list[str] = []
|
||||
int_cols: list[str] = []
|
||||
|
||||
for col_name, dtype in column_types.items():
|
||||
if dtype == DataType.IPv4:
|
||||
for col_name, ct in column_types.items():
|
||||
if ct.name == "IPv4":
|
||||
ip_cols.append(col_name)
|
||||
elif dtype == DataType.STRING:
|
||||
elif ct.name == "字符串":
|
||||
str_cols.append(col_name)
|
||||
elif dtype == DataType.BOOL_ENUM:
|
||||
elif ct.name == "布尔":
|
||||
bool_cols.append(col_name)
|
||||
elif dtype == DataType.HEX:
|
||||
elif ct.name == "字节":
|
||||
byte_cols.append(col_name)
|
||||
elif dtype == DataType.TIMESTAMP:
|
||||
elif ct.name == "时间戳":
|
||||
ts_cols.append(col_name)
|
||||
elif dtype == DataType.FLOAT:
|
||||
elif ct.name in ("浮点", "整数"):
|
||||
# Float/Int are kept as direct numeric features
|
||||
if ct.name == "浮点":
|
||||
float_cols.append(col_name)
|
||||
elif dtype == DataType.INT:
|
||||
else:
|
||||
int_cols.append(col_name)
|
||||
# LAT_LON, ENUM, URL are kept as-is for now (they either
|
||||
# need special distance functions or don't contribute to
|
||||
# numeric features directly)
|
||||
|
||||
result = lf
|
||||
|
||||
@@ -863,11 +892,11 @@ def run_preprocessing_pipeline(
|
||||
# -- Step 3: Normalisation --
|
||||
result = normalize_features(result, feature_cols)
|
||||
|
||||
# -- Step 4: SVD denoise --
|
||||
# -- Step 4: SVD denoise (REPLACES noisy features with cleaned ones) --
|
||||
norm_cols = [f'_norm_{c}' for c in feature_cols]
|
||||
result = svd_denoise(result, norm_cols)
|
||||
result, noise_profile = svd_denoise(result, norm_cols)
|
||||
|
||||
return result
|
||||
return result, noise_profile
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
+3
-415
@@ -1,415 +1,3 @@
|
||||
"""Data type classifier for column-level type inference.
|
||||
|
||||
Provides :class:`DataType` enumeration and functions to classify Polars
|
||||
columns based on name heuristics and value sampling.
|
||||
|
||||
Typical usage::
|
||||
|
||||
from analysis.type_classifier import classify_column, classify_schema, DataType
|
||||
|
||||
# Single column
|
||||
dtype = classify_column('src_ip', series, config_type='ipv4')
|
||||
|
||||
# Full schema
|
||||
type_map = classify_schema(lazyframe, config_overrides={'src_ip': 'ipv4'})
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from enum import Enum
|
||||
|
||||
import polars as pl
|
||||
|
||||
from analysis.tls_ref import get_tls_ref_dict
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TLS version hex-to-label mapping
|
||||
# (Duplicated in views.py for globe visualization — keep both in sync.)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
TLS_HEX_MAP: dict[str, str] = {
|
||||
'0303': 'TLSv1.2',
|
||||
'0304': 'TLSv1.3',
|
||||
'0302': 'TLSv1.1',
|
||||
'0301': 'TLSv1.0',
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DataType enumeration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class DataType(Enum):
|
||||
"""Enumerated column types recognised by the classifier engine."""
|
||||
|
||||
INT = 1
|
||||
FLOAT = 2
|
||||
ENUM = 3
|
||||
HEX = 4
|
||||
URL = 5
|
||||
IPv4 = 6
|
||||
LAT_LON = 7
|
||||
TIMESTAMP = 8
|
||||
BOOL_ENUM = 9
|
||||
STRING = 10
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Detection patterns (raw strings for regex)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
HEX_PATTERN = r'^[0-9a-f]{2}(?:\s[0-9a-f]{2})*$'
|
||||
"""Matches ``"0e a6 3f"``, ``"ab cd ef"`` etc. (space-separated hex pairs)."""
|
||||
|
||||
URL_PATTERN = r'^https?://'
|
||||
"""Matches URLs starting with ``http://`` or ``https://``."""
|
||||
|
||||
IPv4_PATTERN = r'^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$'
|
||||
"""Matches ``"192.168.1.1"`` style IPv4 addresses."""
|
||||
|
||||
MAC_PATTERN = r'^([0-9A-Fa-f]{2}[:.-]){5}[0-9A-Fa-f]{2}$'
|
||||
"""Matches MAC addresses with ``:``, ``-``, or ``.`` separators."""
|
||||
|
||||
BOOL_ENUM_VALUES = (
|
||||
'', ' ', '+', '-', '0', '1',
|
||||
'true', 'false', 'yes', 'no', 't', 'f',
|
||||
)
|
||||
"""String values considered boolean / binary indicators."""
|
||||
|
||||
LAT_LON_NAMES = ('lat', 'latitude', 'lon', 'lng', 'longitude')
|
||||
|
||||
# Name-to-type heuristics (checked after config override, before values)
|
||||
_NAME_TYPE_MAP: list[tuple[re.Pattern, DataType]] = [
|
||||
# User column names - EXACT match only (no generic fuzzy patterns)
|
||||
# IPv4 columns
|
||||
(re.compile(r'^:ips$|^:ipd$|^server.ip$|^client.ip$', re.I), DataType.IPv4),
|
||||
# ENUM columns
|
||||
(re.compile(r'^:prs$|^:prd$', re.I), DataType.ENUM),
|
||||
(re.compile(r'^(?:cnrs|isrs)$', re.I), DataType.BOOL_ENUM),
|
||||
(re.compile(r'^(?:scnt|dcnt|1ipp|4dbn|4srs|eiph)$', re.I), DataType.ENUM),
|
||||
(re.compile(r'^0ver$', re.I), DataType.ENUM),
|
||||
# FLOAT columns
|
||||
(re.compile(r'^4dur$|^8ses$|^2tmo$|^0rnt$', re.I), DataType.FLOAT),
|
||||
# INT columns
|
||||
(re.compile(r'^4ksz$|^8ack$|^8ppk$|^8pak$|^8seq$|^8did$|^8byt$|^row$', re.I), DataType.INT),
|
||||
# HEX columns
|
||||
(re.compile(r'^0cph$|^0crv$', re.I), DataType.HEX),
|
||||
# TIMESTAMP columns
|
||||
(re.compile(r'^8dbd$|^time$|^timestamp$', re.I), DataType.TIMESTAMP),
|
||||
# Dot-notation suffixes for lat/lon (e.g. :ips.latd, :ipd.lond)
|
||||
(re.compile(r'\.latd$', re.I), DataType.LAT_LON),
|
||||
(re.compile(r'\.lond$', re.I), DataType.LAT_LON),
|
||||
# Dot-notation STRING suffixes (e.g. :ips.ispn, :ipd.city)
|
||||
(re.compile(r'\.(ispn|orgn|city|anon)$', re.I), DataType.STRING),
|
||||
# Standalone string columns - exact match
|
||||
(re.compile(r'^(?:cnam|snam|tabl|name|0rnd|source.node|ecdhe.named.curve|crcc|orga|orgu)$', re.I), DataType.STRING),
|
||||
# Special columns with @ prefix
|
||||
(re.compile(r'^@', re.I), DataType.STRING),
|
||||
]
|
||||
|
||||
# Mapping from config.yaml type strings to DataType
|
||||
_CONFIG_TYPE_MAP: dict[str, DataType] = {
|
||||
'int': DataType.INT,
|
||||
'float': DataType.FLOAT,
|
||||
'enum': DataType.ENUM,
|
||||
'hex': DataType.HEX,
|
||||
'url': DataType.URL,
|
||||
'ipv4': DataType.IPv4,
|
||||
'lat_lon': DataType.LAT_LON,
|
||||
'timestamp': DataType.TIMESTAMP,
|
||||
'bool_enum': DataType.BOOL_ENUM,
|
||||
'string': DataType.STRING,
|
||||
}
|
||||
|
||||
# Mapping from tls_ref_data ``data_type`` column values to DataType
|
||||
TLSDB_TYPE_MAP: dict[str, DataType] = {
|
||||
'IPv4': DataType.IPv4,
|
||||
'浮点': DataType.FLOAT,
|
||||
'字符串': DataType.STRING,
|
||||
'整数': DataType.INT,
|
||||
'2小写字母缩写': DataType.ENUM,
|
||||
'XX:YY.Z': DataType.STRING,
|
||||
'2字节': DataType.HEX,
|
||||
'2大写字母': DataType.STRING,
|
||||
'枚举': DataType.ENUM,
|
||||
'布尔': DataType.BOOL_ENUM,
|
||||
'字符串或整数': DataType.STRING,
|
||||
}
|
||||
|
||||
|
||||
def _tlsdb_type_to_datatype(tlsdb_type_str: str) -> DataType | None:
|
||||
"""Convert a TlsDB ``data_type`` string to a :class:`DataType`.
|
||||
|
||||
Returns ``None`` if the type string is unknown, allowing the
|
||||
caller to fall through to value-based detection.
|
||||
"""
|
||||
return TLSDB_TYPE_MAP.get(tlsdb_type_str)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Column classification (single column)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _name_to_type(name: str) -> DataType | None:
|
||||
"""Match column name against known heuristics.
|
||||
|
||||
Returns a :class:`DataType` if the name matches strongly, otherwise
|
||||
``None`` (fall through to value-based detection).
|
||||
"""
|
||||
for pattern, dtype in _NAME_TYPE_MAP:
|
||||
if pattern.search(name):
|
||||
return dtype
|
||||
return None
|
||||
|
||||
|
||||
def _is_port_value(v: str) -> bool:
|
||||
"""Check if *v* is an integer string in the registered port range (1-65535)."""
|
||||
try:
|
||||
n = int(v)
|
||||
return 1 <= n <= 65535
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def _can_float(v: str) -> bool:
|
||||
"""Return ``True`` when *v* can be parsed as a Python float."""
|
||||
try:
|
||||
float(v)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def _values_to_type(series: pl.Series, col_name: str | None = None) -> DataType:
|
||||
"""Infer column type by sampling non-null values.
|
||||
|
||||
Detection order: IPv4 → MAC guard → HEX → URL → PORT → BOOL_ENUM
|
||||
→ LAT_LON → ENUM → FLOAT → STRING
|
||||
"""
|
||||
dtype = series.dtype
|
||||
|
||||
# Boolean dtype ──────────────────────────────────────────────────────
|
||||
if dtype == pl.Boolean:
|
||||
return DataType.BOOL_ENUM
|
||||
|
||||
# Numeric dtypes ─────────────────────────────────────────────────────
|
||||
if dtype in (pl.Float32, pl.Float64):
|
||||
# Check LAT_LON possibility before returning FLOAT
|
||||
non_null = series.drop_nulls()
|
||||
if len(non_null) > 0:
|
||||
sample = non_null.head(100)
|
||||
sample_strs = [str(v).strip().lower() for v in sample]
|
||||
non_blank = [v for v in sample_strs if v not in ('', ' ', '+', '-')]
|
||||
if non_blank:
|
||||
try:
|
||||
floats = [float(v) for v in non_blank]
|
||||
max_abs = max(abs(f) for f in floats)
|
||||
if (
|
||||
all(-180 <= f <= 180 for f in floats)
|
||||
and any('.' in v for v in non_blank)
|
||||
and max_abs > 10
|
||||
):
|
||||
return DataType.LAT_LON
|
||||
except ValueError:
|
||||
pass
|
||||
return DataType.FLOAT
|
||||
|
||||
if dtype in (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
|
||||
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64):
|
||||
return DataType.INT
|
||||
|
||||
# String / categorical types ─────────────────────────────────────────
|
||||
non_null = series.drop_nulls()
|
||||
if len(non_null) == 0:
|
||||
return DataType.STRING
|
||||
|
||||
sample = non_null.head(100)
|
||||
# Convert all to string representations for pattern matching
|
||||
sample_strs = [str(v).strip().lower() for v in sample]
|
||||
|
||||
# ── 1. IPv4 detection ──────────────────────────────────────────────
|
||||
if all(re.match(IPv4_PATTERN, v) for v in sample_strs):
|
||||
if all(_valid_ip_octets(v) for v in sample_strs):
|
||||
return DataType.IPv4
|
||||
|
||||
# ── 2. MAC guard (skip HEX if values look like MAC addresses) ──────
|
||||
is_mac = all(re.match(MAC_PATTERN, v) for v in sample_strs)
|
||||
|
||||
# ── 3. HEX detection ───────────────────────────────────────────────
|
||||
if not is_mac and all(re.match(HEX_PATTERN, v) for v in sample_strs):
|
||||
return DataType.HEX
|
||||
|
||||
# ── 4. URL detection ───────────────────────────────────────────────
|
||||
if all(re.match(URL_PATTERN, v) for v in sample_strs):
|
||||
return DataType.URL
|
||||
|
||||
# ── 5. PORT detection → ENUM ───────────────────────────────────────
|
||||
if all(_is_port_value(v) for v in sample_strs):
|
||||
return DataType.ENUM
|
||||
|
||||
# ── 6. BOOL_ENUM detection ─────────────────────────────────────────
|
||||
if all(v in BOOL_ENUM_VALUES for v in sample_strs):
|
||||
return DataType.BOOL_ENUM
|
||||
|
||||
# ── 7. LAT_LON detection (value-based) ──────────────────────────────
|
||||
non_blank = [v for v in sample_strs if v not in ('', ' ', '+', '-')]
|
||||
if non_blank:
|
||||
try:
|
||||
floats = [float(v) for v in non_blank]
|
||||
max_abs = max(abs(f) for f in floats)
|
||||
if (
|
||||
all(-180 <= f <= 180 for f in floats)
|
||||
and any('.' in v for v in non_blank)
|
||||
and max_abs > 10
|
||||
):
|
||||
return DataType.LAT_LON
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# ── 8. ENUM detection (few unique values) ──────────────────────────
|
||||
unique_count = non_null.n_unique()
|
||||
if unique_count < 20:
|
||||
# Safety: if >80 % of the sample values are unique strings,
|
||||
# the column is high-cardinality — classify as STRING, not ENUM.
|
||||
sample_unique_count = len(set(sample_strs))
|
||||
if sample_unique_count > 0.8 * len(sample_strs):
|
||||
return DataType.STRING
|
||||
return DataType.ENUM
|
||||
|
||||
# ── 9. TLS version hex detection (before STRING fallback) ──────────
|
||||
if col_name is not None and col_name.lower() in ('0ver', '_tlsver', 'tls_version', 'version'):
|
||||
sample_hex = sample_strs[0].replace(' ', '').strip()
|
||||
if sample_hex.upper() in TLS_HEX_MAP:
|
||||
return DataType.ENUM
|
||||
|
||||
# ── 10. FLOAT detection (>80 % can float) ──────────────────────────
|
||||
# Exclude strings starting with '+' from FLOAT detection — these are
|
||||
# text markers (phone codes, IDs, error codes), not real numbers.
|
||||
if any(v.startswith('+') for v in sample_strs):
|
||||
return DataType.STRING
|
||||
|
||||
float_count = sum(1 for v in sample_strs if _can_float(v))
|
||||
if float_count / len(sample_strs) > 0.8:
|
||||
# Guard: values with leading "-" that happen to parse as float
|
||||
# should not force FLOAT when non-float values also exist.
|
||||
if any(v.startswith('-') for v in sample_strs):
|
||||
if float_count < len(sample_strs):
|
||||
return DataType.STRING
|
||||
return DataType.FLOAT
|
||||
|
||||
return DataType.STRING
|
||||
|
||||
|
||||
def _valid_ip_octets(ip_str: str) -> bool:
|
||||
"""Check that each octet in an IPv4 string is 0-255."""
|
||||
parts = ip_str.split('.')
|
||||
if len(parts) != 4:
|
||||
return False
|
||||
for p in parts:
|
||||
try:
|
||||
n = int(p)
|
||||
if n < 0 or n > 255:
|
||||
return False
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def classify_column(
|
||||
name: str,
|
||||
series: pl.Series,
|
||||
config_type: str | None = None,
|
||||
) -> DataType:
|
||||
"""Classify a single column based on its *name* and *series* values.
|
||||
|
||||
Priority
|
||||
--------
|
||||
1. **config_type** — user override from ``config.yaml``.
|
||||
If provided (e.g. ``'ipv4'``), returns the corresponding
|
||||
:class:`DataType` immediately.
|
||||
2. **tls_ref_data** — authoritative type from ``TlsDB.csv`` /
|
||||
``tls_ref_data`` table. Looked up by column name; if found,
|
||||
the declared type is used directly (no value-based guessing).
|
||||
3. **Value-based** pattern matching — samples up to 100 non-null
|
||||
values and tests against IPv4 → HEX → URL → PORT → BOOL_ENUM
|
||||
→ LAT_LON → ENUM → FLOAT in order.
|
||||
4. **Name-based** heuristics — only consulted when value-based
|
||||
detection returned ``STRING`` (i.e. no value pattern matched).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name:
|
||||
Column name (used for name-based heuristics and tls_ref_data
|
||||
lookup).
|
||||
series:
|
||||
Polars :class:`Series` containing the column's data.
|
||||
config_type:
|
||||
Optional user override (case-insensitive, matched to
|
||||
:data:`_CONFIG_TYPE_MAP`).
|
||||
|
||||
Returns
|
||||
-------
|
||||
DataType
|
||||
"""
|
||||
# 1. Config override (highest priority)
|
||||
if config_type is not None:
|
||||
mapped = _CONFIG_TYPE_MAP.get(config_type.lower().strip())
|
||||
if mapped is not None:
|
||||
return mapped
|
||||
|
||||
# 2. tls_ref_data lookup (authoritative type from TlsDB)
|
||||
tls_ref = get_tls_ref_dict()
|
||||
if name in tls_ref:
|
||||
dtype = _tlsdb_type_to_datatype(tls_ref[name]['data_type'])
|
||||
if dtype is not None:
|
||||
return dtype
|
||||
|
||||
# 3. Value-based sampling
|
||||
result = _values_to_type(series, name)
|
||||
|
||||
# 4. Name-based heuristics (only when values returned STRING)
|
||||
if result == DataType.STRING:
|
||||
name_match = _name_to_type(name)
|
||||
if name_match is not None:
|
||||
return name_match
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Schema classification (all columns)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def classify_schema(
|
||||
lf: pl.LazyFrame,
|
||||
config_overrides: dict[str, str] | None = None,
|
||||
) -> dict[str, DataType]:
|
||||
"""Classify all columns in a :class:`LazyFrame`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
lf:
|
||||
A Polars :class:`LazyFrame` whose columns to classify.
|
||||
config_overrides:
|
||||
Optional mapping of ``{col_name: type_string}`` from a
|
||||
``config.yaml`` ``columns:`` section. These take highest
|
||||
priority (see :func:`classify_column`).
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict[str, DataType]
|
||||
Mapping from column name to inferred :class:`DataType`.
|
||||
"""
|
||||
config_overrides = config_overrides or {}
|
||||
df_sample = lf.collect(streaming=True)
|
||||
result: dict[str, DataType] = {}
|
||||
for col_name in df_sample.columns:
|
||||
config_type = config_overrides.get(col_name)
|
||||
result[col_name] = classify_column(
|
||||
name=col_name,
|
||||
series=df_sample[col_name],
|
||||
config_type=config_type,
|
||||
)
|
||||
return result
|
||||
"""Re-export from new types module for backward compatibility."""
|
||||
from analysis.types import ColumnType, TYPE_REGISTRY, TYPE_ALIASES, get_type, classify_column, classify_schema
|
||||
TLS_HEX_MAP = {"0303": "TLSv1.2", "0304": "TLSv1.3", "0302": "TLSv1.1", "0301": "TLSv1.0"}
|
||||
|
||||
@@ -0,0 +1,300 @@
|
||||
"""天璇类型系统 — 仅8种固定类型，面向对象接口。
|
||||
|
||||
每个类型继承 :class:`ColumnType`，实现:
|
||||
- ``distance()`` — 为该类型列计算距离矩阵（追加到 LazyFrame）
|
||||
- ``display()`` — 将值格式化为可读字符串
|
||||
- ``polars_dtype`` — 映射到 Polars 数据类型
|
||||
- ``name`` — 中文名称
|
||||
|
||||
:data:`TYPE_REGISTRY` 是唯一的类型注册表，按名称索引。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Base type — all 8 types inherit from this
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class ColumnType(ABC):
|
||||
"""Base class for all column types.
|
||||
|
||||
Subclasses MUST define:
|
||||
- ``name`` (class attribute, Chinese)
|
||||
- ``polars_dtype`` (class attribute)
|
||||
- ``distance()``
|
||||
- ``display()``
|
||||
"""
|
||||
|
||||
name: str = ""
|
||||
polars_dtype: pl.DataType = pl.Utf8
|
||||
|
||||
@abstractmethod
|
||||
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
||||
"""Compute per-row distances for *columns* of this type.
|
||||
|
||||
Returns a new LazyFrame with distance columns appended
|
||||
(prefixed ``_dist_``).
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def display(self, value: Any) -> str:
|
||||
"""Format *value* for human-readable display."""
|
||||
...
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# 8 concrete types
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class IPv4Type(ColumnType):
|
||||
name = "IPv4"
|
||||
polars_dtype = pl.Utf8
|
||||
|
||||
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
||||
if len(columns) < 2:
|
||||
return lf
|
||||
from analysis.distance import ip_distance
|
||||
return ip_distance(lf, columns)
|
||||
|
||||
def display(self, value: Any) -> str:
|
||||
return str(value) if value else "-"
|
||||
|
||||
|
||||
class StringType(ColumnType):
|
||||
name = "字符串"
|
||||
polars_dtype = pl.Utf8
|
||||
|
||||
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
||||
from analysis.distance import string_distance
|
||||
return string_distance(lf, columns)
|
||||
|
||||
def display(self, value: Any) -> str:
|
||||
return str(value)[:80] if value else "-"
|
||||
|
||||
|
||||
class EnumType(ColumnType):
|
||||
name = "枚举"
|
||||
polars_dtype = pl.Utf8
|
||||
|
||||
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
||||
# 0-1 match distance: compare to mode
|
||||
for col in columns:
|
||||
mode_val = lf.select(pl.col(col).mode().first()).collect(streaming=True).item()
|
||||
lf = lf.with_columns(
|
||||
(pl.col(col) != mode_val).cast(pl.Float64).alias(f'_enum_dist_{col}')
|
||||
)
|
||||
return lf
|
||||
|
||||
def display(self, value: Any) -> str:
|
||||
return str(value) if value else "-"
|
||||
|
||||
|
||||
class FloatType(ColumnType):
|
||||
name = "浮点"
|
||||
polars_dtype = pl.Float64
|
||||
|
||||
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
||||
# Normalised euclidean distance — handled in preprocessing pipeline
|
||||
for col in columns:
|
||||
lf = lf.with_columns(
|
||||
pl.col(col).cast(pl.Float64).alias(f'_raw_{col}')
|
||||
)
|
||||
return lf
|
||||
|
||||
def display(self, value: Any) -> str:
|
||||
try:
|
||||
return f"{float(value):.2f}"
|
||||
except (ValueError, TypeError):
|
||||
return str(value)
|
||||
|
||||
|
||||
class IntType(ColumnType):
|
||||
name = "整数"
|
||||
polars_dtype = pl.Int64
|
||||
|
||||
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
||||
# Normalised manhattan distance — handled in preprocessing pipeline
|
||||
for col in columns:
|
||||
lf = lf.with_columns(
|
||||
pl.col(col).cast(pl.Int64).alias(f'_raw_{col}')
|
||||
)
|
||||
return lf
|
||||
|
||||
def display(self, value: Any) -> str:
|
||||
try:
|
||||
return str(int(value))
|
||||
except (ValueError, TypeError):
|
||||
return str(value)
|
||||
|
||||
|
||||
class BoolType(ColumnType):
|
||||
name = "布尔"
|
||||
polars_dtype = pl.Boolean
|
||||
|
||||
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
||||
# 0-1 match distance
|
||||
for col in columns:
|
||||
lf = lf.with_columns(
|
||||
pl.col(col).cast(pl.Float64).alias(f'_bool_dist_{col}')
|
||||
)
|
||||
return lf
|
||||
|
||||
def display(self, value: Any) -> str:
|
||||
if value is None:
|
||||
return "-"
|
||||
return "是" if str(value).lower() in ('true', '1', '+', 't') else "否"
|
||||
|
||||
|
||||
class BytesType(ColumnType):
|
||||
name = "字节"
|
||||
polars_dtype = pl.Utf8
|
||||
|
||||
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
||||
from analysis.distance import bytes_distance
|
||||
return bytes_distance(lf, columns)
|
||||
|
||||
def display(self, value: Any) -> str:
|
||||
s = str(value).strip() if value else ""
|
||||
if len(s) > 20:
|
||||
return s[:20] + "..."
|
||||
return s or "-"
|
||||
|
||||
|
||||
class TimestampType(ColumnType):
|
||||
name = "时间戳"
|
||||
polars_dtype = pl.Float64
|
||||
|
||||
def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
|
||||
from analysis.distance import timestamp_fft_distance
|
||||
return timestamp_fft_distance(lf, columns)
|
||||
|
||||
def display(self, value: Any) -> str:
|
||||
try:
|
||||
t = float(value)
|
||||
if t < 1e8 or t > 2e12:
|
||||
return str(value)
|
||||
from datetime import datetime
|
||||
return datetime.fromtimestamp(t).strftime("%Y-%m-%d %H:%M")
|
||||
except (ValueError, TypeError, OSError):
|
||||
return str(value) if value else "-"
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Registry — single source of truth
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
TYPE_REGISTRY: dict[str, ColumnType] = {
|
||||
"枚举": EnumType(),
|
||||
"字节": BytesType(),
|
||||
"浮点": FloatType(),
|
||||
"整数": IntType(),
|
||||
"布尔": BoolType(),
|
||||
"字符串": StringType(),
|
||||
"IPv4": IPv4Type(),
|
||||
"时间戳": TimestampType(),
|
||||
}
|
||||
|
||||
# Alias map for config.yaml type strings
|
||||
TYPE_ALIASES: dict[str, str] = {
|
||||
'enum': '枚举', '枚举': '枚举',
|
||||
'bytes': '字节', '字节': '字节', 'hex': '字节',
|
||||
'float': '浮点', '浮点': '浮点',
|
||||
'int': '整数', '整数': '整数',
|
||||
'bool': '布尔', '布尔': '布尔', 'boolean': '布尔',
|
||||
'string': '字符串', '字符串': '字符串', 'str': '字符串',
|
||||
'ipv4': 'IPv4', 'IPv4': 'IPv4',
|
||||
'timestamp': '时间戳', '时间戳': '时间戳',
|
||||
}
|
||||
|
||||
|
||||
def get_type(name: str) -> ColumnType:
|
||||
"""Resolve a type name (Chinese or English alias) to a ColumnType."""
|
||||
resolved = TYPE_ALIASES.get(name.strip().lower() if name.isascii() else name.strip())
|
||||
if resolved is None:
|
||||
return StringType() # fallback
|
||||
return TYPE_REGISTRY[resolved]
|
||||
|
||||
|
||||
def classify_column(name: str, series: pl.Series, config_type: str | None = None) -> ColumnType:
|
||||
"""Classify a single column to a ColumnType.
|
||||
|
||||
Priority: config override → TlsDB lookup → name heuristics → value sampling.
|
||||
"""
|
||||
from analysis.tls_ref import get_tls_ref_dict
|
||||
import re
|
||||
|
||||
# 1. Config override
|
||||
if config_type is not None:
|
||||
t = get_type(config_type)
|
||||
if t.name != "字符串" or config_type.lower().strip() == 'string':
|
||||
return t
|
||||
|
||||
# 2. TlsDB lookup
|
||||
tls_ref = get_tls_ref_dict()
|
||||
if name in tls_ref:
|
||||
t = get_type(tls_ref[name]['data_type'])
|
||||
return t
|
||||
|
||||
# 3. Name heuristics
|
||||
_patterns: list[tuple[re.Pattern, str]] = [
|
||||
(re.compile(r'^(?:src_ip|dst_ip|:ips|:ipd|server.ip|client.ip)$', re.I), 'IPv4'),
|
||||
(re.compile(r'^(?::prs|:prd|scnt|dcnt|1ipp|4dbn|4srs|eiph|name|tabl|0ver|cnrs|isrs)$', re.I), '枚举'),
|
||||
(re.compile(r'^(?:0cph|0crv|cipher|ecdhe)', re.I), '字节'),
|
||||
(re.compile(r'^(?:4dur|8ses|2tmo|0rnt)$', re.I), '浮点'),
|
||||
(re.compile(r'^(?:4ksz|8ack|8ppk|8pak|8seq|8did|8byt|row)$', re.I), '整数'),
|
||||
(re.compile(r'^(?:timestamp|time|8dbd|created|updated)', re.I), '时间戳'),
|
||||
(re.compile(r'\.(ispn|orgn|city|anon|doma|latd|lond)$', re.I), '字符串'),
|
||||
(re.compile(r'^(?:cnam|snam|0rnd|source.node|crcc|orga|orgu|@)', re.I), '字符串'),
|
||||
]
|
||||
for pat, type_name in _patterns:
|
||||
if pat.search(name):
|
||||
return TYPE_REGISTRY[type_name]
|
||||
|
||||
# 4. Value-based sampling
|
||||
non_null = series.drop_nulls()
|
||||
if len(non_null) == 0:
|
||||
return TYPE_REGISTRY["字符串"]
|
||||
|
||||
strs = [str(v).strip().lower() for v in non_null.head(100)]
|
||||
ip4 = re.compile(r'^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$')
|
||||
hex_p = re.compile(r'^[0-9a-f]{2}(?:\s[0-9a-f]{2})*$')
|
||||
|
||||
if all(ip4.match(v) for v in strs):
|
||||
return TYPE_REGISTRY["IPv4"]
|
||||
if all(hex_p.match(v) for v in strs):
|
||||
return TYPE_REGISTRY["字节"]
|
||||
if series.dtype.is_integer():
|
||||
return TYPE_REGISTRY["整数"]
|
||||
if series.dtype == pl.Float64:
|
||||
return TYPE_REGISTRY["浮点"]
|
||||
if series.dtype == pl.Boolean:
|
||||
return TYPE_REGISTRY["布尔"]
|
||||
if non_null.n_unique() < 50:
|
||||
return TYPE_REGISTRY["枚举"]
|
||||
|
||||
return TYPE_REGISTRY["字符串"]
|
||||
|
||||
|
||||
def classify_schema(
|
||||
lf: pl.LazyFrame,
|
||||
config_overrides: dict[str, str] | None = None,
|
||||
) -> dict[str, ColumnType]:
|
||||
"""Classify all columns in a LazyFrame."""
|
||||
config_overrides = config_overrides or {}
|
||||
df_sample = lf.collect(streaming=True)
|
||||
result = {}
|
||||
for col_name in df_sample.columns:
|
||||
config_type = config_overrides.get(col_name)
|
||||
result[col_name] = classify_column(
|
||||
name=col_name,
|
||||
series=df_sample[col_name],
|
||||
config_type=config_type,
|
||||
)
|
||||
return result
|
||||
+21
-43
@@ -176,50 +176,28 @@ def _background_process(run_id, upload_dir):
|
||||
pl.col(lat_lon_col).map_batches(_coerce_to_float_fn, return_dtype=pl.Float64).alias(lat_lon_col)
|
||||
)
|
||||
|
||||
# ── Restore numeric types for Utf8 columns coerced during CSV load ──
|
||||
# All columns were forced to Utf8 for safe SQLite storage. Now try to
|
||||
# recover the original numeric types (Int64 / Float64) from the data.
|
||||
numeric_types = (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
|
||||
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
|
||||
pl.Float32, pl.Float64)
|
||||
# Get current schema to know which columns are still Utf8
|
||||
live_schema = lf.collect_schema()
|
||||
live_names = live_schema.names()
|
||||
live_dtypes = list(live_schema.dtypes())
|
||||
utf8_cols = [
|
||||
live_names[i] for i, dt in enumerate(live_dtypes)
|
||||
if dt == pl.Utf8 and not live_names[i].startswith('_')
|
||||
]
|
||||
if utf8_cols:
|
||||
# Infer types from a small sample to decide Float64 vs Int64
|
||||
sample_df = lf.select(utf8_cols).head(1000).collect(streaming=True)
|
||||
type_map: dict[str, pl.DataType] = {}
|
||||
for col in utf8_cols:
|
||||
try:
|
||||
# Try Float64 first (most general numeric type)
|
||||
casted = sample_df[col].cast(pl.Float64, strict=False)
|
||||
non_null = casted.is_not_null().sum()
|
||||
if non_null > len(casted) * 0.5:
|
||||
# Check if all non-null values are integers → use Int64
|
||||
all_int = True
|
||||
for v in casted.head(100).to_list():
|
||||
if v is not None and v != int(v):
|
||||
all_int = False
|
||||
break
|
||||
type_map[col] = pl.Int64 if all_int else pl.Float64
|
||||
except Exception:
|
||||
logger.error("unknown failed: {}".format(traceback.format_exc()))
|
||||
pass
|
||||
|
||||
if type_map:
|
||||
# Update schema dict to reflect restored types
|
||||
for col, dtype in type_map.items():
|
||||
schema[col] = str(dtype)
|
||||
# Apply casts to the LazyFrame
|
||||
casts = [pl.col(col).cast(dtype, strict=False) for col, dtype in type_map.items()]
|
||||
# ── Restore column types after SQLite storage ──
|
||||
# All columns were stored as Utf8 for SQLite. Recover original types
|
||||
# from the classification system — ONLY the 8 allowed types.
|
||||
from analysis.type_classifier import classify_schema, DataType
|
||||
type_map = classify_schema(lf)
|
||||
# Apply Polars-level type casts based on classification
|
||||
dtypes = {
|
||||
DataType.FLOAT: pl.Float64,
|
||||
DataType.INT: pl.Int64,
|
||||
DataType.BOOL: pl.Boolean,
|
||||
DataType.TIMESTAMP: pl.Float64,
|
||||
}
|
||||
casts = []
|
||||
for col, dt in type_map.items():
|
||||
target = dtypes.get(dt)
|
||||
if target and col in lf.collect_schema().names():
|
||||
casts.append(pl.col(col).cast(target, strict=False))
|
||||
schema[col] = str(target)
|
||||
if casts:
|
||||
lf = lf.with_columns(casts)
|
||||
logger.info('[BACKGROUND] Restored numeric types for %d columns: %s',
|
||||
len(type_map), list(type_map.keys())[:10])
|
||||
restored = [(col, dt.value) for col, dt in type_map.items() if dt in dtypes][:15]
|
||||
logger.info('[BACKGROUND] Type restoration: %s', restored)
|
||||
|
||||
# ── SVD dimensionality reduction ──
|
||||
svd_components = 0
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
"""Rewrite type_classifier.py as thin re-export wrapper."""
|
||||
content = '''"""Re-export from new types module for backward compatibility."""
|
||||
from analysis.types import ColumnType, TYPE_REGISTRY, TYPE_ALIASES, get_type, classify_column, classify_schema
|
||||
TLS_HEX_MAP = {"0303": "TLSv1.2", "0304": "TLSv1.3", "0302": "TLSv1.1", "0301": "TLSv1.0"}
|
||||
'''
|
||||
with open('analysis/type_classifier.py', 'w', encoding='utf-8') as f:
|
||||
f.write(content)
|
||||
print('Rewritten OK')
|
||||
@@ -78,3 +78,12 @@ Restored 1 datasets from disk
|
||||
[24/Jul/2026 12:40:07] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:40:37] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:41:07] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:41:37] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:42:07] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:42:37] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:43:07] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:43:37] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:44:07] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:44:37] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:45:07] "HEAD / HTTP/1.1" 200 0
|
||||
[24/Jul/2026 12:45:37] "HEAD / HTTP/1.1" 200 0
|
||||
|
||||
Reference in New Issue
Block a user