update

Author: czaloom

src/valor_lite/common/ephemeral.py

@@ -6,36 +6,32 @@
 
 
 class MemoryCache:
-    def __init__(self, table: pa.Table):
-        self._table = table
-
-    def count_tables(self) -> int:
-        """Count the number of tables in the cache."""
-        return 1
-
-    def count_rows(self) -> int:
-        """Count the number of rows in the cache."""
-        return self._table.num_rows
-
-
-class MemoryCacheReader(MemoryCache):
     def __init__(
         self,
         table: pa.Table,
         batch_size: int,
     ):
-        super().__init__(table)
-        self._schema = self._table.schema
+        self._table = table
         self._batch_size = batch_size
 
     @property
     def schema(self) -> pa.Schema:
-        return self._schema
+        return self._table.schema
 
     @property
     def batch_size(self) -> int:
         return self._batch_size
 
+    def count_tables(self) -> int:
+        """Count the number of tables in the cache."""
+        return 1
+
+    def count_rows(self) -> int:
+        """Count the number of rows in the cache."""
+        return self._table.num_rows
+
+
+class MemoryCacheReader(MemoryCache):
     def iterate_tables(
         self,
         columns: list[str] | None = None,
@@ -56,9 +52,10 @@ def __init__(
         table: pa.Table,
         batch_size: int,
     ):
-        super().__init__(table)
-        self._schema = table.schema
-        self._batch_size = batch_size
+        super().__init__(
+            table=table,
+            batch_size=batch_size,
+        )
 
         # internal state
         self._buffer = []
@@ -98,7 +95,7 @@ def write_rows(
         """
         if not rows:
             return
-        batch = pa.RecordBatch.from_pylist(rows, schema=self._schema)
+        batch = pa.RecordBatch.from_pylist(rows, schema=self.schema)
         self.write_batch(batch)
 
     def write_columns(
@@ -143,10 +140,10 @@ def write_batch(
             self._buffer.append(batch)
             combined_arrays = [
                 pa.concat_arrays([b.column(name) for b in self._buffer])
-                for name in self._schema.names
+                for name in self.schema.names
             ]
             batch = pa.RecordBatch.from_arrays(
-                combined_arrays, schema=self._schema
+                combined_arrays, schema=self.schema
             )
             self._buffer = []
 
@@ -172,10 +169,10 @@ def flush(self):
         if self._buffer:
             combined_arrays = [
                 pa.concat_arrays([b.column(name) for b in self._buffer])
-                for name in self._schema.names
+                for name in self.schema.names
             ]
             batch = pa.RecordBatch.from_arrays(
-                combined_arrays, schema=self._schema
+                combined_arrays, schema=self.schema
             )
             self._table = pa.concat_tables(
                 [self._table, pa.Table.from_batches([batch])]

src/valor_lite/common/persistent.py

@@ -13,13 +13,40 @@
 
 
 class FileCache:
-    def __init__(self, path: str | Path):
+    def __init__(
+        self,
+        path: str | Path,
+        schema: pa.Schema,
+        batch_size: int,
+        rows_per_file: int,
+        compression: str,
+    ):
         self._path = Path(path)
+        self._schema = schema
+        self._batch_size = batch_size
+        self._rows_per_file = rows_per_file
+        self._compression = compression
 
     @property
     def path(self) -> Path:
         return self._path
 
+    @property
+    def schema(self) -> pa.Schema:
+        return self._schema
+
+    @property
+    def batch_size(self) -> int:
+        return self._batch_size
+
+    @property
+    def rows_per_file(self) -> int:
+        return self._rows_per_file
+
+    @property
+    def compression(self) -> str:
+        return self._compression
+
     @staticmethod
     def _generate_config_path(path: str | Path) -> Path:
         """Generate cache configuration path."""
@@ -84,38 +111,8 @@ def get_dataset_files(self) -> list[Path]:
 
 
 class FileCacheReader(FileCache):
-    def __init__(
-        self,
-        path: str | Path,
-        schema: pa.Schema,
-        batch_size: int,
-        rows_per_file: int,
-        compression: str,
-    ):
-        super().__init__(path)
-        self._schema = schema
-        self._batch_size = batch_size
-        self._rows_per_file = rows_per_file
-        self._compression = compression
-
-    @property
-    def schema(self) -> pa.Schema:
-        return self._schema
-
-    @property
-    def batch_size(self) -> int:
-        return self._batch_size
-
-    @property
-    def rows_per_file(self) -> int:
-        return self._rows_per_file
-
-    @property
-    def compression(self) -> str:
-        return self._compression
-
     @classmethod
-    def load(cls, path: str | Path | FileCache):
+    def load(cls, path: str | Path):
         """
         Load cache from disk.
 
@@ -124,8 +121,6 @@ def load(cls, path: str | Path | FileCache):
         path : str | Path
             Where the cache is stored.
         """
-        if isinstance(path, FileCache):
-            path = path.path
         path = Path(path)
         if not path.exists():
             raise FileNotFoundError(f"Directory does not exist: {path}")
@@ -192,11 +187,13 @@ def __init__(
         rows_per_file: int,
         compression: str,
     ):
-        super().__init__(path)
-        self._schema = schema
-        self._batch_size = batch_size
-        self._rows_per_file = rows_per_file
-        self._compression = compression
+        super().__init__(
+            path=path,
+            schema=schema,
+            batch_size=batch_size,
+            rows_per_file=rows_per_file,
+            compression=compression,
+        )
 
         # internal state
         self._writer = None

tests/common/test_ephemeral_cache.py

@@ -29,15 +29,18 @@ def test_cache_reader():
         )
         writer.write_table(tbl)
         assert writer.count_rows() == 101
+        assert writer.count_tables() == 1
 
     reader = writer.to_reader()
     assert reader.count_rows() == 101
+    assert reader.count_tables() == 1
     for tbl in reader.iterate_tables():
         assert tbl["some_int"].to_pylist() == [i for i in range(101)]
         assert tbl["some_float"].to_pylist() == [i for i in range(101)]
         assert tbl["some_str"].to_pylist() == [f"str{i}" for i in range(101)]
     assert reader.count_rows() == 101
-    assert reader._schema == schema
+    assert reader.count_tables() == 1
+    assert reader.schema == schema
 
 
 def test_cache_write_columns():
@@ -60,9 +63,11 @@ def test_cache_write_columns():
                     "some_str": pa.array([f"str{i}"]),
                 }
             )
+        writer.write_columns({})
 
     reader = writer.to_reader()
     assert reader.count_rows() == 1000
+    assert reader.count_tables() == 1
     for tbl in reader.iterate_tables():
         assert tbl["some_int"].to_pylist() == [i for i in range(1000)]
         assert tbl["some_float"].to_pylist() == [i for i in range(1000)]
@@ -98,6 +103,7 @@ def test_cache_write_rows():
 
     reader = writer.to_reader()
     assert reader.count_rows() == 1000
+    assert reader.count_tables() == 1
     for tbl in reader.iterate_tables():
         assert tbl["some_int"].to_pylist() == [i for i in range(1000)]
         assert tbl["some_float"].to_pylist() == [i for i in range(1000)]
@@ -127,12 +133,15 @@ def test_cache_write_table():
             ]
         )
         assert writer.count_rows() == 0
+        assert writer.count_tables() == 1
 
         writer.write_table(tbl)
         assert writer.count_rows() == 101
+        assert writer.count_tables() == 1
 
         writer.write_table(tbl)
         assert writer.count_rows() == 202
+        assert writer.count_tables() == 1
 
         reader = writer.to_reader()
 

tests/common/test_exceptions.py

@@ -9,6 +9,12 @@ def test_empty_evaluator_exc():
     assert "no data" in str(e)
 
 
+def test_empty_cache_error_exc():
+    with pytest.raises(exc.EmptyCacheError) as e:
+        raise exc.EmptyCacheError()
+    assert "no data" in str(e)
+
+
 def test_cache_error_exc():
     with pytest.raises(exc.InternalCacheError) as e:
         raise exc.InternalCacheError("custom message")

tests/common/test_persistent_cache.py

@@ -13,7 +13,13 @@
 
 
 def test_cache_files_empty(tmp_path: Path):
-    cache = FileCache(tmp_path)
+    cache = FileCache(
+        tmp_path,
+        schema=None,  # type: ignore - testing
+        batch_size=None,  # type: ignore - testing
+        rows_per_file=None,  # type: ignore - testing
+        compression=None,  # type: ignore - testing
+    )
     assert cache._path == tmp_path
     assert cache.get_files() == []
     assert cache.get_dataset_files() == []
@@ -22,7 +28,13 @@ def test_cache_files_empty(tmp_path: Path):
 
 def test_cache_files_does_not_exist(tmp_path: Path):
     path = tmp_path / "does_not_exist"
-    cache = FileCache(path)
+    cache = FileCache(
+        path,
+        schema=None,  # type: ignore - testing
+        batch_size=None,  # type: ignore - testing
+        rows_per_file=None,  # type: ignore - testing
+        compression=None,  # type: ignore - testing
+    )
     assert cache._path == path
     assert cache.get_files() == []
     assert cache.get_dataset_files() == []
@@ -93,6 +105,7 @@ def test_cache_reader(tmp_path: Path):
             ]
 
         assert reader.count_rows() == 202
+        assert reader.count_tables() == 2
         assert reader._schema == schema
 
 
@@ -167,11 +180,13 @@ def test_cache_write_columns(tmp_path: Path):
                     "some_str": pa.array([f"str{i}"]),
                 }
             )
+        writer.write_columns({})
         writer.flush()
         assert len(writer.get_files()) == 11
 
         reader = FileCacheReader.load(writer.path)
         assert reader.count_rows() == 1000
+        assert reader.count_tables() == 10
         for idx, tbl in enumerate(reader.iterate_tables()):
             assert tbl["some_int"].to_pylist() == [
                 i
@@ -221,6 +236,7 @@ def test_cache_write_rows(tmp_path: Path):
 
         reader = FileCacheReader.load(writer.path)
         assert reader.count_rows() == 1000
+        assert reader.count_tables() == 10
         for idx, tbl in enumerate(reader.iterate_tables()):
             assert tbl["some_int"].to_pylist() == [
                 i
@@ -269,6 +285,7 @@ def test_cache_write_table(tmp_path: Path):
 
         reader = FileCacheReader.load(writer.path)
         assert reader.count_rows() == 202
+        assert reader.count_tables() == 2
         for tbl in reader.iterate_tables():
             assert tbl["some_int"].to_pylist() == [i for i in range(101)]
             assert tbl["some_float"].to_pylist() == [i for i in range(101)]
@@ -281,6 +298,33 @@ def test_cache_delete(tmp_path: Path):
     batch_size = 10
     rows_per_file = 100
     path = tmp_path / "cache"
+
+    FileCacheWriter.create(
+        path=path,
+        schema=pa.schema(
+            [
+                ("some_int", pa.int64()),
+                ("some_float", pa.float64()),
+                ("some_str", pa.string()),
+            ]
+        ),
+        batch_size=batch_size,
+        rows_per_file=rows_per_file,
+    )
+    with pytest.raises(FileExistsError):
+        FileCacheWriter.create(
+            path=path,
+            schema=pa.schema(
+                [
+                    ("some_int", pa.int64()),
+                    ("some_float", pa.float64()),
+                    ("some_str", pa.string()),
+                ]
+            ),
+            batch_size=batch_size,
+            rows_per_file=rows_per_file,
+        )
+
     with FileCacheWriter.create(
         path=path,
         schema=pa.schema(
@@ -292,6 +336,7 @@ def test_cache_delete(tmp_path: Path):
         ),
         batch_size=batch_size,
         rows_per_file=rows_per_file,
+        delete_if_exists=True,
     ) as writer:
         tbl = pa.Table.from_pylist(
             [
@@ -311,6 +356,7 @@ def test_cache_delete(tmp_path: Path):
 
         reader = FileCacheReader.load(path=writer.path)
         assert reader.count_rows() == 202
+        assert reader.count_tables() == 2
         for tbl in reader.iterate_tables():
             assert tbl["some_int"].to_pylist() == [i for i in range(101)]
             assert tbl["some_float"].to_pylist() == [i for i in range(101)]