Add test_store_parquet_metadata_modes()

Author: igorborgest

testing/test_awswrangler/test_data_lake.py

@@ -1924,3 +1924,207 @@ def test_to_parquet_modes(database, table, path, external_schema):
     assert len(df3.columns) == 3
     assert len(df3.index) == 4
     assert df3.c1.sum() == 6
+
+
+def test_store_parquet_metadata_modes(database, table, path, external_schema):
+
+    # Round 1 - Warm up
+    df = pd.DataFrame({"c0": [0, None]}, dtype="Int64")
+    paths = wr.s3.to_parquet(df=df, path=path, dataset=True, mode="overwrite")["paths"]
+    wr.s3.wait_objects_exist(paths=paths)
+    wr.s3.store_parquet_metadata(
+        path=path,
+        dataset=True,
+        mode="overwrite",
+        database=database,
+        table=table,
+        description="c0",
+        parameters={"num_cols": str(len(df.columns)), "num_rows": str(len(df.index))},
+        columns_comments={"c0": "0"},
+    )
+    df2 = wr.athena.read_sql_table(table, database)
+    assert df.shape == df2.shape
+    assert df.c0.sum() == df2.c0.sum()
+    parameters = wr.catalog.get_table_parameters(database, table)
+    assert len(parameters) == 5
+    assert parameters["num_cols"] == str(len(df2.columns))
+    assert parameters["num_rows"] == str(len(df2.index))
+    assert wr.catalog.get_table_description(database, table) == "c0"
+    comments = wr.catalog.get_columns_comments(database, table)
+    assert len(comments) == len(df.columns)
+    assert comments["c0"] == "0"
+
+    # Round 2 - Overwrite
+    df = pd.DataFrame({"c1": [None, 1, None]}, dtype="Int16")
+    paths = wr.s3.to_parquet(df=df, path=path, dataset=True, mode="overwrite")["paths"]
+    wr.s3.wait_objects_exist(paths=paths)
+    wr.s3.store_parquet_metadata(
+        path=path,
+        dataset=True,
+        mode="overwrite",
+        database=database,
+        table=table,
+        description="c1",
+        parameters={"num_cols": str(len(df.columns)), "num_rows": str(len(df.index))},
+        columns_comments={"c1": "1"},
+    )
+    df2 = wr.athena.read_sql_table(table, database)
+    assert df.shape == df2.shape
+    assert df.c1.sum() == df2.c1.sum()
+    parameters = wr.catalog.get_table_parameters(database, table)
+    assert len(parameters) == 5
+    assert parameters["num_cols"] == str(len(df2.columns))
+    assert parameters["num_rows"] == str(len(df2.index))
+    assert wr.catalog.get_table_description(database, table) == "c1"
+    comments = wr.catalog.get_columns_comments(database, table)
+    assert len(comments) == len(df.columns)
+    assert comments["c1"] == "1"
+
+    # Round 3 - Append
+    df = pd.DataFrame({"c1": [None, 2, None]}, dtype="Int16")
+    paths = wr.s3.to_parquet(df=df, path=path, dataset=True, mode="append")["paths"]
+    wr.s3.wait_objects_exist(paths=paths)
+    wr.s3.store_parquet_metadata(
+        path=path,
+        dataset=True,
+        mode="append",
+        database=database,
+        table=table,
+        description="c1",
+        parameters={"num_cols": str(len(df.columns)), "num_rows": str(len(df.index) * 2)},
+        columns_comments={"c1": "1"},
+    )
+    df2 = wr.athena.read_sql_table(table, database)
+    assert len(df.columns) == len(df2.columns)
+    assert len(df.index) * 2 == len(df2.index)
+    assert df.c1.sum() + 1 == df2.c1.sum()
+    parameters = wr.catalog.get_table_parameters(database, table)
+    assert len(parameters) == 5
+    assert parameters["num_cols"] == str(len(df2.columns))
+    assert parameters["num_rows"] == str(len(df2.index))
+    assert wr.catalog.get_table_description(database, table) == "c1"
+    comments = wr.catalog.get_columns_comments(database, table)
+    assert len(comments) == len(df.columns)
+    assert comments["c1"] == "1"
+
+    # Round 4 - Append + New Column
+    df = pd.DataFrame({"c2": ["a", None, "b"], "c1": [None, 1, None]})
+    df["c1"] = df["c1"].astype("Int16")
+    paths = wr.s3.to_parquet(df=df, path=path, dataset=True, mode="append")["paths"]
+    wr.s3.wait_objects_exist(paths=paths)
+    wr.s3.store_parquet_metadata(
+        path=path,
+        dataset=True,
+        mode="append",
+        database=database,
+        table=table,
+        description="c1+c2",
+        parameters={"num_cols": "2", "num_rows": "9"},
+        columns_comments={"c1": "1", "c2": "2"},
+    )
+    df2 = wr.athena.read_sql_table(table, database)
+    assert len(df2.columns) == 2
+    assert len(df2.index) == 9
+    assert df2.c1.sum() == 4
+    parameters = wr.catalog.get_table_parameters(database, table)
+    assert len(parameters) == 5
+    assert parameters["num_cols"] == "2"
+    assert parameters["num_rows"] == "9"
+    assert wr.catalog.get_table_description(database, table) == "c1+c2"
+    comments = wr.catalog.get_columns_comments(database, table)
+    assert len(comments) == len(df.columns)
+    assert comments["c1"] == "1"
+    assert comments["c2"] == "2"
+
+    # Round 5 - Overwrite Partitioned
+    df = pd.DataFrame({"c0": ["foo", None], "c1": [0, 1]})
+    paths = wr.s3.to_parquet(df=df, path=path, dataset=True, mode="overwrite", partition_cols=["c1"])["paths"]
+    wr.s3.wait_objects_exist(paths=paths)
+    wr.s3.store_parquet_metadata(
+        path=path,
+        dataset=True,
+        mode="overwrite",
+        database=database,
+        table=table,
+        description="c0+c1",
+        parameters={"num_cols": "2", "num_rows": "2"},
+        columns_comments={"c0": "zero", "c1": "one"},
+    )
+    df2 = wr.athena.read_sql_table(table, database)
+    assert df.shape == df2.shape
+    assert df.c1.sum() == df2.c1.astype(int).sum()
+    parameters = wr.catalog.get_table_parameters(database, table)
+    assert len(parameters) == 5
+    assert parameters["num_cols"] == "2"
+    assert parameters["num_rows"] == "2"
+    assert wr.catalog.get_table_description(database, table) == "c0+c1"
+    comments = wr.catalog.get_columns_comments(database, table)
+    assert len(comments) == len(df.columns)
+    assert comments["c0"] == "zero"
+    assert comments["c1"] == "one"
+
+    # Round 6 - Overwrite Partitions
+    df = pd.DataFrame({"c0": [None, "boo"], "c1": [0, 2]})
+    paths = wr.s3.to_parquet(df=df, path=path, dataset=True, mode="overwrite_partitions", partition_cols=["c1"])[
+        "paths"
+    ]
+    wr.s3.wait_objects_exist(paths=paths)
+    wr.s3.store_parquet_metadata(
+        path=path,
+        dataset=True,
+        mode="append",
+        database=database,
+        table=table,
+        description="c0+c1",
+        parameters={"num_cols": "2", "num_rows": "3"},
+        columns_comments={"c0": "zero", "c1": "one"},
+    )
+    df2 = wr.athena.read_sql_table(table, database)
+    assert len(df2.columns) == 2
+    assert len(df2.index) == 3
+    assert df2.c1.astype(int).sum() == 3
+    parameters = wr.catalog.get_table_parameters(database, table)
+    assert len(parameters) == 5
+    assert parameters["num_cols"] == "2"
+    assert parameters["num_rows"] == "3"
+    assert wr.catalog.get_table_description(database, table) == "c0+c1"
+    comments = wr.catalog.get_columns_comments(database, table)
+    assert len(comments) == len(df.columns)
+    assert comments["c0"] == "zero"
+    assert comments["c1"] == "one"
+
+    # Round 7 - Overwrite Partitions + New Column
+    df = pd.DataFrame({"c0": ["bar", None], "c1": [1, 3], "c2": [True, False]})
+    paths = wr.s3.to_parquet(df=df, path=path, dataset=True, mode="overwrite_partitions", partition_cols=["c1"])[
+        "paths"
+    ]
+    wr.s3.wait_objects_exist(paths=paths)
+    wr.s3.store_parquet_metadata(
+        path=path,
+        dataset=True,
+        mode="append",
+        database=database,
+        table=table,
+        description="c0+c1+c2",
+        parameters={"num_cols": "3", "num_rows": "4"},
+        columns_comments={"c0": "zero", "c1": "one", "c2": "two"},
+    )
+    df2 = wr.athena.read_sql_table(table, database)
+    assert len(df2.columns) == 3
+    assert len(df2.index) == 4
+    assert df2.c1.astype(int).sum() == 6
+    parameters = wr.catalog.get_table_parameters(database, table)
+    assert len(parameters) == 5
+    assert parameters["num_cols"] == "3"
+    assert parameters["num_rows"] == "4"
+    assert wr.catalog.get_table_description(database, table) == "c0+c1+c2"
+    comments = wr.catalog.get_columns_comments(database, table)
+    assert len(comments) == len(df.columns)
+    assert comments["c0"] == "zero"
+    assert comments["c1"] == "one"
+    assert comments["c2"] == "two"
+    engine = wr.catalog.get_engine("aws-data-wrangler-redshift")
+    df3 = wr.db.read_sql_table(con=engine, table=table, schema=external_schema)
+    assert len(df3.columns) == 3
+    assert len(df3.index) == 4
+    assert df3.c1.astype(int).sum() == 6