docs: Document that the write_parquet methods use mkdir in polars (#142)

Author: MoritzPotthoffQC

dataframely/collection/collection.py

@@ -864,9 +864,8 @@ def write_parquet(self, directory: str | Path, **kwargs: Any) -> None:
         members which are not provided in the current collection.
 
         Args:
-            directory: The directory the Parquet files should be written to.
-                If the directory does not exist, it is created automatically,
-                including all of its parents.
+            directory: The directory where the Parquet files should be written to.
+                The `mkdir` kwarg controls whether the directory is created if needed.
             kwargs: Additional keyword arguments passed to :meth:`polars.DataFrame.write_parquet`.
                 `metadata` may only be provided if it is a dictionary.
 

dataframely/filter_result.py

@@ -154,7 +154,8 @@ def write_parquet(self, file: str | Path | IO[bytes], **kwargs: Any) -> None:
         Args:
             file: The file path or writable file-like object to which to write the
                 parquet file. This should be a path to a directory if writing a
-                partitioned dataset.
+                partitioned dataset. The `mkdir` kwarg controls whether the directory
+                is created if needed.
             kwargs: Additional keyword arguments passed directly to
                 :meth:`polars.write_parquet`. ``metadata`` may only be provided if it
                 is a dictionary.

dataframely/schema.py

@@ -884,7 +884,8 @@ def write_parquet(
             df: The data frame to write to the parquet file.
             file: The file path or writable file-like object to which to write the
                 parquet file. This should be a path to a directory if writing a
-                partitioned dataset.
+                partitioned dataset. The `mkdir` kwarg controls whether the directory
+                is created if needed.
             kwargs: Additional keyword arguments passed directly to
                 :meth:`polars.write_parquet`. `metadata` may only be provided if it
                 is a dictionary.

tests/collection/test_storage.py

@@ -498,6 +498,50 @@ def test_read_invalid_parquet_metadata_collection(
     assert collection is None
 
 
+@pytest.mark.parametrize(
+    "any_tmp_path",
+    ["tmp_path", pytest.param("s3_tmp_path", marks=pytest.mark.s3)],
+    indirect=True,
+)
+def test_write_nonexistent_directory(any_tmp_path: str) -> None:
+    # Arrange
+    collection = MyCollection.validate(
+        {
+            "first": pl.LazyFrame({"a": [1, 2, 3]}),
+            "second": pl.LazyFrame({"a": [1, 2], "b": [10, 15]}),
+        },
+        cast=True,
+    )
+    fs: AbstractFileSystem = url_to_fs(any_tmp_path)[0]
+    target_path = fs.sep.join([any_tmp_path, "non_existent_dir"])
+
+    # Act
+    collection.write_parquet(target_path, mkdir=True)
+
+    # Assert
+    out = MyCollection.read_parquet(target_path)
+    assert_frame_equal(collection.first, out.first)
+    assert collection.second is not None
+    assert out.second is not None
+    assert_frame_equal(collection.second, out.second)
+
+
+def test_write_parquet_fails_without_mkdir(tmp_path: str) -> None:
+    # Arrange
+    collection = MyCollection.validate(
+        {
+            "first": pl.LazyFrame({"a": [1, 2, 3]}),
+            "second": pl.LazyFrame({"a": [1, 2], "b": [10, 15]}),
+        },
+        cast=True,
+    )
+    p = f"{tmp_path}/non_existent_dir"
+
+    # Act / Assert
+    with pytest.raises(FileNotFoundError):
+        collection.write_parquet(p)
+
+
 # ---------------------------- DELTA LAKE SPECIFICS ---------------------------------- #
 
 

tests/failure_info/test_storage.py

@@ -133,7 +133,10 @@ def test_invalid_schema_deserialization(
     ["tmp_path", pytest.param("s3_tmp_path", marks=pytest.mark.s3)],
     indirect=True,
 )
-def test_write_parquet_custom_metadata(any_tmp_path: str) -> None:
+@pytest.mark.parametrize("check_non_existent_directory", [True, False])
+def test_write_parquet_custom_metadata(
+    any_tmp_path: str, check_non_existent_directory: bool
+) -> None:
     # Arrange
     df = pl.DataFrame(
         {
@@ -144,10 +147,35 @@ def test_write_parquet_custom_metadata(any_tmp_path: str) -> None:
     _, failure = MySchema.filter(df)
     assert failure._df.height == 4
 
-    # Act
     fs: AbstractFileSystem = url_to_fs(any_tmp_path)[0]
-    p = fs.sep.join([any_tmp_path, "failure.parquet"])
-    failure.write_parquet(p, metadata={"custom": "test"})
+    path_components = (
+        [any_tmp_path]
+        + (["non_existent_dir"] if check_non_existent_directory else [])
+        + ["failure.parquet"]
+    )
+    p = fs.sep.join(path_components)
+
+    # Act
+    if check_non_existent_directory:
+        failure.write_parquet(p, metadata={"custom": "test"}, mkdir=True)
+    else:
+        failure.write_parquet(p, metadata={"custom": "test"})
 
     # Assert
     assert pl.read_parquet_metadata(p)["custom"] == "test"
+
+
+def test_write_parquet_fails_without_mkdir(tmp_path: str) -> None:
+    # Arrange
+    df = pl.DataFrame(
+        {
+            "a": [4, 5, 6, 6, 7, 8],
+            "b": [1, 2, 3, 4, 5, 6],
+        }
+    )
+    _, failure = MySchema.filter(df)
+    p = f"{tmp_path}/non_existent_dir/failure.parquet"
+
+    # Act / Assert
+    with pytest.raises(FileNotFoundError):
+        failure.write_parquet(p)

tests/schema/test_read_write_parquet.py

@@ -5,6 +5,7 @@
 
 import polars as pl
 import pytest
+from fsspec import url_to_fs
 
 import dataframely as dy
 from dataframely._storage.parquet import SCHEMA_METADATA_KEY
@@ -25,3 +26,36 @@ def test_read_invalid_parquet_metadata_schema(
 
     # Assert
     assert schema is None
+
+
+class MySchema(dy.Schema):
+    a = dy.Int64()
+
+
+@pytest.mark.parametrize(
+    "any_tmp_path",
+    ["tmp_path", pytest.param("s3_tmp_path", marks=pytest.mark.s3)],
+    indirect=True,
+)
+def test_write_parquet_non_existing_directory(any_tmp_path: str) -> None:
+    # Arrange
+    df = MySchema.create_empty()
+    fs = url_to_fs(any_tmp_path)[0]
+    file = fs.sep.join([any_tmp_path, "non_existing_dir", "df.parquet"])
+
+    # Act
+    MySchema.write_parquet(df, file=file, mkdir=True)
+
+    # Assert
+    result = MySchema.read_parquet(file)
+    assert result.shape == (0, 1)
+
+
+def test_write_parquet_fails_without_mkdir(tmp_path: str) -> None:
+    # Arrange
+    df = MySchema.create_empty()
+    p = f"{tmp_path}/non_existent_dir/df.parquet"
+
+    # Act / Assert
+    with pytest.raises(FileNotFoundError):
+        MySchema.write_parquet(df, file=p)