Refactoring the parquet crawler

Author: igorborgest

awswrangler/_data_types.py

@@ -381,20 +381,6 @@ def athena_types_from_pyarrow_schema(
     return columns_types, partitions_types
 
 
-def athena_partitions_from_pyarrow_partitions(
-    path: str, partitions: pyarrow.parquet.ParquetPartitions
-) -> Dict[str, List[str]]:
-    """Extract the related Athena partitions values from any PyArrow Partitions."""
-    path = path if path[-1] == "/" else f"{path}/"
-    partitions_values: Dict[str, List[str]] = {}
-    names: List[str] = [p.name for p in partitions]
-    for values in zip(*[p.keys for p in partitions]):
-        suffix: str = "/".join([f"{n}={v}" for n, v in zip(names, values)])
-        suffix = suffix if suffix[-1] == "/" else f"{suffix}/"
-        partitions_values[f"{path}{suffix}"] = list(values)
-    return partitions_values
-
-
 def cast_pandas_with_athena_types(df: pd.DataFrame, dtype: Dict[str, str]) -> pd.DataFrame:
     """Cast columns in a Pandas DataFrame."""
     for col, athena_type in dtype.items():
@@ -412,6 +398,12 @@ def cast_pandas_with_athena_types(df: pd.DataFrame, dtype: Dict[str, str]) -> pd
                     .astype("string")
                     .apply(lambda x: Decimal(str(x)) if str(x) not in ("", "none", " ", "<NA>") else None)
                 )
+            elif pandas_type == "string":
+                curr_type: str = str(df[col].dtypes)
+                if curr_type.startswith("int") or curr_type.startswith("float"):
+                    df[col] = df[col].astype(str).astype("string")
+                else:
+                    df[col] = df[col].astype("string")
             else:
                 df[col] = df[col].astype(pandas_type)
     return df

awswrangler/_utils.py

@@ -3,6 +3,7 @@
 import logging
 import math
 import os
+import random
 from typing import Any, Dict, Generator, List, Optional, Tuple
 
 import boto3  # type: ignore
@@ -11,7 +12,9 @@
 import psycopg2  # type: ignore
 import s3fs  # type: ignore
 
-logger: logging.Logger = logging.getLogger(__name__)
+from awswrangler import exceptions
+
+_logger: logging.Logger = logging.getLogger(__name__)
 
 
 def ensure_session(session: Optional[boto3.Session] = None) -> boto3.Session:
@@ -124,6 +127,8 @@ def chunkify(lst: List[Any], num_chunks: int = 1, max_length: Optional[int] = No
     [[0, 1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]
 
     """
+    if not lst:
+        return []
     n: int = num_chunks if max_length is None else int(math.ceil((float(len(lst)) / float(max_length))))
     np_chunks = np.array_split(lst, n)
     return [arr.tolist() for arr in np_chunks if len(arr) > 0]
@@ -179,3 +184,52 @@ def get_region_from_subnet(subnet_id: str, boto3_session: Optional[boto3.Session
     session: boto3.Session = ensure_session(session=boto3_session)
     client_ec2: boto3.client = client(service_name="ec2", session=session)
     return client_ec2.describe_subnets(SubnetIds=[subnet_id])["Subnets"][0]["AvailabilityZone"][:9]
+
+
+def extract_partitions_from_paths(
+    path: str, paths: List[str]
+) -> Tuple[Optional[Dict[str, str]], Optional[Dict[str, List[str]]]]:
+    """Extract partitions from Amazon S3 paths."""
+    path = path if path.endswith("/") else f"{path}/"
+    partitions_types: Dict[str, str] = {}
+    partitions_values: Dict[str, List[str]] = {}
+    for p in paths:
+        if path not in p:
+            raise exceptions.InvalidArgumentValue(f"Object {p} is not under the root path ({path}).")
+        path_wo_filename: str = p.rpartition("/")[0] + "/"
+        if path_wo_filename not in partitions_values:
+            path_wo_prefix: str = p.replace(f"{path}/", "")
+            dirs: List[str] = [x for x in path_wo_prefix.split("/") if (x != "") and ("=" in x)]
+            if dirs:
+                values_tups: List[Tuple[str, str]] = [tuple(x.split("=")[:2]) for x in dirs]  # type: ignore
+                values_dics: Dict[str, str] = dict(values_tups)
+                p_values: List[str] = list(values_dics.values())
+                p_types: Dict[str, str] = {x: "string" for x in values_dics.keys()}
+                if not partitions_types:
+                    partitions_types = p_types
+                if p_values:
+                    partitions_types = p_types
+                    partitions_values[path_wo_filename] = p_values
+                elif p_types != partitions_types:
+                    raise exceptions.InvalidSchemaConvergence(
+                        f"At least two different partitions schema detected: {partitions_types} and {p_types}"
+                    )
+    if not partitions_types:
+        return None, None
+    return partitions_types, partitions_values
+
+
+def list_sampling(lst: List[Any], sampling: float) -> List[Any]:
+    """Random List sampling."""
+    if sampling > 1.0 or sampling <= 0.0:
+        raise exceptions.InvalidArgumentValue(f"Argument <sampling> must be [0.0 < value <= 1.0]. {sampling} received.")
+    _len: int = len(lst)
+    if _len == 0:
+        return []
+    num_samples: int = int(round(_len * sampling))
+    num_samples = _len if num_samples > _len else num_samples
+    num_samples = 1 if num_samples < 1 else num_samples
+    _logger.debug("_len: %s", _len)
+    _logger.debug("sampling: %s", sampling)
+    _logger.debug("num_samples: %s", num_samples)
+    return random.sample(population=lst, k=num_samples)

awswrangler/athena.py

@@ -282,6 +282,8 @@ def repair_table(
 
     """
     query = f"MSCK REPAIR TABLE `{table}`;"
+    if (database is not None) and (not database.startswith("`")):
+        database = f"`{database}`"
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     query_id = start_query_execution(
         sql=query,
@@ -492,7 +494,7 @@ def read_sql_query(  # pylint: disable=too-many-branches,too-many-locals,too-man
         path: str = f"{_s3_output}/{name}"
         ext_location: str = "\n" if wg_config["enforced"] is True else f",\n    external_location = '{path}'\n"
         sql = (
-            f"CREATE TABLE {name}\n"
+            f'CREATE TABLE "{name}"\n'
             f"WITH(\n"
             f"    format = 'Parquet',\n"
             f"    parquet_compression = 'SNAPPY'"

awswrangler/catalog.py

@@ -240,11 +240,12 @@ def add_parquet_partitions(
     ... )
 
     """
-    inputs: List[Dict[str, Any]] = [
-        _parquet_partition_definition(location=k, values=v, compression=compression)
-        for k, v in partitions_values.items()
-    ]
-    _add_partitions(database=database, table=table, boto3_session=boto3_session, inputs=inputs)
+    if partitions_values:
+        inputs: List[Dict[str, Any]] = [
+            _parquet_partition_definition(location=k, values=v, compression=compression)
+            for k, v in partitions_values.items()
+        ]
+        _add_partitions(database=database, table=table, boto3_session=boto3_session, inputs=inputs)
 
 
 def _parquet_partition_definition(location: str, values: List[str], compression: Optional[str]) -> Dict[str, Any]:

awswrangler/exceptions.py

@@ -83,3 +83,7 @@ class InvalidRedshiftSortkey(Exception):
 
 class InvalidRedshiftPrimaryKeys(Exception):
     """InvalidRedshiftPrimaryKeys exception."""
+
+
+class InvalidSchemaConvergence(Exception):
+    """InvalidSchemaMerge exception."""