Support pyarrow large_list
State: closed

Allow Polars round trip by supporting pyarrow large list.

Fix #6834, fix #6984.

Supersede and close #4800, close #6835, close #6986.

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/datasets/pr_7019). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@albertvillanova really happy to see this fix.

Have you attempted to save a dataset to disk after this? I attempted to utilize your fix in a build from source, and while I can now successfully get a dataset object from a polars df containing a large list, I am getting the following error when attempting to save the resulting dataset to disk:
```
File "/Users/x/VSCodeProjects/HuggingFace/hf.py", line 9, in <module>
    dataset.save_to_disk("data/test.hf")
  File "/Users/x/VSCodeProjects/HuggingFace/datasets/src/datasets/arrow_dataset.py", line 1591, in save_to_disk
    for kwargs in kwargs_per_job:
  File "/Users/x/VSCodeProjects/HuggingFace/datasets/src/datasets/arrow_dataset.py", line 1568, in <genexpr>
    "shard": self.shard(num_shards=num_shards, index=shard_idx, contiguous=True),
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/x/VSCodeProjects/HuggingFace/datasets/src/datasets/arrow_dataset.py", line 4757, in shard
    return self.select(
           ^^^^^^^^^^^^
  File "/Users/x/VSCodeProjects/HuggingFace/datasets/src/datasets/arrow_dataset.py", line 567, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/x/VSCodeProjects/HuggingFace/datasets/src/datasets/fingerprint.py", line 482, in wrapper
    out = func(dataset, *args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/x/VSCodeProjects/HuggingFace/datasets/src/datasets/arrow_dataset.py", line 3892, in select
    return self._select_contiguous(start, length, new_fingerprint=new_fingerprint)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/x/VSCodeProjects/HuggingFace/datasets/src/datasets/arrow_dataset.py", line 567, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/x/VSCodeProjects/HuggingFace/datasets/src/datasets/fingerprint.py", line 482, in wrapper
    out = func(dataset, *args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/x/VSCodeProjects/HuggingFace/datasets/src/datasets/arrow_dataset.py", line 3955, in _select_contiguous
    return Dataset(
           ^^^^^^^^
  File "/Users/x/VSCodeProjects/HuggingFace/datasets/src/datasets/arrow_dataset.py", line 731, in __init__
    raise ValueError(
ValueError: External features info don't match the dataset:
Got
{'0': Value(dtype='int64', id=None), '1': Value(dtype='int64', id=None), '2': Value(dtype='int64', id=None), '3': Value(dtype='int64', id=None), '4': Value(dtype='int64', id=None), '5': Value(dtype='int64', id=None), '6': Value(dtype='int64', id=None), '7': Value(dtype='int64', id=None), '8': Value(dtype='int64', id=None), '9': Value(dtype='int64', id=None), '10': Value(dtype='int64', id=None), '11': Value(dtype='int64', id=None), '12': Value(dtype='int64', id=None), '13': Value(dtype='int64', id=None), '14': Value(dtype='int64', id=None), '15': Value(dtype='int64', id=None), '16': Value(dtype='int64', id=None), '17': Value(dtype='int64', id=None), '18': Value(dtype='int64', id=None), '19': Value(dtype='int64', id=None), 'A': Sequence(feature=Value(dtype='int64', id=None), length=-1, large=False, id=None), 'B': Sequence(feature=Value(dtype='int64', id=None), length=-1, large=False, id=None), 'C': Sequence(feature=Value(dtype='int64', id=None), length=-1, large=False, id=None), 'D': Sequence(feature=Value(dtype='int64', id=None), length=-1, large=False, id=None), '__index_level_0__': Value(dtype='int64', id=None)}
with type
struct<0: int64, 1: int64, 2: int64, 3: int64, 4: int64, 5: int64, 6: int64, 7: int64, 8: int64, 9: int64, 10: int64, 11: int64, 12: int64, 13: int64, 14: int64, 15: int64, 16: int64, 17: int64, 18: int64, 19: int64, A: list<item: int64>, B: list<item: int64>, C: list<item: int64>, D: list<item: int64>, __index_level_0__: int64>

but expected something like
{'0': Value(dtype='int64', id=None), '1': Value(dtype='int64', id=None), '2': Value(dtype='int64', id=None), '3': Value(dtype='int64', id=None), '4': Value(dtype='int64', id=None), '5': Value(dtype='int64', id=None), '6': Value(dtype='int64', id=None), '7': Value(dtype='int64', id=None), '8': Value(dtype='int64', id=None), '9': Value(dtype='int64', id=None), '10': Value(dtype='int64', id=None), '11': Value(dtype='int64', id=None), '12': Value(dtype='int64', id=None), '13': Value(dtype='int64', id=None), '14': Value(dtype='int64', id=None), '15': Value(dtype='int64', id=None), '16': Value(dtype='int64', id=None), '17': Value(dtype='int64', id=None), '18': Value(dtype='int64', id=None), '19': Value(dtype='int64', id=None), 'A': Sequence(feature=Value(dtype='int64', id=None), length=-1, large=True, id=None), 'B': Sequence(feature=Value(dtype='int64', id=None), length=-1, large=True, id=None), 'C': Sequence(feature=Value(dtype='int64', id=None), length=-1, large=True, id=None), 'D': Sequence(feature=Value(dtype='int64', id=None), length=-1, large=True, id=None), '__index_level_0__': Value(dtype='int64', id=None)}
with type
struct<0: int64, 1: int64, 2: int64, 3: int64, 4: int64, 5: int64, 6: int64, 7: int64, 8: int64, 9: int64, 10: int64, 11: int64, 12: int64, 13: int64, 14: int64, 15: int64, 16: int64, 17: int64, 18: int64, 19: int64, A: large_list<item: int64>, B: large_list<item: int64>, C: large_list<item: int64>, D: large_list<item: int64>, __index_level_0__: int64>
```

code to reproduce is actually 2 separate scripts below.

creating test data:
```
import pandas as pd
import numpy as np

df = pd.DataFrame(np.random.randint(0, 100000, size=(100000, 20)))
featureVector = np.random.randint(0, 100000, size=(100000, 1000)).tolist()

df['A'] = featureVector
df['B'] = featureVector
df['C'] = featureVector
df['D'] = featureVector

df.to_parquet('data/train_data.parquet', engine='pyarrow')
```

loading data, converting to HF dataset, attempting to save to disk
```
import datasets
import polars as pl

df = pl.read_parquet('data/train_data.parquet')

dataset = datasets.Dataset.from_polars(df)

dataset.save_to_disk("data/test.hf")
```

If this isn't the appropriate place to put this, let me know. Since it isn't merged yet I didn't think raising an issue was appropriate.

Thanks for your useful review comments, @dakotamurdock. 

I am investigating that issue to fix it in this PR.

Hi @albertvillanova thanks for your work! When is the fix planned to be released?

I tested your feature branch and managed to load from a polars dataframe with the large_list type, persist to disk, load and convert it again. Also asserted that they are both equal.

```
> print(df[:3])

shape: (3, 6)
┌─────────────────┬─────────────────┬──────────┬────────────────┬─────────────────┬────────────────┐
│ plain_text      ┆ title           ┆ language ┆ language_score ┆ plain_text_hash ┆ response_objec │
│ ---             ┆ ---             ┆ ---      ┆ ---            ┆ ---             ┆ t              │
│ str             ┆ str             ┆ str      ┆ f64            ┆ str             ┆ ---            │
│                 ┆                 ┆          ┆                ┆                 ┆ list[struct[3] │
│                 ┆                 ┆          ┆                ┆                 ┆ ]              │
╞═════════════════╪═════════════════╪══════════╪════════════════╪═════════════════╪════════════════╡
│ Royal fans      ┆ Prince Louis    ┆ en       ┆ 0.987661       ┆ 37e438dccb283d1 ┆ [{"Prince Loui │
│ delighted by    ┆ delights crowd  ┆          ┆                ┆ f3be2d9d4bb7ed3 ┆ s","waves","cr │
│ Prince…         ┆ wi…             ┆          ┆                ┆ …               ┆ ow…            │
│ There have been ┆ Reactions After ┆ en       ┆ 0.991371       ┆ 37fafbb69dfcfa5 ┆ [{"David       │
│ diverse reacti… ┆ Davido Alleged… ┆          ┆                ┆ 303d5e3e6917a35 ┆ Adedeji        │
│                 ┆                 ┆          ┆                ┆ …               ┆ Adeleke","is … │
│ Betfred will    ┆ Betfred to pay  ┆ en       ┆ 0.980579       ┆ 922e19e6f598e9b ┆ [{"Betfred","w │
│ pay a £3.25     ┆ £3.25 million   ┆          ┆                ┆ 14cdb6772829cbd ┆ ill            │
│ milli…          ┆ f…              ┆          ┆                ┆ …               ┆ pay","£3.25 …  │
└─────────────────┴─────────────────┴──────────┴────────────────┴─────────────────┴────────────────┘

> Dataset.from_polars(df).save_to_disk('./test')

Saving the dataset (1/1 shards): 100%|██████████| 14997/14997 [00:00<00:00, 225472.09 examples/s]

> another_df = load_from_disk('./test').to_polars()
> print(another_df[:3])

shape: (3, 6)
┌─────────────────┬─────────────────┬──────────┬────────────────┬─────────────────┬────────────────┐
│ plain_text      ┆ title           ┆ language ┆ language_score ┆ plain_text_hash ┆ response_objec │
│ ---             ┆ ---             ┆ ---      ┆ ---            ┆ ---             ┆ t              │
│ str             ┆ str             ┆ str      ┆ f64            ┆ str             ┆ ---            │
│                 ┆                 ┆          ┆                ┆                 ┆ list[struct[3] │
│                 ┆                 ┆          ┆                ┆                 ┆ ]              │
╞═════════════════╪═════════════════╪══════════╪════════════════╪═════════════════╪════════════════╡
│ Royal fans      ┆ Prince Louis    ┆ en       ┆ 0.987661       ┆ 37e438dccb283d1 ┆ [{"Prince Loui │
│ delighted by    ┆ delights crowd  ┆          ┆                ┆ f3be2d9d4bb7ed3 ┆ s","waves","cr │
│ Prince…         ┆ wi…             ┆          ┆                ┆ …               ┆ ow…            │
│ There have been ┆ Reactions After ┆ en       ┆ 0.991371       ┆ 37fafbb69dfcfa5 ┆ [{"David       │
│ diverse reacti… ┆ Davido Alleged… ┆          ┆                ┆ 303d5e3e6917a35 ┆ Adedeji        │
│                 ┆                 ┆          ┆                ┆ …               ┆ Adeleke","is … │
│ Betfred will    ┆ Betfred to pay  ┆ en       ┆ 0.980579       ┆ 922e19e6f598e9b ┆ [{"Betfred","w │
│ pay a £3.25     ┆ £3.25 million   ┆          ┆                ┆ 14cdb6772829cbd ┆ ill            │
│ milli…          ┆ f…              ┆          ┆                ┆ …               ┆ pay","£3.25 …  │
└─────────────────┴─────────────────┴──────────┴────────────────┴─────────────────┴────────────────┘

> another_df.equals(df)

True
```

This is indeed the error I was getting with datasets==2.19.1
```
Dataset.from_polars(df).save_to_disk('./test')
ValueError: Arrow type large_list<item: struct<entity1: large_string, relationship: large_string, entity2: large_string>> does not have a datasets dtype equivalent.
```

@EdoardoLuciani thanks for your feedback!
I think we should make a new release soon: last one was on June 13.
What do you think, @huggingface/datasets?
The only potential problem I see are the breaking changes once we remove all deprecated code...

Your issue was fixed, @dakotamurdock.

I am working in a big refactoring of the approach to support large_list: implement a new `LargeList` type instead of using `Sequence.large` attribute.

There are many feature-functions and most of them are not properly covered by tests.

I am adding tests and fixing these feature-functions.

I think this PR is ready for review, @huggingface/datasets.

<details>
<summary>Show benchmarks</summary>

PyArrow==8.0.0

<details>
<summary>Show updated benchmarks!</summary>

### Benchmark: benchmark_array_xd.json

| metric | read_batch_formatted_as_numpy after write_array2d | read_batch_formatted_as_numpy after write_flattened_sequence | read_batch_formatted_as_numpy after write_nested_sequence | read_batch_unformated after write_array2d | read_batch_unformated after write_flattened_sequence | read_batch_unformated after write_nested_sequence | read_col_formatted_as_numpy after write_array2d | read_col_formatted_as_numpy after write_flattened_sequence | read_col_formatted_as_numpy after write_nested_sequence | read_col_unformated after write_array2d | read_col_unformated after write_flattened_sequence | read_col_unformated after write_nested_sequence | read_formatted_as_numpy after write_array2d | read_formatted_as_numpy after write_flattened_sequence | read_formatted_as_numpy after write_nested_sequence | read_unformated after write_array2d | read_unformated after write_flattened_sequence | read_unformated after write_nested_sequence | write_array2d | write_flattened_sequence | write_nested_sequence |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.005640 / 0.011353 (-0.005713) | 0.003926 / 0.011008 (-0.007083) | 0.063103 / 0.038508 (0.024595) | 0.032088 / 0.023109 (0.008979) | 0.238615 / 0.275898 (-0.037283) | 0.268379 / 0.323480 (-0.055101) | 0.003146 / 0.007986 (-0.004840) | 0.002813 / 0.004328 (-0.001516) | 0.049681 / 0.004250 (0.045431) | 0.044577 / 0.037052 (0.007525) | 0.249782 / 0.258489 (-0.008708) | 0.282548 / 0.293841 (-0.011293) | 0.029986 / 0.128546 (-0.098560) | 0.012474 / 0.075646 (-0.063172) | 0.203347 / 0.419271 (-0.215925) | 0.035950 / 0.043533 (-0.007583) | 0.243410 / 0.255139 (-0.011729) | 0.267056 / 0.283200 (-0.016143) | 0.022086 / 0.141683 (-0.119597) | 1.145513 / 1.452155 (-0.306641) | 1.207583 / 1.492716 (-0.285133) |

### Benchmark: benchmark_getitem\_100B.json

| metric | get_batch_of\_1024\_random_rows | get_batch_of\_1024\_rows | get_first_row | get_last_row |
|--------|---|---|---|---|
| new / old (diff) | 0.095584 / 0.018006 (0.077578) | 0.304264 / 0.000490 (0.303774) | 0.000215 / 0.000200 (0.000015) | 0.000043 / 0.000054 (-0.000011) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.019460 / 0.037411 (-0.017952) | 0.062268 / 0.014526 (0.047742) | 0.074943 / 0.176557 (-0.101613) | 0.121657 / 0.737135 (-0.615478) | 0.075930 / 0.296338 (-0.220408) |

### Benchmark: benchmark_iterating.json

| metric | read 5000 | read 50000 | read_batch 50000 10 | read_batch 50000 100 | read_batch 50000 1000 | read_formatted numpy 5000 | read_formatted pandas 5000 | read_formatted tensorflow 5000 | read_formatted torch 5000 | read_formatted_batch numpy 5000 10 | read_formatted_batch numpy 5000 1000 | shuffled read 5000 | shuffled read 50000 | shuffled read_batch 50000 10 | shuffled read_batch 50000 100 | shuffled read_batch 50000 1000 | shuffled read_formatted numpy 5000 | shuffled read_formatted_batch numpy 5000 10 | shuffled read_formatted_batch numpy 5000 1000 |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.288975 / 0.215209 (0.073766) | 2.869610 / 2.077655 (0.791955) | 1.491057 / 1.504120 (-0.013063) | 1.384160 / 1.541195 (-0.157035) | 1.380977 / 1.468490 (-0.087513) | 0.723181 / 4.584777 (-3.861596) | 2.397960 / 3.745712 (-1.347752) | 2.899919 / 5.269862 (-2.369942) | 1.878714 / 4.565676 (-2.686962) | 0.078162 / 0.424275 (-0.346113) | 0.005115 / 0.007607 (-0.002493) | 0.337599 / 0.226044 (0.111555) | 3.367450 / 2.268929 (1.098522) | 1.823745 / 55.444624 (-53.620880) | 1.540528 / 6.876477 (-5.335949) | 1.546146 / 2.142072 (-0.595927) | 0.796927 / 4.805227 (-4.008300) | 0.134389 / 6.500664 (-6.366275) | 0.042298 / 0.075469 (-0.033172) |

### Benchmark: benchmark_map_filter.json

| metric | filter | map fast-tokenizer batched | map identity | map identity batched | map no-op batched | map no-op batched numpy | map no-op batched pandas | map no-op batched pytorch | map no-op batched tensorflow |
|--------|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.959687 / 1.841788 (-0.882101) | 11.505269 / 8.074308 (3.430961) | 9.631551 / 10.191392 (-0.559841) | 0.142301 / 0.680424 (-0.538123) | 0.013912 / 0.534201 (-0.520289) | 0.314940 / 0.579283 (-0.264343) | 0.263134 / 0.434364 (-0.171229) | 0.352966 / 0.540337 (-0.187372) | 0.440421 / 1.386936 (-0.946515) |

</details>
PyArrow==latest

<details>
<summary>Show updated benchmarks!</summary>

### Benchmark: benchmark_array_xd.json

| metric | read_batch_formatted_as_numpy after write_array2d | read_batch_formatted_as_numpy after write_flattened_sequence | read_batch_formatted_as_numpy after write_nested_sequence | read_batch_unformated after write_array2d | read_batch_unformated after write_flattened_sequence | read_batch_unformated after write_nested_sequence | read_col_formatted_as_numpy after write_array2d | read_col_formatted_as_numpy after write_flattened_sequence | read_col_formatted_as_numpy after write_nested_sequence | read_col_unformated after write_array2d | read_col_unformated after write_flattened_sequence | read_col_unformated after write_nested_sequence | read_formatted_as_numpy after write_array2d | read_formatted_as_numpy after write_flattened_sequence | read_formatted_as_numpy after write_nested_sequence | read_unformated after write_array2d | read_unformated after write_flattened_sequence | read_unformated after write_nested_sequence | write_array2d | write_flattened_sequence | write_nested_sequence |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.005878 / 0.011353 (-0.005475) | 0.003866 / 0.011008 (-0.007142) | 0.051347 / 0.038508 (0.012839) | 0.032662 / 0.023109 (0.009553) | 0.270701 / 0.275898 (-0.005197) | 0.345277 / 0.323480 (0.021797) | 0.004485 / 0.007986 (-0.003501) | 0.002782 / 0.004328 (-0.001546) | 0.048302 / 0.004250 (0.044051) | 0.040355 / 0.037052 (0.003303) | 0.285196 / 0.258489 (0.026707) | 0.320339 / 0.293841 (0.026499) | 0.032937 / 0.128546 (-0.095610) | 0.012298 / 0.075646 (-0.063348) | 0.061579 / 0.419271 (-0.357692) | 0.034129 / 0.043533 (-0.009403) | 0.265985 / 0.255139 (0.010846) | 0.302066 / 0.283200 (0.018867) | 0.018812 / 0.141683 (-0.122871) | 1.175705 / 1.452155 (-0.276450) | 1.197207 / 1.492716 (-0.295510) |

### Benchmark: benchmark_getitem\_100B.json

| metric | get_batch_of\_1024\_random_rows | get_batch_of\_1024\_rows | get_first_row | get_last_row |
|--------|---|---|---|---|
| new / old (diff) | 0.096076 / 0.018006 (0.078070) | 0.312793 / 0.000490 (0.312303) | 0.000228 / 0.000200 (0.000028) | 0.000053 / 0.000054 (-0.000002) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.022858 / 0.037411 (-0.014553) | 0.077160 / 0.014526 (0.062634) | 0.089742 / 0.176557 (-0.086815) | 0.130929 / 0.737135 (-0.606207) | 0.093431 / 0.296338 (-0.202907) |

### Benchmark: benchmark_iterating.json

| metric | read 5000 | read 50000 | read_batch 50000 10 | read_batch 50000 100 | read_batch 50000 1000 | read_formatted numpy 5000 | read_formatted pandas 5000 | read_formatted tensorflow 5000 | read_formatted torch 5000 | read_formatted_batch numpy 5000 10 | read_formatted_batch numpy 5000 1000 | shuffled read 5000 | shuffled read 50000 | shuffled read_batch 50000 10 | shuffled read_batch 50000 100 | shuffled read_batch 50000 1000 | shuffled read_formatted numpy 5000 | shuffled read_formatted_batch numpy 5000 10 | shuffled read_formatted_batch numpy 5000 1000 |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.298884 / 0.215209 (0.083675) | 2.961050 / 2.077655 (0.883395) | 1.620694 / 1.504120 (0.116574) | 1.499331 / 1.541195 (-0.041863) | 1.513118 / 1.468490 (0.044628) | 0.734738 / 4.584777 (-3.850039) | 0.972978 / 3.745712 (-2.772734) | 2.928172 / 5.269862 (-2.341690) | 1.903667 / 4.565676 (-2.662010) | 0.079207 / 0.424275 (-0.345068) | 0.005803 / 0.007607 (-0.001804) | 0.350144 / 0.226044 (0.124099) | 3.519456 / 2.268929 (1.250528) | 1.983809 / 55.444624 (-53.460815) | 1.690527 / 6.876477 (-5.185950) | 1.739301 / 2.142072 (-0.402772) | 0.802045 / 4.805227 (-4.003182) | 0.133041 / 6.500664 (-6.367623) | 0.042112 / 0.075469 (-0.033357) |

### Benchmark: benchmark_map_filter.json

| metric | filter | map fast-tokenizer batched | map identity | map identity batched | map no-op batched | map no-op batched numpy | map no-op batched pandas | map no-op batched pytorch | map no-op batched tensorflow |
|--------|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 1.030056 / 1.841788 (-0.811731) | 12.077692 / 8.074308 (4.003384) | 9.988253 / 10.191392 (-0.203139) | 0.142745 / 0.680424 (-0.537679) | 0.015842 / 0.534201 (-0.518359) | 0.299055 / 0.579283 (-0.280228) | 0.123788 / 0.434364 (-0.310576) | 0.352782 / 0.540337 (-0.187555) | 0.451140 / 1.386936 (-0.935796) |

</details>
</details>

![](https://cml.dev/watermark.png#0cf0be8906063d09456285be9c9f7ce5789726ae "CML watermark")