test: annotate the read/write, blob and fragment suites

tests/test_basic_read_write.py

@@ -3,7 +3,9 @@
 import os
 import sys
 import tempfile
+from collections.abc import Iterator
 from pathlib import Path
+from typing import Any
 
 import lance
 import lance_ray as lr
@@ -24,7 +26,7 @@
 
 
 @pytest.fixture
-def sample_data():
+def sample_data() -> pd.DataFrame:
     """Create sample data for testing."""
     return pd.DataFrame(
         {
@@ -37,22 +39,22 @@ def sample_data():
 
 
 @pytest.fixture
-def temp_dir():
+def temp_dir() -> Iterator[str]:
     """Create a temporary directory for testing."""
     with tempfile.TemporaryDirectory() as temp_dir:
         yield temp_dir
 
 
 @pytest.fixture
-def sample_dataset(sample_data):
+def sample_dataset(sample_data: pd.DataFrame) -> Dataset:
     """Create a Ray Dataset from sample data."""
     return ray.data.from_pandas(sample_data)
 
 
 class TestWriteLance:
     """Test cases for write_lance function."""
 
-    def test_write_lance_basic(self, sample_dataset, temp_dir):
+    def test_write_lance_basic(self, sample_dataset: Dataset, temp_dir: str) -> None:
         """Test basic write functionality."""
         path = Path(temp_dir) / "basic_write.lance"
 
@@ -61,7 +63,9 @@ def test_write_lance_basic(self, sample_dataset, temp_dir):
         assert path.exists()
         assert path.is_dir()
 
-    def test_write_lance_with_stable_row_ids(self, sample_dataset, temp_dir):
+    def test_write_lance_with_stable_row_ids(
+        self, sample_dataset: Dataset, temp_dir: str
+    ) -> None:
         path = Path(temp_dir) / "stable_row_ids.lance"
 
         lr.write_lance(
@@ -72,29 +76,31 @@ def test_write_lance_with_stable_row_ids(self, sample_dataset, temp_dir):
 
         assert lance.dataset(str(path)).has_stable_row_ids
 
-    def test_write_lance_with_schema(self, temp_dir):
+    def test_write_lance_with_schema(self, temp_dir: str) -> None:
         """Test write with explicit schema."""
         path = Path(temp_dir) / "schema_write.lance"
 
         data = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})
         dataset = ray.data.from_pandas(data)
 
-        schema = pa.schema(
-            [pa.field("col1", pa.int64()), pa.field("col2", pa.string())]
-        )
+        schema_fields: list[pa.Field[Any]] = [
+            pa.field("col1", pa.int64()),
+            pa.field("col2", pa.string()),
+        ]
+        schema = pa.schema(schema_fields)
 
         lr.write_lance(dataset, str(path), schema=schema)
         assert path.exists()
 
-    def test_write_lance_invalid_input(self, temp_dir):
+    def test_write_lance_invalid_input(self, temp_dir: str) -> None:
         """Test error handling for invalid inputs."""
         path = Path(temp_dir) / "invalid.lance"
 
         with pytest.raises((ValueError, AttributeError, TypeError)):
-            lr.write_lance(None, str(path))  # type: ignore
+            lr.write_lance(None, str(path))  # type: ignore[arg-type]
 
-    def test_write_with_pandas_map_batches(self, temp_dir):
-        def map_fn(row):
+    def test_write_with_pandas_map_batches(self, temp_dir: str) -> None:
+        def map_fn(row: dict[str, Any]) -> dict[str, Any]:
             return {
                 "id": row["id"],
                 "name": row["name"],
@@ -103,18 +109,17 @@ def map_fn(row):
                 "extra": None,
             }
 
-        def to_pd(batch: pd.DataFrame):
+        def to_pd(batch: pd.DataFrame) -> pd.DataFrame:
             return batch
 
-        schema = pa.schema(
-            [
-                pa.field("id", pa.int64()),
-                pa.field("name", pa.string()),
-                pa.field("age", pa.int32()),
-                pa.field("score", pa.float64()),
-                pa.field("extra", pa.string()),
-            ]
-        )
+        schema_fields: list[pa.Field[Any]] = [
+            pa.field("id", pa.int64()),
+            pa.field("name", pa.string()),
+            pa.field("age", pa.int32()),
+            pa.field("score", pa.float64()),
+            pa.field("extra", pa.string()),
+        ]
+        schema = pa.schema(schema_fields)
         data = pd.DataFrame(
             {
                 "id": [1, 2, 3, 4, 5],
@@ -137,22 +142,21 @@ def to_pd(batch: pd.DataFrame):
         tbl = ds.to_table()
         assert set(data["name"].tolist()) == set(tbl["name"].to_pylist())
 
-    def test_write_lance_preserves_nested_struct_fields(self, temp_dir):
+    def test_write_lance_preserves_nested_struct_fields(self, temp_dir: str) -> None:
         path = Path(temp_dir) / "nested_struct.lance"
-        schema = pa.schema(
-            [
-                pa.field("id", pa.int64()),
-                pa.field(
-                    "meta",
-                    pa.struct(
-                        [
-                            pa.field("userId", pa.string()),
-                            pa.field("a.b", pa.string()),
-                        ]
-                    ),
+        schema_fields: list[pa.Field[Any]] = [
+            pa.field("id", pa.int64()),
+            pa.field(
+                "meta",
+                pa.struct(
+                    [
+                        pa.field("userId", pa.string()),
+                        pa.field("a.b", pa.string()),
+                    ]
                 ),
-            ]
-        )
+            ),
+        ]
+        schema = pa.schema(schema_fields)
         table = pa.Table.from_arrays(
             [
                 pa.array([1, 2], type=pa.int64()),
@@ -190,13 +194,13 @@ class TestReadLance:
     """Test cases for read_lance function."""
 
     @pytest.fixture
-    def lance_dataset_path(self, sample_dataset, temp_dir):
+    def lance_dataset_path(self, sample_dataset: Dataset, temp_dir: str) -> str:
         """Create a Lance dataset for reading tests."""
         path = Path(temp_dir) / "test_dataset.lance"
         lr.write_lance(sample_dataset, str(path))
         return str(path)
 
-    def test_read_lance_basic(self, lance_dataset_path):
+    def test_read_lance_basic(self, lance_dataset_path: str) -> None:
         """Test basic read functionality."""
         dataset = lr.read_lance(lance_dataset_path)
 
@@ -206,23 +210,23 @@ def test_read_lance_basic(self, lance_dataset_path):
         assert len(df) == 5
         assert list(df.columns) == ["id", "name", "age", "score"]
 
-    def test_read_lance_with_columns(self, lance_dataset_path):
+    def test_read_lance_with_columns(self, lance_dataset_path: str) -> None:
         """Test reading specific columns."""
         dataset = lr.read_lance(lance_dataset_path, columns=["id", "name"])
 
         df = dataset.to_pandas()
         assert list(df.columns) == ["id", "name"]
         assert len(df) == 5
 
-    def test_read_lance_with_filter(self, lance_dataset_path):
+    def test_read_lance_with_filter(self, lance_dataset_path: str) -> None:
         """Test reading with filter."""
         dataset = lr.read_lance(lance_dataset_path, filter="age > 30")
 
         df = dataset.to_pandas()
         assert len(df) == 3
         assert all(df["age"] > 30)
 
-    def test_read_lance_columns_and_filter(self, lance_dataset_path):
+    def test_read_lance_columns_and_filter(self, lance_dataset_path: str) -> None:
         """Test reading with both columns and filter."""
         dataset = lr.read_lance(
             lance_dataset_path, columns=["name", "age"], filter="age >= 35"
@@ -233,14 +237,14 @@ def test_read_lance_columns_and_filter(self, lance_dataset_path):
         assert len(df) == 3
         assert all(df["age"] >= 35)
 
-    def test_read_lance_filter_and_count(self, lance_dataset_path):
+    def test_read_lance_filter_and_count(self, lance_dataset_path: str) -> None:
         """Test reading filter and count."""
         dataset = lr.read_lance(
             lance_dataset_path, columns=["name", "age"], filter="age >= 35"
         )
         assert dataset.count() == 3
 
-    def test_read_lance_nonexistent_path(self):
+    def test_read_lance_nonexistent_path(self) -> None:
         """Test reading from non-existent path."""
         with pytest.raises((FileNotFoundError, OSError, Exception)):
             lr.read_lance("/path/that/does/not/exist")
@@ -249,7 +253,9 @@ def test_read_lance_nonexistent_path(self):
 class TestReadWrite:
     """Integration tests for read and write operations."""
 
-    def test_write_then_read_roundtrip(self, sample_data, temp_dir):
+    def test_write_then_read_roundtrip(
+        self, sample_data: pd.DataFrame, temp_dir: str
+    ) -> None:
         """Test writing data and then reading it back."""
         path = Path(temp_dir) / "roundtrip.lance"
 
@@ -267,7 +273,7 @@ def test_write_then_read_roundtrip(self, sample_data, temp_dir):
 
         pd.testing.assert_frame_equal(original_sorted, read_sorted)
 
-    def test_append_mode(self, sample_data, temp_dir):
+    def test_append_mode(self, sample_data: pd.DataFrame, temp_dir: str) -> None:
         """Test append mode with read verification."""
         path = Path(temp_dir) / "append_test.lance"
 
@@ -293,7 +299,7 @@ def test_append_mode(self, sample_data, temp_dir):
 
         assert len(full_df) == 5  # 3 initial + 2 appended
 
-    def test_overwrite_mode(self, sample_dataset, temp_dir):
+    def test_overwrite_mode(self, sample_dataset: Dataset, temp_dir: str) -> None:
         """Test different write modes."""
         path = Path(temp_dir) / "modes_test.lance"
 
@@ -326,7 +332,9 @@ def test_overwrite_mode(self, sample_dataset, temp_dir):
         overwritten_df = overwritten_dataset.to_pandas()
         assert len(overwritten_df) == 2  # Should have 2 rows after overwrite
 
-    def test_read_lance_with_fragment_ids(self, sample_dataset, temp_dir):
+    def test_read_lance_with_fragment_ids(
+        self, sample_dataset: Dataset, temp_dir: str
+    ) -> None:
         """Test reading with fragment IDs."""
         path = Path(temp_dir) / "fragment_ids_test.lance"
         lr.write_lance(
@@ -339,7 +347,7 @@ def test_read_lance_with_fragment_ids(self, sample_dataset, temp_dir):
 class TestAddColumns:
     """Test cases for add_columns function."""
 
-    def test_add_columns_basic(self, sample_dataset, temp_dir):
+    def test_add_columns_basic(self, sample_dataset: Dataset, temp_dir: str) -> None:
         """Test basic add columns functionality."""
         path = Path(temp_dir) / "add_columns_test.lance"
         lr.write_lance(
@@ -370,7 +378,9 @@ def double_score(x: pa.RecordBatch) -> pa.RecordBatch:
 class TestNamespaceReadWrite:
     """Test cases for read/write with DirectoryNamespace."""
 
-    def test_write_and_read_with_directory_namespace(self, sample_data, temp_dir):
+    def test_write_and_read_with_directory_namespace(
+        self, sample_data: pd.DataFrame, temp_dir: str
+    ) -> None:
         """Test write and read using DirectoryNamespace."""
         table_id = ["test_table"]
 
@@ -398,7 +408,7 @@ def test_write_and_read_with_directory_namespace(self, sample_data, temp_dir):
 class TestDatasetOptions:
     """Test cases for dataset options in LanceDataset."""
 
-    def test_dataset_with_version(self, sample_dataset, temp_dir):
+    def test_dataset_with_version(self, sample_dataset: Dataset, temp_dir: str) -> None:
         """Test dataset options like version and block size."""
         path = Path(temp_dir) / "dataset_options_test.lance"
         lr.write_lance(sample_dataset, str(path))
@@ -429,12 +439,12 @@ def test_dataset_with_version(self, sample_dataset, temp_dir):
     from lance import DatasetBasePath, blob_array, blob_field
 except Exception:
 
-    class _Missing:  # type: ignore[no-redef]
+    class _Missing:
         pass
 
-    DatasetBasePath = _Missing
-    blob_array = _Missing
-    blob_field = _Missing
+    DatasetBasePath = _Missing  # type: ignore[assignment,misc]
+    blob_array = _Missing  # type: ignore[assignment]
+    blob_field = _Missing  # type: ignore[assignment]
 
 
 class TestMultiBaseLayout:
@@ -457,7 +467,7 @@ class TestMultiBaseLayout:
         ``Duplicate base path ID 0`` error.
     """
 
-    def test_multiple_initial_bases_without_explicit_id(self, temp_dir):
+    def test_multiple_initial_bases_without_explicit_id(self, temp_dir: str) -> None:
         """Multiple DatasetBasePath objects without explicit id should not collide.
 
         When the user provides two (or more) ``DatasetBasePath`` objects
@@ -494,7 +504,7 @@ def test_multiple_initial_bases_without_explicit_id(self, temp_dir):
         ds = lance.dataset(uri)
         assert ds.count_rows() == 3
 
-        base_paths = ds._ds.base_paths()
+        base_paths = ds._ds.base_paths()  # type: ignore[attr-defined]
         assert len(base_paths) >= 2
         base_ids = list(base_paths.keys())
         assert len(set(base_ids)) == len(base_ids), (
@@ -505,7 +515,7 @@ def test_multiple_initial_bases_without_explicit_id(self, temp_dir):
         bool(missing_fragment_write_options("base_store_params")),
         reason=fragment_write_options_skip_reason("base_store_params"),
     )
-    def test_multiple_initial_bases_with_blob_v2(self, temp_dir):
+    def test_multiple_initial_bases_with_blob_v2(self, temp_dir: str) -> None:
         """Multi-base write/read with blob v2 columns and no explicit IDs.
 
         This is the full end-to-end scenario: blob data lives across
@@ -536,7 +546,7 @@ def test_multiple_initial_bases_with_blob_v2(self, temp_dir):
         )
         ray_ds = ray.data.from_arrow(table)
 
-        base_store_params = {
+        base_store_params: dict[str, dict[str, Any]] = {
             base1_dir.as_uri(): {},
             base2_dir.as_uri(): {},
         }
@@ -557,14 +567,14 @@ def test_multiple_initial_bases_with_blob_v2(self, temp_dir):
         ds = lance.dataset(uri, base_store_params=base_store_params)
         assert ds.count_rows() == 2
 
-        base_paths = ds._ds.base_paths()
+        base_paths = ds._ds.base_paths()  # type: ignore[attr-defined]
         assert len(base_paths) >= 2
         base_ids = list(base_paths.keys())
         assert len(set(base_ids)) == len(base_ids), (
             f"Base path IDs must be unique, got: {base_paths}"
         )
 
-    def test_explicit_ids_are_preserved(self, temp_dir):
+    def test_explicit_ids_are_preserved(self, temp_dir: str) -> None:
         """When the user provides explicit IDs, they must be preserved."""
         base1_dir = Path(temp_dir) / "explicit_base1"
         base2_dir = Path(temp_dir) / "explicit_base2"
@@ -594,11 +604,11 @@ def test_explicit_ids_are_preserved(self, temp_dir):
         ds = lance.dataset(uri)
         assert ds.count_rows() == 2
 
-        base_paths = ds._ds.base_paths()
+        base_paths = ds._ds.base_paths()  # type: ignore[attr-defined]
         assert 5 in base_paths
         assert 10 in base_paths
 
-    def test_mixed_explicit_and_implicit_ids(self, temp_dir):
+    def test_mixed_explicit_and_implicit_ids(self, temp_dir: str) -> None:
         """One base with explicit id, one without — no collision."""
         base1_dir = Path(temp_dir) / "mixed_base1"
         base2_dir = Path(temp_dir) / "mixed_base2"
@@ -628,14 +638,14 @@ def test_mixed_explicit_and_implicit_ids(self, temp_dir):
         ds = lance.dataset(uri)
         assert ds.count_rows() == 2
 
-        base_paths = ds._ds.base_paths()
+        base_paths = ds._ds.base_paths()  # type: ignore[attr-defined]
         assert 3 in base_paths
         base_ids = list(base_paths.keys())
         assert len(set(base_ids)) == len(base_ids), (
             f"Base path IDs must be unique, got: {base_paths}"
         )
 
-    def test_dataset_root_base_gets_id_zero(self, temp_dir):
+    def test_dataset_root_base_gets_id_zero(self, temp_dir: str) -> None:
         """A base with is_dataset_root=True should receive id=0."""
         base1_dir = Path(temp_dir) / "root_base"
         base2_dir = Path(temp_dir) / "extra_base"
@@ -667,5 +677,5 @@ def test_dataset_root_base_gets_id_zero(self, temp_dir):
         ds = lance.dataset(uri)
         assert ds.count_rows() == 2
 
-        base_paths = ds._ds.base_paths()
+        base_paths = ds._ds.base_paths()  # type: ignore[attr-defined]
         assert 0 in base_paths