def _read_text_chunked(
paths: List[str],
chunksize: int,
parser_func: Callable[..., pd.DataFrame],
path_root: Optional[str],
boto3_session: boto3.Session,
pandas_kwargs: Dict[str, Any],
s3_additional_kwargs: Optional[Dict[str, str]],
dataset: bool,
) -> Iterator[pd.DataFrame]:
for path in paths:
_logger.debug("path: %s", path)
fs: s3fs.S3FileSystem = _utils.get_fs(
s3fs_block_size=8_388_608,
session=boto3_session,
s3_additional_kwargs=s3_additional_kwargs, # 8 MB (8 * 2**20)
)
mode, encoding, newline = _get_read_details(
path=path, pandas_kwargs=pandas_kwargs)
with _utils.open_file(fs=fs,
path=path,
mode=mode,
encoding=encoding,
newline=newline) as f:
reader: pandas.io.parsers.TextFileReader = parser_func(
f, chunksize=chunksize, **pandas_kwargs)
for df in reader:
yield _apply_partitions(df=df,
dataset=dataset,
path=path,
path_root=path_root)
def _read_parquet_row_group(
row_group: int,
path: str,
columns: Optional[List[str]],
categories: Optional[List[str]],
boto3_primitives: _utils.Boto3PrimitivesType,
s3_additional_kwargs: Optional[Dict[str, str]],
) -> pa.Table:
boto3_session: boto3.Session = _utils.boto3_from_primitives(
primitives=boto3_primitives)
fs: s3fs.S3FileSystem = _utils.get_fs(
s3fs_block_size=134_217_728,
session=boto3_session,
s3_additional_kwargs=s3_additional_kwargs, # 128 MB (128 * 2**20)
)
with _utils.open_file(fs=fs, path=path, mode="rb") as f:
pq_file: pyarrow.parquet.ParquetFile = pyarrow.parquet.ParquetFile(
source=f, read_dictionary=categories)
num_row_groups: int = pq_file.num_row_groups
_logger.debug("Reading Row Group %s/%s [multi-threaded]",
row_group + 1, num_row_groups)
return pq_file.read_row_group(i=row_group,
columns=columns,
use_threads=False,
use_pandas_metadata=False)
def _read_text_file(
path: str,
parser_func: Callable[..., pd.DataFrame],
path_root: Optional[str],
boto3_session: Union[boto3.Session, Dict[str, Optional[str]]],
pandas_kwargs: Dict[str, Any],
s3_additional_kwargs: Optional[Dict[str, str]],
dataset: bool,
) -> pd.DataFrame:
fs: s3fs.S3FileSystem = _utils.get_fs(
s3fs_block_size=134_217_728,
session=boto3_session,
s3_additional_kwargs=s3_additional_kwargs, # 128 MB (128 * 2**20)
)
mode, encoding, newline = _get_read_details(path=path,
pandas_kwargs=pandas_kwargs)
with _utils.open_file(fs=fs,
path=path,
mode=mode,
encoding=encoding,
newline=newline) as f:
df: pd.DataFrame = parser_func(f, **pandas_kwargs)
return _apply_partitions(df=df,
dataset=dataset,
path=path,
path_root=path_root)
def _read_parquet_init(
path: Union[str, List[str]],
filters: Optional[Union[List[Tuple], List[List[Tuple]]]] = None,
categories: List[str] = None,
validate_schema: bool = True,
dataset: bool = False,
use_threads: bool = True,
boto3_session: Optional[boto3.Session] = None,
s3_additional_kwargs: Optional[Dict[str, str]] = None,
) -> pyarrow.parquet.ParquetDataset:
"""Encapsulate all initialization before the use of the pyarrow.parquet.ParquetDataset."""
session: boto3.Session = _utils.ensure_session(session=boto3_session)
if dataset is False:
path_or_paths: Union[str, List[str]] = path2list(path=path,
boto3_session=session)
elif isinstance(path, str):
path_or_paths = path[:-1] if path.endswith("/") else path
else:
path_or_paths = path
_logger.debug("path_or_paths: %s", path_or_paths)
fs: s3fs.S3FileSystem = _utils.get_fs(
session=session, s3_additional_kwargs=s3_additional_kwargs)
cpus: int = _utils.ensure_cpu_count(use_threads=use_threads)
data: pyarrow.parquet.ParquetDataset = pyarrow.parquet.ParquetDataset(
path_or_paths=path_or_paths,
filesystem=fs,
metadata_nthreads=cpus,
filters=filters,
read_dictionary=categories,
validate_schema=validate_schema,
split_row_groups=False,
use_legacy_dataset=True,
)
return data
def _read_text_full(
parser_func: Callable,
path_root: str,
path: str,
boto3_session: Union[boto3.Session, Dict[str, Optional[str]]],
pandas_args: Dict[str, Any],
s3_additional_kwargs: Optional[Dict[str, str]] = None,
dataset: bool = False,
) -> pd.DataFrame:
fs: s3fs.S3FileSystem = _utils.get_fs(
session=boto3_session, s3_additional_kwargs=s3_additional_kwargs)
if pandas_args.get("compression", "infer") == "infer":
pandas_args["compression"] = infer_compression(path,
compression="infer")
mode: str = "r" if pandas_args.get("compression") is None else "rb"
encoding: Optional[str] = pandas_args.get("encoding", None)
newline: Optional[str] = pandas_args.get("lineterminator", None)
with fs.open(path=path, mode=mode, encoding=encoding,
newline=newline) as f:
df: pd.DataFrame = parser_func(f, **pandas_args)
if dataset is True:
partitions: Dict[str, Any] = _utils.extract_partitions_from_path(
path_root=path_root, path=path)
for column_name, value in partitions.items():
df[column_name] = value
return df
def _read_text_chunksize(
parser_func: Callable,
path_root: str,
paths: List[str],
boto3_session: boto3.Session,
chunksize: int,
pandas_args: Dict[str, Any],
s3_additional_kwargs: Optional[Dict[str, str]] = None,
dataset: bool = False,
) -> Iterator[pd.DataFrame]:
fs: s3fs.S3FileSystem = _utils.get_fs(
session=boto3_session, s3_additional_kwargs=s3_additional_kwargs)
for path in paths:
_logger.debug("path: %s", path)
partitions: Dict[str, Any] = {}
if dataset is True:
partitions = _utils.extract_partitions_from_path(
path_root=path_root, path=path)
if pandas_args.get("compression", "infer") == "infer":
pandas_args["compression"] = infer_compression(path,
compression="infer")
mode: str = "r" if pandas_args.get("compression") is None else "rb"
with fs.open(path, mode) as f:
reader: pandas.io.parsers.TextFileReader = parser_func(
f, chunksize=chunksize, **pandas_args)
for df in reader:
if dataset is True:
for column_name, value in partitions.items():
df[column_name] = value
yield df
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 _get_fs(
boto3_session: Optional[boto3.Session],
s3_additional_kwargs: Optional[Dict[str, str]]) -> s3fs.S3FileSystem:
return _utils.get_fs(
s3fs_block_size=33_554_432, # 32 MB (32 * 2**20)
session=boto3_session,
s3_additional_kwargs=s3_additional_kwargs,
)
def _to_parquet_file(
df: pd.DataFrame,
schema: pa.Schema,
index: bool,
compression: Optional[str],
compression_ext: str,
cpus: int,
dtype: Dict[str, str],
boto3_session: Optional[boto3.Session],
s3_additional_kwargs: Optional[Dict[str, str]],
path: Optional[str] = None,
path_root: Optional[str] = None,
) -> str:
if path is None and path_root is not None:
file_path: str = f"{path_root}{uuid.uuid4().hex}{compression_ext}.parquet"
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.")
_logger.debug("file_path: %s", file_path)
table: pa.Table = pyarrow.Table.from_pandas(df=df,
schema=schema,
nthreads=cpus,
preserve_index=index,
safe=True)
for col_name, col_type in dtype.items():
if col_name in table.column_names:
col_index = table.column_names.index(col_name)
pyarrow_dtype = _data_types.athena2pyarrow(col_type)
field = pa.field(name=col_name, type=pyarrow_dtype)
table = table.set_column(
col_index, field,
table.column(col_name).cast(pyarrow_dtype))
_logger.debug("Casting column %s (%s) to %s (%s)", col_name,
col_index, col_type, pyarrow_dtype)
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)
)
with pyarrow.parquet.ParquetWriter(
where=file_path,
write_statistics=True,
use_dictionary=True,
filesystem=fs,
coerce_timestamps="ms",
compression=compression,
flavor="spark",
schema=table.schema,
) as writer:
writer.write_table(table)
return file_path
def _read_parquet_metadata_file(
path: str,
boto3_session: boto3.Session,
s3_additional_kwargs: Optional[Dict[str, str]],
) -> Dict[str, str]:
fs: s3fs.S3FileSystem = _utils.get_fs(
s3fs_block_size=4_194_304,
session=boto3_session,
s3_additional_kwargs=s3_additional_kwargs # 4 MB (4 * 2**20)
)
with _utils.open_file(fs=fs, path=path, mode="rb") as f:
pq_file: pyarrow.parquet.ParquetFile = pyarrow.parquet.ParquetFile(
source=f)
return _data_types.athena_types_from_pyarrow_schema(
schema=pq_file.schema.to_arrow_schema(), partitions=None)[0]
def _count_row_groups(
path: str,
categories: Optional[List[str]],
boto3_session: boto3.Session,
s3_additional_kwargs: Optional[Dict[str, str]],
) -> int:
fs: s3fs.S3FileSystem = _utils.get_fs(
s3fs_block_size=4_194_304,
session=boto3_session,
s3_additional_kwargs=s3_additional_kwargs # 4 MB (4 * 2**20)
)
with _utils.open_file(fs=fs, path=path, mode="rb") as f:
pq_file: pyarrow.parquet.ParquetFile = pyarrow.parquet.ParquetFile(
source=f, read_dictionary=categories)
return pq_file.num_row_groups
def _read_parquet_file(
path: str,
columns: Optional[List[str]],
categories: Optional[List[str]],
boto3_session: boto3.Session,
s3_additional_kwargs: Optional[Dict[str, str]],
) -> pa.Table:
fs: s3fs.S3FileSystem = _utils.get_fs(
s3fs_block_size=134_217_728,
session=boto3_session,
s3_additional_kwargs=s3_additional_kwargs, # 128 MB (128 * 2**20)
)
with _utils.open_file(fs=fs, path=path, mode="rb") as f:
pq_file: pyarrow.parquet.ParquetFile = pyarrow.parquet.ParquetFile(
source=f, read_dictionary=categories)
return pq_file.read(columns=columns,
use_threads=False,
use_pandas_metadata=False)
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_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 _read_parquet_chunked(
paths: List[str],
chunked: Union[bool, int],
columns: Optional[List[str]],
categories: Optional[List[str]],
validate_schema: bool,
safe: bool,
boto3_session: boto3.Session,
dataset: bool,
path_root: Optional[str],
s3_additional_kwargs: Optional[Dict[str, str]],
use_threads: bool,
) -> Iterator[pd.DataFrame]:
next_slice: Optional[pd.DataFrame] = None
fs: s3fs.S3FileSystem = _utils.get_fs(
s3fs_block_size=8_388_608,
session=boto3_session,
s3_additional_kwargs=s3_additional_kwargs # 8 MB (8 * 2**20)
)
last_schema: Optional[Dict[str, str]] = None
last_path: str = ""
for path in paths:
with _utils.open_file(fs=fs, path=path, mode="rb") as f:
pq_file: pyarrow.parquet.ParquetFile = pyarrow.parquet.ParquetFile(
source=f, read_dictionary=categories)
schema: Dict[str,
str] = _data_types.athena_types_from_pyarrow_schema(
schema=pq_file.schema.to_arrow_schema(),
partitions=None)[0]
if validate_schema is True and last_schema is not None:
if schema != last_schema:
raise exceptions.InvalidSchemaConvergence(
f"Was detect at least 2 different schemas:\n"
f" - {last_path} -> {last_schema}\n"
f" - {path} -> {schema}")
last_schema = schema
last_path = path
num_row_groups: int = pq_file.num_row_groups
_logger.debug("num_row_groups: %s", num_row_groups)
for i in range(num_row_groups):
_logger.debug("Reading Row Group %s...", i)
df: pd.DataFrame = _arrowtable2df(
table=pq_file.read_row_group(i=i,
columns=columns,
use_threads=use_threads,
use_pandas_metadata=False),
categories=categories,
safe=safe,
use_threads=use_threads,
dataset=dataset,
path=path,
path_root=path_root,
)
if chunked is True:
yield df
elif isinstance(chunked, int) and chunked > 0:
if next_slice is not None:
df = pd.concat(objs=[next_slice, df],
ignore_index=True,
sort=False,
copy=False)
while len(df.index) >= chunked:
yield df.iloc[:chunked]
df = df.iloc[chunked:]
if df.empty:
next_slice = None
else:
next_slice = df
else:
raise exceptions.InvalidArgument(f"chunked: {chunked}")
if next_slice is not None:
yield next_slice
