fix(storage): accept LazyFrame in write_dataset

README.md

@@ -2,7 +2,7 @@
 
 Shared research tooling for SIM lab projects: loading weather data, geodata, and
 transferring files to and from object storage. The package is organised into
-topical subpackages, each keeping its heavy dependencies behind an optional
+topical sub-packages, each keeping its heavy dependencies behind an optional
 extra so consumers install only what they use.
 
 ## Installation
@@ -11,13 +11,13 @@ Install straight from GitHub, selecting the extras you need:
 
 ```bash
 # Object storage only
-uv pip install "simlab-tools[storage] @ git+https://github.com/simlab-vs/simlab-tools.git"
+uv add "simlab-tools[storage] @ git+https://github.com/simlab-vs/simlab-tools.git"
 
 # Geodata utilities
-uv pip install "simlab-tools[geo] @ git+https://github.com/simlab-vs/simlab-tools.git"
+uv add "simlab-tools[geo] @ git+https://github.com/simlab-vs/simlab-tools.git"
 
 # Everything
-uv pip install "simlab-tools[all] @ git+https://github.com/simlab-vs/simlab-tools.git"
+uv add "simlab-tools[all] @ git+https://github.com/simlab-vs/simlab-tools.git"
 ```
 
 Available extras: `storage`, `geo`, `smoothing`, `weather`, `all`, and `dev`

src/simlab_tools/storage/dataset.py

@@ -22,8 +22,8 @@
 
 from simlab_tools.storage.credentials import resolve_s3_credentials
 
-# A dataset can be written from either a polars DataFrame or an Arrow table.
-TableLike = pl.DataFrame | pa.Table
+# A dataset can be written from a polars DataFrame/LazyFrame or an Arrow table.
+TableLike = pl.DataFrame | pl.LazyFrame | pa.Table
 
 
 def get_s3_filesystem(
@@ -102,8 +102,9 @@ def write_dataset(
 
     Parameters
     ----------
-    data: polars.DataFrame or pyarrow.Table
-        The table to write. A polars DataFrame is converted to Arrow first.
+    data: polars.DataFrame, polars.LazyFrame, or pyarrow.Table
+        The table to write. A polars DataFrame or LazyFrame is converted to
+        Arrow first (a LazyFrame is collected before conversion).
     filesystem: pyarrow.fs.FileSystem
         Target filesystem, typically from :func:`get_s3_filesystem`.
     bucket: str
@@ -123,7 +124,12 @@ def write_dataset(
         Template for the written file names, e.g. ``"part-{i}.parquet"``.
 
     """
-    table = data if isinstance(data, pa.Table) else data.to_arrow()
+    if isinstance(data, pa.Table):
+        table = data
+    elif isinstance(data, pl.LazyFrame):
+        table = data.collect().to_arrow()
+    else:
+        table = data.to_arrow()
 
     partitioning = None
     if partition_cols:

tests/test_dataset.py

@@ -75,6 +75,21 @@ def test_column_projection(self, sample_df, tmp_path):
 
         assert result.columns == ["value"]
 
+    def test_lazyframe_round_trip(self, sample_df, tmp_path):
+        """A LazyFrame is collected before writing."""
+        fs = LocalFileSystem()
+        write_dataset(
+            sample_df.lazy(), fs, str(tmp_path), "events", partition_cols=["year"]
+        )
+
+        result = read_dataset(fs, str(tmp_path), "events")
+
+        assert_frame_equal(
+            result.select(sample_df.columns).sort("value"),
+            sample_df.sort("value"),
+            check_dtypes=False,
+        )
+
     def test_unpartitioned_round_trip(self, sample_df, tmp_path):
         """Datasets can be written without any partition columns."""
         fs = LocalFileSystem()