diff --git a/python/python/lance/file.py b/python/python/lance/file.py index 011fbe4a01d..7e2b2933aab 100644 --- a/python/python/lance/file.py +++ b/python/python/lance/file.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The Lance Authors +from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Dict, List, Optional, Union @@ -207,6 +208,26 @@ def num_rows(self) -> int: return self._reader.num_rows() +@dataclass +class ListResult: + """ + Result of a non-recursive, delimited list (see + :meth:`LanceFileSession.list_with_delimiter`). + + Attributes + ---------- + common_prefixes : List[str] + The immediate child "directories" of the listed path, relative to the + session's base path. + objects : List[str] + The immediate child files of the listed path, relative to the session's + base path. + """ + + common_prefixes: List[str] + objects: List[str] + + class LanceFileSession: """ A file session for reading and writing Lance files. @@ -344,6 +365,66 @@ def list(self, path: Optional[str] = None) -> List[str]: """ return self._session.list(path) + def list_with_delimiter(self, path: Optional[str] = None) -> ListResult: + """ + Non-recursively list a single directory level (relative to this + session's base path). + + Unlike :meth:`list`, which recurses into the entire subtree, this + returns only the immediate children of ``path``: the child + "directories" as ``common_prefixes`` and the direct child files as + ``objects``. + + Parameters + ---------- + path : str, optional + Path relative to `base_path` to list. If None, lists the base path. + + Returns + ------- + ListResult + The immediate child prefixes and objects of `path`. + """ + common_prefixes, objects = self._session.list_with_delimiter(path) + return ListResult(common_prefixes=common_prefixes, objects=objects) + + def read_range(self, path: str, offset: int, length: int) -> bytes: + """ + Read a byte range from a file (relative to this session's base path). + + Issues a single ranged read. Reading a missing object raises + ``OSError``, consistent with ``download_file``. + + Parameters + ---------- + path : str + Path relative to `base_path` to read from. + offset : int + Byte offset at which to start reading. + length : int + Number of bytes to read. + + Returns + ------- + bytes + The requested byte range. + """ + return self._session.read_range(path, offset, length) + + def delete_file(self, path: str) -> None: + """ + Delete a file (relative to this session's base path). + + Deleting a path that does not exist raises ``OSError``, consistent with + ``download_file``. + + Parameters + ---------- + path : str + Path relative to `base_path` to delete. + """ + self._session.delete_file(path) + def upload_file(self, local_path: Union[str, Path], remote_path: str) -> None: """ Upload a file from local filesystem to the object store. diff --git a/python/python/lance/lance/__init__.pyi b/python/python/lance/lance/__init__.pyi index 26ad75a27b7..49e410459cc 100644 --- a/python/python/lance/lance/__init__.pyi +++ b/python/python/lance/lance/__init__.pyi @@ -135,6 +135,11 @@ class LanceFileSession: ) -> LanceFileWriter: ... def contains(self, path: str) -> bool: ... def list(self, path: Optional[str] = None) -> List[str]: ... + def list_with_delimiter( + self, path: Optional[str] = None + ) -> tuple[List[str], List[str]]: ... + def read_range(self, path: str, offset: int, length: int) -> bytes: ... + def delete_file(self, path: str) -> None: ... def upload_file(self, local_path: str, remote_path: str) -> None: ... def download_file(self, remote_path: str, local_path: str) -> None: ... diff --git a/python/python/tests/test_file.py b/python/python/tests/test_file.py index c71654769b8..d0f0fb0b8b4 100644 --- a/python/python/tests/test_file.py +++ b/python/python/tests/test_file.py @@ -736,6 +736,44 @@ def test_session_list_with_trailing_slash(tmp_path): assert files_no_slash == ["dir/file.lance"] +def test_session_list_with_delimiter(tmp_path): + """Test that LanceFileSession.list_with_delimiter() is non-recursive.""" + session = LanceFileSession(str(tmp_path)) + schema = pa.schema([pa.field("x", pa.int64())]) + + # Two top-level files and two nested subtrees. + with session.open_writer("file1.lance", schema=schema) as writer: + writer.write_batch(pa.table({"x": [1]})) + with session.open_writer("file2.lance", schema=schema) as writer: + writer.write_batch(pa.table({"x": [2]})) + with session.open_writer("subdir/file3.lance", schema=schema) as writer: + writer.write_batch(pa.table({"x": [3]})) + with session.open_writer("subdir/nested/file4.lance", schema=schema) as writer: + writer.write_batch(pa.table({"x": [4]})) + with session.open_writer("other/file5.lance", schema=schema) as writer: + writer.write_batch(pa.table({"x": [5]})) + + # Listing the base path returns only the immediate children: the two + # top-level files and the two child directories (not their contents). + result = session.list_with_delimiter() + assert sorted(result.common_prefixes) == ["other", "subdir"] + assert sorted(result.objects) == ["file1.lance", "file2.lance"] + + # Listing a subdirectory descends exactly one level: the direct file and + # the nested directory, but not the file inside the nested directory. + subdir = session.list_with_delimiter("subdir") + assert subdir.common_prefixes == ["subdir/nested"] + assert subdir.objects == ["subdir/file3.lance"] + + # Trailing slash behaves the same as no trailing slash. + assert session.list_with_delimiter("subdir/") == subdir + + # A non-existent prefix yields empty results rather than erroring. + empty = session.list_with_delimiter("nonexistent") + assert empty.common_prefixes == [] + assert empty.objects == [] + + def test_session_contains(tmp_path): """Test that LanceFileSession.contains() works correctly""" session = LanceFileSession(str(tmp_path)) @@ -759,6 +797,56 @@ def test_session_contains(tmp_path): assert not session.contains("subdir/nonexistent.lance") +def test_session_read_range(tmp_path): + """Test that LanceFileSession.read_range() returns the requested bytes.""" + session = LanceFileSession(str(tmp_path)) + + payload = bytes(range(256)) + local = tmp_path / "src.bin" + local.write_bytes(payload) + session.upload_file(str(local), "data/file.bin") + + # A range in the middle of the file. + assert session.read_range("data/file.bin", 10, 5) == payload[10:15] + # From the start. + assert session.read_range("data/file.bin", 0, 4) == payload[0:4] + # Up to the end. + assert session.read_range("data/file.bin", 250, 6) == payload[250:256] + # A zero-length read yields empty bytes. + assert session.read_range("data/file.bin", 100, 0) == b"" + + # Reading a missing object raises OSError (consistent with download_file). + with pytest.raises(OSError): + session.read_range("data/missing.bin", 0, 4) + + +def test_session_delete_file(tmp_path): + """Test that LanceFileSession.delete_file() removes files and is idempotent.""" + session = LanceFileSession(str(tmp_path)) + schema = pa.schema([pa.field("x", pa.int64())]) + + with session.open_writer("test.lance", schema=schema) as writer: + writer.write_batch(pa.table({"x": [1]})) + with session.open_writer("subdir/nested.lance", schema=schema) as writer: + writer.write_batch(pa.table({"x": [2]})) + + # Deleting an existing file removes it. + assert session.contains("test.lance") + session.delete_file("test.lance") + assert not session.contains("test.lance") + + # Nested paths work too. + assert session.contains("subdir/nested.lance") + session.delete_file("subdir/nested.lance") + assert not session.contains("subdir/nested.lance") + + # Deleting a missing path raises OSError (consistent with download_file). + with pytest.raises(OSError): + session.delete_file("test.lance") + with pytest.raises(OSError): + session.delete_file("never_existed.lance") + + def test_struct_null_regression(): import lance diff --git a/python/src/file.rs b/python/src/file.rs index b0bc20f9d0a..e1906b9f4f2 100644 --- a/python/src/file.rs +++ b/python/src/file.rs @@ -574,6 +574,57 @@ impl LanceFileSession { })? } + /// Non-recursive, delimited list of a single directory level. + /// + /// Returns a tuple `(common_prefixes, objects)` of paths relative to the + /// session's `base_path`, where `common_prefixes` are the immediate child + /// "directories" and `objects` are the immediate child files. Unlike + /// `list`, this does not recurse into the subtree. + #[pyo3(signature=(path=None))] + pub fn list_with_delimiter( + &self, + path: Option, + ) -> PyResult<(Vec, Vec)> { + rt().block_on(None, async { + let list_path = if let Some(prefix) = path { + self.base_path.child_path(&Path::from(prefix)) + } else { + self.base_path.clone() + }; + + let result = self + .object_store + .list_with_delimiter(Some(&list_path)) + .await + .map_err(|e| PyErr::new::(format!("{}", e)))?; + + // Strip the base_path prefix to make each path relative to the session. + let relativize = |location: &Path| -> PyResult { + let relative_parts = location.prefix_match(&self.base_path).ok_or_else(|| { + PyErr::new::(format!( + "Path '{}' does not start with base path '{}'", + location.as_ref(), + self.base_path.as_ref() + )) + })?; + Ok(Path::from_iter(relative_parts).as_ref().to_string()) + }; + + let common_prefixes = result + .common_prefixes + .iter() + .map(relativize) + .collect::>>()?; + let objects = result + .objects + .iter() + .map(|meta| relativize(&meta.location)) + .collect::>>()?; + + Ok((common_prefixes, objects)) + })? + } + /// Upload a file from local filesystem to the object store /// /// Parameters @@ -606,6 +657,61 @@ impl LanceFileSession { })? } + /// Delete a file from the object store. + /// + /// The path is interpreted relative to the session's `base_path`, matching + /// `contains`/`upload_file`/`download_file`. Deleting a missing object + /// raises `OSError`, consistent with `download_file`. + /// + /// Parameters + /// ---------- + /// path : str + /// Path relative to `base_path` to delete. + pub fn delete_file(&self, path: String) -> PyResult<()> { + rt().block_on(None, async { + let full_path = self.base_path.child_path(&Path::from(path)); + self.object_store + .inner + .delete(&full_path) + .await + .map_err(|e| PyIOError::new_err(format!("Failed to delete remote file: {}", e)))?; + Ok(()) + })? + } + + /// Read a byte range from a file in the object store. + /// + /// The path is interpreted relative to the session's `base_path`, matching + /// the other session methods. This issues a single ranged GET. Reading a + /// missing object raises `OSError`, consistent with `download_file`. + /// + /// Parameters + /// ---------- + /// path : str + /// Path relative to `base_path` to read from. + /// offset : int + /// Byte offset at which to start reading. + /// length : int + /// Number of bytes to read. + /// + /// Returns + /// ------- + /// bytes + /// The requested byte range. + pub fn read_range(&self, path: String, offset: usize, length: usize) -> PyResult> { + rt().block_on(None, async { + let full_path = self.base_path.child_path(&Path::from(path)); + let bytes = self + .object_store + .read_one_range(&full_path, offset..offset + length) + .await + .map_err(|e| { + PyIOError::new_err(format!("Failed to read range from remote file: {}", e)) + })?; + Ok(bytes.to_vec()) + })? + } + /// Download a file from object store to local filesystem /// /// Parameters diff --git a/rust/lance-io/src/object_store.rs b/rust/lance-io/src/object_store.rs index 1761dc4b059..31d2e8f997e 100644 --- a/rust/lance-io/src/object_store.rs +++ b/rust/lance-io/src/object_store.rs @@ -25,7 +25,7 @@ use object_store::ObjectStoreExt as OSObjectStoreExt; use object_store::aws::AwsCredentialProvider; #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] use object_store::{ClientOptions, HeaderMap, HeaderValue}; -use object_store::{ObjectMeta, ObjectStore as OSObjectStore, path::Path}; +use object_store::{ListResult, ObjectMeta, ObjectStore as OSObjectStore, path::Path}; use providers::local::FileStoreProvider; use providers::memory::MemoryStoreProvider; use tokio::io::AsyncWriteExt; @@ -849,6 +849,16 @@ impl ObjectStore { .collect()) } + /// Non-recursive, path-segment delimited list of a single directory level. + /// + /// Unlike [`Self::list`], which recurses into the entire subtree, this returns + /// only the immediate children of `prefix`: the child "directories" as + /// [`ListResult::common_prefixes`] and the direct child files as + /// [`ListResult::objects`]. + pub async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result { + Ok(self.inner.list_with_delimiter(prefix).await?) + } + pub fn list( &self, path: Option,