train_test_split - Невозможно работать с массивом Dask с неизвестными размерами фрагментов

Question

У меня есть набор данных для классификации текста, где я использовал паркет dask для экономии места на диске, но столкнулся с проблемой сейчас, когда я хочу разделить набор данных на поезд и протестировать, используя dask_ml.model_selection.train_test_split,

ddf = dd.read_parquet('/storage/data/cleaned')
y = ddf['category'].values
X = ddf.drop('category', axis=1).values
train, test = train_test_split(X, y, test_size=0.2)

В результате чего TypeError: Cannot operate on Dask array with unknown chunk sizes.

Спасибо за помощь.

Answer 1

Вот что я в итоге делал на данный момент:

ddf = dd.read_parquet('/storage/data/cleaned')
ddf = ddf.to_dask_array(lengths=True)
train, test = train_test_split(ddf, test_size=0.2)

Это создаст массив данных dask.array dask.array<array, shape=(3937987, 2), dtype=object, chunksize=(49701, 2)>