Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 9f06e40a41 | |||
| 6cbbc24f14 |
+15
-743
@@ -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
|
When both ``analysis/data_loader.py`` and ``analysis/data_loader/`` exist,
|
||||||
point is :func:`load_csv_directory`, which globs files, detects BOM per file,
|
Python prefers the package directory, so this file is effectively dead code.
|
||||||
validates schema consistency, and returns a merged :class:`pl.LazyFrame`.
|
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
|
from analysis.data_loader._csv import load_csv_directory
|
||||||
import yaml
|
from analysis.data_loader._sqlite import (
|
||||||
|
load_from_db,
|
||||||
from analysis.types import classify_schema, TYPE_REGISTRY
|
load_from_db_lazy,
|
||||||
|
save_to_db,
|
||||||
logger = logging.getLogger(__name__)
|
drop_sqlite_table,
|
||||||
|
)
|
||||||
|
from analysis.data_loader._clean import _coerce_to_float
|
||||||
SUPPORTED_ENCODINGS = ['utf-8', 'utf-8-sig', 'utf-16le', 'utf-16be', 'latin-1']
|
from analysis.data_loader._sqlite import _dtype_to_sqlite
|
||||||
|
|
||||||
|
|
||||||
# ── 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,
|
|
||||||
)
|
|
||||||
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
|
|
||||||
|
|||||||
@@ -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,
|
||||||
|
)
|
||||||
@@ -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)
|
||||||
@@ -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
|
||||||
@@ -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)
|
||||||
@@ -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()
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -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",
|
||||||
|
]
|
||||||
@@ -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
|
||||||
@@ -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
|
||||||
@@ -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()
|
||||||
@@ -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()
|
||||||
@@ -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
|
||||||
@@ -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)
|
||||||
@@ -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
|
||||||
@@ -222,312 +222,7 @@ def cluster_detail(request, display_id, cluster_label):
|
|||||||
@profile
|
@profile
|
||||||
def _run_clustering_pipeline(run, store, ds_id, feature_columns, algorithm,
|
def _run_clustering_pipeline(run, store, ds_id, feature_columns, algorithm,
|
||||||
min_cluster_size, run_umap=True, head=None):
|
min_cluster_size, run_umap=True, head=None):
|
||||||
"""Clustering pipeline on raw rows: clustering → feature extraction → UMAP → ORM save.
|
"""Clustering pipeline on raw rows. Delegates to service layer."""
|
||||||
|
|
||||||
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.
|
|
||||||
"""
|
|
||||||
from analysis.services.clustering import run_clustering_pipeline
|
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)
|
min_cluster_size, run_umap=run_umap, head=head)
|
||||||
|
|||||||
Reference in New Issue
Block a user