HJQ/tianxuan
2
1
Fork 0
Code Issues Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Compare commits

base: HJQ/tianxuan:c54917791ecd4387384d6b7458bfbb44eea99bb4
Branches Tags
HJQ/tianxuan:master HJQ/tianxuan:beta HJQ/tianxuan:beta-clean HJQ/tianxuan:feature/simple-analysis-v2 HJQ/tianxuan:feature/simple-analysis
...
compare: HJQ/tianxuan:9f06e40a4102716e99ecee918778cfedf67ad1c7
Branches Tags
HJQ/tianxuan:master HJQ/tianxuan:beta HJQ/tianxuan:beta-clean HJQ/tianxuan:feature/simple-analysis-v2 HJQ/tianxuan:feature/simple-analysis

2 Commits
c54917791e ... 9f06e40a41

Author SHA1 Message Date
PM-pinou 9f06e40a41 refactor: clean views/clustering.py and split data_loader.py into package 
Task A: Replace _run_clustering_pipeline (533L->228L) — remove dead
docstring code, delegate to analysis.services.clustering.

Task B: Split analysis/data_loader.py (746L) into package:
  analysis/data_loader/
    __init__.py  — re-exports all public symbols
    _csv.py      — load_csv_directory, BOM/encoding detection
    _schema.py   — schema merging and validation
    _clean.py    — numeric coercion helper
    _sqlite.py   — save_to_db, load_from_db, load_from_db_lazy, etc.
  Keep old data_loader.py as thin re-export wrapper.

Verified: manage.py check (0 issues), run_pipeline (3 clusters, OK)
 2026-07-24 13:52:07 +08:00
PM-pinou 6cbbc24f14 refactor: split distance.py (1114 lines) into analysis/distance/ package 
Break monolithic distance.py into 8 focused sub-modules:
- _geo.py: IP distance functions (subnet mask, haversine, geo)
- _text.py: String distance (Levenshtein + enum) + bool distance
- _numeric.py: Bytes/hex popcount + timestamp FFT distance
- _normalize.py: StandardScaler normalization
- _svd.py: svd_denoise (TruncatedSVD, returns noise_profile)
- _umap.py: umap_reduce (3D→2D fallback)
- _cluster_svd.py: cluster_svd_extract (per-cluster SVD)

DELETE: run_preprocessing_pipeline (dead code, replaced by ColumnProcessor)
All public APIs re-exported via __init__.py. Backward compatible.
 2026-07-24 13:50:44 +08:00
16 changed files with 1829 additions and 2164 deletions
Expand all files Collapse all files
analysis/data_loader.py
+14 -742
View File
@@ -1,746 +1,18 @@
"""CSV loading with per-file schema validation, BOM detection, and multi-file merge.
"""Thin re-export wrapper — implementation moved to analysis/data_loader/ package.
Supports multiple encodings and delimiter-separated files. The main entry
point is :func:`load_csv_directory`, which globs files, detects BOM per file,
validates schema consistency, and returns a merged :class:`pl.LazyFrame`.
When both ``analysis/data_loader.py`` and ``analysis/data_loader/`` exist,
Python prefers the package directory, so this file is effectively dead code.
Kept for git history continuity — all real code lives in the package.
If you need to restore the monolithic version, see git history.
"""
import glob
import io
import logging
import os
import sqlite3
from pathlib import Path
from typing import Optional
import polars as pl
import yaml
from analysis.types import classify_schema, TYPE_REGISTRY
logger = logging.getLogger(__name__)
SUPPORTED_ENCODINGS = ['utf-8', 'utf-8-sig', 'utf-16le', 'utf-16be', 'latin-1']
# ── numeric coercion helper ──────────────────────────────────────────────
def _coerce_to_float(series: pl.Series) -> pl.Series:
"""Convert string series to Float64.
Every entry that doesn't parse as a valid float becomes null.
Handles blank strings, standalone '+'/'-', 'N/A', 'null', 'None', etc.
"""
series = series.cast(pl.Utf8).str.strip_chars()
artifacts = ('', '+', '-', '.', '..', 'N/A', 'n/a', 'NA', 'null', 'None', 'NULL', '?', '---', '--', '/')
is_artifact = series.is_in(pl.Series(artifacts).implode())
result = []
for v, bad in zip(series.to_list(), is_artifact.to_list()):
if v is None or bad:
result.append(None)
else:
# Check for "+"/"-" prefix artifacts (e.g., "+abc", "-xyz")
# where the value isn't a genuine numeric string. Values like
# "+5" or "-3.14" are still accepted as valid floats below.
if len(v) > 1 and v[0] in '+-':
try:
float(v[1:]) # validate suffix minus prefix
except ValueError:
result.append(None)
continue
try:
result.append(float(v))
except (ValueError, TypeError):
result.append(None)
return pl.Series(result, dtype=pl.Float64)
# ── BOM detection ───────────────────────────────────────────────────────
def detect_bom(path: str) -> str:
"""Read the first 4 bytes of *path* and return the detected encoding.
Returns ``'utf-8'`` when no BOM is found.
Note: Python 3 ``open()`` fully supports Unicode/Chinese paths on all
platforms (no encoding dance needed).
"""
with open(path, 'rb') as f:
header = f.read(4)
if header[:3] == b'\xef\xbb\xbf':
return 'utf-8-sig'
if header[:2] == b'\xff\xfe':
return 'utf-16le'
if header[:2] == b'\xfe\xff':
return 'utf-16be'
return 'utf-8'
# ── schema utilities ────────────────────────────────────────────────────
def collect_schema(
path: str,
encoding: str = 'utf-8',
delimiter: str = ',',
sample_rows: int = 10000,
) -> dict:
"""Return {column_name: dtype_string} for a single CSV file.
Uses :func:`pl.scan_csv` with ``infer_schema_length=0`` to skip type
inference. All columns are forced to ``Utf8`` for cross-file compatibility,
avoiding type-inference failures when one file's column is ``Int64`` and
another's is ``String``.
"""
lf = pl.scan_csv(
str(path), # str() ensures Unicode/Chinese paths resolve correctly
encoding=encoding,
separator=delimiter,
infer_schema_length=0,
from analysis.data_loader._csv import load_csv_directory
from analysis.data_loader._sqlite import (
load_from_db,
load_from_db_lazy,
save_to_db,
drop_sqlite_table,
)
schema = lf.collect_schema()
result = {}
for name, dtype in zip(schema.names(), schema.dtypes()):
result[name] = 'Utf8' # force Utf8 for cross-file compatibility
return result
def _merge_schema_entries(
schemas: list[dict],
) -> dict:
"""Merge multiple schema dicts, raising on type conflicts.
Returns a single {name: dtype} dict. If a column appears in multiple
files with different dtypes the *last* non-null dtype wins (with a
warning embedded in the result).
"""
merged: dict[str, str] = {}
conflicts: list[str] = []
for s in schemas:
for name, dtype in s.items():
if name in merged and merged[name] != dtype:
conflicts.append(name)
merged[name] = dtype # last wins
return merged, conflicts
def validate_schemas(schemas: list[dict], strict: bool = True) -> None:
"""Compare schemas across files and raise or warn on column mismatch.
Checks that every file has the same set of column names. Type
differences across files are **not** considered an error (Polars can
cast), but differences in *column sets* are.
Args:
schemas: List of ``{column_name: dtype_string}`` dicts, one per file.
strict: If ``True`` (default), raise :class:`ValueError` on any
column mismatch. If ``False``, log a warning instead.
"""
if not schemas:
return
col_sets = [set(s.keys()) for s in schemas]
reference = col_sets[0]
for i, cs in enumerate(col_sets[1:], start=1):
only_in_ref = reference - cs
only_in_current = cs - reference
if only_in_ref or only_in_current:
msg_parts = []
if only_in_ref:
msg_parts.append(f"File #{i} missing columns: {sorted(only_in_ref)}")
if only_in_current:
msg_parts.append(f"File #{i} has extra columns: {sorted(only_in_current)}")
full_msg = f"Schema mismatch: {'; '.join(msg_parts)}"
if strict:
raise ValueError(full_msg)
else:
logger.warning(full_msg)
# ── dtype inference from schema ─────────────────────────────────────────
def detect_schema_dtypes(
lf: pl.LazyFrame,
sample_rows: int = 10000,
) -> dict:
"""Infer Polars dtype strings from a LazyFrame's schema.
The schema is read without collecting the full data (the scan scans
a small sample).
"""
schema = lf.collect_schema()
return {name: str(dtype) for name, dtype in zip(schema.names(), schema.dtypes())}
# ── config loading ──────────────────────────────────────────────────────
def load_config(config_path: Optional[str]) -> dict:
"""Load a YAML config file; returns an empty dict if *config_path* is None."""
if config_path is None:
return {}
path = Path(config_path)
if not path.exists():
raise FileNotFoundError(f"Config file not found: {config_path}")
with open(path, 'r', encoding='utf-8') as f:
return yaml.safe_load(f) or {}
# ── SQLite persistence ──────────────────────────────────────────────────
def _dtype_to_sqlite(dtype: pl.DataType) -> str:
"""Map a Polars dtype to a SQLite column type affinity."""
if dtype in (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64):
return 'INTEGER'
if dtype in (pl.Float32, pl.Float64):
return 'REAL'
if dtype == pl.Boolean:
return 'INTEGER'
return 'TEXT' # String, Date, Datetime, etc. all map to TEXT
def save_to_db(run_id: int, lf: pl.LazyFrame, mode: str = 'replace') -> str:
"""Write a LazyFrame to a SQLite table ``_data_{run_id}``.
Collects the LazyFrame (streaming if possible), creates or replaces the
table, bulk-inserts rows in chunks, and returns the table name.
Args:
run_id: Analysis run ID (used as ``_data_{run_id}``).
lf: LazyFrame to persist.
mode: ``'replace'`` (default) → drop+create+insert.
``'append'`` → skip DDL, only INSERT (table must exist).
This is a *persistence copy* — the original CSV files remain untouched.
"""
from django.conf import settings
table_name = f"_data_{run_id}"
db_path = settings.DATABASES['default']['NAME']
df = lf.collect(streaming=True)
conn = sqlite3.connect(db_path)
try:
if mode == 'replace':
# Drop existing table for retry/idempotency safety
conn.execute(f'DROP TABLE IF EXISTS "{table_name}"')
# Build CREATE TABLE from Polars schema
col_defs = [f'"{name}" {_dtype_to_sqlite(dtype)}'
for name, dtype in zip(df.columns, df.dtypes)]
conn.execute(f'CREATE TABLE "{table_name}" ({", ".join(col_defs)})')
# Bulk insert in chunks
col_names = [f'"{c}"' for c in df.columns]
placeholders = ','.join(['?' for _ in df.columns])
insert_sql = f'INSERT INTO "{table_name}" ({", ".join(col_names)}) VALUES ({placeholders})'
total_rows = len(df)
chunk_size = 10000
for i in range(0, total_rows, chunk_size):
chunk = df.slice(i, chunk_size)
conn.executemany(insert_sql, chunk.rows())
conn.commit()
finally:
conn.close()
logger.info('[SAVE_TO_DB] table=%s rows=%d mode=%s', table_name, total_rows, mode)
return table_name
_STR_TO_DTYPE = {
'Float64': pl.Float64, 'Float32': pl.Float32,
'Int64': pl.Int64, 'Int32': pl.Int32, 'Int16': pl.Int16, 'Int8': pl.Int8,
'UInt64': pl.UInt64, 'UInt32': pl.UInt32, 'UInt16': pl.UInt16, 'UInt8': pl.UInt8,
'Utf8': pl.Utf8, 'String': pl.Utf8, 'Categorical': pl.Categorical,
'Boolean': pl.Boolean, 'Bool': pl.Boolean,
'Date': pl.Date, 'Datetime': pl.Datetime, 'Time': pl.Time,
}
def _str_to_polars_dtype(s: str) -> pl.DataType:
"""Convert a Polars dtype string (e.g. ``'Float64'``, ``'Int32'``) to a ``pl.DataType``.
Strips trailing metadata like ``Datetime(time_unit='us', time_zone=None)``.
"""
base = s.split('(')[0].strip()
return _STR_TO_DTYPE.get(base, pl.Utf8)
def _rows_to_csv_df(rows: list[sqlite3.Row]) -> pl.DataFrame:
"""Convert a batch of sqlite3.Row objects to a DataFrame via CSV fallback.
Loads all columns as Utf8 to handle mixed-type SQLite columns.
"""
import csv as _csv_mod
buf = io.StringIO()
writer = _csv_mod.writer(buf)
writer.writerow(rows[0].keys())
for r in rows:
writer.writerow(str(v) if v is not None else '' for v in r)
buf.seek(0)
return pl.read_csv(buf, infer_schema_length=0)
def load_from_db(
table_name: str,
schema_overrides: dict[str, pl.DataType | str] | None = None,
) -> Optional[pl.LazyFrame]:
"""Load a SQLite table into a LazyFrame.
Args:
table_name: SQLite table name.
schema_overrides: Optional dict of ``{column_name: dtype}`` to cast
columns after loading. Values can be ``pl.DataType`` instances
(e.g. ``pl.Float64``) or dtype strings (e.g. ``'Float64'``).
Use when SQLite ``TEXT`` columns need to be restored to
``Float64`` / ``Int64`` etc.
Returns ``None`` if the table does not exist.
"""
from django.conf import settings
db_path = settings.DATABASES['default']['NAME']
conn = sqlite3.connect(db_path)
try:
# Check table exists
cursor = conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name=?",
(table_name,),
)
if cursor.fetchone() is None:
return None
conn.row_factory = sqlite3.Row
cursor = conn.execute(f'SELECT * FROM "{table_name}"')
# Peek first row to detect if table is empty
first_row = cursor.fetchone()
if first_row is None:
return pl.DataFrame().lazy()
batch_rows = [first_row]
dfs = []
_use_csv_fallback = False
while True:
chunk = cursor.fetchmany(10000)
if not chunk:
# Flush remaining rows
if batch_rows:
if _use_csv_fallback:
dfs.append(_rows_to_csv_df(batch_rows))
else:
try:
dfs.append(pl.from_dicts([dict(r) for r in batch_rows]))
except Exception:
_use_csv_fallback = True
dfs.append(_rows_to_csv_df(batch_rows))
break
batch_rows.extend(chunk)
if len(batch_rows) >= 10000:
if _use_csv_fallback:
dfs.append(_rows_to_csv_df(batch_rows))
else:
try:
dfs.append(pl.from_dicts([dict(r) for r in batch_rows]))
except Exception:
_use_csv_fallback = True
# Re-process current batch with CSV fallback
dfs.append(_rows_to_csv_df(batch_rows))
batch_rows = []
# Concat all batches
if len(dfs) == 1:
df = dfs[0]
else:
df = pl.concat(dfs, how='vertical')
# Apply type overrides to restore original column types
if schema_overrides:
casts = {}
for col, dtype in schema_overrides.items():
if col not in df.columns:
continue
dtype_obj = dtype if isinstance(dtype, pl.DataType) else _str_to_polars_dtype(dtype)
if df[col].dtype != dtype_obj:
try:
casts[col] = pl.col(col).cast(dtype_obj)
except Exception:
# If a cast fails (e.g. "A" → Int64), keep the column as-is
pass
if casts:
df = df.with_columns(list(casts.values()))
return df.lazy()
finally:
conn.close()
def load_from_db_lazy(
table_name: str,
schema_overrides: dict[str, pl.DataType | str] | None = None,
) -> Optional[pl.LazyFrame]:
"""Load a SQLite table into a LazyFrame using Polars native DB reader.
Uses ``pl.read_database`` for efficient bulk reading — an order of
magnitude faster than the manual batch-processing path in
:func:`load_from_db`. Returns a :class:`pl.LazyFrame` so downstream
operations remain lazy (query-plan, not materialised).
Args:
table_name: SQLite table name.
schema_overrides: Optional ``{column_name: dtype}`` dict to cast
columns after loading. Values can be ``pl.DataType`` instances
(e.g. ``pl.Float64``) or dtype strings (e.g. ``'Float64'``).
Returns:
``None`` if the table does not exist or is empty.
"""
from django.conf import settings
db_path = settings.DATABASES['default']['NAME']
conn = sqlite3.connect(db_path)
try:
# Check table existence (fast metadata query)
cursor = conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name=?",
(table_name,),
)
if cursor.fetchone() is None:
return None
# Bulk-read with Polars native engine — leverages the DBAPI2
# connection directly for efficient C-level data transfer.
query = f'SELECT * FROM "{table_name}"'
df = pl.read_database(connection=conn, query=query)
finally:
conn.close()
if df.height == 0:
return None
# Apply type overrides to restore original column types
# (SQLite stores everything as TEXT/INTEGER/REAL/BLOB;
# we need to cast back to Float64 / Int64 / etc.)
if schema_overrides:
casts = {}
for col, dtype in schema_overrides.items():
if col not in df.columns:
continue
dtype_obj = (
dtype if isinstance(dtype, pl.DataType)
else _str_to_polars_dtype(dtype)
)
if df[col].dtype != dtype_obj:
try:
casts[col] = pl.col(col).cast(dtype_obj)
except Exception:
pass # uncastable column — keep as-is
if casts:
df = df.with_columns(list(casts.values()))
rows = df.height
cols = df.width
logger.info(
'[LOAD_FROM_DB_LAZY] table=%s rows=%d cols=%d',
table_name, rows, cols,
)
return df.lazy()
def drop_sqlite_table(table_name: str) -> None:
"""Drop a SQLite data table silently if it exists."""
from django.conf import settings
db_path = settings.DATABASES['default']['NAME']
conn = sqlite3.connect(db_path)
try:
conn.execute(f'DROP TABLE IF EXISTS "{table_name}"')
conn.commit()
finally:
conn.close()
# ── helpers ─────────────────────────────────────────────────────────────
def _file_count(paths: list[str]) -> int:
return len(paths)
def _resolve_encoding(encoding: str, file_encoding: str) -> str:
"""BOM-detected encoding overrides user-supplied encoding.
Polars >= 1.0 requires ``'utf8'`` (not ``'utf-8'``).
"""
resolved = file_encoding if file_encoding != 'utf-8' else encoding
return resolved.replace('-', '') if resolved in ('utf-8', 'utf-8-sig') else resolved
# ── main entry ──────────────────────────────────────────────────────────
def load_csv_directory(
glob_pattern: str,
encoding: str = 'utf-8',
delimiter: str = ',',
config_path: Optional[str] = None,
schema_override: Optional[dict] = None,
sample_rows: int = 10000,
schema_strict: bool = False,
recursive: bool = False,
head: Optional[int] = None,
) -> tuple[pl.LazyFrame, dict, int, int, float]:
"""Load, validate, and merge CSV files matching *glob_pattern*.
Args:
glob_pattern: Glob pattern for CSV files (e.g. ``data/*.csv``).
encoding: Fallback encoding when no BOM is detected.
delimiter: Column delimiter character.
config_path: Optional path to a YAML config (merged into metadata).
schema_override: If given, force these {column: dtype} overrides.
sample_rows: Rows to scan for dtype inference per file.
schema_strict: When True, raise ValueError on column mismatch across
files. When False (default), log warnings and merge with
``diagonal_relaxed`` (missing columns get null).
recursive: When True, use ``glob.glob(..., recursive=True)`` to
support ``**/*.csv`` patterns.
head: Optional row limit for preview scenarios. When set, applies
``lf.head(head)`` before returning.
Returns:
A tuple of ``(lazyframe, schema_dict, total_row_count, file_count,
memory_estimate_mb)``.
Raises:
FileNotFoundError: If no files match the glob pattern.
ValueError: If schemas are inconsistent across files and
*schema_strict* is True.
"""
# Normalise path separators for cross-platform Unicode support
glob_pattern = os.path.normpath(glob_pattern)
# Resolve file list
paths = sorted(glob.glob(glob_pattern, recursive=recursive))
if not paths and not recursive:
# Fallback: try recursive pattern
paths = sorted(glob.glob(glob_pattern, recursive=True))
if not paths:
raise FileNotFoundError(f"No files match glob pattern: {glob_pattern}")
# Load config (user-provided or fall back to project default)
config = load_config(config_path)
if not config:
default_path = Path(__file__).parent.parent / 'config' / 'config.yaml'
if default_path.exists():
config = load_config(str(default_path))
# Collect per-file schema
schemas: list[dict] = []
lazyframes: list[pl.LazyFrame] = []
for path in paths:
file_encoding = detect_bom(path)
actual_encoding = _resolve_encoding(encoding, file_encoding)
lf = pl.scan_csv(
str(path), # str() ensures Unicode/Chinese paths resolve correctly
encoding=actual_encoding,
separator=delimiter,
infer_schema_length=max(sample_rows, 1000),
)
schema = detect_schema_dtypes(lf, sample_rows=sample_rows)
# Apply schema override for this file
if schema_override:
for col, dtype in schema_override.items():
if col in schema:
schema[col] = dtype
schemas.append(schema)
lazyframes.append(lf)
if schema_strict:
# Raise on column mismatch
validate_schemas(schemas)
else:
# Lenient mode: log warnings for column mismatches
col_sets = [set(s.keys()) for s in schemas]
reference = col_sets[0]
for i, cs in enumerate(col_sets[1:], start=1):
only_in_ref = reference - cs
only_in_current = cs - reference
if only_in_ref:
logger.warning(
"File #%d missing columns: %s (will be null)",
i + 1, sorted(only_in_ref),
)
if only_in_current:
logger.warning(
"File #%d has extra columns: %s (will be null in other files)",
i + 1, sorted(only_in_current),
)
# Merge schemas (unify column ordering — first file wins order)
merged_schema, _ = _merge_schema_entries(schemas)
# Concatenate all LazyFrames (union vertically)
merged_lf: pl.LazyFrame = pl.concat(lazyframes, how='diagonal_relaxed')
# Log schema and size estimate
total_row_count_est = len(paths) * (lazyframes[0].collect(streaming=True).height if lazyframes else 0)
logger.info(f'[LOAD] files={len(paths)} total_rows_est={total_row_count_est}')
for col_name, dtype in sorted(merged_schema.items()):
logger.info(f'[LOAD_SCHEMA] col={col_name} dtype={dtype}')
# Apply schema override on merged frame (cast columns)
if schema_override:
casts = {}
for col, dtype_str in schema_override.items():
if col in merged_schema:
try:
polars_dtype = getattr(pl, dtype_str.split('(')[0], None)
if polars_dtype:
casts[col] = pl.col(col).cast(polars_dtype)
except Exception:
pass # skip uncastable overrides silently
if casts:
merged_lf = merged_lf.with_columns(list(casts.values()))
# ── Cleaning: drop outdated columns ──────────────────────────────────
geoip_cfg = config.get('data', {}).get('geoip', {})
drop_cols = geoip_cfg.get(
'drop_columns',
[], # Empty by default — lat/lon columns are now cleaned, not dropped
)
existing_drop = [c for c in drop_cols if c in merged_schema]
if existing_drop:
merged_lf = merged_lf.drop(existing_drop)
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 normalisation ─────────────────────────────────────
try:
type_map = classify_schema(
merged_lf.head(1000),
config_overrides=config.get('columns', {}),
)
except Exception:
type_map = {}
bool_cols = [
c for c, t in type_map.items()
if t.name == "布尔" and c in merged_schema
]
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_cols:
if schema_dtypes.get(c) not in (pl.Utf8, pl.String):
continue
c_str = pl.col(c).cast(pl.Utf8).str.strip_chars()
cleaned = (
pl.when(c_str == "")
.then(pl.lit("False"))
.when(c_str.str.to_lowercase() == "+")
.then(pl.lit("True"))
.when(c_str.str.to_lowercase() == "-")
.then(pl.lit("False"))
.otherwise(c_str)
)
bool_exprs.append(cleaned.alias(c))
cleaned_names.append(c)
if bool_exprs:
merged_lf = merged_lf.with_columns(bool_exprs)
for c in cleaned_names:
merged_schema[c] = 'Utf8'
logger.info('[CLEAN] BOOL columns normalised: %s', cleaned_names)
# ── 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(),
merged_lf.collect_schema().dtypes(),
))
generic_numeric_cols = []
for c, t in type_map.items():
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
cl = c.lower().replace('-', '_').replace(' ', '_')
if any(p in cl for p in _SKIP_NAME_PATTERNS):
continue
generic_numeric_cols.append(c)
if generic_numeric_cols:
num_exprs = []
for c in generic_numeric_cols:
num_exprs.append(
pl.col(c).map_batches(_coerce_to_float, return_dtype=pl.Float64).alias(c)
)
if num_exprs:
merged_lf = merged_lf.with_columns(num_exprs)
for c in generic_numeric_cols:
merged_schema[c] = 'Float64'
logger.info('[CLEAN] Generic numeric coercion applied: %s', generic_numeric_cols)
# ── Cleaning: BYTES (hex) preprocessing ─────────────────────────────
hex_cols = [
c for c, t in type_map.items()
if t.name == "字节" and c in merged_schema
]
if hex_cols:
hex_exprs = []
for c in hex_cols:
hex_exprs.append(
pl.col(c).str.split(" ").list.len().alias(f"{c}_hex_pairs_count")
)
merged_lf = merged_lf.with_columns(hex_exprs)
for c in hex_cols:
merged_schema[f"{c}_hex_pairs_count"] = 'UInt32'
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.
# We estimate from the first file's scanned schema info.
first_rows = lazyframes[0].collect(streaming=True).height if lazyframes else 0
total_row_count = first_rows * len(paths)
# Memory estimate
bytes_per_row = 0
for col_name, dtype_str in merged_schema.items():
numeric_part = dtype_str.split('(')[0].lower()
size_map = {
'int8': 1, 'int16': 2, 'int32': 4, 'int64': 8,
'uint8': 1, 'uint16': 2, 'uint32': 4, 'uint64': 8,
'float32': 4, 'float64': 8,
'boolean': 1, 'bool': 1,
'utf8': 50, 'string': 50, 'str': 50,
'categorical': 8, 'cat': 8,
'date': 4, 'datetime': 8, 'time': 8, 'duration': 8,
}
bytes_per_row += size_map.get(numeric_part, 8)
memory_mb = (bytes_per_row * total_row_count) / (1024 * 1024)
# Apply head limit for preview (keeps LazyFrame lazy — rows are limited on collect)
if head is not None:
merged_lf = merged_lf.head(head)
total_row_count = min(total_row_count, head)
return merged_lf, merged_schema, total_row_count, len(paths), memory_mb
from analysis.data_loader._clean import _coerce_to_float
from analysis.data_loader._sqlite import _dtype_to_sqlite
analysis/data_loader/__init__.py
+30
View File
@@ -0,0 +1,30 @@
"""CSV data loading, schema handling, cleaning, and SQLite persistence.
Package layout:
_csv.py — CSV scanning, BOM/encoding detection, :func:`load_csv_directory`
_schema.py — Schema merging and cross-file validation
_clean.py — Data cleaning helpers (:func:`_coerce_to_float`)
_sqlite.py — SQLite persistence (save/load/drop tables)
"""
from ._csv import (
SUPPORTED_ENCODINGS,
detect_bom,
collect_schema,
detect_schema_dtypes,
load_config,
_file_count,
_resolve_encoding,
load_csv_directory,
)
from ._schema import _merge_schema_entries, validate_schemas
from ._clean import _coerce_to_float
from ._sqlite import (
_dtype_to_sqlite,
_str_to_polars_dtype,
_rows_to_csv_df,
save_to_db,
load_from_db,
load_from_db_lazy,
drop_sqlite_table,
)
analysis/data_loader/_clean.py
+32
View File
@@ -0,0 +1,32 @@
"""Data cleaning helpers: numeric coercion for string columns."""
import polars as pl
def _coerce_to_float(series: pl.Series) -> pl.Series:
"""Convert string series to Float64.
Every entry that doesn't parse as a valid float becomes null.
Handles blank strings, standalone '+'/'-', 'N/A', 'null', 'None', etc.
"""
series = series.cast(pl.Utf8).str.strip_chars()
artifacts = ('', '+', '-', '.', '..', 'N/A', 'n/a', 'NA', 'null', 'None', 'NULL', '?', '---', '--', '/')
is_artifact = series.is_in(pl.Series(artifacts).implode())
result = []
for v, bad in zip(series.to_list(), is_artifact.to_list()):
if v is None or bad:
result.append(None)
else:
# Check for "+"/"-" prefix artifacts (e.g., "+abc", "-xyz")
# where the value isn't a genuine numeric string. Values like
# "+5" or "-3.14" are still accepted as valid floats below.
if len(v) > 1 and v[0] in '+-':
try:
float(v[1:]) # validate suffix minus prefix
except ValueError:
result.append(None)
continue
try:
result.append(float(v))
except (ValueError, TypeError):
result.append(None)
return pl.Series(result, dtype=pl.Float64)
analysis/data_loader/_csv.py
+375
View File
@@ -0,0 +1,375 @@
"""CSV loading: glob, BOM detection, encoding, schema collection, multi-file merge.
Main entry: :func:`load_csv_directory`.
"""
import glob
import logging
import os
from pathlib import Path
from typing import Optional
import polars as pl
import yaml
from analysis.types import classify_schema, TYPE_REGISTRY
from ._clean import _coerce_to_float
from ._schema import _merge_schema_entries, validate_schemas
logger = logging.getLogger(__name__)
SUPPORTED_ENCODINGS = ['utf-8', 'utf-8-sig', 'utf-16le', 'utf-16be', 'latin-1']
# ── BOM detection ───────────────────────────────────────────────────────
def detect_bom(path: str) -> str:
"""Read the first 4 bytes of *path* and return the detected encoding.
Returns ``'utf-8'`` when no BOM is found.
Note: Python 3 ``open()`` fully supports Unicode/Chinese paths on all
platforms (no encoding dance needed).
"""
with open(path, 'rb') as f:
header = f.read(4)
if header[:3] == b'\xef\xbb\xbf':
return 'utf-8-sig'
if header[:2] == b'\xff\xfe':
return 'utf-16le'
if header[:2] == b'\xfe\xff':
return 'utf-16be'
return 'utf-8'
# ── schema utilities ────────────────────────────────────────────────────
def collect_schema(
path: str,
encoding: str = 'utf-8',
delimiter: str = ',',
sample_rows: int = 10000,
) -> dict:
"""Return {column_name: dtype_string} for a single CSV file.
Uses :func:`pl.scan_csv` with ``infer_schema_length=0`` to skip type
inference. All columns are forced to ``Utf8`` for cross-file compatibility,
avoiding type-inference failures when one file's column is ``Int64`` and
another's is ``String``.
"""
lf = pl.scan_csv(
str(path), # str() ensures Unicode/Chinese paths resolve correctly
encoding=encoding,
separator=delimiter,
infer_schema_length=0,
)
schema = lf.collect_schema()
result = {}
for name, dtype in zip(schema.names(), schema.dtypes()):
result[name] = 'Utf8' # force Utf8 for cross-file compatibility
return result
# ── dtype inference from schema ─────────────────────────────────────────
def detect_schema_dtypes(
lf: pl.LazyFrame,
sample_rows: int = 10000,
) -> dict:
"""Infer Polars dtype strings from a LazyFrame's schema.
The schema is read without collecting the full data (the scan scans
a small sample).
"""
schema = lf.collect_schema()
return {name: str(dtype) for name, dtype in zip(schema.names(), schema.dtypes())}
# ── config loading ──────────────────────────────────────────────────────
def load_config(config_path: Optional[str]) -> dict:
"""Load a YAML config file; returns an empty dict if *config_path* is None."""
if config_path is None:
return {}
path = Path(config_path)
if not path.exists():
raise FileNotFoundError(f"Config file not found: {config_path}")
with open(path, 'r', encoding='utf-8') as f:
return yaml.safe_load(f) or {}
# ── helpers ─────────────────────────────────────────────────────────────
def _file_count(paths: list[str]) -> int:
return len(paths)
def _resolve_encoding(encoding: str, file_encoding: str) -> str:
"""BOM-detected encoding overrides user-supplied encoding.
Polars >= 1.0 requires ``'utf8'`` (not ``'utf-8'``).
"""
resolved = file_encoding if file_encoding != 'utf-8' else encoding
return resolved.replace('-', '') if resolved in ('utf-8', 'utf-8-sig') else resolved
# ── main entry ──────────────────────────────────────────────────────────
def load_csv_directory(
glob_pattern: str,
encoding: str = 'utf-8',
delimiter: str = ',',
config_path: Optional[str] = None,
schema_override: Optional[dict] = None,
sample_rows: int = 10000,
schema_strict: bool = False,
recursive: bool = False,
head: Optional[int] = None,
) -> tuple[pl.LazyFrame, dict, int, int, float]:
"""Load, validate, and merge CSV files matching *glob_pattern*.
Args:
glob_pattern: Glob pattern for CSV files (e.g. ``data/*.csv``).
encoding: Fallback encoding when no BOM is detected.
delimiter: Column delimiter character.
config_path: Optional path to a YAML config (merged into metadata).
schema_override: If given, force these {column: dtype} overrides.
sample_rows: Rows to scan for dtype inference per file.
schema_strict: When True, raise ValueError on column mismatch across
files. When False (default), log warnings and merge with
``diagonal_relaxed`` (missing columns get null).
recursive: When True, use ``glob.glob(..., recursive=True)`` to
support ``**/*.csv`` patterns.
head: Optional row limit for preview scenarios. When set, applies
``lf.head(head)`` before returning.
Returns:
A tuple of ``(lazyframe, schema_dict, total_row_count, file_count,
memory_estimate_mb)``.
Raises:
FileNotFoundError: If no files match the glob pattern.
ValueError: If schemas are inconsistent across files and
*schema_strict* is True.
"""
# Normalise path separators for cross-platform Unicode support
glob_pattern = os.path.normpath(glob_pattern)
# Resolve file list
paths = sorted(glob.glob(glob_pattern, recursive=recursive))
if not paths and not recursive:
# Fallback: try recursive pattern
paths = sorted(glob.glob(glob_pattern, recursive=True))
if not paths:
raise FileNotFoundError(f"No files match glob pattern: {glob_pattern}")
# Load config (user-provided or fall back to project default)
config = load_config(config_path)
if not config:
default_path = Path(__file__).parent.parent.parent / 'config' / 'config.yaml'
if default_path.exists():
config = load_config(str(default_path))
# Collect per-file schema
schemas: list[dict] = []
lazyframes: list[pl.LazyFrame] = []
for path in paths:
file_encoding = detect_bom(path)
actual_encoding = _resolve_encoding(encoding, file_encoding)
lf = pl.scan_csv(
str(path), # str() ensures Unicode/Chinese paths resolve correctly
encoding=actual_encoding,
separator=delimiter,
infer_schema_length=max(sample_rows, 1000),
)
schema = detect_schema_dtypes(lf, sample_rows=sample_rows)
# Apply schema override for this file
if schema_override:
for col, dtype in schema_override.items():
if col in schema:
schema[col] = dtype
schemas.append(schema)
lazyframes.append(lf)
if schema_strict:
# Raise on column mismatch
validate_schemas(schemas)
else:
# Lenient mode: log warnings for column mismatches
col_sets = [set(s.keys()) for s in schemas]
reference = col_sets[0]
for i, cs in enumerate(col_sets[1:], start=1):
only_in_ref = reference - cs
only_in_current = cs - reference
if only_in_ref:
logger.warning(
"File #%d missing columns: %s (will be null)",
i + 1, sorted(only_in_ref),
)
if only_in_current:
logger.warning(
"File #%d has extra columns: %s (will be null in other files)",
i + 1, sorted(only_in_current),
)
# Merge schemas (unify column ordering — first file wins order)
merged_schema, _ = _merge_schema_entries(schemas)
# Concatenate all LazyFrames (union vertically)
merged_lf: pl.LazyFrame = pl.concat(lazyframes, how='diagonal_relaxed')
# Log schema and size estimate
total_row_count_est = len(paths) * (lazyframes[0].collect(streaming=True).height if lazyframes else 0)
logger.info(f'[LOAD] files={len(paths)} total_rows_est={total_row_count_est}')
for col_name, dtype in sorted(merged_schema.items()):
logger.info(f'[LOAD_SCHEMA] col={col_name} dtype={dtype}')
# Apply schema override on merged frame (cast columns)
if schema_override:
casts = {}
for col, dtype_str in schema_override.items():
if col in merged_schema:
try:
polars_dtype = getattr(pl, dtype_str.split('(')[0], None)
if polars_dtype:
casts[col] = pl.col(col).cast(polars_dtype)
except Exception:
pass # skip uncastable overrides silently
if casts:
merged_lf = merged_lf.with_columns(list(casts.values()))
# ── Cleaning: drop outdated columns ──────────────────────────────────
geoip_cfg = config.get('data', {}).get('geoip', {})
drop_cols = geoip_cfg.get(
'drop_columns',
[], # Empty by default — lat/lon columns are now cleaned, not dropped
)
existing_drop = [c for c in drop_cols if c in merged_schema]
if existing_drop:
merged_lf = merged_lf.drop(existing_drop)
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 normalisation ─────────────────────────────────────
try:
type_map = classify_schema(
merged_lf.head(1000),
config_overrides=config.get('columns', {}),
)
except Exception:
type_map = {}
bool_cols = [
c for c, t in type_map.items()
if t.name == "布尔" and c in merged_schema
]
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_cols:
if schema_dtypes.get(c) not in (pl.Utf8, pl.String):
continue
c_str = pl.col(c).cast(pl.Utf8).str.strip_chars()
cleaned = (
pl.when(c_str == "")
.then(pl.lit("False"))
.when(c_str.str.to_lowercase() == "+")
.then(pl.lit("True"))
.when(c_str.str.to_lowercase() == "-")
.then(pl.lit("False"))
.otherwise(c_str)
)
bool_exprs.append(cleaned.alias(c))
cleaned_names.append(c)
if bool_exprs:
merged_lf = merged_lf.with_columns(bool_exprs)
for c in cleaned_names:
merged_schema[c] = 'Utf8'
logger.info('[CLEAN] BOOL columns normalised: %s', cleaned_names)
# ── 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(),
merged_lf.collect_schema().dtypes(),
))
generic_numeric_cols = []
for c, t in type_map.items():
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
cl = c.lower().replace('-', '_').replace(' ', '_')
if any(p in cl for p in _SKIP_NAME_PATTERNS):
continue
generic_numeric_cols.append(c)
if generic_numeric_cols:
num_exprs = []
for c in generic_numeric_cols:
num_exprs.append(
pl.col(c).map_batches(_coerce_to_float, return_dtype=pl.Float64).alias(c)
)
if num_exprs:
merged_lf = merged_lf.with_columns(num_exprs)
for c in generic_numeric_cols:
merged_schema[c] = 'Float64'
logger.info('[CLEAN] Generic numeric coercion applied: %s', generic_numeric_cols)
# ── Cleaning: BYTES (hex) preprocessing ─────────────────────────────
hex_cols = [
c for c, t in type_map.items()
if t.name == "字节" and c in merged_schema
]
if hex_cols:
hex_exprs = []
for c in hex_cols:
hex_exprs.append(
pl.col(c).str.split(" ").list.len().alias(f"{c}_hex_pairs_count")
)
merged_lf = merged_lf.with_columns(hex_exprs)
for c in hex_cols:
merged_schema[f"{c}_hex_pairs_count"] = 'UInt32'
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.
# We estimate from the first file's scanned schema info.
first_rows = lazyframes[0].collect(streaming=True).height if lazyframes else 0
total_row_count = first_rows * len(paths)
# Memory estimate
bytes_per_row = 0
for col_name, dtype_str in merged_schema.items():
numeric_part = dtype_str.split('(')[0].lower()
size_map = {
'int8': 1, 'int16': 2, 'int32': 4, 'int64': 8,
'uint8': 1, 'uint16': 2, 'uint32': 4, 'uint64': 8,
'float32': 4, 'float64': 8,
'boolean': 1, 'bool': 1,
'utf8': 50, 'string': 50, 'str': 50,
'categorical': 8, 'cat': 8,
'date': 4, 'datetime': 8, 'time': 8, 'duration': 8,
}
bytes_per_row += size_map.get(numeric_part, 8)
memory_mb = (bytes_per_row * total_row_count) / (1024 * 1024)
# Apply head limit for preview (keeps LazyFrame lazy — rows are limited on collect)
if head is not None:
merged_lf = merged_lf.head(head)
total_row_count = min(total_row_count, head)
return merged_lf, merged_schema, total_row_count, len(paths), memory_mb
analysis/data_loader/_schema.py
+57
View File
@@ -0,0 +1,57 @@
"""Schema detection, validation, and merging across CSV files."""
import logging
logger = logging.getLogger(__name__)
def _merge_schema_entries(
schemas: list[dict],
) -> tuple[dict[str, str], list[str]]:
"""Merge multiple schema dicts, raising on type conflicts.
Returns a single {name: dtype} dict. If a column appears in multiple
files with different dtypes the *last* non-null dtype wins (with a
warning embedded in the result).
"""
merged: dict[str, str] = {}
conflicts: list[str] = []
for s in schemas:
for name, dtype in s.items():
if name in merged and merged[name] != dtype:
conflicts.append(name)
merged[name] = dtype # last wins
return merged, conflicts
def validate_schemas(schemas: list[dict], strict: bool = True) -> None:
"""Compare schemas across files and raise or warn on column mismatch.
Checks that every file has the same set of column names. Type
differences across files are **not** considered an error (Polars can
cast), but differences in *column sets* are.
Args:
schemas: List of ``{column_name: dtype_string}`` dicts, one per file.
strict: If ``True`` (default), raise :class:`ValueError` on any
column mismatch. If ``False``, log a warning instead.
"""
if not schemas:
return
col_sets = [set(s.keys()) for s in schemas]
reference = col_sets[0]
for i, cs in enumerate(col_sets[1:], start=1):
only_in_ref = reference - cs
only_in_current = cs - reference
if only_in_ref or only_in_current:
msg_parts = []
if only_in_ref:
msg_parts.append(f"File #{i} missing columns: {sorted(only_in_ref)}")
if only_in_current:
msg_parts.append(f"File #{i} has extra columns: {sorted(only_in_current)}")
full_msg = f"Schema mismatch: {'; '.join(msg_parts)}"
if strict:
raise ValueError(full_msg)
else:
logger.warning(full_msg)
analysis/data_loader/_sqlite.py
+292
View File
@@ -0,0 +1,292 @@
"""SQLite persistence layer for LazyFrame data."""
import csv as _csv_mod
import io
import logging
import sqlite3
from typing import Optional
import polars as pl
logger = logging.getLogger(__name__)
# ── dtype mapping ─────────────────────────────────────────────────────────
_STR_TO_DTYPE = {
'Float64': pl.Float64, 'Float32': pl.Float32,
'Int64': pl.Int64, 'Int32': pl.Int32, 'Int16': pl.Int16, 'Int8': pl.Int8,
'UInt64': pl.UInt64, 'UInt32': pl.UInt32, 'UInt16': pl.UInt16, 'UInt8': pl.UInt8,
'Utf8': pl.Utf8, 'String': pl.Utf8, 'Categorical': pl.Categorical,
'Boolean': pl.Boolean, 'Bool': pl.Boolean,
'Date': pl.Date, 'Datetime': pl.Datetime, 'Time': pl.Time,
}
def _dtype_to_sqlite(dtype: pl.DataType) -> str:
"""Map a Polars dtype to a SQLite column type affinity."""
if dtype in (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64):
return 'INTEGER'
if dtype in (pl.Float32, pl.Float64):
return 'REAL'
if dtype == pl.Boolean:
return 'INTEGER'
return 'TEXT' # String, Date, Datetime, etc. all map to TEXT
def _str_to_polars_dtype(s: str) -> pl.DataType:
"""Convert a Polars dtype string (e.g. ``'Float64'``, ``'Int32'``) to a ``pl.DataType``.
Strips trailing metadata like ``Datetime(time_unit='us', time_zone=None)``.
"""
base = s.split('(')[0].strip()
return _STR_TO_DTYPE.get(base, pl.Utf8)
def _rows_to_csv_df(rows: list[sqlite3.Row]) -> pl.DataFrame:
"""Convert a batch of sqlite3.Row objects to a DataFrame via CSV fallback.
Loads all columns as Utf8 to handle mixed-type SQLite columns.
"""
buf = io.StringIO()
writer = _csv_mod.writer(buf)
writer.writerow(rows[0].keys())
for r in rows:
writer.writerow(str(v) if v is not None else '' for v in r)
buf.seek(0)
return pl.read_csv(buf, infer_schema_length=0)
# ── public API ───────────────────────────────────────────────────────────
def save_to_db(run_id: int, lf: pl.LazyFrame, mode: str = 'replace') -> str:
"""Write a LazyFrame to a SQLite table ``_data_{run_id}``.
Collects the LazyFrame (streaming if possible), creates or replaces the
table, bulk-inserts rows in chunks, and returns the table name.
Args:
run_id: Analysis run ID (used as ``_data_{run_id}``).
lf: LazyFrame to persist.
mode: ``'replace'`` (default) → drop+create+insert.
``'append'`` → skip DDL, only INSERT (table must exist).
This is a *persistence copy* — the original CSV files remain untouched.
"""
from django.conf import settings
table_name = f"_data_{run_id}"
db_path = settings.DATABASES['default']['NAME']
df = lf.collect(streaming=True)
conn = sqlite3.connect(db_path)
try:
if mode == 'replace':
# Drop existing table for retry/idempotency safety
conn.execute(f'DROP TABLE IF EXISTS "{table_name}"')
# Build CREATE TABLE from Polars schema
col_defs = [f'"{name}" {_dtype_to_sqlite(dtype)}'
for name, dtype in zip(df.columns, df.dtypes)]
conn.execute(f'CREATE TABLE "{table_name}" ({", ".join(col_defs)})')
# Bulk insert in chunks
col_names = [f'"{c}"' for c in df.columns]
placeholders = ','.join(['?' for _ in df.columns])
insert_sql = f'INSERT INTO "{table_name}" ({", ".join(col_names)}) VALUES ({placeholders})'
total_rows = len(df)
chunk_size = 10000
for i in range(0, total_rows, chunk_size):
chunk = df.slice(i, chunk_size)
conn.executemany(insert_sql, chunk.rows())
conn.commit()
finally:
conn.close()
logger.info('[SAVE_TO_DB] table=%s rows=%d mode=%s', table_name, total_rows, mode)
return table_name
def load_from_db(
table_name: str,
schema_overrides: dict[str, pl.DataType | str] | None = None,
) -> Optional[pl.LazyFrame]:
"""Load a SQLite table into a LazyFrame.
Args:
table_name: SQLite table name.
schema_overrides: Optional dict of ``{column_name: dtype}`` to cast
columns after loading. Values can be ``pl.DataType`` instances
(e.g. ``pl.Float64``) or dtype strings (e.g. ``'Float64'``).
Use when SQLite ``TEXT`` columns need to be restored to
``Float64`` / ``Int64`` etc.
Returns ``None`` if the table does not exist.
"""
from django.conf import settings
db_path = settings.DATABASES['default']['NAME']
conn = sqlite3.connect(db_path)
try:
# Check table exists
cursor = conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name=?",
(table_name,),
)
if cursor.fetchone() is None:
return None
conn.row_factory = sqlite3.Row
cursor = conn.execute(f'SELECT * FROM "{table_name}"')
# Peek first row to detect if table is empty
first_row = cursor.fetchone()
if first_row is None:
return pl.DataFrame().lazy()
batch_rows = [first_row]
dfs = []
_use_csv_fallback = False
while True:
chunk = cursor.fetchmany(10000)
if not chunk:
# Flush remaining rows
if batch_rows:
if _use_csv_fallback:
dfs.append(_rows_to_csv_df(batch_rows))
else:
try:
dfs.append(pl.from_dicts([dict(r) for r in batch_rows]))
except Exception:
_use_csv_fallback = True
dfs.append(_rows_to_csv_df(batch_rows))
break
batch_rows.extend(chunk)
if len(batch_rows) >= 10000:
if _use_csv_fallback:
dfs.append(_rows_to_csv_df(batch_rows))
else:
try:
dfs.append(pl.from_dicts([dict(r) for r in batch_rows]))
except Exception:
_use_csv_fallback = True
# Re-process current batch with CSV fallback
dfs.append(_rows_to_csv_df(batch_rows))
batch_rows = []
# Concat all batches
if len(dfs) == 1:
df = dfs[0]
else:
df = pl.concat(dfs, how='vertical')
# Apply type overrides to restore original column types
if schema_overrides:
casts = {}
for col, dtype in schema_overrides.items():
if col not in df.columns:
continue
dtype_obj = dtype if isinstance(dtype, pl.DataType) else _str_to_polars_dtype(dtype)
if df[col].dtype != dtype_obj:
try:
casts[col] = pl.col(col).cast(dtype_obj)
except Exception:
# If a cast fails (e.g. "A" → Int64), keep the column as-is
pass
if casts:
df = df.with_columns(list(casts.values()))
return df.lazy()
finally:
conn.close()
def load_from_db_lazy(
table_name: str,
schema_overrides: dict[str, pl.DataType | str] | None = None,
) -> Optional[pl.LazyFrame]:
"""Load a SQLite table into a LazyFrame using Polars native DB reader.
Uses ``pl.read_database`` for efficient bulk reading — an order of
magnitude faster than the manual batch-processing path in
:func:`load_from_db`. Returns a :class:`pl.LazyFrame` so downstream
operations remain lazy (query-plan, not materialised).
Args:
table_name: SQLite table name.
schema_overrides: Optional ``{column_name: dtype}`` dict to cast
columns after loading. Values can be ``pl.DataType`` instances
(e.g. ``pl.Float64``) or dtype strings (e.g. ``'Float64'``).
Returns:
``None`` if the table does not exist or is empty.
"""
from django.conf import settings
db_path = settings.DATABASES['default']['NAME']
conn = sqlite3.connect(db_path)
try:
# Check table existence (fast metadata query)
cursor = conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name=?",
(table_name,),
)
if cursor.fetchone() is None:
return None
# Bulk-read with Polars native engine — leverages the DBAPI2
# connection directly for efficient C-level data transfer.
query = f'SELECT * FROM "{table_name}"'
df = pl.read_database(connection=conn, query=query)
finally:
conn.close()
if df.height == 0:
return None
# Apply type overrides to restore original column types
# (SQLite stores everything as TEXT/INTEGER/REAL/BLOB;
# we need to cast back to Float64 / Int64 / etc.)
if schema_overrides:
casts = {}
for col, dtype in schema_overrides.items():
if col not in df.columns:
continue
dtype_obj = (
dtype if isinstance(dtype, pl.DataType)
else _str_to_polars_dtype(dtype)
)
if df[col].dtype != dtype_obj:
try:
casts[col] = pl.col(col).cast(dtype_obj)
except Exception:
pass # uncastable column — keep as-is
if casts:
df = df.with_columns(list(casts.values()))
rows = df.height
cols = df.width
logger.info(
'[LOAD_FROM_DB_LAZY] table=%s rows=%d cols=%d',
table_name, rows, cols,
)
return df.lazy()
def drop_sqlite_table(table_name: str) -> None:
"""Drop a SQLite data table silently if it exists."""
from django.conf import settings
db_path = settings.DATABASES['default']['NAME']
conn = sqlite3.connect(db_path)
try:
conn.execute(f'DROP TABLE IF EXISTS "{table_name}"')
conn.commit()
finally:
conn.close()
analysis/distance.py
-1114
View File
File diff suppressed because it is too large Load Diff
analysis/distance/__init__.py
+37
View File
@@ -0,0 +1,37 @@
"""Distance computation package for TianXuan analysis pipeline.
Provides distance functions for various data types (IP, string, boolean,
bytes, timestamp), plus normalisation, SVD denoising, UMAP reduction,
and per-cluster SVD feature extraction.
Public API:
ip_distance — subnet-mask + Haversine geographic distance for IP pairs
string_distance — Levenshtein / enum-like distance for string columns
bool_distance — 0/1 match-to-mode distance for boolean columns
bytes_distance — Hamming (popcount) distance for hex byte columns
timestamp_fft_distance — FFT phase distance for timestamp columns
normalize_features — StandardScaler normalisation (μ=0, σ=1)
svd_denoise — TruncatedSVD noise removal (returns noise_profile)
umap_reduce — UMAP dimensionality reduction (3D→2D fallback)
cluster_svd_extract — Per-cluster SVD distinguishing feature extraction
"""
from analysis.distance._geo import ip_distance
from analysis.distance._text import bool_distance, string_distance
from analysis.distance._numeric import bytes_distance, timestamp_fft_distance
from analysis.distance._normalize import normalize_features
from analysis.distance._svd import svd_denoise
from analysis.distance._umap import umap_reduce
from analysis.distance._cluster_svd import cluster_svd_extract
__all__ = [
"ip_distance",
"string_distance",
"bool_distance",
"bytes_distance",
"timestamp_fft_distance",
"normalize_features",
"svd_denoise",
"umap_reduce",
"cluster_svd_extract",
]
analysis/distance/_cluster_svd.py
+87
View File
@@ -0,0 +1,87 @@
"""Per-cluster SVD feature extraction."""
import numpy as np
import polars as pl
from sklearn.decomposition import TruncatedSVD
def cluster_svd_extract(
lf: pl.LazyFrame,
labels: list[int],
feature_cols: list[str],
) -> dict[int, dict]:
"""Per-cluster SVD-based distinguishing feature extraction.
For each cluster label:
1. Extracts rows belonging to that cluster.
2. Fits ``TruncatedSVD(n_components=min(N, len(feature_cols)))``.
3. The absolute loadings on the first SVD component serve as
*feature_strength* — higher values mean more distinguishing.
Returns a mapping ``{cluster_label: {...}}`` where each value is::
{
'features': [...], # top-10 feature names by strength
'explained_variance_ratio': [...], # per-component ratios
'feature_strength': [...], # per-feature loading magnitude
}
Parameters
----------
lf:
LazyFrame with rows aligned to *labels*.
labels:
Cluster labels (same length as *lf* row count, ``-1`` = noise).
feature_cols:
Numeric column names to analyse.
Returns
-------
dict[int, dict]
Per-cluster SVD extraction results.
"""
available = [c for c in feature_cols if c in lf.collect_schema().names()]
if not available:
return {}
df = lf.select(available).collect(streaming=True)
mat = df.to_numpy().astype(np.float64)
mat = np.nan_to_num(mat, nan=0.0, posinf=0.0, neginf=0.0)
if len(labels) != len(mat):
labels = labels[:len(mat)]
unique_labels = sorted(set(labels))
result: dict[int, dict] = {}
for label in unique_labels:
mask = np.array(labels) == label
n = int(mask.sum())
if n < 2:
continue
cluster_mat = mat[mask]
n_comp = min(n, len(available))
svd = TruncatedSVD(n_components=n_comp, random_state=42)
svd.fit(cluster_mat)
# First component loadings → feature strength
if n_comp > 0 and svd.components_.shape[0] > 0:
feature_strength = np.abs(svd.components_[0]).astype(np.float64)
else:
feature_strength = np.zeros(len(available), dtype=np.float64)
total_strength = float(feature_strength.sum())
if total_strength > 1e-10:
feature_strength = feature_strength / total_strength
top_indices = np.argsort(-feature_strength)
top_k = min(10, len(available))
result[int(label)] = {
'features': [available[i] for i in top_indices[:top_k]],
'explained_variance_ratio': svd.explained_variance_ratio_.tolist(),
'feature_strength': feature_strength.tolist(),
}
return result
analysis/distance/_geo.py
+170
View File
@@ -0,0 +1,170 @@
"""IP distance functions for Polars LazyFrames.
Provides subnet-mask distance and Haversine geographic distance
between pairs of IP columns. All distances are normalised to [0, 1];
invalid / non-geolocatable IPs produce ``NaN``.
"""
from __future__ import annotations
import math
from typing import Optional
import polars as pl
from analysis.geoip import ip_to_int, lookup
# ---------------------------------------------------------------------------
# Earth constant
# ---------------------------------------------------------------------------
_HALF_EARTH_CIRCUMFERENCE_KM = 20_037.508 # equatorial / 2
# ---------------------------------------------------------------------------
# Core distance helpers
# ---------------------------------------------------------------------------
def _subnet_mask_distance(ip1: Optional[str], ip2: Optional[str], mask: int) -> float:
"""Normalised subnet-mask distance between two IP strings.
Computes ``abs(int(ip1)/2**mask - int(ip2)/2**mask) / 2**(32-mask)``.
Returns ``NaN`` when either IP is ``None`` or unparseable.
"""
if ip1 is None or ip2 is None:
return float("nan")
if not isinstance(ip1, str):
ip1 = str(ip1)
if not isinstance(ip2, str):
ip2 = str(ip2)
int1 = ip_to_int(ip1)
int2 = ip_to_int(ip2)
if int1 is None or int2 is None:
return float("nan")
divisor = 2.0 ** mask
max_val = 2.0 ** (32 - mask)
return abs(int1 / divisor - int2 / divisor) / max_val
def _haversine_km(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
"""Great-circle distance (km) via the Haversine formula."""
dlat = math.radians(lat2 - lat1)
dlon = math.radians(lon2 - lon1)
a = (
math.sin(dlat / 2) ** 2
+ math.cos(math.radians(lat1))
* math.cos(math.radians(lat2))
* math.sin(dlon / 2) ** 2
)
return 6371.0 * 2.0 * math.asin(math.sqrt(a))
def _geo_distance(ip1: Optional[str], ip2: Optional[str]) -> float:
"""Normalised [0, 1] geographic distance between two IP strings.
Uses ``analysis.geoip.lookup()`` to obtain lat/lon, then Haversine.
Returns ``NaN`` when either IP cannot be geolocated.
"""
if ip1 is None or ip2 is None:
return float("nan")
if not isinstance(ip1, str):
ip1 = str(ip1)
if not isinstance(ip2, str):
ip2 = str(ip2)
geo1 = lookup(ip1)
geo2 = lookup(ip2)
if geo1 is None or geo2 is None:
return float("nan")
km = _haversine_km(geo1["lat"], geo1["lon"], geo2["lat"], geo2["lon"])
return km / _HALF_EARTH_CIRCUMFERENCE_KM
# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------
def ip_distance(
lf: pl.LazyFrame,
ip_columns: list[str],
subnet_masks: Optional[list[int]] = None,
) -> pl.LazyFrame:
"""Append IP distance columns to a LazyFrame.
For each **consecutive pair** of *ip_columns*
(``ip_columns[0]`` vs ``ip_columns[1]``,
``ip_columns[2]`` vs ``ip_columns[3]``, …) computes:
* one subnet-mask distance per entry in *subnet_masks*
* one Haversine geographic distance
All distances are normalised to [0, 1]. Invalid IPs or IPs not
found in the GeoIP database produce ``NaN``.
Parameters
----------
lf:
LazyFrame containing the IP columns.
ip_columns:
List of IP column names (used in consecutive pairs).
subnet_masks:
Subnet masks for network distance (default ``[24, 28]``).
Returns
-------
pl.LazyFrame
Input frame with ``_ip_dist_0``, ``_ip_dist_1``, …
columns appended.
Examples
--------
>>> lf = pl.LazyFrame({
... "src_ip": ["8.8.8.8", "1.1.1.1"],
... "dst_ip": ["8.8.4.4", "1.0.0.1"],
... })
>>> result = ip_distance(lf, ["src_ip", "dst_ip"], subnet_masks=[24])
>>> result.collect().columns
['src_ip', 'dst_ip', '_ip_dist_0', '_ip_dist_1']
"""
if subnet_masks is None:
subnet_masks = [24, 28]
result = lf
dist_idx = 0
# Walk consecutive pairs
for i in range(0, len(ip_columns) - 1, 2):
col_a = ip_columns[i]
col_b = ip_columns[i + 1]
# -- subnet-mask distances --
for mask in subnet_masks:
name = f"_ip_dist_{dist_idx}"
result = result.with_columns(
pl.struct([col_a, col_b])
.map_elements(
lambda s, _mask=mask, _ca=col_a, _cb=col_b: (
_subnet_mask_distance(s[_ca], s[_cb], _mask)
),
return_dtype=pl.Float64,
)
.alias(name)
)
dist_idx += 1
# -- geographic distance --
name = f"_ip_dist_{dist_idx}"
result = result.with_columns(
pl.struct([col_a, col_b])
.map_elements(
lambda s, _ca=col_a, _cb=col_b: _geo_distance(s[_ca], s[_cb]),
return_dtype=pl.Float64,
)
.alias(name)
)
dist_idx += 1
return result
analysis/distance/_normalize.py
+52
View File
@@ -0,0 +1,52 @@
"""StandardScaler normalisation (step 2 of preprocessing pipeline)."""
import numpy as np
import polars as pl
from sklearn.preprocessing import StandardScaler
def normalize_features(lf: pl.LazyFrame, feature_cols: list[str]) -> pl.LazyFrame:
"""Apply StandardScaler normalisation to numeric feature columns.
Collects *feature_cols* into a numpy array, replaces ``NaN`` with
zero, centres to μ=0 σ=1 via ``sklearn.preprocessing.StandardScaler``,
then returns the result as *lf* extended with ``_norm_*`` prefixed
columns.
Parameters
----------
lf:
LazyFrame containing *feature_cols*.
feature_cols:
Numeric column names to normalise. Columns not present in
*lf* are silently skipped; non-numeric columns raise an error.
Returns
-------
pl.LazyFrame
Input frame with additional ``_norm_{col}`` columns.
Memory
------
Peak memory ≈ collected numpy matrix + scaler workspace
(well under the 800 MB budget for typical inputs).
"""
available = [c for c in feature_cols if c in lf.collect_schema().names()]
if not available:
return lf
df = lf.select(available).collect(streaming=True)
mat = df.to_numpy().astype(np.float64)
mat = np.nan_to_num(mat, nan=0.0)
scaler = StandardScaler()
scaled = scaler.fit_transform(mat)
# Build new columns with _norm_ prefix
df_result = lf.collect(streaming=True)
for i, col in enumerate(available):
df_result = df_result.with_columns(
pl.Series(f'_norm_{col}', scaled[:, i])
)
return df_result.lazy()
analysis/distance/_numeric.py
+278
View File
@@ -0,0 +1,278 @@
"""Bytes/hex distance (Hamming/popcount) and timestamp FFT phase distance."""
import numpy as np
import polars as pl
# ---------------------------------------------------------------------------
# Bytes Hamming distance
# ---------------------------------------------------------------------------
_HEX_POPCOUNT: dict[str, int] = {
'0': 0, '1': 1, '2': 1, '3': 2,
'4': 1, '5': 2, '6': 2, '7': 3,
'8': 1, '9': 2,
'a': 2, 'b': 3, 'c': 2, 'd': 3, 'e': 3, 'f': 4,
'A': 2, 'B': 3, 'C': 2, 'D': 3, 'E': 3, 'F': 4,
}
def _hex_popcount_normalized(val: str | None) -> float:
"""Return the proportion of 1-bits in a hex byte string (0.0–1.0).
This is the Hamming distance from the all-zero bit string,
normalised by total bits. Returns ``NaN`` for null or
non-hex input.
"""
if val is None:
return float("nan")
s = str(val).replace(' ', '').replace('\t', '').replace('\n', '')
if not s:
return 0.0
total_ones = 0
for ch in s:
ones = _HEX_POPCOUNT.get(ch)
if ones is None:
return float("nan") # non-hex character
total_ones += ones
total_bits = len(s) * 4
return total_ones / total_bits
def bytes_distance(
lf: pl.LazyFrame,
byte_columns: list[str],
) -> pl.LazyFrame:
"""Append normalised Hamming-distance columns for hex byte columns.
Each byte column is converted from hex to binary, and the
proportion of 1-bits (Hamming distance from the all-zero
bit vector) is computed as a ``_bytes_dist_N`` Float64 column
normalised to [0, 1]. Non-hex / null values produce ``NaN``.
Parameters
----------
lf:
LazyFrame containing the byte columns.
byte_columns:
Column names whose values are hex byte strings
(with optional whitespace separators, e.g. ``"A3 F2"``).
Returns
-------
pl.LazyFrame
Input frame with ``_bytes_dist_0``, ``_bytes_dist_1``, …
appended. Columns not present in *lf* are silently skipped.
Examples
--------
>>> df = pl.DataFrame({"payload": ["A3F2", "FFFF", None]})
>>> result = bytes_distance(df.lazy(), ["payload"])
>>> result.collect()
shape: (3, 2)
┌─────────┬────────────────┐
│ payload ┆ _bytes_dist_0 │
│ --- ┆ --- │
│ str ┆ f64 │
╞═════════╪════════════════╡
│ A3F2 ┆ 0.625 │
│ FFFF ┆ 1.0 │
│ null ┆ NaN │
└─────────┴────────────────┘
"""
existing = set(lf.collect_schema().names())
result = lf
for i, col_name in enumerate(byte_columns):
if col_name not in existing:
continue
result = result.with_columns(
pl.col(col_name)
.cast(pl.Utf8)
.map_elements(_hex_popcount_normalized, return_dtype=pl.Float64)
.alias(f'_bytes_dist_{i}')
)
return result
# ---------------------------------------------------------------------------
# Timestamp FFT phase distance
# ---------------------------------------------------------------------------
# ── threshold: minimum fraction of repeated intervals to use FFT path ──
_FFT_PERIODICITY_THRESHOLD = 0.05
# ── minimum samples before FFT is attempted ──
_FFT_MIN_SAMPLES = 4
def _parse_timestamp_to_epoch(series: pl.Series) -> pl.Series:
"""Convert a timestamp series to Unix epoch seconds (float64).
Handles int/float (assumed already epoch), pl.Datetime/pl.Date
(ns→seconds), and str/pl.Utf8 (strptime→ns→seconds, with direct
float-cast fallback on parse failure).
"""
dtype = series.dtype
if dtype in (pl.Utf8, pl.String):
try:
parsed = series.str.strptime(pl.Datetime, format=None, strict=False)
except Exception:
return series.cast(pl.Float64)
return parsed.cast(pl.Int64).cast(pl.Float64) / 1_000_000_000.0
elif dtype in (pl.Datetime, pl.Date):
return series.cast(pl.Int64).cast(pl.Float64) / 1_000_000_000.0
else:
return series.cast(pl.Float64)
def _compute_fft_distances(values: np.ndarray) -> np.ndarray:
"""Compute per-row phase distances for a 1-D array.
Sorts *values*, checks for periodicity, then either:
* FFT path: ``np.fft.fft()`` → ``np.angle()``, mapped back to
original (unsorted) row order via rank lookup.
* Statistical fallback: interval z-scores mapped to row order.
Returns an ``(n,)`` float64 array, normalised to zero-mean
unit-variance.
"""
n = len(values)
if n < _FFT_MIN_SAMPLES:
return np.full(n, np.nan, dtype=np.float64)
# Sort once; record rank→position mapping for unsorting later
sorted_order = np.argsort(values)
ranks = np.empty(n, dtype=np.int64)
ranks[sorted_order] = np.arange(n)
sorted_vals = values[sorted_order]
# ── periodicity check ──────────────────────────────────────────
intervals = np.diff(sorted_vals)
unique_int, counts = np.unique(
np.round(intervals, 6), return_counts=True
)
repeat_ratio = float(np.sum(counts[counts > 1])) / max(len(intervals), 1)
if repeat_ratio >= _FFT_PERIODICITY_THRESHOLD:
# ── FFT phase path ─────────────────────────────────────────
centered = sorted_vals - np.mean(sorted_vals)
fft_result = np.fft.fft(centered)
phases = np.angle(fft_result)
dist_values = phases[ranks]
else:
# ── statistical fallback ───────────────────────────────────
mean_int = float(np.mean(intervals))
std_int = float(np.std(intervals))
if std_int > 1e-10:
z_scores = (intervals - mean_int) / std_int
else:
z_scores = np.zeros_like(intervals)
# Row at sorted position r gets the z-score of interval r-1
# (preceding interval); row at position 0 reuses last interval.
dist_values = np.empty(n, dtype=np.float64)
for r in range(n):
if r > 0:
dist_values[ranks[r]] = z_scores[r - 1]
else:
dist_values[ranks[0]] = (
z_scores[-1] if n > 1 else 0.0
)
# Normalise to zero-mean unit-variance
d_std = float(np.std(dist_values))
d_mean = float(np.mean(dist_values))
if d_std > 1e-10:
dist_values = (dist_values - d_mean) / d_std
else:
dist_values = dist_values - d_mean
return dist_values
def timestamp_fft_distance(
lf: pl.LazyFrame, ts_columns: list[str]
) -> pl.LazyFrame:
"""Append FFT-based phase-distance columns for timestamp columns.
For each column in *ts_columns* that exists in the schema:
1. Parse values to Unix epoch float seconds (handles int, float,
datetime strings, and ``pl.Datetime``).
2. Collect the column to numpy, sort, then run ``np.fft.fft``
followed by ``np.angle`` to extract per-row phase values.
3. If fewer than 5 % of timestamp intervals repeat (weak
periodicity), fall back to statistical features derived
from interval mean, standard deviation, and z-scores.
Columns are appended as ``_ts_dist_0``, ``_ts_dist_1``, … in the
order *ts_columns* are supplied. Rows whose timestamp is null
receive ``NaN``.
Parameters
----------
lf : pl.LazyFrame
The input LazyFrame with at least one timestamp column.
ts_columns : list[str]
Column names to process.
Returns
-------
pl.LazyFrame
A new LazyFrame with ``_ts_dist_N`` columns appended.
Examples
--------
>>> import polars as pl
>>> from analysis.distance import timestamp_fft_distance
>>> df = pl.DataFrame({
... 'ts': [0, 60, 120, 180, 240, 300],
... })
>>> result = timestamp_fft_distance(df.lazy(), ['ts'])
>>> result.collect().columns
['ts', '_ts_dist_0']
"""
schema = lf.collect_schema()
available = [c for c in ts_columns if c in schema.names()]
if not available:
return lf
# Collect all needed columns once
df = lf.select(available).collect(streaming=True)
n_rows = len(df)
dist_arrays: list[np.ndarray] = []
for col in available:
series = df[col]
ts_numeric = _parse_timestamp_to_epoch(series)
# Identify valid (non-null) rows
valid = ~ts_numeric.is_null()
valid_indices = np.where(valid.to_numpy())[0]
values = ts_numeric.filter(valid).to_numpy().astype(np.float64)
if len(values) < _FFT_MIN_SAMPLES:
continue
phase_values = _compute_fft_distances(values)
# Build full-length array; invalid rows stay NaN
full_dist = np.full(n_rows, np.nan, dtype=np.float64)
full_dist[valid_indices] = phase_values
dist_arrays.append(full_dist)
if not dist_arrays:
return lf
# Re-materialise the original frame and bolt on distance columns
result_df = lf.collect(streaming=True)
for i, arr in enumerate(dist_arrays):
result_df = result_df.with_columns(
pl.Series(f'_ts_dist_{i}', arr)
)
return result_df.lazy()
analysis/distance/_svd.py
+116
View File
@@ -0,0 +1,116 @@
"""SVD noise removal (step 3 of preprocessing pipeline)."""
import numpy as np
import polars as pl
def svd_denoise(
lf: pl.LazyFrame,
feature_cols: list[str],
variance_threshold: float = 0.95,
) -> tuple[pl.LazyFrame, dict]:
"""Remove noise from numeric features via TruncatedSVD.
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. 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 signal and treats 5 % as noise.
Returns
-------
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, noise_profile
df = lf.select(available).collect(streaming=True)
mat = df.to_numpy().astype(np.float64)
mat = np.nan_to_num(mat, nan=0.0)
n, d = mat.shape
max_components = min(n, d - 1)
if max_components < 1:
return lf, noise_profile
from sklearn.decomposition import TruncatedSVD
svd = TruncatedSVD(n_components=max_components, random_state=42)
_ = svd.fit_transform(mat)
total_variance = float(np.sum(svd.explained_variance_))
if total_variance <= 1e-10:
return lf, noise_profile
cumsum = np.cumsum(svd.explained_variance_)
for k in range(1, max_components + 1):
if cumsum[k - 1] / total_variance >= variance_threshold:
break
else:
k = max_components
# 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(col, reconstructed[:, i])
)
return df_result.lazy(), noise_profile
analysis/distance/_text.py
+164
View File
@@ -0,0 +1,164 @@
"""String distance, enum detection, and boolean distance functions."""
import polars as pl
# ---------------------------------------------------------------------------
# Enum-like detection helper
# ---------------------------------------------------------------------------
def _is_enum_like(s: pl.Series) -> bool:
"""Check if ≥50% of rows have values appearing ≥2 times."""
total = len(s)
if total == 0:
return False
non_null = s.drop_nulls()
if len(non_null) == 0:
return False
vc = non_null.value_counts()
if 'count' not in vc.columns:
return False
duplicated = vc.filter(pl.col('count') >= 2)
if len(duplicated) == 0:
return False
duplicated_count: int = duplicated['count'].sum() # type: ignore[union-attr]
return duplicated_count >= 0.5 * total
def _levenshtein_normalized(s1: str | None, s2: str) -> float:
"""Normalised Levenshtein distance between *s1* and *s2*.
Empty or ``None`` values return 1.0 (maximum distance).
"""
if s1 is None or s1 == '':
return 1.0
if s2 is None or s2 == '':
return 1.0
s1_str = str(s1)
s2_str = str(s2)
max_len = max(len(s1_str), len(s2_str))
if max_len == 0:
return 0.0
import Levenshtein
return Levenshtein.distance(s1_str, s2_str) / max_len
# ---------------------------------------------------------------------------
# String distance
# ---------------------------------------------------------------------------
def string_distance(lf: pl.LazyFrame, str_columns: list[str]) -> pl.LazyFrame:
"""Compute string distance columns.
For each column in *str_columns*:
* If **≥50%** of rows have values appearing ≥2 times, the column is
treated as an enum: 0 when the value equals the mode, 1 otherwise.
* Otherwise, the normalised Levenshtein distance to the mode is
computed for each row.
Empty or ``None`` values always produce the maximum distance (1.0).
Parameters
----------
lf:
Input LazyFrame.
str_columns:
List of string column names to compute distances for.
Returns
-------
pl.LazyFrame
Input LazyFrame extended with ``_str_dist_0``, ``_str_dist_1``,
... columns (one per entry in *str_columns*).
"""
df = lf.collect(streaming=True)
dist_exprs: list[pl.Expr] = []
for idx, col in enumerate(str_columns):
if col not in df.columns:
continue
series = df[col].cast(pl.Utf8)
enum_like = _is_enum_like(series)
non_null = series.drop_nulls()
if len(non_null) > 0:
mode_val = non_null.value_counts().sort('count', descending=True).row(0)[0]
mode_str = str(mode_val)
else:
mode_str = ''
col_name = f'_str_dist_{idx}'
if enum_like:
dist_exprs.append(
pl.when(pl.col(col).is_null())
.then(pl.lit(1.0))
.when(pl.col(col).cast(pl.Utf8) == pl.lit(mode_str))
.then(pl.lit(0.0))
.otherwise(pl.lit(1.0))
.alias(col_name)
)
else:
# Capture in closure for map_elements
_mode = mode_str
dist_exprs.append(
pl.col(col).map_elements(
lambda v: _levenshtein_normalized(v, _mode),
return_dtype=pl.Float64,
).alias(col_name)
)
return lf.with_columns(dist_exprs)
# ---------------------------------------------------------------------------
# Boolean distance
# ---------------------------------------------------------------------------
def bool_distance(lf: pl.LazyFrame, bool_columns: list[str]) -> pl.LazyFrame:
"""Compute boolean distance columns.
For each column in *bool_columns*, each row gets:
0 if its value matches the column mode, 1 otherwise.
``None`` values always produce 1.0 (maximum distance).
Parameters
----------
lf:
Input LazyFrame.
bool_columns:
List of boolean column names to compute distances for.
Returns
-------
pl.LazyFrame
Input LazyFrame extended with ``_bool_dist_0``, ``_bool_dist_1``,
... columns (one per entry in *bool_columns*).
"""
df = lf.collect(streaming=True)
dist_exprs: list[pl.Expr] = []
for idx, col in enumerate(bool_columns):
if col not in df.columns:
continue
series = df[col]
non_null = series.drop_nulls()
col_name = f'_bool_dist_{idx}'
if len(non_null) == 0:
dist_exprs.append(pl.lit(0.0).alias(col_name))
else:
mode_val = non_null.value_counts().sort('count', descending=True).row(0)[0]
dist_exprs.append(
pl.when(pl.col(col).is_null())
.then(pl.lit(1.0))
.when(pl.col(col) == pl.lit(mode_val))
.then(pl.lit(0.0))
.otherwise(pl.lit(1.0))
.alias(col_name)
)
return lf.with_columns(dist_exprs)
analysis/distance/_umap.py
+122
View File
@@ -0,0 +1,122 @@
"""UMAP dimensionality reduction (3D → 2D fallback)."""
import logging
import numpy as np
import polars as pl
from sklearn.preprocessing import StandardScaler
import umap
def umap_reduce(
lf: pl.LazyFrame,
feature_cols: list[str],
n_components: int = 3,
) -> tuple[pl.LazyFrame, dict]:
"""UMAP dimensionality reduction with automatic 3D→2D fallback.
Trains on a ≤10K-row sample (when *lf* exceeds that limit), then
batch-transforms the remaining rows in 1K-row batches.
If ``n_components=3`` fails (e.g. too few rows or numerical issues),
falls back to ``n_components=2`` automatically.
Parameters
----------
lf:
LazyFrame containing the numeric feature columns.
feature_cols:
Numeric column names to use for UMAP projection. Columns not
present in *lf* are silently skipped.
n_components:
Target dimensionality (2 or 3). Default 3 with fallback to 2.
Returns
-------
tuple[pl.LazyFrame, dict]
- LazyFrame: input data extended with ``_umap_x``, ``_umap_y``
(and ``_umap_z`` when ``n_components ≥ 3``).
- dict: metadata with keys ``n_components`` and ``reducer``.
"""
logger = logging.getLogger(__name__)
MAX_TRAIN = 10000
BATCH_SIZE = 1000
available = [c for c in feature_cols if c in lf.collect_schema().names()]
if not available:
return lf, {'n_components': 0, 'reducer': 'umap', 'error': 'no feature columns available'}
df = lf.collect(streaming=True)
n_total = len(df)
if n_total < 2:
return lf, {'n_components': 0, 'reducer': 'umap', 'error': 'too few rows'}
mat = df.select(available).to_numpy().astype(np.float64)
mat = np.nan_to_num(mat, nan=0.0, posinf=0.0, neginf=0.0)
scaler = StandardScaler()
if n_total > MAX_TRAIN:
rng = np.random.RandomState(42)
idx_sample = rng.choice(n_total, size=MAX_TRAIN, replace=False)
mat_sample = mat[idx_sample]
scaler.fit(mat_sample)
mat_sample_scaled = np.nan_to_num(scaler.transform(mat_sample), nan=0.0)
else:
scaler.fit(mat)
mat_sample_scaled = np.nan_to_num(scaler.transform(mat), nan=0.0)
# ── Train UMAP; fallback 3→2 on failure ──
reducer = None
coords = None
effective_n = n_components
# Try requested n_components first, then fall back to 2 (no duplicates)
attempts = sorted(set((n_components, 2)), reverse=True)
for attempt_n in attempts:
if attempt_n != n_components:
logger.warning(f'UMAP {n_components}D failed, falling back to 2D')
effective_n = attempt_n
try:
reducer = umap.UMAP(
n_components=effective_n, random_state=42,
n_neighbors=15, min_dist=0.1, metric='euclidean',
)
if n_total > MAX_TRAIN:
reducer.fit(mat_sample_scaled)
else:
coords = reducer.fit_transform(mat_sample_scaled)
break
except Exception as e:
if attempt_n == 2:
raise
logger.warning(f'UMAP {attempt_n}D failed: {e}')
# ── Batch transform if dataset was larger than training sample ──
if n_total > MAX_TRAIN and reducer is not None:
coords_list: list[np.ndarray] = []
for start in range(0, n_total, BATCH_SIZE):
end = min(start + BATCH_SIZE, n_total)
mat_batch = scaler.transform(mat[start:end])
mat_batch = np.nan_to_num(mat_batch, nan=0.0, posinf=0.0, neginf=0.0)
coords_list.append(reducer.transform(mat_batch))
coords = np.vstack(coords_list)
if coords is None:
return lf, {'n_components': 0, 'reducer': 'umap', 'error': 'UMAP produced no output'}
# ── Build result LazyFrame with coordinate columns ──
coord_series = [
pl.Series('_umap_x', coords[:, 0]),
pl.Series('_umap_y', coords[:, 1]),
]
if effective_n >= 3 and coords.shape[1] >= 3:
coord_series.append(pl.Series('_umap_z', coords[:, 2]))
df_result = df.with_columns(coord_series)
metadata: dict = {
'n_components': effective_n,
'reducer': 'umap',
}
return df_result.lazy(), metadata
analysis/views/clustering.py
+2 -307
View File
@@ -222,312 +222,7 @@ def cluster_detail(request, display_id, cluster_label):
@profile
def _run_clustering_pipeline(run, store, ds_id, feature_columns, algorithm,
min_cluster_size, run_umap=True, head=None):
"""Clustering pipeline on raw rows: clustering → feature extraction → UMAP → ORM save.
Operates directly on raw data rows (no entity aggregation).
Per-row results stored in EntityProfile with row index as identifier.
Args:
run: AnalysisRun ORM object.
store: SessionStore instance.
ds_id: dataset ID in SessionStore pointing to raw data.
feature_columns: list of feature column names (None = auto-detect all numeric).
algorithm: clustering algorithm.
min_cluster_size: int for HDBSCAN.
run_umap: bool, whether to compute and save UMAP-2D embeddings.
head: optional int, downsample to at most this many rows (min 100).
"""
import warnings
warnings.filterwarnings('ignore', category=RuntimeWarning, module='sklearn')
try:
from analysis.tool_registry import _handle_run_clustering, _handle_extract_features
import polars as pl
entry = store.get_dataset(ds_id)
if entry is None:
run.status = 'failed'
run.error_message = '数据集在会话中不存在'
run.save(update_fields=['status', 'error_message'])
return
schema = entry.get('schema', {})
# ── Downsampling ──
if head is not None and head > 0:
lf = entry['lazyframe']
try:
row_count = lf.select(pl.len()).collect(streaming=True).item()
except Exception:
logger.error("unknown failed: {}".format(traceback.format_exc()))
row_count = 0
if row_count > head:
head = max(head, 100) # ensure minimum rows for clustering
run.progress_msg = f'正在采样 {head}/{row_count} 行...'
run.save(update_fields=['progress_msg'])
lf = lf.head(head)
store.store_dataset(
ds_id, lf,
schema=schema,
metadata=entry.get('metadata', {}),
)
entry = store.get_dataset(ds_id)
row_count = head
# Resolve feature columns — use actual LazyFrame dtypes (not stored schema dict
# which may be stale after Utf8 coercion in _background_process).
numeric_types = (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
pl.Float32, pl.Float64)
lf = entry['lazyframe']
live_schema = lf.collect_schema()
live_names = live_schema.names()
live_dtypes = list(live_schema.dtypes())
if feature_columns:
feature_cols = [
c for c in feature_columns
if c in live_names and live_dtypes[live_names.index(c)] in numeric_types
]
else:
feature_cols = [
live_names[i] for i, dt in enumerate(live_dtypes)
if dt in numeric_types and not live_names[i].startswith('_')
]
# ── Step 1: Clustering ────────────────────────────────────────────────
# If no numeric columns detected, pass empty list to let
# _handle_run_clustering apply its Utf8→Float64 coercion fallback.
run.status = 'clustering'
run.progress_pct = 70
run.progress_msg = '正在进行聚类分析...'
run.save(update_fields=['status', 'progress_pct', 'progress_msg'])
result = asyncio.run(_handle_run_clustering(
dataset_id=ds_id,
cluster_columns=feature_cols,
algorithm=algorithm,
params={'min_cluster_size': min_cluster_size},
random_state=RANDOM_SEED,
))
if 'error' in result:
err = result['error']
if 'numeric' in err.lower():
available = [
f"{k}({v})" for k, v in schema.items()
if v.split('(')[0].strip() in numeric_types and not k.startswith('_')
]
err += f"。可用数值列: {available}"
raise Exception(err)
cluster_id = result.get('cluster_result_id', '')
run.cluster_count = result.get('n_clusters', 0)
run.save(update_fields=['cluster_count'])
# ── Step 2: Feature extraction ──────────────────────────────────────
run.status = 'extracting'
run.progress_pct = 90
run.progress_msg = '正在提取聚类特征...'
run.save(update_fields=['status', 'progress_pct', 'progress_msg'])
feat_result = asyncio.run(_handle_extract_features(
dataset_id=ds_id,
cluster_result_id=cluster_id,
top_k=10, method='zscore', save_to_db=True,
run_id=run.id,
))
if 'error' in feat_result:
raise Exception(feat_result['error'])
# ── Step 2b: Cluster SVD feature extraction ──────────────────────
run.progress_msg = '正在进行聚类SVD特征提取...'
run.save(update_fields=['progress_msg'])
try:
from analysis.distance import cluster_svd_extract
lf = entry['lazyframe']
cluster_entry = store.get_cluster_result(cluster_id)
labels_list = cluster_entry.get('labels', []) if cluster_entry else []
if labels_list and feature_cols:
svd_results = cluster_svd_extract(lf, labels_list, feature_cols)
# Persist SVD features to ClusterFeature (method='cluster_svd')
from analysis.models import ClusterFeature as CF
from analysis.models import ClusterResult as CR
for label, svd_info in svd_results.items():
cr = CR.objects.filter(run=run, cluster_label=label).first()
if cr is None:
continue
# Delete old zscore entries for this cluster, replace with SVD
CF.objects.filter(cluster=cr, distinguishing_method='zscore').delete()
for feat_name in svd_info.get('features', [])[:10]:
idx = svd_info['features'].index(feat_name) if feat_name in svd_info['features'] else -1
strength = (svd_info.get('feature_strength', [])[idx]
if idx >= 0 and idx < len(svd_info.get('feature_strength', []))
else None)
CF.objects.get_or_create(
cluster=cr,
feature_name=feat_name,
defaults={
'mean': None, 'std': None, 'median': None,
'p25': None, 'p75': None, 'missing_rate': None,
'distinguishing_score': float(strength) if strength is not None else None,
'distinguishing_method': 'cluster_svd',
},
)
except Exception as svd_err:
logger.warning(f'Cluster SVD extraction skipped: {svd_err}')
# ── Step 2c: Save per-row cluster labels to EntityProfile ──
# Always persist row→cluster mapping so cluster_detail can show data rows
# even when UMAP is skipped or fails.
try:
cluster_entry_save = store.get_cluster_result(cluster_id)
labels_list_save = cluster_entry_save.get('labels', []) if cluster_entry_save else []
if labels_list_save:
from analysis.models import EntityProfile as EP_save
from analysis.models import ClusterResult as CR_save
# Delete stale EP records for this run first
EP_save.objects.filter(run=run).delete()
cr_cache_save = {}
row_profiles = []
for r_idx in range(len(labels_list_save)):
lbl = labels_list_save[r_idx]
ck = f'{run.id}_{lbl}'
if ck not in cr_cache_save:
cr_cache_save[ck] = CR_save.objects.filter(run=run, cluster_label=lbl).first()
row_profiles.append(EP_save(
entity_value=f'row_{r_idx}', run=run, cluster_label=lbl,
cluster=cr_cache_save[ck],
embedding_x=None, embedding_y=None, embedding_z=None,
))
if row_profiles:
EP_save.objects.bulk_create(row_profiles, ignore_conflicts=True, batch_size=1000)
run.entity_count = len(row_profiles)
run.save(update_fields=['entity_count'])
except Exception as save_err:
logger.warning(f'Row labels save skipped: {save_err}')
# ── Step 3: UMAP-3D embedding (2D fallback) ───────────────────────
if run_umap:
try:
lf = entry['lazyframe']
# UMAP: sample max 10K for training, batch-transform rest in 1K batches
MAX_UMAP_TRAIN = UMAP_TRAIN_SAMPLE
BATCH_SIZE = UMAP_BATCH_SIZE
try:
n_total = lf.select(pl.len()).collect(streaming=True).item()
except Exception:
logger.error("unknown failed: {}".format(traceback.format_exc()))
n_total = 0
from sklearn.preprocessing import StandardScaler
import umap
import numpy as np
live_schema = lf.collect_schema()
num_cols = [name for name, dt in zip(live_schema.names(), live_schema.dtypes())
if dt in (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
pl.Float32, pl.Float64) and not name.startswith('_')]
if len(num_cols) >= 2:
if n_total > MAX_UMAP_TRAIN:
run.progress_msg = f'UMAP 采样 {MAX_UMAP_TRAIN}/{n_total} 行训练...'
run.save(update_fields=['progress_msg'])
# LazyFrame.sample() is unavailable in Polars 1.42 — collect then sample
df_sample = df_umap.sample(n=MAX_UMAP_TRAIN, seed=RANDOM_SEED)
df_sample = df_umap
umap_components = 3 # try 3D first
coords = None
reducer = None
if len(df_umap) > 2:
for attempt_n in (3, 2):
try:
if n_total > MAX_UMAP_TRAIN:
# Train UMAP on 10K sample, batch-transform full dataset
mat_sample = df_sample.select(num_cols).to_numpy()
scaler = StandardScaler().fit(mat_sample)
mat_sample_scaled = np.nan_to_num(
scaler.transform(mat_sample), nan=0.0)
reducer = umap.UMAP(n_components=attempt_n,
random_state=RANDOM_SEED,
n_neighbors=15, min_dist=0.1,
metric='euclidean')
reducer.fit(mat_sample_scaled)
# Batch transform in 1K-row batches
coords_list = []
for start in range(0, len(df_umap), BATCH_SIZE):
end = min(start + BATCH_SIZE, len(df_umap))
mat_batch = scaler.transform(
df_umap[start:end].select(num_cols).to_numpy())
mat_batch = np.nan_to_num(mat_batch, nan=0.0)
coords_list.append(reducer.transform(mat_batch))
coords = np.vstack(coords_list)
else:
mat = np.nan_to_num(StandardScaler().fit_transform(
df_umap.select(num_cols).to_numpy()), nan=0.0)
reducer = umap.UMAP(n_components=attempt_n,
random_state=RANDOM_SEED,
n_neighbors=15, min_dist=0.1,
metric='euclidean')
coords = reducer.fit_transform(mat)
umap_components = attempt_n
break
except Exception as e_umap:
if attempt_n == 2:
raise
logger.warning(f'UMAP 3D failed ({e_umap}), falling back to 2D')
if coords is None:
raise Exception('UMAP produced no output')
cluster_entry = store.get_cluster_result(cluster_id)
labels_list = cluster_entry.get('labels', []) if cluster_entry else []
from analysis.models import EntityProfile as EP
# Update existing EP records with UMAP coordinates
for i in range(len(df_umap)):
ev = f'row_{i}'
has_z = umap_components >= 3 and coords.shape[1] >= 3
update_fields = {
'embedding_x': float(coords[i, 0]) if i < len(coords) else None,
'embedding_y': float(coords[i, 1]) if i < len(coords) else None,
}
if has_z:
update_fields['embedding_z'] = float(coords[i, 2]) if i < len(coords) else None
EP.objects.filter(run=run, entity_value=ev).update(**update_fields)
run.entity_count = len(df_umap)
run.save(update_fields=['entity_count'])
except Exception as umap_err:
logger.warning(f'UMAP embedding skipped: {umap_err}')
# ── Done ──
run.status = 'completed'
run.progress_pct = 100
run.progress_msg = '分析完成'
run.save(update_fields=['status', 'progress_pct', 'progress_msg', 'entity_count'])
except Exception:
tb = traceback.format_exc()
logger.error(tb)
try:
run.status = 'failed'
run.error_message = tb
run.progress_msg = f'失败: {tb}'
run.run_log += f'\n[ERROR] {tb}'
run.save(update_fields=['status', 'error_message', 'progress_msg', 'run_log'])
except Exception as e:
logger.warning('save failed after pipeline error: %s', e)
"""
Delegates to the service layer.
"""
"""Clustering pipeline on raw rows. Delegates to service layer."""
from analysis.services.clustering import run_clustering_pipeline
run_clustering_pipeline(run, store, feature_columns, algorithm,
run_clustering_pipeline(run, store, ds_id, feature_columns, algorithm,
min_cluster_size, run_umap=run_umap, head=head)