fn_registry/python/functions/datascience/extract_null_mask.py

"""extract_null_mask — extrae la mascara de nulos (1=falta / 0=presente) de una tabla.

Lector read-only inyectado: recibe `query_fn(sql) -> dict` con el mismo contrato
que duckdb_query_readonly / pg_query (y que el `_q` de profile_table):
`{"status": "ok", "rows": [{col: val, ...}, ...]}`. Esta funcion NO abre ninguna
conexion por su cuenta — solo usa `query_fn`. Construye UNA sola query que, por
cada columna pedida, evalua `CASE WHEN "col" IS NULL THEN 1 ELSE 0 END` y devuelve
una muestra de filas con esos bits. El resultado es un dict `mask` con una lista
0/1 por columna, alineada por fila (1 = el valor falta / IS NULL, 0 = presente),
listo para alimentar el capitulo de calidad / patron de nulos de AutomaticEDA sin
que el capitulo toque la base de datos.

Estilo dict-no-throw del grupo `eda`: nunca lanza; captura cualquier excepcion y
degrada a `{"status": "error", "error": str, ...}`.
"""


def _to_bit(value):
    """Coacciona el valor 0/1 del CASE a int de forma defensiva.

    El SQL ya devuelve 0 (presente) o 1 (falta). Por si una celda llega como None
    (el CASE no se aplico o el backend la nulifico), se cuenta como 1 (falta). El
    resto se reduce a int: un entero distinto de 0 cuenta como 1 (falta), 0 como
    presente. Un valor no convertible se trata como presente (0) — nunca lanza.
    """
    if value is None:
        return 1
    try:
        return 1 if int(value) != 0 else 0
    except (TypeError, ValueError):
        return 0


def extract_null_mask(query_fn, table, columns, max_rows=5000):
    """Extrae la mascara de nulos (1=falta / 0=presente) de una muestra de la tabla.

    Args:
        query_fn: callable lector read-only del backend activo. Recibe un string
            SQL y devuelve un dict {"status": "ok", "rows": [{col: val, ...}]}
            (mismo contrato que duckdb_query_readonly / el `_q` de profile_table).
            No se abre ninguna conexion aqui: toda la lectura pasa por query_fn.
        table: nombre de la tabla. Se escapa con comillas dobles en la query.
        columns: lista de nombres de columna a evaluar. Cada una produce una
            entrada en `mask` con una lista 0/1 paralela por fila. Vacia o None ->
            status error.
        max_rows: limite de filas a muestrear (clausula LIMIT). Default 5000.

    Returns:
        dict (nunca lanza):
            {
              "status": "ok" | "error",
              "error": str,                 # solo si status == "error"
              "table": str,
              "columns": [str, ...],        # columnas efectivamente leidas, en orden
              "mask": {col: [int 0/1, ...], ...},  # alineada por fila, 1=falta, 0=presente
              "n": int                      # nº de filas muestreadas
            }
        Todas las listas de `mask` tienen la misma longitud (= n).
    """
    base = {"status": "ok", "table": table, "columns": [], "mask": {}, "n": 0}
    try:
        if query_fn is None:
            return {**base, "status": "error", "error": "query_fn es None"}
        if not table:
            return {**base, "status": "error", "error": "table es obligatorio"}
        if not columns:
            return {**base, "status": "error", "error": "columns vacío"}

        # Identificadores escapados con comillas dobles (como hace profile_table)
        # para tolerar nombres con mayusculas/espacios/palabras reservadas. Cada
        # columna se proyecta como su propio bit IS NULL conservando el alias.
        select_sql = ", ".join(
            f'(CASE WHEN "{c}" IS NULL THEN 1 ELSE 0 END) AS "{c}"' for c in columns
        )
        sql = f'SELECT {select_sql} FROM "{table}" LIMIT {int(max_rows)}'

        q = query_fn(sql)
        if not isinstance(q, dict) or q.get("status") != "ok":
            err = (
                q.get("error", "query_fn fallo")
                if isinstance(q, dict)
                else "query_fn no devolvio un dict"
            )
            return {**base, "status": "error", "error": err}

        rows = q.get("rows", []) or []
        mask = {c: [] for c in columns}
        for row in rows:
            for c in columns:
                # row.get tolera filas que no traigan la columna (None -> falta).
                mask[c].append(_to_bit(row.get(c) if isinstance(row, dict) else None))

        return {
            "status": "ok",
            "table": table,
            "columns": list(columns),
            "mask": mask,
            "n": len(rows),
        }
    except Exception as e:  # noqa: BLE001 - dict-no-throw: degradar, nunca lanzar
        return {**base, "status": "error", "error": str(e)}