Hi,
this behavior is expected.
dataset[column] loads the entire dataset column into memory (this is why this part seems so slow to you), so operating directly on this data will be faster than operating on the dataset created with load_dataset, which is memory-mapped, i. e., stored in an arrow file by default. However, in many cases, a dataset or dataset transforms are too big to fit in RAM, so having them in memory is not an option and we’d rather pay a penalty (in respect to read/write operations) of having it stored in a file.
If you are sure that you have enough RAM, you can load the dataset as follows to have it in memory:
load_dataset("squad", split="validation", keep_in_memory=True)