fix: not supported

Author: skoschik

awswrangler/s3/_write_deltalake.py

@@ -193,8 +193,6 @@ def to_deltalake_streaming(
     boto3_session: boto3.Session | None = None,
     s3_additional_kwargs: dict[str, str] | None = None,
     batch_size: int | None = None,
-    max_open_files: int | None = None,
-    max_rows_per_file: int | None = None,
     target_file_size: int | None = None,
 ) -> None:
     dtype = dtype or {}
@@ -224,7 +222,5 @@ def to_deltalake_streaming(
         mode=mode,
         schema_mode=schema_mode,
         storage_options=storage_options,
-        max_open_files=max_open_files,
-        max_rows_per_file=max_rows_per_file,
         target_file_size=target_file_size,
     )

tests/unit/test_s3_deltalake.py

@@ -1,24 +1,36 @@
 from __future__ import annotations
 
-from typing import Any, Iterator
+from typing import Any, Iterator, Iterable
 
 import boto3
+import pyarrow as pa
 import pytest
+from pandas.testing import assert_frame_equal
 
 import awswrangler as wr
 import awswrangler.pandas as pd
+from awswrangler.s3._write_deltalake import _df_iter_to_record_batch_reader
 
 from .._utils import (
     get_time_str_with_random_suffix,
 )
 
 
+def assert_df_equal_unordered(left: pd.DataFrame, right: pd.DataFrame, by: list[str]) -> None:
+    """Compare two dataframes ignoring row order and dtypes."""
+    l2 = left.sort_values(by).reset_index(drop=True)
+    r2 = right.sort_values(by).reset_index(drop=True)
+
+    assert_frame_equal(l2, r2, check_dtype=False, check_like=True)
+
+
 @pytest.fixture(scope="session")
 def lock_dynamodb_table() -> Iterator[str]:
     name = f"deltalake_lock_{get_time_str_with_random_suffix()}"
     print(f"Table name: {name}")
 
     dynamodb_client = boto3.client("dynamodb")
+
     dynamodb_client.create_table(
         TableName=name,
         BillingMode="PAY_PER_REQUEST",
@@ -94,3 +106,165 @@ def test_read_deltalake_partitions(path: str, lock_settings: dict[str, Any]) ->
 
     df2 = wr.s3.read_deltalake(path=path, columns=["c0"], partitions=[("par0", "=", "foo"), ("par1", "=", "1")])
     assert df2.shape == (1, 1)
+
+
+@pytest.mark.parametrize("chunksize", [2, 10])
+def test_to_deltalake_streaming_single_commit_overwrite(
+    path: str,
+    lock_settings: dict[str, Any],
+    chunksize: int,
+) -> None:
+    df1 = pd.DataFrame({"c0": [1, 1], "c1": [10, 11], "v": [100, 200]})
+    df2 = pd.DataFrame({"c0": [2, 2], "c1": [12, 13], "v": [300, 400]})
+
+    def dfs() -> Iterable[pd.DataFrame]:
+        yield df1
+        yield df2
+
+    wr.s3.to_deltalake_streaming(
+        dfs=dfs(),
+        path=path,
+        mode="overwrite",
+        partition_cols=["c0", "c1"],
+        **lock_settings,
+    )
+
+    out = wr.s3.read_deltalake(path=path)
+
+    expected = pd.concat([df1, df2], ignore_index=True)
+    assert_df_equal_unordered(expected, out, by=["c0", "c1", "v"])
+
+
+def test_to_deltalake_streaming_creates_one_version_per_run(
+    path: str,
+    lock_settings: dict[str, Any],
+) -> None:
+    df_run1_a = pd.DataFrame({"c0": [1], "c1": [10], "v": [111]})
+    df_run1_b = pd.DataFrame({"c0": [1], "c1": [11], "v": [112]})
+
+    wr.s3.to_deltalake_streaming(
+        dfs=[df_run1_a, df_run1_b],
+        path=path,
+        mode="overwrite",
+        partition_cols=["c0", "c1"],
+        **lock_settings,
+    )
+
+    run1_expected = pd.concat([df_run1_a, df_run1_b], ignore_index=True)
+    latest_v0 = wr.s3.read_deltalake(path=path)
+    assert_df_equal_unordered(run1_expected, latest_v0, by=["c0", "c1", "v"])
+
+    df_run2_a = pd.DataFrame({"c0": [2], "c1": [12], "v": [221]})
+    df_run2_b = pd.DataFrame({"c0": [2], "c1": [13], "v": [222]})
+
+    wr.s3.to_deltalake_streaming(
+        dfs=[df_run2_a, df_run2_b],
+        path=path,
+        mode="overwrite",
+        partition_cols=["c0", "c1"],
+        **lock_settings,
+    )
+
+    v0 = wr.s3.read_deltalake(path=path, version=0)
+    v1 = wr.s3.read_deltalake(path=path, version=1)
+    run2_expected = pd.concat([df_run2_a, df_run2_b], ignore_index=True)
+
+    assert_df_equal_unordered(run1_expected, v0, by=["c0", "c1", "v"])
+    assert_df_equal_unordered(run2_expected, v1, by=["c0", "c1", "v"])
+
+
+def test_to_deltalake_streaming_partitions_and_filters(
+    path: str,
+    lock_settings: dict[str, Any],
+) -> None:
+    df1 = pd.DataFrame({"c0": [1, 1, 2], "c1": [10, 11, 12], "v": [1, 2, 3]})
+    df2 = pd.DataFrame({"c0": [2, 3, 3], "c1": [13, 14, 15], "v": [4, 5, 6]})
+
+    wr.s3.to_deltalake_streaming(
+        dfs=[df1, df2],
+        path=path,
+        mode="overwrite",
+        partition_cols=["c0", "c1"],
+        **lock_settings,
+    )
+
+    only_c02 = wr.s3.read_deltalake(
+        path=path,
+        partitions=[("c0", "=", "2")],
+        columns=["v", "c1"],
+    )
+    assert set(only_c02["c1"].tolist()) == {12, 13}
+    assert sorted(only_c02["v"].tolist()) == [3, 4]
+
+
+def test_to_deltalake_streaming_empty_iterator_is_noop(
+    path: str,
+    lock_settings: dict[str, Any],
+) -> None:
+    wr.s3.to_deltalake_streaming(
+        dfs=[pd.DataFrame({"c0": [1], "c1": [1], "v": [1]})],
+        path=path,
+        mode="overwrite",
+        partition_cols=["c0", "c1"],
+        **lock_settings,
+    )
+    baseline = wr.s3.read_deltalake(path=path)
+
+    def empty() -> Iterator[pd.DataFrame]:
+        if False:
+            yield pd.DataFrame()  # pragma: no cover
+
+    wr.s3.to_deltalake_streaming(
+        dfs=empty(),
+        path=path,
+        mode="overwrite",
+        partition_cols=["c0", "c1"],
+        **lock_settings,
+    )
+    after = wr.s3.read_deltalake(path=path)
+    assert after.equals(baseline)
+
+
+def test_df_iter_to_record_batch_reader_schema_and_rows() -> None:
+    df_empty = pd.DataFrame({"a": [], "b": []})
+    df1 = pd.DataFrame({"a": [1, 2], "b": ["x", "y"]})
+    df2 = pd.DataFrame({"a": [3], "b": ["z"]})
+
+    reader, schema = _df_iter_to_record_batch_reader(
+        df_iter=[df_empty, df1, df2],
+        index=False,
+        dtype={},
+        target_schema=None,
+        batch_size=None,
+    )
+
+    assert isinstance(schema, pa.Schema)
+    assert {f.name for f in schema} == {"a", "b"}
+
+    table: pa.Table = reader.read_all()
+    pdf = table.to_pandas()
+    assert len(pdf) == 3
+    assert sorted(pdf["a"].tolist()) == [1, 2, 3]
+    assert set(pdf["b"].tolist()) == {"x", "y", "z"}
+
+
+def test_df_iter_to_record_batch_reader_respects_batch_size() -> None:
+    df1 = pd.DataFrame({"a": list(range(5)), "b": ["x"] * 5})
+    df2 = pd.DataFrame({"a": list(range(5, 9)), "b": ["y"] * 4})
+
+    reader, _ = _df_iter_to_record_batch_reader(
+        df_iter=[df1, df2],
+        index=False,
+        dtype={},
+        target_schema=None,
+        batch_size=3,
+    )
+
+    batch_count = 0
+    row_count = 0
+    for batch in reader:
+        batch_count += 1
+        row_count += batch.num_rows
+
+    assert batch_count >= 3
+    assert row_count == 9