process_text_na_dataframe corrupts data due to column reordering by OrdinalEncoder

[shaigilat]

Describe the bug

The function process_text_na_dataframe in tabpfn/utils.py calculates the indices of string columns based on the input DataFrame X. However, it applies the NA mask using these indices to X_encoded, which is the output of a ColumnTransformer (the ord_encoder).

Since ColumnTransformer reorders columns (typically grouping categorical features first), the indices from the input X do not match the indices in X_encoded. This results in:

Silent Data Corruption: The NA mask is applied to the wrong columns (often overwriting valid numerical data with NaN).

Steps/Code to Reproduce

import pandas as pd
from tabpfn import TabPFNClassifier
url = "https://raw.githubusercontent.com/datasciencedojo/datasets/master/titanic.csv"
print(f"Loading dataset from: {url}")
df = pd.read_csv(url)

# Select a mix of features to match your requirements:
# Select a mix of features (Numeric, Categorical, some with NaNs)
features = ['Pclass', 'Sex', 'Age', 'Embarked', 'Ticket', 'Name', 'Cabin']
target = 'Survived'

X = df[features]
y = df[target]
clf = TabPFNClassifier(device="auto", n_estimators=2)
clf.fit(X, y)

print(X.head())
_, X_processed, y_processed = clf._initialize_dataset_preprocessing(X, y, rng=42)
print(X_processed[:5])

Expected Results

The model should correctly handle missing values in string columns and encode the dataset without altering valid numerical data in other columns. The resulting training matrix passed to the model should contain all original numerical feature values.

Actual Results

The code produces silent data corruption. It does not crash, but it overwrites valid numerical data with NaN values.

This happens because the NA mask is calculated using indices from the input DataFrame but applied to the output encoded array. Since ColumnTransformer reorders columns (moving categorical features to the front), the mask is applied to the wrong columns. Consequently, numerical features that have shifted into those indices are erroneously overwritten with NaN.

Evidence: In the Titanic dataset example (see images), the Age or Fare numerical columns (which appear at the end of the encoded array) are overwritten with NaNs because they now occupy the indices that originally belonged to string columns with missing values (like Cabin).

Versions

Collecting system and dependency information...
PyTorch version: 2.9.1
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 15.7.3 (arm64)
GCC version: Could not collect
Clang version: 17.0.0 (clang-1700.6.3.2)
CMake version: version 3.26.4
Libc version: N/A

Python version: 3.12.6 (v3.12.6:a4a2d2b0d85, Sep  6 2024, 16:08:03) [Clang 13.0.0 (clang-1300.0.29.30)] (64-bit runtime)
Python platform: macOS-15.7.3-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Pro

Dependency Versions:
--------------------
tabpfn: 6.3.0
torch: 2.9.1
numpy: 2.4.0
scipy: 1.16.3
pandas: 2.3.3
scikit-learn: 1.7.2
typing_extensions: 4.15.0
einops: 0.8.1
huggingface-hub: 1.2.4

[bejaeger]

Thanks for reporting this issue @shaigilat

We are aware of it and will implement a fix to the main branch soon.

As a temporary workaround you can do the following:

You can adapt the get_ordinal_encoder() function so it returns a OrderPreservingColumnTransformer instead of a ColumnTransformer. The latter does not keep the column order which creates the issue you identified.