Subselect features in Dask Dataframe

I have a dask dataframe ddf with a matrix ddf['X'] and a list of indices indices. I want to select the features (columns) of ddf['X'] at the indices. My current implementation is

def subselect_variables(df):
    subset = df.iloc[:, indices]
    return subset
ddf_X = (
        ddf['X']
        .map_partitions(subselect_variables, meta={col: 'f4'for col in range(len(indices))})
    )
ddf_X.to_parquet(
    my_path,
    engine='pyarrow',
    schema=my_schema,
    write_metadata_file=True,
    row_group_size=my_row_group_size
    )

But it results in the error pandas.errors.IndexingError: Too many indexers. Can someone help?

I also tried to directly select the features

ddf_X = (
        ddf['X']
        .map_partitions(lambda df: df.iloc[:, indices], meta={col: 'f4'for col in range(len(indices))})
    )

Which resulted in the same error.
I also tried replacing : with slice(None), which also resulted in the same error.

Thanks for your suggestions! It led me in the right direction. Indeed, if ddf['X'] is a Series, it must be treated as 1-dimensional. What you also need to consider is the meta assignment and the external function, as you might run out of memory. Here is the solution that worked:

def subselect_variables(df):
    subset = df.map(lambda x: [x[i] for i in indices])
    return subset

ddf_X = (
        ddf['X']
        .map_partitions(subselect_series, meta=('X', 'f4'))
    )

To write it into a parquet file, you also need to cast it to a dask DataFrame, e.g. like

if isinstance(ddf_X, dd.Series):
        ddf_X = ddf_X.to_frame(name='X')

You are trying to index a one-dimensional thing (Series) with two dimensions or indexing. You may think it is 2D because each element is a list, but to pandas this just looks like a one-dimensional set of objects whose internals pandas knows nothing about. This has nothing to do with dask.

You need to figure out how you would do this indexing in pandas before trying it in dask. Pandas is not able to index into lists in an object series. They may be a way to do that more directly with arrow, awkward, (or even numpy?) or by first expanding out the lists into columns (explode?). However the following will work, if slowly and inefficiently.

Something like this grabs just the values you are after, but still keeps them in lists:

ddf_X = (
    ddf['X']
    .map(lambda value: [v for i, v in enumerate(value) if i in indices])
) 

Maybe you want

ddf_X = (
    ddf['X']
    .map_partitions(
         lambda s: pd.DataFrame(np.array(s.tolist())[:, indices]),
         meta={col: 'f4'for col in range(len(indices))
    )
)