Add polars compatibility
State: closed | Author: psmyth94 (CONTRIBUTOR)

Hey there,

I've just finished adding support to convert and format to `polars.DataFrame`. This was in response to the open issue about integrating Polars [#3334](https://github.com/huggingface/datasets/issues/3334). Datasets can be switched to Polars format via `Dataset.set_format("polars")`. I've also included `to_polars` and `from_polars`. All polars functions are checked via config.POLARS_AVAILABLE. 

A few notes:
This only supports `DataFrames` and not `LazyFrames`. This probably could be integrated fairly easily via `is_lazy` args in `set_format`, and `to_polars`.

Let me know your feedbacks.

Comments: Hi ! thanks for adding polars support :)

You added from_polars in arrow_dataset.py but not to_polars, is this on purpose ?

Also no need to touch table.py imo, which is for arrow-only logic (tables are just wrappers of pyarrow.Table with the exact same methods + optimization to existing methods + separation between in-memory and memory-mapped)

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

Hi @lhoestq, thanks for pointing out the missing `to_polars` method.

I see your point about `table.py` so I removed them.

I also added tests in `test_arrow_dataset.py`, `test_dataset_dict.py`, and `test_formatting.py`. Let me know if I am missing any. 

duckdb index files were deleted yesterday in dataset_with_script@ref/convert/parquet so I changed the hash to reflect the new SHA.

<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.004993 / 0.011353 (-0.006360) | 0.003658 / 0.011008 (-0.007350) | 0.063868 / 0.038508 (0.025360) | 0.030022 / 0.023109 (0.006912) | 0.246359 / 0.275898 (-0.029539) | 0.273409 / 0.323480 (-0.050070) | 0.003091 / 0.007986 (-0.004894) | 0.003383 / 0.004328 (-0.000945) | 0.050666 / 0.004250 (0.046415) | 0.040609 / 0.037052 (0.003557) | 0.267250 / 0.258489 (0.008761) | 0.289823 / 0.293841 (-0.004018) | 0.027635 / 0.128546 (-0.100911) | 0.010786 / 0.075646 (-0.064860) | 0.208442 / 0.419271 (-0.210830) | 0.036627 / 0.043533 (-0.006906) | 0.254116 / 0.255139 (-0.001023) | 0.274368 / 0.283200 (-0.008832) | 0.018222 / 0.141683 (-0.123460) | 1.184472 / 1.452155 (-0.267683) | 1.194309 / 1.492716 (-0.298407) |

### 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.092861 / 0.018006 (0.074855) | 0.304736 / 0.000490 (0.304246) | 0.000219 / 0.000200 (0.000019) | 0.000175 / 0.000054 (0.000121) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.019378 / 0.037411 (-0.018034) | 0.062342 / 0.014526 (0.047817) | 0.074107 / 0.176557 (-0.102450) | 0.121746 / 0.737135 (-0.615390) | 0.075657 / 0.296338 (-0.220681) |

### 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.286474 / 0.215209 (0.071265) | 2.832043 / 2.077655 (0.754389) | 1.453520 / 1.504120 (-0.050600) | 1.324714 / 1.541195 (-0.216480) | 1.335439 / 1.468490 (-0.133051) | 0.571753 / 4.584777 (-4.013024) | 2.427361 / 3.745712 (-1.318352) | 2.899838 / 5.269862 (-2.370024) | 1.775754 / 4.565676 (-2.789922) | 0.064177 / 0.424275 (-0.360098) | 0.004978 / 0.007607 (-0.002629) | 0.343585 / 0.226044 (0.117541) | 3.368494 / 2.268929 (1.099565) | 1.819825 / 55.444624 (-53.624800) | 1.502633 / 6.876477 (-5.373844) | 1.549182 / 2.142072 (-0.592891) | 0.658245 / 4.805227 (-4.146983) | 0.120052 / 6.500664 (-6.380612) | 0.043051 / 0.075469 (-0.032419) |

### 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.977055 / 1.841788 (-0.864733) | 11.595567 / 8.074308 (3.521259) | 9.450951 / 10.191392 (-0.740441) | 0.141060 / 0.680424 (-0.539364) | 0.014359 / 0.534201 (-0.519842) | 0.289938 / 0.579283 (-0.289345) | 0.266035 / 0.434364 (-0.168329) | 0.326802 / 0.540337 (-0.213536) | 0.431913 / 1.386936 (-0.955023) |

</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.005391 / 0.011353 (-0.005961) | 0.003724 / 0.011008 (-0.007284) | 0.050432 / 0.038508 (0.011924) | 0.029904 / 0.023109 (0.006794) | 0.270870 / 0.275898 (-0.005028) | 0.296773 / 0.323480 (-0.026706) | 0.004265 / 0.007986 (-0.003721) | 0.002751 / 0.004328 (-0.001577) | 0.050366 / 0.004250 (0.046116) | 0.046415 / 0.037052 (0.009363) | 0.283272 / 0.258489 (0.024783) | 0.320188 / 0.293841 (0.026347) | 0.029827 / 0.128546 (-0.098719) | 0.010736 / 0.075646 (-0.064910) | 0.059541 / 0.419271 (-0.359731) | 0.057080 / 0.043533 (0.013548) | 0.270653 / 0.255139 (0.015514) | 0.291235 / 0.283200 (0.008035) | 0.018590 / 0.141683 (-0.123093) | 1.129402 / 1.452155 (-0.322752) | 1.194499 / 1.492716 (-0.298217) |

### 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.102220 / 0.018006 (0.084214) | 0.302176 / 0.000490 (0.301686) | 0.000229 / 0.000200 (0.000029) | 0.000056 / 0.000054 (0.000002) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.022809 / 0.037411 (-0.014602) | 0.076054 / 0.014526 (0.061528) | 0.087466 / 0.176557 (-0.089091) | 0.128495 / 0.737135 (-0.608640) | 0.089933 / 0.296338 (-0.206406) |

### 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.296546 / 0.215209 (0.081337) | 2.898693 / 2.077655 (0.821039) | 1.605002 / 1.504120 (0.100883) | 1.468370 / 1.541195 (-0.072825) | 1.503541 / 1.468490 (0.035051) | 0.577233 / 4.584777 (-4.007544) | 2.460154 / 3.745712 (-1.285558) | 2.755651 / 5.269862 (-2.514211) | 1.777711 / 4.565676 (-2.787966) | 0.063137 / 0.424275 (-0.361138) | 0.005056 / 0.007607 (-0.002551) | 0.350189 / 0.226044 (0.124145) | 3.485473 / 2.268929 (1.216545) | 1.952553 / 55.444624 (-53.492072) | 1.669108 / 6.876477 (-5.207369) | 1.788504 / 2.142072 (-0.353569) | 0.672869 / 4.805227 (-4.132359) | 0.117717 / 6.500664 (-6.382948) | 0.040499 / 0.075469 (-0.034970) |

### 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.048187 / 1.841788 (-0.793601) | 12.663229 / 8.074308 (4.588921) | 10.316487 / 10.191392 (0.125095) | 0.142537 / 0.680424 (-0.537887) | 0.016024 / 0.534201 (-0.518177) | 0.292735 / 0.579283 (-0.286548) | 0.273294 / 0.434364 (-0.161069) | 0.327636 / 0.540337 (-0.212701) | 0.443062 / 1.386936 (-0.943874) |

</details>
</details>

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


I'm so excited I tweeted about it: https://x.com/qlhoest/status/1766135995513082086?s=20 I hope it's fine !

Thanks @lhoestq for the support and totally fine with the share! Happy to see people excited for this 😃