Colunas de entrada polars saem do DataFrame original

CHANGELOG.md

@@ -35,6 +35,8 @@ e este projeto adere ao [Versionamento Semântico](https://semver.org/lang/pt-BR
 
 ### Corrigido
 
+- 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).
 - 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/core.py

@@ -1425,7 +1425,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)
+        return from_pandas(df_pandas, conversion_info, status_col, expected_columns)
 
     # 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 +1600,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)
+    return from_pandas(df_pandas, conversion_info, status_col, expected_columns)
 
 
 def _setup_columns(  # noqa: C901, PLR0912, PLR0913 (uma família de colunas por recurso ligado)

src/dataframeit/utils.py

@@ -28,9 +28,10 @@
 from pydantic import BaseModel
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Hashable
+    from collections.abc import Callable, Collection, Hashable
 
     from polars import DataFrame as PolarsDataFrame
+    from polars import Series as PolarsSeries
 
 # Import opcional de Polars
 try:
@@ -90,6 +91,8 @@ 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)
+    # DataFrame polars de entrada, de onde a saída tira as colunas que a execução não escreve
+    polars_frame: PolarsDataFrame | None = None
 
 
 def parse_json(resposta: str) -> dict:
@@ -152,7 +155,9 @@ 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 data.to_pandas(), ConversionInfo(
+            original_type=ORIGINAL_TYPE_POLARS_DF, polars_frame=data
+        )
 
     # pandas Series
     if isinstance(data, pd.Series):
@@ -199,6 +204,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] = (),
 ) -> pd.DataFrame | PolarsDataFrame:
     """Converte DataFrame pandas de volta para o formato original.
 
@@ -210,6 +216,8 @@ 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.
 
     Returns:
         Dados no formato original (DataFrame, Series, list ou dict).
@@ -245,51 +253,73 @@ 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)
+        untouched = _untouched_columns(conversion_info.polars_frame, status_col, model_columns)
+        converted = pl.from_pandas(df[[col for col in df.columns if col not in untouched]])
+        return pl.DataFrame(
+            [untouched[col] if col in untouched else converted[col] for col in df.columns]
+        )
     return df
 
 
-def _reorder_columns(df: pd.DataFrame, status_col: str = "_dataframeit_status") -> pd.DataFrame:
-    """Reordena colunas para que internas fiquem no final.
+def _untouched_columns(
+    source: PolarsDataFrame | None,
+    status_col: str,
+    model_columns: Collection[str],
+) -> dict[str, PolarsSeries]:
+    """Colunas da entrada polars que a execução não escreve, pelo nome.
 
-    Ordem final das colunas:
-    1. Colunas do usuário (originais + campos do modelo)
-    2. Colunas de trace (_trace_*)
-    3. Colunas de busca (_search_credits)
-    4. Colunas de tokens (_input_tokens, _cached_input_tokens, _output_tokens,
-       _reasoning_tokens)
-    5. Colunas de controle (_dataframeit_status, _error_details)
+    São as colunas do usuário que não são campo do modelo. A saída as tira de
+    `source`, o DataFrame polars de entrada: convertidas de volta do pandas,
+    elas trocariam de tipo (Date por Datetime, Enum por Categorical) e trariam
+    arredondado o inteiro acima de 2**53 que tem nulo na coluna ou está dentro
+    de Struct e List, porque o `to_pandas()` do polars o leva a float. As linhas da saída estão na ordem das de `source`, porque a
+    execução grava por rótulo e não remove nem reordena linhas.
 
-    Args:
-        df: DataFrame a reordenar.
-        status_col: Nome da coluna de status.
+    Sem `source` (entrada polars.Series, que só traz o texto), todas as colunas
+    são da execução.
+    """
+    if source is None:
+        return {}
+    return {
+        col: source[col]
+        for col in source.columns
+        if col not in model_columns and _column_rank(col, status_col) == _USER_COLUMN
+    }
 
-    Returns:
-        DataFrame com colunas reordenadas.
+
+# Posição da coluna do usuário em _column_rank; as demais são as que a execução acrescenta.
+_USER_COLUMN = (0, 0)
+
+
+def _column_rank(col: Hashable, status_col: str) -> tuple[int, int]:
+    """Família da coluna e posição dentro dela, na ordem em que a saída as traz.
+
+    1. Colunas do usuário (originais + campos do modelo)
+    2. Colunas de trace (_trace e _trace_*)
+    3. Colunas de busca (_search_credits)
+    4. Colunas de tokens, na ordem de TOKEN_COLUMNS
+    5. Colunas de controle (status e _error_details)
     """
-    # Separar colunas por categoria
-    user_cols = []
-    trace_cols = []
-    search_cols = []
-    token_cols = [col for col in TOKEN_COLUMNS if col in df.columns]
-    status_cols = []
+    if col == "_trace" or str(col).startswith("_trace_"):
+        return (1, 0)
+    if col == "_search_credits":
+        return (2, 0)
+    if col in TOKEN_COLUMNS:
+        return (3, TOKEN_COLUMNS.index(col))
+    if col in (status_col, "_error_details"):
+        return (4, 0)
+    return _USER_COLUMN
 
-    for col in df.columns:
-        if str(col).startswith("_trace_"):
-            trace_cols.append(col)
-        elif col == "_search_credits":
-            search_cols.append(col)
-        elif col in TOKEN_COLUMNS:
-            continue
-        elif col in (status_col, "_error_details"):
-            status_cols.append(col)
-        else:
-            user_cols.append(col)
 
-    # Ordem final: user -> trace -> search -> tokens -> status
-    new_order = user_cols + trace_cols + search_cols + token_cols + status_cols
+def _reorder_columns(df: pd.DataFrame, status_col: str = "_dataframeit_status") -> pd.DataFrame:
+    """Reordena as colunas para que as da execução fiquem no final.
 
-    return df[new_order]
+    A ordem é a de `_column_rank`; dentro de cada família, menos a de tokens,
+    as colunas ficam na ordem em que estavam. A seleção é por posição, para que
+    um rótulo repetido não traga a coluna mais de uma vez.
+    """
+    order = sorted(range(df.shape[1]), key=lambda i: _column_rank(df.columns[i], status_col))
+    return df.iloc[:, order]
 
 
 # =============================================================================