Add suffix filters for read_parquet_table() #495

Author: igorborgest

awswrangler/s3/_read_parquet.py

@@ -614,6 +614,8 @@ def read_parquet(
 def read_parquet_table(
     table: str,
     database: str,
+    filename_suffix: Union[str, List[str], None] = None,
+    filename_ignore_suffix: Union[str, List[str], None] = None,
     catalog_id: Optional[str] = None,
     partition_filter: Optional[Callable[[Dict[str, str]], bool]] = None,
     columns: Optional[List[str]] = None,
@@ -655,6 +657,12 @@ def read_parquet_table(
         AWS Glue Catalog table name.
     database : str
         AWS Glue Catalog database name.
+    path_suffix: Union[str, List[str], None]
+        Suffix or List of suffixes to be read (e.g. [".gz.parquet", ".snappy.parquet"]).
+        If None, will try to read all files. (default)
+    path_ignore_suffix: Union[str, List[str], None]
+        Suffix or List of suffixes for S3 keys to be ignored.(e.g. [".csv", "_SUCCESS"]).
+        If None, will try to read all files. (default)
     catalog_id : str, optional
         The ID of the Data Catalog from which to retrieve Databases.
         If none is provided, the AWS account ID is used by default.
@@ -741,6 +749,8 @@ def read_parquet_table(
     return _data_types.cast_pandas_with_athena_types(
         df=read_parquet(
             path=path,
+            path_suffix=filename_suffix,
+            path_ignore_suffix=filename_ignore_suffix,
             partition_filter=partition_filter,
             columns=columns,
             validate_schema=validate_schema,

tests/test_athena_parquet.py

@@ -707,3 +707,32 @@ def test_to_parquet_nested_structs(glue_database, glue_table, path):
     wr.s3.to_parquet(df=df, path=path, dataset=True, database=glue_database, table=glue_table)
     df3 = wr.athena.read_sql_query(sql=f"SELECT * FROM {glue_table}", database=glue_database)
     assert df3.shape == (2, 2)
+
+
+def test_ignore_empty_files(glue_database, glue_table, path):
+    df = pd.DataFrame({"c0": [0, 1], "c1": ["foo", "boo"]})
+    bucket, directory = wr._utils.parse_path(path)
+    wr.s3.to_parquet(df=df, path=path, dataset=True, database=glue_database, table=glue_table)
+    boto3.client("s3").put_object(Body=b"", Bucket=bucket, Key=f"{directory}to_be_ignored")
+    df2 = wr.athena.read_sql_query(sql=f"SELECT * FROM {glue_table}", database=glue_database)
+    assert df2.shape == df.shape
+    df3 = wr.s3.read_parquet_table(database=glue_database, table=glue_table)
+    assert df3.shape == df.shape
+
+
+def test_suffix(glue_database, glue_table, path):
+    df = pd.DataFrame({"c0": [0, 1], "c1": ["foo", "boo"]})
+    bucket, directory = wr._utils.parse_path(path)
+    wr.s3.to_parquet(df=df, path=path, dataset=True, database=glue_database, table=glue_table)
+    boto3.client("s3").put_object(Body=b"garbage", Bucket=bucket, Key=f"{directory}to_be_ignored")
+    df2 = wr.s3.read_parquet_table(database=glue_database, table=glue_table, filename_suffix=".parquet")
+    assert df2.shape == df.shape
+
+
+def test_ignore_suffix(glue_database, glue_table, path):
+    df = pd.DataFrame({"c0": [0, 1], "c1": ["foo", "boo"]})
+    bucket, directory = wr._utils.parse_path(path)
+    wr.s3.to_parquet(df=df, path=path, dataset=True, database=glue_database, table=glue_table)
+    boto3.client("s3").put_object(Body=b"garbage", Bucket=bucket, Key=f"{directory}to_be_ignored")
+    df2 = wr.s3.read_parquet_table(database=glue_database, table=glue_table, filename_ignore_suffix="ignored")
+    assert df2.shape == df.shape