Preservar tipos do polars nas colunas de entrada

CHANGELOG.md

@@ -35,6 +35,7 @@ e este projeto adere ao [Versionamento Semântico](https://semver.org/lang/pt-BR
 
 ### Corrigido
 
+- Com entrada polars, as colunas que o dataframeit não toca voltam com o tipo de entrada. Um `Int64` com nulo deixa de virar `Float64`, o que arredondava ids acima de 2**53, e um `Date` deixa de virar `Datetime` (#195).
 - O operando de `in` e `not_in` em `condition` é conferido antes de processar, com `ValueError` quando não é lista, tupla, set ou frozenset. Antes, um operando que não era coleção, como `{"in": 5}`, levantava `TypeError` em cada linha depois das chamadas dos campos anteriores, e um texto, como `{"in": "pfj"}`, fazia busca de substring e aceitava `"pf"` (#182).
 - A retomada de um modelo com `strict=True` deixa de acusar como incompatíveis os campos que o checkpoint grava em outra forma, como a data em texto, a tupla como lista, o `Enum` pelo valor e o inteiro que a coluna com linha pendente relê como float. A linha recusada pela validação configurada é validada de novo sem strict, e cada campo recusado só é aceito quando o valor validado diz o mesmo que o relido; `"3"` num campo `int` strict continua acusado. A volta ao tipo declarado passa a valer também para os campos escalares: data, `datetime`, `Enum`, `int` e `bool` relidos do CSV, do XLSX ou do parquet deixam de misturar o tipo da execução com o do arquivo. Os campos de `Enum` de valores em texto são relidos como texto do CSV e do XLSX, como os de `Literal`, e um valor como `"1"` deixa de voltar como número. O checkpoint grava o `Enum` escalar pelo valor; antes, o CSV e o XLSX gravavam `"Cor.AZUL"`, que a retomada recusava, e o parquet falhava ao gravar (#178).
 - A forma de `condition` em `json_schema_extra` é conferida antes de processar, com `ValueError` que diz o que é aceito: dict com `field` em texto e um operador (`equals`, `not_equals`, `in`, `not_in` ou `exists`), ou callable. Antes, `condition={}` e os demais valores falsos, como `False` ou `[]`, rodavam o campo sem aviso, embora `evaluate_condition` trate o dict vazio como condição falsa; `field` que não era texto levantava `AttributeError`; e as demais formas eram avisadas a cada linha, com o campo pulado (#163).

src/dataframeit/utils.py

@@ -90,6 +90,7 @@ class ConversionInfo:
     original_type: str
     original_index: pd.Index | list | None = None  # Índice/chaves originais para reconversão
     series_name: Hashable | None = None  # Nome original da Series (se aplicável)
+    polars_schema: dict[str, Any] | None = None  # Tipos do DataFrame polars de entrada
 
 
 def parse_json(resposta: str) -> dict:
@@ -152,7 +153,10 @@ def to_pandas(data: object) -> tuple[pd.DataFrame, ConversionInfo]:
 
     # polars DataFrame
     if pl is not None and isinstance(data, pl.DataFrame):
-        return data.to_pandas(), ConversionInfo(original_type=ORIGINAL_TYPE_POLARS_DF)
+        return _polars_to_pandas(data), ConversionInfo(
+            original_type=ORIGINAL_TYPE_POLARS_DF,
+            polars_schema=dict(data.schema),
+        )
 
     # pandas Series
     if isinstance(data, pd.Series):
@@ -245,10 +249,40 @@ def from_pandas(
 
     df = _reorder_columns(df, status_col)
     if original_type in _POLARS_INPUTS and pl is not None:
-        return pl.from_pandas(df)
+        result = pl.from_pandas(df)
+        if conversion_info.polars_schema:
+            result = _restore_polars_dtypes(result, conversion_info.polars_schema)
+        return result
+    return df
+
+
+def _polars_to_pandas(data: pl.DataFrame) -> pd.DataFrame:
+    """Converte para pandas mantendo inteiros com nulo como inteiros anuláveis.
+
+    `to_pandas()` leva um inteiro com nulo a float64, o que arredonda ids acima de 2**53.
+    """
+    df = data.to_pandas()
+    for name, dtype in data.schema.items():
+        if dtype.is_integer() and data[name].null_count() > 0:
+            df[name] = pd.array(data[name].to_list(), dtype=str(dtype))
     return df
 
 
+def _restore_polars_dtypes(result: pl.DataFrame, schema: dict[str, Any]) -> pl.DataFrame:
+    """Devolve às colunas de entrada o tipo polars que elas tinham.
+
+    A volta pelo pandas troca `Date` por `Datetime`, por exemplo. Uma coluna que não aceita
+    o tipo original fica como está.
+    """
+    for name, dtype in schema.items():
+        if name in result.columns and result.schema[name] != dtype:
+            try:
+                result = result.with_columns(pl.col(name).cast(dtype))
+            except pl.exceptions.PolarsError:
+                continue
+    return result
+
+
 def _reorder_columns(df: pd.DataFrame, status_col: str = "_dataframeit_status") -> pd.DataFrame:
     """Reordena colunas para que internas fiquem no final.
 

tests/test_multiple_data_types.py

@@ -7,6 +7,7 @@
 - dict
 """
 
+import datetime
 import importlib.util
 import sys
 
@@ -30,6 +31,7 @@
     ORIGINAL_TYPE_POLARS_DF,
     ORIGINAL_TYPE_POLARS_SERIES,
     ConversionInfo,
+    _restore_polars_dtypes,
     from_pandas,
     to_pandas,
 )
@@ -304,3 +306,43 @@ def test_sem_polars_a_entrada_pandas_continua_funcionando(monkeypatch):
     assert copia.pl is None
     assert info.original_type == ORIGINAL_TYPE_PANDAS_DF
     assert copia.from_pandas(convertido, info).equals(df)
+
+
+@pytest.mark.skipif(not HAS_POLARS, reason="polars não instalado")
+class TestPolarsDtypeRoundTrip:
+    """Colunas do polars que o dataframeit não toca voltam com o tipo de entrada."""
+
+    def test_integer_with_null_keeps_precision_and_dtype(self):
+        ids = [9007199254740993, None, 1234567890123456789]
+        data = pl.DataFrame({"texto": ["a", "b", "c"], "id": pl.Series(ids, dtype=pl.Int64)})
+
+        df, info = to_pandas(data)
+        df["resultado"] = ["x", "y", "z"]
+        result = from_pandas(df, info)
+
+        assert result["id"].dtype == pl.Int64
+        assert result["id"].to_list() == ids
+
+    def test_integer_without_null_is_unchanged(self):
+        data = pl.DataFrame({"id": pl.Series([1, 2], dtype=pl.UInt8)})
+
+        df, info = to_pandas(data)
+
+        assert df["id"].tolist() == [1, 2]
+        assert from_pandas(df, info)["id"].dtype == pl.UInt8
+
+    def test_date_column_stays_date(self):
+        data = pl.DataFrame({"texto": ["a"], "dia": [datetime.date(2024, 1, 2)]})
+
+        df, info = to_pandas(data)
+        result = from_pandas(df, info)
+
+        assert result["dia"].dtype == pl.Date
+        assert result["dia"].to_list() == [datetime.date(2024, 1, 2)]
+
+    def test_column_that_rejects_the_original_dtype_is_left_alone(self):
+        result = pl.DataFrame({"id": ["abc"]})
+
+        restored = _restore_polars_dtypes(result, {"id": pl.Int64})
+
+        assert restored["id"].dtype == pl.String