feat: torch cache-able dataset, with sampling support #1591
Conversation
python/python/lance/sampler.py
Outdated
|
|
||
| if n >= len(dataset): | ||
| # Dont have enough data in the dataset. Just do a full scan | ||
| dataset.to_batches(columns=columns, batch_size=batch_size) |
There was a problem hiding this comment.
This should have a return/yeild here?
There was a problem hiding this comment.
And maybe also add this condition in test_sampler, where. n>len(ds)
There was a problem hiding this comment.
Good catch. Wrote too much rust recently
eddyxu
changed the title
Co-authored-by: Will Jones <[email protected]>
Co-authored-by: Will Jones <[email protected]>
This can be done with shuffle in [torch.dataset.DataLoader(https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader) or tf.data.pipeline Sampling was mostly used by training IVF kmeans in GPU. If you are concern about the api capablity, i can mark it as private API. |
python/python/lance/cache.py
Outdated
| def __del__(self): | ||
| if self.cache_dir is not None: | ||
| self.cache_dir.cleanup() | ||
| self.cache_dir = None |
There was a problem hiding this comment.
I don't this this is garueeteed to be called since this is only called during GC.
There was a problem hiding this comment.
maybe we can use import atexit instead?
There was a problem hiding this comment.
atexit is strictly later than GC?
There was a problem hiding this comment.
GC isn't garueeteed to happen on an object. But atexit is
There was a problem hiding this comment.
Ok, add ti both to atext and __exit__
| self.cache_dir.cleanup() | ||
| self.cache_dir = None | ||
|
|
||
| def __iter__(self): |
There was a problem hiding this comment.
nit: do we wannt check and make sure the second stream doesn't iter faster than the first stream?
| writer = pa.ipc.new_stream(str(self.cache_file), batch.schema) | ||
| writer.write(batch) |
There was a problem hiding this comment.
does this append or overwrite?
There was a problem hiding this comment.
Pytorch Dataset