diff --git a/README.md b/README.md index 92d6ebf..5094017 100644 --- a/README.md +++ b/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` diff --git a/src/simlab_tools/storage/dataset.py b/src/simlab_tools/storage/dataset.py index 58f9ff5..ba7c583 100644 --- a/src/simlab_tools/storage/dataset.py +++ b/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: diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 1bb43bf..0049f2e 100644 --- a/tests/test_dataset.py +++ b/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()