fix: Schema evolution for to_csv and to_json (#2104)

Author: LeonLuttenberger

awswrangler/catalog/_create.py

@@ -252,6 +252,29 @@ def _overwrite_table_parameters(
     return parameters
 
 
+def _update_table_input(table_input: Dict[str, Any], columns_types: Dict[str, str], allow_reorder: bool = True) -> bool:
+    column_updated = False
+
+    catalog_cols: Dict[str, str] = {x["Name"]: x["Type"] for x in table_input["StorageDescriptor"]["Columns"]}
+
+    if not allow_reorder:
+        for catalog_key, frame_key in zip(catalog_cols, columns_types):
+            if catalog_key != frame_key:
+                raise exceptions.InvalidArgumentValue(f"Column {frame_key} is out of order.")
+
+    for c, t in columns_types.items():
+        if c not in catalog_cols:
+            _logger.debug("New column %s with type %s.", c, t)
+            table_input["StorageDescriptor"]["Columns"].append({"Name": c, "Type": t})
+            column_updated = True
+        elif t != catalog_cols[c]:  # Data type change detected!
+            raise exceptions.InvalidArgumentValue(
+                f"Data type change detected on column {c} (Old type: {catalog_cols[c]} / New type {t})."
+            )
+
+    return column_updated
+
+
 def _create_parquet_table(
     database: str,
     table: str,
@@ -282,19 +305,14 @@ def _create_parquet_table(
     table = sanitize_table_name(table=table)
     partitions_types = {} if partitions_types is None else partitions_types
     _logger.debug("catalog_table_input: %s", catalog_table_input)
+
     table_input: Dict[str, Any]
     if (catalog_table_input is not None) and (mode in ("append", "overwrite_partitions")):
         table_input = catalog_table_input
-        catalog_cols: Dict[str, str] = {x["Name"]: x["Type"] for x in table_input["StorageDescriptor"]["Columns"]}
-        for c, t in columns_types.items():
-            if c not in catalog_cols:
-                _logger.debug("New column %s with type %s.", c, t)
-                table_input["StorageDescriptor"]["Columns"].append({"Name": c, "Type": t})
-                mode = "update"
-            elif t != catalog_cols[c]:  # Data type change detected!
-                raise exceptions.InvalidArgumentValue(
-                    f"Data type change detected on column {c} (Old type: {catalog_cols[c]} / New type {t})."
-                )
+
+        is_table_updated = _update_table_input(table_input, columns_types)
+        if is_table_updated:
+            mode = "update"
     else:
         table_input = _parquet_table_definition(
             table=table,
@@ -368,11 +386,18 @@ def _create_csv_table(  # pylint: disable=too-many-arguments,too-many-locals
     table = sanitize_table_name(table=table)
     partitions_types = {} if partitions_types is None else partitions_types
     _logger.debug("catalog_table_input: %s", catalog_table_input)
-    table_input: Dict[str, Any]
+
     if schema_evolution is False:
         _utils.check_schema_changes(columns_types=columns_types, table_input=catalog_table_input, mode=mode)
+
+    table_input: Dict[str, Any]
     if (catalog_table_input is not None) and (mode in ("append", "overwrite_partitions")):
         table_input = catalog_table_input
+
+        is_table_updated = _update_table_input(table_input, columns_types, allow_reorder=False)
+        if is_table_updated:
+            mode = "update"
+
     else:
         table_input = _csv_table_definition(
             table=table,
@@ -415,7 +440,7 @@ def _create_csv_table(  # pylint: disable=too-many-arguments,too-many-locals
     )
 
 
-def _create_json_table(  # pylint: disable=too-many-arguments
+def _create_json_table(  # pylint: disable=too-many-arguments,too-many-locals
     database: str,
     table: str,
     path: str,
@@ -453,6 +478,11 @@ def _create_json_table(  # pylint: disable=too-many-arguments
         _utils.check_schema_changes(columns_types=columns_types, table_input=catalog_table_input, mode=mode)
     if (catalog_table_input is not None) and (mode in ("append", "overwrite_partitions")):
         table_input = catalog_table_input
+
+        is_table_updated = _update_table_input(table_input, columns_types)
+        if is_table_updated:
+            mode = "update"
+
     else:
         table_input = _json_table_definition(
             table=table,

tests/test_s3_parquet.py

@@ -1,7 +1,7 @@
 import itertools
 import logging
 import math
-from datetime import datetime, timedelta, timezone
+from datetime import date, datetime, timedelta, timezone
 
 import boto3
 import numpy as np
@@ -571,6 +571,65 @@ def test_read_parquet_versioned(path) -> None:
         assert version_id == wr.s3.describe_objects(path=path_file, version_id=version_id)[path_file]["VersionId"]
 
 
+def test_parquet_schema_evolution(path, glue_database, glue_table):
+    df = pd.DataFrame(
+        {
+            "id": [1, 2],
+            "value": ["foo", "boo"],
+        }
+    )
+    wr.s3.to_parquet(
+        df=df,
+        path=path,
+        dataset=True,
+        mode="overwrite",
+        database=glue_database,
+        table=glue_table,
+    )
+
+    df2 = pd.DataFrame(
+        {"id": [3, 4], "value": ["bar", None], "date": [date(2020, 1, 3), date(2020, 1, 4)], "flag": [True, False]}
+    )
+    wr.s3.to_parquet(
+        df=df2,
+        path=path,
+        dataset=True,
+        mode="append",
+        database=glue_database,
+        table=glue_table,
+        schema_evolution=True,
+        catalog_versioning=True,
+    )
+
+    column_types = wr.catalog.get_table_types(glue_database, glue_table)
+    assert len(column_types) == len(df2.columns)
+
+
+def test_to_parquet_schema_evolution_out_of_order(path, glue_database, glue_table) -> None:
+    df = pd.DataFrame({"c0": [0, 1, 2], "c1": ["a", "b", "c"]})
+    wr.s3.to_parquet(df=df, path=path, dataset=True, database=glue_database, table=glue_table)
+
+    df2 = df.copy()
+    df2["c2"] = ["x", "y", "z"]
+
+    wr.s3.to_parquet(
+        df=df2,
+        path=path,
+        dataset=True,
+        database=glue_database,
+        table=glue_table,
+        mode="append",
+        schema_evolution=True,
+        catalog_versioning=True,
+    )
+
+    df_out = wr.s3.read_parquet(path=path, dataset=True)
+    df_expected = pd.concat([df, df2], ignore_index=True)
+
+    assert len(df_out) == len(df_expected)
+    assert list(df_out.columns) == list(df_expected.columns)
+
+
 def test_read_parquet_schema_validation_with_index_column(path) -> None:
     path_file = f"{path}file.parquet"
     df = pd.DataFrame({"idx": [1], "col": [2]})

tests/test_s3_text.py

@@ -359,27 +359,86 @@ def test_read_csv_versioned(path) -> None:
         assert version_id == wr.s3.describe_objects(path=path_file, version_id=version_id)[path_file]["VersionId"]
 
 
-def test_to_csv_schema_evolution(path, glue_database, glue_table) -> None:
-    path_file = f"{path}0.csv"
+@pytest.mark.parametrize("mode", ["append", "overwrite"])
+def test_to_csv_schema_evolution(path, glue_database, glue_table, mode) -> None:
     df = pd.DataFrame({"c0": [0, 1, 2], "c1": [3, 4, 5]})
-    wr.s3.to_csv(df=df, path=path_file, dataset=True, database=glue_database, table=glue_table)
+    wr.s3.to_csv(df=df, path=path, dataset=True, database=glue_database, table=glue_table, index=False)
+
     df["c2"] = [6, 7, 8]
     wr.s3.to_csv(
         df=df,
-        path=path_file,
+        path=path,
         dataset=True,
         database=glue_database,
         table=glue_table,
-        mode="overwrite",
+        mode=mode,
         schema_evolution=True,
+        index=False,
     )
+
+    column_types = wr.catalog.get_table_types(glue_database, glue_table)
+    assert len(column_types) == len(df.columns)
+
     df["c3"] = [9, 10, 11]
+    with pytest.raises(wr.exceptions.InvalidArgumentValue):
+        wr.s3.to_csv(df=df, path=path, dataset=True, database=glue_database, table=glue_table, schema_evolution=False)
+
+
+@pytest.mark.parametrize("schema_evolution", [False, True])
+def test_to_csv_schema_evolution_out_of_order(path, glue_database, glue_table, schema_evolution) -> None:
+    df = pd.DataFrame({"c0": [0, 1, 2], "c1": [3, 4, 5]})
+    wr.s3.to_csv(df=df, path=path, dataset=True, database=glue_database, table=glue_table, index=False)
+
+    df["c2"] = [6, 7, 8]
+    df = df[["c0", "c2", "c1"]]
+
     with pytest.raises(wr.exceptions.InvalidArgumentValue):
         wr.s3.to_csv(
-            df=df, path=path_file, dataset=True, database=glue_database, table=glue_table, schema_evolution=False
+            df=df,
+            path=path,
+            dataset=True,
+            database=glue_database,
+            table=glue_table,
+            mode="append",
+            schema_evolution=schema_evolution,
+            index=False,
         )
 
 
+@pytest.mark.parametrize("mode", ["append", "overwrite"])
+def test_to_json_schema_evolution(path, glue_database, glue_table, mode) -> None:
+    df = pd.DataFrame({"c0": [0, 1, 2], "c1": [3, 4, 5]})
+    wr.s3.to_json(
+        df=df,
+        path=path,
+        dataset=True,
+        database=glue_database,
+        table=glue_table,
+        orient="split",
+        index=False,
+    )
+
+    df["c2"] = [6, 7, 8]
+    wr.s3.to_json(
+        df=df,
+        path=path,
+        dataset=True,
+        database=glue_database,
+        table=glue_table,
+        mode=mode,
+        schema_evolution=True,
+        orient="split",
+        index=False,
+    )
+
+    column_types = wr.catalog.get_table_types(glue_database, glue_table)
+    assert len(column_types) == len(df.columns)
+
+    df["c3"] = [9, 10, 11]
+    with pytest.raises(wr.exceptions.InvalidArgumentValue):
+        wr.s3.to_json(df=df, path=path, dataset=True, database=glue_database, table=glue_table, schema_evolution=False)
+
+
 def test_exceptions(path):
     with pytest.raises(wr.exceptions.EmptyDataFrame):
         wr.s3.to_json(df=pd.DataFrame(), path=path)