PM-pinou
279 lines
12 KiB
Python
279 lines
12 KiB
Python
"""Background pipeline: CSV loading, SVD, and pipeline worker."""
|
|
import os
|
|
import threading
|
|
import traceback
|
|
import logging
|
|
|
|
from analysis.appdata import ensure_appdata_dir
|
|
from analysis.constants import MAX_CLUSTER_FEATURES, RANDOM_SEED
|
|
from analysis.models import AnalysisRun
|
|
from analysis.profile_util import profile
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _run_pipeline_worker(run_id, analysis_fn, **ctx):
|
|
"""Unified background worker: setup / ORM / progress / error handling."""
|
|
import django
|
|
os.environ['DJANGO_SETTINGS_MODULE'] = 'tianxuan.settings'
|
|
os.environ['DJANGO_ALLOW_ASYNC_UNSAFE'] = 'true'
|
|
django.setup()
|
|
|
|
from analysis.db_utils import retry_on_lock
|
|
from django.db import close_old_connections
|
|
|
|
close_old_connections()
|
|
|
|
@retry_on_lock()
|
|
def _get_run():
|
|
return AnalysisRun.objects.get(id=run_id)
|
|
|
|
run = _get_run()
|
|
try:
|
|
analysis_fn(run, ctx)
|
|
except Exception:
|
|
tb = traceback.format_exc()
|
|
logger.error(tb)
|
|
run.status = 'failed'
|
|
run.error_message = tb
|
|
run.run_log += f'\n[FATAL] {tb}'
|
|
retry_on_lock()(lambda: run.save(update_fields=['status', 'error_message', 'run_log']))()
|
|
|
|
|
|
@profile
|
|
def _background_process(run_id, upload_dir):
|
|
"""Background: stream-load CSVs → save to SQLite → ready for analysis."""
|
|
import django
|
|
os.environ['DJANGO_SETTINGS_MODULE'] = 'tianxuan.settings'
|
|
django.setup()
|
|
from analysis.models import AnalysisRun
|
|
from analysis.session_store import SessionStore
|
|
|
|
run = AnalysisRun.objects.get(id=run_id)
|
|
store = SessionStore()
|
|
try:
|
|
# Check upload directory still exists (may have been temp-cleaned)
|
|
if not upload_dir.is_dir():
|
|
msg = '上传数据目录已不存在，请重新上传'
|
|
logger.error('[BACKGROUND] %s (run_id=%s, upload_dir=%s)', msg, run_id, upload_dir)
|
|
run.status = 'failed'
|
|
run.error_message = msg
|
|
run.save(update_fields=['status', 'error_message'])
|
|
return
|
|
|
|
import glob
|
|
import polars as pl
|
|
csv_paths = sorted(glob.glob(str(upload_dir / '*.csv')))
|
|
if not csv_paths:
|
|
msg = '没有找到 CSV 文件'
|
|
logger.error('[BACKGROUND] %s (run_id=%s)', msg, run_id)
|
|
run.status = 'failed'
|
|
run.error_message = msg
|
|
run.save(update_fields=['status', 'error_message'])
|
|
return
|
|
|
|
# ── Detect lat/lon columns for schema override ──
|
|
import csv as _csv
|
|
_LAT_LON_KW = ('latd', 'lond', 'latitude', 'longitude')
|
|
lat_lon_overrides = {}
|
|
try:
|
|
with open(csv_paths[0], 'r', encoding='utf-8') as _f:
|
|
_reader = _csv.reader(_f)
|
|
_headers = next(_reader, [])
|
|
for _col in _headers:
|
|
_norm = _col.lower().replace('-', '_').replace('.', '_').replace(' ', '_')
|
|
if any(kw in _norm for kw in _LAT_LON_KW):
|
|
lat_lon_overrides[_col] = pl.Utf8
|
|
logger.info('[BACKGROUND] Forcing lat/lon column %r to Utf8 for safe parsing', _col)
|
|
except Exception:
|
|
logger.warning('[BACKGROUND] Failed to detect lat/lon columns', exc_info=True)
|
|
|
|
total_files = len(csv_paths)
|
|
logger.info('[BACKGROUND] Found %d CSV files in %s', total_files, upload_dir)
|
|
# Adaptive batch size: aim for ~20 batches total, cap at 500 files per batch
|
|
batch_size = max(1, min(500, total_files // 20 or 1))
|
|
total_batches = (total_files + batch_size - 1) // batch_size
|
|
|
|
run.status = 'loading'; run.progress_pct = 10
|
|
run.progress_msg = f'正在流式加载 {total_files} 个文件（共 {total_batches} 批）...'
|
|
run.save(update_fields=['status', 'progress_pct', 'progress_msg'])
|
|
|
|
from analysis.data_loader import save_to_db, drop_sqlite_table, load_from_db, _dtype_to_sqlite
|
|
|
|
table_name = f'_data_{run_id}'
|
|
|
|
# ── Phase 1: Pre-scan all files for full union column set ──
|
|
all_columns = []
|
|
for _path in csv_paths:
|
|
with open(_path, 'r', encoding='utf-8') as _f:
|
|
_reader = _csv.reader(_f)
|
|
_headers = next(_reader, [])
|
|
for _h in _headers:
|
|
if _h not in all_columns:
|
|
all_columns.append(_h)
|
|
reference_columns = sorted(all_columns, key=str.lower)
|
|
schema = {col: 'Utf8' for col in all_columns}
|
|
|
|
# ── Phase 2: Create SQLite table once with all columns as TEXT ──
|
|
import sqlite3 as _sqlite3
|
|
from django.conf import settings as _dj_settings
|
|
_db_path = _dj_settings.DATABASES['default']['NAME']
|
|
_conn = _sqlite3.connect(_db_path)
|
|
try:
|
|
_col_defs = [f'"{_col}" TEXT' for _col in reference_columns]
|
|
_conn.execute(f'DROP TABLE IF EXISTS "{table_name}"')
|
|
_conn.execute(f'CREATE TABLE "{table_name}" ({", ".join(_col_defs)})')
|
|
_conn.commit()
|
|
finally:
|
|
_conn.close()
|
|
|
|
# ── Phase 3: Batched collection → SQLite append ──
|
|
schema_overrides = {col: pl.Utf8 for col in reference_columns}
|
|
completed_batches = 0
|
|
total_rows = 0
|
|
for i in range(0, total_files, batch_size):
|
|
batch = csv_paths[i:i+batch_size]
|
|
# Scan each file INDIVIDUALLY — batch scan fails when files have different schemas
|
|
dfs = []
|
|
for fp in batch:
|
|
lf = pl.scan_csv(fp, infer_schema_length=0,
|
|
schema_overrides=schema_overrides)
|
|
df = lf.collect(streaming=True)
|
|
# Fill missing columns with NULL
|
|
for col in reference_columns:
|
|
if col not in df.columns:
|
|
df = df.with_columns(pl.lit(None).alias(col))
|
|
# Select all reference_columns in canonical order
|
|
df = df.select(reference_columns)
|
|
dfs.append(df)
|
|
df_batch = pl.concat(dfs, how='diagonal_relaxed')
|
|
|
|
# Append to pre-created SQLite table
|
|
save_to_db(run.id, df_batch.lazy(), mode='append')
|
|
|
|
# Delete temp files after successful SQLite write
|
|
for f in batch:
|
|
try:
|
|
os.unlink(f)
|
|
except OSError:
|
|
pass
|
|
|
|
completed_batches += 1
|
|
run.progress_pct = int((completed_batches / total_batches) * 80 + 10)
|
|
run.progress_msg = f'正在加载数据... ({completed_batches}/{total_batches} 批)'
|
|
run.save(update_fields=['progress_pct', 'progress_msg'])
|
|
|
|
# ── Phase 3: Reload from SQLite with correct types ──
|
|
lf = load_from_db(table_name, schema_overrides=schema)
|
|
if lf is None:
|
|
raise RuntimeError(f'SQLite 表 {table_name} 加载失败')
|
|
|
|
# Apply _coerce_to_float on lat/lon columns (restore Float64 from Utf8 storage)
|
|
from analysis.data_loader import _coerce_to_float as _coerce_to_float_fn
|
|
for lat_lon_col in lat_lon_overrides:
|
|
if lat_lon_col in lf.columns:
|
|
lf = lf.with_columns(
|
|
pl.col(lat_lon_col).map_batches(_coerce_to_float_fn, return_dtype=pl.Float64).alias(lat_lon_col)
|
|
)
|
|
|
|
# ── Restore column types after SQLite storage ──
|
|
# All columns were stored as Utf8 for SQLite. Recover original types
|
|
# from the classification system — ONLY the 8 allowed types.
|
|
from analysis.type_classifier import classify_schema, DataType
|
|
type_map = classify_schema(lf)
|
|
# Apply Polars-level type casts based on classification
|
|
dtypes = {
|
|
DataType.FLOAT: pl.Float64,
|
|
DataType.INT: pl.Int64,
|
|
DataType.BOOL: pl.Boolean,
|
|
DataType.TIMESTAMP: pl.Float64,
|
|
}
|
|
casts = []
|
|
for col, dt in type_map.items():
|
|
target = dtypes.get(dt)
|
|
if target and col in lf.collect_schema().names():
|
|
casts.append(pl.col(col).cast(target, strict=False))
|
|
schema[col] = str(target)
|
|
if casts:
|
|
lf = lf.with_columns(casts)
|
|
restored = [(col, dt.value) for col, dt in type_map.items() if dt in dtypes][:15]
|
|
logger.info('[BACKGROUND] Type restoration: %s', restored)
|
|
|
|
# ── SVD dimensionality reduction ──
|
|
svd_components = 0
|
|
# Only operate on columns with known numeric types in the schema
|
|
numeric_dtypes = {'Int64', 'Float64', 'Int32', 'Float32',
|
|
'Int8', 'Int16', 'UInt8', 'UInt16', 'UInt32', 'UInt64'}
|
|
numeric_cols = [c for c, dt in schema.items() if dt in numeric_dtypes]
|
|
n_numeric = len(numeric_cols)
|
|
|
|
if n_numeric >= 2:
|
|
n_components = min(MAX_CLUSTER_FEATURES, n_numeric - 1)
|
|
try:
|
|
from sklearn.decomposition import TruncatedSVD # fmt: skip
|
|
import numpy as np # fmt: skip
|
|
|
|
# Collect numeric data for SVD (full materialise → transform → add back)
|
|
df_full = lf.select(numeric_cols).collect(streaming=True)
|
|
X = df_full.to_numpy()
|
|
X = np.nan_to_num(X, nan=0.0, posinf=0.0, neginf=0.0)
|
|
|
|
svd = TruncatedSVD(n_components=n_components, random_state=RANDOM_SEED)
|
|
X_svd = svd.fit_transform(X)
|
|
svd_components = n_components
|
|
|
|
# Build a new LazyFrame that includes _svd_* columns
|
|
# Strategy: collect full LazyFrame once, add SVD columns, re-wrap as lazy
|
|
df_all = lf.collect(streaming=True)
|
|
for i in range(n_components):
|
|
col_name = f'_svd_{i}'
|
|
schema[col_name] = 'Float64'
|
|
df_all = df_all.with_columns(
|
|
pl.Series(col_name, X_svd[:, i])
|
|
)
|
|
lf = df_all.lazy()
|
|
|
|
logger.info('[BACKGROUND] SVD: %d components from %d numeric columns '
|
|
'(explained variance ratio sum=%.4f)',
|
|
svd_components, n_numeric,
|
|
float(svd.explained_variance_ratio_.sum()))
|
|
except Exception as e:
|
|
logger.warning('[BACKGROUND] SVD failed (non-fatal): %s', e)
|
|
elif n_numeric == 1:
|
|
logger.info('[BACKGROUND] SVD skipped: only 1 numeric column (%s)',
|
|
numeric_cols[0])
|
|
|
|
# Get row count
|
|
df_count = lf.select(pl.len()).collect(streaming=True)
|
|
total_rows = df_count[0, 0] if df_count.height > 0 else 0
|
|
|
|
run.sqlite_table = table_name
|
|
run.total_flows = total_rows
|
|
run.save(update_fields=['sqlite_table', 'total_flows'])
|
|
|
|
ds_id = f'upload_{run.display_id}'
|
|
store.store_dataset(ds_id, lf, schema=schema, metadata={
|
|
'row_count': total_rows, 'file_count': total_files, 'upload_dir': str(upload_dir),
|
|
'csv_glob': str(upload_dir / '*.csv'),
|
|
'sqlite_table': table_name,
|
|
'svd_components': svd_components,
|
|
})
|
|
|
|
# ── Phase 3: Data ready for filtering/clustering via MCP tools ──
|
|
run.status = 'ready'; run.progress_pct = 60; run.progress_msg = '数据准备就绪'
|
|
run.save(update_fields=['status', 'progress_pct', 'progress_msg'])
|
|
except Exception:
|
|
tb = traceback.format_exc()
|
|
logger.error(tb)
|
|
# Drop partial SQLite table on failure (idempotent)
|
|
try:
|
|
from analysis.data_loader import drop_sqlite_table # fmt: skip
|
|
drop_sqlite_table(f'_data_{run_id}')
|
|
except Exception:
|
|
logger.error("unknown failed: {}".format(traceback.format_exc()))
|
|
pass # Expected: drop_sqlite_table may fail if table already gone
|
|
run.status = 'failed'
|
|
run.error_message = tb
|
|
run.run_log += f'\n[ERROR] {tb}'
|
|
run.save(update_fields=['status', 'error_message', 'run_log'])
|