Add distribution and sorting configuration in Redshift integration.

Author: igorborgest

awswrangler/exceptions.py

@@ -36,3 +36,19 @@ class LineTerminatorNotFound(Exception):
 
 class MissingBatchDetected(Exception):
     pass
+
+
+class InvalidRedshiftDiststyle(Exception):
+    pass
+
+
+class InvalidRedshiftDistkey(Exception):
+    pass
+
+
+class InvalidRedshiftSortstyle(Exception):
+    pass
+
+
+class InvalidRedshiftSortkey(Exception):
+    pass

awswrangler/pandas.py

@@ -760,6 +760,10 @@ def to_redshift(
             schema,
             table,
             iam_role,
+            diststyle="AUTO",
+            distkey=None,
+            sortstyle="COMPOUND",
+            sortkey=None,
             preserve_index=False,
             mode="append",
     ):
@@ -771,6 +775,12 @@ def to_redshift(
         :param schema: The Redshift Schema for the table
         :param table: The name of the desired Redshift table
         :param iam_role: AWS IAM role with the related permissions
+        :param diststyle: Redshift distribution styles. Must be in ["AUTO", "EVEN", "ALL", "KEY"]
+               https://docs.aws.amazon.com/redshift/latest/dg/t_Distributing_data.html
+        :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 preserve_index: Should we preserve the Dataframe index?
         :param mode: append or overwrite
         :return: None
@@ -779,7 +789,7 @@ def to_redshift(
             path += "/"
         self._session.s3.delete_objects(path=path)
         num_rows = len(dataframe.index)
-        logger.info(f"Number of rows: {num_rows}")
+        logger.debug(f"Number of rows: {num_rows}")
         if num_rows < MIN_NUMBER_OF_ROWS_TO_DISTRIBUTE:
             num_partitions = 1
         else:
@@ -808,6 +818,10 @@ def to_redshift(
             preserve_index=False,
             num_files=num_partitions,
             iam_role=iam_role,
+            diststyle=diststyle,
+            distkey=distkey,
+            sortstyle=sortstyle,
+            sortkey=sortkey,
             mode=mode,
         )
         self._session.s3.delete_objects(path=path)

awswrangler/redshift.py

@@ -7,10 +7,26 @@
     RedshiftLoadError,
     UnsupportedType,
     InvalidDataframeType,
+    InvalidRedshiftDiststyle,
+    InvalidRedshiftDistkey,
+    InvalidRedshiftSortstyle,
+    InvalidRedshiftSortkey,
 )
 
 logger = logging.getLogger(__name__)
 
+DISTSTYLES = [
+    "AUTO",
+    "EVEN",
+    "ALL",
+    "KEY",
+]
+
+SORTSTYLES = [
+    "COMPOUND",
+    "INTERLEAVED",
+]
+
 
 class Redshift:
     def __init__(self, session):
@@ -87,24 +103,48 @@ def load_table(
             redshift_conn,
             num_files,
             iam_role,
+            diststyle="AUTO",
+            distkey=None,
+            sortstyle="COMPOUND",
+            sortkey=None,
             mode="append",
             preserve_index=False,
     ):
+        """
+        Load Parquet files into a Redshift table using a manifest file.
+        Creates the table if necessary.
+        :param dataframe: Pandas or Spark Dataframe
+        :param dataframe_type: "pandas" or "spark"
+        :param manifest_path: S3 path for manifest file (E.g. S3://...)
+        :param schema_name: Redshift schema
+        :param table_name: Redshift table name
+        :param redshift_conn: A PEP 249 compatible connection (Can be generated with Redshift.generate_connection())
+        :param num_files: Number of files to be loaded
+        :param iam_role: AWS IAM role with the related permissions
+        :param diststyle: Redshift distribution styles. Must be in ["AUTO", "EVEN", "ALL", "KEY"]
+               https://docs.aws.amazon.com/redshift/latest/dg/t_Distributing_data.html
+        :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 preserve_index: Should we preserve the Dataframe index? (ONLY for Pandas Dataframe)
+        :return: None
+        """
         cursor = redshift_conn.cursor()
         if mode == "overwrite":
-            cursor.execute("-- AWS DATA WRANGLER\n"
-                           f"DROP TABLE IF EXISTS {schema_name}.{table_name}")
-            schema = Redshift._get_redshift_schema(
+            Redshift._create_table(
+                cursor=cursor,
                 dataframe=dataframe,
                 dataframe_type=dataframe_type,
+                schema_name=schema_name,
+                table_name=table_name,
+                diststyle=diststyle,
+                distkey=distkey,
+                sortstyle=sortstyle,
+                sortkey=sortkey,
                 preserve_index=preserve_index,
             )
-            cols_str = "".join([f"{col[0]} {col[1]},\n" for col in schema])[:-2]
-            sql = (
-                "-- AWS DATA WRANGLER\n"
-                f"CREATE TABLE IF NOT EXISTS {schema_name}.{table_name} (\n{cols_str}"
-                ") DISTSTYLE AUTO")
-            cursor.execute(sql)
         sql = ("-- AWS DATA WRANGLER\n"
                f"COPY {schema_name}.{table_name} FROM '{manifest_path}'\n"
                f"IAM_ROLE '{iam_role}'\n"
@@ -129,6 +169,111 @@ def load_table(
         redshift_conn.commit()
         cursor.close()
 
+    @staticmethod
+    def _create_table(
+            cursor,
+            dataframe,
+            dataframe_type,
+            schema_name,
+            table_name,
+            diststyle="AUTO",
+            distkey=None,
+            sortstyle="COMPOUND",
+            sortkey=None,
+            preserve_index=False,
+    ):
+        """
+        Creates Redshift table.
+        :param cursor: A PEP 249 compatible cursor
+        :param dataframe: Pandas or Spark Dataframe
+        :param dataframe_type: "pandas" or "spark"
+        :param schema_name: Redshift schema
+        :param table_name: Redshift table name
+        :param diststyle: Redshift distribution styles. Must be in ["AUTO", "EVEN", "ALL", "KEY"]
+               https://docs.aws.amazon.com/redshift/latest/dg/t_Distributing_data.html
+        :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 preserve_index: Should we preserve the Dataframe index? (ONLY for Pandas Dataframe)
+        :return: None
+        """
+        sql = f"-- AWS DATA WRANGLER\n" \
+              f"DROP TABLE IF EXISTS {schema_name}.{table_name}"
+        logger.debug(f"Drop table query:\n{sql}")
+        cursor.execute(sql)
+        schema = Redshift._get_redshift_schema(
+            dataframe=dataframe,
+            dataframe_type=dataframe_type,
+            preserve_index=preserve_index,
+        )
+        if diststyle:
+            diststyle = diststyle.upper()
+        else:
+            diststyle = "AUTO"
+        if sortstyle:
+            sortstyle = sortstyle.upper()
+        else:
+            sortstyle = "COMPOUND"
+        Redshift._validate_parameters(schema=schema,
+                                      diststyle=diststyle,
+                                      distkey=distkey,
+                                      sortstyle=sortstyle,
+                                      sortkey=sortkey)
+        cols_str = "".join([f"{col[0]} {col[1]},\n" for col in schema])[:-2]
+        distkey_str = ""
+        if distkey and diststyle == "KEY":
+            distkey_str = f"\nDISTKEY({distkey})"
+        sortkey_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")\nDISTSTYLE {diststyle}"
+               f"{distkey_str}"
+               f"{sortkey_str}")
+        logger.debug(f"Create table query:\n{sql}")
+        cursor.execute(sql)
+
+    @staticmethod
+    def _validate_parameters(schema, diststyle, distkey, sortstyle, sortkey):
+        """
+        Validates the sanity of Redshift's parameters
+        :param schema: List of tuples (column name, column type)
+        :param diststyle: Redshift distribution styles. Must be in ["AUTO", "EVEN", "ALL", "KEY"]
+               https://docs.aws.amazon.com/redshift/latest/dg/t_Distributing_data.html
+        :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
+        :return: None
+        """
+        if diststyle not in DISTSTYLES:
+            raise InvalidRedshiftDiststyle(
+                f"diststyle must be in {DISTSTYLES}")
+        cols = [x[0] for x in schema]
+        logger.debug(f"Redshift columns: {cols}")
+        if (diststyle == "KEY") and (not distkey):
+            raise InvalidRedshiftDistkey(
+                "You must pass a distkey if you intend to use KEY diststyle")
+        if distkey and distkey not in cols:
+            raise InvalidRedshiftDistkey(
+                f"distkey ({distkey}) must be in the columns list: {cols})")
+        if sortstyle and sortstyle not in SORTSTYLES:
+            raise InvalidRedshiftSortstyle(
+                f"sortstyle must be in {SORTSTYLES}")
+        if sortkey:
+            if type(sortkey) != list:
+                raise InvalidRedshiftSortkey(
+                    f"sortkey must be a List of items in the columns list: {cols}. "
+                    f"Currently value: {sortkey}")
+            for key in sortkey:
+                if key not in cols:
+                    raise InvalidRedshiftSortkey(
+                        f"sortkey must be a List of items in the columns list: {cols}. "
+                        f"Currently value: {key}")
+
     @staticmethod
     def _get_redshift_schema(dataframe, dataframe_type, preserve_index=False):
         schema_built = []

awswrangler/spark.py

@@ -28,9 +28,31 @@ def to_redshift(
             schema,
             table,
             iam_role,
+            diststyle="AUTO",
+            distkey=None,
+            sortstyle="COMPOUND",
+            sortkey=None,
             min_num_partitions=200,
             mode="append",
     ):
+        """
+        Load Spark Dataframe as a Table on Amazon Redshift
+        :param dataframe: Pandas Dataframe
+        :param path: S3 path to write temporary files (E.g. s3://BUCKET_NAME/ANY_NAME/)
+        :param connection: A PEP 249 compatible connection (Can be generated with Redshift.generate_connection())
+        :param schema: The Redshift Schema for the table
+        :param table: The name of the desired Redshift table
+        :param iam_role: AWS IAM role with the related permissions
+        :param diststyle: Redshift distribution styles. Must be in ["AUTO", "EVEN", "ALL", "KEY"]
+               https://docs.aws.amazon.com/redshift/latest/dg/t_Distributing_data.html
+        :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 min_num_partitions: Minimal number of partitions
+        :param mode: append or overwrite
+        :return: None
+        """
         logger.debug(f"Minimum number of partitions : {min_num_partitions}")
         self._session.s3.delete_objects(path=path)
         num_slices = self._session.redshift.get_number_of_slices(
@@ -90,6 +112,10 @@ def write(pandas_dataframe):
             preserve_index=False,
             num_files=num_partitions,
             iam_role=iam_role,
+            diststyle=diststyle,
+            distkey=distkey,
+            sortstyle=sortstyle,
+            sortkey=sortkey,
             mode=mode,
         )
         dataframe.unpersist()

docs/pandas-to-redshift/pandas-to-redshift-flow.jpg

@@ -0,0 +1,2 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<mxfile modified="2019-08-12T13:40:33.978Z" host="www.draw.io" agent="Mozilla/5.0 (Macintosh; Intel Mac OS X 10_14_6) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/76.0.3809.100 Safari/537.36" etag="Og3SVW1iHjwL1TZe0QbU" version="11.1.4" type="google"><diagram id="uhJXyVCDerINJ9iByMH9" name="Page-1">7V3Rcpu6Fv2azNzz4AwIEPgxiZMmc9smJ06n7VNGBtmmxeCC3Nj9+isJBAhhh9QG0/rmITFCCKG9lvbeS8I5M64W63cxWs4/RB4OzoDmrc+M0RkAuq4b9A8r2aQlNgRpwSz2vaxSUTD2f+GsUMtKV76HE6kiiaKA+Eu50I3CELtEKkNxHL3I1aZRIN91iWZYKRi7KFBLP/semaelDrCL8lvsz+bizjocpmcWSFTOniSZIy96KRUZ12fGVRxFJP20WF/hgA2eGJf0upstZ/OOxTgkTS4Ibx/0FdLvH759Ag/Jx9vRlDwP9MwaCdmIJ8YeHYDsMIxC+ucyjlahh1k7Gj2KYjKPZlGIgvdRtKSFOi38hgnZZOZDKxLRojlZBNlZ2sV484Vdf26Jw69Zc/xgtJaONtlRQlBMLpgZi87wshs/CLI66jhkQ5NEq9jFOx7ezPCE4hkmuwYpa5CNTOkO2TC/w9EC007TCjEOEPF/ytBBGQJneb3CSPRDZqc32Czr9k8UrLI7vaPdB1q4WkxwTD9EU/b0ge9S4gAY0Ee7nNByOGOf/JCefMReMven7KqrYJUQdlkFBrLNX+Y+weMl4uP5Qqku23erDX7imOD1zkHLzkLAsMEvyqYKw3LS45eCeLpg07xEOqi1NNJ2J+RoF+RWQ5DDXmHcUjA+wgEmmPWQYmpD5n44Y42ECfUP7PnmmHfOxbRvngp7ghfLKEb8EcYGH0uKqWODHlgy5IFuK5AHdg3kDbOlgYfmCTsE2JArTq+4AhWuXFHzkJwUsl/gIREjjjhfZcoDCj3EqowQQdMYLfDxaSLgLmjiqDQxhJcu0wQMWxpzEZ6WBj0fSTGQomAacWMU4wd/rCJxYpBwilzQCrq1XBcnRSsfI9EO7WjalNw8LS7dsgvurn3ypfT5a0FjelQQlx10w9umfNzIuUVjemYtPUQ+t6SIWPRhNWJxHLmRdCbJrivH5qIpUTGaThNMFFjmfd8DqU6noNDeAIp+OgGnoRPQrUN7AX4pfS60KVVYMqQk21Gob0VhAaS0zYPCCnQ2AX4tspjfmQGpiyEyVGNMb4kmvAKDSjbCtLZ1eWaNaAkK/FlIC1wKHuo5jUvmqnwXBRfZiYXveez6ywBNcHCJ3O8zzqCrKIhifl9jyn9qkbiTp4pXzLWNrMtnZfmgzlsOtHNzaNoHmeMGQIJWrhe0P285W/G1jPEWhE2RKwPsFgc/MbNcGVYNgAgYENNpRJvk1h24qXlZlXg2+Q+P3qkt6ANqlc//1DjymjCsGnvRv2jB4qeg/HTF+Q93H58/fvpwef34fH/z/Hj/efz8dP88uhs/Pd5dfnq6rm0xL7mpK1SotIVh0qhX48F5tJiskm5iQXW+M1SdwDBrkiZda0soMLpRCvrpMIcNHabRLxltqEwwn2OfKwxsgID2gOIfKy6sTf2AFSw5e+kIEZ/4EcufaKZEf894JTZN5CkXVxi+441gO0buPGvofDdHsxzsfEzraZ9ZWy5itwojko4Mb5OVBH5CBmhKHdTgJe84sxe7N0WO0jE/STuVnB8/oXM0icKmVaN71El9TlsMFusUf7TWJ1LgV6k47BUTRbdLVJSd5YPgXJ3L5H68UihfPs6k8N4p3CZUYd+twq137bgkt1V4saM4Lr2pNt4zuqji+J500RUQdM0NqxrUHZ8aAHZLDb1P1ACgITX6IYJAR5fQAwxQsX8bEghQeHjheQlPA4vk70ZKPQvZXcMseKwykfM3XZpK48fe+SxdM45NTKsbYgpFM1cxv5bO/EmKJjCOReb97KxmaP+XGPeRGAVxDiIxZpjaV19U1BwBwvYVRnDKsW/jOUF4ub5MCtt14V4tvP4xc4J9wDlBt6AcKxxkjqhGIHkC1MEUYWxFW7LkupwKLh6UDZI0KmPwCqN4gQIVYKOIuR1ty86kamQ453sy0FuiS+OmBNW0v30Q8IfVKd+u2efn1On3emv6fUeL5L0MD42mqqFh9MoVGKpsmG6ETRnA5/dME+e8YrlVlVWI34bp5wSzikyhP75SaIAqRayme2GNtnY8mapIdEIUaRotGf1KoQzVf4lFrpQkCxT6U5yQ828JJ0rBlhefixBVgkSTbzQm4K6IL37RZvxYPA/beMt5d3wKAcesUAiKsKG8xqR3SSFLU4blhCjUVGw3+6W2G6rafioUghUC2VAlkFVDILOtMM3sSPnrZcZuNpXkTbNXBDLV/ektZez3FNvZVoi/P3E3DynmmaYQW/fM1YGlJOtAbqO9ZN3qTBq6WC5Lqze92ZffS7/feNqymgqNJQ16qNky2qw9t/a3j1IxHuXXdea4tHfsw/3oWvX9HQo0BpA3Z9lOzYpf3f5Ks7Xc85TlGeHOX/f7dr/8vvqe8ujx/oGWPF1cvqcQ1+5u6K/rL3fjp/HRg12oVUA/zD1ZGfaw7iWztt7FtDpeoOrVDhSB5sP5jY5gLxZUyq9jhgyIeSqYv2HpIcIEfbJZpisBbNNu6dX8yuk0W3RjjErZZxrlZVlnce2R2WTB6qKNY5k1fAJ1bsRoK320TlnCFM75dTfSMz6pu0DGbFDEahiJUZggN9uR3y8aVBUU22iooOhtRVLwlCkgPMXrLqVfEqTod4kC7yPk5T4gW7TqHfirCrwjtgQeD/7glOF/cFTXa0A0K5fMPrQqAfKWt+IPpgOpHqO8LvxjhflT3o2OThHDkOWLXM4opxwabIcgaH17ZXprAy6/XP8Y/vfx17/wYrCbHwUVrovS49GlFs6/+c0vNQu8tQPUlSOovfnuNxJat87b0sNjWMc8ppuGXWwWk6awgEYBfMGwfhtLT/d/VYLi3D28FhVAqy3DnbK8CJuuy8N+yYtQXZdPvyIuyfkgxcWlF7PLOy778TVwsPqOjT08dqSsylgtLaY9YrKK+Z6JleviJJmuAvaS+++vrb3JehQjQWkV1rOw45mcYHH0HZfOOGBiQHiouA/I9tZ11d662K8kLbG0RqfOXrZ5mrOZrEhUWSp7TWe3mM0taxcvUynnOPafOi523Tr7TxzLtHZOtW9YbNDlDInyHaj2r/WBbfHd7kYYateXiW/K7DrDtW25hZYTXPGY7XP1ijvSv3oqhlUqKkSEdWsU9nbM7JXi7f4akr8iAd+eVzeJQq2Ogs5dvZTW9hYLn4tL5XWI/ePJmm1dyrDu2MFhq29t10QYoCbCaE1a2r13syVkHxCh+2oL9a4EOnbFUFUXfzhnsv2xJEg/RkEgvjSqR6CGoLqeMLTrviPkMJimh8W/ZkjHuvgHF8b1/wA=</diagram></mxfile>