Retomada de saída polars arredonda campo inteiro acima de 2**53 e devolve a coluna como Float64 #222

Quando a saída polars de uma execução interrompida volta como entrada com resume=True, o campo int do modelo que já tinha valor acima de 2**53 sai arredondado, e a coluna inteira sai Float64. Nenhum erro, nenhum aviso.

from unittest.mock import patch
import warnings

from pydantic import BaseModel
import polars as pl

from dataframeit import ProviderUsageLimitError, dataframeit


def rodar(df, modelo, respostas, **opcoes):
    fila = iter(respostas)

    def llm(*args, **kwargs):
        resposta = next(fila)
        if isinstance(resposta, Exception):
            raise resposta
        return {"data": resposta, "usage": None}

    with patch("dataframeit.core.call_langchain", side_effect=llm), \
         patch("dataframeit.core.validate_provider_dependencies"), \
         warnings.catch_warnings():
        warnings.simplefilter("ignore")
        return dataframeit(df, questions=modelo, prompt="{texto}", track_tokens=False, max_retries=1, **opcoes)


class M(BaseModel):
    nota: int


GRANDE = 2**53 + 1
entrada = pl.DataFrame({"texto": ["a", "b", "c"]})
parcial = rodar(entrada, M, [{"nota": GRANDE}, ProviderUsageLimitError("cota")])
print(parcial.schema["nota"], parcial["nota"].to_list())
# Int64 [9007199254740993, None, None]

final = rodar(parcial, M, [{"nota": 7}, {"nota": 8}], resume=True)
print(final.schema["nota"], final["nota"].to_list())
# Float64 [9007199254740992.0, 7.0, 8.0]

O 9007199254740993 gravado na primeira execução vira 9007199254740992.0, e as respostas novas, 7 e 8, saem 7.0 e 8.0 num campo declarado int.

Causa

A coluna do modelo volta com nulo nas linhas pendentes. to_pandas (utils.py:158) converte o frame com data.to_pandas(), e uma coluna polars Int64 com nulo chega ao pandas como float64. O valor já entra arredondado, antes de qualquer conferência da retomada.

O #217 tirou desse caminho as colunas de entrada do usuário, que saem do frame original. A coluna do modelo não tem essa saída, porque a execução escreve nela.

Três correções pela entrada foram tentadas durante o #217, olhando só o dado: inteiro anulável do pandas (Int64), coluna object com inteiros Python e float só quando o valor cabe exato. Cada uma quebrou outro caso de retomada, e nenhuma entrou. Entre as quebras: o campo bool relido de uma coluna Int64 de 0 e 1 com nulo, o campo float sobre inteiro grande, que levantou ArrowInvalid na volta, e a coluna UInt64, que levantou OverflowError.

Proposta

Tratar só a coluna do modelo, e só quando a anotação do campo é int: ler essa coluna do frame polars de entrada como inteiros Python, em vez de passar pelo float64. A anotação diz qual tipo esperar, informação que as tentativas acima não usavam. O teste é o trecho acima, mais os casos que quebraram: campo bool sobre Int64 de 0 e 1 com nulo, campo float sobre inteiro acima de 2**53 e coluna UInt64.

Relacionada: #196, que reúne outros casos em que a retomada altera valor já gravado.