def _to_text(
file_format: str,
df: pd.DataFrame,
use_threads: bool,
boto3_session: Optional[boto3.Session],
s3_additional_kwargs: Optional[Dict[str, str]],
path: Optional[str] = None,
path_root: Optional[str] = None,
**pandas_kwargs: Any,
) -> List[str]:
if df.empty is True:
raise exceptions.EmptyDataFrame()
if path is None and path_root is not None:
file_path: str = f"{path_root}{uuid.uuid4().hex}.{file_format}"
elif path is not None and path_root is None:
file_path = path
else:
raise RuntimeError("path and path_root received at the same time.")
encoding: Optional[str] = pandas_kwargs.get("encoding", None)
with open_s3_object(
path=file_path,
mode="w",
use_threads=use_threads,
s3_additional_kwargs=s3_additional_kwargs,
boto3_session=boto3_session,
encoding=encoding,
newline=None,
) as f:
_logger.debug("pandas_kwargs: %s", pandas_kwargs)
if file_format == "csv":
df.to_csv(f, **pandas_kwargs)
elif file_format == "json":
df.to_json(f, **pandas_kwargs)
return [file_path]
def _to_text(
file_format: str,
df: pd.DataFrame,
boto3_session: Optional[boto3.Session],
s3_additional_kwargs: Optional[Dict[str, str]],
path: Optional[str] = None,
path_root: Optional[str] = None,
**pandas_kwargs,
) -> str:
if df.empty is True:
raise exceptions.EmptyDataFrame()
if path is None and path_root is not None:
file_path: str = f"{path_root}{uuid.uuid4().hex}.{file_format}"
elif path is not None and path_root is None:
file_path = path
else:
raise RuntimeError("path and path_root received at the same time.")
fs: s3fs.S3FileSystem = _utils.get_fs(
s3fs_block_size=33_554_432,
session=boto3_session,
s3_additional_kwargs=s3_additional_kwargs, # 32 MB (32 * 2**20)
)
encoding: Optional[str] = pandas_kwargs.get("encoding", None)
newline: Optional[str] = pandas_kwargs.get("line_terminator", None)
with _utils.open_file(fs=fs,
path=file_path,
mode="w",
encoding=encoding,
newline=newline) as f:
_logger.debug("pandas_kwargs: %s", pandas_kwargs)
if file_format == "csv":
df.to_csv(f, **pandas_kwargs)
elif file_format == "json":
df.to_json(f, **pandas_kwargs)
return file_path
def _validate_args(
df: pd.DataFrame,
table: Optional[str],
dataset: bool,
path: str,
partition_cols: Optional[List[str]],
mode: Optional[str],
description: Optional[str],
parameters: Optional[Dict[str, str]],
columns_comments: Optional[Dict[str, str]],
) -> None:
if df.empty is True:
raise exceptions.EmptyDataFrame()
if dataset is False:
if path.endswith("/"):
raise exceptions.InvalidArgumentValue(
"If <dataset=False>, the argument <path> should be a object path, not a directory."
)
if partition_cols:
raise exceptions.InvalidArgumentCombination(
"Please, pass dataset=True to be able to use partition_cols.")
if mode is not None:
raise exceptions.InvalidArgumentCombination(
"Please pass dataset=True to be able to use mode.")
if any(arg is not None
for arg in (table, description, parameters, columns_comments)):
raise exceptions.InvalidArgumentCombination(
"Please pass dataset=True to be able to use any one of these "
"arguments: database, table, description, parameters, "
"columns_comments.")
def _validate_args(
df: pd.DataFrame,
table: Optional[str],
database: Optional[str],
dataset: bool,
path: str,
partition_cols: Optional[List[str]],
mode: Optional[str],
description: Optional[str],
parameters: Optional[Dict[str, str]],
columns_comments: Optional[Dict[str, str]],
) -> None:
if df.empty is True:
raise exceptions.EmptyDataFrame()
if dataset is False:
if path.endswith("/"):
raise exceptions.InvalidArgumentValue(
"If <dataset=False>, the argument <path> should be a file path, not a directory."
)
if partition_cols:
raise exceptions.InvalidArgumentCombination(
"Please, pass dataset=True to be able to use partition_cols.")
if mode is not None:
raise exceptions.InvalidArgumentCombination(
"Please pass dataset=True to be able to use mode.")
if any(arg is not None
for arg in (table, description, parameters, columns_comments)):
raise exceptions.InvalidArgumentCombination(
"Please pass dataset=True to be able to use any one of these "
"arguments: database, table, description, parameters, "
"columns_comments.")
elif (database is None) != (table is None):
raise exceptions.InvalidArgumentCombination(
"Arguments database and table must be passed together. If you want to store your dataset metadata in "
"the Glue Catalog, please ensure you are passing both.")
def _validate_args(
df: pd.DataFrame,
table: Optional[str],
database: Optional[str],
dataset: bool,
path: Optional[str],
partition_cols: Optional[List[str]],
bucketing_info: Optional[Tuple[List[str], int]],
mode: Optional[str],
description: Optional[str],
parameters: Optional[Dict[str, str]],
columns_comments: Optional[Dict[str, str]],
) -> None:
if df.empty is True:
raise exceptions.EmptyDataFrame("DataFrame cannot be empty.")
if dataset is False:
if path is None:
raise exceptions.InvalidArgumentValue(
"If dataset is False, the `path` argument must be passed.")
if path.endswith("/"):
raise exceptions.InvalidArgumentValue(
"If <dataset=False>, the argument <path> should be a key, not a prefix."
)
if partition_cols:
raise exceptions.InvalidArgumentCombination(
"Please, pass dataset=True to be able to use partition_cols.")
if bucketing_info:
raise exceptions.InvalidArgumentCombination(
"Please, pass dataset=True to be able to use bucketing_info.")
if mode is not None:
raise exceptions.InvalidArgumentCombination(
"Please pass dataset=True to be able to use mode.")
if any(arg is not None
for arg in (table, description, parameters, columns_comments)):
raise exceptions.InvalidArgumentCombination(
"Please pass dataset=True to be able to use any one of these "
"arguments: database, table, description, parameters, "
"columns_comments.")
elif (database is None) != (table is None):
raise exceptions.InvalidArgumentCombination(
"Arguments database and table must be passed together. If you want to store your dataset metadata in "
"the Glue Catalog, please ensure you are passing both.")
elif all(x is None for x in [path, database, table]):
raise exceptions.InvalidArgumentCombination(
"You must specify a `path` if dataset is True and database/table are not enabled."
)
elif bucketing_info and bucketing_info[1] <= 0:
raise exceptions.InvalidArgumentValue(
"Please pass a value greater than 1 for the number of buckets for bucketing."
)
def _to_text(
file_format: str,
df: pd.DataFrame,
path: str,
fs: Optional[s3fs.S3FileSystem] = None,
boto3_session: Optional[boto3.Session] = None,
s3_additional_kwargs: Optional[Dict[str, str]] = None,
**pandas_kwargs,
) -> None:
if df.empty is True: # pragma: no cover
raise exceptions.EmptyDataFrame()
if fs is None:
fs = _utils.get_fs(session=boto3_session, s3_additional_kwargs=s3_additional_kwargs)
encoding: Optional[str] = pandas_kwargs.get("encoding", None)
newline: Optional[str] = pandas_kwargs.get("line_terminator", None)
with fs.open(path=path, mode="w", encoding=encoding, newline=newline) as f:
if file_format == "csv":
df.to_csv(f, **pandas_kwargs)
elif file_format == "json":
df.to_json(f, **pandas_kwargs)
def _to_text(
file_format: str,
df: pd.DataFrame,
use_threads: bool,
boto3_session: Optional[boto3.Session],
s3_additional_kwargs: Optional[Dict[str, str]],
path: Optional[str] = None,
path_root: Optional[str] = None,
filename: Optional[str] = None,
**pandas_kwargs: Any,
) -> List[str]:
if df.empty is True:
raise exceptions.EmptyDataFrame()
if path is None and path_root is not None:
if filename is None:
filename = uuid.uuid4().hex
file_path: str = (
f"{path_root}{filename}.{file_format}{_COMPRESSION_2_EXT.get(pandas_kwargs.get('compression'))}"
)
elif path is not None and path_root is None:
file_path = path
else:
raise RuntimeError("path and path_root received at the same time.")
mode, encoding, newline = _get_write_details(path=file_path,
pandas_kwargs=pandas_kwargs)
with open_s3_object(
path=file_path,
mode=mode,
use_threads=use_threads,
s3_additional_kwargs=s3_additional_kwargs,
boto3_session=boto3_session,
encoding=encoding,
newline=newline,
) as f:
_logger.debug("pandas_kwargs: %s", pandas_kwargs)
if file_format == "csv":
df.to_csv(f, mode=mode, **pandas_kwargs)
elif file_format == "json":
df.to_json(f, **pandas_kwargs)
return [file_path]
def to_sql(
df: pd.DataFrame,
con: pg8000.Connection,
table: str,
schema: str,
mode: str = "append",
index: bool = False,
dtype: Optional[Dict[str, str]] = None,
varchar_lengths: Optional[Dict[str, int]] = None,
) -> None:
"""Write records stored in a DataFrame into PostgreSQL.
Parameters
----------
df : pandas.DataFrame
Pandas DataFrame https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html
con : pg8000.Connection
Use pg8000.connect() to use "
"credentials directly or wr.postgresql.connect() to fetch it from the Glue Catalog.
table : str
Table name
schema : str
Schema name
mode : str
Append or overwrite.
index : bool
True to store the DataFrame index as a column in the table,
otherwise False to ignore it.
dtype: Dict[str, str], optional
Dictionary of columns names and PostgreSQL types to be casted.
Useful when you have columns with undetermined or mixed data types.
(e.g. {'col name': 'TEXT', 'col2 name': 'FLOAT'})
varchar_lengths : Dict[str, int], optional
Dict of VARCHAR length by columns. (e.g. {"col1": 10, "col5": 200}).
Returns
-------
None
None.
Examples
--------
Writing to PostgreSQL using a Glue Catalog Connections
>>> import awswrangler as wr
>>> con = wr.postgresql.connect("MY_GLUE_CONNECTION")
>>> wr.postgresql.to_sql(
... df=df
... table="my_table",
... schema="public",
... con=con
... )
>>> con.close()
"""
if df.empty is True:
raise exceptions.EmptyDataFrame()
_validate_connection(con=con)
try:
with con.cursor() as cursor:
_create_table(
df=df,
cursor=cursor,
table=table,
schema=schema,
mode=mode,
index=index,
dtype=dtype,
varchar_lengths=varchar_lengths,
)
if index:
df.reset_index(level=df.index.names, inplace=True)
placeholders: str = ", ".join(["%s"] * len(df.columns))
sql: str = f'INSERT INTO "{schema}"."{table}" VALUES ({placeholders})'
_logger.debug("sql: %s", sql)
parameters: List[List[Any]] = _db_utils.extract_parameters(df=df)
cursor.executemany(sql, parameters)
con.commit()
except Exception as ex:
con.rollback()
_logger.error(ex)
raise
def to_sql(df: pd.DataFrame, con: sqlalchemy.engine.Engine,
**pandas_kwargs: Any) -> None:
"""Write records stored in a DataFrame to a SQL database.
Support for **Redshift**, **PostgreSQL** and **MySQL**.
Support for all pandas to_sql() arguments:
https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.to_sql.html
Note
----
Redshift: For large DataFrames (1MM+ rows) consider the function **wr.db.copy_to_redshift()**.
Note
----
Redshift: `index=False` will be forced.
Parameters
----------
df : pandas.DataFrame
Pandas DataFrame https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html
con : sqlalchemy.engine.Engine
SQLAlchemy Engine. Please use,
wr.db.get_engine(), wr.db.get_redshift_temp_engine() or wr.catalog.get_engine()
pandas_kwargs
KEYWORD arguments forwarded to pandas.DataFrame.to_sql(). You can NOT pass `pandas_kwargs` explicit, just add
valid Pandas arguments in the function call and Wrangler will accept it.
e.g. wr.db.to_sql(df, con=con, name="table_name", schema="schema_name", if_exists="replace", index=False)
https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.to_sql.html
Returns
-------
None
None.
Examples
--------
Writing to Redshift with temporary credentials
>>> import awswrangler as wr
>>> import pandas as pd
>>> wr.db.to_sql(
... df=pd.DataFrame({'col': [1, 2, 3]}),
... con=wr.db.get_redshift_temp_engine(cluster_identifier="...", user="******"),
... name="table_name",
... schema="schema_name"
... )
Writing to Redshift with temporary credentials and using pandas_kwargs
>>> import awswrangler as wr
>>> import pandas as pd
>>> wr.db.to_sql(
... df=pd.DataFrame({'col': [1, 2, 3]}),
... con=wr.db.get_redshift_temp_engine(cluster_identifier="...", user="******"),
... name="table_name",
... schema="schema_name",
... if_exists="replace",
... index=False,
... )
Writing to Redshift from Glue Catalog Connections
>>> import awswrangler as wr
>>> import pandas as pd
>>> wr.db.to_sql(
... df=pd.DataFrame({'col': [1, 2, 3]}),
... con=wr.catalog.get_engine(connection="..."),
... name="table_name",
... schema="schema_name"
... )
"""
if "pandas_kwargs" in pandas_kwargs:
raise exceptions.InvalidArgument(
"You can NOT pass `pandas_kwargs` explicit, just add valid "
"Pandas arguments in the function call and Wrangler will accept it."
"e.g. wr.db.to_sql(df, con, name='...', schema='...', if_exists='replace')"
)
if df.empty is True:
raise exceptions.EmptyDataFrame()
if not isinstance(con, sqlalchemy.engine.Engine):
raise exceptions.InvalidConnection(
"Invalid 'con' argument, please pass a "
"SQLAlchemy Engine. Use wr.db.get_engine(), "
"wr.db.get_redshift_temp_engine() or wr.catalog.get_engine()")
if "dtype" in pandas_kwargs:
cast_columns: Dict[str, VisitableType] = pandas_kwargs["dtype"]
else:
cast_columns = {}
dtypes: Dict[str,
VisitableType] = _data_types.sqlalchemy_types_from_pandas(
df=df, db_type=con.name, dtype=cast_columns)
pandas_kwargs["dtype"] = dtypes
pandas_kwargs["con"] = con
if pandas_kwargs["con"].name.lower(
) == "redshift": # Redshift does not accept index
pandas_kwargs["index"] = False
_utils.try_it(f=df.to_sql,
ex=sqlalchemy.exc.InternalError,
**pandas_kwargs)
def to_parquet( # pylint: disable=too-many-arguments,too-many-locals
df: pd.DataFrame,
path: str,
index: bool = False,
compression: Optional[str] = "snappy",
use_threads: bool = True,
boto3_session: Optional[boto3.Session] = None,
s3_additional_kwargs: Optional[Dict[str, str]] = None,
sanitize_columns: bool = False,
dataset: bool = False,
partition_cols: Optional[List[str]] = None,
mode: Optional[str] = None,
catalog_versioning: bool = False,
database: Optional[str] = None,
table: Optional[str] = None,
dtype: Optional[Dict[str, str]] = None,
description: Optional[str] = None,
parameters: Optional[Dict[str, str]] = None,
columns_comments: Optional[Dict[str, str]] = None,
regular_partitions: bool = True,
projection_enabled: bool = False,
projection_types: Optional[Dict[str, str]] = None,
projection_ranges: Optional[Dict[str, str]] = None,
projection_values: Optional[Dict[str, str]] = None,
projection_intervals: Optional[Dict[str, str]] = None,
projection_digits: Optional[Dict[str, str]] = None,
) -> Dict[str, Union[List[str], Dict[str, List[str]]]]:
"""Write Parquet file or dataset on Amazon S3.
The concept of Dataset goes beyond the simple idea of files and enable more
complex features like partitioning, casting and catalog integration (Amazon Athena/AWS Glue Catalog).
Note
----
If `dataset=True` The table name and all column names will be automatically sanitized using
`wr.catalog.sanitize_table_name` and `wr.catalog.sanitize_column_name`.
Please, pass `sanitize_columns=True` to force the same behaviour for `dataset=False`.
Note
----
On `append` mode, the `parameters` will be upsert on an existing table.
Note
----
In case of `use_threads=True` the number of threads that will be spawned will be get from os.cpu_count().
Parameters
----------
df: pandas.DataFrame
Pandas DataFrame https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html
path : str
S3 path (for file e.g. ``s3://bucket/prefix/filename.parquet``) (for dataset e.g. ``s3://bucket/prefix``).
index : bool
True to store the DataFrame index in file, otherwise False to ignore it.
compression: str, optional
Compression style (``None``, ``snappy``, ``gzip``).
use_threads : bool
True to enable concurrent requests, False to disable multiple threads.
If enabled os.cpu_count() will be used as the max number of threads.
boto3_session : boto3.Session(), optional
Boto3 Session. The default boto3 session will be used if boto3_session receive None.
s3_additional_kwargs:
Forward to s3fs, useful for server side encryption
https://s3fs.readthedocs.io/en/latest/#serverside-encryption
sanitize_columns : bool
True to sanitize columns names or False to keep it as is.
True value is forced if `dataset=True`.
dataset : bool
If True store a parquet dataset instead of a single file.
If True, enable all follow arguments:
partition_cols, mode, database, table, description, parameters, columns_comments, .
partition_cols: List[str], optional
List of column names that will be used to create partitions. Only takes effect if dataset=True.
mode: str, optional
``append`` (Default), ``overwrite``, ``overwrite_partitions``. Only takes effect if dataset=True.
catalog_versioning : bool
If True and `mode="overwrite"`, creates an archived version of the table catalog before updating it.
database : str, optional
Glue/Athena catalog: Database name.
table : str, optional
Glue/Athena catalog: Table name.
dtype : Dict[str, str], optional
Dictionary of columns names and Athena/Glue types to be casted.
Useful when you have columns with undetermined or mixed data types.
(e.g. {'col name': 'bigint', 'col2 name': 'int'})
description : str, optional
Glue/Athena catalog: Table description
parameters : Dict[str, str], optional
Glue/Athena catalog: Key/value pairs to tag the table.
columns_comments : Dict[str, str], optional
Glue/Athena catalog:
Columns names and the related comments (e.g. {'col0': 'Column 0.', 'col1': 'Column 1.', 'col2': 'Partition.'}).
regular_partitions : bool
Create regular partitions (Non projected partitions) on Glue Catalog.
Disable when you will work only with Partition Projection.
Keep enabled even when working with projections is useful to keep
Redshift Spectrum working with the regular partitions.
projection_enabled : bool
Enable Partition Projection on Athena (https://docs.aws.amazon.com/athena/latest/ug/partition-projection.html)
projection_types : Optional[Dict[str, str]]
Dictionary of partitions names and Athena projections types.
Valid types: "enum", "integer", "date", "injected"
https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html
(e.g. {'col_name': 'enum', 'col2_name': 'integer'})
projection_ranges: Optional[Dict[str, str]]
Dictionary of partitions names and Athena projections ranges.
https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html
(e.g. {'col_name': '0,10', 'col2_name': '-1,8675309'})
projection_values: Optional[Dict[str, str]]
Dictionary of partitions names and Athena projections values.
https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html
(e.g. {'col_name': 'A,B,Unknown', 'col2_name': 'foo,boo,bar'})
projection_intervals: Optional[Dict[str, str]]
Dictionary of partitions names and Athena projections intervals.
https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html
(e.g. {'col_name': '1', 'col2_name': '5'})
projection_digits: Optional[Dict[str, str]]
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'})
Returns
-------
Dict[str, Union[List[str], Dict[str, List[str]]]]
Dictionary with:
'paths': List of all stored files paths on S3.
'partitions_values': Dictionary of partitions added with keys as S3 path locations
and values as a list of partitions values as str.
Examples
--------
Writing single file
>>> import awswrangler as wr
>>> import pandas as pd
>>> wr.s3.to_parquet(
... df=pd.DataFrame({'col': [1, 2, 3]}),
... path='s3://bucket/prefix/my_file.parquet',
... )
{
'paths': ['s3://bucket/prefix/my_file.parquet'],
'partitions_values': {}
}
Writing single file encrypted with a KMS key
>>> import awswrangler as wr
>>> import pandas as pd
>>> wr.s3.to_parquet(
... df=pd.DataFrame({'col': [1, 2, 3]}),
... path='s3://bucket/prefix/my_file.parquet',
... s3_additional_kwargs={
... 'ServerSideEncryption': 'aws:kms',
... 'SSEKMSKeyId': 'YOUR_KMY_KEY_ARN'
... }
... )
{
'paths': ['s3://bucket/prefix/my_file.parquet'],
'partitions_values': {}
}
Writing partitioned dataset
>>> import awswrangler as wr
>>> import pandas as pd
>>> wr.s3.to_parquet(
... df=pd.DataFrame({
... 'col': [1, 2, 3],
... 'col2': ['A', 'A', 'B']
... }),
... path='s3://bucket/prefix',
... dataset=True,
... partition_cols=['col2']
... )
{
'paths': ['s3://.../col2=A/x.parquet', 's3://.../col2=B/y.parquet'],
'partitions_values: {
's3://.../col2=A/': ['A'],
's3://.../col2=B/': ['B']
}
}
Writing dataset to S3 with metadata on Athena/Glue Catalog.
>>> import awswrangler as wr
>>> import pandas as pd
>>> wr.s3.to_parquet(
... df=pd.DataFrame({
... 'col': [1, 2, 3],
... 'col2': ['A', 'A', 'B']
... }),
... path='s3://bucket/prefix',
... dataset=True,
... partition_cols=['col2'],
... database='default', # Athena/Glue database
... table='my_table' # Athena/Glue table
... )
{
'paths': ['s3://.../col2=A/x.parquet', 's3://.../col2=B/y.parquet'],
'partitions_values: {
's3://.../col2=A/': ['A'],
's3://.../col2=B/': ['B']
}
}
Writing dataset casting empty column data type
>>> import awswrangler as wr
>>> import pandas as pd
>>> wr.s3.to_parquet(
... df=pd.DataFrame({
... 'col': [1, 2, 3],
... 'col2': ['A', 'A', 'B'],
... 'col3': [None, None, None]
... }),
... path='s3://bucket/prefix',
... dataset=True,
... database='default', # Athena/Glue database
... table='my_table' # Athena/Glue table
... dtype={'col3': 'date'}
... )
{
'paths': ['s3://.../x.parquet'],
'partitions_values: {}
}
"""
if (database is None) ^ (table is None):
raise exceptions.InvalidArgumentCombination(
"Please pass database and table arguments to be able to store the metadata into the Athena/Glue Catalog."
)
if df.empty is True:
raise exceptions.EmptyDataFrame()
partition_cols = partition_cols if partition_cols else []
dtype = dtype if dtype else {}
partitions_values: Dict[str, List[str]] = {}
# Sanitize table to respect Athena's standards
if (sanitize_columns is True) or (dataset is True):
df = catalog.sanitize_dataframe_columns_names(df=df)
partition_cols = [catalog.sanitize_column_name(p) for p in partition_cols]
dtype = {catalog.sanitize_column_name(k): v.lower() for k, v in dtype.items()}
catalog.drop_duplicated_columns(df=df)
session: boto3.Session = _utils.ensure_session(session=boto3_session)
cpus: int = _utils.ensure_cpu_count(use_threads=use_threads)
fs: s3fs.S3FileSystem = _utils.get_fs(session=session, s3_additional_kwargs=s3_additional_kwargs)
compression_ext: Optional[str] = _COMPRESSION_2_EXT.get(compression, None)
if compression_ext is None:
raise exceptions.InvalidCompression(f"{compression} is invalid, please use None, snappy or gzip.")
if dataset is False:
if path.endswith("/"): # pragma: no cover
raise exceptions.InvalidArgumentValue(
"If <dataset=False>, the argument <path> should be a object path, not a directory."
)
if partition_cols:
raise exceptions.InvalidArgumentCombination("Please, pass dataset=True to be able to use partition_cols.")
if mode is not None:
raise exceptions.InvalidArgumentCombination("Please pass dataset=True to be able to use mode.")
if any(arg is not None for arg in (database, table, description, parameters)):
raise exceptions.InvalidArgumentCombination(
"Please pass dataset=True to be able to use any one of these "
"arguments: database, table, description, parameters, "
"columns_comments."
)
df = _data_types.cast_pandas_with_athena_types(df=df, dtype=dtype)
schema: pa.Schema = _data_types.pyarrow_schema_from_pandas(
df=df, index=index, ignore_cols=partition_cols, dtype=dtype
)
_logger.debug("schema: \n%s", schema)
paths = [
_to_parquet_file(
df=df, path=path, schema=schema, index=index, compression=compression, cpus=cpus, fs=fs, dtype=dtype
)
]
else:
mode = "append" if mode is None else mode
if (
(mode in ("append", "overwrite_partitions")) and (database is not None) and (table is not None)
): # Fetching Catalog Types
catalog_types: Optional[Dict[str, str]] = catalog.get_table_types(
database=database, table=table, boto3_session=session
)
if catalog_types is not None:
for k, v in catalog_types.items():
dtype[k] = v
paths, partitions_values = _to_parquet_dataset(
df=df,
path=path,
index=index,
compression=compression,
compression_ext=compression_ext,
cpus=cpus,
fs=fs,
use_threads=use_threads,
partition_cols=partition_cols,
dtype=dtype,
mode=mode,
boto3_session=session,
)
if (database is not None) and (table is not None):
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(
database=database,
table=table,
path=path,
columns_types=columns_types,
partitions_types=partitions_types,
compression=compression,
description=description,
parameters=parameters,
columns_comments=columns_comments,
boto3_session=session,
mode=mode,
catalog_versioning=catalog_versioning,
projection_enabled=projection_enabled,
projection_types=projection_types,
projection_ranges=projection_ranges,
projection_values=projection_values,
projection_intervals=projection_intervals,
projection_digits=projection_digits,
)
if partitions_values and (regular_partitions is True):
_logger.debug("partitions_values:\n%s", partitions_values)
catalog.add_parquet_partitions(
database=database,
table=table,
partitions_values=partitions_values,
compression=compression,
boto3_session=session,
)
return {"paths": paths, "partitions_values": partitions_values}
def to_csv( # pylint: disable=too-many-arguments,too-many-locals
df: pd.DataFrame,
path: str,
sep: str = ",",
index: bool = True,
columns: Optional[List[str]] = None,
use_threads: bool = True,
boto3_session: Optional[boto3.Session] = None,
s3_additional_kwargs: Optional[Dict[str, str]] = None,
sanitize_columns: bool = False,
dataset: bool = False,
partition_cols: Optional[List[str]] = None,
mode: Optional[str] = None,
catalog_versioning: bool = False,
database: Optional[str] = None,
table: Optional[str] = None,
dtype: Optional[Dict[str, str]] = None,
description: Optional[str] = None,
parameters: Optional[Dict[str, str]] = None,
columns_comments: Optional[Dict[str, str]] = None,
regular_partitions: bool = True,
projection_enabled: bool = False,
projection_types: Optional[Dict[str, str]] = None,
projection_ranges: Optional[Dict[str, str]] = None,
projection_values: Optional[Dict[str, str]] = None,
projection_intervals: Optional[Dict[str, str]] = None,
projection_digits: Optional[Dict[str, str]] = None,
**pandas_kwargs,
) -> Dict[str, Union[List[str], Dict[str, List[str]]]]:
"""Write CSV file or dataset on Amazon S3.
The concept of Dataset goes beyond the simple idea of files and enable more
complex features like partitioning, casting and catalog integration (Amazon Athena/AWS Glue Catalog).
Note
----
If `dataset=True` The table name and all column names will be automatically sanitized using
`wr.catalog.sanitize_table_name` and `wr.catalog.sanitize_column_name`.
Please, pass `sanitize_columns=True` to force the same behaviour for `dataset=False`.
Note
----
If `dataset=True`, `pandas_kwargs` will be ignored due
restrictive quoting, date_format, escapechar, encoding, etc required by Athena/Glue Catalog.
Note
----
By now Pandas does not support in-memory CSV compression.
https://github.com/pandas-dev/pandas/issues/22555
So the `compression` will not be supported on Wrangler too.
Note
----
On `append` mode, the `parameters` will be upsert on an existing table.
Note
----
In case of `use_threads=True` the number of threads that will be spawned will be get from os.cpu_count().
Parameters
----------
df: pandas.DataFrame
Pandas DataFrame https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html
path : str
Amazon S3 path (e.g. s3://bucket/filename.csv).
sep : str
String of length 1. Field delimiter for the output file.
index : bool
Write row names (index).
columns : List[str], optional
Columns to write.
use_threads : bool
True to enable concurrent requests, False to disable multiple threads.
If enabled os.cpu_count() will be used as the max number of threads.
boto3_session : boto3.Session(), optional
Boto3 Session. The default boto3 Session will be used if boto3_session receive None.
s3_additional_kwargs:
Forward to s3fs, useful for server side encryption
https://s3fs.readthedocs.io/en/latest/#serverside-encryption
sanitize_columns : bool
True to sanitize columns names or False to keep it as is.
True value is forced if `dataset=True`.
dataset : bool
If True store a parquet dataset instead of a single file.
If True, enable all follow arguments:
partition_cols, mode, database, table, description, parameters, columns_comments, .
partition_cols: List[str], optional
List of column names that will be used to create partitions. Only takes effect if dataset=True.
mode : str, optional
``append`` (Default), ``overwrite``, ``overwrite_partitions``. Only takes effect if dataset=True.
catalog_versioning : bool
If True and `mode="overwrite"`, creates an archived version of the table catalog before updating it.
database : str, optional
Glue/Athena catalog: Database name.
table : str, optional
Glue/Athena catalog: Table name.
dtype : Dict[str, str], optional
Dictionary of columns names and Athena/Glue types to be casted.
Useful when you have columns with undetermined or mixed data types.
(e.g. {'col name': 'bigint', 'col2 name': 'int'})
description : str, optional
Glue/Athena catalog: Table description
parameters : Dict[str, str], optional
Glue/Athena catalog: Key/value pairs to tag the table.
columns_comments : Dict[str, str], optional
Glue/Athena catalog:
Columns names and the related comments (e.g. {'col0': 'Column 0.', 'col1': 'Column 1.', 'col2': 'Partition.'}).
regular_partitions : bool
Create regular partitions (Non projected partitions) on Glue Catalog.
Disable when you will work only with Partition Projection.
Keep enabled even when working with projections is useful to keep
Redshift Spectrum working with the regular partitions.
projection_enabled : bool
Enable Partition Projection on Athena (https://docs.aws.amazon.com/athena/latest/ug/partition-projection.html)
projection_types : Optional[Dict[str, str]]
Dictionary of partitions names and Athena projections types.
Valid types: "enum", "integer", "date", "injected"
https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html
(e.g. {'col_name': 'enum', 'col2_name': 'integer'})
projection_ranges: Optional[Dict[str, str]]
Dictionary of partitions names and Athena projections ranges.
https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html
(e.g. {'col_name': '0,10', 'col2_name': '-1,8675309'})
projection_values: Optional[Dict[str, str]]
Dictionary of partitions names and Athena projections values.
https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html
(e.g. {'col_name': 'A,B,Unknown', 'col2_name': 'foo,boo,bar'})
projection_intervals: Optional[Dict[str, str]]
Dictionary of partitions names and Athena projections intervals.
https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html
(e.g. {'col_name': '1', 'col2_name': '5'})
projection_digits: Optional[Dict[str, str]]
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'})
pandas_kwargs :
keyword arguments forwarded to pandas.DataFrame.to_csv()
https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.to_csv.html
Returns
-------
None
None.
Examples
--------
Writing single file
>>> import awswrangler as wr
>>> import pandas as pd
>>> wr.s3.to_csv(
... df=pd.DataFrame({'col': [1, 2, 3]}),
... path='s3://bucket/prefix/my_file.csv',
... )
{
'paths': ['s3://bucket/prefix/my_file.csv'],
'partitions_values': {}
}
Writing single file encrypted with a KMS key
>>> import awswrangler as wr
>>> import pandas as pd
>>> wr.s3.to_csv(
... df=pd.DataFrame({'col': [1, 2, 3]}),
... path='s3://bucket/prefix/my_file.csv',
... s3_additional_kwargs={
... 'ServerSideEncryption': 'aws:kms',
... 'SSEKMSKeyId': 'YOUR_KMY_KEY_ARN'
... }
... )
{
'paths': ['s3://bucket/prefix/my_file.csv'],
'partitions_values': {}
}
Writing partitioned dataset
>>> import awswrangler as wr
>>> import pandas as pd
>>> wr.s3.to_csv(
... df=pd.DataFrame({
... 'col': [1, 2, 3],
... 'col2': ['A', 'A', 'B']
... }),
... path='s3://bucket/prefix',
... dataset=True,
... partition_cols=['col2']
... )
{
'paths': ['s3://.../col2=A/x.csv', 's3://.../col2=B/y.csv'],
'partitions_values: {
's3://.../col2=A/': ['A'],
's3://.../col2=B/': ['B']
}
}
Writing dataset to S3 with metadata on Athena/Glue Catalog.
>>> import awswrangler as wr
>>> import pandas as pd
>>> wr.s3.to_csv(
... df=pd.DataFrame({
... 'col': [1, 2, 3],
... 'col2': ['A', 'A', 'B']
... }),
... path='s3://bucket/prefix',
... dataset=True,
... partition_cols=['col2'],
... database='default', # Athena/Glue database
... table='my_table' # Athena/Glue table
... )
{
'paths': ['s3://.../col2=A/x.csv', 's3://.../col2=B/y.csv'],
'partitions_values: {
's3://.../col2=A/': ['A'],
's3://.../col2=B/': ['B']
}
}
Writing dataset casting empty column data type
>>> import awswrangler as wr
>>> import pandas as pd
>>> wr.s3.to_csv(
... df=pd.DataFrame({
... 'col': [1, 2, 3],
... 'col2': ['A', 'A', 'B'],
... 'col3': [None, None, None]
... }),
... path='s3://bucket/prefix',
... dataset=True,
... database='default', # Athena/Glue database
... table='my_table' # Athena/Glue table
... dtype={'col3': 'date'}
... )
{
'paths': ['s3://.../x.csv'],
'partitions_values: {}
}
"""
if (database is None) ^ (table is None):
raise exceptions.InvalidArgumentCombination(
"Please pass database and table arguments to be able to store the metadata into the Athena/Glue Catalog."
)
if df.empty is True:
raise exceptions.EmptyDataFrame()
partition_cols = partition_cols if partition_cols else []
dtype = dtype if dtype else {}
partitions_values: Dict[str, List[str]] = {}
# Sanitize table to respect Athena's standards
if (sanitize_columns is True) or (dataset is True):
df = catalog.sanitize_dataframe_columns_names(df=df)
partition_cols = [catalog.sanitize_column_name(p) for p in partition_cols]
dtype = {catalog.sanitize_column_name(k): v.lower() for k, v in dtype.items()}
catalog.drop_duplicated_columns(df=df)
session: boto3.Session = _utils.ensure_session(session=boto3_session)
fs: s3fs.S3FileSystem = _utils.get_fs(session=session, s3_additional_kwargs=s3_additional_kwargs)
if dataset is False:
if partition_cols:
raise exceptions.InvalidArgumentCombination("Please, pass dataset=True to be able to use partition_cols.")
if mode is not None:
raise exceptions.InvalidArgumentCombination("Please pass dataset=True to be able to use mode.")
if columns_comments:
raise exceptions.InvalidArgumentCombination("Please pass dataset=True to be able to use columns_comments.")
if any(arg is not None for arg in (database, table, description, parameters)):
raise exceptions.InvalidArgumentCombination(
"Please pass dataset=True to be able to use any one of these "
"arguments: database, table, description, parameters, "
"columns_comments."
)
pandas_kwargs["sep"] = sep
pandas_kwargs["index"] = index
pandas_kwargs["columns"] = columns
_to_text(file_format="csv", df=df, path=path, fs=fs, **pandas_kwargs)
paths = [path]
else:
mode = "append" if mode is None else mode
if columns:
df = df[columns]
if (
(mode in ("append", "overwrite_partitions")) and (database is not None) and (table is not None)
): # Fetching Catalog Types
catalog_types: Optional[Dict[str, str]] = catalog.get_table_types(
database=database, table=table, boto3_session=session
)
if catalog_types is not None:
for k, v in catalog_types.items():
dtype[k] = v
paths, partitions_values = _to_csv_dataset(
df=df,
path=path,
index=index,
sep=sep,
fs=fs,
use_threads=use_threads,
partition_cols=partition_cols,
dtype=dtype,
mode=mode,
boto3_session=session,
)
if (database is not None) and (table is not None):
columns_types, partitions_types = _data_types.athena_types_from_pandas_partitioned(
df=df, index=index, partition_cols=partition_cols, dtype=dtype, index_left=True
)
catalog.create_csv_table(
database=database,
table=table,
path=path,
columns_types=columns_types,
partitions_types=partitions_types,
description=description,
parameters=parameters,
columns_comments=columns_comments,
boto3_session=session,
mode=mode,
catalog_versioning=catalog_versioning,
sep=sep,
projection_enabled=projection_enabled,
projection_types=projection_types,
projection_ranges=projection_ranges,
projection_values=projection_values,
projection_intervals=projection_intervals,
projection_digits=projection_digits,
)
if partitions_values and (regular_partitions is True):
_logger.debug("partitions_values:\n%s", partitions_values)
catalog.add_csv_partitions(
database=database, table=table, partitions_values=partitions_values, boto3_session=session, sep=sep
)
return {"paths": paths, "partitions_values": partitions_values}
def to_sql(df: pd.DataFrame, con: sqlalchemy.engine.Engine,
**pandas_kwargs) -> None:
"""Write records stored in a DataFrame to a SQL database.
Support for **Redshift**, **PostgreSQL** and **MySQL**.
Support for all pandas to_sql() arguments:
https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.to_sql.html
Note
----
Redshift: For large DataFrames (1MM+ rows) consider the function **wr.db.copy_to_redshift()**.
Parameters
----------
df : pandas.DataFrame
Pandas DataFrame https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html
con : sqlalchemy.engine.Engine
SQLAlchemy Engine. Please use,
wr.db.get_engine(), wr.db.get_redshift_temp_engine() or wr.catalog.get_engine()
pandas_kwargs
keyword arguments forwarded to pandas.DataFrame.to_csv()
https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.to_sql.html
Returns
-------
None
None.
Examples
--------
Writing to Redshift with temporary credentials
>>> import awswrangler as wr
>>> import pandas as pd
>>> wr.db.to_sql(
... df=pd.DataFrame({'col': [1, 2, 3]}),
... con=wr.db.get_redshift_temp_engine(cluster_identifier="...", user="******"),
... name="table_name",
... schema="schema_name"
... )
Writing to Redshift from Glue Catalog Connections
>>> import awswrangler as wr
>>> import pandas as pd
>>> wr.db.to_sql(
... df=pd.DataFrame({'col': [1, 2, 3]}),
... con=wr.catalog.get_engine(connection="..."),
... name="table_name",
... schema="schema_name"
... )
"""
if df.empty is True: # pragma: no cover
raise exceptions.EmptyDataFrame()
if not isinstance(con, sqlalchemy.engine.Engine): # pragma: no cover
raise exceptions.InvalidConnection(
"Invalid 'con' argument, please pass a "
"SQLAlchemy Engine. Use wr.db.get_engine(), "
"wr.db.get_redshift_temp_engine() or wr.catalog.get_engine()")
if "dtype" in pandas_kwargs:
cast_columns: Dict[str, VisitableType] = pandas_kwargs["dtype"]
else:
cast_columns = {}
dtypes: Dict[str,
VisitableType] = _data_types.sqlalchemy_types_from_pandas(
df=df, db_type=con.name, dtype=cast_columns)
pandas_kwargs["dtype"] = dtypes
pandas_kwargs["con"] = con
df.to_sql(**pandas_kwargs)
def to_sql(
df: pd.DataFrame,
con: "pyodbc.Connection",
table: str,
schema: str,
mode: str = "append",
index: bool = False,
dtype: Optional[Dict[str, str]] = None,
varchar_lengths: Optional[Dict[str, int]] = None,
use_column_names: bool = False,
chunksize: int = 200,
) -> None:
"""Write records stored in a DataFrame into Microsoft SQL Server.
Parameters
----------
df : pandas.DataFrame
Pandas DataFrame https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html
con : pyodbc.Connection
Use pyodbc.connect() to use credentials directly or wr.sqlserver.connect() to fetch it from the Glue Catalog.
table : str
Table name
schema : str
Schema name
mode : str
Append or overwrite.
index : bool
True to store the DataFrame index as a column in the table,
otherwise False to ignore it.
dtype: Dict[str, str], optional
Dictionary of columns names and Microsoft SQL Server types to be casted.
Useful when you have columns with undetermined or mixed data types.
(e.g. {'col name': 'TEXT', 'col2 name': 'FLOAT'})
varchar_lengths : Dict[str, int], optional
Dict of VARCHAR length by columns. (e.g. {"col1": 10, "col5": 200}).
use_column_names: bool
If set to True, will use the column names of the DataFrame for generating the INSERT SQL Query.
E.g. If the DataFrame has two columns `col1` and `col3` and `use_column_names` is True, data will only be
inserted into the database columns `col1` and `col3`.
chunksize: int
Number of rows which are inserted with each SQL query. Defaults to inserting 200 rows per query.
Returns
-------
None
None.
Examples
--------
Writing to Microsoft SQL Server using a Glue Catalog Connections
>>> import awswrangler as wr
>>> con = wr.sqlserver.connect(connection="MY_GLUE_CONNECTION", odbc_driver_version=17)
>>> wr.sqlserver.to_sql(
... df=df,
... table="table",
... schema="dbo",
... con=con
... )
>>> con.close()
"""
if df.empty is True:
raise exceptions.EmptyDataFrame()
_validate_connection(con=con)
try:
with con.cursor() as cursor:
_create_table(
df=df,
cursor=cursor,
table=table,
schema=schema,
mode=mode,
index=index,
dtype=dtype,
varchar_lengths=varchar_lengths,
)
if index:
df.reset_index(level=df.index.names, inplace=True)
column_placeholders: str = ", ".join(["?"] * len(df.columns))
table_identifier = _get_table_identifier(schema, table)
insertion_columns = ""
if use_column_names:
insertion_columns = f"({', '.join(df.columns)})"
placeholder_parameter_pair_generator = _db_utils.generate_placeholder_parameter_pairs(
df=df,
column_placeholders=column_placeholders,
chunksize=chunksize)
for placeholders, parameters in placeholder_parameter_pair_generator:
sql: str = f"INSERT INTO {table_identifier} {insertion_columns} VALUES {placeholders}"
_logger.debug("sql: %s", sql)
cursor.executemany(sql, (parameters, ))
con.commit()
except Exception as ex:
con.rollback()
_logger.error(ex)
raise
def to_sql(
df: pd.DataFrame,
con: pymysql.connections.Connection,
table: str,
schema: str,
mode: str = "append",
index: bool = False,
dtype: Optional[Dict[str, str]] = None,
varchar_lengths: Optional[Dict[str, int]] = None,
use_column_names: bool = False,
chunksize: int = 200,
) -> None:
"""Write records stored in a DataFrame into MySQL.
Parameters
----------
df : pandas.DataFrame
Pandas DataFrame https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html
con : pymysql.connections.Connection
Use pymysql.connect() to use credentials directly or wr.mysql.connect() to fetch it from the Glue Catalog.
table : str
Table name
schema : str
Schema name
mode : str
Append, overwrite, upsert_duplicate_key, upsert_replace_into, upsert_distinct.
append: Inserts new records into table
overwrite: Drops table and recreates
upsert_duplicate_key: Performs an upsert using `ON DUPLICATE KEY` clause. Requires table schema to have
defined keys, otherwise duplicate records will be inserted.
upsert_replace_into: Performs upsert using `REPLACE INTO` clause. Less efficient and still requires the
table schema to have keys or else duplicate records will be inserted
upsert_distinct: Inserts new records, including duplicates, then recreates the table and inserts `DISTINCT`
records from old table. This is the least efficient approach but handles scenarios where there are no
keys on table.
index : bool
True to store the DataFrame index as a column in the table,
otherwise False to ignore it.
dtype: Dict[str, str], optional
Dictionary of columns names and MySQL types to be casted.
Useful when you have columns with undetermined or mixed data types.
(e.g. {'col name': 'TEXT', 'col2 name': 'FLOAT'})
varchar_lengths : Dict[str, int], optional
Dict of VARCHAR length by columns. (e.g. {"col1": 10, "col5": 200}).
use_column_names: bool
If set to True, will use the column names of the DataFrame for generating the INSERT SQL Query.
E.g. If the DataFrame has two columns `col1` and `col3` and `use_column_names` is True, data will only be
inserted into the database columns `col1` and `col3`.
chunksize: int
Number of rows which are inserted with each SQL query. Defaults to inserting 200 rows per query.
Returns
-------
None
None.
Examples
--------
Writing to MySQL using a Glue Catalog Connections
>>> import awswrangler as wr
>>> con = wr.mysql.connect("MY_GLUE_CONNECTION")
>>> wr.mysql.to_sql(
... df=df,
... table="my_table",
... schema="test",
... con=con
... )
>>> con.close()
"""
if df.empty is True:
raise exceptions.EmptyDataFrame()
mode = mode.strip().lower()
modes = [
"append",
"overwrite",
"upsert_replace_into",
"upsert_duplicate_key",
"upsert_distinct",
]
if mode not in modes:
raise exceptions.InvalidArgumentValue(
f"mode must be one of {', '.join(modes)}")
_validate_connection(con=con)
try:
with con.cursor() as cursor:
_create_table(
df=df,
cursor=cursor,
table=table,
schema=schema,
mode=mode,
index=index,
dtype=dtype,
varchar_lengths=varchar_lengths,
)
if index:
df.reset_index(level=df.index.names, inplace=True)
column_placeholders: str = ", ".join(["%s"] * len(df.columns))
insertion_columns = ""
upsert_columns = ""
upsert_str = ""
if use_column_names:
insertion_columns = f"({', '.join(df.columns)})"
if mode == "upsert_duplicate_key":
upsert_columns = ", ".join(
df.columns.map(
lambda column: f"`{column}`=VALUES(`{column}`)"))
upsert_str = f" ON DUPLICATE KEY UPDATE {upsert_columns}"
placeholder_parameter_pair_generator = _db_utils.generate_placeholder_parameter_pairs(
df=df,
column_placeholders=column_placeholders,
chunksize=chunksize)
sql: str
for placeholders, parameters in placeholder_parameter_pair_generator:
if mode == "upsert_replace_into":
sql = f"REPLACE INTO `{schema}`.`{table}` {insertion_columns} VALUES {placeholders}"
else:
sql = f"INSERT INTO `{schema}`.`{table}` {insertion_columns} VALUES {placeholders}{upsert_str}"
_logger.debug("sql: %s", sql)
cursor.executemany(sql, (parameters, ))
con.commit()
if mode == "upsert_distinct":
temp_table = f"{table}_{uuid.uuid4().hex}"
cursor.execute(
f"CREATE TABLE `{schema}`.`{temp_table}` LIKE `{schema}`.`{table}`"
)
cursor.execute(
f"INSERT INTO `{schema}`.`{temp_table}` SELECT DISTINCT * FROM `{schema}`.`{table}`"
)
cursor.execute(f"DROP TABLE IF EXISTS `{schema}`.`{table}`")
cursor.execute(
f"ALTER TABLE `{schema}`.`{temp_table}` RENAME TO `{table}`"
)
con.commit()
except Exception as ex:
con.rollback()
_logger.error(ex)
raise
def to_sql(
df: pd.DataFrame,
con: redshift_connector.Connection,
table: str,
schema: str,
mode: str = "append",
index: bool = False,
dtype: Optional[Dict[str, str]] = None,
diststyle: str = "AUTO",
distkey: Optional[str] = None,
sortstyle: str = "COMPOUND",
sortkey: Optional[List[str]] = None,
primary_keys: Optional[List[str]] = None,
varchar_lengths_default: int = 256,
varchar_lengths: Optional[Dict[str, int]] = None,
) -> None:
"""Write records stored in a DataFrame into Redshift.
Note
----
For large DataFrames (1K+ rows) consider the function **wr.redshift.copy()**.
Parameters
----------
df : pandas.DataFrame
Pandas DataFrame https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html
con : redshift_connector.Connection
Use redshift_connector.connect() to use "
"credentials directly or wr.redshift.connect() to fetch it from the Glue Catalog.
table : str
Table name
schema : str
Schema name
mode : str
Append, overwrite or upsert.
index : bool
True to store the DataFrame index as a column in the table,
otherwise False to ignore it.
dtype: Dict[str, str], optional
Dictionary of columns names and Redshift types to be casted.
Useful when you have columns with undetermined or mixed data types.
(e.g. {'col name': 'VARCHAR(10)', 'col2 name': 'FLOAT'})
diststyle : str
Redshift distribution styles. Must be in ["AUTO", "EVEN", "ALL", "KEY"].
https://docs.aws.amazon.com/redshift/latest/dg/t_Distributing_data.html
distkey : str, optional
Specifies a column name or positional number for the distribution key.
sortstyle : str
Sorting can be "COMPOUND" or "INTERLEAVED".
https://docs.aws.amazon.com/redshift/latest/dg/t_Sorting_data.html
sortkey : List[str], optional
List of columns to be sorted.
primary_keys : List[str], optional
Primary keys.
varchar_lengths_default : int
The size that will be set for all VARCHAR columns not specified with varchar_lengths.
varchar_lengths : Dict[str, int], optional
Dict of VARCHAR length by columns. (e.g. {"col1": 10, "col5": 200}).
Returns
-------
None
None.
Examples
--------
Writing to Redshift using a Glue Catalog Connections
>>> import awswrangler as wr
>>> con = wr.redshift.connect("MY_GLUE_CONNECTION")
>>> wr.redshift.to_sql(
... df=df
... table="my_table",
... schema="public",
... con=con
... )
>>> con.close()
"""
if df.empty is True:
raise exceptions.EmptyDataFrame()
_validate_connection(con=con)
con.autocommit = False
try:
with con.cursor() as cursor:
created_table, created_schema = _create_table(
df=df,
path=None,
cursor=cursor,
table=table,
schema=schema,
mode=mode,
index=index,
dtype=dtype,
diststyle=diststyle,
sortstyle=sortstyle,
distkey=distkey,
sortkey=sortkey,
primary_keys=primary_keys,
varchar_lengths_default=varchar_lengths_default,
varchar_lengths=varchar_lengths,
)
if index:
df.reset_index(level=df.index.names, inplace=True)
placeholders: str = ", ".join(["%s"] * len(df.columns))
schema_str = f"{created_schema}." if created_schema else ""
sql: str = f"INSERT INTO {schema_str}{created_table} VALUES ({placeholders})"
_logger.debug("sql: %s", sql)
parameters: List[List[Any]] = _db_utils.extract_parameters(df=df)
cursor.executemany(sql, parameters)
if table != created_table: # upsert
_upsert(cursor=cursor,
schema=schema,
table=table,
temp_table=created_table,
primary_keys=primary_keys)
con.commit()
except Exception as ex:
con.rollback()
_logger.error(ex)
raise
def to_sql(
df: pd.DataFrame,
con: pg8000.Connection,
table: str,
schema: str,
mode: str = "append",
index: bool = False,
dtype: Optional[Dict[str, str]] = None,
varchar_lengths: Optional[Dict[str, int]] = None,
use_column_names: bool = False,
chunksize: int = 200,
upsert_conflict_columns: Optional[List[str]] = None,
insert_conflict_columns: Optional[List[str]] = None,
) -> None:
"""Write records stored in a DataFrame into PostgreSQL.
Parameters
----------
df : pandas.DataFrame
Pandas DataFrame https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html
con : pg8000.Connection
Use pg8000.connect() to use credentials directly or wr.postgresql.connect() to fetch it from the Glue Catalog.
table : str
Table name
schema : str
Schema name
mode : str
Append, overwrite or upsert.
append: Inserts new records into table.
overwrite: Drops table and recreates.
upsert: Perform an upsert which checks for conflicts on columns given by `upsert_conflict_columns` and
sets the new values on conflicts. Note that `upsert_conflict_columns` is required for this mode.
index : bool
True to store the DataFrame index as a column in the table,
otherwise False to ignore it.
dtype: Dict[str, str], optional
Dictionary of columns names and PostgreSQL types to be casted.
Useful when you have columns with undetermined or mixed data types.
(e.g. {'col name': 'TEXT', 'col2 name': 'FLOAT'})
varchar_lengths : Dict[str, int], optional
Dict of VARCHAR length by columns. (e.g. {"col1": 10, "col5": 200}).
use_column_names: bool
If set to True, will use the column names of the DataFrame for generating the INSERT SQL Query.
E.g. If the DataFrame has two columns `col1` and `col3` and `use_column_names` is True, data will only be
inserted into the database columns `col1` and `col3`.
chunksize: int
Number of rows which are inserted with each SQL query. Defaults to inserting 200 rows per query.
upsert_conflict_columns: List[str], optional
This parameter is only supported if `mode` is set top `upsert`. In this case conflicts for the given columns are
checked for evaluating the upsert.
insert_conflict_columns: List[str], optional
This parameter is only supported if `mode` is set top `append`. In this case conflicts for the given columns are
checked for evaluating the insert 'ON CONFLICT DO NOTHING'.
Returns
-------
None
None.
Examples
--------
Writing to PostgreSQL using a Glue Catalog Connections
>>> import awswrangler as wr
>>> con = wr.postgresql.connect("MY_GLUE_CONNECTION")
>>> wr.postgresql.to_sql(
... df=df,
... table="my_table",
... schema="public",
... con=con
... )
>>> con.close()
"""
if df.empty is True:
raise exceptions.EmptyDataFrame("DataFrame cannot be empty.")
mode = mode.strip().lower()
allowed_modes = ["append", "overwrite", "upsert"]
_db_utils.validate_mode(mode=mode, allowed_modes=allowed_modes)
if mode == "upsert" and not upsert_conflict_columns:
raise exceptions.InvalidArgumentValue(
"<upsert_conflict_columns> needs to be set when using upsert mode."
)
_validate_connection(con=con)
try:
with con.cursor() as cursor:
_create_table(
df=df,
cursor=cursor,
table=table,
schema=schema,
mode=mode,
index=index,
dtype=dtype,
varchar_lengths=varchar_lengths,
)
if index:
df.reset_index(level=df.index.names, inplace=True)
column_placeholders: str = ", ".join(["%s"] * len(df.columns))
insertion_columns = ""
upsert_str = ""
if use_column_names:
insertion_columns = f"({', '.join(df.columns)})"
if mode == "upsert":
upsert_columns = ", ".join(
df.columns.map(
lambda column: f"{column}=EXCLUDED.{column}"))
conflict_columns = ", ".join(
upsert_conflict_columns) # type: ignore
upsert_str = f" ON CONFLICT ({conflict_columns}) DO UPDATE SET {upsert_columns}"
if mode == "append" and insert_conflict_columns:
conflict_columns = ", ".join(
insert_conflict_columns) # type: ignore
upsert_str = f" ON CONFLICT ({conflict_columns}) DO NOTHING"
placeholder_parameter_pair_generator = _db_utils.generate_placeholder_parameter_pairs(
df=df,
column_placeholders=column_placeholders,
chunksize=chunksize)
for placeholders, parameters in placeholder_parameter_pair_generator:
sql: str = f'INSERT INTO "{schema}"."{table}" {insertion_columns} VALUES {placeholders}{upsert_str}'
_logger.debug("sql: %s", sql)
cursor.executemany(sql, (parameters, ))
con.commit()
except Exception as ex:
con.rollback()
_logger.error(ex)
raise