diff --git a/python/python/lance/blob.py b/python/python/lance/blob.py index 6d1af6797dc..f4a92633807 100644 --- a/python/python/lance/blob.py +++ b/python/python/lance/blob.py @@ -383,6 +383,28 @@ def read_range(self, offset: int, length: int) -> bytes: """Read a blob-local byte range without changing the current cursor.""" return self.inner.read_range(offset, length) + def read_ranges(self, ranges: list[tuple[int, int]]) -> list[bytes]: + """ + Read multiple blob-local byte ranges without changing the current cursor. + + Each range is an ``(offset, length)`` pair, matching + :py:meth:`read_range`. The underlying physical reads may be reordered, + coalesced, or split for efficiency. For every range, offset plus length + must fit in an unsigned 64-bit integer and must not extend beyond the + blob size. + + Parameters + ---------- + ranges : List[Tuple[int, int]] + The ``(offset, length)`` byte ranges to read. + + Returns + ------- + data : List[bytes] + One payload per requested range, in input order. + """ + return self.inner.read_ranges(ranges) + def readinto(self, b: bytearray) -> int: return self.inner.read_into(b) diff --git a/python/python/lance/lance/__init__.pyi b/python/python/lance/lance/__init__.pyi index 27ccf0c6f9a..3b2e1181bca 100644 --- a/python/python/lance/lance/__init__.pyi +++ b/python/python/lance/lance/__init__.pyi @@ -311,6 +311,8 @@ class LanceBlobFile: def tell(self) -> int: ... def size(self) -> int: ... def readall(self) -> bytes: ... + def read_range(self, offset: int, length: int) -> bytes: ... + def read_ranges(self, ranges: List[Tuple[int, int]]) -> List[bytes]: ... def read_into(self, b: bytearray) -> int: ... class _Dataset: diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index 2906fa3674b..9f5267a78aa 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -756,6 +756,61 @@ def test_blob_file_seek(tmp_path, dataset_with_blobs): assert f.read(1) == b"a" +@pytest.mark.parametrize( + "ranges", + [ + pytest.param([], id="empty_list"), + pytest.param([(0, 0), (4, 0)], id="empty_ranges"), + pytest.param([(2, 1), (0, 1), (1, 1)], id="non_monotonic"), + pytest.param([(1, 2), (0, 4), (1, 2), (0, 0), (2, 2)], id="dup_overlap"), + ], +) +def test_blob_file_read_ranges_matches_read_range(dataset_with_blobs, ranges): + row_ids = _blob_row_ids(dataset_with_blobs) + blobs = dataset_with_blobs.take_blobs("blobs", ids=row_ids) + with blobs[4] as f: + expected = [f.read_range(offset, length) for offset, length in ranges] + assert f.read_ranges(ranges) == expected + + +def test_blob_file_read_ranges_preserves_input_order(dataset_with_blobs): + row_ids = _blob_row_ids(dataset_with_blobs) + blobs = dataset_with_blobs.take_blobs("blobs", ids=row_ids) + with blobs[1] as f: + assert f.read_ranges([(2, 1), (0, 1), (1, 2)]) == [b"r", b"b", b"ar"] + + +def test_blob_file_read_ranges_does_not_change_cursor(dataset_with_blobs): + row_ids = _blob_row_ids(dataset_with_blobs) + blobs = dataset_with_blobs.take_blobs("blobs", ids=row_ids) + with blobs[1] as f: + assert f.tell() == 0 + assert f.read_ranges([(0, 3)]) == [b"bar"] + assert f.tell() == 0 + f.seek(2) + assert f.read_ranges([(2, 1), (0, 3)]) == [b"r", b"bar"] + assert f.tell() == 2 + assert f.read(1) == b"r" + + +@pytest.mark.parametrize( + ("ranges", "message"), + [ + pytest.param([(2**64 - 1, 2)], "offset \\+ length overflowed", id="overflow"), + pytest.param([(0, 1), (1, 100)], "exceeds blob size", id="out_of_bounds"), + ], +) +def test_blob_file_read_ranges_rejects_invalid_ranges( + dataset_with_blobs, ranges, message +): + row_ids = _blob_row_ids(dataset_with_blobs) + blobs = dataset_with_blobs.take_blobs("blobs", ids=row_ids) + with blobs[0] as f: + with pytest.raises(ValueError, match=message): + f.read_ranges(ranges) + assert f.tell() == 0 + + def test_null_blobs(tmp_path): table = pa.table( { diff --git a/python/src/dataset/blob.rs b/python/src/dataset/blob.rs index 1f6272075f2..7952ce76fb1 100644 --- a/python/src/dataset/blob.rs +++ b/python/src/dataset/blob.rs @@ -79,6 +79,32 @@ impl LanceBlobFile { Ok(PyBytes::new(py, &data)) } + /// Read multiple blob-local `(offset, length)` ranges without changing the current cursor. + pub fn read_ranges<'py>( + &self, + py: Python<'py>, + ranges: Vec<(u64, u64)>, + ) -> PyResult>> { + let ranges = ranges + .into_iter() + .enumerate() + .map(|(i, (offset, length))| { + let end = offset.checked_add(length).ok_or_else(|| { + PyValueError::new_err(format!( + "Blob range request {i} offset + length overflowed u64: \ + offset={offset}, length={length}" + )) + })?; + Ok(offset..end) + }) + .collect::>>()?; + let inner = self.inner.clone(); + let data = rt() + .block_on(Some(py), inner.read_ranges(&ranges))? + .infer_error()?; + Ok(data.iter().map(|bytes| PyBytes::new(py, bytes)).collect()) + } + pub fn read_into(&self, dst: Bound<'_, PyByteArray>) -> PyResult { let inner = self.inner.clone();