tianxuan/analysis/data_loader/_sqlite.py

"""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()