diff --git a/analysis/data_loader.py b/analysis/data_loader.py index 59f2e40..858dc6a 100644 --- a/analysis/data_loader.py +++ b/analysis/data_loader.py @@ -1,746 +1,18 @@ -"""CSV loading with per-file schema validation, BOM detection, and multi-file merge. +"""Thin re-export wrapper — implementation moved to analysis/data_loader/ package. -Supports multiple encodings and delimiter-separated files. The main entry -point is :func:`load_csv_directory`, which globs files, detects BOM per file, -validates schema consistency, and returns a merged :class:`pl.LazyFrame`. +When both ``analysis/data_loader.py`` and ``analysis/data_loader/`` exist, +Python prefers the package directory, so this file is effectively dead code. +Kept for git history continuity — all real code lives in the package. + +If you need to restore the monolithic version, see git history. """ -import glob -import io -import logging -import os -import sqlite3 -from pathlib import Path -from typing import Optional -import polars as pl -import yaml - -from analysis.types import classify_schema, TYPE_REGISTRY - -logger = logging.getLogger(__name__) - - -SUPPORTED_ENCODINGS = ['utf-8', 'utf-8-sig', 'utf-16le', 'utf-16be', 'latin-1'] - - -# ── numeric coercion helper ────────────────────────────────────────────── - -def _coerce_to_float(series: pl.Series) -> pl.Series: - """Convert string series to Float64. - - Every entry that doesn't parse as a valid float becomes null. - Handles blank strings, standalone '+'/'-', 'N/A', 'null', 'None', etc. - """ - series = series.cast(pl.Utf8).str.strip_chars() - artifacts = ('', '+', '-', '.', '..', 'N/A', 'n/a', 'NA', 'null', 'None', 'NULL', '?', '---', '--', '/') - is_artifact = series.is_in(pl.Series(artifacts).implode()) - result = [] - for v, bad in zip(series.to_list(), is_artifact.to_list()): - if v is None or bad: - result.append(None) - else: - # Check for "+"/"-" prefix artifacts (e.g., "+abc", "-xyz") - # where the value isn't a genuine numeric string. Values like - # "+5" or "-3.14" are still accepted as valid floats below. - if len(v) > 1 and v[0] in '+-': - try: - float(v[1:]) # validate suffix minus prefix - except ValueError: - result.append(None) - continue - try: - result.append(float(v)) - except (ValueError, TypeError): - result.append(None) - return pl.Series(result, dtype=pl.Float64) - - -# ── BOM detection ─────────────────────────────────────────────────────── - -def detect_bom(path: str) -> str: - """Read the first 4 bytes of *path* and return the detected encoding. - - Returns ``'utf-8'`` when no BOM is found. - - Note: Python 3 ``open()`` fully supports Unicode/Chinese paths on all - platforms (no encoding dance needed). - """ - with open(path, 'rb') as f: - header = f.read(4) - - if header[:3] == b'\xef\xbb\xbf': - return 'utf-8-sig' - if header[:2] == b'\xff\xfe': - return 'utf-16le' - if header[:2] == b'\xfe\xff': - return 'utf-16be' - return 'utf-8' - - -# ── schema utilities ──────────────────────────────────────────────────── - -def collect_schema( - path: str, - encoding: str = 'utf-8', - delimiter: str = ',', - sample_rows: int = 10000, -) -> dict: - """Return {column_name: dtype_string} for a single CSV file. - - Uses :func:`pl.scan_csv` with ``infer_schema_length=0`` to skip type - inference. All columns are forced to ``Utf8`` for cross-file compatibility, - avoiding type-inference failures when one file's column is ``Int64`` and - another's is ``String``. - """ - lf = pl.scan_csv( - str(path), # str() ensures Unicode/Chinese paths resolve correctly - encoding=encoding, - separator=delimiter, - infer_schema_length=0, - ) - schema = lf.collect_schema() - result = {} - for name, dtype in zip(schema.names(), schema.dtypes()): - result[name] = 'Utf8' # force Utf8 for cross-file compatibility - return result - - -def _merge_schema_entries( - schemas: list[dict], -) -> dict: - """Merge multiple schema dicts, raising on type conflicts. - - Returns a single {name: dtype} dict. If a column appears in multiple - files with different dtypes the *last* non-null dtype wins (with a - warning embedded in the result). - """ - merged: dict[str, str] = {} - conflicts: list[str] = [] - for s in schemas: - for name, dtype in s.items(): - if name in merged and merged[name] != dtype: - conflicts.append(name) - merged[name] = dtype # last wins - return merged, conflicts - - -def validate_schemas(schemas: list[dict], strict: bool = True) -> None: - """Compare schemas across files and raise or warn on column mismatch. - - Checks that every file has the same set of column names. Type - differences across files are **not** considered an error (Polars can - cast), but differences in *column sets* are. - - Args: - schemas: List of ``{column_name: dtype_string}`` dicts, one per file. - strict: If ``True`` (default), raise :class:`ValueError` on any - column mismatch. If ``False``, log a warning instead. - """ - if not schemas: - return - - col_sets = [set(s.keys()) for s in schemas] - reference = col_sets[0] - - for i, cs in enumerate(col_sets[1:], start=1): - only_in_ref = reference - cs - only_in_current = cs - reference - if only_in_ref or only_in_current: - msg_parts = [] - if only_in_ref: - msg_parts.append(f"File #{i} missing columns: {sorted(only_in_ref)}") - if only_in_current: - msg_parts.append(f"File #{i} has extra columns: {sorted(only_in_current)}") - full_msg = f"Schema mismatch: {'; '.join(msg_parts)}" - if strict: - raise ValueError(full_msg) - else: - logger.warning(full_msg) - - -# ── dtype inference from schema ───────────────────────────────────────── - -def detect_schema_dtypes( - lf: pl.LazyFrame, - sample_rows: int = 10000, -) -> dict: - """Infer Polars dtype strings from a LazyFrame's schema. - - The schema is read without collecting the full data (the scan scans - a small sample). - """ - schema = lf.collect_schema() - return {name: str(dtype) for name, dtype in zip(schema.names(), schema.dtypes())} - - -# ── config loading ────────────────────────────────────────────────────── - -def load_config(config_path: Optional[str]) -> dict: - """Load a YAML config file; returns an empty dict if *config_path* is None.""" - if config_path is None: - return {} - path = Path(config_path) - if not path.exists(): - raise FileNotFoundError(f"Config file not found: {config_path}") - with open(path, 'r', encoding='utf-8') as f: - return yaml.safe_load(f) or {} - - -# ── SQLite persistence ────────────────────────────────────────────────── - -def _dtype_to_sqlite(dtype: pl.DataType) -> str: - """Map a Polars dtype to a SQLite column type affinity.""" - if dtype in (pl.Int8, pl.Int16, pl.Int32, pl.Int64, - pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64): - return 'INTEGER' - if dtype in (pl.Float32, pl.Float64): - return 'REAL' - if dtype == pl.Boolean: - return 'INTEGER' - return 'TEXT' # String, Date, Datetime, etc. all map to TEXT - - -def save_to_db(run_id: int, lf: pl.LazyFrame, mode: str = 'replace') -> str: - """Write a LazyFrame to a SQLite table ``_data_{run_id}``. - - Collects the LazyFrame (streaming if possible), creates or replaces the - table, bulk-inserts rows in chunks, and returns the table name. - - Args: - run_id: Analysis run ID (used as ``_data_{run_id}``). - lf: LazyFrame to persist. - mode: ``'replace'`` (default) → drop+create+insert. - ``'append'`` → skip DDL, only INSERT (table must exist). - - This is a *persistence copy* — the original CSV files remain untouched. - """ - from django.conf import settings - - table_name = f"_data_{run_id}" - db_path = settings.DATABASES['default']['NAME'] - - df = lf.collect(streaming=True) - conn = sqlite3.connect(db_path) - - try: - if mode == 'replace': - # Drop existing table for retry/idempotency safety - conn.execute(f'DROP TABLE IF EXISTS "{table_name}"') - - # Build CREATE TABLE from Polars schema - col_defs = [f'"{name}" {_dtype_to_sqlite(dtype)}' - for name, dtype in zip(df.columns, df.dtypes)] - conn.execute(f'CREATE TABLE "{table_name}" ({", ".join(col_defs)})') - - # Bulk insert in chunks - col_names = [f'"{c}"' for c in df.columns] - placeholders = ','.join(['?' for _ in df.columns]) - insert_sql = f'INSERT INTO "{table_name}" ({", ".join(col_names)}) VALUES ({placeholders})' - - total_rows = len(df) - chunk_size = 10000 - for i in range(0, total_rows, chunk_size): - chunk = df.slice(i, chunk_size) - conn.executemany(insert_sql, chunk.rows()) - - conn.commit() - finally: - conn.close() - - logger.info('[SAVE_TO_DB] table=%s rows=%d mode=%s', table_name, total_rows, mode) - return table_name - - -_STR_TO_DTYPE = { - 'Float64': pl.Float64, 'Float32': pl.Float32, - 'Int64': pl.Int64, 'Int32': pl.Int32, 'Int16': pl.Int16, 'Int8': pl.Int8, - 'UInt64': pl.UInt64, 'UInt32': pl.UInt32, 'UInt16': pl.UInt16, 'UInt8': pl.UInt8, - 'Utf8': pl.Utf8, 'String': pl.Utf8, 'Categorical': pl.Categorical, - 'Boolean': pl.Boolean, 'Bool': pl.Boolean, - 'Date': pl.Date, 'Datetime': pl.Datetime, 'Time': pl.Time, -} - - -def _str_to_polars_dtype(s: str) -> pl.DataType: - """Convert a Polars dtype string (e.g. ``'Float64'``, ``'Int32'``) to a ``pl.DataType``. - - Strips trailing metadata like ``Datetime(time_unit='us', time_zone=None)``. - """ - base = s.split('(')[0].strip() - return _STR_TO_DTYPE.get(base, pl.Utf8) - - -def _rows_to_csv_df(rows: list[sqlite3.Row]) -> pl.DataFrame: - """Convert a batch of sqlite3.Row objects to a DataFrame via CSV fallback. - - Loads all columns as Utf8 to handle mixed-type SQLite columns. - """ - import csv as _csv_mod - - buf = io.StringIO() - writer = _csv_mod.writer(buf) - writer.writerow(rows[0].keys()) - for r in rows: - writer.writerow(str(v) if v is not None else '' for v in r) - buf.seek(0) - return pl.read_csv(buf, infer_schema_length=0) - - -def load_from_db( - table_name: str, - schema_overrides: dict[str, pl.DataType | str] | None = None, -) -> Optional[pl.LazyFrame]: - """Load a SQLite table into a LazyFrame. - - Args: - table_name: SQLite table name. - schema_overrides: Optional dict of ``{column_name: dtype}`` to cast - columns after loading. Values can be ``pl.DataType`` instances - (e.g. ``pl.Float64``) or dtype strings (e.g. ``'Float64'``). - Use when SQLite ``TEXT`` columns need to be restored to - ``Float64`` / ``Int64`` etc. - - Returns ``None`` if the table does not exist. - """ - from django.conf import settings - - db_path = settings.DATABASES['default']['NAME'] - conn = sqlite3.connect(db_path) - - try: - # Check table exists - cursor = conn.execute( - "SELECT name FROM sqlite_master WHERE type='table' AND name=?", - (table_name,), - ) - if cursor.fetchone() is None: - return None - - conn.row_factory = sqlite3.Row - cursor = conn.execute(f'SELECT * FROM "{table_name}"') - - # Peek first row to detect if table is empty - first_row = cursor.fetchone() - if first_row is None: - return pl.DataFrame().lazy() - - batch_rows = [first_row] - dfs = [] - _use_csv_fallback = False - - while True: - chunk = cursor.fetchmany(10000) - if not chunk: - # Flush remaining rows - if batch_rows: - if _use_csv_fallback: - dfs.append(_rows_to_csv_df(batch_rows)) - else: - try: - dfs.append(pl.from_dicts([dict(r) for r in batch_rows])) - except Exception: - _use_csv_fallback = True - dfs.append(_rows_to_csv_df(batch_rows)) - break - - batch_rows.extend(chunk) - if len(batch_rows) >= 10000: - if _use_csv_fallback: - dfs.append(_rows_to_csv_df(batch_rows)) - else: - try: - dfs.append(pl.from_dicts([dict(r) for r in batch_rows])) - except Exception: - _use_csv_fallback = True - # Re-process current batch with CSV fallback - dfs.append(_rows_to_csv_df(batch_rows)) - batch_rows = [] - - # Concat all batches - if len(dfs) == 1: - df = dfs[0] - else: - df = pl.concat(dfs, how='vertical') - - # Apply type overrides to restore original column types - if schema_overrides: - casts = {} - for col, dtype in schema_overrides.items(): - if col not in df.columns: - continue - dtype_obj = dtype if isinstance(dtype, pl.DataType) else _str_to_polars_dtype(dtype) - if df[col].dtype != dtype_obj: - try: - casts[col] = pl.col(col).cast(dtype_obj) - except Exception: - # If a cast fails (e.g. "A" → Int64), keep the column as-is - pass - if casts: - df = df.with_columns(list(casts.values())) - - return df.lazy() - finally: - conn.close() - - -def load_from_db_lazy( - table_name: str, - schema_overrides: dict[str, pl.DataType | str] | None = None, -) -> Optional[pl.LazyFrame]: - """Load a SQLite table into a LazyFrame using Polars native DB reader. - - Uses ``pl.read_database`` for efficient bulk reading — an order of - magnitude faster than the manual batch-processing path in - :func:`load_from_db`. Returns a :class:`pl.LazyFrame` so downstream - operations remain lazy (query-plan, not materialised). - - Args: - table_name: SQLite table name. - schema_overrides: Optional ``{column_name: dtype}`` dict to cast - columns after loading. Values can be ``pl.DataType`` instances - (e.g. ``pl.Float64``) or dtype strings (e.g. ``'Float64'``). - - Returns: - ``None`` if the table does not exist or is empty. - """ - from django.conf import settings - - db_path = settings.DATABASES['default']['NAME'] - conn = sqlite3.connect(db_path) - - try: - # Check table existence (fast metadata query) - cursor = conn.execute( - "SELECT name FROM sqlite_master WHERE type='table' AND name=?", - (table_name,), - ) - if cursor.fetchone() is None: - return None - - # Bulk-read with Polars native engine — leverages the DBAPI2 - # connection directly for efficient C-level data transfer. - query = f'SELECT * FROM "{table_name}"' - df = pl.read_database(connection=conn, query=query) - finally: - conn.close() - - if df.height == 0: - return None - - # Apply type overrides to restore original column types - # (SQLite stores everything as TEXT/INTEGER/REAL/BLOB; - # we need to cast back to Float64 / Int64 / etc.) - if schema_overrides: - casts = {} - for col, dtype in schema_overrides.items(): - if col not in df.columns: - continue - dtype_obj = ( - dtype if isinstance(dtype, pl.DataType) - else _str_to_polars_dtype(dtype) - ) - if df[col].dtype != dtype_obj: - try: - casts[col] = pl.col(col).cast(dtype_obj) - except Exception: - pass # uncastable column — keep as-is - if casts: - df = df.with_columns(list(casts.values())) - - rows = df.height - cols = df.width - logger.info( - '[LOAD_FROM_DB_LAZY] table=%s rows=%d cols=%d', - table_name, rows, cols, - ) - return df.lazy() - - -def drop_sqlite_table(table_name: str) -> None: - """Drop a SQLite data table silently if it exists.""" - from django.conf import settings - - db_path = settings.DATABASES['default']['NAME'] - conn = sqlite3.connect(db_path) - try: - conn.execute(f'DROP TABLE IF EXISTS "{table_name}"') - conn.commit() - finally: - conn.close() - - -# ── helpers ───────────────────────────────────────────────────────────── - -def _file_count(paths: list[str]) -> int: - return len(paths) - - -def _resolve_encoding(encoding: str, file_encoding: str) -> str: - """BOM-detected encoding overrides user-supplied encoding. - - Polars >= 1.0 requires ``'utf8'`` (not ``'utf-8'``). - """ - resolved = file_encoding if file_encoding != 'utf-8' else encoding - return resolved.replace('-', '') if resolved in ('utf-8', 'utf-8-sig') else resolved - - -# ── main entry ────────────────────────────────────────────────────────── - -def load_csv_directory( - glob_pattern: str, - encoding: str = 'utf-8', - delimiter: str = ',', - config_path: Optional[str] = None, - schema_override: Optional[dict] = None, - sample_rows: int = 10000, - schema_strict: bool = False, - recursive: bool = False, - head: Optional[int] = None, -) -> tuple[pl.LazyFrame, dict, int, int, float]: - """Load, validate, and merge CSV files matching *glob_pattern*. - - Args: - glob_pattern: Glob pattern for CSV files (e.g. ``data/*.csv``). - encoding: Fallback encoding when no BOM is detected. - delimiter: Column delimiter character. - config_path: Optional path to a YAML config (merged into metadata). - schema_override: If given, force these {column: dtype} overrides. - sample_rows: Rows to scan for dtype inference per file. - schema_strict: When True, raise ValueError on column mismatch across - files. When False (default), log warnings and merge with - ``diagonal_relaxed`` (missing columns get null). - recursive: When True, use ``glob.glob(..., recursive=True)`` to - support ``**/*.csv`` patterns. - head: Optional row limit for preview scenarios. When set, applies - ``lf.head(head)`` before returning. - - Returns: - A tuple of ``(lazyframe, schema_dict, total_row_count, file_count, - memory_estimate_mb)``. - - Raises: - FileNotFoundError: If no files match the glob pattern. - ValueError: If schemas are inconsistent across files and - *schema_strict* is True. - """ - # Normalise path separators for cross-platform Unicode support - glob_pattern = os.path.normpath(glob_pattern) - - # Resolve file list - paths = sorted(glob.glob(glob_pattern, recursive=recursive)) - if not paths and not recursive: - # Fallback: try recursive pattern - paths = sorted(glob.glob(glob_pattern, recursive=True)) - if not paths: - raise FileNotFoundError(f"No files match glob pattern: {glob_pattern}") - - # Load config (user-provided or fall back to project default) - config = load_config(config_path) - if not config: - default_path = Path(__file__).parent.parent / 'config' / 'config.yaml' - if default_path.exists(): - config = load_config(str(default_path)) - - # Collect per-file schema - schemas: list[dict] = [] - lazyframes: list[pl.LazyFrame] = [] - for path in paths: - file_encoding = detect_bom(path) - actual_encoding = _resolve_encoding(encoding, file_encoding) - - lf = pl.scan_csv( - str(path), # str() ensures Unicode/Chinese paths resolve correctly - encoding=actual_encoding, - separator=delimiter, - infer_schema_length=max(sample_rows, 1000), - ) - schema = detect_schema_dtypes(lf, sample_rows=sample_rows) - - # Apply schema override for this file - if schema_override: - for col, dtype in schema_override.items(): - if col in schema: - schema[col] = dtype - - schemas.append(schema) - lazyframes.append(lf) - - if schema_strict: - # Raise on column mismatch - validate_schemas(schemas) - else: - # Lenient mode: log warnings for column mismatches - col_sets = [set(s.keys()) for s in schemas] - reference = col_sets[0] - for i, cs in enumerate(col_sets[1:], start=1): - only_in_ref = reference - cs - only_in_current = cs - reference - if only_in_ref: - logger.warning( - "File #%d missing columns: %s (will be null)", - i + 1, sorted(only_in_ref), - ) - if only_in_current: - logger.warning( - "File #%d has extra columns: %s (will be null in other files)", - i + 1, sorted(only_in_current), - ) - - # Merge schemas (unify column ordering — first file wins order) - merged_schema, _ = _merge_schema_entries(schemas) - - # Concatenate all LazyFrames (union vertically) - merged_lf: pl.LazyFrame = pl.concat(lazyframes, how='diagonal_relaxed') - - # Log schema and size estimate - total_row_count_est = len(paths) * (lazyframes[0].collect(streaming=True).height if lazyframes else 0) - logger.info(f'[LOAD] files={len(paths)} total_rows_est={total_row_count_est}') - for col_name, dtype in sorted(merged_schema.items()): - logger.info(f'[LOAD_SCHEMA] col={col_name} dtype={dtype}') - - # Apply schema override on merged frame (cast columns) - if schema_override: - casts = {} - for col, dtype_str in schema_override.items(): - if col in merged_schema: - try: - polars_dtype = getattr(pl, dtype_str.split('(')[0], None) - if polars_dtype: - casts[col] = pl.col(col).cast(polars_dtype) - except Exception: - pass # skip uncastable overrides silently - if casts: - merged_lf = merged_lf.with_columns(list(casts.values())) - - # ── Cleaning: drop outdated columns ────────────────────────────────── - geoip_cfg = config.get('data', {}).get('geoip', {}) - drop_cols = geoip_cfg.get( - 'drop_columns', - [], # Empty by default — lat/lon columns are now cleaned, not dropped - ) - existing_drop = [c for c in drop_cols if c in merged_schema] - if existing_drop: - merged_lf = merged_lf.drop(existing_drop) - merged_schema = {k: v for k, v in merged_schema.items() if k not in existing_drop} - logger.info('[CLEAN] dropped columns: %s', existing_drop) - - # ── Cleaning: BOOL normalisation ───────────────────────────────────── - try: - type_map = classify_schema( - merged_lf.head(1000), - config_overrides=config.get('columns', {}), - ) - except Exception: - type_map = {} - - bool_cols = [ - c for c, t in type_map.items() - if t.name == "布尔" and c in merged_schema - ] - if bool_cols: - schema_dtypes = dict(zip( - merged_lf.collect_schema().names(), - merged_lf.collect_schema().dtypes(), - )) - bool_exprs = [] - cleaned_names: list[str] = [] - for c in bool_cols: - if schema_dtypes.get(c) not in (pl.Utf8, pl.String): - continue - c_str = pl.col(c).cast(pl.Utf8).str.strip_chars() - cleaned = ( - pl.when(c_str == "") - .then(pl.lit("False")) - .when(c_str.str.to_lowercase() == "+") - .then(pl.lit("True")) - .when(c_str.str.to_lowercase() == "-") - .then(pl.lit("False")) - .otherwise(c_str) - ) - bool_exprs.append(cleaned.alias(c)) - cleaned_names.append(c) - if bool_exprs: - merged_lf = merged_lf.with_columns(bool_exprs) - for c in cleaned_names: - merged_schema[c] = 'Utf8' - logger.info('[CLEAN] BOOL columns normalised: %s', cleaned_names) - - # ── Generic numeric coercion ──────────────────────────────────────── - # For remaining string columns (not IPv4/BYTES/ENUM/BOOL etc.), - # try to coerce to Float64. - _SKIP_NUM_COERCE = {"IPv4", "字节", "枚举", "布尔"} - _SKIP_NAME_PATTERNS = ('mac', 'mac_', '_mac', 'addr', 'address') - schema_dtypes_gn = dict(zip( - merged_lf.collect_schema().names(), - merged_lf.collect_schema().dtypes(), - )) - generic_numeric_cols = [] - for c, t in type_map.items(): - if t.name in _SKIP_NUM_COERCE: - continue - if c not in merged_schema: - continue - dtype = schema_dtypes_gn.get(c) - if dtype not in (pl.Utf8, pl.String): - continue - cl = c.lower().replace('-', '_').replace(' ', '_') - if any(p in cl for p in _SKIP_NAME_PATTERNS): - continue - generic_numeric_cols.append(c) - - if generic_numeric_cols: - num_exprs = [] - for c in generic_numeric_cols: - num_exprs.append( - pl.col(c).map_batches(_coerce_to_float, return_dtype=pl.Float64).alias(c) - ) - if num_exprs: - merged_lf = merged_lf.with_columns(num_exprs) - for c in generic_numeric_cols: - merged_schema[c] = 'Float64' - logger.info('[CLEAN] Generic numeric coercion applied: %s', generic_numeric_cols) - - # ── Cleaning: BYTES (hex) preprocessing ───────────────────────────── - hex_cols = [ - c for c, t in type_map.items() - if t.name == "字节" and c in merged_schema - ] - if hex_cols: - hex_exprs = [] - for c in hex_cols: - hex_exprs.append( - pl.col(c).str.split(" ").list.len().alias(f"{c}_hex_pairs_count") - ) - merged_lf = merged_lf.with_columns(hex_exprs) - for c in hex_cols: - merged_schema[f"{c}_hex_pairs_count"] = 'UInt32' - logger.info('[CLEAN] BYTES columns added hex_pairs_count: %s', hex_cols) - - # Row count estimate (fast path: use first file row count × files) - # For accurate count we'd need to collect, but that defeats lazy. - # We estimate from the first file's scanned schema info. - first_rows = lazyframes[0].collect(streaming=True).height if lazyframes else 0 - total_row_count = first_rows * len(paths) - - # Memory estimate - bytes_per_row = 0 - for col_name, dtype_str in merged_schema.items(): - numeric_part = dtype_str.split('(')[0].lower() - size_map = { - 'int8': 1, 'int16': 2, 'int32': 4, 'int64': 8, - 'uint8': 1, 'uint16': 2, 'uint32': 4, 'uint64': 8, - 'float32': 4, 'float64': 8, - 'boolean': 1, 'bool': 1, - 'utf8': 50, 'string': 50, 'str': 50, - 'categorical': 8, 'cat': 8, - 'date': 4, 'datetime': 8, 'time': 8, 'duration': 8, - } - bytes_per_row += size_map.get(numeric_part, 8) - memory_mb = (bytes_per_row * total_row_count) / (1024 * 1024) - - # Apply head limit for preview (keeps LazyFrame lazy — rows are limited on collect) - if head is not None: - merged_lf = merged_lf.head(head) - total_row_count = min(total_row_count, head) - - return merged_lf, merged_schema, total_row_count, len(paths), memory_mb +from analysis.data_loader._csv import load_csv_directory +from analysis.data_loader._sqlite import ( + load_from_db, + load_from_db_lazy, + save_to_db, + drop_sqlite_table, +) +from analysis.data_loader._clean import _coerce_to_float +from analysis.data_loader._sqlite import _dtype_to_sqlite diff --git a/analysis/data_loader/__init__.py b/analysis/data_loader/__init__.py new file mode 100644 index 0000000..15ef9ff --- /dev/null +++ b/analysis/data_loader/__init__.py @@ -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, +) diff --git a/analysis/data_loader/_clean.py b/analysis/data_loader/_clean.py new file mode 100644 index 0000000..0fe6904 --- /dev/null +++ b/analysis/data_loader/_clean.py @@ -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) diff --git a/analysis/data_loader/_csv.py b/analysis/data_loader/_csv.py new file mode 100644 index 0000000..b7ca34c --- /dev/null +++ b/analysis/data_loader/_csv.py @@ -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 diff --git a/analysis/data_loader/_schema.py b/analysis/data_loader/_schema.py new file mode 100644 index 0000000..536bbe2 --- /dev/null +++ b/analysis/data_loader/_schema.py @@ -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) diff --git a/analysis/data_loader/_sqlite.py b/analysis/data_loader/_sqlite.py new file mode 100644 index 0000000..40ae823 --- /dev/null +++ b/analysis/data_loader/_sqlite.py @@ -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() diff --git a/analysis/views/clustering.py b/analysis/views/clustering.py index b4f6c47..f9c7361 100644 --- a/analysis/views/clustering.py +++ b/analysis/views/clustering.py @@ -222,312 +222,7 @@ def cluster_detail(request, display_id, cluster_label): @profile def _run_clustering_pipeline(run, store, ds_id, feature_columns, algorithm, min_cluster_size, run_umap=True, head=None): - """Clustering pipeline on raw rows: clustering → feature extraction → UMAP → ORM save. - - Operates directly on raw data rows (no entity aggregation). - Per-row results stored in EntityProfile with row index as identifier. - - Args: - run: AnalysisRun ORM object. - store: SessionStore instance. - ds_id: dataset ID in SessionStore pointing to raw data. - feature_columns: list of feature column names (None = auto-detect all numeric). - algorithm: clustering algorithm. - min_cluster_size: int for HDBSCAN. - run_umap: bool, whether to compute and save UMAP-2D embeddings. - head: optional int, downsample to at most this many rows (min 100). - """ - import warnings - warnings.filterwarnings('ignore', category=RuntimeWarning, module='sklearn') - - try: - from analysis.tool_registry import _handle_run_clustering, _handle_extract_features - import polars as pl - - entry = store.get_dataset(ds_id) - if entry is None: - run.status = 'failed' - run.error_message = '数据集在会话中不存在' - run.save(update_fields=['status', 'error_message']) - return - - schema = entry.get('schema', {}) - - # ── Downsampling ── - if head is not None and head > 0: - lf = entry['lazyframe'] - try: - row_count = lf.select(pl.len()).collect(streaming=True).item() - except Exception: - logger.error("unknown failed: {}".format(traceback.format_exc())) - row_count = 0 - - if row_count > head: - head = max(head, 100) # ensure minimum rows for clustering - run.progress_msg = f'正在采样 {head}/{row_count} 行...' - run.save(update_fields=['progress_msg']) - lf = lf.head(head) - store.store_dataset( - ds_id, lf, - schema=schema, - metadata=entry.get('metadata', {}), - ) - entry = store.get_dataset(ds_id) - row_count = head - - # Resolve feature columns — use actual LazyFrame dtypes (not stored schema dict - # which may be stale after Utf8 coercion in _background_process). - numeric_types = (pl.Int8, pl.Int16, pl.Int32, pl.Int64, - pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64, - pl.Float32, pl.Float64) - - lf = entry['lazyframe'] - live_schema = lf.collect_schema() - live_names = live_schema.names() - live_dtypes = list(live_schema.dtypes()) - - if feature_columns: - feature_cols = [ - c for c in feature_columns - if c in live_names and live_dtypes[live_names.index(c)] in numeric_types - ] - else: - feature_cols = [ - live_names[i] for i, dt in enumerate(live_dtypes) - if dt in numeric_types and not live_names[i].startswith('_') - ] - - # ── Step 1: Clustering ──────────────────────────────────────────────── - # If no numeric columns detected, pass empty list to let - # _handle_run_clustering apply its Utf8→Float64 coercion fallback. - run.status = 'clustering' - run.progress_pct = 70 - run.progress_msg = '正在进行聚类分析...' - run.save(update_fields=['status', 'progress_pct', 'progress_msg']) - - result = asyncio.run(_handle_run_clustering( - dataset_id=ds_id, - cluster_columns=feature_cols, - algorithm=algorithm, - params={'min_cluster_size': min_cluster_size}, - random_state=RANDOM_SEED, - )) - - if 'error' in result: - err = result['error'] - if 'numeric' in err.lower(): - available = [ - f"{k}({v})" for k, v in schema.items() - if v.split('(')[0].strip() in numeric_types and not k.startswith('_') - ] - err += f"。可用数值列: {available}" - raise Exception(err) - - cluster_id = result.get('cluster_result_id', '') - run.cluster_count = result.get('n_clusters', 0) - run.save(update_fields=['cluster_count']) - - # ── Step 2: Feature extraction ────────────────────────────────────── - run.status = 'extracting' - run.progress_pct = 90 - run.progress_msg = '正在提取聚类特征...' - run.save(update_fields=['status', 'progress_pct', 'progress_msg']) - - feat_result = asyncio.run(_handle_extract_features( - dataset_id=ds_id, - cluster_result_id=cluster_id, - top_k=10, method='zscore', save_to_db=True, - run_id=run.id, - )) - - if 'error' in feat_result: - raise Exception(feat_result['error']) - - # ── Step 2b: Cluster SVD feature extraction ────────────────────── - run.progress_msg = '正在进行聚类SVD特征提取...' - run.save(update_fields=['progress_msg']) - try: - from analysis.distance import cluster_svd_extract - lf = entry['lazyframe'] - cluster_entry = store.get_cluster_result(cluster_id) - labels_list = cluster_entry.get('labels', []) if cluster_entry else [] - - if labels_list and feature_cols: - svd_results = cluster_svd_extract(lf, labels_list, feature_cols) - - # Persist SVD features to ClusterFeature (method='cluster_svd') - from analysis.models import ClusterFeature as CF - from analysis.models import ClusterResult as CR - for label, svd_info in svd_results.items(): - cr = CR.objects.filter(run=run, cluster_label=label).first() - if cr is None: - continue - # Delete old zscore entries for this cluster, replace with SVD - CF.objects.filter(cluster=cr, distinguishing_method='zscore').delete() - for feat_name in svd_info.get('features', [])[:10]: - idx = svd_info['features'].index(feat_name) if feat_name in svd_info['features'] else -1 - strength = (svd_info.get('feature_strength', [])[idx] - if idx >= 0 and idx < len(svd_info.get('feature_strength', [])) - else None) - CF.objects.get_or_create( - cluster=cr, - feature_name=feat_name, - defaults={ - 'mean': None, 'std': None, 'median': None, - 'p25': None, 'p75': None, 'missing_rate': None, - 'distinguishing_score': float(strength) if strength is not None else None, - 'distinguishing_method': 'cluster_svd', - }, - ) - except Exception as svd_err: - logger.warning(f'Cluster SVD extraction skipped: {svd_err}') - - # ── Step 2c: Save per-row cluster labels to EntityProfile ── - # Always persist row→cluster mapping so cluster_detail can show data rows - # even when UMAP is skipped or fails. - try: - cluster_entry_save = store.get_cluster_result(cluster_id) - labels_list_save = cluster_entry_save.get('labels', []) if cluster_entry_save else [] - if labels_list_save: - from analysis.models import EntityProfile as EP_save - from analysis.models import ClusterResult as CR_save - # Delete stale EP records for this run first - EP_save.objects.filter(run=run).delete() - cr_cache_save = {} - row_profiles = [] - for r_idx in range(len(labels_list_save)): - lbl = labels_list_save[r_idx] - ck = f'{run.id}_{lbl}' - if ck not in cr_cache_save: - cr_cache_save[ck] = CR_save.objects.filter(run=run, cluster_label=lbl).first() - row_profiles.append(EP_save( - entity_value=f'row_{r_idx}', run=run, cluster_label=lbl, - cluster=cr_cache_save[ck], - embedding_x=None, embedding_y=None, embedding_z=None, - )) - if row_profiles: - EP_save.objects.bulk_create(row_profiles, ignore_conflicts=True, batch_size=1000) - run.entity_count = len(row_profiles) - run.save(update_fields=['entity_count']) - except Exception as save_err: - logger.warning(f'Row labels save skipped: {save_err}') - - # ── Step 3: UMAP-3D embedding (2D fallback) ─────────────────────── - if run_umap: - try: - lf = entry['lazyframe'] - # UMAP: sample max 10K for training, batch-transform rest in 1K batches - MAX_UMAP_TRAIN = UMAP_TRAIN_SAMPLE - BATCH_SIZE = UMAP_BATCH_SIZE - try: - n_total = lf.select(pl.len()).collect(streaming=True).item() - except Exception: - logger.error("unknown failed: {}".format(traceback.format_exc())) - n_total = 0 - from sklearn.preprocessing import StandardScaler - import umap - import numpy as np - - live_schema = lf.collect_schema() - num_cols = [name for name, dt in zip(live_schema.names(), live_schema.dtypes()) - if dt in (pl.Int8, pl.Int16, pl.Int32, pl.Int64, - pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64, - pl.Float32, pl.Float64) and not name.startswith('_')] - - if len(num_cols) >= 2: - if n_total > MAX_UMAP_TRAIN: - run.progress_msg = f'UMAP 采样 {MAX_UMAP_TRAIN}/{n_total} 行训练...' - run.save(update_fields=['progress_msg']) - # LazyFrame.sample() is unavailable in Polars 1.42 — collect then sample - df_sample = df_umap.sample(n=MAX_UMAP_TRAIN, seed=RANDOM_SEED) - df_sample = df_umap - - umap_components = 3 # try 3D first - coords = None - reducer = None - - if len(df_umap) > 2: - for attempt_n in (3, 2): - try: - if n_total > MAX_UMAP_TRAIN: - # Train UMAP on 10K sample, batch-transform full dataset - mat_sample = df_sample.select(num_cols).to_numpy() - scaler = StandardScaler().fit(mat_sample) - mat_sample_scaled = np.nan_to_num( - scaler.transform(mat_sample), nan=0.0) - reducer = umap.UMAP(n_components=attempt_n, - random_state=RANDOM_SEED, - n_neighbors=15, min_dist=0.1, - metric='euclidean') - reducer.fit(mat_sample_scaled) - # Batch transform in 1K-row batches - coords_list = [] - for start in range(0, len(df_umap), BATCH_SIZE): - end = min(start + BATCH_SIZE, len(df_umap)) - mat_batch = scaler.transform( - df_umap[start:end].select(num_cols).to_numpy()) - mat_batch = np.nan_to_num(mat_batch, nan=0.0) - coords_list.append(reducer.transform(mat_batch)) - coords = np.vstack(coords_list) - else: - mat = np.nan_to_num(StandardScaler().fit_transform( - df_umap.select(num_cols).to_numpy()), nan=0.0) - reducer = umap.UMAP(n_components=attempt_n, - random_state=RANDOM_SEED, - n_neighbors=15, min_dist=0.1, - metric='euclidean') - coords = reducer.fit_transform(mat) - umap_components = attempt_n - break - except Exception as e_umap: - if attempt_n == 2: - raise - logger.warning(f'UMAP 3D failed ({e_umap}), falling back to 2D') - - if coords is None: - raise Exception('UMAP produced no output') - - cluster_entry = store.get_cluster_result(cluster_id) - labels_list = cluster_entry.get('labels', []) if cluster_entry else [] - - from analysis.models import EntityProfile as EP - - # Update existing EP records with UMAP coordinates - for i in range(len(df_umap)): - ev = f'row_{i}' - has_z = umap_components >= 3 and coords.shape[1] >= 3 - update_fields = { - 'embedding_x': float(coords[i, 0]) if i < len(coords) else None, - 'embedding_y': float(coords[i, 1]) if i < len(coords) else None, - } - if has_z: - update_fields['embedding_z'] = float(coords[i, 2]) if i < len(coords) else None - EP.objects.filter(run=run, entity_value=ev).update(**update_fields) - run.entity_count = len(df_umap) - run.save(update_fields=['entity_count']) - except Exception as umap_err: - logger.warning(f'UMAP embedding skipped: {umap_err}') - - # ── Done ── - run.status = 'completed' - run.progress_pct = 100 - run.progress_msg = '分析完成' - run.save(update_fields=['status', 'progress_pct', 'progress_msg', 'entity_count']) - - except Exception: - tb = traceback.format_exc() - logger.error(tb) - try: - run.status = 'failed' - run.error_message = tb - run.progress_msg = f'失败: {tb}' - run.run_log += f'\n[ERROR] {tb}' - run.save(update_fields=['status', 'error_message', 'progress_msg', 'run_log']) - except Exception as e: - logger.warning('save failed after pipeline error: %s', e) - """ - Delegates to the service layer. - """ + """Clustering pipeline on raw rows. Delegates to service layer.""" from analysis.services.clustering import run_clustering_pipeline - run_clustering_pipeline(run, store, feature_columns, algorithm, + run_clustering_pipeline(run, store, ds_id, feature_columns, algorithm, min_cluster_size, run_umap=run_umap, head=head)