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

refactor: clean views/clustering.py and split data_loader.py into package

Browse Source 
Task A: Replace _run_clustering_pipeline (533L->228L) — remove dead
docstring code, delegate to analysis.services.clustering.

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

Verified: manage.py check (0 issues), run_pipeline (3 clusters, OK)
This commit is contained in:
PM-pinou
2026-07-24 13:52:07 +08:00
parent 6cbbc24f14
commit 9f06e40a41
7 changed files with 803 additions and 1050 deletions
Download Patch File Download Diff File Expand all files Collapse all files
analysis/data_loader.py
+15 -743
View File
@@ -1,746 +1,18 @@
"""CSV loading with per-file schema validation, BOM detection, and multi-file merge. """Thin re-export wrapper — implementation moved to analysis/data_loader/ package.
Supports multiple encodings and delimiter-separated files. The main entry 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
analysis/data_loader/__init__.py
+30
View File
@@ -0,0 +1,30 @@
"""CSV data loading, schema handling, cleaning, and SQLite persistence.
Package layout:
_csv.py — CSV scanning, BOM/encoding detection, :func:`load_csv_directory`
_schema.py — Schema merging and cross-file validation
_clean.py — Data cleaning helpers (:func:`_coerce_to_float`)
_sqlite.py — SQLite persistence (save/load/drop tables)
"""
from ._csv import (
SUPPORTED_ENCODINGS,
detect_bom,
collect_schema,
detect_schema_dtypes,
load_config,
_file_count,
_resolve_encoding,
load_csv_directory,
)
from ._schema import _merge_schema_entries, validate_schemas
from ._clean import _coerce_to_float
from ._sqlite import (
_dtype_to_sqlite,
_str_to_polars_dtype,
_rows_to_csv_df,
save_to_db,
load_from_db,
load_from_db_lazy,
drop_sqlite_table,
)
analysis/data_loader/_clean.py
+32
View File
@@ -0,0 +1,32 @@
"""Data cleaning helpers: numeric coercion for string columns."""
import polars as pl
def _coerce_to_float(series: pl.Series) -> pl.Series:
"""Convert string series to Float64.
Every entry that doesn't parse as a valid float becomes null.
Handles blank strings, standalone '+'/'-', 'N/A', 'null', 'None', etc.
"""
series = series.cast(pl.Utf8).str.strip_chars()
artifacts = ('', '+', '-', '.', '..', 'N/A', 'n/a', 'NA', 'null', 'None', 'NULL', '?', '---', '--', '/')
is_artifact = series.is_in(pl.Series(artifacts).implode())
result = []
for v, bad in zip(series.to_list(), is_artifact.to_list()):
if v is None or bad:
result.append(None)
else:
# Check for "+"/"-" prefix artifacts (e.g., "+abc", "-xyz")
# where the value isn't a genuine numeric string. Values like
# "+5" or "-3.14" are still accepted as valid floats below.
if len(v) > 1 and v[0] in '+-':
try:
float(v[1:]) # validate suffix minus prefix
except ValueError:
result.append(None)
continue
try:
result.append(float(v))
except (ValueError, TypeError):
result.append(None)
return pl.Series(result, dtype=pl.Float64)
analysis/data_loader/_csv.py
+375
View File
@@ -0,0 +1,375 @@
"""CSV loading: glob, BOM detection, encoding, schema collection, multi-file merge.
Main entry: :func:`load_csv_directory`.
"""
import glob
import logging
import os
from pathlib import Path
from typing import Optional
import polars as pl
import yaml
from analysis.types import classify_schema, TYPE_REGISTRY
from ._clean import _coerce_to_float
from ._schema import _merge_schema_entries, validate_schemas
logger = logging.getLogger(__name__)
SUPPORTED_ENCODINGS = ['utf-8', 'utf-8-sig', 'utf-16le', 'utf-16be', 'latin-1']
# ── BOM detection ───────────────────────────────────────────────────────
def detect_bom(path: str) -> str:
"""Read the first 4 bytes of *path* and return the detected encoding.
Returns ``'utf-8'`` when no BOM is found.
Note: Python 3 ``open()`` fully supports Unicode/Chinese paths on all
platforms (no encoding dance needed).
"""
with open(path, 'rb') as f:
header = f.read(4)
if header[:3] == b'\xef\xbb\xbf':
return 'utf-8-sig'
if header[:2] == b'\xff\xfe':
return 'utf-16le'
if header[:2] == b'\xfe\xff':
return 'utf-16be'
return 'utf-8'
# ── schema utilities ────────────────────────────────────────────────────
def collect_schema(
path: str,
encoding: str = 'utf-8',
delimiter: str = ',',
sample_rows: int = 10000,
) -> dict:
"""Return {column_name: dtype_string} for a single CSV file.
Uses :func:`pl.scan_csv` with ``infer_schema_length=0`` to skip type
inference. All columns are forced to ``Utf8`` for cross-file compatibility,
avoiding type-inference failures when one file's column is ``Int64`` and
another's is ``String``.
"""
lf = pl.scan_csv(
str(path), # str() ensures Unicode/Chinese paths resolve correctly
encoding=encoding,
separator=delimiter,
infer_schema_length=0,
)
schema = lf.collect_schema()
result = {}
for name, dtype in zip(schema.names(), schema.dtypes()):
result[name] = 'Utf8' # force Utf8 for cross-file compatibility
return result
# ── dtype inference from schema ─────────────────────────────────────────
def detect_schema_dtypes(
lf: pl.LazyFrame,
sample_rows: int = 10000,
) -> dict:
"""Infer Polars dtype strings from a LazyFrame's schema.
The schema is read without collecting the full data (the scan scans
a small sample).
"""
schema = lf.collect_schema()
return {name: str(dtype) for name, dtype in zip(schema.names(), schema.dtypes())}
# ── config loading ──────────────────────────────────────────────────────
def load_config(config_path: Optional[str]) -> dict:
"""Load a YAML config file; returns an empty dict if *config_path* is None."""
if config_path is None:
return {}
path = Path(config_path)
if not path.exists():
raise FileNotFoundError(f"Config file not found: {config_path}")
with open(path, 'r', encoding='utf-8') as f:
return yaml.safe_load(f) or {}
# ── helpers ─────────────────────────────────────────────────────────────
def _file_count(paths: list[str]) -> int:
return len(paths)
def _resolve_encoding(encoding: str, file_encoding: str) -> str:
"""BOM-detected encoding overrides user-supplied encoding.
Polars >= 1.0 requires ``'utf8'`` (not ``'utf-8'``).
"""
resolved = file_encoding if file_encoding != 'utf-8' else encoding
return resolved.replace('-', '') if resolved in ('utf-8', 'utf-8-sig') else resolved
# ── main entry ──────────────────────────────────────────────────────────
def load_csv_directory(
glob_pattern: str,
encoding: str = 'utf-8',
delimiter: str = ',',
config_path: Optional[str] = None,
schema_override: Optional[dict] = None,
sample_rows: int = 10000,
schema_strict: bool = False,
recursive: bool = False,
head: Optional[int] = None,
) -> tuple[pl.LazyFrame, dict, int, int, float]:
"""Load, validate, and merge CSV files matching *glob_pattern*.
Args:
glob_pattern: Glob pattern for CSV files (e.g. ``data/*.csv``).
encoding: Fallback encoding when no BOM is detected.
delimiter: Column delimiter character.
config_path: Optional path to a YAML config (merged into metadata).
schema_override: If given, force these {column: dtype} overrides.
sample_rows: Rows to scan for dtype inference per file.
schema_strict: When True, raise ValueError on column mismatch across
files. When False (default), log warnings and merge with
``diagonal_relaxed`` (missing columns get null).
recursive: When True, use ``glob.glob(..., recursive=True)`` to
support ``**/*.csv`` patterns.
head: Optional row limit for preview scenarios. When set, applies
``lf.head(head)`` before returning.
Returns:
A tuple of ``(lazyframe, schema_dict, total_row_count, file_count,
memory_estimate_mb)``.
Raises:
FileNotFoundError: If no files match the glob pattern.
ValueError: If schemas are inconsistent across files and
*schema_strict* is True.
"""
# Normalise path separators for cross-platform Unicode support
glob_pattern = os.path.normpath(glob_pattern)
# Resolve file list
paths = sorted(glob.glob(glob_pattern, recursive=recursive))
if not paths and not recursive:
# Fallback: try recursive pattern
paths = sorted(glob.glob(glob_pattern, recursive=True))
if not paths:
raise FileNotFoundError(f"No files match glob pattern: {glob_pattern}")
# Load config (user-provided or fall back to project default)
config = load_config(config_path)
if not config:
default_path = Path(__file__).parent.parent.parent / 'config' / 'config.yaml'
if default_path.exists():
config = load_config(str(default_path))
# Collect per-file schema
schemas: list[dict] = []
lazyframes: list[pl.LazyFrame] = []
for path in paths:
file_encoding = detect_bom(path)
actual_encoding = _resolve_encoding(encoding, file_encoding)
lf = pl.scan_csv(
str(path), # str() ensures Unicode/Chinese paths resolve correctly
encoding=actual_encoding,
separator=delimiter,
infer_schema_length=max(sample_rows, 1000),
)
schema = detect_schema_dtypes(lf, sample_rows=sample_rows)
# Apply schema override for this file
if schema_override:
for col, dtype in schema_override.items():
if col in schema:
schema[col] = dtype
schemas.append(schema)
lazyframes.append(lf)
if schema_strict:
# Raise on column mismatch
validate_schemas(schemas)
else:
# Lenient mode: log warnings for column mismatches
col_sets = [set(s.keys()) for s in schemas]
reference = col_sets[0]
for i, cs in enumerate(col_sets[1:], start=1):
only_in_ref = reference - cs
only_in_current = cs - reference
if only_in_ref:
logger.warning(
"File #%d missing columns: %s (will be null)",
i + 1, sorted(only_in_ref),
)
if only_in_current:
logger.warning(
"File #%d has extra columns: %s (will be null in other files)",
i + 1, sorted(only_in_current),
)
# Merge schemas (unify column ordering — first file wins order)
merged_schema, _ = _merge_schema_entries(schemas)
# Concatenate all LazyFrames (union vertically)
merged_lf: pl.LazyFrame = pl.concat(lazyframes, how='diagonal_relaxed')
# Log schema and size estimate
total_row_count_est = len(paths) * (lazyframes[0].collect(streaming=True).height if lazyframes else 0)
logger.info(f'[LOAD] files={len(paths)} total_rows_est={total_row_count_est}')
for col_name, dtype in sorted(merged_schema.items()):
logger.info(f'[LOAD_SCHEMA] col={col_name} dtype={dtype}')
# Apply schema override on merged frame (cast columns)
if schema_override:
casts = {}
for col, dtype_str in schema_override.items():
if col in merged_schema:
try:
polars_dtype = getattr(pl, dtype_str.split('(')[0], None)
if polars_dtype:
casts[col] = pl.col(col).cast(polars_dtype)
except Exception:
pass # skip uncastable overrides silently
if casts:
merged_lf = merged_lf.with_columns(list(casts.values()))
# ── Cleaning: drop outdated columns ──────────────────────────────────
geoip_cfg = config.get('data', {}).get('geoip', {})
drop_cols = geoip_cfg.get(
'drop_columns',
[], # Empty by default — lat/lon columns are now cleaned, not dropped
)
existing_drop = [c for c in drop_cols if c in merged_schema]
if existing_drop:
merged_lf = merged_lf.drop(existing_drop)
merged_schema = {k: v for k, v in merged_schema.items() if k not in existing_drop}
logger.info('[CLEAN] dropped columns: %s', existing_drop)
# ── Cleaning: BOOL normalisation ─────────────────────────────────────
try:
type_map = classify_schema(
merged_lf.head(1000),
config_overrides=config.get('columns', {}),
)
except Exception:
type_map = {}
bool_cols = [
c for c, t in type_map.items()
if t.name == "布尔" and c in merged_schema
]
if bool_cols:
schema_dtypes = dict(zip(
merged_lf.collect_schema().names(),
merged_lf.collect_schema().dtypes(),
))
bool_exprs = []
cleaned_names: list[str] = []
for c in bool_cols:
if schema_dtypes.get(c) not in (pl.Utf8, pl.String):
continue
c_str = pl.col(c).cast(pl.Utf8).str.strip_chars()
cleaned = (
pl.when(c_str == "")
.then(pl.lit("False"))
.when(c_str.str.to_lowercase() == "+")
.then(pl.lit("True"))
.when(c_str.str.to_lowercase() == "-")
.then(pl.lit("False"))
.otherwise(c_str)
)
bool_exprs.append(cleaned.alias(c))
cleaned_names.append(c)
if bool_exprs:
merged_lf = merged_lf.with_columns(bool_exprs)
for c in cleaned_names:
merged_schema[c] = 'Utf8'
logger.info('[CLEAN] BOOL columns normalised: %s', cleaned_names)
# ── Generic numeric coercion ────────────────────────────────────────
# For remaining string columns (not IPv4/BYTES/ENUM/BOOL etc.),
# try to coerce to Float64.
_SKIP_NUM_COERCE = {"IPv4", "字节", "枚举", "布尔"}
_SKIP_NAME_PATTERNS = ('mac', 'mac_', '_mac', 'addr', 'address')
schema_dtypes_gn = dict(zip(
merged_lf.collect_schema().names(),
merged_lf.collect_schema().dtypes(),
))
generic_numeric_cols = []
for c, t in type_map.items():
if t.name in _SKIP_NUM_COERCE:
continue
if c not in merged_schema:
continue
dtype = schema_dtypes_gn.get(c)
if dtype not in (pl.Utf8, pl.String):
continue
cl = c.lower().replace('-', '_').replace(' ', '_')
if any(p in cl for p in _SKIP_NAME_PATTERNS):
continue
generic_numeric_cols.append(c)
if generic_numeric_cols:
num_exprs = []
for c in generic_numeric_cols:
num_exprs.append(
pl.col(c).map_batches(_coerce_to_float, return_dtype=pl.Float64).alias(c)
)
if num_exprs:
merged_lf = merged_lf.with_columns(num_exprs)
for c in generic_numeric_cols:
merged_schema[c] = 'Float64'
logger.info('[CLEAN] Generic numeric coercion applied: %s', generic_numeric_cols)
# ── Cleaning: BYTES (hex) preprocessing ─────────────────────────────
hex_cols = [
c for c, t in type_map.items()
if t.name == "字节" and c in merged_schema
]
if hex_cols:
hex_exprs = []
for c in hex_cols:
hex_exprs.append(
pl.col(c).str.split(" ").list.len().alias(f"{c}_hex_pairs_count")
)
merged_lf = merged_lf.with_columns(hex_exprs)
for c in hex_cols:
merged_schema[f"{c}_hex_pairs_count"] = 'UInt32'
logger.info('[CLEAN] BYTES columns added hex_pairs_count: %s', hex_cols)
# Row count estimate (fast path: use first file row count × files)
# For accurate count we'd need to collect, but that defeats lazy.
# We estimate from the first file's scanned schema info.
first_rows = lazyframes[0].collect(streaming=True).height if lazyframes else 0
total_row_count = first_rows * len(paths)
# Memory estimate
bytes_per_row = 0
for col_name, dtype_str in merged_schema.items():
numeric_part = dtype_str.split('(')[0].lower()
size_map = {
'int8': 1, 'int16': 2, 'int32': 4, 'int64': 8,
'uint8': 1, 'uint16': 2, 'uint32': 4, 'uint64': 8,
'float32': 4, 'float64': 8,
'boolean': 1, 'bool': 1,
'utf8': 50, 'string': 50, 'str': 50,
'categorical': 8, 'cat': 8,
'date': 4, 'datetime': 8, 'time': 8, 'duration': 8,
}
bytes_per_row += size_map.get(numeric_part, 8)
memory_mb = (bytes_per_row * total_row_count) / (1024 * 1024)
# Apply head limit for preview (keeps LazyFrame lazy — rows are limited on collect)
if head is not None:
merged_lf = merged_lf.head(head)
total_row_count = min(total_row_count, head)
return merged_lf, merged_schema, total_row_count, len(paths), memory_mb
analysis/data_loader/_schema.py
+57
View File
@@ -0,0 +1,57 @@
"""Schema detection, validation, and merging across CSV files."""
import logging
logger = logging.getLogger(__name__)
def _merge_schema_entries(
schemas: list[dict],
) -> tuple[dict[str, str], list[str]]:
"""Merge multiple schema dicts, raising on type conflicts.
Returns a single {name: dtype} dict. If a column appears in multiple
files with different dtypes the *last* non-null dtype wins (with a
warning embedded in the result).
"""
merged: dict[str, str] = {}
conflicts: list[str] = []
for s in schemas:
for name, dtype in s.items():
if name in merged and merged[name] != dtype:
conflicts.append(name)
merged[name] = dtype # last wins
return merged, conflicts
def validate_schemas(schemas: list[dict], strict: bool = True) -> None:
"""Compare schemas across files and raise or warn on column mismatch.
Checks that every file has the same set of column names. Type
differences across files are **not** considered an error (Polars can
cast), but differences in *column sets* are.
Args:
schemas: List of ``{column_name: dtype_string}`` dicts, one per file.
strict: If ``True`` (default), raise :class:`ValueError` on any
column mismatch. If ``False``, log a warning instead.
"""
if not schemas:
return
col_sets = [set(s.keys()) for s in schemas]
reference = col_sets[0]
for i, cs in enumerate(col_sets[1:], start=1):
only_in_ref = reference - cs
only_in_current = cs - reference
if only_in_ref or only_in_current:
msg_parts = []
if only_in_ref:
msg_parts.append(f"File #{i} missing columns: {sorted(only_in_ref)}")
if only_in_current:
msg_parts.append(f"File #{i} has extra columns: {sorted(only_in_current)}")
full_msg = f"Schema mismatch: {'; '.join(msg_parts)}"
if strict:
raise ValueError(full_msg)
else:
logger.warning(full_msg)
analysis/data_loader/_sqlite.py
+292
View File
@@ -0,0 +1,292 @@
"""SQLite persistence layer for LazyFrame data."""
import csv as _csv_mod
import io
import logging
import sqlite3
from typing import Optional
import polars as pl
logger = logging.getLogger(__name__)
# ── dtype mapping ─────────────────────────────────────────────────────────
_STR_TO_DTYPE = {
'Float64': pl.Float64, 'Float32': pl.Float32,
'Int64': pl.Int64, 'Int32': pl.Int32, 'Int16': pl.Int16, 'Int8': pl.Int8,
'UInt64': pl.UInt64, 'UInt32': pl.UInt32, 'UInt16': pl.UInt16, 'UInt8': pl.UInt8,
'Utf8': pl.Utf8, 'String': pl.Utf8, 'Categorical': pl.Categorical,
'Boolean': pl.Boolean, 'Bool': pl.Boolean,
'Date': pl.Date, 'Datetime': pl.Datetime, 'Time': pl.Time,
}
def _dtype_to_sqlite(dtype: pl.DataType) -> str:
"""Map a Polars dtype to a SQLite column type affinity."""
if dtype in (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64):
return 'INTEGER'
if dtype in (pl.Float32, pl.Float64):
return 'REAL'
if dtype == pl.Boolean:
return 'INTEGER'
return 'TEXT' # String, Date, Datetime, etc. all map to TEXT
def _str_to_polars_dtype(s: str) -> pl.DataType:
"""Convert a Polars dtype string (e.g. ``'Float64'``, ``'Int32'``) to a ``pl.DataType``.
Strips trailing metadata like ``Datetime(time_unit='us', time_zone=None)``.
"""
base = s.split('(')[0].strip()
return _STR_TO_DTYPE.get(base, pl.Utf8)
def _rows_to_csv_df(rows: list[sqlite3.Row]) -> pl.DataFrame:
"""Convert a batch of sqlite3.Row objects to a DataFrame via CSV fallback.
Loads all columns as Utf8 to handle mixed-type SQLite columns.
"""
buf = io.StringIO()
writer = _csv_mod.writer(buf)
writer.writerow(rows[0].keys())
for r in rows:
writer.writerow(str(v) if v is not None else '' for v in r)
buf.seek(0)
return pl.read_csv(buf, infer_schema_length=0)
# ── public API ───────────────────────────────────────────────────────────
def save_to_db(run_id: int, lf: pl.LazyFrame, mode: str = 'replace') -> str:
"""Write a LazyFrame to a SQLite table ``_data_{run_id}``.
Collects the LazyFrame (streaming if possible), creates or replaces the
table, bulk-inserts rows in chunks, and returns the table name.
Args:
run_id: Analysis run ID (used as ``_data_{run_id}``).
lf: LazyFrame to persist.
mode: ``'replace'`` (default) → drop+create+insert.
``'append'`` → skip DDL, only INSERT (table must exist).
This is a *persistence copy* — the original CSV files remain untouched.
"""
from django.conf import settings
table_name = f"_data_{run_id}"
db_path = settings.DATABASES['default']['NAME']
df = lf.collect(streaming=True)
conn = sqlite3.connect(db_path)
try:
if mode == 'replace':
# Drop existing table for retry/idempotency safety
conn.execute(f'DROP TABLE IF EXISTS "{table_name}"')
# Build CREATE TABLE from Polars schema
col_defs = [f'"{name}" {_dtype_to_sqlite(dtype)}'
for name, dtype in zip(df.columns, df.dtypes)]
conn.execute(f'CREATE TABLE "{table_name}" ({", ".join(col_defs)})')
# Bulk insert in chunks
col_names = [f'"{c}"' for c in df.columns]
placeholders = ','.join(['?' for _ in df.columns])
insert_sql = f'INSERT INTO "{table_name}" ({", ".join(col_names)}) VALUES ({placeholders})'
total_rows = len(df)
chunk_size = 10000
for i in range(0, total_rows, chunk_size):
chunk = df.slice(i, chunk_size)
conn.executemany(insert_sql, chunk.rows())
conn.commit()
finally:
conn.close()
logger.info('[SAVE_TO_DB] table=%s rows=%d mode=%s', table_name, total_rows, mode)
return table_name
def load_from_db(
table_name: str,
schema_overrides: dict[str, pl.DataType | str] | None = None,
) -> Optional[pl.LazyFrame]:
"""Load a SQLite table into a LazyFrame.
Args:
table_name: SQLite table name.
schema_overrides: Optional dict of ``{column_name: dtype}`` to cast
columns after loading. Values can be ``pl.DataType`` instances
(e.g. ``pl.Float64``) or dtype strings (e.g. ``'Float64'``).
Use when SQLite ``TEXT`` columns need to be restored to
``Float64`` / ``Int64`` etc.
Returns ``None`` if the table does not exist.
"""
from django.conf import settings
db_path = settings.DATABASES['default']['NAME']
conn = sqlite3.connect(db_path)
try:
# Check table exists
cursor = conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name=?",
(table_name,),
)
if cursor.fetchone() is None:
return None
conn.row_factory = sqlite3.Row
cursor = conn.execute(f'SELECT * FROM "{table_name}"')
# Peek first row to detect if table is empty
first_row = cursor.fetchone()
if first_row is None:
return pl.DataFrame().lazy()
batch_rows = [first_row]
dfs = []
_use_csv_fallback = False
while True:
chunk = cursor.fetchmany(10000)
if not chunk:
# Flush remaining rows
if batch_rows:
if _use_csv_fallback:
dfs.append(_rows_to_csv_df(batch_rows))
else:
try:
dfs.append(pl.from_dicts([dict(r) for r in batch_rows]))
except Exception:
_use_csv_fallback = True
dfs.append(_rows_to_csv_df(batch_rows))
break
batch_rows.extend(chunk)
if len(batch_rows) >= 10000:
if _use_csv_fallback:
dfs.append(_rows_to_csv_df(batch_rows))
else:
try:
dfs.append(pl.from_dicts([dict(r) for r in batch_rows]))
except Exception:
_use_csv_fallback = True
# Re-process current batch with CSV fallback
dfs.append(_rows_to_csv_df(batch_rows))
batch_rows = []
# Concat all batches
if len(dfs) == 1:
df = dfs[0]
else:
df = pl.concat(dfs, how='vertical')
# Apply type overrides to restore original column types
if schema_overrides:
casts = {}
for col, dtype in schema_overrides.items():
if col not in df.columns:
continue
dtype_obj = dtype if isinstance(dtype, pl.DataType) else _str_to_polars_dtype(dtype)
if df[col].dtype != dtype_obj:
try:
casts[col] = pl.col(col).cast(dtype_obj)
except Exception:
# If a cast fails (e.g. "A" → Int64), keep the column as-is
pass
if casts:
df = df.with_columns(list(casts.values()))
return df.lazy()
finally:
conn.close()
def load_from_db_lazy(
table_name: str,
schema_overrides: dict[str, pl.DataType | str] | None = None,
) -> Optional[pl.LazyFrame]:
"""Load a SQLite table into a LazyFrame using Polars native DB reader.
Uses ``pl.read_database`` for efficient bulk reading — an order of
magnitude faster than the manual batch-processing path in
:func:`load_from_db`. Returns a :class:`pl.LazyFrame` so downstream
operations remain lazy (query-plan, not materialised).
Args:
table_name: SQLite table name.
schema_overrides: Optional ``{column_name: dtype}`` dict to cast
columns after loading. Values can be ``pl.DataType`` instances
(e.g. ``pl.Float64``) or dtype strings (e.g. ``'Float64'``).
Returns:
``None`` if the table does not exist or is empty.
"""
from django.conf import settings
db_path = settings.DATABASES['default']['NAME']
conn = sqlite3.connect(db_path)
try:
# Check table existence (fast metadata query)
cursor = conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name=?",
(table_name,),
)
if cursor.fetchone() is None:
return None
# Bulk-read with Polars native engine — leverages the DBAPI2
# connection directly for efficient C-level data transfer.
query = f'SELECT * FROM "{table_name}"'
df = pl.read_database(connection=conn, query=query)
finally:
conn.close()
if df.height == 0:
return None
# Apply type overrides to restore original column types
# (SQLite stores everything as TEXT/INTEGER/REAL/BLOB;
# we need to cast back to Float64 / Int64 / etc.)
if schema_overrides:
casts = {}
for col, dtype in schema_overrides.items():
if col not in df.columns:
continue
dtype_obj = (
dtype if isinstance(dtype, pl.DataType)
else _str_to_polars_dtype(dtype)
)
if df[col].dtype != dtype_obj:
try:
casts[col] = pl.col(col).cast(dtype_obj)
except Exception:
pass # uncastable column — keep as-is
if casts:
df = df.with_columns(list(casts.values()))
rows = df.height
cols = df.width
logger.info(
'[LOAD_FROM_DB_LAZY] table=%s rows=%d cols=%d',
table_name, rows, cols,
)
return df.lazy()
def drop_sqlite_table(table_name: str) -> None:
"""Drop a SQLite data table silently if it exists."""
from django.conf import settings
db_path = settings.DATABASES['default']['NAME']
conn = sqlite3.connect(db_path)
try:
conn.execute(f'DROP TABLE IF EXISTS "{table_name}"')
conn.commit()
finally:
conn.close()
analysis/views/clustering.py
+2 -307
View File
@@ -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)