Add arguments to set VARCHAR lengths for Redshift and Aurora tables

Author: igorborgest

awswrangler/aurora.py

@@ -141,7 +141,9 @@ def load_table(dataframe: pd.DataFrame,
                    mode: str = "append",
                    preserve_index: bool = False,
                    engine: str = "mysql",
-                   region: str = "us-east-1"):
+                   region: str = "us-east-1",
+                   varchar_default_length: int = 256,
+                   varchar_lengths: Optional[Dict[str, int]] = None) -> None:
         """
         Load text/CSV files into a Aurora table using a manifest file.
         Creates the table if necessary.
@@ -158,6 +160,8 @@ def load_table(dataframe: pd.DataFrame,
         :param preserve_index: Should we preserve the Dataframe index? (ONLY for Pandas Dataframe)
         :param engine: "mysql" or "postgres"
         :param region: AWS S3 bucket region (Required only for postgres engine)
+        :param varchar_default_length: The size that will be set for all VARCHAR columns not specified with varchar_lengths
+        :param varchar_lengths: Dict of VARCHAR length by columns. (e.g. {"col1": 10, "col5": 200})
         :return: None
         """
         if "postgres" in engine.lower():
@@ -170,7 +174,9 @@ def load_table(dataframe: pd.DataFrame,
                                        mode=mode,
                                        preserve_index=preserve_index,
                                        region=region,
-                                       columns=columns)
+                                       columns=columns,
+                                       varchar_default_length=varchar_default_length,
+                                       varchar_lengths=varchar_lengths)
         elif "mysql" in engine.lower():
             Aurora.load_table_mysql(dataframe=dataframe,
                                     dataframe_type=dataframe_type,
@@ -181,7 +187,9 @@ def load_table(dataframe: pd.DataFrame,
                                     mode=mode,
                                     preserve_index=preserve_index,
                                     num_files=num_files,
-                                    columns=columns)
+                                    columns=columns,
+                                    varchar_default_length=varchar_default_length,
+                                    varchar_lengths=varchar_lengths)
         else:
             raise InvalidEngine(f"{engine} is not a valid engine. Please use 'mysql' or 'postgres'!")
 
@@ -195,7 +203,9 @@ def load_table_postgres(dataframe: pd.DataFrame,
                             mode: str = "append",
                             preserve_index: bool = False,
                             region: str = "us-east-1",
-                            columns: Optional[List[str]] = None):
+                            columns: Optional[List[str]] = None,
+                            varchar_default_length: int = 256,
+                            varchar_lengths: Optional[Dict[str, int]] = None):
         """
         Load text/CSV files into a Aurora table using a manifest file.
         Creates the table if necessary.
@@ -210,6 +220,8 @@ def load_table_postgres(dataframe: pd.DataFrame,
         :param preserve_index: Should we preserve the Dataframe index? (ONLY for Pandas Dataframe)
         :param region: AWS S3 bucket region (Required only for postgres engine)
         :param columns: List of columns to load
+        :param varchar_default_length: The size that will be set for all VARCHAR columns not specified with varchar_lengths
+        :param varchar_lengths: Dict of VARCHAR length by columns. (e.g. {"col1": 10, "col5": 200})
         :return: None
         """
         with connection.cursor() as cursor:
@@ -221,7 +233,9 @@ def load_table_postgres(dataframe: pd.DataFrame,
                                      table_name=table_name,
                                      preserve_index=preserve_index,
                                      engine="postgres",
-                                     columns=columns)
+                                     columns=columns,
+                                     varchar_default_length=varchar_default_length,
+                                     varchar_lengths=varchar_lengths)
                 connection.commit()
                 logger.debug("CREATE TABLE committed.")
         for path in load_paths:
@@ -266,7 +280,9 @@ def load_table_mysql(dataframe: pd.DataFrame,
                          num_files: int,
                          mode: str = "append",
                          preserve_index: bool = False,
-                         columns: Optional[List[str]] = None):
+                         columns: Optional[List[str]] = None,
+                         varchar_default_length: int = 256,
+                         varchar_lengths: Optional[Dict[str, int]] = None):
         """
         Load text/CSV files into a Aurora table using a manifest file.
         Creates the table if necessary.
@@ -281,6 +297,8 @@ def load_table_mysql(dataframe: pd.DataFrame,
         :param mode: append or overwrite
         :param preserve_index: Should we preserve the Dataframe index? (ONLY for Pandas Dataframe)
         :param columns: List of columns to load
+        :param varchar_default_length: The size that will be set for all VARCHAR columns not specified with varchar_lengths
+        :param varchar_lengths: Dict of VARCHAR length by columns. (e.g. {"col1": 10, "col5": 200})
         :return: None
         """
         with connection.cursor() as cursor:
@@ -292,7 +310,9 @@ def load_table_mysql(dataframe: pd.DataFrame,
                                      table_name=table_name,
                                      preserve_index=preserve_index,
                                      engine="mysql",
-                                     columns=columns)
+                                     columns=columns,
+                                     varchar_default_length=varchar_default_length,
+                                     varchar_lengths=varchar_lengths)
             sql = Aurora._get_load_sql(path=manifest_path,
                                        schema_name=schema_name,
                                        table_name=table_name,
@@ -368,7 +388,9 @@ def _create_table(cursor,
                       table_name,
                       preserve_index=False,
                       engine: str = "mysql",
-                      columns: Optional[List[str]] = None):
+                      columns: Optional[List[str]] = None,
+                      varchar_default_length: int = 256,
+                      varchar_lengths: Optional[Dict[str, int]] = None) -> None:
         """
         Creates Aurora table.
 
@@ -380,6 +402,8 @@ def _create_table(cursor,
         :param preserve_index: Should we preserve the Dataframe index? (ONLY for Pandas Dataframe)
         :param engine: "mysql" or "postgres"
         :param columns: List of columns to load
+        :param varchar_default_length: The size that will be set for all VARCHAR columns not specified with varchar_lengths
+        :param varchar_lengths: Dict of VARCHAR length by columns. (e.g. {"col1": 10, "col5": 200})
         :return: None
         """
         sql: str = f"-- AWS DATA WRANGLER\n" \
@@ -397,7 +421,9 @@ def _create_table(cursor,
                                     dataframe_type=dataframe_type,
                                     preserve_index=preserve_index,
                                     engine=engine,
-                                    columns=columns)
+                                    columns=columns,
+                                    varchar_default_length=varchar_default_length,
+                                    varchar_lengths=varchar_lengths)
         cols_str: str = "".join([f"{col[0]} {col[1]},\n" for col in schema])[:-2]
         sql = f"-- AWS DATA WRANGLER\n" f"CREATE TABLE IF NOT EXISTS {schema_name}.{table_name} (\n" f"{cols_str})"
         logger.debug(f"Create table query:\n{sql}")
@@ -408,7 +434,10 @@ def _get_schema(dataframe,
                     dataframe_type: str,
                     preserve_index: bool,
                     engine: str = "mysql",
-                    columns: Optional[List[str]] = None) -> List[Tuple[str, str]]:
+                    columns: Optional[List[str]] = None,
+                    varchar_default_length: int = 256,
+                    varchar_lengths: Optional[Dict[str, int]] = None) -> List[Tuple[str, str]]:
+        varchar_lengths = {} if varchar_lengths is None else varchar_lengths
         schema_built: List[Tuple[str, str]] = []
         if "postgres" in engine.lower():
             convert_func = data_types.pyarrow2postgres
@@ -421,7 +450,8 @@ def _get_schema(dataframe,
                 dataframe=dataframe, preserve_index=preserve_index, indexes_position="right")
             for name, dtype in pyarrow_schema:
                 if columns is None or name in columns:
-                    aurora_type: str = convert_func(dtype)
+                    varchar_len = varchar_lengths.get(name, varchar_default_length)
+                    aurora_type: str = convert_func(dtype=dtype, varchar_length=varchar_len)
                     schema_built.append((name, aurora_type))
         else:
             raise InvalidDataframeType(f"{dataframe_type} is not a valid DataFrame type. Please use 'pandas'!")

awswrangler/data_types.py

@@ -81,7 +81,7 @@ def athena2python(dtype: str) -> Optional[type]:
         raise UnsupportedType(f"Unsupported Athena type: {dtype}")
 
 
-def athena2redshift(dtype: str) -> str:
+def athena2redshift(dtype: str, varchar_length: int = 256) -> str:
     dtype = dtype.lower()
     if dtype == "smallint":
         return "SMALLINT"
@@ -96,7 +96,7 @@ def athena2redshift(dtype: str) -> str:
     elif dtype in ("boolean", "bool"):
         return "BOOL"
     elif dtype in ("string", "char", "varchar", "array", "row", "map"):
-        return "VARCHAR(256)"
+        return f"VARCHAR({varchar_length})"
     elif dtype == "timestamp":
         return "TIMESTAMP"
     elif dtype == "date":
@@ -125,7 +125,7 @@ def pandas2athena(dtype: str) -> str:
         raise UnsupportedType(f"Unsupported Pandas type: {dtype}")
 
 
-def pandas2redshift(dtype: str) -> str:
+def pandas2redshift(dtype: str, varchar_length: int = 256) -> str:
     dtype = dtype.lower()
     if dtype == "int32":
         return "INTEGER"
@@ -138,7 +138,7 @@ def pandas2redshift(dtype: str) -> str:
     elif dtype == "bool":
         return "BOOLEAN"
     elif dtype == "object" and isinstance(dtype, str):
-        return "VARCHAR(256)"
+        return f"VARCHAR({varchar_length})"
     elif dtype[:10] == "datetime64":
         return "TIMESTAMP"
     else:
@@ -177,7 +177,7 @@ def pyarrow2athena(dtype: pa.types) -> str:
         raise UnsupportedType(f"Unsupported Pyarrow type: {dtype}")
 
 
-def pyarrow2redshift(dtype: pa.types) -> str:
+def pyarrow2redshift(dtype: pa.types, varchar_length: int = 256) -> str:
     dtype_str = str(dtype).lower()
     if dtype_str == "int16":
         return "SMALLINT"
@@ -192,7 +192,7 @@ def pyarrow2redshift(dtype: pa.types) -> str:
     elif dtype_str == "bool":
         return "BOOLEAN"
     elif dtype_str == "string":
-        return "VARCHAR(256)"
+        return f"VARCHAR({varchar_length})"
     elif dtype_str.startswith("timestamp"):
         return "TIMESTAMP"
     elif dtype_str.startswith("date"):
@@ -203,7 +203,7 @@ def pyarrow2redshift(dtype: pa.types) -> str:
         raise UnsupportedType(f"Unsupported Pyarrow type: {dtype}")
 
 
-def pyarrow2postgres(dtype: pa.types) -> str:
+def pyarrow2postgres(dtype: pa.types, varchar_length: int = 256) -> str:
     dtype_str = str(dtype).lower()
     if dtype_str == "int16":
         return "SMALLINT"
@@ -218,7 +218,7 @@ def pyarrow2postgres(dtype: pa.types) -> str:
     elif dtype_str == "bool":
         return "BOOLEAN"
     elif dtype_str == "string":
-        return "VARCHAR(256)"
+        return f"VARCHAR({varchar_length})"
     elif dtype_str.startswith("timestamp"):
         return "TIMESTAMP"
     elif dtype_str.startswith("date"):
@@ -229,7 +229,7 @@ def pyarrow2postgres(dtype: pa.types) -> str:
         raise UnsupportedType(f"Unsupported Pyarrow type: {dtype}")
 
 
-def pyarrow2mysql(dtype: pa.types) -> str:
+def pyarrow2mysql(dtype: pa.types, varchar_length: int = 256) -> str:
     dtype_str = str(dtype).lower()
     if dtype_str == "int16":
         return "SMALLINT"
@@ -244,7 +244,7 @@ def pyarrow2mysql(dtype: pa.types) -> str:
     elif dtype_str == "bool":
         return "BOOLEAN"
     elif dtype_str == "string":
-        return "VARCHAR(256)"
+        return f"VARCHAR({varchar_length})"
     elif dtype_str.startswith("timestamp"):
         return "TIMESTAMP"
     elif dtype_str.startswith("date"):
@@ -321,7 +321,7 @@ def redshift2pyarrow(dtype: str) -> str:
         raise UnsupportedType(f"Unsupported Redshift type: {dtype_str}")
 
 
-def spark2redshift(dtype: str) -> str:
+def spark2redshift(dtype: str, varchar_length: int = 256) -> str:
     dtype = dtype.lower()
     if dtype == "smallint":
         return "SMALLINT"
@@ -340,7 +340,7 @@ def spark2redshift(dtype: str) -> str:
     elif dtype == "date":
         return "DATE"
     elif dtype == "string":
-        return "VARCHAR(256)"
+        return f"VARCHAR({varchar_length})"
     elif dtype.startswith("decimal"):
         return dtype.replace(" ", "").upper()
     else: