src/query/executor/tabular.py

"""TabularExecutor — runs compiled pandas/polars chain on a Parquet file.

Picks engine by file size:
  ≤ 100 MB    → eager pandas
  100 MB-1 GB → pyarrow with predicate pushdown
  > 1 GB      → polars lazy scan

Initial scope ships eager pandas only; the others are added when a real
file is too big.
"""

from __future__ import annotations

import asyncio
import io
import time
from collections.abc import Callable, Coroutine
from typing import Any

import pandas as pd

from ...catalog.models import Catalog, Source, Table
from ...storage.parquet import parquet_blob_name
from ...middlewares.logging import get_logger
from ..compiler.pandas import CompiledPandas, PandasCompiler
from ..ir.models import QueryIR
from .base import BaseExecutor, QueryResult

logger = get_logger("tabular_executor")

_AZ_BLOB_PREFIX = "az_blob://"
# Go's Supabase S3 data plane writes location_ref with this prefix instead of az_blob://.
# Both encode the same path structure after the prefix: {user_id}/{document_id}.
_OBJECT_STORAGE_PREFIX = "object_storage://"
_LOCATION_REF_PREFIXES = (_AZ_BLOB_PREFIX, _OBJECT_STORAGE_PREFIX)
_ROW_HARD_CAP = 10_000

# Largest Parquet blob we will pull into memory (F-13, 2026-07-23).
#
# The whole object is buffered and then `pd.read_parquet`-ed, so peak RSS is the file
# plus the decompressed frame — often 3-5x for string-heavy data. Filtering and
# `_ROW_HARD_CAP` both happen strictly AFTER the frame exists, so they bound the
# RESULT, never the working set. An OOM here is not catchable: it kills the process
# and every concurrent request with it, which is the one failure mode in this service
# that escapes every never-throw seam.
#
# Deliberately a SAFETY NET, not a tight cap: 500 MB of Parquet is far beyond any
# upload seen so far, so no real user should ever meet it. Every rejection logs, and
# that log is the signal to revisit the number (or to implement the size-tiered
# strategy this module's docstring already sketches: pyarrow pushdown, polars lazy).
_MAX_BLOB_BYTES = 500 * 1024 * 1024


class TabularExecutor(BaseExecutor):
    """Executes compiled pandas chain on a Parquet blob.

    `fetch_blob` is injectable for tests — defaults to the storage backend
    selected by `settings.storage_provider` (azure_blob | supabase_s3).
    """

    def __init__(
        self,
        catalog: Catalog,
        fetch_blob: Callable[[str], Coroutine[Any, Any, bytes]] | None = None,
        blob_size: Callable[[str], Coroutine[Any, Any, int | None]] | None = None,
    ) -> None:
        self._catalog = catalog
        self._compiler = PandasCompiler(catalog)
        self._fetch_blob = fetch_blob or self._default_fetch_blob
        # Injected alongside `fetch_blob` so a test that fakes the download can also
        # fake the size probe. When only `fetch_blob` is injected the probe is skipped
        # and the post-download check still applies.
        self._blob_size = blob_size or (
            self._default_blob_size if fetch_blob is None else None
        )

    @staticmethod
    async def _default_blob_size(blob_name: str) -> int | None:
        from ...config.settings import settings

        provider = (settings.storage_provider or "").strip().lower()
        if provider == "supabase_s3":
            from ...storage.object_storage import object_storage

            return await object_storage.object_size(blob_name)

        from ...storage.az_blob.az_blob import blob_storage

        return await blob_storage.object_size(blob_name)

    @staticmethod
    async def _default_fetch_blob(blob_name: str) -> bytes:
        # Pick the storage backend by the same toggle the Go data plane uses.
        # Blank/unknown falls back to Azure Blob (the pre-migration default).
        from ...config.settings import settings

        provider = (settings.storage_provider or "").strip().lower()
        if provider == "supabase_s3":
            from ...storage.object_storage import object_storage

            return await object_storage.download_file(blob_name)

        from ...storage.az_blob.az_blob import blob_storage

        return await blob_storage.download_file(blob_name)

    async def run(self, ir: QueryIR) -> QueryResult:
        started = time.perf_counter()
        table_name = ""
        source_name = ""
        try:
            source, table = self._lookup(ir)
            table_name = table.name
            source_name = source.name
            if source.source_type != "tabular":
                raise ValueError(
                    f"TabularExecutor cannot run on source_type={source.source_type!r}; "
                    "expected 'tabular'"
                )

            compiled = self._compiler.compile(ir)
            rendered_query = _render_query(ir, {c.column_id: c for c in table.columns})
            logger.info("pandas query", query=rendered_query)
            blob_name = _resolve_blob_name(source, table)

            # Refuse an oversized blob BEFORE buffering it — see _MAX_BLOB_BYTES.
            # A None size means the backend couldn't tell us (HEAD failed, or a test
            # injected only `fetch_blob`); we proceed and rely on the post-download
            # check below, which still beats no bound at all. (F-13)
            if self._blob_size is not None:
                size = await self._blob_size(blob_name)
                if size is not None and size > _MAX_BLOB_BYTES:
                    raise ValueError(
                        f"file is too large to analyse ({size / 1024 / 1024:.0f} MB; "
                        f"limit {_MAX_BLOB_BYTES / 1024 / 1024:.0f} MB) — "
                        "filter it down or split it before uploading"
                    )

            blob_bytes = await self._fetch_blob(blob_name)

            # Backstop for the case where the size probe was unavailable. The bytes
            # are already resident here, so this cannot prevent the download — but it
            # does stop `pd.read_parquet` from multiplying them into a frame several
            # times larger, which is where the OOM actually happens. (F-13)
            if len(blob_bytes) > _MAX_BLOB_BYTES:
                raise ValueError(
                    f"file is too large to analyse ({len(blob_bytes) / 1024 / 1024:.0f} MB; "
                    f"limit {_MAX_BLOB_BYTES / 1024 / 1024:.0f} MB) — "
                    "filter it down or split it before uploading"
                )

            result_df = await asyncio.to_thread(_load_and_apply, blob_bytes, compiled)

            truncated = len(result_df) > _ROW_HARD_CAP
            capped = result_df.head(_ROW_HARD_CAP)

            # `output_columns` is derived from the SELECT list, not from the frame the
            # compiler actually produced, and the rows below are mapped BY NAME
            # (`data_access._retrieve_data` does `row.get(c)`). A declared name that is
            # absent from the frame silently becomes a column of None — a real-looking
            # table with a fabricated column, which is worse than an error. F-17's
            # validator fix closes the one known way in (mixed select, no group_by);
            # this is the backstop that turns any future divergence into an honest
            # failure instead of a wrong answer.
            #
            # Presence, not order: the grouped path builds its frame via
            # `reset_index()`, which emits the group columns first regardless of where
            # they sat in the select list, so a positional comparison would reject
            # perfectly good queries. Name-mapping makes order irrelevant. (F-17)
            produced = set(capped.columns)
            fabricated = [c for c in compiled.output_columns if c not in produced]
            if fabricated:
                raise ValueError(
                    f"pandas compiler did not produce declared column(s) {fabricated!r} "
                    f"(frame has {list(capped.columns)!r}) — refusing to return a "
                    "result whose columns would be fabricated as null"
                )

            columns = compiled.output_columns
            rows = capped.to_dict(orient="records")
            elapsed_ms = int((time.perf_counter() - started) * 1000)
            logger.info(
                "tabular query complete",
                source_id=ir.source_id,
                rows=len(rows),
                truncated=truncated,
                elapsed_ms=elapsed_ms,
            )
            return QueryResult(
                source_id=ir.source_id,
                backend="tabular",
                columns=columns,
                rows=rows,
                row_count=len(rows),
                truncated=truncated,
                elapsed_ms=elapsed_ms,
                table_id=ir.table_id,
                table_name=table_name,
                source_name=source_name,
                query=rendered_query,  # rendered pandas chain, for traceability (KM-691)
            )

        except Exception as e:
            elapsed_ms = int((time.perf_counter() - started) * 1000)
            logger.error(
                "tabular executor failed",
                degraded_seam="tabular_execute",
                source_id=ir.source_id,
                error=repr(e),
                elapsed_ms=elapsed_ms,
            )
            return QueryResult(
                source_id=ir.source_id,
                backend="tabular",
                elapsed_ms=elapsed_ms,
                # See db.py: fall back to repr only when str() is empty, so no
                # existing user-facing error text changes. (F-26)
                error=str(e) or repr(e),
                table_id=ir.table_id,
                table_name=table_name,
                source_name=source_name,
            )

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _lookup(self, ir: QueryIR) -> tuple[Source, Table]:
        source = next(
            (s for s in self._catalog.sources if s.source_id == ir.source_id), None
        )
        if source is None:
            raise ValueError(f"source_id {ir.source_id!r} not in catalog")
        table = next(
            (t for t in source.tables if t.table_id == ir.table_id), None
        )
        if table is None:
            raise ValueError(f"table_id {ir.table_id!r} not in source {ir.source_id!r}")
        return source, table


# ---------------------------------------------------------------------------
# Module-level helpers (pure functions — easier to test in isolation)
# ---------------------------------------------------------------------------

def _resolve_blob_name(source: Source, table: Table) -> str:
    """Map source.location_ref + table → the Parquet blob name to download.

    Delegates to ``parquet_service.parquet_blob_name`` so the same naming
    convention (and ``_safe_sheet_name`` sanitization) is used on both the
    write side (ingestion) and the read side (query execution).

      CSV / Parquet → ``{user_id}/{document_id}.parquet``
      XLSX          → ``{user_id}/{document_id}__{safe_sheet}.parquet``
                      (writer always uploads with sheet suffix for XLSX,
                      regardless of sheet count — see processing_service
                      `_build_excel_documents`)

    XLSX is detected via ``Source.name`` (the original filename). This relies
    on the upload pipeline preserving the file extension, which it does today
    because `Document.filename` is set once at upload and never renamed.
    """
    matched_prefix = next(
        (p for p in _LOCATION_REF_PREFIXES if source.location_ref.startswith(p)),
        None,
    )
    if matched_prefix is None:
        raise ValueError(
            f"TabularExecutor expects 'az_blob://...' or 'object_storage://...' "
            f"location_ref, got {source.location_ref!r}"
        )
    path = source.location_ref[len(matched_prefix):]
    parts = path.split("/", 1)
    if len(parts) != 2 or not parts[0] or not parts[1]:
        raise ValueError(f"Malformed location_ref: {source.location_ref!r}")
    user_id, document_id = parts
    is_xlsx = source.name.lower().endswith(".xlsx")
    sheet_name = table.name if is_xlsx else None
    return parquet_blob_name(user_id, document_id, sheet_name)


def _render_value(v: object, cap: int = 5) -> str:
    """Render a filter value for the human-readable query string, truncating a
    long list so a value-handoff `in`/`not_in` set doesn't dump thousands of ids
    into the logs AND the user-facing traceability record (KM-691)."""
    if isinstance(v, list | tuple) and len(v) > cap:
        head = ", ".join(repr(x) for x in v[:cap])
        return f"[{head}, … +{len(v) - cap} more]"
    return repr(v)


def _render_query(ir: QueryIR, cols_by_id: dict) -> str:
    from ..ir.models import AggSelect, ColumnSelect
    parts = ["df"]
    if ir.filters:
        conds = " & ".join(
            f'(df["{cols_by_id[f.column_id].name}"] {f.op} {_render_value(f.value)})'
            for f in ir.filters
        )
        parts.append(f"[{conds}]")
    aggs = [s for s in ir.select if isinstance(s, AggSelect)]
    cols = [s for s in ir.select if isinstance(s, ColumnSelect)]
    if aggs:
        col_names = [cols_by_id[s.column_id].name for s in cols]
        if ir.group_by:
            group_names = [cols_by_id[g].name for g in ir.group_by]
            parts.append(f'.groupby({group_names})')
        for agg in aggs:
            col = f'["{cols_by_id[agg.column_id].name}"]' if agg.column_id else ""
            fn_map = {"count": "count()", "count_distinct": "nunique()", "sum": "sum()", "avg": "mean()", "min": "min()", "max": "max()"}
            parts.append(f'{col}.{fn_map.get(agg.fn, agg.fn + "()")}')
    elif cols:
        col_names = [cols_by_id[s.column_id].name for s in cols]
        parts.append(f'[{col_names}]')
    if ir.limit:
        parts.append(f'.head({ir.limit})')
    return "".join(parts)


def _load_and_apply(blob_bytes: bytes, compiled: CompiledPandas) -> pd.DataFrame:
    """Load Parquet bytes into a DataFrame and apply the compiled op chain."""
    df = pd.read_parquet(io.BytesIO(blob_bytes))
    return compiled.apply(df)