From ecbfb102398772aac768dd5b6f20ff5e0b87214b Mon Sep 17 00:00:00 2001 From: Bruno Oliveira Date: Tue, 13 Oct 2026 15:54:55 -0300 Subject: [PATCH 1/3] =?UTF-8?q?Sa=C3=ADda=20polars:=20campo=20dict=20em=20?= =?UTF-8?q?texto=20JSON=20e=20coluna=20sem=20valor=20com=20o=20tipo=20da?= =?UTF-8?q?=20anota=C3=A7=C3=A3o?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Na saída polars, um campo dict saía Struct com as chaves de todas as linhas: a chave ausente ficava igual à chave com valor nulo, e o dict vazio deixava de ser vazio. O campo anotado como dict, inclusive dentro de Optional, passa a sair como texto JSON, serializado como no checkpoint, com nulo na linha sem resposta. O dict com valores de tipos diferentes, que saía Object com aviso, sai do mesmo jeito. A saída pandas não muda. O tipo da coluna de um campo vinha só das respostas, e duas execuções do mesmo modelo podiam devolver schemas diferentes: um Optional[int] todo nulo saía String, e um list[str] só com listas vazias, List(Null). A coluna sem nenhum valor recebe o tipo da anotação quando ela é str, int, float, bool, date, datetime ou lista de um deles. Com outra anotação, ou com algum valor na coluna, fica o tipo inferido. A entrada polars vazia sai com o schema da entrada com linhas. from_pandas recebe a anotação de cada campo do modelo, e _json_default vai para utils.py, de onde core.py o importa. Closes #219 --- CHANGELOG.md | 3 +- docs/en/reference/api.md | 2 + docs/reference/api.md | 2 + src/dataframeit/core.py | 21 +-- src/dataframeit/utils.py | 90 ++++++++++- tests/test_entrada_e_saida.py | 257 +++++++++++++++++++++++++++++- tests/test_multiple_data_types.py | 11 +- 7 files changed, 360 insertions(+), 26 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1260cea5..392895e3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -38,7 +38,8 @@ e este projeto adere ao [Versionamento Semântico](https://semver.org/lang/pt-BR ### Corrigido -- Com entrada polars, a saída deixa de falhar depois de todas as chamadas quando um campo do modelo guarda valores que o polars não reúne numa coluna tipada, como `Union[int, str]` com número numa linha e texto em outra, `tuple[int, str]`, `Any` e `dict[str, Any]` com tipos diferentes, lista de `Enum`, inteiro acima de 2**53 ao lado de float, inteiro acima do `Int64` e `datetime` com fuso de deslocamento, como `+05:30`. A coluna sai com o tipo `Object`, com os objetos Python que a entrada pandas devolveria, e um aviso a nomeia. O campo `Enum` sai pelo valor do membro, inclusive na coluna `Object` (#218). +- Com entrada polars, a saída deixa de falhar depois de todas as chamadas quando um campo do modelo guarda valores que o polars não reúne numa coluna tipada, como `Union[int, str]` com número numa linha e texto em outra, `tuple[int, str]`, `Any` com tipos diferentes, lista de `Enum`, inteiro acima de 2**53 ao lado de float, inteiro acima do `Int64` e `datetime` com fuso de deslocamento, como `+05:30`. A coluna sai com o tipo `Object`, com os objetos Python que a entrada pandas devolveria, e um aviso a nomeia. O campo `Enum` sai pelo valor do membro, inclusive na coluna `Object` (#218). +- Na saída polars, um campo `dict` sai como texto JSON, como no checkpoint. Antes, saía `Struct` com as chaves de todas as linhas: a chave ausente ficava igual à chave com valor nulo, e o dict vazio deixava de ser vazio. A coluna de um campo sem nenhum valor, toda nula ou só com listas vazias, sai com o tipo da anotação quando ela é `str`, `int`, `float`, `bool`, `date`, `datetime` ou lista de um deles. Antes, o tipo vinha só das respostas: um `Optional[int]` todo nulo saía `String`, e um `list[str]` só com listas vazias, `List(Null)` (#219). - A coluna `_trace` do agente único fica com as demais colunas de trace na ordem da saída, depois dos campos do modelo. Antes, numa entrada que já a trazia, os campos novos saíam depois dela. - Com entrada polars, as colunas que a execução não escreve saem do DataFrame de entrada, com o tipo e os valores de lá. Antes, a saída inteira era convertida de volta do pandas: um `Int64` com nulo voltava `Float64`, o que arredondava ids acima de 2**53, um `Date` voltava `Datetime`, e `Enum`, `Decimal`, `Array`, `Null` e o inteiro dentro de `Struct` e `List` também trocavam de tipo; uma coluna `Object` também trocava de tipo ou, com objetos que o Arrow não infere, levantava `ArrowInvalid` depois das chamadas (#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). diff --git a/docs/en/reference/api.md b/docs/en/reference/api.md index 095e6ba1..5ad34e58 100644 --- a/docs/en/reference/api.md +++ b/docs/en/reference/api.md @@ -115,6 +115,8 @@ Returns data in the same format as input with extracted columns added. With polars input, the input columns that are not model fields and do not have the name of one of the [added columns](#added-columns) come back as they went in, with the same type and the same values. In polars output, an `Enum` field comes out as the member's value. When the answers of a field do not fit a typed polars column, such as a `Union[int, str]` that received a number in one row and text in another, the column comes out with the `Object` type, holding the Python objects the model returned, except an `Enum` member outside a list or dict, which comes out as its value there too, and a warning names the columns; convert them before writing to a file. +A field annotated as `dict`, such as `dict[str, int]` or `Optional[dict[str, Any]]`, comes out in polars output as JSON text, with null where the row has no answer; `json.loads` returns the dict of each row. The type of the other model columns comes from the answers. A column with no value at all, all null or holding only empty lists, comes out with the type of the annotation when it is `str`, `int`, `float`, `bool`, `date`, `datetime`, or a list of one of them: `String`, `Int64`, `Float64`, `Boolean`, `Date`, `Datetime`, or a `List` of one of these. With any other annotation, the column with no value keeps the type polars infers. + ### Added Columns | Column | When it exists | Content | diff --git a/docs/reference/api.md b/docs/reference/api.md index d1153f7a..7a466b8d 100644 --- a/docs/reference/api.md +++ b/docs/reference/api.md @@ -115,6 +115,8 @@ Retorna dados no mesmo formato da entrada com colunas extraídas adicionadas. Com entrada polars, as colunas de entrada que não são campo do modelo nem têm nome de uma das [colunas adicionadas](#colunas-adicionadas) voltam como entraram, com o mesmo tipo e os mesmos valores. Na saída polars, um campo `Enum` sai pelo valor do membro. Quando as respostas de um campo não cabem numa coluna tipada do polars, como um `Union[int, str]` que recebeu número numa linha e texto em outra, a coluna sai com o tipo `Object`, com os objetos Python que o modelo devolveu, exceto o membro de `Enum` fora de lista ou dict, que também nela sai pelo valor, e um aviso diz quais colunas são; converta-as antes de gravar em arquivo. +Um campo anotado como `dict`, como `dict[str, int]` ou `Optional[dict[str, Any]]`, sai na saída polars como texto JSON, com nulo onde a linha não tem resposta; `json.loads` devolve o dict de cada linha. O tipo das demais colunas do modelo vem das respostas. A coluna sem nenhum valor, toda nula ou só com listas vazias, sai com o tipo da anotação quando ela é `str`, `int`, `float`, `bool`, `date`, `datetime` ou lista de um deles: `String`, `Int64`, `Float64`, `Boolean`, `Date`, `Datetime` ou `List` de um desses. Com outra anotação, a coluna sem valor fica com o tipo que o polars infere. + ### Colunas Adicionadas | Coluna | Quando existe | Conteúdo | diff --git a/src/dataframeit/core.py b/src/dataframeit/core.py index 301857d4..f7fc2ae3 100644 --- a/src/dataframeit/core.py +++ b/src/dataframeit/core.py @@ -60,6 +60,7 @@ ORIGINAL_TYPE_POLARS_DF, TOKEN_COLUMNS, ConversionInfo, + _json_default, from_pandas, get_complex_fields, normalize_complex_columns, @@ -1275,6 +1276,7 @@ def dataframeit( # noqa: C901, PLR0912, PLR0913, PLR0915, PLR0917 (API pública if not expected_columns: msg = "Modelo Pydantic não pode estar vazio" raise ValueError(msg) + annotations = {name: field.annotation for name, field in questions.model_fields.items()} # Cada campo extraído é gravado numa coluna de mesmo nome, que não pode ser # a do texto de entrada. @@ -1354,7 +1356,7 @@ def dataframeit( # noqa: C901, PLR0912, PLR0913, PLR0915, PLR0917 (API pública trace_mode=trace_mode, pydantic_model=questions, ) - return from_pandas(df_pandas, conversion_info, status_col) + return from_pandas(df_pandas, conversion_info, status_col, annotations) # Verificar conflitos de colunas existing_cols = [col for col in expected_columns if col in df_pandas.columns] @@ -1425,7 +1427,7 @@ def dataframeit( # noqa: C901, PLR0912, PLR0913, PLR0915, PLR0917 (API pública _apply_processed_values(df_pandas, processed_values) if complex_fields: normalize_complex_columns(df_pandas, complex_fields) - return from_pandas(df_pandas, conversion_info, status_col, expected_columns) + return from_pandas(df_pandas, conversion_info, status_col, annotations) # Os resultados são gravados por rótulo (df.at); com rótulo repetido, o # resultado de uma linha cairia em outra. Um checkpoint já concluído, que @@ -1600,7 +1602,7 @@ def dataframeit( # noqa: C901, PLR0912, PLR0913, PLR0915, PLR0917 (API pública ) # Retornar no formato original (remove colunas de status/erro se não houver erros) - return from_pandas(df_pandas, conversion_info, status_col, expected_columns) + return from_pandas(df_pandas, conversion_info, status_col, annotations) def _setup_columns( # noqa: C901, PLR0912, PLR0913 (uma família de colunas por recurso ligado) @@ -2016,19 +2018,6 @@ def _array_as_list(value: object) -> object: return value -def _json_default(value: object) -> object: - """O que json.dumps não serializa: numpy pelo tolist, Enum pelo valor, o resto texto. - - O tolist dá lista ao array e número ou bool do Python ao escalar do numpy, - que pelo texto voltaria "5" ou "True". O Enum vai pelo valor porque é o que - a validação do campo aceita de volta; o texto dele ("Cor.AZUL") não valida. - """ - if isinstance(value, Enum): - return value.value - tolist = getattr(value, "tolist", None) - return tolist() if callable(tolist) else str(value) - - def _structures_as_json(df: pd.DataFrame, columns: Collection[str] | None = None) -> pd.DataFrame: """Cópia com listas, dicts, tuplas e arrays do numpy serializados como JSON. diff --git a/src/dataframeit/utils.py b/src/dataframeit/utils.py index 8332dcdd..b8185f0f 100644 --- a/src/dataframeit/utils.py +++ b/src/dataframeit/utils.py @@ -11,6 +11,7 @@ from __future__ import annotations import ast +import datetime import functools import json import operator @@ -29,7 +30,7 @@ from pydantic import BaseModel if TYPE_CHECKING: - from collections.abc import Callable, Collection, Hashable + from collections.abc import Callable, Collection, Hashable, Mapping from polars import DataFrame as PolarsDataFrame from polars import Series as PolarsSeries @@ -205,7 +206,7 @@ def from_pandas( df: pd.DataFrame, conversion_info: ConversionInfo | bool, # noqa: FBT001 (bool aceito por retrocompatibilidade) status_col: str = "_dataframeit_status", - model_columns: Collection[str] = (), + model_columns: Mapping[str, object] = types.MappingProxyType({}), ) -> pd.DataFrame | PolarsDataFrame: """Converte DataFrame pandas de volta para o formato original. @@ -217,8 +218,9 @@ def from_pandas( conversion_info: ConversionInfo com metadados da conversão original, ou bool para retrocompatibilidade (True = era polars). status_col: Nome da coluna de status (o `status_column` do usuário). - model_columns: Campos do modelo que a execução gravou. Na saída polars, - só eles e as colunas de controle são convertidos do pandas. + model_columns: Anotação de cada campo do modelo que a execução gravou, pelo + nome. Na saída polars, só eles e as colunas de controle são + convertidos do pandas, e a anotação entra em `_to_polars`. Returns: Dados no formato original (DataFrame, Series, list ou dict). @@ -255,7 +257,7 @@ def from_pandas( # OverflowError, ComputeError), sai Object em vez de derrubar a saída. # O pânico do polars não é Exception e passa por aqui. try: - columns.append(pl.from_pandas(values)) + columns.append(_to_polars(values, model_columns.get(col))) except Exception: # noqa: BLE001 (a recusa muda de tipo com os valores) # Há versões do polars em que o construtor de Series inspeciona os # valores mesmo com dtype=pl.Object: quando o primeiro valor não nulo @@ -313,6 +315,84 @@ def _enums_by_value(values: pd.Series) -> pd.Series: ) +def _json_default(value: object) -> object: + """O que json.dumps não serializa: numpy pelo tolist, Enum pelo valor, o resto texto. + + O tolist dá lista ao array e número ou bool do Python ao escalar do numpy, + que pelo texto voltaria "5" ou "True". O Enum vai pelo valor porque é o que + a validação do campo aceita de volta; o texto dele ("Cor.AZUL") não valida. + """ + if isinstance(value, Enum): + return value.value + tolist = getattr(value, "tolist", None) + return tolist() if callable(tolist) else str(value) + + +def _without_optional(annotation: object) -> object: + """A anotação sem o Optional e sem o Annotated de dentro dele. + + A união de mais de um tipo além de None fica como está. + """ + args = [arg for arg in get_args(annotation) if arg is not type(None)] + if get_origin(annotation) in (typing.Union, types.UnionType) and len(args) == 1: + annotation = args[0] + if get_origin(annotation) is typing.Annotated: + annotation = get_args(annotation)[0] + return annotation + + +def _to_polars(values: pd.Series, annotation: object) -> PolarsSeries: + """Converte a coluna para polars; `annotation` é a do campo do modelo, ou None. + + O campo dict sai como texto JSON, como o checkpoint o grava. Como Struct, a + coluna teria em cada linha as chaves de todas, e a chave ausente não se + distinguiria da chave com valor nulo. O dict com chave que o JSON não + aceita, como um Enum, levanta TypeError, e `from_pandas` devolve a coluna + como Object. + + A coluna sem valor, toda nula ou só com listas vazias, não dá ao polars de + onde inferir o tipo, e o schema da saída mudaria com as respostas. Ela + recebe o tipo da anotação quando esta é um dos tipos de `scalars` ou uma + lista de um deles. A busca em `scalars` é por identidade, e não por chave + de dict, porque a anotação pode não ter hash, como o item de + list[Annotated[int, {...}]]. + """ + # Só é chamada com o polars instalado; o Any dispensa conferir o None a cada uso. + polars: Any = pl + scalars = ( + (str, polars.String), + (int, polars.Int64), + (float, polars.Float64), + (bool, polars.Boolean), + (datetime.date, polars.Date), + (datetime.datetime, polars.Datetime), + ) + annotation = _without_optional(annotation) + origin, args = get_origin(annotation), get_args(annotation) + if (origin or annotation) is dict: + # dtype=object porque a coluna só de NaN, como a do checkpoint CSV relido + # sem nenhuma resposta, seria inferida como float. + values = pd.Series( + [ + json.dumps(value, ensure_ascii=False, default=_json_default) + if isinstance(value, dict) + else value + for value in values + ], + name=values.name, + dtype=object, + ) + series = polars.from_pandas(values) + item = args[0] if origin is list and args else annotation + dtype = next((dtype for python_type, dtype in scalars if item is python_type), None) + # Da lista vazia o explode tira um nulo ou nada, conforme a versão do polars, + # e a coluna que não é lista ele devolve como está. + leaves = series.explode() + if dtype is None or leaves.null_count() < len(leaves): + return series + return series.cast(polars.List(dtype) if origin is list else dtype) + + def _untouched_columns( source: PolarsDataFrame | None, status_col: str, diff --git a/tests/test_entrada_e_saida.py b/tests/test_entrada_e_saida.py index 5ddc4bda..d94b00b4 100644 --- a/tests/test_entrada_e_saida.py +++ b/tests/test_entrada_e_saida.py @@ -1,18 +1,20 @@ """Entrada e saída de dataframeit(): colunas de controle, texto ausente, índice e avisos.""" +import json import linecache import logging +import typing import warnings from datetime import date, datetime, timedelta, timezone from decimal import Decimal from enum import Enum -from typing import Any, Optional, Union +from typing import Annotated, Any, Literal, Optional, Union from unittest.mock import patch import numpy as np import pandas as pd import pytest -from pydantic import BaseModel, ConfigDict, create_model +from pydantic import BaseModel, ConfigDict, Field, create_model from dataframeit import ProviderAbortError, ProviderError, core, dataframeit from dataframeit.core import _format_token_stats @@ -897,7 +899,7 @@ def test_retomada_polars_de_campo_bool_gravado_como_inteiro(modelo): "uniao_com_numero_primeiro": (Union[int, str], [1, "dois"]), "uniao_com_texto_primeiro": (Union[int, str], ["dois", 1]), "any_com_tipos_diferentes": (Any, [1, "x"]), - "dict_com_valores_de_tipos_diferentes": (dict[str, Any], [{"a": 1}, {"a": [1]}]), + "dict_com_chave_que_o_json_nao_aceita": (dict[Cor, int], [{Cor.AZUL: 1}, {Cor.VERDE: 2}]), "lista_de_uniao": (list[Union[int, str]], [[1, "a"], [2]]), "lista_de_enum": (list[Cor], [[Cor.AZUL], [Cor.VERDE]]), "inteiro_que_o_float_arredonda_com_float": (Union[int, float], [_ID_GRANDE, 2.5]), @@ -1086,3 +1088,252 @@ def test_saida_polars_com_coluna_object_e_retomada(): assert llm.call_count == 1 assert retomada["campo"].to_list() == [1, "dois", 3] + + +# ============================================================================= +# Saída polars: campo dict em texto JSON +# ============================================================================= + + +def test_campo_dict_sai_como_texto_json_na_saida_polars(): + pl = pytest.importorskip("polars") + modelo = create_model("ModeloDoCaso", campo=(dict[str, int], ...)) + respostas = [{"a": 1}, {"b": 2}, {}] + pandas, _ = _rodar( + pd.DataFrame({"texto": ["a", "b", "c"]}), + questions=modelo, + llm=_responde_por_linha("campo", respostas), + ) + + with warnings.catch_warnings(): + warnings.simplefilter("error") + resultado, _ = _rodar( + pl.DataFrame({"texto": ["a", "b", "c"]}), + questions=modelo, + llm=_responde_por_linha("campo", respostas), + ) + + assert resultado["campo"].dtype == pl.Utf8 + assert resultado["campo"].to_list() == ['{"a": 1}', '{"b": 2}', "{}"] + assert [json.loads(texto) for texto in resultado["campo"]] == pandas["campo"].tolist() + + +# Anotação de dict, as respostas de duas linhas e o texto de cada uma na saída. +_CAMPOS_DICT = { + "dict_sem_parametro": (dict, [{"a": 1}, {}], ['{"a": 1}', "{}"]), + "optional_com_linha_nula": (Optional[dict[str, int]], [{"a": 1}, None], ['{"a": 1}', None]), + "uniao_com_none": (dict[str, int] | None, [{"a": 1}, None], ['{"a": 1}', None]), + "annotated_dentro_de_optional": ( + Optional[Annotated[dict[str, int], Field(min_length=1)]], + [{"a": 1}, None], + ['{"a": 1}', None], + ), + "valores_de_tipos_diferentes": ( + dict[str, Any], + [{"a": 1}, {"a": [1]}], + ['{"a": 1}', '{"a": [1]}'], + ), + "texto_acentuado": (dict[str, str], [{"órgão": "saúde"}, {}], ['{"órgão": "saúde"}', "{}"]), + "valor_que_o_json_nao_serializa": ( + dict[str, date], + [{"a": date(2024, 1, 2)}, {}], + ['{"a": "2024-01-02"}', "{}"], + ), +} + + +@pytest.mark.parametrize("caso", _CAMPOS_DICT) +def test_campo_dict_sai_como_texto_json_em_cada_forma_de_anotacao(caso): + pl = pytest.importorskip("polars") + tipo, respostas, textos = _CAMPOS_DICT[caso] + + with warnings.catch_warnings(): + warnings.simplefilter("error") + resultado, _ = _rodar( + pl.DataFrame({"texto": ["a", "b"]}), + questions=create_model("ModeloDoCaso", campo=(tipo, ...)), + llm=_responde_por_linha("campo", respostas), + ) + + assert resultado["campo"].dtype == pl.Utf8 + assert resultado["campo"].to_list() == textos + + +def test_retomada_de_saida_polars_com_campo_dict_em_texto_json(): + pl = pytest.importorskip("polars") + modelo = create_model("ModeloDoCaso", campo=(dict[str, int], ...)) + with pytest.warns(UserWarning, match="interrompida"): + interrompida, _ = _rodar( + pl.DataFrame({"texto": ["a", "b", "c"]}), + questions=modelo, + llm=_responde_em_sequencia({"campo": {"a": 1}}, ProviderAbortError("limite")), + ) + assert interrompida["campo"].to_list() == ['{"a": 1}', None, None] + + retomada, llm = _rodar( + interrompida, + questions=modelo, + llm=_responde_em_sequencia({"campo": {"b": 2}}, {"campo": {}}), + resume=True, + ) + + assert llm.call_count == 2 + assert retomada["campo"].to_list() == ['{"a": 1}', '{"b": 2}', "{}"] + + +def test_modelo_aninhado_continua_struct_na_saida_polars(): + pl = pytest.importorskip("polars") + + resultado, _ = _rodar( + pl.DataFrame({"texto": ["a", "b"]}), + questions=create_model("ModeloDoCaso", campo=(Modelo, ...)), + llm=_responde_por_linha("campo", [{"x": "p"}, {"x": "q"}]), + ) + + assert resultado["campo"].dtype == pl.Struct({"x": pl.Utf8}) + assert resultado["campo"].to_list() == [{"x": "p"}, {"x": "q"}] + + +def test_saida_polars_ja_concluida_mantem_o_campo_dict_em_texto_json(): + pl = pytest.importorskip("polars") + entrada = pl.DataFrame( + {"texto": ["a"], "campo": ['{"a": 1}'], "_dataframeit_status": ["processed"]} + ) + + resultado, llm = _rodar( + entrada, + questions=create_model("ModeloDoCaso", campo=(dict[str, int], ...)), + resume=True, + ) + + llm.assert_not_called() + assert resultado["campo"].to_list() == ['{"a": 1}'] + + +# ============================================================================= +# Saída polars: tipo da coluna sem valor +# ============================================================================= + +# Anotação, um valor dela e o tipo polars da coluna, com ou sem valor. +_TIPOS_PELA_ANOTACAO = { + "str": (str, "a", lambda pl: pl.Utf8), + "int": (int, 1, lambda pl: pl.Int64), + "float": (float, 1.5, lambda pl: pl.Float64), + "bool": (bool, True, lambda pl: pl.Boolean), + "date": (date, date(2024, 1, 2), lambda pl: pl.Date), + "datetime": ( + datetime, + datetime.fromisoformat("2024-01-02T03:04"), + lambda pl: pl.Datetime("us"), + ), + "lista_de_str": (list[str], ["a"], lambda pl: pl.List(pl.Utf8)), + "lista_de_date": (list[date], [date(2024, 1, 2)], lambda pl: pl.List(pl.Date)), + "int_com_restricao": (Annotated[int, Field(ge=0)], 1, lambda pl: pl.Int64), +} + + +@pytest.mark.parametrize("caso", _TIPOS_PELA_ANOTACAO) +def test_coluna_toda_nula_sai_com_o_tipo_da_anotacao(caso): + pl = pytest.importorskip("polars") + tipo, valor, esperado = _TIPOS_PELA_ANOTACAO[caso] + modelo = create_model("ModeloDoCaso", campo=(Optional[tipo], None)) + entrada = pl.DataFrame({"texto": ["a", "b"]}) + + com_valor, _ = _rodar( + entrada, questions=modelo, llm=_responde_por_linha("campo", [valor, None]) + ) + sem_valor, _ = _rodar(entrada, questions=modelo, llm=_responde(campo=None)) + + assert com_valor["campo"].dtype == esperado(pl) + assert sem_valor["campo"].dtype == esperado(pl) + assert sem_valor["campo"].to_list() == [None, None] + + +@pytest.mark.parametrize("respostas", [[[], []], [[], None]]) +def test_lista_so_com_listas_vazias_sai_com_o_tipo_da_anotacao(respostas): + pl = pytest.importorskip("polars") + + resultado, _ = _rodar( + pl.DataFrame({"texto": ["a", "b"]}), + questions=create_model("ModeloDoCaso", campo=(Optional[list[str]], None)), + llm=_responde_por_linha("campo", respostas), + ) + + assert resultado["campo"].dtype == pl.List(pl.Utf8) + assert resultado["campo"].to_list() == respostas + + +# A resposta do dublê não passa pela validação do modelo, e um inteiro no campo +# float mostra que a coluna com algum valor fica com o tipo dos valores. +@pytest.mark.parametrize( + ("tipo", "respostas", "esperado"), + [ + (float, [7, None], lambda pl: pl.Int64), + (list[float], [[7], []], lambda pl: pl.List(pl.Int64)), + ], +) +def test_coluna_com_algum_valor_fica_com_o_tipo_inferido(tipo, respostas, esperado): + pl = pytest.importorskip("polars") + + resultado, _ = _rodar( + pl.DataFrame({"texto": ["a", "b"]}), + questions=create_model("ModeloDoCaso", campo=(Optional[tipo], None)), + llm=_responde_por_linha("campo", respostas), + ) + + assert resultado["campo"].dtype == esperado(pl) + assert resultado["campo"].to_list() == respostas + + +# Anotações que não dão tipo à coluna sem valor, com respostas que não trazem valor. +_ANOTACOES_SEM_TIPO_PARA_A_COLUNA = { + "decimal": (Optional[Decimal], [None, None]), + "literal": (Optional[Literal[1, 2]], [None, None]), + "enum": (Optional[Cor], [None, None]), + "uniao_de_dois_tipos": (Union[int, str, None], [None, None]), + "modelo_aninhado": (Optional[Modelo], [None, None]), + "tupla": (Optional[tuple[int, str]], [None, None]), + "lista_de_decimal": (list[Decimal], [[], []]), + "lista_sem_item": (typing.List, [[], []]), # noqa: UP006 (tem origem list e nenhum argumento) + "lista_de_item_sem_hash": (list[Annotated[int, {"unidade": "dias"}]], [[], []]), +} + + +@pytest.mark.parametrize("caso", _ANOTACOES_SEM_TIPO_PARA_A_COLUNA) +def test_coluna_sem_valor_de_outra_anotacao_fica_com_o_tipo_inferido(caso): + pl = pytest.importorskip("polars") + tipo, respostas = _ANOTACOES_SEM_TIPO_PARA_A_COLUNA[caso] + + with warnings.catch_warnings(): + warnings.simplefilter("error") + resultado, _ = _rodar( + pl.DataFrame({"texto": ["a", "b"]}), + questions=create_model("ModeloDoCaso", campo=(tipo, ...)), + llm=_responde_por_linha("campo", respostas), + ) + + assert resultado["campo"].dtype == pl.from_pandas(pd.Series(respostas, dtype=object)).dtype + assert resultado["campo"].to_list() == respostas + + +def test_entrada_polars_vazia_sai_com_o_schema_da_entrada_com_linhas(): + pl = pytest.importorskip("polars") + modelo = create_model( + "ModeloDoCaso", n=(Optional[int], None), tags=(list[str], ...), campo=(dict[str, int], ...) + ) + com_linhas, _ = _rodar( + pl.DataFrame({"texto": ["a"]}), + questions=modelo, + llm=_responde(n=1, tags=["t"], campo={"a": 1}), + ) + + vazia, llm = _rodar(pl.DataFrame({"texto": pl.Series([], dtype=pl.Utf8)}), questions=modelo) + + llm.assert_not_called() + assert dict(vazia.schema) == dict(com_linhas.schema) + assert dict(vazia.schema) == { + "texto": pl.Utf8, + "n": pl.Int64, + "tags": pl.List(pl.Utf8), + "campo": pl.Utf8, + } diff --git a/tests/test_multiple_data_types.py b/tests/test_multiple_data_types.py index ac0439f1..1938228e 100644 --- a/tests/test_multiple_data_types.py +++ b/tests/test_multiple_data_types.py @@ -314,7 +314,16 @@ def test_campo_do_modelo_que_estava_na_entrada_sai_do_pandas(self): df["nota"] = [7.5] assert from_pandas(df, info)["nota"].to_list() == [7] - assert from_pandas(df, info, model_columns=["nota"])["nota"].to_list() == [7.5] + assert from_pandas(df, info, model_columns={"nota": float})["nota"].to_list() == [7.5] + + def test_campo_dict_so_com_nan_fica_com_o_tipo_inferido(self): + # O checkpoint CSV relido traz NaN na linha sem resposta. + df, info = to_pandas(pl.DataFrame({"texto": ["a"]})) + df["campo"] = pd.Series([float("nan")], dtype=object) + + result = from_pandas(df, info, model_columns={"campo": dict[str, int]}) + + assert result["campo"].dtype == pl.from_pandas(df["campo"]).dtype @pytest.mark.parametrize( "coluna", From f9949e12a251287c76afeb5520a5890556aa46a7 Mon Sep 17 00:00:00 2001 From: Bruno Oliveira Date: Tue, 13 Oct 2026 16:14:25 -0300 Subject: [PATCH 2/3] =?UTF-8?q?Campo=20dict=20sem=20nenhuma=20resposta=20s?= =?UTF-8?q?ai=20como=20texto=20na=20sa=C3=ADda=20polars?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A coluna de texto JSON do campo dict segue a regra do campo str quando não tem valor: sai String, e não o tipo que o polars infere de uma coluna só de nulos, como a do checkpoint CSV relido sem nenhuma resposta. --- src/dataframeit/utils.py | 5 ++--- tests/test_multiple_data_types.py | 5 +++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/dataframeit/utils.py b/src/dataframeit/utils.py index b8185f0f..aef84156 100644 --- a/src/dataframeit/utils.py +++ b/src/dataframeit/utils.py @@ -370,8 +370,6 @@ def _to_polars(values: pd.Series, annotation: object) -> PolarsSeries: annotation = _without_optional(annotation) origin, args = get_origin(annotation), get_args(annotation) if (origin or annotation) is dict: - # dtype=object porque a coluna só de NaN, como a do checkpoint CSV relido - # sem nenhuma resposta, seria inferida como float. values = pd.Series( [ json.dumps(value, ensure_ascii=False, default=_json_default) @@ -380,8 +378,9 @@ def _to_polars(values: pd.Series, annotation: object) -> PolarsSeries: for value in values ], name=values.name, - dtype=object, ) + # A coluna de texto JSON sem nenhuma resposta segue a regra do campo str. + annotation = str series = polars.from_pandas(values) item = args[0] if origin is list and args else annotation dtype = next((dtype for python_type, dtype in scalars if item is python_type), None) diff --git a/tests/test_multiple_data_types.py b/tests/test_multiple_data_types.py index 1938228e..69438041 100644 --- a/tests/test_multiple_data_types.py +++ b/tests/test_multiple_data_types.py @@ -316,14 +316,15 @@ def test_campo_do_modelo_que_estava_na_entrada_sai_do_pandas(self): assert from_pandas(df, info)["nota"].to_list() == [7] assert from_pandas(df, info, model_columns={"nota": float})["nota"].to_list() == [7.5] - def test_campo_dict_so_com_nan_fica_com_o_tipo_inferido(self): + def test_campo_dict_sem_nenhuma_resposta_sai_como_texto(self): # O checkpoint CSV relido traz NaN na linha sem resposta. df, info = to_pandas(pl.DataFrame({"texto": ["a"]})) df["campo"] = pd.Series([float("nan")], dtype=object) result = from_pandas(df, info, model_columns={"campo": dict[str, int]}) - assert result["campo"].dtype == pl.from_pandas(df["campo"]).dtype + assert result["campo"].dtype == pl.String + assert result["campo"].to_list() == [None] @pytest.mark.parametrize( "coluna", From d337bb645ecead5aa72e1b5ca88628ed78141aa1 Mon Sep 17 00:00:00 2001 From: Bruno Oliveira Date: Wed, 14 Oct 2026 09:43:32 -0300 Subject: [PATCH 3/3] Entrada polars vazia volta a manter o tipo da coluna do modelo MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit O retorno de entrada vazia passava as anotações a from_pandas, e a coluna do modelo que a entrada de 0 linhas já trazia era convertida de novo do pandas: um Struct ou um Decimal saía String, e um Datetime com fuso perdia o fuso. O retorno volta a ser o da main, e um teste prende o tipo da entrada no Struct e no Datetime com fuso. A referência deixa de dizer que json.loads devolve o dict de cada linha: o texto é o que o checkpoint grava, e o valor que o JSON não representa sai como texto. Ela também diz o tipo da coluna do modelo na entrada sem linhas. A entrada da #219 no CHANGELOG passa para "Alterado", porque troca o tipo de uma coluna que o código do usuário lia. _to_polars devolve a série antes do explode quando a anotação não dá tipo à coluna. O caso do Enum como valor de dict entra no teste parametrizado do campo dict. --- CHANGELOG.md | 2 +- docs/en/reference/api.md | 2 +- docs/reference/api.md | 2 +- src/dataframeit/core.py | 2 +- src/dataframeit/utils.py | 4 +++- tests/test_entrada_e_saida.py | 45 +++++++++++++++++------------------ 6 files changed, 29 insertions(+), 28 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 392895e3..a9beeb43 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,11 +35,11 @@ e este projeto adere ao [Versionamento Semântico](https://semver.org/lang/pt-BR - Com `search_per_field=True`, o modo por grupo também devolve `search_provider` no uso, e o modo por campo devolve os campos na ordem do modelo, como o modo por grupo. O campo pulado por condição é registrado uma vez no log, e não duas. - O Python mínimo sobe para 3.11. O 3.10 saiu de suporte em outubro de 2026, e as versões atuais do numpy e do pandas já não o aceitam (#229). - O piso do pyarrow nos extras `polars` e `all` sobe para 10.0.1, a primeira versão com wheel para Python 3.11 (#229). +- Na saída polars, um campo `dict` sai como texto JSON, como no checkpoint. Antes, saía `Struct` com as chaves de todas as linhas: a chave ausente ficava igual à chave com valor nulo, e o dict vazio deixava de ser vazio. A coluna de um campo sem nenhum valor, toda nula ou só com listas vazias, sai com o tipo da anotação quando ela é `str`, `int`, `float`, `bool`, `date`, `datetime` ou lista de um deles. Antes, o tipo vinha só das respostas: um `Optional[int]` todo nulo saía `String`, e um `list[str]` só com listas vazias, `List(Null)` (#219). ### Corrigido - Com entrada polars, a saída deixa de falhar depois de todas as chamadas quando um campo do modelo guarda valores que o polars não reúne numa coluna tipada, como `Union[int, str]` com número numa linha e texto em outra, `tuple[int, str]`, `Any` com tipos diferentes, lista de `Enum`, inteiro acima de 2**53 ao lado de float, inteiro acima do `Int64` e `datetime` com fuso de deslocamento, como `+05:30`. A coluna sai com o tipo `Object`, com os objetos Python que a entrada pandas devolveria, e um aviso a nomeia. O campo `Enum` sai pelo valor do membro, inclusive na coluna `Object` (#218). -- Na saída polars, um campo `dict` sai como texto JSON, como no checkpoint. Antes, saía `Struct` com as chaves de todas as linhas: a chave ausente ficava igual à chave com valor nulo, e o dict vazio deixava de ser vazio. A coluna de um campo sem nenhum valor, toda nula ou só com listas vazias, sai com o tipo da anotação quando ela é `str`, `int`, `float`, `bool`, `date`, `datetime` ou lista de um deles. Antes, o tipo vinha só das respostas: um `Optional[int]` todo nulo saía `String`, e um `list[str]` só com listas vazias, `List(Null)` (#219). - A coluna `_trace` do agente único fica com as demais colunas de trace na ordem da saída, depois dos campos do modelo. Antes, numa entrada que já a trazia, os campos novos saíam depois dela. - Com entrada polars, as colunas que a execução não escreve saem do DataFrame de entrada, com o tipo e os valores de lá. Antes, a saída inteira era convertida de volta do pandas: um `Int64` com nulo voltava `Float64`, o que arredondava ids acima de 2**53, um `Date` voltava `Datetime`, e `Enum`, `Decimal`, `Array`, `Null` e o inteiro dentro de `Struct` e `List` também trocavam de tipo; uma coluna `Object` também trocava de tipo ou, com objetos que o Arrow não infere, levantava `ArrowInvalid` depois das chamadas (#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). diff --git a/docs/en/reference/api.md b/docs/en/reference/api.md index 5ad34e58..f5e17f72 100644 --- a/docs/en/reference/api.md +++ b/docs/en/reference/api.md @@ -115,7 +115,7 @@ Returns data in the same format as input with extracted columns added. With polars input, the input columns that are not model fields and do not have the name of one of the [added columns](#added-columns) come back as they went in, with the same type and the same values. In polars output, an `Enum` field comes out as the member's value. When the answers of a field do not fit a typed polars column, such as a `Union[int, str]` that received a number in one row and text in another, the column comes out with the `Object` type, holding the Python objects the model returned, except an `Enum` member outside a list or dict, which comes out as its value there too, and a warning names the columns; convert them before writing to a file. -A field annotated as `dict`, such as `dict[str, int]` or `Optional[dict[str, Any]]`, comes out in polars output as JSON text, with null where the row has no answer; `json.loads` returns the dict of each row. The type of the other model columns comes from the answers. A column with no value at all, all null or holding only empty lists, comes out with the type of the annotation when it is `str`, `int`, `float`, `bool`, `date`, `datetime`, or a list of one of them: `String`, `Int64`, `Float64`, `Boolean`, `Date`, `Datetime`, or a `List` of one of these. With any other annotation, the column with no value keeps the type polars infers. +A field annotated as `dict`, such as `dict[str, int]` or `Optional[dict[str, Any]]`, comes out in polars output as JSON text, with null where the row has no answer. The text is the same the checkpoint writes, and values that JSON does not represent come out in it as text. The type of the other model columns comes from the answers. A column with no value at all, all null or holding only empty lists, comes out with the type of the annotation when it is `str`, `int`, `float`, `bool`, `date`, `datetime`, or a list of one of them: `String`, `Int64`, `Float64`, `Boolean`, `Date`, `Datetime`, or a `List` of one of these. With any other annotation, the column with no value keeps the type polars infers. With input that has no rows, a model column the input already has keeps its type, and one it does not have comes out as `String`. ### Added Columns diff --git a/docs/reference/api.md b/docs/reference/api.md index 7a466b8d..b562df9f 100644 --- a/docs/reference/api.md +++ b/docs/reference/api.md @@ -115,7 +115,7 @@ Retorna dados no mesmo formato da entrada com colunas extraídas adicionadas. Com entrada polars, as colunas de entrada que não são campo do modelo nem têm nome de uma das [colunas adicionadas](#colunas-adicionadas) voltam como entraram, com o mesmo tipo e os mesmos valores. Na saída polars, um campo `Enum` sai pelo valor do membro. Quando as respostas de um campo não cabem numa coluna tipada do polars, como um `Union[int, str]` que recebeu número numa linha e texto em outra, a coluna sai com o tipo `Object`, com os objetos Python que o modelo devolveu, exceto o membro de `Enum` fora de lista ou dict, que também nela sai pelo valor, e um aviso diz quais colunas são; converta-as antes de gravar em arquivo. -Um campo anotado como `dict`, como `dict[str, int]` ou `Optional[dict[str, Any]]`, sai na saída polars como texto JSON, com nulo onde a linha não tem resposta; `json.loads` devolve o dict de cada linha. O tipo das demais colunas do modelo vem das respostas. A coluna sem nenhum valor, toda nula ou só com listas vazias, sai com o tipo da anotação quando ela é `str`, `int`, `float`, `bool`, `date`, `datetime` ou lista de um deles: `String`, `Int64`, `Float64`, `Boolean`, `Date`, `Datetime` ou `List` de um desses. Com outra anotação, a coluna sem valor fica com o tipo que o polars infere. +Um campo anotado como `dict`, como `dict[str, int]` ou `Optional[dict[str, Any]]`, sai na saída polars como texto JSON, com nulo onde a linha não tem resposta. O texto é o mesmo que o checkpoint grava, e os valores que o JSON não representa saem nele como texto. O tipo das demais colunas do modelo vem das respostas. A coluna sem nenhum valor, toda nula ou só com listas vazias, sai com o tipo da anotação quando ela é `str`, `int`, `float`, `bool`, `date`, `datetime` ou lista de um deles: `String`, `Int64`, `Float64`, `Boolean`, `Date`, `Datetime` ou `List` de um desses. Com outra anotação, a coluna sem valor fica com o tipo que o polars infere. Na entrada sem linhas, a coluna do modelo que a entrada já traz mantém o tipo dela, e a que ela não traz sai `String`. ### Colunas Adicionadas diff --git a/src/dataframeit/core.py b/src/dataframeit/core.py index f7fc2ae3..728a09b8 100644 --- a/src/dataframeit/core.py +++ b/src/dataframeit/core.py @@ -1356,7 +1356,7 @@ def dataframeit( # noqa: C901, PLR0912, PLR0913, PLR0915, PLR0917 (API pública trace_mode=trace_mode, pydantic_model=questions, ) - return from_pandas(df_pandas, conversion_info, status_col, annotations) + return from_pandas(df_pandas, conversion_info, status_col) # Verificar conflitos de colunas existing_cols = [col for col in expected_columns if col in df_pandas.columns] diff --git a/src/dataframeit/utils.py b/src/dataframeit/utils.py index aef84156..71362aaa 100644 --- a/src/dataframeit/utils.py +++ b/src/dataframeit/utils.py @@ -384,10 +384,12 @@ def _to_polars(values: pd.Series, annotation: object) -> PolarsSeries: series = polars.from_pandas(values) item = args[0] if origin is list and args else annotation dtype = next((dtype for python_type, dtype in scalars if item is python_type), None) + if dtype is None: + return series # Da lista vazia o explode tira um nulo ou nada, conforme a versão do polars, # e a coluna que não é lista ele devolve como está. leaves = series.explode() - if dtype is None or leaves.null_count() < len(leaves): + if leaves.null_count() < len(leaves): return series return series.cast(polars.List(dtype) if origin is list else dtype) diff --git a/tests/test_entrada_e_saida.py b/tests/test_entrada_e_saida.py index d94b00b4..f16c2f72 100644 --- a/tests/test_entrada_e_saida.py +++ b/tests/test_entrada_e_saida.py @@ -844,6 +844,27 @@ def test_entrada_polars_vazia_mantem_o_tipo_das_colunas(): assert resultado["dia"].dtype == pl.Date +def test_entrada_polars_vazia_mantem_o_tipo_da_coluna_do_modelo_que_ja_traz(): + # Sem linhas não há resposta a converter, e a coluna do modelo sai da + # entrada; convertida de volta do pandas, ela perderia o Struct e o fuso. + pl = pytest.importorskip("polars") + modelo = create_model( + "ModeloDoCaso", sub=(Optional[Modelo], None), quando=(Optional[datetime], None) + ) + entrada = pl.DataFrame( + { + "texto": pl.Series([], dtype=pl.Utf8), + "sub": pl.Series([], dtype=pl.Struct({"x": pl.Utf8})), + "quando": pl.Series([], dtype=pl.Datetime("us", "UTC")), + } + ) + + resultado, llm = _rodar(entrada, questions=modelo) + + llm.assert_not_called() + assert dict(resultado.schema) == dict(entrada.schema) + + def test_retomada_de_saida_polars_traz_os_tokens_das_linhas_novas(): pl = pytest.importorskip("polars") uso = {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5} @@ -1139,6 +1160,7 @@ def test_campo_dict_sai_como_texto_json_na_saida_polars(): [{"a": date(2024, 1, 2)}, {}], ['{"a": "2024-01-02"}', "{}"], ), + "enum_como_valor": (dict[str, Cor], [{"a": Cor.AZUL}, {}], ['{"a": "azul"}', "{}"]), } @@ -1314,26 +1336,3 @@ def test_coluna_sem_valor_de_outra_anotacao_fica_com_o_tipo_inferido(caso): assert resultado["campo"].dtype == pl.from_pandas(pd.Series(respostas, dtype=object)).dtype assert resultado["campo"].to_list() == respostas - - -def test_entrada_polars_vazia_sai_com_o_schema_da_entrada_com_linhas(): - pl = pytest.importorskip("polars") - modelo = create_model( - "ModeloDoCaso", n=(Optional[int], None), tags=(list[str], ...), campo=(dict[str, int], ...) - ) - com_linhas, _ = _rodar( - pl.DataFrame({"texto": ["a"]}), - questions=modelo, - llm=_responde(n=1, tags=["t"], campo={"a": 1}), - ) - - vazia, llm = _rodar(pl.DataFrame({"texto": pl.Series([], dtype=pl.Utf8)}), questions=modelo) - - llm.assert_not_called() - assert dict(vazia.schema) == dict(com_linhas.schema) - assert dict(vazia.schema) == { - "texto": pl.Utf8, - "n": pl.Int64, - "tags": pl.List(pl.Utf8), - "campo": pl.Utf8, - }