OPD has now a working 'remove_from_disk()' method.

Author: yohplala

oups/store/ordered_parquet_dataset/ordered_parquet_dataset/base.py

@@ -285,14 +285,6 @@ def __init__(
             self._release_lock()
             raise
 
-    def _release_lock(self):
-        """
-        Release lock with reference counting.
-        """
-        self._lock._ref_count -= 1
-        if self._lock._ref_count <= 0:
-            self._lock.unlock(unconditionally=True)
-
     def __del__(self):
         """
         Release lock when object is garbage collected.
@@ -411,7 +403,7 @@ def max_file_id(self):
         # Get max 'file_id' from 'self.row_group_stats'.
         return -1 if self.row_group_stats.empty else int(self.row_group_stats[KEY_FILE_IDS].max())
 
-    def remove_from_disk(self, preserve_metadata: bool = False):
+    def remove_from_disk(self, preserve_metadata: bool = False, release_lock: bool = True) -> None:
         """
         Remove all dataset files from disk and update in-memory state.
 
@@ -420,9 +412,19 @@ def remove_from_disk(self, preserve_metadata: bool = False):
         preserve_metadata : bool, default False
             If True, keep user metadata accessible but clear row_group_stats.
             If False, reset both row_group_stats and key_value_metadata.
+        release_lock : bool, default True
+            If True, release the lock after removal. Set to False if you plan
+            to continue using this OPD instance after removal.
+
+        Notes
+        -----
+        After calling this method with ``release_lock=True``, the OPD instance
+        should not be used for file operations, though metadata access remains
+        available if ``preserve_metadata=True``.
 
         """
         if not self._is_newly_initialized:
+            self._lock.refresh(unconditionally=True)
             # Remove dataset directory and all its contents
             remove_dir(self.dirpath)
             get_md_filepath(self.dirpath).unlink()
@@ -432,6 +434,8 @@ def remove_from_disk(self, preserve_metadata: bool = False):
         self._row_group_stats = DataFrame(columns=RGS_STATS_COLUMNS).astype(RGS_STATS_BASE_DTYPES)
         if not preserve_metadata:
             self._key_value_metadata = {}
+        if release_lock:
+            self._release_lock()
 
     def to_pandas(self) -> DataFrame:
         """
@@ -524,6 +528,15 @@ def _align_file_ids(self):
         # Set new ids.
         self._row_group_stats.loc[mask_ids_to_rename, KEY_FILE_IDS] = new_ids
 
+    def _release_lock(self):
+        """
+        Release lock with reference counting.
+        """
+        if self._lock._ref_count > 0:
+            self._lock._ref_count -= 1
+            if self._lock._ref_count == 0:
+                self._lock.unlock(unconditionally=True)
+
     def _remove_row_group_files(
         self,
         file_ids: List[int],

oups/store/ordered_parquet_dataset/ordered_parquet_dataset/read_only.py

@@ -174,6 +174,12 @@ def __setattr__(self, name: str, value: Any):
         """
         raise PermissionError(f"cannot set attribute '{name}' on a read-only dataset.")
 
+    def remove_from_disk(self, preserve_metadata=False):
+        """
+        Block dataset removal.
+        """
+        raise PermissionError("cannot call 'remove_from_disk' on a read-only dataset.")
+
     def write(self, **kwargs):
         """
         Block write operations.

oups/store/store/store.py

@@ -239,6 +239,11 @@ def __delitem__(self, key: dataclass):
             Key specifying the location where to delete the data. It has to be
             an instance of the dataclass provided at Store instantiation.
 
+        Raises
+        ------
+        KeyError
+            If the key is not found in the store.
+
         """
         if key in self.keys:
             # Get OPD instance and remove its files
@@ -250,6 +255,8 @@ def __delitem__(self, key: dataclass):
             while (upper_dir != self.basepath) and (not list(upper_dir.iterdir())):
                 upper_dir.rmdir()
                 upper_dir = upper_dir.parent
+        else:
+            raise KeyError(f"key '{key}' not found in store.")
 
     def __getitem__(self, key: dataclass) -> OrderedParquetDataset:
         """
@@ -286,7 +293,8 @@ def get(self, key: dataclass, **kwargs) -> OrderedParquetDataset:
         Returns
         -------
         OrderedParquetDataset
-            The OrderedParquetDataset instance corresponding to key.
+            The OrderedParquetDataset instance corresponding to key. Creates a
+            new OrderedParquetDataset object if the key doesn't exist.
 
         """
         opd = OrderedParquetDataset(self.basepath / key.to_path, **kwargs)

tests/test_store/test_ordered_parquet_dataset/test_opd_read_only.py

@@ -149,11 +149,18 @@ def test_exception_readonly_permission_errors(tmp_path):
     ):
         readonly_opd.ordered_on = "different_column"
 
-    # Test write() - should prevent writing data
+    # Test remove_from_disk() - should prevent clearing dataset from disk.
+    with pytest.raises(
+        PermissionError,
+        match="cannot call 'remove_from_disk' on a read-only dataset",
+    ):
+        readonly_opd.remove_from_disk()
+
+    # Test write() - should prevent writing data.
     with pytest.raises(PermissionError, match="cannot call 'write' on a read-only dataset"):
         readonly_opd.write(df=df_ref)
 
-    # Test _align_file_ids() - should prevent file ID alignment
+    # Test _align_file_ids() - should prevent file ID alignment.
     with pytest.raises(
         PermissionError,
         match="cannot call '_align_file_ids' on a read-only dataset",
@@ -167,14 +174,14 @@ def test_exception_readonly_permission_errors(tmp_path):
     ):
         readonly_opd._remove_row_group_files(file_ids=[0])
 
-    # Test _sort_row_groups() - should prevent sorting
+    # Test _sort_row_groups() - should prevent sorting.
     with pytest.raises(
         PermissionError,
         match="cannot call '_sort_row_groups' on a read-only dataset",
     ):
         readonly_opd._sort_row_groups()
 
-    # Test _write_metadata_file() - should prevent metadata writing
+    # Test _write_metadata_file() - should prevent metadata writing.
     with pytest.raises(
         PermissionError,
         match="cannot call '_write_metadata_file' on a read-only dataset",
@@ -188,7 +195,7 @@ def test_exception_readonly_permission_errors(tmp_path):
     ):
         readonly_opd._write_row_group_files([df_ref])
 
-    # Verify that read operations still work
+    # Verify that read operations still work.
     df_result = readonly_opd.to_pandas()
     assert len(df_result) == 2  # Should have 2 rows from the 2-row groups subset (indices 1 and 2)
     assert readonly_opd.ordered_on == "timestamp"

tests/test_store/test_store/test_store.py

@@ -250,3 +250,33 @@ class WeatherEntry:
     assert res[0][we2].equals(df2.iloc[:2].reset_index(drop=True))
     assert res[1][we2].equals(df2.iloc[2:2].reset_index(drop=True))
     assert res[2][we2].equals(df2.iloc[2:].reset_index(drop=True))
+
+
+def test_exception_store_delitem(tmp_path):
+    """
+    Test that Store raises KeyError for non-existent keys in get, __getitem__, and
+    __delitem__.
+    """
+
+    @toplevel
+    class Indexer:
+        category: str
+        item: str
+
+    store = Store(tmp_path, Indexer)
+    # Create a valid key and add some data
+    valid_key = Indexer(category="data", item="dataset1")
+    df = DataFrame({"timestamp": [1, 2, 3], "value": [10, 20, 30]})
+    store[valid_key].write(df=df, ordered_on="timestamp")
+    # Create a new key not in store.
+    invalid_key = Indexer(category="missing", item="notfound")
+    # Test that __delitem__ raises KeyError.
+    with pytest.raises(KeyError, match="not found"):
+        del store[invalid_key]
+    # Verify valid key still works.
+    assert valid_key in store
+    assert len(store[valid_key].row_group_stats) > 0
+    # Test after deletion, accessing should raise KeyError.
+    del store[valid_key]
+    with pytest.raises(KeyError, match="not found"):
+        del store[valid_key]