feat: add tf.data APIs for reading batches #1488
Conversation
|
ACTION NEEDED Lance follows the Conventional Commits specification for release automation. The PR title and description are used as the merge commit message. Please update your PR title and description to match the specification. For details on the error please inspect the "PR Title Check" action. |
wjones127
force-pushed
the
wjones127/tf-take
branch
from
83d614c
to
3cacb62
Compare
wjones127
changed the title
wjones127
force-pushed
the
wjones127/tf-take
branch
from
3cacb62
to
34dea38
Compare
| dataset = lance.dataset(dataset) | ||
| num_rows = dataset.count_rows() | ||
| num_batches = (num_rows + batch_size - 1) // batch_size | ||
| indices = tf.data.Dataset.range(num_batches, dtype=tf.int64) |
There was a problem hiding this comment.
say a 1B dataset, this one can be 1-million batches, so a few MB batch ids?
There was a problem hiding this comment.
Yeah the user sets the batch size, so they control the memory use here.
python/python/lance/tf/data.py
Outdated
| return (start, end) | ||
|
|
||
|
|
||
| def lance_batches( |
There was a problem hiding this comment.
does tf.data has a from_batches, can we name this to from_batch to be more consist with the tf style
Users can make the batches as small as they want, and can always call the |
Closes #1499