🤗 Dataset Select Columns

Some datasets arrive with seventeen columns when you need three. 🤗 Dataset Select Columns is the whitelist answer to that: you type the columns you want to keep and everything else disappears. If Remove Columns is the "drop the heavy stuff" tool, this is the "give me exactly this, nothing else" tool - and for a wide table it's a lot easier to type the three you want than the fourteen you don't.

How it works

It calls datasets.Dataset.select_columns. Two inputs: the dataset, and columns - a comma-separated list of the ones to keep, like text, label. One output: the dataset restricted to exactly those columns, in the order you listed them (which is itself a nice side effect - you can reorder columns just by listing them in the order you want). Rows and values are untouched.

It works on both fully-loaded and streaming datasets, so you can slim a lazy IterableDataset before it's ever iterated - which matters when the columns you're dropping hold images or audio and every unnecessary decode is wasted work.

Remove vs. Select - pick your poison

The pack gives you both, and the choice is genuinely about your table's shape:

• Select Columns - know what you want (or the table is wide). Whitelist, type the keepers.
• Remove Columns - know what you don't want (or the table is narrow). Blacklist, type the drop list.

If you're materializing thousands of rows for a prompt or export loop, doing either one before To Data List / To LIST is one of the cheapest wins in the pack - the later columns get dropped, the smaller each row dict is.

Gotchas

If you name a column that doesn't exist, you get an error - no silent "keep nothing" surprises. And the operation is a projection, not a filter: it never touches rows, so if you wanted a subset of rows this is the wrong node (that's Filter / Select Rows / Take). It also won't reach inside a nested column to pick one field - Select operates on top-level columns, and unwrapping nested structure is what Flatten is for.

Installing it

Part of StableLlama/ComfyUI-huggingface_dataset. ComfyUI-Manager → search "Hugging Face dataset", or:

cd ComfyUI/custom_nodes
git clone https://github.com/StableLlama/ComfyUI-huggingface_dataset
pip install -r requirements.txt

Restart ComfyUI. Only runtime dependency is datasets, installed by Manager.