IterableDataset Arrow formatting #5821
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Show benchmarksPyArrow==8.0.0 Show updated benchmarks!Benchmark: benchmark_array_xd.json
Benchmark: benchmark_getitem_100B.json
Benchmark: benchmark_indices_mapping.json
Benchmark: benchmark_iterating.json
Benchmark: benchmark_map_filter.json
Show updated benchmarks!Benchmark: benchmark_array_xd.json
Benchmark: benchmark_getitem_100B.json
Benchmark: benchmark_indices_mapping.json
Benchmark: benchmark_iterating.json
Benchmark: benchmark_map_filter.json
|
|
The documentation is not available anymore as the PR was closed or merged. |
Show benchmarksPyArrow==8.0.0 Show updated benchmarks!Benchmark: benchmark_array_xd.json
Benchmark: benchmark_getitem_100B.json
Benchmark: benchmark_indices_mapping.json
Benchmark: benchmark_iterating.json
Benchmark: benchmark_map_filter.json
Show updated benchmarks!Benchmark: benchmark_array_xd.json
Benchmark: benchmark_getitem_100B.json
Benchmark: benchmark_indices_mapping.json
Benchmark: benchmark_iterating.json
Benchmark: benchmark_map_filter.json
|
Show benchmarksPyArrow==8.0.0 Show updated benchmarks!Benchmark: benchmark_array_xd.json
Benchmark: benchmark_getitem_100B.json
Benchmark: benchmark_indices_mapping.json
Benchmark: benchmark_iterating.json
Benchmark: benchmark_map_filter.json
Show updated benchmarks!Benchmark: benchmark_array_xd.json
Benchmark: benchmark_getitem_100B.json
Benchmark: benchmark_indices_mapping.json
Benchmark: benchmark_iterating.json
Benchmark: benchmark_map_filter.json
|
This was referenced May 12, 2023
|
Will need to take #5810 into account if it gets merged before this one |
Show benchmarksPyArrow==8.0.0 Show updated benchmarks!Benchmark: benchmark_array_xd.json
Benchmark: benchmark_getitem_100B.json
Benchmark: benchmark_indices_mapping.json
Benchmark: benchmark_iterating.json
Benchmark: benchmark_map_filter.json
Show updated benchmarks!Benchmark: benchmark_array_xd.json
Benchmark: benchmark_getitem_100B.json
Benchmark: benchmark_indices_mapping.json
Benchmark: benchmark_iterating.json
Benchmark: benchmark_map_filter.json
|
Show benchmarksPyArrow==8.0.0 Show updated benchmarks!Benchmark: benchmark_array_xd.json
Benchmark: benchmark_getitem_100B.json
Benchmark: benchmark_indices_mapping.json
Benchmark: benchmark_iterating.json
Benchmark: benchmark_map_filter.json
Show updated benchmarks!Benchmark: benchmark_array_xd.json
Benchmark: benchmark_getitem_100B.json
Benchmark: benchmark_indices_mapping.json
Benchmark: benchmark_iterating.json
Benchmark: benchmark_map_filter.json
|
| ex_iterable = ExamplesIterable(Dataset._generate_examples_from_shards, kwargs={"shards": shards}) | ||
| ex_iterable = ArrowExamplesIterable( | ||
| Dataset._generate_tables_from_shards, | ||
| kwargs={"shards": shards, "batch_size": config.DEFAULT_MAX_BATCH_SIZE}, |
There was a problem hiding this comment.
I am wondering if we should support users to pass a custom batch_size.
There was a problem hiding this comment.
yea I'm not sure - we can wait for some feedback on this and improve later imo
src/datasets/iterable_dataset.py
Outdated
| def _convert_to_arrow( | ||
| iterable: Iterable[Tuple[Key, dict]], | ||
| batch_size: int, | ||
| drop_last_batch=False, |
There was a problem hiding this comment.
Missing type hint here.
src/datasets/iterable_dataset.py
Outdated
| batch_size: int, | ||
| drop_last_batch=False, | ||
| ) -> Iterator[Tuple[Key, pa.Table]]: | ||
| """Iterate over sub-tables of size `batch_size`. |
There was a problem hiding this comment.
The definition in the docstring is the same as the one for _batch_arrow_tables below. Maybe we should add to the description they expect different iterables...
Maybe also the naming of both functions could be aligned: now they have very different names whereas their functionality is analogous...
There was a problem hiding this comment.
I like the idea of having different names because one is expensive (contains "convert" in the name) while the other one is not (simply batches data)
|
I fixed the docstring and type hint |
Show benchmarksPyArrow==8.0.0 Show updated benchmarks!Benchmark: benchmark_array_xd.json
Benchmark: benchmark_getitem_100B.json
Benchmark: benchmark_indices_mapping.json
Benchmark: benchmark_iterating.json
Benchmark: benchmark_map_filter.json
Show updated benchmarks!Benchmark: benchmark_array_xd.json
Benchmark: benchmark_getitem_100B.json
Benchmark: benchmark_indices_mapping.json
Benchmark: benchmark_iterating.json
Benchmark: benchmark_map_filter.json
|
|
let me know if it sounds good for you now @albertvillanova :) |
Comment on lines
+134
to
+165
| keys_buffer = [] | ||
| chunks_buffer = [] | ||
| chunks_buffer_size = 0 | ||
| for key, pa_table in iterable: | ||
| for chunk in pa_table.to_reader(max_chunksize=batch_size): | ||
| if len(chunk) == 0: | ||
| continue | ||
| elif chunks_buffer_size + len(chunk) < batch_size: | ||
| keys_buffer.append(key) | ||
| chunks_buffer.append(chunk) | ||
| chunks_buffer_size += len(chunk) | ||
| continue | ||
| elif chunks_buffer_size + len(chunk) == batch_size: | ||
| keys_buffer.append(key) | ||
| chunks_buffer.append(chunk) | ||
| new_key = "_".join(str(_key) for _key in keys_buffer) | ||
| yield new_key, pa.Table.from_batches(chunks_buffer) | ||
| keys_buffer = [] | ||
| chunks_buffer = [] | ||
| chunks_buffer_size = 0 | ||
| else: | ||
| cropped_chunk_length = batch_size - chunks_buffer_size | ||
| keys_buffer.append(f"{key}[:{cropped_chunk_length}]") | ||
| chunks_buffer.append(chunk.slice(0, cropped_chunk_length)) | ||
| new_key = "_".join(str(_key) for _key in keys_buffer) | ||
| yield new_key, pa.Table.from_batches(chunks_buffer) | ||
| keys_buffer = [f"{key}[{cropped_chunk_length}:]"] | ||
| chunks_buffer = [chunk.slice(cropped_chunk_length, len(chunk) - cropped_chunk_length)] | ||
| chunks_buffer_size = len(chunk) - cropped_chunk_length | ||
| if not drop_last_batch and chunks_buffer: | ||
| new_key = "_".join(str(_key) for _key in keys_buffer) | ||
| yield new_key, pa.Table.from_batches(chunks_buffer) |
There was a problem hiding this comment.
This function just rebatches, from the original batch size returned by self.generate_tables_fn(**self.kwargs), to the target batch_size. And to do that, it generates batches from Tables (using Table.to_reader) and then Tables from batches (using Table.from_batches).
I am just wondering it the implementation could be simpler... Not sure though.
What do you think about a naive approach like?
iterator = iter(pa_table.slice(i, 1) for _, pa_table in iterable for i in range(pa_table.num_rows))
batch = True
key = 0
while batch:
batch = list(islice(iterator, batch_size))
if batch:
if drop_last_batch and len(batch) < batch_size:
continue
yield key, pa.concat_tables(batch)
key += 1There was a problem hiding this comment.
I think iterating on batches of size 1 and grouping them again is slower. It may also return tables with too many record batches (1 batch per row) to be efficiently used
There was a problem hiding this comment.
Then, feel free to merge as it is...
Show benchmarksPyArrow==8.0.0 Show updated benchmarks!Benchmark: benchmark_array_xd.json
Benchmark: benchmark_getitem_100B.json
Benchmark: benchmark_indices_mapping.json
Benchmark: benchmark_iterating.json
Benchmark: benchmark_map_filter.json
Show updated benchmarks!Benchmark: benchmark_array_xd.json
Benchmark: benchmark_getitem_100B.json
Benchmark: benchmark_indices_mapping.json
Benchmark: benchmark_iterating.json
Benchmark: benchmark_map_filter.json
|
There was a problem hiding this comment.
After your commit 028822a, you say with_format only supports "arrow" or None.
Maybe we should fix all the tests in test_iterable_dataset.py that contain .with_format("torch")?
Otherwise, feel free to merge as it is.
they're updated in #5852 |
Show benchmarksPyArrow==8.0.0 Show updated benchmarks!Benchmark: benchmark_array_xd.json
Benchmark: benchmark_getitem_100B.json
Benchmark: benchmark_indices_mapping.json
Benchmark: benchmark_iterating.json
Benchmark: benchmark_map_filter.json
Show updated benchmarks!Benchmark: benchmark_array_xd.json
Benchmark: benchmark_getitem_100B.json
Benchmark: benchmark_indices_mapping.json
Benchmark: benchmark_iterating.json
Benchmark: benchmark_map_filter.json
|
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Adding an optional
.iter_arrowto examples iterable. This allows to use Arrow formatting in map/filter.This will also be useful for torch formatting, since we can reuse the TorchFormatter that converts Arrow data to torch tensors
Related to #5793 and #3444
Required for #5852
Example:
Speed x10 in map
Implementation details
I added an optional
iter_arrowmethod to examples iterable. If an example iterable has this method, then it can be used to iterate on the examples by batch of arrow tables.