Improve table creation flow.

Author: igorborgest

awswrangler/catalog/__init__.py

@@ -2,6 +2,8 @@
 
 from awswrangler.catalog._add import add_csv_partitions, add_parquet_partitions  # noqa
 from awswrangler.catalog._create import (  # noqa
+    _create_csv_table,
+    _create_parquet_table,
     create_csv_table,
     create_database,
     create_parquet_table,
@@ -10,6 +12,7 @@
 )
 from awswrangler.catalog._delete import delete_database, delete_table_if_exists  # noqa
 from awswrangler.catalog._get import (  # noqa
+    _get_table_input,
     databases,
     get_columns_comments,
     get_connection,
@@ -20,8 +23,10 @@
     get_partitions,
     get_table_description,
     get_table_location,
+    get_table_number_of_versions,
     get_table_parameters,
     get_table_types,
+    get_table_versions,
     get_tables,
     search_tables,
     table,

awswrangler/catalog/_create.py

@@ -12,7 +12,7 @@
 
 from awswrangler import _utils, exceptions
 from awswrangler._config import apply_configs
-from awswrangler.catalog._utils import _extract_dtypes_from_table_details
+from awswrangler.catalog._utils import _catalog_id, _extract_dtypes_from_table_details
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
@@ -21,16 +21,12 @@ def _get_table_input(
     database: str, table: str, boto3_session: Optional[boto3.Session], catalog_id: Optional[str] = None
 ) -> Optional[Dict[str, Any]]:
     client_glue: boto3.client = _utils.client(service_name="glue", session=boto3_session)
-    args: Dict[str, str] = {}
-    if catalog_id is not None:
-        args["CatalogId"] = catalog_id
-    args["DatabaseName"] = database
-    args["Name"] = table
     try:
-        response: Dict[str, Any] = client_glue.get_table(**args)
+        response: Dict[str, Any] = client_glue.get_table(
+            **_catalog_id(catalog_id=catalog_id, DatabaseName=database, Name=table)
+        )
     except client_glue.exceptions.EntityNotFoundException:
         return None
-
     table_input: Dict[str, Any] = {}
     for k, v in response["Table"].items():
         if k in [
@@ -49,7 +45,6 @@ def _get_table_input(
             "TargetTable",
         ]:
             table_input[k] = v
-
     return table_input
 
 
@@ -162,10 +157,7 @@ def get_databases(
     """
     client_glue: boto3.client = _utils.client(service_name="glue", session=boto3_session)
     paginator = client_glue.get_paginator("get_databases")
-    if catalog_id is None:
-        response_iterator: Iterator = paginator.paginate()
-    else:
-        response_iterator = paginator.paginate(CatalogId=catalog_id)
+    response_iterator = paginator.paginate(**_catalog_id(catalog_id=catalog_id))
     for page in response_iterator:
         for db in page["DatabaseList"]:
             yield db
@@ -436,10 +428,7 @@ def table(
 
     """
     client_glue: boto3.client = _utils.client(service_name="glue", session=boto3_session)
-    if catalog_id is None:
-        tbl: Dict[str, Any] = client_glue.get_table(DatabaseName=database, Name=table)["Table"]
-    else:
-        tbl = client_glue.get_table(CatalogId=catalog_id, DatabaseName=database, Name=table)["Table"]
+    tbl = client_glue.get_table(**_catalog_id(catalog_id=catalog_id, DatabaseName=database, Name=table))["Table"]
     df_dict: Dict[str, List] = {"Column Name": [], "Type": [], "Partition": [], "Comment": []}
     for col in tbl["StorageDescriptor"]["Columns"]:
         df_dict["Column Name"].append(col["Name"])
@@ -522,10 +511,7 @@ def get_connection(
 
     """
     client_glue: boto3.client = _utils.client(service_name="glue", session=boto3_session)
-    args: Dict[str, Any] = {"Name": name, "HidePassword": False}
-    if catalog_id is not None:
-        args["CatalogId"] = catalog_id
-    return client_glue.get_connection(**args)["Connection"]
+    return client_glue.get_connection(**_catalog_id(catalog_id=catalog_id, Name=name, HidePassword=False))["Connection"]
 
 
 def get_engine(
@@ -812,12 +798,9 @@ def get_table_parameters(
 
     """
     client_glue: boto3.client = _utils.client(service_name="glue", session=boto3_session)
-    args: Dict[str, str] = {}
-    if catalog_id is not None:
-        args["CatalogId"] = catalog_id
-    args["DatabaseName"] = database
-    args["Name"] = table
-    response: Dict[str, Any] = client_glue.get_table(**args)
+    response: Dict[str, Any] = client_glue.get_table(
+        **_catalog_id(catalog_id=catalog_id, DatabaseName=database, Name=table)
+    )
     parameters: Dict[str, str] = response["Table"]["Parameters"]
     return parameters
 
@@ -851,16 +834,14 @@ def get_table_description(
 
     """
     client_glue: boto3.client = _utils.client(service_name="glue", session=boto3_session)
-    args: Dict[str, str] = {}
-    if catalog_id is not None:
-        args["CatalogId"] = catalog_id
-    args["DatabaseName"] = database
-    args["Name"] = table
-    response: Dict[str, Any] = client_glue.get_table(**args)
+    response: Dict[str, Any] = client_glue.get_table(
+        **_catalog_id(catalog_id=catalog_id, DatabaseName=database, Name=table)
+    )
     desc: Optional[str] = response["Table"].get("Description", None)
     return desc
 
 
+@apply_configs
 def get_columns_comments(
     database: str, table: str, catalog_id: Optional[str] = None, boto3_session: Optional[boto3.Session] = None
 ) -> Dict[str, str]:
@@ -890,16 +871,91 @@ def get_columns_comments(
 
     """
     client_glue: boto3.client = _utils.client(service_name="glue", session=boto3_session)
-    args: Dict[str, str] = {}
-    if catalog_id is not None:
-        args["CatalogId"] = catalog_id
-    args["DatabaseName"] = database
-    args["Name"] = table
-    response: Dict[str, Any] = client_glue.get_table(**args)
+    response: Dict[str, Any] = client_glue.get_table(
+        **_catalog_id(catalog_id=catalog_id, DatabaseName=database, Name=table)
+    )
     comments: Dict[str, str] = {}
     for c in response["Table"]["StorageDescriptor"]["Columns"]:
         comments[c["Name"]] = c["Comment"]
     if "PartitionKeys" in response["Table"]:
         for p in response["Table"]["PartitionKeys"]:
             comments[p["Name"]] = p["Comment"]
     return comments
+
+
+@apply_configs
+def get_table_versions(
+    database: str, table: str, catalog_id: Optional[str] = None, boto3_session: Optional[boto3.Session] = None
+) -> List[Dict[str, Any]]:
+    """Get all versions.
+
+    Parameters
+    ----------
+    database : str
+        Database name.
+    table : str
+        Table name.
+    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.
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 session will be used if boto3_session receive None.
+
+    Returns
+    -------
+    List[Dict[str, Any]
+        List of table inputs:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/glue.html#Glue.Client.get_table_versions
+
+    Examples
+    --------
+    >>> import awswrangler as wr
+    >>> tables_versions = wr.catalog.get_table_versions(database="...", table="...")
+
+    """
+    client_glue: boto3.client = _utils.client(service_name="glue", session=boto3_session)
+    paginator = client_glue.get_paginator("get_table_versions")
+    versions: List[Dict[str, Any]] = []
+    response_iterator = paginator.paginate(**_catalog_id(DatabaseName=database, TableName=table, catalog_id=catalog_id))
+    for page in response_iterator:
+        for tbl in page["TableVersions"]:
+            versions.append(tbl)
+    return versions
+
+
+@apply_configs
+def get_table_number_of_versions(
+    database: str, table: str, catalog_id: Optional[str] = None, boto3_session: Optional[boto3.Session] = None
+) -> int:
+    """Get tatal number of versions.
+
+    Parameters
+    ----------
+    database : str
+        Database name.
+    table : str
+        Table name.
+    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.
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 session will be used if boto3_session receive None.
+
+    Returns
+    -------
+    int
+        Total number of versions.
+
+    Examples
+    --------
+    >>> import awswrangler as wr
+    >>> num = wr.catalog.get_table_number_of_versions(database="...", table="...")
+
+    """
+    client_glue: boto3.client = _utils.client(service_name="glue", session=boto3_session)
+    paginator = client_glue.get_paginator("get_table_versions")
+    count: int = 0
+    response_iterator = paginator.paginate(**_catalog_id(DatabaseName=database, TableName=table, catalog_id=catalog_id))
+    for page in response_iterator:
+        count += len(page["TableVersions"])
+    return count

awswrangler/catalog/_get.py

@@ -1,9 +1,8 @@
 """Amazon CSV S3 Write Module (PRIVATE)."""
 
 import logging
-from typing import Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
-import boto3  # type: ignore
 import pandas as pd  # type: ignore
 
 from awswrangler import _data_types, _utils, catalog, exceptions
@@ -13,21 +12,25 @@
 _COMPRESSION_2_EXT: Dict[Optional[str], str] = {None: "", "gzip": ".gz", "snappy": ".snappy"}
 
 
+def _extract_dtypes_from_table_input(table_input: Dict[str, Any]) -> Dict[str, str]:
+    dtypes: Dict[str, str] = {}
+    for col in table_input["StorageDescriptor"]["Columns"]:
+        dtypes[col["Name"]] = col["Type"]
+    if "PartitionKeys" in table_input:
+        for par in table_input["PartitionKeys"]:
+            dtypes[par["Name"]] = par["Type"]
+    return dtypes
+
+
 def _apply_dtype(
-    df: pd.DataFrame,
-    mode: str,
-    database: Optional[str],
-    table: Optional[str],
-    dtype: Dict[str, str],
-    boto3_session: boto3.Session,
+    df: pd.DataFrame, dtype: Dict[str, str], catalog_table_input: Optional[Dict[str, Any]], mode: str
 ) -> pd.DataFrame:
-    if (mode in ("append", "overwrite_partitions")) and (database is not None) and (table is not None):
-        catalog_types: Optional[Dict[str, str]] = catalog.get_table_types(
-            database=database, table=table, boto3_session=boto3_session
-        )
-        if catalog_types is not None:
-            for k, v in catalog_types.items():
-                dtype[k] = v
+    if mode in ("append", "overwrite_partitions"):
+        if catalog_table_input is not None:
+            catalog_types: Optional[Dict[str, str]] = _extract_dtypes_from_table_input(table_input=catalog_table_input)
+            if catalog_types is not None:
+                for k, v in catalog_types.items():
+                    dtype[k] = v
     df = _data_types.cast_pandas_with_athena_types(df=df, dtype=dtype)
     return df
 

awswrangler/s3/_write.py

@@ -37,7 +37,7 @@ def write(self, func: Callable, boto3_session: boto3.Session, **func_kwargs) ->
         if self._exec is not None:
             _logger.debug("Submitting: %s", func)
             future = self._exec.submit(
-                fn=_WriteProxy._caller,
+                _WriteProxy._caller,
                 func=func,
                 boto3_primitives=_utils.boto3_to_primitives(boto3_session=boto3_session),
                 func_kwargs=func_kwargs,

awswrangler/s3/_write_concurrent.py

@@ -2,7 +2,7 @@
 
 import logging
 import uuid
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import boto3  # type: ignore
 import pandas as pd  # type: ignore
@@ -95,6 +95,7 @@ def to_parquet(  # pylint: disable=too-many-arguments,too-many-locals
     projection_values: Optional[Dict[str, str]] = None,
     projection_intervals: Optional[Dict[str, str]] = None,
     projection_digits: Optional[Dict[str, str]] = None,
+    catalog_id: Optional[str] = None,
 ) -> Dict[str, Union[List[str], Dict[str, List[str]]]]:
     """Write Parquet file or dataset on Amazon S3.
 
@@ -196,6 +197,9 @@ def to_parquet(  # pylint: disable=too-many-arguments,too-many-locals
         Dictionary of partitions names and Athena projections digits.
         https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html
         (e.g. {'col_name': '1', 'col2_name': '2'})
+    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.
 
     Returns
     -------
@@ -333,7 +337,12 @@ def to_parquet(  # pylint: disable=too-many-arguments,too-many-locals
         df, dtype, partition_cols = _sanitize(df=df, dtype=dtype, partition_cols=partition_cols)
 
     # Evaluating dtype
-    df = _apply_dtype(df=df, mode=mode, database=database, table=table, dtype=dtype, boto3_session=session)
+    catalog_table_input: Optional[Dict[str, Any]] = None
+    if database is not None and table is not None:
+        catalog_table_input: Optional[Dict[str, Any]] = catalog._get_table_input(  # pylint: disable=protected-access
+            database=database, table=table, boto3_session=session, catalog_id=catalog_id
+        )
+    df = _apply_dtype(df=df, dtype=dtype, catalog_table_input=catalog_table_input, mode=mode)
     schema: pa.Schema = _data_types.pyarrow_schema_from_pandas(
         df=df, index=index, ignore_cols=partition_cols, dtype=dtype
     )
@@ -376,7 +385,7 @@ def to_parquet(  # pylint: disable=too-many-arguments,too-many-locals
             columns_types, partitions_types = _data_types.athena_types_from_pandas_partitioned(
                 df=df, index=index, partition_cols=partition_cols, dtype=dtype
             )
-            catalog.create_parquet_table(
+            catalog._create_parquet_table(  # pylint: disable=protected-access
                 database=database,
                 table=table,
                 path=path,
@@ -395,6 +404,8 @@ def to_parquet(  # pylint: disable=too-many-arguments,too-many-locals
                 projection_values=projection_values,
                 projection_intervals=projection_intervals,
                 projection_digits=projection_digits,
+                catalog_id=catalog_id,
+                catalog_table_input=catalog_table_input,
             )
             if partitions_values and (regular_partitions is True):
                 _logger.debug("partitions_values:\n%s", partitions_values)