fix(io): prune empty local directories after their files are deleted

rust/lance-io/src/local.rs

@@ -22,7 +22,7 @@ use lance_core::{Error, Result};
 use object_store::path::Path;
 use tokio::io::AsyncSeekExt;
 use tokio::sync::OnceCell;
-use tracing::instrument;
+use tracing::{instrument, warn};
 
 use crate::object_reader::stream_local_range;
 use crate::object_store::DEFAULT_LOCAL_IO_PARALLELISM;
@@ -49,6 +49,31 @@ pub fn remove_dir_all(path: &Path) -> Result<()> {
     Ok(())
 }
 
+/// Remove `path`'s ancestor directories that are now empty, stopping at the first one
+/// that is not. Best effort.
+///
+/// `LocalFileSystem::with_automatic_cleanup` cannot do this: it derives its stop boundary
+/// from a rootless `file:///` URL, which `Url::to_file_path` rejects on Windows.
+pub fn prune_empty_parent_dirs(path: &Path) {
+    let local_path = to_local_path(path);
+    let mut parent = std::path::Path::new(&local_path).parent();
+    while let Some(dir) = parent {
+        if dir.as_os_str().is_empty() {
+            break;
+        }
+        if let Err(err) = std::fs::remove_dir(dir) {
+            if !matches!(
+                err.kind(),
+                ErrorKind::DirectoryNotEmpty | ErrorKind::NotFound
+            ) {
+                warn!("failed to prune empty directory {}: {}", dir.display(), err);
+            }
+            break;
+        }
+        parent = dir.parent();
+    }
+}
+
 /// Copy a file from one location to another, supporting cross-filesystem copies.
 ///
 /// Unlike hard links, this function works across filesystem boundaries.

rust/lance-io/src/object_store.rs

@@ -590,6 +590,11 @@ impl ObjectStore {
         self.scheme == "file" || self.scheme == "file+uring"
     }
 
+    /// Wider than [`Self::is_local`]: `file-object-store` also stores real directories.
+    fn backed_by_local_fs(&self) -> bool {
+        self.is_local() || self.scheme == "file-object-store"
+    }
+
     pub fn is_cloud(&self) -> bool {
         if self.is_local() || self.scheme == "memory" || self.scheme == "shared-memory" {
             return false;
@@ -822,6 +827,9 @@ impl ObjectStore {
 
     pub async fn delete(&self, path: &Path) -> Result<()> {
         self.inner.delete(path).await?;
+        if self.backed_by_local_fs() {
+            super::local::prune_empty_parent_dirs(path);
+        }
         Ok(())
     }
 
@@ -947,11 +955,15 @@ impl ObjectStore {
         locations: BoxStream<'a, Result<Path>>,
     ) -> BoxStream<'a, Result<Path>> {
         let store = Arc::clone(&self.inner);
+        let prune_empty_dirs = self.backed_by_local_fs();
         locations
             .and_then(move |location| {
                 let store = Arc::clone(&store);
                 async move {
                     store.delete(&location).await?;
+                    if prune_empty_dirs {
+                        super::local::prune_empty_parent_dirs(&location);
+                    }
                     Ok(location)
                 }
             })
@@ -1429,15 +1441,7 @@ mod tests {
             "delete",
         )
         .unwrap();
-        let file_url = Url::from_directory_path(&path).unwrap();
-        let url = if scheme.is_empty() {
-            file_url
-        } else {
-            let mut url = Url::parse(&format!("{scheme}:///")).unwrap();
-            // Use the file:// URL's normalized path so this works on Windows too.
-            url.set_path(file_url.path());
-            url
-        };
+        let url = dir_url(&path, scheme);
         let (store, base) = ObjectStore::from_uri(url.as_ref()).await.unwrap();
         store
             .remove_dir_all(base.clone().join("foo"))
@@ -1447,6 +1451,86 @@ mod tests {
         assert!(!path.join("foo").exists());
     }
 
+    fn dir_url(path: impl AsRef<std::path::Path>, scheme: &str) -> Url {
+        let file_url = Url::from_directory_path(path).unwrap();
+        if scheme.is_empty() {
+            return file_url;
+        }
+        let mut url = Url::parse(&format!("{scheme}:///")).unwrap();
+        // Use the file:// URL's normalized path so this works on Windows too.
+        url.set_path(file_url.path());
+        url
+    }
+
+    #[tokio::test]
+    async fn test_delete_prunes_empty_dirs_local_store() {
+        test_delete_prunes_empty_dirs("").await;
+    }
+
+    #[tokio::test]
+    async fn test_delete_prunes_empty_dirs_file_object_store() {
+        test_delete_prunes_empty_dirs("file-object-store").await;
+    }
+
+    async fn test_delete_prunes_empty_dirs(scheme: &str) {
+        let (path, store, file) = prune_fixture(scheme).await;
+        store.delete(&file).await.unwrap();
+        assert_pruned(&path);
+    }
+
+    #[tokio::test]
+    async fn test_remove_stream_prunes_empty_dirs_local_store() {
+        test_remove_stream_prunes_empty_dirs("").await;
+    }
+
+    #[tokio::test]
+    async fn test_remove_stream_prunes_empty_dirs_file_object_store() {
+        test_remove_stream_prunes_empty_dirs("file-object-store").await;
+    }
+
+    async fn test_remove_stream_prunes_empty_dirs(scheme: &str) {
+        let (path, store, file) = prune_fixture(scheme).await;
+        store
+            .remove_stream(futures::stream::once(async { Ok(file) }).boxed())
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        assert_pruned(&path);
+    }
+
+    /// A `foo/bar/nested/test_file` to delete, next to a `foo/zoo/keep_file` that must
+    /// stop the upward walk at `foo/`.
+    async fn prune_fixture(scheme: &str) -> (TempStdDir, Arc<ObjectStore>, Path) {
+        let path = TempStdDir::default();
+        let nested = path.join("foo").join("bar").join("nested");
+        create_dir_all(&nested).unwrap();
+        create_dir_all(path.join("foo").join("zoo")).unwrap();
+        write_to_file(nested.join("test_file").to_str().unwrap(), "delete").unwrap();
+        write_to_file(
+            path.join("foo")
+                .join("zoo")
+                .join("keep_file")
+                .to_str()
+                .unwrap(),
+            "keep",
+        )
+        .unwrap();
+
+        let url = dir_url(&path, scheme);
+        let (store, base) = ObjectStore::from_uri(url.as_ref()).await.unwrap();
+        let file = base
+            .join("foo")
+            .join("bar")
+            .join("nested")
+            .join("test_file");
+        (path, store, file)
+    }
+
+    fn assert_pruned(path: &TempStdDir) {
+        assert!(!path.join("foo").join("bar").exists());
+        assert!(path.join("foo").join("zoo").join("keep_file").exists());
+    }
+
     #[derive(Debug)]
     struct TestWrapper {
         called: AtomicBool,

rust/lance/src/dataset/cleanup.rs

@@ -1891,6 +1891,17 @@ mod tests {
             Ok(file_count)
         }
 
+        /// Absolute on-disk path of an index's `_indices/<uuid>/` directory.
+        /// Used to assert the directory itself (not just its files) is removed,
+        /// which the object-store `exists()` check cannot observe (an empty
+        /// directory has no object key).
+        fn index_dir_on_disk(&self, uuid: Uuid) -> std::path::PathBuf {
+            std::path::Path::new(self._tmpdir.as_str())
+                .join("my_db")
+                .join("_indices")
+                .join(uuid.to_string())
+        }
+
         async fn count_blob_files(&self) -> Result<usize> {
             let registry = Arc::new(ObjectStoreRegistry::default());
             let (os, path) =
@@ -2761,6 +2772,12 @@ mod tests {
                 .await
                 .unwrap()
         );
+        assert!(
+            !fixture.index_dir_on_disk(seg_a).exists(),
+            "empty _indices/<uuid> directory left behind after cleanup"
+        );
+        assert!(fixture.index_dir_on_disk(seg_b).exists());
+        assert!(fixture.index_dir_on_disk(seg_c).exists());
     }
 
     #[tokio::test]
@@ -2818,6 +2835,8 @@ mod tests {
                 .await
                 .unwrap()
         );
+        assert!(!fixture.index_dir_on_disk(staging_uuid).exists());
+        assert!(!fixture.index_dir_on_disk(built_segment_uuid).exists());
     }
 
     #[tokio::test]