# Querying column is slow for datasets with indices mapping

**URL:** <https://discuss.huggingface.co/t/querying-column-is-slow-for-datasets-with-indices-mapping/6180>\
**Category:** 🤗Datasets\
**Created:** [May 14, 2021, 11:39am UTC](https://discuss.huggingface.co/t/querying-column-is-slow-for-datasets-with-indices-mapping/6180 "2021-05-14T11:39:07Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![mozharovsky](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/mozharovsky/32/1480_2.png) [@mozharovsky](https://discuss.huggingface.co/u/mozharovsky)\
**Post date:** [May 14, 2021, 11:39am UTC](https://discuss.huggingface.co/t/querying-column-is-slow-for-datasets-with-indices-mapping/6180/1 "2021-05-14T11:39:07Z")

</div>

Hi,

I have a large dataset and wanna query values for a specific column for batch packing. It turns out that querying is super slow if I use datasets with indices mapping (e.g., after `train_test_split` or `select` operations).

The code below reproduces the issue:

```python
import time
import numpy as np
from datasets import Dataset

input_ids = np.random.randint(0, 60_000, (500_000, 128)).tolist()
length = np.random.randint(3, 128, (500_000)).tolist()

dataset = Dataset.from_dict({"input_ids": input_ids, "length": length})
dataset_dict = dataset.train_test_split(test_size=0.1)

# ---------------------------------------------------------------
# Original dataset
start = time.time()
_ = dataset["length"]
print(f"Operation took {time.time() - start:.2f} seconds")
# Operation took 0.15 seconds
# ---------------------------------------------------------------

# ---------------------------------------------------------------
# Dataset with indices mapping
start = time.time()
_ = dataset_dict["train"]["length"]
print(f"Operation took {time.time() - start:.2f} seconds")
# Operation took 5.74 seconds
# ---------------------------------------------------------------

```

It takes forever to load values for my 500+ Gb dataset, and I have to use `flatten_indices` on each dataset split to get deep copies with super fast querying performance. Anyway, `flatten_indices` is no faster and also takes too much time. So preprocessing becomes super painful to achieve acceptable performance.

Is there a way to achieve acceptable performance without flattening indices?

---

<div class="post-metadata">

**Author:** ![thomwolf](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/thomwolf/32/2198_2.png) [@thomwolf](https://discuss.huggingface.co/u/thomwolf)\
**Post date:** [May 15, 2021, 4:27pm UTC](https://discuss.huggingface.co/t/querying-column-is-slow-for-datasets-with-indices-mapping/6180/2 "2021-05-15T16:27:08Z")

</div>

Which version are you using? The latest versions give dramatic improvements on such querying operations.

---

<div class="post-metadata">

**Author:** ![mozharovsky](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/mozharovsky/32/1480_2.png) [@mozharovsky](https://discuss.huggingface.co/u/mozharovsky)\
**Post date:** [May 15, 2021, 8:43pm UTC](https://discuss.huggingface.co/t/querying-column-is-slow-for-datasets-with-indices-mapping/6180/3 "2021-05-15T20:43:07Z")

</div>

I’m using the latest published version (1.6.2).

---

<div class="post-metadata">

**Author:** ![lhoestq](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/lhoestq/32/52888_2.png) [@lhoestq](https://discuss.huggingface.co/u/lhoestq)\
**Post date:** [May 17, 2021, 8:56am UTC](https://discuss.huggingface.co/t/querying-column-is-slow-for-datasets-with-indices-mapping/6180/4 "2021-05-17T08:56:40Z")

</div>

Hi ! Since you used train\_test\_split, then the new train dataset is now shuffled. Querying a shuffled dataset takes some time since it has to re-order all the elements. So the bigger the dataset, the longer it takes to query all the data of a shuffled column.

If you want more speed, either don’t shuffle (`shuffle=False` in train\_test\_split), or use `flatten_indices` once and for all (it re-order the rows correctly).