diff --git a/CHANGELOG.md b/CHANGELOG.md index f0257bcd..cd886d98 100644 --- a/CHANGELOG.md +++ b/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). diff --git a/src/dataframeit/utils.py b/src/dataframeit/utils.py index 93bce6ee..af50dcbf 100644 --- a/src/dataframeit/utils.py +++ b/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. diff --git a/tests/test_multiple_data_types.py b/tests/test_multiple_data_types.py index 06201395..3d53718c 100644 --- a/tests/test_multiple_data_types.py +++ b/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