dask / dask/dask-ml

KFold cross validation fails with dask dataframes

Open
#956 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
951
Forks
262
PR merge metrics
No merged PRs in 30d

Description

**Describe the issue**:

`KFold.split` doesn't support dask dataframes. With the recent integrations of dask in e.g., xgboost, optuna, it would be very useful if it did. The error message acknowledges that dataframe are not supported and should be converted to dask arrays. With modern ML workflows, this isn't ideal since datasets commonly contain fields of many types (float, int, bool, categorical).

**Minimal Complete Verifiable Example**:

```python
import dask.dataframe as dd
from dask_ml.model_selection import train_test_split, KFold
ddf = dd.demo.make_timeseries()
train, test = train_test_split(ddf) # works

k_folder = KFold(n_splits=5)
for train, test in k_folder.split(ddf): # fails
pass
```

traceback

```
TypeError Traceback (most recent call last)
Cell In[12], line 1
----> 1 for train, test in k_folder.split(ddf):
2 pass

File ~/mambaforge/envs/ml-example/lib/python3.10/site-packages/dask_ml/model_selection/_split.py:241, in KFold.split(self, X, y, groups)
240 def split(self, X, y=None, groups=None):
--> 241 X = check_array(X)
242 n_samples = X.shape[0]
243 n_splits = self.n_splits

File ~/mambaforge/envs/ml-example/lib/python3.10/site-packages/dask_ml/utils.py:197, in check_array(array, accept_dask_array, accept_dask_dataframe, accept_unknown_chunks, accept_multiple_blocks, preserve_pandas_dataframe, remove_zero_chunks, *args, **kwargs)
195 elif isinstance(array, dd.DataFrame):
196 if not accept_dask_dataframe:
--> 197 raise TypeError(
198 "This estimator does not support dask dataframes. "
199 "This might be resolved with one of\n\n"
200 " 1. ddf.to_dask_array(lengths=True)\n"
201 " 2. ddf.to_dask_array() # may cause other issues because "
202 "of unknown chunk sizes"
203 )
204 # TODO: sample?
205 return array

TypeError: This estimator does not support dask dataframes. This might be resolved with one of

1. ddf.to_dask_array(lengths=True)
2. ddf.to_dask_array() # may cause other issues because of unknown chunk sizes
```

**Anything else we need to know?**:

We recently worked around this limitation with the following:

```python
def _make_cv(df, num_folds):
frac = [1 / num_folds]*num_folds
splits = df.random_split(frac, shuffle=True)
for i in range(num_folds):
train = [splits[j] for j in range(num_folds) if j != i]
test = splits[i]
yield train, test

for i, (train, test) in enumerate(_make_cv(ddf, n_splits)):
pass
```

Contributor guide

Open the contributing guide

Research direction

Reproduce the failure with the provided dask.dataframe example, then read dask_ml/model_selection/_split.py around KFold.split and dask_ml/utils.py around check_array. Determine how KFold can accept dataframe columns with mixed types while retaining its existing split behavior; done means the example runs without the dataframe rejection and regression coverage verifies the supported input.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data, machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.