Allow loading/saving of FAISS index using fsspec
State: closed | Author: Dref360 (CONTRIBUTOR)

Fixes #5428

Allow loading/saving of FAISS index using fsspec:
1. Simply use BufferedIOWriter/Reader to Read/Write indices on fsspec stream.
2. Needed `mockfs` in the test, so I took it out of the `TestCase`. Let me know if that makes sense.


I can work on the documentation once the code changes are approved.


Comments: _The documentation is not available anymore as the PR was closed or merged._

Thanks for the quick review! I updated the code with your suggestion

Thanks for the quick review @albertvillanova! I updated the code with your suggestions

<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.008577 / 0.011353 (-0.002776) | 0.005714 / 0.011008 (-0.005294) | 0.114718 / 0.038508 (0.076210) | 0.039799 / 0.023109 (0.016690) | 0.387530 / 0.275898 (0.111632) | 0.395739 / 0.323480 (0.072259) | 0.006775 / 0.007986 (-0.001211) | 0.006280 / 0.004328 (0.001952) | 0.086470 / 0.004250 (0.082220) | 0.054424 / 0.037052 (0.017371) | 0.361989 / 0.258489 (0.103500) | 0.424678 / 0.293841 (0.130837) | 0.043081 / 0.128546 (-0.085465) | 0.013903 / 0.075646 (-0.061743) | 0.397625 / 0.419271 (-0.021647) | 0.059789 / 0.043533 (0.016256) | 0.375195 / 0.255139 (0.120056) | 0.403724 / 0.283200 (0.120524) | 0.121470 / 0.141683 (-0.020213) | 1.734496 / 1.452155 (0.282341) | 1.820479 / 1.492716 (0.327763) |

### 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.239672 / 0.018006 (0.221665) | 0.499373 / 0.000490 (0.498883) | 0.005034 / 0.000200 (0.004834) | 0.000099 / 0.000054 (0.000045) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.033000 / 0.037411 (-0.004411) | 0.130930 / 0.014526 (0.116404) | 0.151690 / 0.176557 (-0.024866) | 0.211839 / 0.737135 (-0.525296) | 0.148727 / 0.296338 (-0.147612) |

### 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.480592 / 0.215209 (0.265382) | 4.809700 / 2.077655 (2.732046) | 2.232414 / 1.504120 (0.728294) | 2.035432 / 1.541195 (0.494237) | 2.115991 / 1.468490 (0.647501) | 0.817841 / 4.584777 (-3.766936) | 4.718035 / 3.745712 (0.972323) | 4.107102 / 5.269862 (-1.162759) | 2.166838 / 4.565676 (-2.398839) | 0.102207 / 0.424275 (-0.322068) | 0.014686 / 0.007607 (0.007079) | 0.599922 / 0.226044 (0.373877) | 5.985840 / 2.268929 (3.716912) | 2.769199 / 55.444624 (-52.675425) | 2.427095 / 6.876477 (-4.449382) | 2.586666 / 2.142072 (0.444593) | 0.987650 / 4.805227 (-3.817578) | 0.199419 / 6.500664 (-6.301245) | 0.076710 / 0.075469 (0.001240) |

### 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.454509 / 1.841788 (-0.387278) | 18.267849 / 8.074308 (10.193541) | 16.701880 / 10.191392 (6.510488) | 0.204225 / 0.680424 (-0.476199) | 0.020295 / 0.534201 (-0.513906) | 0.504254 / 0.579283 (-0.075029) | 0.535071 / 0.434364 (0.100707) | 0.611825 / 0.540337 (0.071488) | 0.697289 / 1.386936 (-0.689647) |

</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.009141 / 0.011353 (-0.002211) | 0.005987 / 0.011008 (-0.005021) | 0.092003 / 0.038508 (0.053495) | 0.043239 / 0.023109 (0.020130) | 0.400425 / 0.275898 (0.124527) | 0.464849 / 0.323480 (0.141369) | 0.008256 / 0.007986 (0.000270) | 0.006251 / 0.004328 (0.001923) | 0.095263 / 0.004250 (0.091013) | 0.057899 / 0.037052 (0.020847) | 0.402899 / 0.258489 (0.144410) | 0.477411 / 0.293841 (0.183570) | 0.044122 / 0.128546 (-0.084424) | 0.014158 / 0.075646 (-0.061489) | 0.116354 / 0.419271 (-0.302917) | 0.061045 / 0.043533 (0.017512) | 0.411635 / 0.255139 (0.156497) | 0.466281 / 0.283200 (0.183082) | 0.129423 / 0.141683 (-0.012260) | 1.799790 / 1.452155 (0.347635) | 2.004578 / 1.492716 (0.511862) |

### 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.224012 / 0.018006 (0.206006) | 0.502972 / 0.000490 (0.502482) | 0.003560 / 0.000200 (0.003360) | 0.000110 / 0.000054 (0.000056) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.034794 / 0.037411 (-0.002618) | 0.139646 / 0.014526 (0.125120) | 0.144330 / 0.176557 (-0.032226) | 0.202528 / 0.737135 (-0.534607) | 0.151561 / 0.296338 (-0.144777) |

### 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.504343 / 0.215209 (0.289133) | 5.050690 / 2.077655 (2.973035) | 2.433107 / 1.504120 (0.928987) | 2.197443 / 1.541195 (0.656248) | 2.331225 / 1.468490 (0.862734) | 0.834066 / 4.584777 (-3.750711) | 4.837648 / 3.745712 (1.091936) | 4.105672 / 5.269862 (-1.164189) | 2.281557 / 4.565676 (-2.284120) | 0.102257 / 0.424275 (-0.322018) | 0.014425 / 0.007607 (0.006818) | 0.629290 / 0.226044 (0.403245) | 6.251513 / 2.268929 (3.982585) | 2.959012 / 55.444624 (-52.485613) | 2.570031 / 6.876477 (-4.306446) | 2.657525 / 2.142072 (0.515453) | 1.002861 / 4.805227 (-3.802367) | 0.199326 / 6.500664 (-6.301338) | 0.078428 / 0.075469 (0.002958) |

### 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.579587 / 1.841788 (-0.262201) | 18.567509 / 8.074308 (10.493201) | 17.162144 / 10.191392 (6.970752) | 0.193460 / 0.680424 (-0.486964) | 0.020819 / 0.534201 (-0.513382) | 0.501929 / 0.579283 (-0.077354) | 0.508039 / 0.434364 (0.073675) | 0.582656 / 0.540337 (0.042319) | 0.693624 / 1.386936 (-0.693312) |

</details>
</details>

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