Add pandas.to_redshift upsert

Author: igorborgest

awswrangler/__init__.py

@@ -23,13 +23,7 @@ def __init__(self, service):
         self._service = service
 
     def __getattr__(self, name):
-        return getattr(
-            getattr(
-                DynamicInstantiate.__default_session,
-                self._service
-            ),
-            name
-        )
+        return getattr(getattr(DynamicInstantiate.__default_session, self._service), name)
 
 
 if importlib.util.find_spec("pyspark"):  # type: ignore

awswrangler/exceptions.py

@@ -58,6 +58,10 @@ class InvalidRedshiftSortkey(Exception):
     pass
 
 
+class InvalidRedshiftPrimaryKeys(Exception):
+    pass
+
+
 class EmptyDataframe(Exception):
     pass
 

awswrangler/pandas.py

@@ -1096,6 +1096,7 @@ def to_redshift(
             distkey: Optional[str] = None,
             sortstyle: str = "COMPOUND",
             sortkey: Optional[str] = None,
+            primary_keys: Optional[str] = None,
             preserve_index: bool = False,
             mode: str = "append",
             cast_columns: Optional[Dict[str, str]] = None,
@@ -1113,6 +1114,7 @@ def to_redshift(
         :param distkey: Specifies a column name or positional number for the distribution key
         :param sortstyle: Sorting can be "COMPOUND" or "INTERLEAVED" (https://docs.aws.amazon.com/redshift/latest/dg/t_Sorting_data.html)
         :param sortkey: List of columns to be sorted
+        :param primary_keys: Primary keys
         :param preserve_index: Should we preserve the Dataframe index?
         :param mode: append, overwrite or upsert
         :param cast_columns: Dictionary of columns names and Redshift types to be casted. (E.g. {"col name": "SMALLINT", "col2 name": "FLOAT4"})
@@ -1159,6 +1161,7 @@ def to_redshift(
             distkey=distkey,
             sortstyle=sortstyle,
             sortkey=sortkey,
+            primary_keys=primary_keys,
             mode=mode,
             cast_columns=cast_columns,
         )
@@ -1344,14 +1347,23 @@ def _read_parquet_path(session_primitives: Any,
         :param filters: List of filters to apply, like ``[[('x', '=', 0), ...], ...]``.
         :param procs_cpu_bound: Number of cores used for CPU bound tasks
         """
-        path = path[:-1] if path[-1] == "/" else path
+        session = session_primitives.session
+        is_file: bool = session.s3.does_object_exists(path=path)
+        if is_file is False:
+            path = path[:-1] if path[-1] == "/" else path
         procs_cpu_bound = procs_cpu_bound if procs_cpu_bound is not None else session_primitives.procs_cpu_bound if session_primitives.procs_cpu_bound is not None else 1
         use_threads: bool = True if procs_cpu_bound > 1 else False
-        fs: S3FileSystem = s3.get_fs(session_primitives=session_primitives)
-        fs.invalidate_cache()
-        fs = pa.filesystem._ensure_filesystem(fs)
-        logger.debug(f"Reading Parquet table: {path}")
-        table = pq.read_table(source=path, columns=columns, filters=filters, filesystem=fs, use_threads=use_threads)
+        logger.debug(f"Reading Parquet: {path}")
+        if is_file is True:
+            client_s3 = session.boto3_session.client(service_name="s3", use_ssl=True, config=session.botocore_config)
+            bucket, key = path.replace("s3://", "").split("/", 1)
+            obj = client_s3.get_object(Bucket=bucket, Key=key)
+            table = pq.ParquetFile(source=BytesIO(obj["Body"].read())).read(columns=columns, use_threads=use_threads)
+        else:
+            fs: S3FileSystem = s3.get_fs(session_primitives=session_primitives)
+            fs = pa.filesystem._ensure_filesystem(fs)
+            fs.invalidate_cache()
+            table = pq.read_table(source=path, columns=columns, filters=filters, filesystem=fs, use_threads=use_threads)
         # Check if we lose some integer during the conversion (Happens when has some null value)
         integers = [field.name for field in table.schema if str(field.type).startswith("int")]
         logger.debug(f"Converting to Pandas: {path}")

awswrangler/redshift.py

@@ -3,16 +3,12 @@
 import logging
 
 import pg8000  # type: ignore
+import pyarrow as pa  # type: ignore
 
 from awswrangler import data_types
-from awswrangler.exceptions import (
-    RedshiftLoadError,
-    InvalidDataframeType,
-    InvalidRedshiftDiststyle,
-    InvalidRedshiftDistkey,
-    InvalidRedshiftSortstyle,
-    InvalidRedshiftSortkey,
-)
+from awswrangler.exceptions import (RedshiftLoadError, InvalidDataframeType, InvalidRedshiftDiststyle,
+                                    InvalidRedshiftDistkey, InvalidRedshiftSortstyle, InvalidRedshiftSortkey,
+                                    InvalidRedshiftPrimaryKeys)
 
 logger = logging.getLogger(__name__)
 
@@ -165,6 +161,7 @@ def load_table(dataframe,
                    distkey=None,
                    sortstyle="COMPOUND",
                    sortkey=None,
+                   primary_keys: Optional[List[str]] = None,
                    mode="append",
                    preserve_index=False,
                    cast_columns=None):
@@ -184,11 +181,14 @@ def load_table(dataframe,
         :param distkey: Specifies a column name or positional number for the distribution key
         :param sortstyle: Sorting can be "COMPOUND" or "INTERLEAVED" (https://docs.aws.amazon.com/redshift/latest/dg/t_Sorting_data.html)
         :param sortkey: List of columns to be sorted
-        :param mode: append or overwrite
+        :param primary_keys: Primary keys
+        :param mode: append, overwrite or upsert
         :param preserve_index: Should we preserve the Dataframe index? (ONLY for Pandas Dataframe)
         :param cast_columns: Dictionary of columns names and Redshift types to be casted. (E.g. {"col name": "INT", "col2 name": "FLOAT"})
         :return: None
         """
+        final_table_name: Optional[str] = None
+        temp_table_name: Optional[str] = None
         cursor = redshift_conn.cursor()
         if mode == "overwrite":
             Redshift._create_table(cursor=cursor,
@@ -200,13 +200,27 @@ def load_table(dataframe,
                                    distkey=distkey,
                                    sortstyle=sortstyle,
                                    sortkey=sortkey,
+                                   primary_keys=primary_keys,
                                    preserve_index=preserve_index,
                                    cast_columns=cast_columns)
+            table_name = f"{schema_name}.{table_name}"
+        elif mode == "upsert":
+            guid: str = pa.compat.guid()
+            temp_table_name = f"temp_redshift_{guid}"
+            final_table_name = table_name
+            table_name = temp_table_name
+            sql: str = f"CREATE TEMPORARY TABLE {temp_table_name} (LIKE {schema_name}.{final_table_name})"
+            logger.debug(sql)
+            cursor.execute(sql)
+        else:
+            table_name = f"{schema_name}.{table_name}"
+
         sql = ("-- AWS DATA WRANGLER\n"
-               f"COPY {schema_name}.{table_name} FROM '{manifest_path}'\n"
+               f"COPY {table_name} FROM '{manifest_path}'\n"
                f"IAM_ROLE '{iam_role}'\n"
                "MANIFEST\n"
                "FORMAT AS PARQUET")
+        logger.debug(sql)
         cursor.execute(sql)
         cursor.execute("-- AWS DATA WRANGLER\n SELECT pg_last_copy_id() AS query_id")
         query_id = cursor.fetchall()[0][0]
@@ -219,6 +233,23 @@ def load_table(dataframe,
             cursor.close()
             raise RedshiftLoadError(
                 f"Redshift load rollbacked. {num_files_loaded} files counted. {num_files} expected.")
+
+        if (mode == "upsert") and (final_table_name is not None):
+            if not primary_keys:
+                primary_keys = Redshift.get_primary_keys(connection=redshift_conn,
+                                                         schema=schema_name,
+                                                         table=final_table_name)
+            if not primary_keys:
+                raise InvalidRedshiftPrimaryKeys()
+            equals_clause = f"{final_table_name}.%s = {temp_table_name}.%s"
+            join_clause = " AND ".join([equals_clause % (pk, pk) for pk in primary_keys])
+            sql = f"DELETE FROM {schema_name}.{final_table_name} USING {temp_table_name} WHERE {join_clause}"
+            logger.debug(sql)
+            cursor.execute(sql)
+            sql = f"INSERT INTO {schema_name}.{final_table_name} SELECT * FROM {temp_table_name}"
+            logger.debug(sql)
+            cursor.execute(sql)
+
         redshift_conn.commit()
         cursor.close()
 
@@ -232,6 +263,7 @@ def _create_table(cursor,
                       distkey=None,
                       sortstyle="COMPOUND",
                       sortkey=None,
+                      primary_keys: List[str] = None,
                       preserve_index=False,
                       cast_columns=None):
         """
@@ -246,6 +278,7 @@ def _create_table(cursor,
         :param distkey: Specifies a column name or positional number for the distribution key
         :param sortstyle: Sorting can be "COMPOUND" or "INTERLEAVED" (https://docs.aws.amazon.com/redshift/latest/dg/t_Sorting_data.html)
         :param sortkey: List of columns to be sorted
+        :param primary_keys: Primary keys
         :param preserve_index: Should we preserve the Dataframe index? (ONLY for Pandas Dataframe)
         :param cast_columns: Dictionary of columns names and Redshift types to be casted. (E.g. {"col name": "INT", "col2 name": "FLOAT"})
         :return: None
@@ -273,22 +306,43 @@ def _create_table(cursor,
                                       distkey=distkey,
                                       sortstyle=sortstyle,
                                       sortkey=sortkey)
-        cols_str = "".join([f"{col[0]} {col[1]},\n" for col in schema])[:-2]
-        distkey_str = ""
+        cols_str: str = "".join([f"{col[0]} {col[1]},\n" for col in schema])[:-2]
+        primary_keys_str: str = ""
+        if primary_keys:
+            primary_keys_str = f",\nPRIMARY KEY ({', '.join(primary_keys)})"
+        distkey_str: str = ""
         if distkey and diststyle == "KEY":
             distkey_str = f"\nDISTKEY({distkey})"
-        sortkey_str = ""
+        sortkey_str: str = ""
         if sortkey:
             sortkey_str = f"\n{sortstyle} SORTKEY({','.join(sortkey)})"
         sql = (f"-- AWS DATA WRANGLER\n"
                f"CREATE TABLE IF NOT EXISTS {schema_name}.{table_name} (\n"
                f"{cols_str}"
+               f"{primary_keys_str}"
                f")\nDISTSTYLE {diststyle}"
                f"{distkey_str}"
                f"{sortkey_str}")
         logger.debug(f"Create table query:\n{sql}")
         cursor.execute(sql)
 
+    @staticmethod
+    def get_primary_keys(connection, schema, table):
+        """
+        Get PKs
+        :param connection: A PEP 249 compatible connection (Can be generated with Redshift.generate_connection())
+        :param schema: Schema name
+        :param table: Redshift table name
+        :return: PKs list List[str]
+        """
+        cursor = connection.cursor()
+        cursor.execute(f"SELECT indexdef FROM pg_indexes WHERE schemaname = '{schema}' AND tablename = '{table}'")
+        result = cursor.fetchall()[0][0]
+        rfields = result.split('(')[1].strip(')').split(',')
+        fields = [field.strip().strip('"') for field in rfields]
+        cursor.close()
+        return fields
+
     @staticmethod
     def _validate_parameters(schema, diststyle, distkey, sortstyle, sortkey):
         """
@@ -347,8 +401,8 @@ def _get_redshift_schema(dataframe, dataframe_type, preserve_index=False, cast_c
             raise InvalidDataframeType(dataframe_type)
         return schema_built
 
-    @staticmethod
-    def to_parquet(sql: str,
+    def to_parquet(self,
+                   sql: str,
                    path: str,
                    iam_role: str,
                    connection: Any,
@@ -366,8 +420,11 @@ def to_parquet(sql: str,
         path = path if path[-1] == "/" else path + "/"
         cursor: Any = connection.cursor()
         partition_str: str = ""
+        manifest_str: str = ""
         if partition_cols is not None:
             partition_str = f"PARTITION BY ({','.join([x for x in partition_cols])})\n"
+        else:
+            manifest_str = "\nmanifest"
         query: str = f"-- AWS DATA WRANGLER\n" \
                      f"UNLOAD ('{sql}')\n" \
                      f"TO '{path}'\n" \
@@ -376,7 +433,8 @@ def to_parquet(sql: str,
                      f"PARALLEL ON\n" \
                      f"ENCRYPTED \n" \
                      f"{partition_str}" \
-                     f"FORMAT PARQUET;"
+                     f"FORMAT PARQUET" \
+                     f"{manifest_str};"
         logger.debug(f"query:\n{query}")
         cursor.execute(query)
         query = "-- AWS DATA WRANGLER\nSELECT pg_last_query_id() AS query_id"
@@ -391,4 +449,8 @@ def to_parquet(sql: str,
         logger.debug(f"paths: {paths}")
         connection.commit()
         cursor.close()
+        if manifest_str != "":
+            self._session.s3.wait_object_exists(path=f"{path}manifest")
+        for p in paths:
+            self._session.s3.wait_object_exists(path=p)
         return paths

awswrangler/s3.py

@@ -1,7 +1,8 @@
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 import multiprocessing as mp
 from math import ceil
 import logging
+from time import sleep
 
 from botocore.exceptions import ClientError, HTTPClientError  # type: ignore
 import s3fs  # type: ignore
@@ -21,7 +22,7 @@ def mkdir_if_not_exists(fs, path):
 
 
 def get_fs(session_primitives=None):
-    aws_access_key_id, aws_secret_access_key, profile_name, config, s3_additional_kwargs = None, None, None, None, None
+    aws_access_key_id, aws_secret_access_key, profile_name = None, None, None
     args = {}
 
     if session_primitives is not None:
@@ -42,17 +43,49 @@ def get_fs(session_primitives=None):
         args["key"] = aws_access_key_id,
         args["secret"] = aws_secret_access_key
 
+    args["default_cache_type"] = "none"
+    args["default_fill_cache"] = False
     fs = s3fs.S3FileSystem(**args)
-    fs.invalidate_cache(path=None)
     return fs
 
 
 class S3:
     def __init__(self, session):
         self._session = session
+        self._client_s3 = session.boto3_session.client(service_name="s3", use_ssl=True, config=session.botocore_config)
+
+    def does_object_exists(self, path: str) -> bool:
+        """
+        Check if object exists on S3
+
+        :param path: S3 path (e.g. s3://...)
+        :return: boolean
+        """
+        bucket: str
+        key: str
+        bucket, key = path.replace("s3://", "").split("/", 1)
+        try:
+            self._client_s3.head_object(Bucket=bucket, Key=key)
+            return True
+        except ClientError as ex:
+            if ex.response["ResponseMetadata"]["HTTPStatusCode"] == 404:
+                return False
+            raise ex
+
+    def wait_object_exists(self, path: str, polling_sleep: float = 0.1) -> None:
+        """
+        Wait object exists on S3
+
+        :param path: S3 path (e.g. s3://...)
+        :param polling_sleep: Milliseconds
+        :return: None
+        """
+        while self.does_object_exists(path=path) is False:
+            sleep(polling_sleep)
 
     @staticmethod
-    def parse_path(path):
+    def parse_path(path: str) -> Tuple[str, str]:
+        bucket: str
         bucket, path = path.replace("s3://", "").split("/", 1)
         if not path:
             path = ""

awswrangler/sagemaker.py

@@ -23,15 +23,15 @@ def get_job_outputs(self, path: str) -> Any:
         if key.split("/")[-1] != "model.tar.gz":
             key = f"{key}/model.tar.gz"
         body = self._client_s3.get_object(Bucket=bucket, Key=key)["Body"].read()
-        body = tarfile.io.BytesIO(body)
+        body = tarfile.io.BytesIO(body)  # type: ignore
         tar = tarfile.open(fileobj=body)
 
         results = []
         for member in tar.getmembers():
             f = tar.extractfile(member)
             file_type = member.name.split(".")[-1]
 
-            if file_type == "pkl":
+            if (file_type == "pkl") and (f is not None):
                 f = pickle.load(f)
 
             results.append(f)

awswrangler/session.py

@@ -15,7 +15,6 @@
 from awswrangler.emr import EMR
 from awswrangler.sagemaker import SageMaker
 
-
 PYSPARK_INSTALLED = False
 if importlib.util.find_spec("pyspark"):  # type: ignore
     PYSPARK_INSTALLED = True

docs/source/api/awswrangler.rst

@@ -15,6 +15,7 @@ Submodules
    awswrangler.pandas
    awswrangler.redshift
    awswrangler.s3
+   awswrangler.sagemaker
    awswrangler.session
    awswrangler.spark
    awswrangler.utils

docs/source/api/awswrangler.sagemaker.rst

@@ -0,0 +1,7 @@
+awswrangler.sagemaker module
+============================
+
+.. automodule:: awswrangler.sagemaker
+   :members:
+   :undoc-members:
+   :show-inheritance: