PM-pinou
701 lines
27 KiB
Python
701 lines
27 KiB
Python
"""CSV loading with per-file schema validation, BOM detection, and multi-file merge.
|
||
|
||
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`.
|
||
"""
|
||
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.type_classifier import classify_schema, DataType
|
||
|
||
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 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_ENUM normalisation ────────────────────────────────
|
||
try:
|
||
type_map = classify_schema(
|
||
merged_lf.head(1000),
|
||
config_overrides=config.get('columns', {}),
|
||
)
|
||
except Exception:
|
||
type_map = {}
|
||
|
||
bool_enum_cols = [
|
||
c for c, t in type_map.items()
|
||
if t == DataType.BOOL_ENUM and c in merged_schema
|
||
]
|
||
if bool_enum_cols:
|
||
# Only clean string-typed BOOL_ENUM columns; native Boolean is already clean
|
||
schema_dtypes = dict(zip(
|
||
merged_lf.collect_schema().names(),
|
||
merged_lf.collect_schema().dtypes(),
|
||
))
|
||
bool_exprs = []
|
||
cleaned_names: list[str] = []
|
||
for c in bool_enum_cols:
|
||
if schema_dtypes.get(c) not in (pl.Utf8, pl.String):
|
||
continue # native Boolean / numeric — skip
|
||
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_ENUM columns normalised: %s', cleaned_names)
|
||
|
||
# ── Cleaning: LAT_LON normalisation ────────────────────────────────
|
||
latlon_cols = [
|
||
c for c, t in type_map.items()
|
||
if t == DataType.LAT_LON and c in merged_schema
|
||
]
|
||
if latlon_cols:
|
||
schema_dtypes_ll = dict(zip(
|
||
merged_lf.collect_schema().names(),
|
||
merged_lf.collect_schema().dtypes(),
|
||
))
|
||
latlon_exprs = []
|
||
for c in latlon_cols:
|
||
dtype = schema_dtypes_ll.get(c)
|
||
if dtype in (pl.Utf8, pl.String):
|
||
latlon_exprs.append(
|
||
pl.col(c).map_batches(_coerce_to_float, return_dtype=pl.Float64).alias(c)
|
||
)
|
||
merged_schema[c] = 'Float64'
|
||
if latlon_exprs:
|
||
merged_lf = merged_lf.with_columns(latlon_exprs)
|
||
logger.info('[CLEAN] LAT_LON columns normalised to Float64: %s', latlon_cols)
|
||
|
||
# ── Cleaning: Generic numeric coercion ────────────────────────────
|
||
# For any remaining string column (not IPv4/URL/HEX/ENUM/etc.),
|
||
# try to coerce to Float64. Whatever can't parse becomes null.
|
||
# This prevents "mean on str column" crashes for ANY column.
|
||
_SKIP_NUM_COERCE = (DataType.IPv4, DataType.URL, DataType.HEX,
|
||
DataType.ENUM, DataType.LAT_LON, DataType.BOOL_ENUM)
|
||
# Also protect columns with these name patterns (MAC addresses, etc.)
|
||
_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 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 # already numeric
|
||
# Skip columns with protected name patterns
|
||
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: HEX preprocessing (hex_pairs_count) ────────────────────
|
||
hex_cols = [
|
||
c for c, t in type_map.items()
|
||
if t == DataType.HEX 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] HEX 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
|