feat: Add ignore null store parquet metadata (#2450)

* feat: add ignore_null flag to store_parquet_metadata
fix: address line-too-long pylint issues

* test: ignore_null flag for store_parquet_metadata

* refactor: fix formatting in test_athena_parquet.py

* test: use the new ignore_null flag in the tests for it

* test: think we need to provide dtype for to_parquet method if some columns are null, else reading schemas doesn't work properly later

* test: shouldn't attempt to partition on an entirely null column

* test: make sure dataset=True for test compatibility

Author: raaidarshad

awswrangler/s3/_write_parquet.py

@@ -432,9 +432,10 @@ def to_parquet(  # pylint: disable=too-many-arguments,too-many-locals,too-many-b
         Useful when you have columns with undetermined or mixed data types.
         (e.g. {'col name': 'bigint', 'col2 name': 'int'})
     athena_partition_projection_settings: typing.AthenaPartitionProjectionSettings, optional
-        Parameters of the Athena Partition Projection (https://docs.aws.amazon.com/athena/latest/ug/partition-projection.html).
-        AthenaPartitionProjectionSettings is a `TypedDict`, meaning the passed parameter can be instantiated either as an
-        instance of AthenaPartitionProjectionSettings or as a regular Python dict.
+        Parameters of the Athena Partition Projection
+        (https://docs.aws.amazon.com/athena/latest/ug/partition-projection.html).
+        AthenaPartitionProjectionSettings is a `TypedDict`, meaning the passed parameter can be instantiated either as
+        an instance of AthenaPartitionProjectionSettings or as a regular Python dict.
 
         Following projection parameters are supported:
 
@@ -758,6 +759,7 @@ def store_parquet_metadata(  # pylint: disable=too-many-arguments,too-many-local
     path_suffix: Optional[str] = None,
     path_ignore_suffix: Union[str, List[str], None] = None,
     ignore_empty: bool = True,
+    ignore_null: bool = False,
     dtype: Optional[Dict[str, str]] = None,
     sampling: float = 1.0,
     dataset: bool = False,
@@ -810,6 +812,8 @@ def store_parquet_metadata(  # pylint: disable=too-many-arguments,too-many-local
         Suffix or List of suffixes for S3 keys to be ignored.
     ignore_empty: bool
         Ignore files with 0 bytes.
+    ignore_null: bool
+        Ignore columns with null type.
     dtype : Dict[str, str], optional
         Dictionary of columns names and Athena/Glue types to be casted.
         Useful when you have columns with undetermined data types as partitions columns.
@@ -844,9 +848,10 @@ def store_parquet_metadata(  # pylint: disable=too-many-arguments,too-many-local
         Keep enabled even when working with projections is useful to keep
         Redshift Spectrum working with the regular partitions.
     athena_partition_projection_settings: typing.AthenaPartitionProjectionSettings, optional
-        Parameters of the Athena Partition Projection (https://docs.aws.amazon.com/athena/latest/ug/partition-projection.html).
-        AthenaPartitionProjectionSettings is a `TypedDict`, meaning the passed parameter can be instantiated either as an
-        instance of AthenaPartitionProjectionSettings or as a regular Python dict.
+        Parameters of the Athena Partition Projection
+        (https://docs.aws.amazon.com/athena/latest/ug/partition-projection.html).
+        AthenaPartitionProjectionSettings is a `TypedDict`, meaning the passed parameter can be instantiated either as
+        an instance of AthenaPartitionProjectionSettings or as a regular Python dict.
 
         Following projection parameters are supported:
 
@@ -934,7 +939,7 @@ def store_parquet_metadata(  # pylint: disable=too-many-arguments,too-many-local
         path_suffix=path_suffix,
         path_ignore_suffix=path_ignore_suffix,
         ignore_empty=ignore_empty,
-        ignore_null=False,
+        ignore_null=ignore_null,
         use_threads=use_threads,
         s3_additional_kwargs=s3_additional_kwargs,
         boto3_session=boto3_session,

tests/unit/test_athena_parquet.py

@@ -373,6 +373,38 @@ def test_store_metadata_partitions_sample_dataset(glue_database, glue_table, pat
     assert df.c2.sum() * num_files == df2.c2.sum()
 
 
+def test_store_metadata_ignore_null_columns(glue_database, glue_table, path):
+    df = pd.DataFrame({"c0": [0, 1, 2], "c1": [3, 4, 5], "c2_null": [None, None, None], "c3_null": [None, None, None]})
+    wr.s3.to_parquet(df=df, path=path, dataset=True, dtype={"c2_null": "int", "c3_null": "int"})
+    wr.s3.store_parquet_metadata(
+        path=path,
+        database=glue_database,
+        table=glue_table,
+        ignore_null=True,
+        dataset=True,
+        dtype={"c2_null": "int", "c3_null": "int"},
+    )
+
+
+@pytest.mark.parametrize("partition_cols", [None, ["c0"], ["c0", "c1"]])
+def test_store_metadata_ignore_null_columns_partitions(glue_database, glue_table, path, partition_cols):
+    # only partition on non-null columns
+    num_files = 10
+    df = pd.DataFrame({"c0": [0, 1, 2], "c1": [3, 4, 5], "c2_null": [None, None, None], "c3_null": [None, None, None]})
+    for _ in range(num_files):
+        wr.s3.to_parquet(
+            df=df, path=path, dataset=True, dtype={"c2_null": "int", "c3_null": "int"}, partition_cols=partition_cols
+        )
+    wr.s3.store_parquet_metadata(
+        path=path,
+        database=glue_database,
+        table=glue_table,
+        ignore_null=True,
+        dtype={"c2_null": "int", "c3_null": "int"},
+        dataset=True,
+    )
+
+
 @pytest.mark.parametrize("partition_cols", [None, ["c1"], ["c2"], ["c1", "c2"], ["c2", "c1"]])
 def test_to_parquet_reverse_partitions(glue_database, glue_table, path, partition_cols):
     df = pd.DataFrame({"c0": [0, 1, 2], "c1": [3, 4, 5], "c2": [6, 7, 8]})