Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions python/python/lance/blob.py
Original file line number Diff line number Diff line change
Expand Up @@ -534,14 +534,16 @@ def readable(self) -> bool:

def seek(self, offset: int, whence: int = io.SEEK_SET) -> int:
if whence == io.SEEK_SET:
self.inner.seek(offset)
position = offset
elif whence == io.SEEK_CUR:
self.inner.seek(self.inner.tell() + offset)
position = self.inner.tell() + offset
elif whence == io.SEEK_END:
self.inner.seek(self.inner.size() + offset)
position = self.inner.size() + offset
else:
raise ValueError(f"Invalid whence: {whence}")

if position < 0:
raise ValueError(f"negative seek value {position}")
self.inner.seek(position)
return self.inner.tell()

def seekable(self) -> bool:
Expand Down
34 changes: 34 additions & 0 deletions python/python/tests/test_blob.py
Original file line number Diff line number Diff line change
Expand Up @@ -843,6 +843,40 @@ def test_blob_file_seek(tmp_path, dataset_with_blobs):
with blobs[1] as f:
assert f.seek(1) == 1
assert f.read(1) == b"a"
assert f.seek(-1, io.SEEK_CUR) == 1
assert f.seek(-1, io.SEEK_END) == 2


@pytest.mark.parametrize(
"whence",
[
pytest.param(io.SEEK_SET, id="set"),
pytest.param(io.SEEK_CUR, id="cur"),
pytest.param(io.SEEK_END, id="end"),
],
)
def test_blob_file_negative_seek_raises_value_error(dataset_with_blobs, whence):
blob = dataset_with_blobs.take_blobs("blobs", indices=[1])[0]
offset = -(blob.size() + 1) if whence == io.SEEK_END else -1
with pytest.raises(ValueError, match="negative seek value -1"):
blob.seek(offset, whence)
assert blob.tell() == 0


def test_blob_file_negative_seek_does_not_move_cursor(dataset_with_blobs):
blob = dataset_with_blobs.take_blobs("blobs", indices=[1])[0]
blob.seek(2)
with pytest.raises(ValueError, match="negative seek value -1"):
blob.seek(-1)
assert blob.tell() == 2


def test_blob_file_read_past_eof_leaves_cursor(dataset_with_blobs):
with dataset_with_blobs.take_blobs("blobs", indices=[1])[0] as blob:
past_eof = blob.size() + 1
assert blob.seek(past_eof) == past_eof
assert blob.read() == b""
assert blob.tell() == past_eof


@pytest.mark.parametrize(
Expand Down
17 changes: 14 additions & 3 deletions rust/lance/src/dataset/blob.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1795,8 +1795,7 @@ impl BlobFile {
/// Read the entire blob file from the current cursor position
/// to the end of the file
///
/// After this call the cursor will be pointing to the end of
/// the file.
/// Advances the cursor by the number of bytes returned.
pub async fn read(&self) -> Result<bytes::Bytes> {
let mut state = self.state.lock().await;
match state.deref_mut() {
Expand All @@ -1807,7 +1806,6 @@ impl BlobFile {
cursor, prefetch, ..
} => {
if *cursor >= self.size {
*cursor = self.size;
*prefetch = None;
return Ok(Bytes::new());
}
Expand Down Expand Up @@ -9783,6 +9781,19 @@ mod tests {
assert_eq!(blob.range_submission_count(), after_fill + 1);
}

#[tokio::test]
async fn read_past_eof_leaves_the_cursor() {
let payload = b"abcdef";
let (_dir, dataset) = write_blob_v2_dataset(&[payload.as_slice()]).await;
let blobs = dataset.take_blobs_by_indices(&[0], "blob").await.unwrap();
let blob = blobs[0].as_ref().unwrap();
let past_eof = blob.size() + 1;

blob.seek(past_eof).await.unwrap();
assert!(blob.read().await.unwrap().is_empty());
assert_eq!(blob.tell().await.unwrap(), past_eof);
}

#[tokio::test]
async fn changing_buffer_size_drops_unused_prefetch() {
let payload: Vec<u8> = (0..40).map(|i| i as u8).collect();
Expand Down
Loading