diff --git a/.github/workflows/static-checking.yml b/.github/workflows/static-checking.yml index c1d9ff0cb..a23a74d99 100644 --- a/.github/workflows/static-checking.yml +++ b/.github/workflows/static-checking.yml @@ -24,15 +24,12 @@ jobs: uses: actions/setup-python@v1 with: python-version: ${{ matrix.python-version }} - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install -r requirements.txt - pip install -r requirements-dev.txt + - name: Setup Environment + run: ./setup-dev-env.sh - name: CloudFormation Lint run: cfn-lint -t testing/cloudformation.yaml - name: Documentation Lint - run: pydocstyle awswrangler/ --add-ignore=D204 + run: pydocstyle awswrangler/ --add-ignore=D204,D403 - name: mypy check run: mypy awswrangler - name: Flake8 Lint diff --git a/.gitignore b/.gitignore index 1947b2e87..8a3474a30 100644 --- a/.gitignore +++ b/.gitignore @@ -138,6 +138,8 @@ testing/*parameters-*.properties testing/*requirements*.txt testing/coverage/* building/*requirements*.txt +building/arrow +building/lambda/arrow /docs/coverage/ /docs/build/ /docs/source/_build/ diff --git a/.pylintrc b/.pylintrc index 132ce213a..4f41cb3fb 100644 --- a/.pylintrc +++ b/.pylintrc @@ -141,7 +141,8 @@ disable=print-statement, comprehension-escape, C0330, C0103, - W1202 + W1202, + too-few-public-methods # Enable the message, report, category or checker with the given id(s). You can # either give multiple identifier separated by comma (,) or put this option diff --git a/README.md b/README.md index 608297330..424808bf8 100644 --- a/README.md +++ b/README.md @@ -5,19 +5,18 @@ **NOTE** -We just released a new major version `1.0` with breaking changes. Please make sure that all your old projects has dependencies frozen on the desired version (e.g. `pip install awswrangler==0.3.2`). +Due the new major version `1.*.*` with breaking changes, please make sure that all your old projects has dependencies frozen on the desired version (e.g. `pip install awswrangler==0.3.2`). --- ![AWS Data Wrangler](docs/source/_static/logo2.png?raw=true "AWS Data Wrangler") -[![Release](https://img.shields.io/badge/release-1.0.4-brightgreen.svg)](https://pypi.org/project/awswrangler/) +[![Release](https://img.shields.io/badge/release-1.1.0-brightgreen.svg)](https://pypi.org/project/awswrangler/) [![Python Version](https://img.shields.io/badge/python-3.6%20%7C%203.7%20%7C%203.8-brightgreen.svg)](https://anaconda.org/conda-forge/awswrangler) [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) -[![Checked with mypy](http://www.mypy-lang.org/static/mypy_badge.svg)](http://mypy-lang.org/) -[![Average time to resolve an issue](http://isitmaintained.com/badge/resolution/awslabs/aws-data-wrangler.svg)](http://isitmaintained.com/project/awslabs/aws-data-wrangler "Average time to resolve an issue") +[![Checked with mypy](http://www.mypy-lang.org/static/mypy_badge.svg)](http://mypy-lang.org/) [![Coverage](https://img.shields.io/badge/coverage-100%25-brightgreen.svg)](https://pypi.org/project/awswrangler/) ![Static Checking](https://github.com/awslabs/aws-data-wrangler/workflows/Static%20Checking/badge.svg?branch=master) [![Documentation Status](https://readthedocs.org/projects/aws-data-wrangler/badge/?version=latest)](https://aws-data-wrangler.readthedocs.io/?badge=latest) @@ -85,6 +84,9 @@ df = wr.db.read_sql_query("SELECT * FROM external_schema.my_table", con=engine) - [11 - CSV Datasets](https://github.com/awslabs/aws-data-wrangler/blob/master/tutorials/11%20-%20CSV%20Datasets.ipynb) - [12 - CSV Crawler](https://github.com/awslabs/aws-data-wrangler/blob/master/tutorials/12%20-%20CSV%20Crawler.ipynb) - [13 - Merging Datasets on S3](https://github.com/awslabs/aws-data-wrangler/blob/master/tutorials/13%20-%20Merging%20Datasets%20on%20S3.ipynb) + - [14 - PyTorch](https://github.com/awslabs/aws-data-wrangler/blob/master/tutorials/14%20-%20PyTorch.ipynb) + - [15 - EMR](https://github.com/awslabs/aws-data-wrangler/blob/dev/tutorials/15%20-%20EMR.ipynb) + - [16 - EMR & Docker](https://github.com/awslabs/aws-data-wrangler/blob/dev/tutorials/16%20-%20EMR%20%26%20Docker.ipynb) - [**API Reference**](https://aws-data-wrangler.readthedocs.io/en/latest/api.html) - [Amazon S3](https://aws-data-wrangler.readthedocs.io/en/latest/api.html#amazon-s3) - [AWS Glue Catalog](https://aws-data-wrangler.readthedocs.io/en/latest/api.html#aws-glue-catalog) diff --git a/awswrangler/__init__.py b/awswrangler/__init__.py index ce11c7ad5..78299541e 100644 --- a/awswrangler/__init__.py +++ b/awswrangler/__init__.py @@ -6,8 +6,13 @@ """ import logging +from importlib.util import find_spec from awswrangler import athena, catalog, cloudwatch, db, emr, exceptions, s3 # noqa from awswrangler.__metadata__ import __description__, __license__, __title__, __version__ # noqa +from awswrangler._utils import get_account_id # noqa + +if find_spec("torch") and find_spec("torchvision") and find_spec("torchaudio") and find_spec("PIL"): + from awswrangler import torch # noqa logging.getLogger("awswrangler").addHandler(logging.NullHandler()) diff --git a/awswrangler/__metadata__.py b/awswrangler/__metadata__.py index 724b626da..cfc9336b9 100644 --- a/awswrangler/__metadata__.py +++ b/awswrangler/__metadata__.py @@ -7,5 +7,5 @@ __title__ = "awswrangler" __description__ = "Pandas on AWS." -__version__ = "1.0.4" +__version__ = "1.1.0" __license__ = "Apache License 2.0" diff --git a/awswrangler/_data_types.py b/awswrangler/_data_types.py index 62928e816..fac82a37b 100644 --- a/awswrangler/_data_types.py +++ b/awswrangler/_data_types.py @@ -1,8 +1,9 @@ """Internal (private) Data Types Module.""" import logging +import re from decimal import Decimal -from typing import Dict, List, Optional, Tuple +from typing import Any, Dict, List, Match, Optional, Sequence, Tuple import pandas as pd # type: ignore import pyarrow as pa # type: ignore @@ -139,8 +140,10 @@ def pyarrow2athena(dtype: pa.DataType) -> str: # pylint: disable=too-many-branc return f"decimal({dtype.precision},{dtype.scale})" if pa.types.is_list(dtype): return f"array<{pyarrow2athena(dtype=dtype.value_type)}>" - if pa.types.is_struct(dtype): # pragma: no cover - return f"struct<{', '.join([f'{f.name}: {pyarrow2athena(dtype=f.type)}' for f in dtype])}>" + if pa.types.is_struct(dtype): + return f"struct<{', '.join([f'{f.name}:{pyarrow2athena(dtype=f.type)}' for f in dtype])}>" + if pa.types.is_map(dtype): # pragma: no cover + return f"map<{pyarrow2athena(dtype=dtype.key_type)},{pyarrow2athena(dtype=dtype.item_type)}>" if dtype == pa.null(): raise exceptions.UndetectedType("We can not infer the data type from an entire null object column") raise exceptions.UnsupportedType(f"Unsupported Pyarrow type: {dtype}") # pragma: no cover @@ -167,7 +170,7 @@ def pyarrow2pandas_extension( # pylint: disable=too-many-branches,too-many-retu def pyarrow2sqlalchemy( # pylint: disable=too-many-branches,too-many-return-statements dtype: pa.DataType, db_type: str -) -> VisitableType: +) -> Optional[VisitableType]: """Pyarrow to Athena data types conversion.""" if pa.types.is_int8(dtype): return sqlalchemy.types.SmallInteger @@ -207,14 +210,14 @@ def pyarrow2sqlalchemy( # pylint: disable=too-many-branches,too-many-return-sta return sqlalchemy.types.Date if pa.types.is_binary(dtype): if db_type == "redshift": - raise exceptions.UnsupportedType(f"Binary columns are not supported for Redshift.") # pragma: no cover + raise exceptions.UnsupportedType("Binary columns are not supported for Redshift.") # pragma: no cover return sqlalchemy.types.Binary if pa.types.is_decimal(dtype): return sqlalchemy.types.Numeric(precision=dtype.precision, scale=dtype.scale) if pa.types.is_dictionary(dtype): return pyarrow2sqlalchemy(dtype=dtype.value_type, db_type=db_type) if dtype == pa.null(): # pragma: no cover - raise exceptions.UndetectedType("We can not infer the data type from an entire null object column") + return None raise exceptions.UnsupportedType(f"Unsupported Pyarrow type: {dtype}") # pragma: no cover @@ -243,12 +246,23 @@ def pyarrow_types_from_pandas( else: cols.append(name) - # Filling cols_dtypes and indexes + # Filling cols_dtypes + for col in cols: + _logger.debug("Inferring PyArrow type from column: %s", col) + try: + schema: pa.Schema = pa.Schema.from_pandas(df=df[[col]], preserve_index=False) + except pa.ArrowInvalid as ex: # pragma: no cover + cols_dtypes[col] = process_not_inferred_dtype(ex) + else: + cols_dtypes[col] = schema.field(col).type + + # Filling indexes indexes: List[str] = [] - for field in pa.Schema.from_pandas(df=df[cols], preserve_index=index): - name = str(field.name) - cols_dtypes[name] = field.type - if (name not in df.columns) and (index is True): + if index is True: + for field in pa.Schema.from_pandas(df=df[[]], preserve_index=True): + name = str(field.name) + _logger.debug("Inferring PyArrow type from index: %s", name) + cols_dtypes[name] = field.type indexes.append(name) # Merging Index @@ -257,10 +271,43 @@ def pyarrow_types_from_pandas( # Filling schema columns_types: Dict[str, pa.DataType] columns_types = {n: cols_dtypes[n] for n in sorted_cols} - _logger.debug(f"columns_types: {columns_types}") + _logger.debug("columns_types: %s", columns_types) return columns_types +def process_not_inferred_dtype(ex: pa.ArrowInvalid) -> pa.DataType: + """Infer data type from PyArrow inference exception.""" + ex_str = str(ex) + _logger.debug("PyArrow was not able to infer data type:\n%s", ex_str) + match: Optional[Match] = re.search( + pattern="Could not convert (.*) with type (.*): did not recognize " + "Python value type when inferring an Arrow data type", + string=ex_str, + ) + if match is None: + raise ex # pragma: no cover + groups: Optional[Sequence[str]] = match.groups() + if groups is None: + raise ex # pragma: no cover + if len(groups) != 2: + raise ex # pragma: no cover + _logger.debug("groups: %s", groups) + type_str: str = groups[1] + if type_str == "UUID": + return pa.string() + raise ex # pragma: no cover + + +def process_not_inferred_array(ex: pa.ArrowInvalid, values: Any) -> pa.Array: + """Infer `pyarrow.array` from PyArrow inference exception.""" + dtype = process_not_inferred_dtype(ex=ex) + if dtype == pa.string(): + array: pa.Array = pa.array(obj=[str(x) for x in values], type=dtype, safe=True) + else: + raise ex # pragma: no cover + return array + + def athena_types_from_pandas( df: pd.DataFrame, index: bool, dtype: Optional[Dict[str, str]] = None, index_left: bool = False ) -> Dict[str, str]: @@ -275,7 +322,7 @@ def athena_types_from_pandas( athena_columns_types[k] = casts[k] else: athena_columns_types[k] = pyarrow2athena(dtype=v) - _logger.debug(f"athena_columns_types: {athena_columns_types}") + _logger.debug("athena_columns_types: %s", athena_columns_types) return athena_columns_types @@ -315,7 +362,7 @@ def pyarrow_schema_from_pandas( if (k in df.columns) and (k not in ignore): columns_types[k] = athena2pyarrow(v) columns_types = {k: v for k, v in columns_types.items() if v is not None} - _logger.debug(f"columns_types: {columns_types}") + _logger.debug("columns_types: %s", columns_types) return pa.schema(fields=columns_types) @@ -324,11 +371,11 @@ def athena_types_from_pyarrow_schema( ) -> Tuple[Dict[str, str], Optional[Dict[str, str]]]: """Extract the related Athena data types from any PyArrow Schema considering possible partitions.""" columns_types: Dict[str, str] = {str(f.name): pyarrow2athena(dtype=f.type) for f in schema} - _logger.debug(f"columns_types: {columns_types}") + _logger.debug("columns_types: %s", columns_types) partitions_types: Optional[Dict[str, str]] = None if partitions is not None: partitions_types = {p.name: pyarrow2athena(p.dictionary.type) for p in partitions} - _logger.debug(f"partitions_types: {partitions_types}") + _logger.debug("partitions_types: %s", partitions_types) return columns_types, partitions_types @@ -372,7 +419,7 @@ def sqlalchemy_types_from_pandas( df: pd.DataFrame, db_type: str, dtype: Optional[Dict[str, VisitableType]] = None ) -> Dict[str, VisitableType]: """Extract the related SQLAlchemy data types from any Pandas DataFrame.""" - casts: Dict[str, VisitableType] = dtype if dtype else {} + casts: Dict[str, VisitableType] = dtype if dtype is not None else {} pa_columns_types: Dict[str, Optional[pa.DataType]] = pyarrow_types_from_pandas( df=df, index=False, ignore_cols=list(casts.keys()) ) @@ -382,5 +429,5 @@ def sqlalchemy_types_from_pandas( sqlalchemy_columns_types[k] = casts[k] else: sqlalchemy_columns_types[k] = pyarrow2sqlalchemy(dtype=v, db_type=db_type) - _logger.debug(f"sqlalchemy_columns_types: {sqlalchemy_columns_types}") + _logger.debug("sqlalchemy_columns_types: %s", sqlalchemy_columns_types) return sqlalchemy_columns_types diff --git a/awswrangler/_utils.py b/awswrangler/_utils.py index 21a27d37e..df168bdb9 100644 --- a/awswrangler/_utils.py +++ b/awswrangler/_utils.py @@ -166,3 +166,16 @@ def ensure_postgresql_casts(): def get_directory(path: str) -> str: """Extract directory path.""" return path.rsplit(sep="/", maxsplit=1)[0] + "/" + + +def get_account_id(boto3_session: Optional[boto3.Session] = None) -> str: + """Get Account ID.""" + session: boto3.Session = ensure_session(session=boto3_session) + return client(service_name="sts", session=session).get_caller_identity().get("Account") + + +def get_region_from_subnet(subnet_id: str, boto3_session: Optional[boto3.Session] = None) -> str: + """Extract region from Subnet ID.""" + session: boto3.Session = ensure_session(session=boto3_session) + client_ec2: boto3.client = client(service_name="ec2", session=session) + return client_ec2.describe_subnets(SubnetIds=[subnet_id])["Subnets"][0]["AvailabilityZone"][:9] diff --git a/awswrangler/athena.py b/awswrangler/athena.py index 1933606ad..671dabd42 100644 --- a/awswrangler/athena.py +++ b/awswrangler/athena.py @@ -2,6 +2,7 @@ import csv import logging +import pprint import time from decimal import Decimal from typing import Any, Dict, Iterator, List, Optional, Tuple, Union @@ -68,7 +69,7 @@ def create_athena_bucket(boto3_session: Optional[boto3.Session] = None) -> str: """ session: boto3.Session = _utils.ensure_session(session=boto3_session) - account_id: str = _utils.client(service_name="sts", session=session).get_caller_identity().get("Account") + account_id: str = _utils.get_account_id(boto3_session=session) region_name: str = str(session.region_name).lower() s3_output = f"s3://aws-athena-query-results-{account_id}-{region_name}/" s3_resource = session.resource("s3") @@ -120,19 +121,49 @@ def start_query_execution( >>> query_exec_id = wr.athena.start_query_execution(sql='...', database='...') """ + session: boto3.Session = _utils.ensure_session(session=boto3_session) + wg_config: Dict[str, Union[bool, Optional[str]]] = _get_workgroup_config(session=session, workgroup=workgroup) + return _start_query_execution( + sql=sql, + wg_config=wg_config, + database=database, + s3_output=s3_output, + workgroup=workgroup, + encryption=encryption, + kms_key=kms_key, + boto3_session=session, + ) + + +def _start_query_execution( + sql: str, + wg_config: Dict[str, Union[Optional[bool], Optional[str]]], + database: Optional[str] = None, + s3_output: Optional[str] = None, + workgroup: Optional[str] = None, + encryption: Optional[str] = None, + kms_key: Optional[str] = None, + boto3_session: Optional[boto3.Session] = None, +) -> str: args: Dict[str, Any] = {"QueryString": sql} session: boto3.Session = _utils.ensure_session(session=boto3_session) # s3_output - if s3_output is None: # pragma: no cover - s3_output = create_athena_bucket(boto3_session=session) - args["ResultConfiguration"] = {"OutputLocation": s3_output} + args["ResultConfiguration"] = { + "OutputLocation": _get_s3_output(s3_output=s3_output, wg_config=wg_config, boto3_session=session) + } # encryption - if encryption is not None: - args["ResultConfiguration"]["EncryptionConfiguration"] = {"EncryptionOption": encryption} - if kms_key is not None: - args["ResultConfiguration"]["EncryptionConfiguration"]["KmsKey"] = kms_key + if wg_config["enforced"] is True: + if wg_config["encryption"] is not None: + args["ResultConfiguration"]["EncryptionConfiguration"] = {"EncryptionOption": wg_config["encryption"]} + if wg_config["kms_key"] is not None: + args["ResultConfiguration"]["EncryptionConfiguration"]["KmsKey"] = wg_config["kms_key"] + else: + if encryption is not None: + args["ResultConfiguration"]["EncryptionConfiguration"] = {"EncryptionOption": encryption} + if kms_key is not None: + args["ResultConfiguration"]["EncryptionConfiguration"]["KmsKey"] = kms_key # database if database is not None: @@ -143,10 +174,25 @@ def start_query_execution( args["WorkGroup"] = workgroup client_athena: boto3.client = _utils.client(service_name="athena", session=session) + _logger.debug("args: \n%s", pprint.pformat(args)) response = client_athena.start_query_execution(**args) return response["QueryExecutionId"] +def _get_s3_output( + s3_output: Optional[str], wg_config: Dict[str, Union[bool, Optional[str]]], boto3_session: boto3.Session +) -> str: + if s3_output is None: + _s3_output: Optional[str] = wg_config["s3_output"] # type: ignore + if _s3_output is not None: + s3_output = _s3_output + else: + s3_output = create_athena_bucket(boto3_session=boto3_session) + elif wg_config["enforced"] is True: + s3_output = wg_config["s3_output"] # type: ignore + return s3_output + + def wait_query(query_execution_id: str, boto3_session: Optional[boto3.Session] = None) -> Dict[str, Any]: """Wait for the query end. @@ -176,8 +222,8 @@ def wait_query(query_execution_id: str, boto3_session: Optional[boto3.Session] = time.sleep(_QUERY_WAIT_POLLING_DELAY) response = client_athena.get_query_execution(QueryExecutionId=query_execution_id) state = response["QueryExecution"]["Status"]["State"] - _logger.debug(f"state: {state}") - _logger.debug(f"StateChangeReason: {response['QueryExecution']['Status'].get('StateChangeReason')}") + _logger.debug("state: %s", state) + _logger.debug("StateChangeReason: %s", response["QueryExecution"]["Status"].get("StateChangeReason")) if state == "FAILED": raise exceptions.QueryFailed(response["QueryExecution"]["Status"].get("StateChangeReason")) if state == "CANCELLED": @@ -265,7 +311,7 @@ def _get_query_metadata( cols_types: Dict[str, str] = get_query_columns_types( query_execution_id=query_execution_id, boto3_session=boto3_session ) - _logger.debug(f"cols_types: {cols_types}") + _logger.debug("cols_types: %s", cols_types) dtype: Dict[str, str] = {} parse_timestamps: List[str] = [] parse_dates: List[str] = [] @@ -298,11 +344,11 @@ def _get_query_metadata( converters[col_name] = lambda x: Decimal(str(x)) if str(x) not in ("", "none", " ", "") else None else: dtype[col_name] = pandas_type - _logger.debug(f"dtype: {dtype}") - _logger.debug(f"parse_timestamps: {parse_timestamps}") - _logger.debug(f"parse_dates: {parse_dates}") - _logger.debug(f"converters: {converters}") - _logger.debug(f"binaries: {binaries}") + _logger.debug("dtype: %s", dtype) + _logger.debug("parse_timestamps: %s", parse_timestamps) + _logger.debug("parse_dates: %s", parse_dates) + _logger.debug("converters: %s", converters) + _logger.debug("binaries: %s", binaries) return dtype, parse_timestamps, parse_dates, converters, binaries @@ -324,16 +370,18 @@ def _fix_csv_types(df: pd.DataFrame, parse_dates: List[str], binaries: List[str] return df -def read_sql_query( # pylint: disable=too-many-branches,too-many-locals +def read_sql_query( # pylint: disable=too-many-branches,too-many-locals,too-many-return-statements,too-many-statements sql: str, database: str, ctas_approach: bool = True, categories: List[str] = None, - chunksize: Optional[int] = None, + chunksize: Optional[Union[int, bool]] = None, s3_output: Optional[str] = None, workgroup: Optional[str] = None, encryption: Optional[str] = None, kms_key: Optional[str] = None, + keep_files: bool = True, + ctas_temp_table_name: Optional[str] = None, use_threads: bool = True, boto3_session: Optional[boto3.Session] = None, ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]: @@ -355,18 +403,31 @@ def read_sql_query( # pylint: disable=too-many-branches,too-many-locals Note ---- - If `chunksize` is passed, then a Generator of DataFrames is returned. + Valid encryption modes: [None, 'SSE_S3', 'SSE_KMS']. - Note - ---- - If `ctas_approach` is True, `chunksize` will return non deterministic chunks sizes, - but it still useful to overcome memory limitation. + `P.S. 'CSE_KMS' is not supported.` Note ---- Create the default Athena bucket if it doesn't exist and s3_output is None. + (E.g. s3://aws-athena-query-results-ACCOUNT-REGION/) + Note + ---- + ``Batching`` (`chunksize` argument) (Memory Friendly): + + Will anable the function to return a Iterable of DataFrames instead of a regular DataFrame. + + There are two batching strategies on Wrangler: + + - If **chunksize=True**, a new DataFrame will be returned for each file in the query result. + + - If **chunked=INTEGER**, Wrangler will iterate on the data by number of rows igual the received INTEGER. + + `P.S.` `chunksize=True` if faster and uses less memory while `chunksize=INTEGER` is more precise + in number of rows for each Dataframe. + Note ---- In case of `use_threads=True` the number of threads that will be spawned will be get from os.cpu_count(). @@ -383,16 +444,24 @@ def read_sql_query( # pylint: disable=too-many-branches,too-many-locals categories: List[str], optional List of columns names that should be returned as pandas.Categorical. Recommended for memory restricted environments. - chunksize: int, optional - If specified, return an generator where chunksize is the number of rows to include in each chunk. + chunksize : Union[int, bool], optional + If passed will split the data in a Iterable of DataFrames (Memory friendly). + If `True` wrangler will iterate on the data by files in the most efficient way without guarantee of chunksize. + If an `INTEGER` is passed Wrangler will iterate on the data by number of rows igual the received INTEGER. s3_output : str, optional AWS S3 path. workgroup : str, optional Athena workgroup. encryption : str, optional - None, 'SSE_S3', 'SSE_KMS', 'CSE_KMS'. + Valid values: [None, 'SSE_S3', 'SSE_KMS']. Notice: 'CSE_KMS' is not supported. kms_key : str, optional - For SSE-KMS and CSE-KMS , this is the KMS key ARN or ID. + For SSE-KMS, this is the KMS key ARN or ID. + keep_files : bool + Should Wrangler delete or keep the staging files produced by Athena? + ctas_temp_table_name : str, optional + The name of the temporary table and also the directory name on S3 where the CTAS result is stored. + If None, it will use the follow random pattern: `f"temp_table_{pyarrow.compat.guid()}"`. + On S3 this directory will be under under the pattern: `f"{s3_output}/{ctas_temp_table_name}/"`. 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. @@ -411,31 +480,30 @@ def read_sql_query( # pylint: disable=too-many-branches,too-many-locals """ session: boto3.Session = _utils.ensure_session(session=boto3_session) - wg_s3_output, _, _ = _ensure_workgroup(session=session, workgroup=workgroup) - if s3_output is None: - if wg_s3_output is None: - _s3_output: str = create_athena_bucket(boto3_session=session) - else: - _s3_output = wg_s3_output - else: - _s3_output = s3_output + wg_config: Dict[str, Union[bool, Optional[str]]] = _get_workgroup_config(session=session, workgroup=workgroup) + _s3_output: str = _get_s3_output(s3_output=s3_output, wg_config=wg_config, boto3_session=session) _s3_output = _s3_output[:-1] if _s3_output[-1] == "/" else _s3_output name: str = "" if ctas_approach is True: - name = f"temp_table_{pa.compat.guid()}" + if ctas_temp_table_name is not None: + name = catalog.sanitize_table_name(ctas_temp_table_name) + else: + name = f"temp_table_{pa.compat.guid()}" path: str = f"{_s3_output}/{name}" + ext_location: str = "\n" if wg_config["enforced"] is True else f",\n external_location = '{path}'\n" sql = ( f"CREATE TABLE {name}\n" f"WITH(\n" f" format = 'Parquet',\n" - f" parquet_compression = 'SNAPPY',\n" - f" external_location = '{path}'\n" + f" parquet_compression = 'SNAPPY'" + f"{ext_location}" f") AS\n" f"{sql}" ) - _logger.debug(f"sql: {sql}") - query_id: str = start_query_execution( + _logger.debug("sql: %s", sql) + query_id: str = _start_query_execution( sql=sql, + wg_config=wg_config, database=database, s3_output=_s3_output, workgroup=workgroup, @@ -443,36 +511,48 @@ def read_sql_query( # pylint: disable=too-many-branches,too-many-locals kms_key=kms_key, boto3_session=session, ) - _logger.debug(f"query_id: {query_id}") + _logger.debug("query_id: %s", query_id) query_response: Dict[str, Any] = wait_query(query_execution_id=query_id, boto3_session=session) if query_response["QueryExecution"]["Status"]["State"] in ["FAILED", "CANCELLED"]: # pragma: no cover reason: str = query_response["QueryExecution"]["Status"]["StateChangeReason"] message_error: str = f"Query error: {reason}" raise exceptions.AthenaQueryError(message_error) - dfs: Union[pd.DataFrame, Iterator[pd.DataFrame]] + ret: Union[pd.DataFrame, Iterator[pd.DataFrame]] if ctas_approach is True: catalog.delete_table_if_exists(database=database, table=name, boto3_session=session) manifest_path: str = f"{_s3_output}/tables/{query_id}-manifest.csv" + metadata_path: str = f"{_s3_output}/tables/{query_id}.metadata" + _logger.debug("manifest_path: %s", manifest_path) + _logger.debug("metadata_path: %s", metadata_path) + s3.wait_objects_exist(paths=[manifest_path, metadata_path], use_threads=False, boto3_session=session) paths: List[str] = _extract_ctas_manifest_paths(path=manifest_path, boto3_session=session) - chunked: bool = chunksize is not None - _logger.debug(f"chunked: {chunked}") + chunked: Union[bool, int] = False if chunksize is None else chunksize + _logger.debug("chunked: %s", chunked) if not paths: if chunked is False: - dfs = pd.DataFrame() - else: - dfs = _utils.empty_generator() - else: - s3.wait_objects_exist(paths=paths, use_threads=False, boto3_session=session) - dfs = s3.read_parquet( - path=paths, use_threads=use_threads, boto3_session=session, chunked=chunked, categories=categories - ) - return dfs + return pd.DataFrame() + return _utils.empty_generator() + s3.wait_objects_exist(paths=paths, use_threads=False, boto3_session=session) + ret = s3.read_parquet( + path=paths, use_threads=use_threads, boto3_session=session, chunked=chunked, categories=categories + ) + paths_delete: List[str] = paths + [manifest_path, metadata_path] + _logger.debug(type(ret)) + if chunked is False: + if keep_files is False: + s3.delete_objects(path=paths_delete, use_threads=use_threads, boto3_session=session) + return ret + if keep_files is False: + return _delete_after_iterate(dfs=ret, paths=paths_delete, use_threads=use_threads, boto3_session=session) + return ret dtype, parse_timestamps, parse_dates, converters, binaries = _get_query_metadata( query_execution_id=query_id, categories=categories, boto3_session=session ) path = f"{_s3_output}/{query_id}.csv" s3.wait_objects_exist(paths=[path], use_threads=False, boto3_session=session) - _logger.debug(f"Start CSV reading from {path}") + _logger.debug("Start CSV reading from %s", path) + _chunksize: Optional[int] = chunksize if isinstance(chunksize, int) else None + _logger.debug("_chunksize: %s", _chunksize) ret = s3.read_csv( path=[path], dtype=dtype, @@ -481,16 +561,32 @@ def read_sql_query( # pylint: disable=too-many-branches,too-many-locals quoting=csv.QUOTE_ALL, keep_default_na=False, na_values=[""], - chunksize=chunksize, + chunksize=_chunksize, skip_blank_lines=False, use_threads=False, boto3_session=session, ) _logger.debug("Start type casting...") - if chunksize is None: - return _fix_csv_types(df=ret, parse_dates=parse_dates, binaries=binaries) _logger.debug(type(ret)) - return _fix_csv_types_generator(dfs=ret, parse_dates=parse_dates, binaries=binaries) + if chunksize is None: + df = _fix_csv_types(df=ret, parse_dates=parse_dates, binaries=binaries) + if keep_files is False: + s3.delete_objects(path=[path, f"{path}.metadata"], use_threads=use_threads, boto3_session=session) + return df + dfs = _fix_csv_types_generator(dfs=ret, parse_dates=parse_dates, binaries=binaries) + if keep_files is False: + return _delete_after_iterate( + dfs=dfs, paths=[path, f"{path}.metadata"], use_threads=use_threads, boto3_session=session + ) + return dfs + + +def _delete_after_iterate( + dfs: Iterator[pd.DataFrame], paths: List[str], use_threads: bool, boto3_session: boto3.Session +) -> Iterator[pd.DataFrame]: + for df in dfs: + yield df + s3.delete_objects(path=paths, use_threads=use_threads, boto3_session=boto3_session) def stop_query_execution(query_execution_id: str, boto3_session: Optional[boto3.Session] = None) -> None: @@ -545,19 +641,27 @@ def get_work_group(workgroup: str, boto3_session: Optional[boto3.Session] = None return client_athena.get_work_group(WorkGroup=workgroup) -def _ensure_workgroup( +def _get_workgroup_config( session: boto3.Session, workgroup: Optional[str] = None -) -> Tuple[Optional[str], Optional[str], Optional[str]]: +) -> Dict[str, Union[bool, Optional[str]]]: if workgroup is not None: res: Dict[str, Any] = get_work_group(workgroup=workgroup, boto3_session=session) + enforced: bool = res["WorkGroup"]["Configuration"]["EnforceWorkGroupConfiguration"] config: Dict[str, Any] = res["WorkGroup"]["Configuration"]["ResultConfiguration"] wg_s3_output: Optional[str] = config.get("OutputLocation") encrypt_config: Optional[Dict[str, str]] = config.get("EncryptionConfiguration") wg_encryption: Optional[str] = None if encrypt_config is None else encrypt_config.get("EncryptionOption") wg_kms_key: Optional[str] = None if encrypt_config is None else encrypt_config.get("KmsKey") else: - wg_s3_output, wg_encryption, wg_kms_key = None, None, None - return wg_s3_output, wg_encryption, wg_kms_key + enforced, wg_s3_output, wg_encryption, wg_kms_key = False, None, None, None + wg_config: Dict[str, Union[bool, Optional[str]]] = { + "enforced": enforced, + "s3_output": wg_s3_output, + "encryption": wg_encryption, + "kms_key": wg_kms_key, + } + _logger.debug("wg_config: \n%s", pprint.pformat(wg_config)) + return wg_config def read_sql_table( @@ -565,11 +669,13 @@ def read_sql_table( database: str, ctas_approach: bool = True, categories: List[str] = None, - chunksize: Optional[int] = None, + chunksize: Optional[Union[int, bool]] = None, s3_output: Optional[str] = None, workgroup: Optional[str] = None, encryption: Optional[str] = None, kms_key: Optional[str] = None, + keep_files: bool = True, + ctas_temp_table_name: Optional[str] = None, use_threads: bool = True, boto3_session: Optional[boto3.Session] = None, ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]: @@ -591,18 +697,31 @@ def read_sql_table( Note ---- - If `chunksize` is passed, then a Generator of DataFrames is returned. + Valid encryption modes: [None, 'SSE_S3', 'SSE_KMS']. - Note - ---- - If `ctas_approach` is True, `chunksize` will return non deterministic chunks sizes, - but it still useful to overcome memory limitation. + `P.S. 'CSE_KMS' is not supported.` Note ---- Create the default Athena bucket if it doesn't exist and s3_output is None. + (E.g. s3://aws-athena-query-results-ACCOUNT-REGION/) + Note + ---- + ``Batching`` (`chunksize` argument) (Memory Friendly): + + Will anable the function to return a Iterable of DataFrames instead of a regular DataFrame. + + There are two batching strategies on Wrangler: + + - If **chunksize=True**, a new DataFrame will be returned for each file in the query result. + + - If **chunked=INTEGER**, Wrangler will iterate on the data by number of rows igual the received INTEGER. + + `P.S.` `chunksize=True` if faster and uses less memory while `chunksize=INTEGER` is more precise + in number of rows for each Dataframe. + Note ---- In case of `use_threads=True` the number of threads that will be spawned will be get from os.cpu_count(). @@ -619,8 +738,10 @@ def read_sql_table( categories: List[str], optional List of columns names that should be returned as pandas.Categorical. Recommended for memory restricted environments. - chunksize: int, optional - If specified, return an generator where chunksize is the number of rows to include in each chunk. + chunksize : Union[int, bool], optional + If passed will split the data in a Iterable of DataFrames (Memory friendly). + If `True` wrangler will iterate on the data by files in the most efficient way without guarantee of chunksize. + If an `INTEGER` is passed Wrangler will iterate on the data by number of rows igual the received INTEGER. s3_output : str, optional AWS S3 path. workgroup : str, optional @@ -629,6 +750,12 @@ def read_sql_table( None, 'SSE_S3', 'SSE_KMS', 'CSE_KMS'. kms_key : str, optional For SSE-KMS and CSE-KMS , this is the KMS key ARN or ID. + keep_files : bool + Should Wrangler delete or keep the staging files produced by Athena? + ctas_temp_table_name : str, optional + The name of the temporary table and also the directory name on S3 where the CTAS result is stored. + If None, it will use the follow random pattern: `f"temp_table_{pyarrow.compat.guid()}"`. + On S3 this directory will be under under the pattern: `f"{s3_output}/{ctas_temp_table_name}/"`. 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. @@ -646,6 +773,7 @@ def read_sql_table( >>> df = wr.athena.read_sql_table(table='...', database='...') """ + table = catalog.sanitize_table_name(table=table) return read_sql_query( sql=f'SELECT * FROM "{table}"', database=database, @@ -656,6 +784,8 @@ def read_sql_table( workgroup=workgroup, encryption=encryption, kms_key=kms_key, + keep_files=keep_files, + ctas_temp_table_name=ctas_temp_table_name, use_threads=use_threads, boto3_session=boto3_session, ) diff --git a/awswrangler/catalog.py b/awswrangler/catalog.py index 8a53d4370..9ef55066c 100644 --- a/awswrangler/catalog.py +++ b/awswrangler/catalog.py @@ -93,6 +93,7 @@ def create_parquet_table( parameters: Optional[Dict[str, str]] = None, columns_comments: Optional[Dict[str, str]] = None, mode: str = "overwrite", + catalog_versioning: bool = False, boto3_session: Optional[boto3.Session] = None, ) -> None: """Create a Parquet Table (Metadata Only) in the AWS Glue Catalog. @@ -121,6 +122,8 @@ def create_parquet_table( Columns names and the related comments (e.g. {'col0': 'Column 0.', 'col1': 'Column 1.', 'col2': 'Partition.'}). mode: str 'overwrite' to recreate any possible existing table or 'append' to keep any possible existing table. + catalog_versioning : bool + If True and `mode="overwrite"`, creates an archived version of the table catalog before updating it. boto3_session : boto3.Session(), optional Boto3 Session. The default boto3 session will be used if boto3_session receive None. @@ -157,6 +160,7 @@ def create_parquet_table( parameters=parameters, columns_comments=columns_comments, mode=mode, + catalog_versioning=catalog_versioning, boto3_session=boto3_session, table_input=table_input, ) @@ -766,7 +770,7 @@ def drop_duplicated_columns(df: pd.DataFrame) -> pd.DataFrame: duplicated_cols = df.columns.duplicated() duplicated_cols_names: List[str] = list(df.columns[duplicated_cols]) if len(duplicated_cols_names) > 0: - _logger.warning(f"Dropping repeated columns: {duplicated_cols_names}") + _logger.warning("Dropping repeated columns: %s", duplicated_cols_names) return df.loc[:, ~duplicated_cols] @@ -865,6 +869,7 @@ def create_csv_table( parameters: Optional[Dict[str, str]] = None, columns_comments: Optional[Dict[str, str]] = None, mode: str = "overwrite", + catalog_versioning: bool = False, sep: str = ",", boto3_session: Optional[boto3.Session] = None, ) -> None: @@ -884,16 +889,18 @@ def create_csv_table( Dictionary with keys as column names and vales as data types (e.g. {'col0': 'bigint', 'col1': 'double'}). partitions_types: Dict[str, str], optional Dictionary with keys as partition names and values as data types (e.g. {'col2': 'date'}). - compression: str, optional + compression : str, optional Compression style (``None``, ``gzip``, etc). - description: str, optional + description : str, optional Table description - parameters: Dict[str, str], optional + parameters : Dict[str, str], optional Key/value pairs to tag the table. columns_comments: Dict[str, str], optional Columns names and the related comments (e.g. {'col0': 'Column 0.', 'col1': 'Column 1.', 'col2': 'Partition.'}). - mode: str + mode : str 'overwrite' to recreate any possible axisting table or 'append' to keep any possible axisting table. + catalog_versioning : bool + If True and `mode="overwrite"`, creates an archived version of the table catalog before updating it. sep : str String of length 1. Field delimiter for the output file. boto3_session : boto3.Session(), optional @@ -937,6 +944,7 @@ def create_csv_table( parameters=parameters, columns_comments=columns_comments, mode=mode, + catalog_versioning=catalog_versioning, boto3_session=boto3_session, table_input=table_input, ) @@ -949,6 +957,7 @@ def _create_table( parameters: Optional[Dict[str, str]], columns_comments: Optional[Dict[str, str]], mode: str, + catalog_versioning: bool, boto3_session: Optional[boto3.Session], table_input: Dict[str, Any], ): @@ -967,11 +976,15 @@ def _create_table( if name in columns_comments: par["Comment"] = columns_comments[name] session: boto3.Session = _utils.ensure_session(session=boto3_session) - - if mode == "overwrite": - delete_table_if_exists(database=database, table=table, boto3_session=session) client_glue: boto3.client = _utils.client(service_name="glue", session=session) - client_glue.create_table(DatabaseName=database, TableInput=table_input) + exist: bool = does_table_exist(database=database, table=table, boto3_session=session) + if mode not in ("overwrite", "append"): # pragma: no cover + raise exceptions.InvalidArgument(f"{mode} is not a valid mode. It must be 'overwrite' or 'append'.") + if (exist is True) and (mode == "overwrite"): + skip_archive: bool = not catalog_versioning + client_glue.update_table(DatabaseName=database, TableInput=table_input, SkipArchive=skip_archive) + elif exist is False: + client_glue.create_table(DatabaseName=database, TableInput=table_input) def _csv_table_definition( diff --git a/awswrangler/cloudwatch.py b/awswrangler/cloudwatch.py index e0a01f066..c36fab70b 100644 --- a/awswrangler/cloudwatch.py +++ b/awswrangler/cloudwatch.py @@ -56,11 +56,11 @@ def start_query( ... ) """ - _logger.debug(f"log_group_names: {log_group_names}") + _logger.debug("log_group_names: %s", log_group_names) start_timestamp: int = int(1000 * start_time.timestamp()) end_timestamp: int = int(1000 * end_time.timestamp()) - _logger.debug(f"start_timestamp: {start_timestamp}") - _logger.debug(f"end_timestamp: {end_timestamp}") + _logger.debug("start_timestamp: %s", start_timestamp) + _logger.debug("end_timestamp: %s", end_timestamp) args: Dict[str, Any] = { "logGroupNames": log_group_names, "startTime": start_timestamp, @@ -109,7 +109,7 @@ def wait_query(query_id: str, boto3_session: Optional[boto3.Session] = None) -> time.sleep(_QUERY_WAIT_POLLING_DELAY) response = client_logs.get_query_results(queryId=query_id) status = response["status"] - _logger.debug(f"status: {status}") + _logger.debug("status: %s", status) if status == "Failed": # pragma: no cover raise exceptions.QueryFailed(f"query ID: {query_id}") if status == "Cancelled": diff --git a/awswrangler/db.py b/awswrangler/db.py index 491fe7784..f5e90c78e 100644 --- a/awswrangler/db.py +++ b/awswrangler/db.py @@ -2,6 +2,7 @@ import json import logging +import time from typing import Any, Dict, Iterator, List, Optional, Tuple, Union from urllib.parse import quote_plus @@ -91,7 +92,16 @@ def to_sql(df: pd.DataFrame, con: sqlalchemy.engine.Engine, **pandas_kwargs) -> ) pandas_kwargs["dtype"] = dtypes pandas_kwargs["con"] = con - df.to_sql(**pandas_kwargs) + max_attempts: int = 3 + for attempt in range(max_attempts): + try: + df.to_sql(**pandas_kwargs) + except sqlalchemy.exc.InternalError as ex: # pragma: no cover + if attempt == (max_attempts - 1): + raise ex + time.sleep(1) + else: + break def read_sql_query( @@ -155,29 +165,15 @@ def read_sql_query( ... ) """ - 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()" - ) + _validate_engine(con=con) with con.connect() as _con: args = _convert_params(sql, params) cursor = _con.execute(*args) if chunksize is None: return _records2df(records=cursor.fetchall(), cols_names=cursor.keys(), index=index_col, dtype=dtype) - return _iterate_cursor(cursor=cursor, chunksize=chunksize, index=index_col, dtype=dtype) - - -def _iterate_cursor( - cursor, chunksize: int, index: Optional[Union[str, List[str]]], dtype: Optional[Dict[str, pa.DataType]] = None -) -> Iterator[pd.DataFrame]: - while True: - records = cursor.fetchmany(chunksize) - if not records: - break - df: pd.DataFrame = _records2df(records=records, cols_names=cursor.keys(), index=index, dtype=dtype) - yield df + return _iterate_cursor( + cursor=cursor, chunksize=chunksize, cols_names=cursor.keys(), index=index_col, dtype=dtype + ) def _records2df( @@ -189,7 +185,10 @@ def _records2df( arrays: List[pa.Array] = [] for col_values, col_name in zip(tuple(zip(*records)), cols_names): # Transposing if (dtype is None) or (col_name not in dtype): - array: pa.Array = pa.array(obj=col_values, safe=True) # Creating Arrow array + try: + array: pa.Array = pa.array(obj=col_values, safe=True) # Creating Arrow array + except pa.ArrowInvalid as ex: + array = _data_types.process_not_inferred_array(ex, values=col_values) # Creating Arrow array else: array = pa.array(obj=col_values, type=dtype[col_name], safe=True) # Creating Arrow array with dtype arrays.append(array) @@ -207,6 +206,20 @@ def _records2df( return df +def _iterate_cursor( + cursor: Any, + chunksize: int, + cols_names: List[str], + index: Optional[Union[str, List[str]]], + dtype: Optional[Dict[str, pa.DataType]] = None, +) -> Iterator[pd.DataFrame]: + while True: + records = cursor.fetchmany(chunksize) + if not records: + break + yield _records2df(records=records, cols_names=cols_names, index=index, dtype=dtype) + + def _convert_params(sql: str, params: Optional[Union[List, Tuple, Dict]]) -> List[Any]: args: List[Any] = [sql] if params is not None: @@ -646,7 +659,7 @@ def copy_files_to_redshift( # pylint: disable=too-many-locals,too-many-argument athena_types, _ = s3.read_parquet_metadata( path=paths, dataset=False, use_threads=use_threads, boto3_session=session ) - _logger.debug(f"athena_types: {athena_types}") + _logger.debug("athena_types: %s", athena_types) redshift_types: Dict[str, str] = {} for col_name, col_type in athena_types.items(): length: int = _varchar_lengths[col_name] if col_name in _varchar_lengths else varchar_lengths_default @@ -680,7 +693,7 @@ def copy_files_to_redshift( # pylint: disable=too-many-locals,too-many-argument def _rs_upsert(con: Any, table: str, temp_table: str, schema: str, primary_keys: Optional[List[str]] = None) -> None: if not primary_keys: primary_keys = _rs_get_primary_keys(con=con, schema=schema, table=table) - _logger.debug(f"primary_keys: {primary_keys}") + _logger.debug("primary_keys: %s", primary_keys) if not primary_keys: # pragma: no cover raise exceptions.InvalidRedshiftPrimaryKeys() equals_clause: str = f"{table}.%s = {temp_table}.%s" @@ -735,7 +748,7 @@ def _rs_create_table( f"{distkey_str}" f"{sortkey_str}" ) - _logger.debug(f"Create table query:\n{sql}") + _logger.debug("Create table query:\n%s", sql) con.execute(sql) return table, schema @@ -746,7 +759,7 @@ def _rs_validate_parameters( if diststyle not in _RS_DISTSTYLES: raise exceptions.InvalidRedshiftDiststyle(f"diststyle must be in {_RS_DISTSTYLES}") cols = list(redshift_types.keys()) - _logger.debug(f"Redshift columns: {cols}") + _logger.debug("Redshift columns: %s", cols) if (diststyle == "KEY") and (not distkey): raise exceptions.InvalidRedshiftDistkey("You must pass a distkey if you intend to use KEY diststyle") if distkey and distkey not in cols: @@ -775,13 +788,13 @@ def _rs_copy( sql: str = ( f"COPY {table_name} FROM '{manifest_path}'\n" f"IAM_ROLE '{iam_role}'\n" "MANIFEST\n" "FORMAT AS PARQUET" ) - _logger.debug(f"copy query:\n{sql}") + _logger.debug("copy query:\n%s", sql) con.execute(sql) sql = "SELECT pg_last_copy_id() AS query_id" query_id: int = con.execute(sql).fetchall()[0][0] sql = f"SELECT COUNT(DISTINCT filename) as num_files_loaded " f"FROM STL_LOAD_COMMITS WHERE query = {query_id}" num_files_loaded: int = con.execute(sql).fetchall()[0][0] - _logger.debug(f"{num_files_loaded} files counted. {num_files} expected.") + _logger.debug("%s files counted. %s expected.", num_files_loaded, num_files) if num_files_loaded != num_files: # pragma: no cover raise exceptions.RedshiftLoadError( f"Redshift load rollbacked. {num_files_loaded} files counted. {num_files} expected." @@ -846,17 +859,17 @@ def write_redshift_copy_manifest( payload: str = json.dumps(manifest) bucket: str bucket, key = _utils.parse_path(manifest_path) - _logger.debug(f"payload: {payload}") + _logger.debug("payload: %s", payload) client_s3: boto3.client = _utils.client(service_name="s3", session=session) - _logger.debug(f"bucket: {bucket}") - _logger.debug(f"key: {key}") + _logger.debug("bucket: %s", bucket) + _logger.debug("key: %s", key) client_s3.put_object(Body=payload, Bucket=bucket, Key=key) return manifest def _rs_drop_table(con: Any, schema: str, table: str) -> None: sql = f"DROP TABLE IF EXISTS {schema}.{table}" - _logger.debug(f"Drop table query:\n{sql}") + _logger.debug("Drop table query:\n%s", sql) con.execute(sql) @@ -887,8 +900,11 @@ def unload_redshift( path: str, con: sqlalchemy.engine.Engine, iam_role: str, + region: Optional[str] = None, + max_file_size: Optional[float] = None, + kms_key_id: Optional[str] = None, categories: List[str] = None, - chunked: bool = False, + chunked: Union[bool, int] = False, keep_files: bool = False, use_threads: bool = True, boto3_session: Optional[boto3.Session] = None, @@ -906,6 +922,22 @@ def unload_redshift( https://docs.aws.amazon.com/redshift/latest/dg/r_UNLOAD.html + Note + ---- + ``Batching`` (`chunked` argument) (Memory Friendly): + + Will anable the function to return a Iterable of DataFrames instead of a regular DataFrame. + + There are two batching strategies on Wrangler: + + - If **chunked=True**, a new DataFrame will be returned for each file in your path/dataset. + + - If **chunked=INTEGER**, Wrangler will iterate on the data by number of rows igual the received INTEGER. + + `P.S.` `chunked=True` if faster and uses less memory while `chunked=INTEGER` is more precise + in number of rows for each Dataframe. + + Note ---- In case of `use_threads=True` the number of threads that will be spawned will be get from os.cpu_count(). @@ -921,14 +953,28 @@ def unload_redshift( wr.db.get_engine(), wr.db.get_redshift_temp_engine() or wr.catalog.get_engine() iam_role : str AWS IAM role with the related permissions. + region : str, optional + Specifies the AWS Region where the target Amazon S3 bucket is located. + REGION is required for UNLOAD to an Amazon S3 bucket that isn't in the + same AWS Region as the Amazon Redshift cluster. By default, UNLOAD + assumes that the target Amazon S3 bucket is located in the same AWS + Region as the Amazon Redshift cluster. + max_file_size : float, optional + Specifies the maximum size (MB) of files that UNLOAD creates in Amazon S3. + Specify a decimal value between 5.0 MB and 6200.0 MB. If None, the default + maximum file size is 6200.0 MB. + kms_key_id : str, optional + Specifies the key ID for an AWS Key Management Service (AWS KMS) key to be + used to encrypt data files on Amazon S3. categories: List[str], optional List of columns names that should be returned as pandas.Categorical. Recommended for memory restricted environments. keep_files : bool Should keep the stage files? - chunked : bool - If True will break the data in smaller DataFrames (Non deterministic number of lines). - Otherwise return a single DataFrame with the whole data. + chunked : Union[int, bool] + If passed will split the data in a Iterable of DataFrames (Memory friendly). + If `True` wrangler will iterate on the data by files in the most efficient way without guarantee of chunksize. + If an `INTEGER` is passed Wrangler will iterate on the data by number of rows igual the received INTEGER. 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. @@ -956,7 +1002,15 @@ def unload_redshift( """ session: boto3.Session = _utils.ensure_session(session=boto3_session) paths: List[str] = unload_redshift_to_files( - sql=sql, path=path, con=con, iam_role=iam_role, use_threads=use_threads, boto3_session=session + sql=sql, + path=path, + con=con, + iam_role=iam_role, + region=region, + max_file_size=max_file_size, + kms_key_id=kms_key_id, + use_threads=use_threads, + boto3_session=session, ) s3.wait_objects_exist(paths=paths, use_threads=False, boto3_session=session) if chunked is False: @@ -979,6 +1033,7 @@ def unload_redshift( return _read_parquet_iterator( paths=paths, categories=categories, + chunked=chunked, use_threads=use_threads, boto3_session=session, s3_additional_kwargs=s3_additional_kwargs, @@ -991,13 +1046,14 @@ def _read_parquet_iterator( keep_files: bool, use_threads: bool, categories: List[str] = None, + chunked: Union[bool, int] = True, boto3_session: Optional[boto3.Session] = None, s3_additional_kwargs: Optional[Dict[str, str]] = None, ) -> Iterator[pd.DataFrame]: dfs: Iterator[pd.DataFrame] = s3.read_parquet( path=paths, categories=categories, - chunked=True, + chunked=chunked, dataset=False, use_threads=use_threads, boto3_session=boto3_session, @@ -1013,6 +1069,9 @@ def unload_redshift_to_files( path: str, con: sqlalchemy.engine.Engine, iam_role: str, + region: Optional[str] = None, + max_file_size: Optional[float] = None, + kms_key_id: Optional[str] = None, use_threads: bool = True, manifest: bool = False, partition_cols: Optional[List] = None, @@ -1037,6 +1096,19 @@ def unload_redshift_to_files( wr.db.get_engine(), wr.db.get_redshift_temp_engine() or wr.catalog.get_engine() iam_role : str AWS IAM role with the related permissions. + region : str, optional + Specifies the AWS Region where the target Amazon S3 bucket is located. + REGION is required for UNLOAD to an Amazon S3 bucket that isn't in the + same AWS Region as the Amazon Redshift cluster. By default, UNLOAD + assumes that the target Amazon S3 bucket is located in the same AWS + Region as the Amazon Redshift cluster. + max_file_size : float, optional + Specifies the maximum size (MB) of files that UNLOAD creates in Amazon S3. + Specify a decimal value between 5.0 MB and 6200.0 MB. If None, the default + maximum file size is 6200.0 MB. + kms_key_id : str, optional + Specifies the key ID for an AWS Key Management Service (AWS KMS) key to be + used to encrypt data files on Amazon S3. 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. @@ -1067,23 +1139,39 @@ def unload_redshift_to_files( session: boto3.Session = _utils.ensure_session(session=boto3_session) s3.delete_objects(path=path, use_threads=use_threads, boto3_session=session) with con.connect() as _con: - partition_str: str = f"PARTITION BY ({','.join(partition_cols)})\n" if partition_cols else "" + partition_str: str = f"\nPARTITION BY ({','.join(partition_cols)})" if partition_cols else "" manifest_str: str = "\nmanifest" if manifest is True else "" + region_str: str = f"\nREGION AS '{region}'" if region is not None else "" + max_file_size_str: str = f"\nMAXFILESIZE AS {max_file_size} MB" if max_file_size is not None else "" + kms_key_id_str: str = f"\nKMS_KEY_ID '{kms_key_id}'" if kms_key_id is not None else "" sql = ( f"UNLOAD ('{sql}')\n" f"TO '{path}'\n" f"IAM_ROLE '{iam_role}'\n" "ALLOWOVERWRITE\n" "PARALLEL ON\n" - "ENCRYPTED\n" + "FORMAT PARQUET\n" + "ENCRYPTED" + f"{kms_key_id_str}" f"{partition_str}" - "FORMAT PARQUET" + f"{region_str}" + f"{max_file_size_str}" f"{manifest_str};" ) + _logger.debug("sql: \n%s", sql) _con.execute(sql) sql = "SELECT pg_last_query_id() AS query_id" query_id: int = _con.execute(sql).fetchall()[0][0] sql = f"SELECT path FROM STL_UNLOAD_LOG WHERE query={query_id};" paths = [x[0].replace(" ", "") for x in _con.execute(sql).fetchall()] - _logger.debug(f"paths: {paths}") + _logger.debug("paths: %s", paths) return paths + + +def _validate_engine(con: sqlalchemy.engine.Engine) -> None: # pragma: no cover + 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()" + ) diff --git a/awswrangler/emr.py b/awswrangler/emr.py index aee470621..5a93d752d 100644 --- a/awswrangler/emr.py +++ b/awswrangler/emr.py @@ -1,18 +1,75 @@ """EMR (Elastic Map Reduce) module.""" # pylint: disable=line-too-long -import json import logging +import pprint from typing import Any, Dict, List, Optional, Union import boto3 # type: ignore -from awswrangler import _utils +from awswrangler import _utils, exceptions _logger: logging.Logger = logging.getLogger(__name__) +def _get_default_logging_path( + subnet_id: Optional[str] = None, + account_id: Optional[str] = None, + region: Optional[str] = None, + boto3_session: Optional[boto3.Session] = None, +) -> str: + """Get EMR default logging path. + + E.g. "s3://aws-logs-{account_id}-{region}/elasticmapreduce/" + + Parameters + ---------- + subnet_id : str, optional + Subnet ID. If not provided, you must pass `account_id` and `region` explicit. + account_id: str, optional + Account ID. + region: str, optional + Region e.g. 'us-east-1' + boto3_session : boto3.Session(), optional + Boto3 Session. The default boto3 session will be used if boto3_session receive None. + + Returns + ------- + str + Default logging path. + E.g. "s3://aws-logs-{account_id}-{region}/elasticmapreduce/" + + Examples + -------- + >>> import awswrangler as wr + >>> state = wr.emr._get_default_logging_path("subnet-id") + 's3://aws-logs-{account_id}-{region}/elasticmapreduce/' + + """ + if account_id is None: + boto3_session = _utils.ensure_session(session=boto3_session) + _account_id: str = _utils.get_account_id(boto3_session=boto3_session) + else: + _account_id = account_id + if (region is None) and (subnet_id is not None): + boto3_session = _utils.ensure_session(session=boto3_session) + _region: str = _utils.get_region_from_subnet(subnet_id=subnet_id, boto3_session=boto3_session) + elif (region is None) and (subnet_id is None): + raise exceptions.InvalidArgumentCombination("You must pass region or subnet_id or both.") + else: + _region = region # type: ignore + return f"s3://aws-logs-{_account_id}-{_region}/elasticmapreduce/" + + def _build_cluster_args(**pars): # pylint: disable=too-many-branches,too-many-statements + account_id: str = _utils.get_account_id(boto3_session=pars["boto3_session"]) + region: str = _utils.get_region_from_subnet(subnet_id=pars["subnet_id"], boto3_session=pars["boto3_session"]) + + # S3 Logging path + if pars.get("logging_s3_path") is None: + pars["logging_s3_path"] = _get_default_logging_path( + subnet_id=None, account_id=account_id, region=region, boto3_session=pars["boto3_session"] + ) spark_env: Optional[Dict[str, str]] = None yarn_env: Optional[Dict[str, str]] = None @@ -20,25 +77,25 @@ def _build_cluster_args(**pars): # pylint: disable=too-many-branches,too-many-s if pars["spark_pyarrow"] is True: if pars["spark_defaults"] is None: - pars["spark_defaults"]: Dict[str, str] = {"spark.sql.execution.arrow.enabled": "true"} + pars["spark_defaults"] = {"spark.sql.execution.arrow.enabled": "true"} else: # pragma: no cover - pars["spark_defaults"]["spark.sql.execution.arrow.enabled"]: str = "true" + pars["spark_defaults"]["spark.sql.execution.arrow.enabled"] = "true" spark_env = {"ARROW_PRE_0_15_IPC_FORMAT": "1"} yarn_env = {"ARROW_PRE_0_15_IPC_FORMAT": "1"} livy_env = {"ARROW_PRE_0_15_IPC_FORMAT": "1"} if pars["python3"] is True: if spark_env is None: - spark_env: Dict[str, str] = {"PYSPARK_PYTHON": "/usr/bin/python3"} # pragma: no cover + spark_env = {"PYSPARK_PYTHON": "/usr/bin/python3"} # pragma: no cover else: - spark_env["PYSPARK_PYTHON"]: str = "/usr/bin/python3" + spark_env["PYSPARK_PYTHON"] = "/usr/bin/python3" if pars["spark_jars_path"] is not None: paths: str = ",".join(pars["spark_jars_path"]) if pars["spark_defaults"] is None: # pragma: no cover - pars["spark_defaults"]: Dict[str, str] = {"spark.jars": paths} + pars["spark_defaults"] = {"spark.jars": paths} else: - pars["spark_defaults"]["spark.jars"]: str = paths + pars["spark_defaults"]["spark.jars"] = paths args: Dict[str, Any] = { "Name": pars["cluster_name"], @@ -72,9 +129,32 @@ def _build_cluster_args(**pars): # pylint: disable=too-many-branches,too-many-s args["Instances"]["ServiceAccessSecurityGroup"] = pars["security_group_service_access"] # Configurations - args["Configurations"]: List[Dict[str, Any]] = [ + args["Configurations"] = [ {"Classification": "spark-log4j", "Properties": {"log4j.rootCategory": f"{pars['spark_log_level']}, console"}} ] + if pars["docker"] is True: + if pars.get("extra_registries") is None: + extra_registries: List[str] = [] + else: # pragma: no cover + extra_registries = pars["extra_registries"] + registries: str = f"local,centos,{account_id}.dkr.ecr.{region}.amazonaws.com,{','.join(extra_registries)}" + registries = registries[:-1] if registries.endswith(",") else registries + args["Configurations"].append( + { + "Classification": "container-executor", + "Properties": {}, + "Configurations": [ + { + "Classification": "docker", + "Properties": { + "docker.privileged-containers.registries": registries, + "docker.trusted.registries": registries, + }, + "Configurations": [], + } + ], + } + ) if spark_env is not None: args["Configurations"].append( { @@ -110,15 +190,11 @@ def _build_cluster_args(**pars): # pylint: disable=too-many-branches,too-many-s } ) if pars["hive_glue_catalog"] is True: - args["Configurations"].append( - { - "Classification": "hive-site", - "Properties": { - "hive.metastore.client.factory.class": "com.amazonaws.glue.catalog.metastore.AWSGlueDataCatalogHiveClientFactory" # noqa - }, - "Configurations": [], - } - ) + hive_conf: Optional[Dict[str, Any]] = {"Classification": "hive-site", "Properties": {}, "Configurations": []} + hive_conf["Properties"][ + "hive.metastore.client.factory.class" + ] = "com.amazonaws.glue.catalog.metastore.AWSGlueDataCatalogHiveClientFactory" + args["Configurations"].append(hive_conf) if pars["presto_glue_catalog"] is True: args["Configurations"].append( { @@ -147,20 +223,21 @@ def _build_cluster_args(**pars): # pylint: disable=too-many-branches,too-many-s "Properties": pars["spark_defaults"], } args["Configurations"].append(spark_defaults) + if pars.get("custom_classifications") is not None: + for c in pars["custom_classifications"]: + args["Configurations"].append(c) # Applications if pars["applications"]: - args["Applications"]: List[Dict[str, str]] = [{"Name": x} for x in pars["applications"]] + args["Applications"] = [{"Name": x} for x in pars["applications"]] # Bootstraps if pars["bootstraps_paths"]: # pragma: no cover - args["BootstrapActions"]: List[Dict] = [ - {"Name": x, "ScriptBootstrapAction": {"Path": x}} for x in pars["bootstraps_paths"] - ] + args["BootstrapActions"] = [{"Name": x, "ScriptBootstrapAction": {"Path": x}} for x in pars["bootstraps_paths"]] # Debugging and Steps if (pars["debugging"] is True) or (pars["steps"] is not None): - args["Steps"]: List[Dict[str, Any]] = [] + args["Steps"] = [] if pars["debugging"] is True: args["Steps"].append( { @@ -199,7 +276,7 @@ def _build_cluster_args(**pars): # pylint: disable=too-many-branches,too-many-s ], } if pars["instance_num_spot_master"] > 0: # pragma: no cover - fleet_master["LaunchSpecifications"]: Dict = { + fleet_master["LaunchSpecifications"] = { "SpotSpecification": { "TimeoutDurationMinutes": pars["spot_provisioning_timeout_master"], "TimeoutAction": timeout_action_master, @@ -236,7 +313,7 @@ def _build_cluster_args(**pars): # pylint: disable=too-many-branches,too-many-s ], } if pars["instance_num_spot_core"] > 0: - fleet_core["LaunchSpecifications"]: Dict = { + fleet_core["LaunchSpecifications"] = { "SpotSpecification": { "TimeoutDurationMinutes": pars["spot_provisioning_timeout_core"], "TimeoutAction": timeout_action_core, @@ -275,7 +352,7 @@ def _build_cluster_args(**pars): # pylint: disable=too-many-branches,too-many-s ], } if pars["instance_num_spot_task"] > 0: - fleet_task["LaunchSpecifications"]: Dict = { + fleet_task["LaunchSpecifications"] = { "SpotSpecification": { "TimeoutDurationMinutes": pars["spot_provisioning_timeout_task"], "TimeoutAction": timeout_action_task, @@ -287,35 +364,35 @@ def _build_cluster_args(**pars): # pylint: disable=too-many-branches,too-many-s if pars["tags"] is not None: args["Tags"] = [{"Key": k, "Value": v} for k, v in pars["tags"].items()] - _logger.info(f"args: \n{json.dumps(args, default=str, indent=4)}") + _logger.debug("args: \n%s", pprint.pformat(args)) return args def create_cluster( # pylint: disable=too-many-arguments,too-many-locals,unused-argument - cluster_name: str, - logging_s3_path: str, - emr_release: str, subnet_id: str, - emr_ec2_role: str, - emr_role: str, - instance_type_master: str, - instance_type_core: str, - instance_type_task: str, - instance_ebs_size_master: int, - instance_ebs_size_core: int, - instance_ebs_size_task: int, - instance_num_on_demand_master: int, - instance_num_on_demand_core: int, - instance_num_on_demand_task: int, - instance_num_spot_master: int, - instance_num_spot_core: int, - instance_num_spot_task: int, - spot_bid_percentage_of_on_demand_master: int, - spot_bid_percentage_of_on_demand_core: int, - spot_bid_percentage_of_on_demand_task: int, - spot_provisioning_timeout_master: int, - spot_provisioning_timeout_core: int, - spot_provisioning_timeout_task: int, + cluster_name: str = "my-emr-cluster", + logging_s3_path: Optional[str] = None, + emr_release: str = "emr-6.0.0", + emr_ec2_role: str = "EMR_EC2_DefaultRole", + emr_role: str = "EMR_DefaultRole", + instance_type_master: str = "r5.xlarge", + instance_type_core: str = "r5.xlarge", + instance_type_task: str = "r5.xlarge", + instance_ebs_size_master: int = 64, + instance_ebs_size_core: int = 64, + instance_ebs_size_task: int = 64, + instance_num_on_demand_master: int = 1, + instance_num_on_demand_core: int = 0, + instance_num_on_demand_task: int = 0, + instance_num_spot_master: int = 0, + instance_num_spot_core: int = 0, + instance_num_spot_task: int = 0, + spot_bid_percentage_of_on_demand_master: int = 100, + spot_bid_percentage_of_on_demand_core: int = 100, + spot_bid_percentage_of_on_demand_task: int = 100, + spot_provisioning_timeout_master: int = 5, + spot_provisioning_timeout_core: int = 5, + spot_provisioning_timeout_task: int = 5, spot_timeout_to_on_demand_master: bool = True, spot_timeout_to_on_demand_core: bool = True, spot_timeout_to_on_demand_task: bool = True, @@ -337,10 +414,13 @@ def create_cluster( # pylint: disable=too-many-arguments,too-many-locals,unused security_group_slave: Optional[str] = None, security_groups_slave_additional: Optional[List[str]] = None, security_group_service_access: Optional[str] = None, + docker: bool = False, + extra_public_registries: Optional[List[str]] = None, spark_log_level: str = "WARN", spark_jars_path: Optional[List[str]] = None, spark_defaults: Optional[Dict[str, str]] = None, spark_pyarrow: bool = False, + custom_classifications: Optional[List[Dict[str, Any]]] = None, maximize_resource_allocation: bool = False, steps: Optional[List[Dict[str, Any]]] = None, keep_cluster_alive_when_no_steps: bool = True, @@ -354,18 +434,19 @@ def create_cluster( # pylint: disable=too-many-arguments,too-many-locals,unused Parameters ---------- + subnet_id : str + VPC subnet ID. cluster_name : str Cluster name. - logging_s3_path : str + logging_s3_path : str, optional Logging s3 path (e.g. s3://BUCKET_NAME/DIRECTORY_NAME/). + If None, the default is `s3://aws-logs-{AccountId}-{RegionId}/elasticmapreduce/` emr_release : str EMR release (e.g. emr-5.28.0). emr_ec2_role : str IAM role name. emr_role : str IAM role name. - subnet_id : str - VPC subnet ID. instance_type_master : str EC2 instance type. instance_type_core : str @@ -448,6 +529,7 @@ def create_cluster( # pylint: disable=too-many-arguments,too-many-locals,unused Debugging enabled? applications : List[str], optional List of applications (e.g ["Hadoop", "Spark", "Ganglia", "Hive"]). + If None, ["Spark"] will be considered. visible_to_all_users : bool True or False. key_pair_name : str, optional @@ -465,6 +547,10 @@ def create_cluster( # pylint: disable=too-many-arguments,too-many-locals,unused security_group_service_access : str, optional The identifier of the Amazon EC2 security group for the Amazon EMR service to access clusters in VPC private subnets. + docker : bool + Enable Docker Hub and ECR registries access. + extra_public_registries: List[str], optional + Additional docker registries. spark_log_level : str log4j.rootCategory log level (ALL, DEBUG, INFO, WARN, ERROR, FATAL, OFF, TRACE). spark_jars_path : List[str], optional @@ -475,6 +561,8 @@ def create_cluster( # pylint: disable=too-many-arguments,too-many-locals,unused spark_pyarrow : bool Enable PySpark to use PyArrow behind the scenes. P.S. You must install pyarrow by your self via bootstrap + custom_classifications: List[Dict[str, Any]], optional + Extra classifications. maximize_resource_allocation : bool Configure your executors to utilize the maximum resources possible https://docs.aws.amazon.com/emr/latest/ReleaseGuide/emr-spark-configure.html#emr-spark-maximizeresourceallocation @@ -500,6 +588,30 @@ def create_cluster( # pylint: disable=too-many-arguments,too-many-locals,unused Examples -------- + Minimal Example + + >>> import awswrangler as wr + >>> cluster_id = wr.emr.create_cluster("SUBNET_ID") + + Minimal Example With Custom Classification + + >>> import awswrangler as wr + >>> cluster_id = wr.emr.create_cluster( + >>> subnet_id="SUBNET_ID", + >>> custom_classifications=[ + >>> { + >>> "Classification": "livy-conf", + >>> "Properties": { + >>> "livy.spark.master": "yarn", + >>> "livy.spark.deploy-mode": "cluster", + >>> "livy.server.session.timeout": "16h", + >>> }, + >>> } + >>> ], + >>> ) + + Full Example + >>> import awswrangler as wr >>> cluster_id = wr.emr.create_cluster( ... cluster_name="wrangler_cluster", @@ -548,10 +660,12 @@ def create_cluster( # pylint: disable=too-many-arguments,too-many-locals,unused ... }) """ + applications = ["Spark"] if applications is None else applications + boto3_session = _utils.ensure_session(session=boto3_session) args: Dict[str, Any] = _build_cluster_args(**locals()) client_emr: boto3.client = _utils.client(service_name="emr", session=boto3_session) response: Dict[str, Any] = client_emr.run_job_flow(**args) - _logger.debug(f"response: \n{json.dumps(response, default=str, indent=4)}") + _logger.debug("response: \n%s", pprint.pformat(response)) return response["JobFlowId"] @@ -582,7 +696,7 @@ def get_cluster_state(cluster_id: str, boto3_session: Optional[boto3.Session] = """ client_emr: boto3.client = _utils.client(service_name="emr", session=boto3_session) response: Dict[str, Any] = client_emr.describe_cluster(ClusterId=cluster_id) - _logger.debug(f"response: \n{json.dumps(response, default=str, indent=4)}") + _logger.debug("response: \n%s", pprint.pformat(response)) return response["Cluster"]["Status"]["State"] @@ -609,7 +723,7 @@ def terminate_cluster(cluster_id: str, boto3_session: Optional[boto3.Session] = """ client_emr: boto3.client = _utils.client(service_name="emr", session=boto3_session) response: Dict[str, Any] = client_emr.terminate_job_flows(JobFlowIds=[cluster_id]) - _logger.debug(f"response: \n{json.dumps(response, default=str, indent=4)}") + _logger.debug("response: \n%s", pprint.pformat(response)) def submit_steps( @@ -641,14 +755,14 @@ def submit_steps( """ client_emr: boto3.client = _utils.client(service_name="emr", session=boto3_session) response: Dict[str, Any] = client_emr.add_job_flow_steps(JobFlowId=cluster_id, Steps=steps) - _logger.debug(f"response: \n{json.dumps(response, default=str, indent=4)}") + _logger.debug("response: \n%s", pprint.pformat(response)) return response["StepIds"] def submit_step( cluster_id: str, - name: str, command: str, + name: str = "my-step", action_on_failure: str = "CONTINUE", script: bool = False, boto3_session: Optional[boto3.Session] = None, @@ -659,11 +773,11 @@ def submit_step( ---------- cluster_id : str Cluster ID. - name : str - Step name. command : str e.g. 'echo "Hello!"' e.g. for script 's3://.../script.sh arg1 arg2' + name : str, optional + Step name. action_on_failure : str 'TERMINATE_JOB_FLOW', 'TERMINATE_CLUSTER', 'CANCEL_AND_WAIT', 'CONTINUE' script : bool @@ -693,31 +807,34 @@ def submit_step( ) client_emr: boto3.client = _utils.client(service_name="emr", session=session) response: Dict[str, Any] = client_emr.add_job_flow_steps(JobFlowId=cluster_id, Steps=[step]) - _logger.debug(f"response: \n{json.dumps(response, default=str, indent=4)}") + _logger.debug("response: \n%s", pprint.pformat(response)) return response["StepIds"][0] def build_step( - name: str, command: str, + name: str = "my-step", action_on_failure: str = "CONTINUE", script: bool = False, + region: Optional[str] = None, boto3_session: Optional[boto3.Session] = None, ) -> Dict[str, Any]: """Build the Step structure (dictionary). Parameters ---------- - name : str - Step name. command : str e.g. 'echo "Hello!"' e.g. for script 's3://.../script.sh arg1 arg2' + name : str, optional + Step name. action_on_failure : str 'TERMINATE_JOB_FLOW', 'TERMINATE_CLUSTER', 'CANCEL_AND_WAIT', 'CONTINUE' script : bool - True for raw command or False for script runner. + False for raw command or True for script runner. https://docs.aws.amazon.com/emr/latest/ReleaseGuide/emr-commandrunner.html + region: str, optional + Region name to not get it from boto3.Session. (e.g. `us-east-1`) boto3_session : boto3.Session(), optional Boto3 Session. The default boto3 session will be used if boto3_session receive None. @@ -734,14 +851,17 @@ def build_step( >>> wr.emr.submit_steps(cluster_id="cluster-id", steps=steps) """ - session: boto3.Session = _utils.ensure_session(session=boto3_session) jar: str = "command-runner.jar" if script is True: - if session.region_name is not None: - region: str = session.region_name - else: # pragma: no cover - region = "us-east-1" - jar = f"s3://{region}.elasticmapreduce/libs/script-runner/script-runner.jar" + if region is not None: # pragma: no cover + _region: str = region + else: + session: boto3.Session = _utils.ensure_session(session=boto3_session) + if session.region_name is not None: + _region = session.region_name + else: # pragma: no cover + _region = "us-east-1" + jar = f"s3://{_region}.elasticmapreduce/libs/script-runner/script-runner.jar" step: Dict[str, Any] = { "Name": name, "ActionOnFailure": action_on_failure, @@ -778,5 +898,195 @@ def get_step_state(cluster_id: str, step_id: str, boto3_session: Optional[boto3. """ client_emr: boto3.client = _utils.client(service_name="emr", session=boto3_session) response: Dict[str, Any] = client_emr.describe_step(ClusterId=cluster_id, StepId=step_id) - _logger.debug(f"response: \n{json.dumps(response, default=str, indent=4)}") + _logger.debug("response: \n%s", pprint.pformat(response)) return response["Step"]["Status"]["State"] + + +def submit_ecr_credentials_refresh( + cluster_id: str, path: str, action_on_failure: str = "CONTINUE", boto3_session: Optional[boto3.Session] = None +) -> str: + """Update internal ECR credentials. + + Parameters + ---------- + cluster_id : str + Cluster ID. + path : str + Amazon S3 path where Wrangler will stage the script ecr_credentials_refresh.py (e.g. s3://bucket/emr/) + action_on_failure : str + 'TERMINATE_JOB_FLOW', 'TERMINATE_CLUSTER', 'CANCEL_AND_WAIT', 'CONTINUE' + boto3_session : boto3.Session(), optional + Boto3 Session. The default boto3 session will be used if boto3_session receive None. + + Returns + ------- + str + Step ID. + + Examples + -------- + >>> import awswrangler as wr + >>> step_id = wr.emr.submit_ecr_credentials_refresh("cluster_id", "s3://bucket/emr/") + + """ + path = path[:-1] if path.endswith("/") else path + path_script: str = f"{path}/ecr_credentials_refresh.py" + session: boto3.Session = _utils.ensure_session(session=boto3_session) + client_s3: boto3.client = _utils.client(service_name="s3", session=session) + bucket, key = _utils.parse_path(path=path_script) + client_s3.put_object(Body=_get_ecr_credentials_refresh_content().encode(encoding="utf-8"), Bucket=bucket, Key=key) + command: str = f"spark-submit --deploy-mode cluster {path_script}" + name: str = "ECR Credentials Refresh" + step: Dict[str, Any] = build_step( + name=name, command=command, action_on_failure=action_on_failure, script=False, boto3_session=session + ) + client_emr: boto3.client = _utils.client(service_name="emr", session=session) + response: Dict[str, Any] = client_emr.add_job_flow_steps(JobFlowId=cluster_id, Steps=[step]) + _logger.debug("response: \n%s", pprint.pformat(response)) + return response["StepIds"][0] + + +def _get_ecr_credentials_refresh_content() -> str: + return """ +import subprocess +from pyspark.sql import SparkSession +spark = SparkSession.builder.appName("ECR Setup Job").getOrCreate() + +COMMANDS = [ + "sudo -s eval $(aws ecr get-login --region us-east-1 --no-include-email)", + "sudo hdfs dfs -put -f /root/.docker/config.json /user/hadoop/" +] + +for command in COMMANDS: + subprocess.run(command.split(" "), timeout=6.0, check=True) + +print("done!") + """ + + +def build_spark_step( + path: str, + deploy_mode: str = "cluster", + docker_image: Optional[str] = None, + name: str = "my-step", + action_on_failure: str = "CONTINUE", + region: Optional[str] = None, + boto3_session: Optional[boto3.Session] = None, +) -> Dict[str, Any]: + """Build the Step structure (dictionary). + + Parameters + ---------- + path : str + Script path. (e.g. s3://bucket/app.py) + deploy_mode : str + "cluster" | "client" + docker_image : str, optional + e.g. "{ACCOUNT_ID}.dkr.ecr.{REGION}.amazonaws.com/{IMAGE_NAME}:{TAG}" + name : str, optional + Step name. + action_on_failure : str + 'TERMINATE_JOB_FLOW', 'TERMINATE_CLUSTER', 'CANCEL_AND_WAIT', 'CONTINUE' + region: str, optional + Region name to not get it from boto3.Session. (e.g. `us-east-1`) + boto3_session : boto3.Session(), optional + Boto3 Session. The default boto3 session will be used if boto3_session receive None. + + Returns + ------- + Dict[str, Any] + Step structure. + + Examples + -------- + >>> import awswrangler as wr + >>> step_id = wr.emr.submit_steps( + >>> cluster_id="cluster-id", + >>> steps=[ + >>> wr.emr.build_spark_step(path="s3://bucket/app.py") + >>> ] + >>> ) + + """ + if docker_image is None: # pragma: no cover + cmd: str = f"spark-submit --deploy-mode {deploy_mode} {path}" + else: + config: str = "hdfs:///user/hadoop/config.json" + cmd = ( + f"spark-submit --deploy-mode cluster " + f"--conf spark.executorEnv.YARN_CONTAINER_RUNTIME_TYPE=docker " + f"--conf spark.executorEnv.YARN_CONTAINER_RUNTIME_DOCKER_IMAGE={docker_image} " + f"--conf spark.executorEnv.YARN_CONTAINER_RUNTIME_DOCKER_CLIENT_CONFIG={config} " + f"--conf spark.executorEnv.YARN_CONTAINER_RUNTIME_DOCKER_MOUNTS=/etc/passwd:/etc/passwd:ro " + f"--conf spark.yarn.appMasterEnv.YARN_CONTAINER_RUNTIME_TYPE=docker " + f"--conf spark.yarn.appMasterEnv.YARN_CONTAINER_RUNTIME_DOCKER_IMAGE={docker_image} " + f"--conf spark.yarn.appMasterEnv.YARN_CONTAINER_RUNTIME_DOCKER_CLIENT_CONFIG={config} " + f"--conf spark.yarn.appMasterEnv.YARN_CONTAINER_RUNTIME_DOCKER_MOUNTS=/etc/passwd:/etc/passwd:ro " + f"{path}" + ) + return build_step( + command=cmd, + name=name, + action_on_failure=action_on_failure, + script=False, + region=region, + boto3_session=boto3_session, + ) + + +def submit_spark_step( + cluster_id: str, + path: str, + deploy_mode: str = "cluster", + docker_image: Optional[str] = None, + name: str = "my-step", + action_on_failure: str = "CONTINUE", + region: Optional[str] = None, + boto3_session: Optional[boto3.Session] = None, +) -> str: + """Submit Spark Step. + + Parameters + ---------- + cluster_id : str + Cluster ID. + path : str + Script path. (e.g. s3://bucket/app.py) + deploy_mode : str + "cluster" | "client" + docker_image : str, optional + e.g. "{ACCOUNT_ID}.dkr.ecr.{REGION}.amazonaws.com/{IMAGE_NAME}:{TAG}" + name : str, optional + Step name. + action_on_failure : str + 'TERMINATE_JOB_FLOW', 'TERMINATE_CLUSTER', 'CANCEL_AND_WAIT', 'CONTINUE' + region: str, optional + Region name to not get it from boto3.Session. (e.g. `us-east-1`) + boto3_session : boto3.Session(), optional + Boto3 Session. The default boto3 session will be used if boto3_session receive None. + + Returns + ------- + str + Step ID. + + Examples + -------- + >>> import awswrangler as wr + >>> step_id = wr.emr.submit_spark_step( + >>> cluster_id="cluster-id", + >>> path="s3://bucket/emr/app.py" + >>> ) + + """ + session: boto3.Session = _utils.ensure_session(session=boto3_session) + step = build_spark_step( + path=path, + deploy_mode=deploy_mode, + docker_image=docker_image, + name=name, + action_on_failure=action_on_failure, + region=region, + boto3_session=session, + ) + return submit_steps(cluster_id=cluster_id, steps=[step], boto3_session=session)[0] diff --git a/awswrangler/s3.py b/awswrangler/s3.py index 45459caf2..f4d1e5746 100644 --- a/awswrangler/s3.py +++ b/awswrangler/s3.py @@ -56,10 +56,10 @@ def get_bucket_region(bucket: str, boto3_session: Optional[boto3.Session] = None """ client_s3: boto3.client = _utils.client(service_name="s3", session=boto3_session) - _logger.debug(f"bucket: {bucket}") + _logger.debug("bucket: %s", bucket) region: str = client_s3.get_bucket_location(Bucket=bucket)["LocationConstraint"] region = "us-east-1" if region is None else region - _logger.debug(f"region: {region}") + _logger.debug("region: %s", region) return region @@ -111,7 +111,7 @@ def does_object_exist(path: str, boto3_session: Optional[boto3.Session] = None) raise ex # pragma: no cover -def list_objects(path: str, boto3_session: Optional[boto3.Session] = None) -> List[str]: +def list_directories(path: str, boto3_session: Optional[boto3.Session] = None) -> List[str]: """List Amazon S3 objects from a prefix. Parameters @@ -130,6 +130,42 @@ def list_objects(path: str, boto3_session: Optional[boto3.Session] = None) -> Li -------- Using the default boto3 session + >>> import awswrangler as wr + >>> wr.s3.list_objects('s3://bucket/prefix/') + ['s3://bucket/prefix/dir0', 's3://bucket/prefix/dir1', 's3://bucket/prefix/dir2'] + + Using a custom boto3 session + + >>> import boto3 + >>> import awswrangler as wr + >>> wr.s3.list_objects('s3://bucket/prefix/', boto3_session=boto3.Session()) + ['s3://bucket/prefix/dir0', 's3://bucket/prefix/dir1', 's3://bucket/prefix/dir2'] + + """ + return _list_objects(path=path, delimiter="/", boto3_session=boto3_session) + + +def list_objects(path: str, suffix: Optional[str] = None, boto3_session: Optional[boto3.Session] = None) -> List[str]: + """List Amazon S3 objects from a prefix. + + Parameters + ---------- + path : str + S3 path (e.g. s3://bucket/prefix). + suffix: str, optional + Suffix for filtering S3 keys. + boto3_session : boto3.Session(), optional + Boto3 Session. The default boto3 session will be used if boto3_session receive None. + + Returns + ------- + List[str] + List of objects paths. + + Examples + -------- + Using the default boto3 session + >>> import awswrangler as wr >>> wr.s3.list_objects('s3://bucket/prefix') ['s3://bucket/prefix0', 's3://bucket/prefix1', 's3://bucket/prefix2'] @@ -142,28 +178,49 @@ def list_objects(path: str, boto3_session: Optional[boto3.Session] = None) -> Li ['s3://bucket/prefix0', 's3://bucket/prefix1', 's3://bucket/prefix2'] """ + return _list_objects(path=path, delimiter=None, suffix=suffix, boto3_session=boto3_session) + + +def _list_objects( + path: str, + delimiter: Optional[str] = None, + suffix: Optional[str] = None, + boto3_session: Optional[boto3.Session] = None, +) -> List[str]: client_s3: boto3.client = _utils.client(service_name="s3", session=boto3_session) paginator = client_s3.get_paginator("list_objects_v2") bucket: str prefix: str bucket, prefix = _utils.parse_path(path=path) - response_iterator = paginator.paginate(Bucket=bucket, Prefix=prefix, PaginationConfig={"PageSize": 1000}) + args: Dict[str, Any] = {"Bucket": bucket, "Prefix": prefix, "PaginationConfig": {"PageSize": 1000}} + if delimiter is not None: + args["Delimiter"] = delimiter + response_iterator = paginator.paginate(**args) paths: List[str] = [] - for page in response_iterator: - contents: Optional[List] = page.get("Contents") - if contents is not None: - for content in contents: - if (content is not None) and ("Key" in content): - key: str = content["Key"] - paths.append(f"s3://{bucket}/{key}") + for page in response_iterator: # pylint: disable=too-many-nested-blocks + if delimiter is None: + contents: Optional[List] = page.get("Contents") + if contents is not None: + for content in contents: + if (content is not None) and ("Key" in content): + key: str = content["Key"] + if (suffix is None) or key.endswith(suffix): + paths.append(f"s3://{bucket}/{key}") + else: + prefixes: Optional[List[Optional[Dict[str, str]]]] = page.get("CommonPrefixes") + if prefixes is not None: + for pfx in prefixes: + if (pfx is not None) and ("Prefix" in pfx): + key = pfx["Prefix"] + paths.append(f"s3://{bucket}/{key}") return paths -def _path2list(path: Union[str, List[str]], boto3_session: Optional[boto3.Session]) -> List[str]: +def _path2list(path: object, boto3_session: boto3.Session, suffix: str = None) -> List[str]: if isinstance(path, str): # prefix paths: List[str] = list_objects(path=path, boto3_session=boto3_session) elif isinstance(path, list): - paths = path + paths = path if suffix is None else [x for x in path if x.endswith(suffix)] else: raise exceptions.InvalidArgumentType(f"{type(path)} is not a valid path type. Please, use str or List[str].") return paths @@ -229,7 +286,7 @@ def _split_paths_by_bucket(paths: List[str]) -> Dict[str, List[str]]: def _delete_objects(bucket: str, keys: List[str], client_s3: boto3.client) -> None: - _logger.debug(f"len(keys): {len(keys)}") + _logger.debug("len(keys): %s", len(keys)) batch: List[Dict[str, str]] = [{"Key": key} for key in keys] client_s3.delete_objects(Bucket=bucket, Delete={"Objects": batch}) @@ -309,7 +366,7 @@ def _describe_object( break except botocore.exceptions.ClientError as e: # pragma: no cover if e.response["ResponseMetadata"]["HTTPStatusCode"] == 404: # Not Found - _logger.debug(f"Object not found. {i} seconds remaining to wait.") + _logger.debug("Object not found. %s seconds remaining to wait.", i) if i == 1: # Last try, there is no more need to sleep break time.sleep(1) @@ -376,6 +433,7 @@ def to_csv( # pylint: disable=too-many-arguments 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, @@ -426,6 +484,8 @@ def to_csv( # pylint: disable=too-many-arguments 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 @@ -620,10 +680,11 @@ def to_csv( # pylint: disable=too-many-arguments columns_comments=columns_comments, boto3_session=session, mode="overwrite", + catalog_versioning=catalog_versioning, sep=sep, ) if partitions_values: - _logger.debug(f"partitions_values:\n{partitions_values}") + _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 ) @@ -652,7 +713,7 @@ def _to_csv_dataset( if (mode == "overwrite") or ((mode == "overwrite_partitions") and (not partition_cols)): delete_objects(path=path, use_threads=use_threads, boto3_session=boto3_session) df = _data_types.cast_pandas_with_athena_types(df=df, dtype=dtype) - _logger.debug(f"dtypes: {df.dtypes}") + _logger.debug("dtypes: %s", df.dtypes) if not partition_cols: file_path: str = f"{path}{uuid.uuid4().hex}.csv" _to_text( @@ -789,6 +850,7 @@ def to_parquet( # pylint: disable=too-many-arguments 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, @@ -836,6 +898,8 @@ def to_parquet( # pylint: disable=too-many-arguments 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 @@ -1035,9 +1099,10 @@ def to_parquet( # pylint: disable=too-many-arguments columns_comments=columns_comments, boto3_session=session, mode="overwrite", + catalog_versioning=catalog_versioning, ) if partitions_values: - _logger.debug(f"partitions_values:\n{partitions_values}") + _logger.debug("partitions_values:\n%s", partitions_values) catalog.add_parquet_partitions( database=database, table=table, @@ -1075,7 +1140,7 @@ def _to_parquet_dataset( schema: pa.Schema = _data_types.pyarrow_schema_from_pandas( df=df, index=index, ignore_cols=partition_cols, dtype=dtype ) - _logger.debug(f"schema: {schema}") + _logger.debug("schema: \n%s", schema) if not partition_cols: file_path: str = f"{path}{uuid.uuid4().hex}{compression_ext}.parquet" _to_parquet_file( @@ -1123,7 +1188,7 @@ def _to_parquet_file( 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(f"Casting column {col_name} ({col_index}) to {col_type} ({pyarrow_dtype})") + _logger.debug("Casting column %s (%s) to %s (%s)", col_name, col_index, col_type, pyarrow_dtype) pyarrow.parquet.write_table( table=table, where=path, @@ -1451,7 +1516,7 @@ def _read_text_chunksize( ) -> Iterator[pd.DataFrame]: fs: s3fs.S3FileSystem = _utils.get_fs(session=boto3_session, s3_additional_kwargs=s3_additional_kwargs) for path in paths: - _logger.debug(f"path: {path}") + _logger.debug("path: %s", path) if pandas_args.get("compression", "infer") == "infer": pandas_args["compression"] = infer_compression(path, compression="infer") with fs.open(path, "rb") as f: @@ -1491,7 +1556,7 @@ def _read_parquet_init( path_or_paths = path[:-1] if path.endswith("/") else path else: path_or_paths = path - _logger.debug(f"path_or_paths: {path_or_paths}") + _logger.debug("path_or_paths: %s", path_or_paths) fs: s3fs.S3FileSystem = _utils.get_fs(session=boto3_session, s3_additional_kwargs=s3_additional_kwargs) cpus: int = _utils.ensure_cpu_count(use_threads=use_threads) data: pyarrow.parquet.ParquetDataset = pyarrow.parquet.ParquetDataset( @@ -1501,6 +1566,7 @@ def _read_parquet_init( filters=filters, read_dictionary=categories, validate_schema=validate_schema, + split_row_groups=False, ) return data @@ -1510,7 +1576,7 @@ def read_parquet( filters: Optional[Union[List[Tuple], List[List[Tuple]]]] = None, columns: Optional[List[str]] = None, validate_schema: bool = True, - chunked: bool = False, + chunked: Union[bool, int] = False, dataset: bool = False, categories: List[str] = None, use_threads: bool = True, @@ -1522,6 +1588,22 @@ def read_parquet( The concept of Dataset goes beyond the simple idea of files and enable more complex features like partitioning and catalog integration (AWS Glue Catalog). + Note + ---- + ``Batching`` (`chunked` argument) (Memory Friendly): + + Will anable the function to return a Iterable of DataFrames instead of a regular DataFrame. + + There are two batching strategies on Wrangler: + + - If **chunked=True**, a new DataFrame will be returned for each file in your path/dataset. + + - If **chunked=INTEGER**, Wrangler will iterate on the data by number of rows igual the received INTEGER. + + `P.S.` `chunked=True` if faster and uses less memory while `chunked=INTEGER` is more precise + in number of rows for each Dataframe. + + Note ---- In case of `use_threads=True` the number of threads that will be spawned will be get from os.cpu_count(). @@ -1538,11 +1620,12 @@ def read_parquet( Check that individual file schemas are all the same / compatible. Schemas within a folder prefix should all be the same. Disable if you have schemas that are different and want to disable this check. - chunked : bool - If True will break the data in smaller DataFrames (Non deterministic number of lines). - Otherwise return a single DataFrame with the whole data. + chunked : Union[int, bool] + If passed will split the data in a Iterable of DataFrames (Memory friendly). + If `True` wrangler will iterate on the data by files in the most efficient way without guarantee of chunksize. + If an `INTEGER` is passed Wrangler will iterate on the data by number of rows igual the received INTEGER. dataset: bool - If True read a parquet dataset instead of simple file(s) loading all the related partitions as columns. + If `True` read a parquet dataset instead of simple file(s) loading all the related partitions as columns. categories: List[str], optional List of columns names that should be returned as pandas.Categorical. Recommended for memory restricted environments. @@ -1583,29 +1666,38 @@ def read_parquet( >>> import awswrangler as wr >>> df = wr.s3.read_parquet(path=['s3://bucket/filename0.parquet', 's3://bucket/filename1.parquet']) - Reading in chunks + Reading in chunks (Chunk by file) >>> import awswrangler as wr >>> dfs = wr.s3.read_parquet(path=['s3://bucket/filename0.csv', 's3://bucket/filename1.csv'], chunked=True) >>> for df in dfs: >>> print(df) # Smaller Pandas DataFrame + Reading in chunks (Chunk by 1MM rows) + + >>> import awswrangler as wr + >>> dfs = wr.s3.read_parquet(path=['s3://bucket/filename0.csv', 's3://bucket/filename1.csv'], chunked=1_000_000) + >>> for df in dfs: + >>> print(df) # 1MM Pandas DataFrame + """ data: pyarrow.parquet.ParquetDataset = _read_parquet_init( path=path, filters=filters, dataset=dataset, categories=categories, + validate_schema=validate_schema, use_threads=use_threads, boto3_session=boto3_session, s3_additional_kwargs=s3_additional_kwargs, - validate_schema=validate_schema, ) if chunked is False: return _read_parquet( data=data, columns=columns, categories=categories, use_threads=use_threads, validate_schema=validate_schema ) - return _read_parquet_chunked(data=data, columns=columns, categories=categories, use_threads=use_threads) + return _read_parquet_chunked( + data=data, columns=columns, categories=categories, chunked=chunked, use_threads=use_threads + ) def _read_parquet( @@ -1639,22 +1731,45 @@ def _read_parquet_chunked( data: pyarrow.parquet.ParquetDataset, columns: Optional[List[str]] = None, categories: List[str] = None, + chunked: Union[bool, int] = True, use_threads: bool = True, ) -> Iterator[pd.DataFrame]: + next_slice: Optional[pd.DataFrame] = None for piece in data.pieces: - table: pa.Table = piece.read( - columns=columns, use_threads=use_threads, partitions=data.partitions, use_pandas_metadata=False - ) - yield table.to_pandas( - use_threads=use_threads, - split_blocks=True, - self_destruct=True, - integer_object_nulls=False, - date_as_object=True, - ignore_metadata=True, + df: pd.DataFrame = _table2df( + table=piece.read( + columns=columns, use_threads=use_threads, partitions=data.partitions, use_pandas_metadata=False + ), categories=categories, - types_mapper=_data_types.pyarrow2pandas_extension, + use_threads=use_threads, ) + if chunked is True: + yield df + else: + if next_slice is not None: + df = pd.concat(objs=[next_slice, df], ignore_index=True, sort=False) + while len(df.index) >= chunked: + yield df.iloc[:chunked] + df = df.iloc[chunked:] + if df.empty: + next_slice = None + else: + next_slice = df + if next_slice is not None: + yield next_slice + + +def _table2df(table: pa.Table, categories: List[str] = None, use_threads: bool = True) -> pd.DataFrame: + return table.to_pandas( + use_threads=use_threads, + split_blocks=True, + self_destruct=True, + integer_object_nulls=False, + date_as_object=True, + ignore_metadata=True, + categories=categories, + types_mapper=_data_types.pyarrow2pandas_extension, + ) def read_parquet_metadata( @@ -1731,6 +1846,7 @@ def store_parquet_metadata( columns_comments: Optional[Dict[str, str]] = None, compression: Optional[str] = None, mode: str = "overwrite", + catalog_versioning: bool = False, boto3_session: Optional[boto3.Session] = None, ) -> Tuple[Dict[str, str], Optional[Dict[str, str]], Optional[Dict[str, List[str]]]]: """Infer and store parquet metadata on AWS Glue Catalog. @@ -1772,6 +1888,8 @@ def store_parquet_metadata( Compression style (``None``, ``snappy``, ``gzip``, etc). mode: str 'overwrite' to recreate any possible existing table or 'append' to keep any possible existing table. + catalog_versioning : bool + If True and `mode="overwrite"`, creates an archived version of the table catalog before updating it. boto3_session : boto3.Session(), optional Boto3 Session. The default boto3 session will be used if boto3_session receive None. @@ -1817,6 +1935,7 @@ def store_parquet_metadata( parameters=parameters, columns_comments=columns_comments, mode=mode, + catalog_versioning=catalog_versioning, boto3_session=session, ) partitions_values: Dict[str, List[str]] = _data_types.athena_partitions_from_pyarrow_partitions( @@ -1976,13 +2095,30 @@ def read_parquet_table( filters: Optional[Union[List[Tuple], List[List[Tuple]]]] = None, columns: Optional[List[str]] = None, categories: List[str] = None, - chunked: bool = False, + chunked: Union[bool, int] = False, use_threads: bool = True, boto3_session: Optional[boto3.Session] = None, s3_additional_kwargs: Optional[Dict[str, str]] = None, ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]: """Read Apache Parquet table registered on AWS Glue Catalog. + Note + ---- + ``Batching`` (`chunked` argument) (Memory Friendly): + + Will anable the function to return a Iterable of DataFrames instead of a regular DataFrame. + + There are two batching strategies on Wrangler: + + - If **chunked=True**, a new DataFrame will be returned for each file in your path/dataset. + + - If **chunked=INTEGER**, Wrangler will paginate through files slicing and concatenating + to return DataFrames with the number of row igual the received INTEGER. + + `P.S.` `chunked=True` if faster and uses less memory while `chunked=INTEGER` is more precise + in number of rows for each Dataframe. + + Note ---- In case of `use_threads=True` the number of threads that will be spawned will be get from os.cpu_count(). @@ -2036,13 +2172,20 @@ def read_parquet_table( ... } ... ) - Reading Parquet Table in chunks + Reading Parquet Table in chunks (Chunk by file) >>> import awswrangler as wr >>> dfs = wr.s3.read_parquet_table(database='...', table='...', chunked=True) >>> for df in dfs: >>> print(df) # Smaller Pandas DataFrame + Reading in chunks (Chunk by 1MM rows) + + >>> import awswrangler as wr + >>> dfs = wr.s3.read_parquet(path=['s3://bucket/filename0.csv', 's3://bucket/filename1.csv'], chunked=1_000_000) + >>> for df in dfs: + >>> print(df) # 1MM Pandas DataFrame + """ path: str = catalog.get_table_location(database=database, table=table, boto3_session=boto3_session) return read_parquet( @@ -2112,12 +2255,12 @@ def merge_datasets( session: boto3.Session = _utils.ensure_session(session=boto3_session) paths: List[str] = list_objects(path=f"{source_path}/", boto3_session=session) - _logger.debug(f"len(paths): {len(paths)}") + _logger.debug("len(paths): %s", len(paths)) if len(paths) < 1: return [] if mode == "overwrite": - _logger.debug(f"Deleting to overwrite: {target_path}/") + _logger.debug("Deleting to overwrite: %s/", target_path) delete_objects(path=f"{target_path}/", use_threads=use_threads, boto3_session=session) elif mode == "overwrite_partitions": paths_wo_prefix: List[str] = [x.replace(f"{source_path}/", "") for x in paths] @@ -2125,7 +2268,7 @@ def merge_datasets( partitions_paths: List[str] = list(set(paths_wo_filename)) target_partitions_paths = [f"{target_path}/{x}" for x in partitions_paths] for path in target_partitions_paths: - _logger.debug(f"Deleting to overwrite_partitions: {path}") + _logger.debug("Deleting to overwrite_partitions: %s", path) delete_objects(path=path, use_threads=use_threads, boto3_session=session) elif mode != "append": raise exceptions.InvalidArgumentValue(f"{mode} is a invalid mode option.") @@ -2133,7 +2276,7 @@ def merge_datasets( new_objects: List[str] = copy_objects( paths=paths, source_path=source_path, target_path=target_path, use_threads=use_threads, boto3_session=session ) - _logger.debug(f"len(new_objects): {len(new_objects)}") + _logger.debug("len(new_objects): %s", len(new_objects)) return new_objects @@ -2141,6 +2284,7 @@ def copy_objects( paths: List[str], source_path: str, target_path: str, + replace_filenames: Optional[Dict[str, str]] = None, use_threads: bool = True, boto3_session: Optional[boto3.Session] = None, ) -> List[str]: @@ -2180,7 +2324,7 @@ def copy_objects( ["s3://bucket1/dir1/key0", "s3://bucket1/dir1/key1"] """ - _logger.debug(f"len(paths): {len(paths)}") + _logger.debug("len(paths): %s", len(paths)) if len(paths) < 1: return [] source_path = source_path[:-1] if source_path[-1] == "/" else source_path @@ -2191,15 +2335,24 @@ def copy_objects( for path in paths: path_wo_prefix: str = path.replace(f"{source_path}/", "") path_final: str = f"{target_path}/{path_wo_prefix}" + if replace_filenames is not None: + parts: List[str] = path_final.rsplit(sep="/", maxsplit=1) + if len(parts) == 2: + path_wo_filename: str = parts[0] + filename: str = parts[1] + if filename in replace_filenames: + new_filename: str = replace_filenames[filename] + _logger.debug("Replacing filename: %s -> %s", filename, new_filename) + path_final = f"{path_wo_filename}/{new_filename}" new_objects.append(path_final) batch.append((path, path_final)) - _logger.debug(f"len(new_objects): {len(new_objects)}") + _logger.debug("len(new_objects): %s", len(new_objects)) _copy_objects(batch=batch, use_threads=use_threads, boto3_session=session) return new_objects def _copy_objects(batch: List[Tuple[str, str]], use_threads: bool, boto3_session: boto3.Session) -> None: - _logger.debug(f"len(batch): {len(batch)}") + _logger.debug("len(batch): %s", len(batch)) client_s3: boto3.client = _utils.client(service_name="s3", session=boto3_session) resource_s3: boto3.resource = _utils.resource(service_name="s3", session=boto3_session) for source, target in batch: diff --git a/awswrangler/torch.py b/awswrangler/torch.py new file mode 100644 index 000000000..70df93f34 --- /dev/null +++ b/awswrangler/torch.py @@ -0,0 +1,473 @@ +"""PyTorch Module.""" +import io +import logging +import os +import pathlib +import re +from collections.abc import Iterable +from io import BytesIO +from typing import Any, Callable, Iterator, List, Optional, Tuple, Union + +import boto3 # type: ignore +import numpy as np # type: ignore +import sqlalchemy # type: ignore +import torch # type: ignore +import torchaudio # type: ignore +from PIL import Image # type: ignore +from torch.utils.data.dataset import Dataset, IterableDataset # type: ignore +from torchvision.transforms.functional import to_tensor # type: ignore + +from awswrangler import _utils, db, s3 + +_logger: logging.Logger = logging.getLogger(__name__) + + +class _BaseS3Dataset: + """PyTorch Amazon S3 Map-Style Dataset.""" + + def __init__( + self, path: Union[str, List[str]], suffix: Optional[str] = None, boto3_session: Optional[boto3.Session] = None + ): + r"""PyTorch Map-Style S3 Dataset. + + Parameters + ---------- + path : Union[str, List[str]] + S3 prefix (e.g. s3://bucket/prefix) or list of S3 objects paths (e.g. [s3://bucket/key0, s3://bucket/key1]). + suffix: str, optional + S3 suffix filtering of object keys (i.e. suffix=".png" -> s3://\*.png). + boto3_session : boto3.Session(), optional + Boto3 Session. The default boto3 session will be used if boto3_session receive None. + + Returns + ------- + torch.utils.data.Dataset + + """ + super().__init__() + self._session = _utils.ensure_session(session=boto3_session) + self._paths: List[str] = s3._path2list( # pylint: disable=protected-access + path=path, suffix=suffix, boto3_session=self._session + ) + + def _fetch_data(self, path: str) -> Any: + """Add parquet and csv support.""" + bucket, key = _utils.parse_path(path=path) + buff = BytesIO() + client_s3: boto3.client = _utils.client(service_name="s3", session=self._session) + client_s3.download_fileobj(Bucket=bucket, Key=key, Fileobj=buff) + buff.seek(0) + return buff + + @staticmethod + def _load_data(data: io.BytesIO, path: str) -> Any: + if path.endswith(".pt"): + data = torch.load(data) + elif path.endswith(".tar.gz") or path.endswith(".tgz"): # pragma: no cover + raise NotImplementedError("Tar loader not implemented!") + # tarfile.open(fileobj=data) + # tar = tarfile.open(fileobj=data) + # for member in tar.getmembers(): + else: # pragma: no cover + raise NotImplementedError() + + return data + + +class _ListS3Dataset(_BaseS3Dataset, Dataset): + """PyTorch Amazon S3 Map-Style List Dataset.""" + + def __getitem__(self, index): + path = self._paths[index] + data = self._fetch_data(path) + return [self._data_fn(data), self._label_fn(path)] + + def __len__(self): + return len(self._paths) + + def _data_fn(self, data) -> Any: # pragma: no cover + raise NotImplementedError() + + def _label_fn(self, path: str) -> Any: # pragma: no cover + raise NotImplementedError() + + +class _S3PartitionedDataset(_ListS3Dataset): + """PyTorch Amazon S3 Map-Style Partitioned Dataset.""" + + def _label_fn(self, path: str) -> torch.Tensor: + label = int(re.findall(r"/(.*?)=(.*?)/", path)[-1][1]) + return torch.tensor([label]) # pylint: disable=not-callable + + def _data_fn(self, data) -> Any: # pragma: no cover + raise NotImplementedError() + + +# class S3FilesDataset(_BaseS3Dataset, Dataset): +# """PyTorch Amazon S3 Files Map-Style Dataset.""" +# +# def __init__( +# self, path: Union[str, List[str]], suffix: Optional[str] = None, boto3_session: Optional[boto3.Session] = None +# ): +# """PyTorch S3 Files Map-Style Dataset. +# +# Each file under Amazon S3 path would be handled as a tensor or batch of tensors. +# +# Note +# ---- +# All files will be loaded to memory since random access is needed. +# +# Parameters +# ---------- +# path : Union[str, List[str]] +# S3 prefix (e.g. s3://bucket/prefix) or +# list of S3 objects paths (e.g. [s3://bucket/key0, s3://bucket/key1]). +# boto3_session : boto3.Session(), optional +# Boto3 Session. The default boto3 session will be used if boto3_session receive None. +# +# Returns +# ------- +# torch.utils.data.Dataset +# +# """ +# super(S3FilesDataset, self).__init__(path, suffix, boto3_session) +# self._download_files() +# +# def _download_files(self) -> None: +# self._data = [] +# for path in self._paths: +# data = self._fetch_data(path) +# data = self._load_data(data, path) +# self._data.append(data) +# +# self.data = torch.cat(self._data, dim=0) +# +# def __getitem__(self, index): +# return self._data[index] +# +# def __len__(self): +# return len(self._data) + + +class LambdaS3Dataset(_ListS3Dataset): + """PyTorch Amazon S3 Lambda Map-Style Dataset.""" + + def __init__( + self, + path: Union[str, List[str]], + data_fn: Callable, + label_fn: Callable, + suffix: Optional[str] = None, + boto3_session: Optional[boto3.Session] = None, + ): + r"""PyTorch Amazon S3 Lambda Dataset. + + Parameters + ---------- + path : Union[str, List[str]] + S3 prefix (e.g. s3://bucket/prefix) or list of S3 objects paths (e.g. [s3://bucket/key0, s3://bucket/key1]). + data_fn: Callable + Function that receives a io.BytesIO object and returns a torch.Tensor + label_fn: Callable + Function that receives object path (str) and return a torch.Tensor + suffix: str, optional + S3 suffix filtering of object keys (i.e. suffix=".png" -> s3://\*.png). + boto3_session : boto3.Session(), optional + Boto3 Session. The default boto3 session will be used if boto3_session receive None. + + Returns + ------- + torch.utils.data.Dataset + + Examples + -------- + >>> import re + >>> import torch + >>> import awswrangler as wr + >>> ds = wr.torch.LambdaS3Dataset( + >>> 's3://bucket/path', + >>> data_fn=lambda x: torch.load(x), + >>> label_fn=lambda x: torch.Tensor(int(re.findall(r"/class=(.*?)/", x)[-1])), + >>> ) + + """ + super(LambdaS3Dataset, self).__init__(path, suffix, boto3_session) + self._data_func = data_fn + self._label_func = label_fn + + def _label_fn(self, path: str) -> torch.Tensor: + return self._label_func(path) + + def _data_fn(self, data) -> torch.Tensor: + return self._data_func(data) + + +class AudioS3Dataset(_S3PartitionedDataset): + """PyTorch S3 Audio Dataset.""" + + def __init__( + self, + path: Union[str, List[str]], + cache_dir: str = "/tmp/", + suffix: Optional[str] = None, + boto3_session: Optional[boto3.Session] = None, + ): + r"""PyTorch Amazon S3 Audio Dataset. + + Read individual WAV audio files stores in Amazon S3 and return + them as torch tensors. + + Note + ---- + This dataset assumes audio files are stored with the following structure: + + + :: + + bucket + ├── class=0 + │ ├── audio0.wav + │ └── audio1.wav + └── class=1 + ├── audio2.wav + └── audio3.wav + + Parameters + ---------- + path : Union[str, List[str]] + S3 prefix (e.g. s3://bucket/prefix) or list of S3 objects paths (e.g. [s3://bucket/key0, s3://bucket/key1]). + suffix: str, optional + S3 suffix filtering of object keys (i.e. suffix=".png" -> s3://\*.png). + boto3_session : boto3.Session(), optional + Boto3 Session. The default boto3 session will be used if boto3_session receive None. + + Returns + ------- + torch.utils.data.Dataset + + Examples + -------- + Create a Audio S3 Dataset + + >>> import awswrangler as wr + >>> ds = wr.torch.AudioS3Dataset('s3://bucket/path') + + + Training a Model + + >>> criterion = CrossEntropyLoss().to(device) + >>> opt = SGD(model.parameters(), 0.025) + >>> loader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers) + >>> + >>> for epoch in range(epochs): + >>> + >>> correct = 0 + >>> model.train() + >>> for i, (inputs, labels) in enumerate(loader): + >>> + >>> # Forward Pass + >>> outputs = model(inputs) + >>> + >>> # Backward Pass + >>> loss = criterion(outputs, labels) + >>> loss.backward() + >>> opt.step() + >>> opt.zero_grad() + >>> + >>> # Accuracy + >>> _, predicted = torch.max(outputs.data, 1) + >>> correct += (predicted == labels).sum().item() + >>> accuracy = 100 * correct / ((i+1) * batch_size) + >>> print(f'batch: {i} loss: {loss.mean().item():.4f} acc: {accuracy:.2f}') + + """ + super(AudioS3Dataset, self).__init__(path, suffix, boto3_session) + self._cache_dir: str = cache_dir[:-1] if cache_dir.endswith("/") else cache_dir + + def _data_fn(self, filename: str) -> Tuple[Any, Any]: # pylint: disable=arguments-differ + waveform, sample_rate = torchaudio.load(filename) + os.remove(path=filename) + return waveform, sample_rate + + def _fetch_data(self, path: str) -> str: + bucket, key = _utils.parse_path(path=path) + filename: str = f"{self._cache_dir}/{bucket}/{key}" + pathlib.Path(filename).parent.mkdir(parents=True, exist_ok=True) + client_s3 = _utils.client(service_name="s3", session=self._session) + client_s3.download_file(Bucket=bucket, Key=key, Filename=filename) + return filename + + +class ImageS3Dataset(_S3PartitionedDataset): + """PyTorch Amazon S3 Image Dataset.""" + + def __init__(self, path: Union[str, List[str]], suffix: str, boto3_session: boto3.Session): + r"""PyTorch Amazon S3 Image Dataset. + + ImageS3Dataset assumes images are patitioned (within class= folders) in Amazon S3. + Each lisited object will be loaded by default Pillow library. + + Note + ---- + Assumes Images are stored with the following structure: + + + :: + + bucket + ├── class=0 + │ ├── img0.jpeg + │ └── img1.jpeg + └── class=1 + ├── img2.jpeg + └── img3.jpeg + + Parameters + ---------- + path : Union[str, List[str]] + S3 prefix (e.g. s3://bucket/prefix) or list of S3 objects paths (e.g. [s3://bucket/key0, s3://bucket/key1]). + suffix: str, optional + S3 suffix filtering of object keys (i.e. suffix=".png" -> s3://\*.png). + boto3_session : boto3.Session(), optional + Boto3 Session. The default boto3 session will be used if boto3_session receive None. + + Returns + ------- + torch.utils.data.Dataset + + Examples + -------- + >>> import awswrangler as wr + >>> ds = wr.torch.ImageS3Dataset('s3://bucket/path') + + """ + super(ImageS3Dataset, self).__init__(path, suffix, boto3_session) + + def _data_fn(self, data: io.BytesIO) -> Any: + image = Image.open(data) + tensor = to_tensor(image) + return tensor + + +class S3IterableDataset(IterableDataset, _BaseS3Dataset): # pylint: disable=abstract-method + r"""PyTorch Amazon S3 Iterable Dataset. + + Parameters + ---------- + path : Union[str, List[str]] + S3 prefix (e.g. s3://bucket/prefix) or list of S3 objects paths (e.g. [s3://bucket/key0, s3://bucket/key1]). + suffix: str, optional + S3 suffix filtering of object keys (i.e. suffix=".png" -> s3://\*.png). + boto3_session : boto3.Session(), optional + Boto3 Session. The default boto3 session will be used if boto3_session receive None. + + Returns + ------- + torch.utils.data.Dataset + + Examples + -------- + >>> import awswrangler as wr + >>> ds = wr.torch.S3IterableDataset('s3://bucket/path') + + """ + + def __iter__(self) -> Union[Iterator[torch.Tensor], Iterator[Tuple[torch.Tensor, torch.Tensor]]]: + """Iterate over data returning tensors or expanding Iterables.""" + for path in self._paths: + data = self._fetch_data(path) + data = self._load_data(data, path) + + if isinstance(data, torch.Tensor): + pass + elif isinstance(data, Iterable) and all([isinstance(d, torch.Tensor) for d in data]): + data = zip(*data) + else: # pragma: no cover + raise NotImplementedError(f"ERROR: Type: {type(data)} has not been implemented!") + for d in data: + yield d + + +class SQLDataset(IterableDataset): # pylint: disable=too-few-public-methods,abstract-method + """Pytorch Iterable SQL Dataset.""" + + def __init__( + self, + sql: str, + con: sqlalchemy.engine.Engine, + label_col: Optional[Union[int, str]] = None, + chunksize: Optional[int] = None, + ): + """Pytorch Iterable SQL Dataset. + + Support for **Redshift**, **PostgreSQL** and **MySQL**. + + Parameters + ---------- + sql : str + 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() + label_col : int, optional + Label column number. + chunksize : int, optional + The chunksize determines que number of rows to be retrived from the database at each time. + + Returns + ------- + torch.utils.data.dataset.IterableDataset + + Examples + -------- + >>> import awswrangler as wr + >>> con = wr.catalog.get_engine("aws-data-wrangler-postgresql") + >>> ds = wr.torch.SQLDataset('select * from public.tutorial', con=con) + + """ + super().__init__() + self._sql = sql + self._con = con + self._label_col = label_col + self._chunksize = chunksize + + def __iter__(self) -> Union[Iterator[torch.Tensor], Iterator[Tuple[torch.Tensor, torch.Tensor]]]: + """Iterate over the Dataset.""" + if torch.utils.data.get_worker_info() is not None: # type: ignore + raise NotImplementedError() # pragma: no cover + db._validate_engine(con=self._con) # pylint: disable=protected-access + with self._con.connect() as con: + cursor: Any = con.execute(self._sql) + if (self._label_col is not None) and isinstance(self._label_col, str): + label_col: Optional[int] = list(cursor.keys()).index(self._label_col) + else: + label_col = self._label_col + _logger.debug("label_col: %s", label_col) + if self._chunksize is None: + return SQLDataset._records2tensor(records=cursor.fetchall(), label_col=label_col) + return self._iterate_cursor(cursor=cursor, chunksize=self._chunksize, label_col=label_col) + + @staticmethod + def _iterate_cursor( + cursor: Any, chunksize: int, label_col: Optional[int] = None + ) -> Union[Iterator[torch.Tensor], Iterator[Tuple[torch.Tensor, torch.Tensor]]]: + while True: + records = cursor.fetchmany(chunksize) + if not records: + break + yield from SQLDataset._records2tensor(records=records, label_col=label_col) + + @staticmethod + def _records2tensor( + records: List[Tuple[Any]], label_col: Optional[int] = None + ) -> Union[Iterator[torch.Tensor], Iterator[Tuple[torch.Tensor, torch.Tensor]]]: # pylint: disable=unused-argument + for row in records: + if label_col is None: + arr_data: np.ndarray = np.array(row, dtype=np.float) + yield torch.as_tensor(arr_data, dtype=torch.float) # pylint: disable=no-member + else: + arr_data = np.array(row[:label_col] + row[label_col + 1 :], dtype=np.float) # noqa: E203 + arr_label: np.ndarray = np.array(row[label_col], dtype=np.long) + ts_data: torch.Tensor = torch.as_tensor(arr_data, dtype=torch.float) # pylint: disable=no-member + ts_label: torch.Tensor = torch.as_tensor(arr_label, dtype=torch.long) # pylint: disable=no-member + yield ts_data, ts_label diff --git a/building/build-docs.sh b/building/build-docs.sh index c32c20aa0..8c807b485 100755 --- a/building/build-docs.sh +++ b/building/build-docs.sh @@ -4,4 +4,4 @@ set -ex pushd .. rm -rf docs/build docs/source/stubs make -C docs/ html -doc8 --ignore D005 docs/source +doc8 --ignore D005,D002 docs/source diff --git a/docs/source/api.rst b/docs/source/api.rst index 897fc7a3e..c92e735da 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -3,6 +3,19 @@ API Reference ============= +PyTorch +------- + +.. currentmodule:: awswrangler.torch + +.. autosummary:: + :toctree: stubs + + AudioS3Dataset + ImageS3Dataset + S3IterableDataset + SQLDataset + Amazon S3 --------- @@ -16,6 +29,7 @@ Amazon S3 does_object_exist get_bucket_region list_objects + list_directories read_csv read_fwf read_json @@ -112,8 +126,11 @@ EMR get_cluster_state terminate_cluster submit_step + submit_spark_step + submit_ecr_credentials_refresh submit_steps build_step + build_spark_step get_step_state CloudWatch Logs diff --git a/docs/source/index.rst b/docs/source/index.rst index 2528a6032..0a9059392 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -1,4 +1,4 @@ -.. note:: We just released a new major version `1.0` with breaking changes. Please make sure that all your old projects has dependencies frozen on the desired version (e.g. `pip install awswrangler==0.3.2`). +.. note:: Due the new major version `1.*.*` with breaking changes, please make sure that all your old projects has dependencies frozen on the desired version (e.g. `pip install awswrangler==0.3.2`). Quick Start ----------- diff --git a/requirements-dev.txt b/requirements-dev.txt index 3fdd3cdf3..e6e788815 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,20 +1,21 @@ black~=19.3b0 -pylint~=2.4.4 +pylint~=2.5.2 flake8~=3.7.9 mypy~=0.770 isort~=4.3.21 pydocstyle~=5.0.2 doc8~=0.8.0 -tox~=3.14.6 +tox~=3.15.0 pytest~=5.4.1 pytest-cov~=2.8.1 -pytest-xdist~=1.31.0 +pytest-xdist~=1.32.0 scikit-learn~=0.22.1 -awscli>=1.18.22 -cfn-lint~=0.29.5 -cfn-flip~=1.2.2 +awscli>=1.18.0 +cfn-lint~=0.30.1 +cfn-flip~=1.2.3 twine~=3.1.1 wheel~=0.34.2 -sphinx~=3.0.1 +sphinx~=3.0.3 sphinx_bootstrap_theme~=0.7.1 -moto~=1.3.14 \ No newline at end of file +moto~=1.3.14 +jupyterlab~=2.1.2 \ No newline at end of file diff --git a/requirements-torch.txt b/requirements-torch.txt new file mode 100644 index 000000000..20f8cdba9 --- /dev/null +++ b/requirements-torch.txt @@ -0,0 +1,4 @@ +torch~=1.5.0 +torchvision~=0.6.0 +torchaudio~=0.5.0 +Pillow~=7.1.0 diff --git a/requirements.txt b/requirements.txt index ec72c05b7..c6ff840d4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,10 +1,10 @@ -numpy~=1.18.1 -pandas~=1.0.3 -pyarrow~=0.16.0 -boto3>=1.12.22 -botocore>=1.15.22 -s3fs~=0.4.2 -psycopg2-binary~=2.8.5 -pymysql~=0.9.3 -SQLAlchemy==1.3.13 -sqlalchemy-redshift~=0.7.7 \ No newline at end of file +boto3>=1.12.0 +botocore>=1.15.0 +numpy~=1.18.0 +pandas~=1.0.0 +pyarrow~=0.17.0 +s3fs~=0.4.0 +psycopg2-binary~=2.8.0 +pymysql~=0.9.0 +sqlalchemy-redshift~=0.7.0 +SQLAlchemy==1.3.13 \ No newline at end of file diff --git a/setup-dev-env.sh b/setup-dev-env.sh index 692724ee0..c9c2e9902 100755 --- a/setup-dev-env.sh +++ b/setup-dev-env.sh @@ -3,5 +3,4 @@ set -ex pip install --upgrade pip pip install -r requirements-dev.txt -pip install -r requirements.txt -pip install -e . +pip install -e ".[torch]" diff --git a/setup.py b/setup.py index dbd7baa5d..f9fdc6107 100644 --- a/setup.py +++ b/setup.py @@ -24,4 +24,7 @@ include_package_data=True, python_requires=">=3.6, <3.9", install_requires=[open("requirements.txt").read().strip().split("\n")], + extras_require={ + "torch": open("requirements-torch.txt").read().strip().split("\n") + } ) diff --git a/testing/cloudformation.yaml b/testing/cloudformation.yaml index d709f1861..1b4b90f08 100644 --- a/testing/cloudformation.yaml +++ b/testing/cloudformation.yaml @@ -96,6 +96,15 @@ Resources: PolicyDocument: Version: 2012-10-17 Statement: + - Effect: Allow + Action: + - kms:Encrypt + - kms:Decrypt + - kms:GenerateDataKey + Resource: + - Fn::GetAtt: + - KmsKey + - Arn - Effect: Allow Action: - s3:Get* diff --git a/testing/run-validations.sh b/testing/run-validations.sh index 966038ec9..d32fc7808 100755 --- a/testing/run-validations.sh +++ b/testing/run-validations.sh @@ -9,7 +9,7 @@ mv temp.yaml cloudformation.yaml pushd .. black --line-length 120 --target-version py36 awswrangler testing/test_awswrangler isort -rc --line-width 120 awswrangler testing/test_awswrangler -pydocstyle awswrangler/ --add-ignore=D204 +pydocstyle awswrangler/ --add-ignore=D204,D403 mypy awswrangler flake8 setup.py awswrangler testing/test_awswrangler pylint -j 0 awswrangler diff --git a/testing/test_awswrangler/test_cloudwatch.py b/testing/test_awswrangler/test_cloudwatch.py index f59b8b3dd..eced7a754 100644 --- a/testing/test_awswrangler/test_cloudwatch.py +++ b/testing/test_awswrangler/test_cloudwatch.py @@ -48,7 +48,7 @@ def loggroup(cloudformation_outputs): def test_query_cancelled(loggroup): client_logs = boto3.client("logs") query_id = wr.cloudwatch.start_query( - log_group_names=[loggroup], query="fields @timestamp, @message | sort @timestamp desc | limit 5" + log_group_names=[loggroup], query="fields @timestamp, @message | sort @timestamp desc" ) client_logs.stop_query(queryId=query_id) with pytest.raises(exceptions.QueryCancelled): diff --git a/testing/test_awswrangler/test_data_lake.py b/testing/test_awswrangler/test_data_lake.py index afa2a8307..f95c691fb 100644 --- a/testing/test_awswrangler/test_data_lake.py +++ b/testing/test_awswrangler/test_data_lake.py @@ -3,6 +3,7 @@ import gzip import logging import lzma +import math from io import BytesIO, TextIOWrapper import boto3 @@ -73,10 +74,7 @@ def workgroup0(bucket): client.create_work_group( Name=wkg_name, Configuration={ - "ResultConfiguration": { - "OutputLocation": f"s3://{bucket}/athena_workgroup0/", - "EncryptionConfiguration": {"EncryptionOption": "SSE_S3"}, - }, + "ResultConfiguration": {"OutputLocation": f"s3://{bucket}/athena_workgroup0/"}, "EnforceWorkGroupConfiguration": True, "PublishCloudWatchMetricsEnabled": True, "BytesScannedCutoffPerQuery": 100_000_000, @@ -97,7 +95,10 @@ def workgroup1(bucket): client.create_work_group( Name=wkg_name, Configuration={ - "ResultConfiguration": {"OutputLocation": f"s3://{bucket}/athena_workgroup1/"}, + "ResultConfiguration": { + "OutputLocation": f"s3://{bucket}/athena_workgroup1/", + "EncryptionConfiguration": {"EncryptionOption": "SSE_S3"}, + }, "EnforceWorkGroupConfiguration": True, "PublishCloudWatchMetricsEnabled": True, "BytesScannedCutoffPerQuery": 100_000_000, @@ -108,7 +109,57 @@ def workgroup1(bucket): yield wkg_name +@pytest.fixture(scope="module") +def workgroup2(bucket, kms_key): + wkg_name = "awswrangler_test_2" + client = boto3.client("athena") + wkgs = client.list_work_groups() + wkgs = [x["Name"] for x in wkgs["WorkGroups"]] + if wkg_name not in wkgs: + client.create_work_group( + Name=wkg_name, + Configuration={ + "ResultConfiguration": { + "OutputLocation": f"s3://{bucket}/athena_workgroup2/", + "EncryptionConfiguration": {"EncryptionOption": "SSE_KMS", "KmsKey": kms_key}, + }, + "EnforceWorkGroupConfiguration": False, + "PublishCloudWatchMetricsEnabled": True, + "BytesScannedCutoffPerQuery": 100_000_000, + "RequesterPaysEnabled": False, + }, + Description="AWS Data Wrangler Test WorkGroup Number 2", + ) + yield wkg_name + + +@pytest.fixture(scope="module") +def workgroup3(bucket, kms_key): + wkg_name = "awswrangler_test_3" + client = boto3.client("athena") + wkgs = client.list_work_groups() + wkgs = [x["Name"] for x in wkgs["WorkGroups"]] + if wkg_name not in wkgs: + client.create_work_group( + Name=wkg_name, + Configuration={ + "ResultConfiguration": { + "OutputLocation": f"s3://{bucket}/athena_workgroup3/", + "EncryptionConfiguration": {"EncryptionOption": "SSE_KMS", "KmsKey": kms_key}, + }, + "EnforceWorkGroupConfiguration": True, + "PublishCloudWatchMetricsEnabled": True, + "BytesScannedCutoffPerQuery": 100_000_000, + "RequesterPaysEnabled": False, + }, + Description="AWS Data Wrangler Test WorkGroup Number 3", + ) + yield wkg_name + + def test_athena_ctas(bucket, database, kms_key): + wr.s3.delete_objects(path=f"s3://{bucket}/test_athena_ctas/") + wr.s3.delete_objects(path=f"s3://{bucket}/test_athena_ctas_result/") df = get_df_list() columns_types, partitions_types = wr.catalog.extract_athena_types(df=df, partition_cols=["par0", "par1"]) assert len(columns_types) == 16 @@ -127,6 +178,9 @@ def test_athena_ctas(bucket, database, kms_key): partition_cols=["par0", "par1"], )["paths"] wr.s3.wait_objects_exist(paths=paths) + dirs = wr.s3.list_directories(path=f"s3://{bucket}/test_athena_ctas/") + for d in dirs: + assert d.startswith(f"s3://{bucket}/test_athena_ctas/par0=") df = wr.s3.read_parquet_table(table="test_athena_ctas", database=database) assert len(df.index) == 3 ensure_data_types(df=df, has_list=True) @@ -137,14 +191,51 @@ def test_athena_ctas(bucket, database, kms_key): encryption="SSE_KMS", kms_key=kms_key, s3_output=f"s3://{bucket}/test_athena_ctas_result", + keep_files=False, ) assert len(df.index) == 3 ensure_data_types(df=df, has_list=True) + temp_table = "test_athena_ctas2" + s3_output = f"s3://{bucket}/s3_output/" + final_destination = f"{s3_output}{temp_table}/" + + # keep_files=False + wr.s3.delete_objects(path=s3_output) + dfs = wr.athena.read_sql_query( + sql=f"SELECT * FROM test_athena_ctas", + database=database, + ctas_approach=True, + chunksize=1, + keep_files=False, + ctas_temp_table_name=temp_table, + s3_output=s3_output, + ) + assert wr.catalog.does_table_exist(database=database, table=temp_table) is False + assert len(wr.s3.list_objects(path=s3_output)) > 2 + assert len(wr.s3.list_objects(path=final_destination)) > 0 + for df in dfs: + ensure_data_types(df=df, has_list=True) + assert len(wr.s3.list_objects(path=s3_output)) == 0 + + # keep_files=True + wr.s3.delete_objects(path=s3_output) dfs = wr.athena.read_sql_query( - sql=f"SELECT * FROM test_athena_ctas", database=database, ctas_approach=True, chunksize=1 + sql=f"SELECT * FROM test_athena_ctas", + database=database, + ctas_approach=True, + chunksize=2, + keep_files=True, + ctas_temp_table_name=temp_table, + s3_output=s3_output, ) + assert wr.catalog.does_table_exist(database=database, table=temp_table) is False + assert len(wr.s3.list_objects(path=s3_output)) > 2 + assert len(wr.s3.list_objects(path=final_destination)) > 0 for df in dfs: ensure_data_types(df=df, has_list=True) + assert len(wr.s3.list_objects(path=s3_output)) > 2 + + # Cleaning Up wr.catalog.delete_table_if_exists(database=database, table="test_athena_ctas") wr.s3.delete_objects(path=paths) wr.s3.wait_objects_not_exist(paths=paths) @@ -173,12 +264,17 @@ def test_athena(bucket, database, kms_key, workgroup0, workgroup1): encryption="SSE_KMS", kms_key=kms_key, workgroup=workgroup0, + keep_files=False, ) for df2 in dfs: print(df2) ensure_data_types(df=df2) df = wr.athena.read_sql_query( - sql="SELECT * FROM __test_athena", database=database, ctas_approach=False, workgroup=workgroup1 + sql="SELECT * FROM __test_athena", + database=database, + ctas_approach=False, + workgroup=workgroup1, + keep_files=False, ) assert len(df.index) == 3 ensure_data_types(df=df) @@ -252,13 +348,12 @@ def test_fwf(bucket): def test_parquet(bucket): - wr.s3.delete_objects(path=f"s3://{bucket}/test_parquet_file") - wr.s3.delete_objects(path=f"s3://{bucket}/test_parquet_dataset") + wr.s3.delete_objects(path=f"s3://{bucket}/test_parquet/") df_file = pd.DataFrame({"id": [1, 2, 3]}) - path_file = f"s3://{bucket}/test_parquet_file.parquet" + path_file = f"s3://{bucket}/test_parquet/test_parquet_file.parquet" df_dataset = pd.DataFrame({"id": [1, 2, 3], "partition": ["A", "A", "B"]}) df_dataset["partition"] = df_dataset["partition"].astype("category") - path_dataset = f"s3://{bucket}/test_parquet_dataset" + path_dataset = f"s3://{bucket}/test_parquet/test_parquet_dataset" with pytest.raises(wr.exceptions.InvalidArgumentCombination): wr.s3.to_parquet(df=df_file, path=path_file, mode="append") with pytest.raises(wr.exceptions.InvalidCompression): @@ -288,8 +383,7 @@ def test_parquet(bucket): wr.s3.to_parquet( df=df_dataset, path=path_dataset, dataset=True, partition_cols=["partition"], mode="overwrite_partitions" ) - wr.s3.delete_objects(path=f"s3://{bucket}/test_parquet_file") - wr.s3.delete_objects(path=f"s3://{bucket}/test_parquet_dataset") + wr.s3.delete_objects(path=f"s3://{bucket}/test_parquet/") def test_parquet_catalog(bucket, database): @@ -704,7 +798,7 @@ def test_parquet_validate_schema(bucket, database): df2 = pd.DataFrame({"id2": [1, 2, 3], "val": ["foo", "boo", "bar"]}) path_file2 = f"s3://{bucket}/test_parquet_file_validate/1.parquet" wr.s3.to_parquet(df=df2, path=path_file2) - wr.s3.wait_objects_exist(paths=[path_file2]) + wr.s3.wait_objects_exist(paths=[path_file2], use_threads=False) df3 = wr.s3.read_parquet(path=path, validate_schema=False) assert len(df3.index) == 6 assert len(df3.columns) == 3 @@ -1084,3 +1178,192 @@ def test_copy(bucket): wr.s3.delete_objects(path=path) wr.s3.delete_objects(path=path2) + + +@pytest.mark.parametrize("col2", [[1, 1, 1, 1, 1], [1, 2, 3, 4, 5], [1, 1, 1, 1, 2], [1, 2, 2, 2, 2]]) +@pytest.mark.parametrize("chunked", [True, 1, 2, 100]) +def test_parquet_chunked(bucket, database, col2, chunked): + table = f"test_parquet_chunked_{chunked}_{''.join([str(x) for x in col2])}" + path = f"s3://{bucket}/{table}/" + wr.s3.delete_objects(path=path) + values = list(range(5)) + df = pd.DataFrame({"col1": values, "col2": col2}) + paths = wr.s3.to_parquet( + df, path, index=False, dataset=True, database=database, table=table, partition_cols=["col2"], mode="overwrite" + )["paths"] + wr.s3.wait_objects_exist(paths=paths) + + dfs = list(wr.s3.read_parquet(path=path, dataset=True, chunked=chunked)) + assert sum(values) == pd.concat(dfs, ignore_index=True).col1.sum() + if chunked is not True: + assert len(dfs) == int(math.ceil(len(df) / chunked)) + for df2 in dfs[:-1]: + assert chunked == len(df2) + assert chunked >= len(dfs[-1]) + else: + assert len(dfs) == len(set(col2)) + + dfs = list(wr.athena.read_sql_table(database=database, table=table, chunksize=chunked)) + assert sum(values) == pd.concat(dfs, ignore_index=True).col1.sum() + if chunked is not True: + assert len(dfs) == int(math.ceil(len(df) / chunked)) + for df2 in dfs[:-1]: + assert chunked == len(df2) + assert chunked >= len(dfs[-1]) + + wr.s3.delete_objects(path=paths) + assert wr.catalog.delete_table_if_exists(database=database, table=table) is True + + +@pytest.mark.parametrize("workgroup", [None, 0, 1, 2, 3]) +@pytest.mark.parametrize("encryption", [None, "SSE_S3", "SSE_KMS"]) +def test_athena_encryption( + bucket, database, kms_key, encryption, workgroup, workgroup0, workgroup1, workgroup2, workgroup3 +): + kms_key = None if (encryption == "SSE_S3") or (encryption is None) else kms_key + if workgroup == 0: + workgroup = workgroup0 + elif workgroup == 1: + workgroup = workgroup1 + elif workgroup == 2: + workgroup = workgroup2 + elif workgroup == 3: + workgroup = workgroup3 + table = f"test_athena_encryption_{str(encryption).lower()}_{str(workgroup).lower()}" + path = f"s3://{bucket}/{table}/" + wr.s3.delete_objects(path=path) + df = pd.DataFrame({"a": [1, 2], "b": ["foo", "boo"]}) + paths = wr.s3.to_parquet( + df=df, path=path, dataset=True, mode="overwrite", database=database, table=table, s3_additional_kwargs=None + )["paths"] + wr.s3.wait_objects_exist(paths=paths, use_threads=False) + temp_table = table + "2" + s3_output = f"s3://{bucket}/encryptio_s3_output/" + final_destination = f"{s3_output}{temp_table}/" + wr.s3.delete_objects(path=final_destination) + df2 = wr.athena.read_sql_table( + table=table, + ctas_approach=True, + database=database, + encryption=encryption, + workgroup=workgroup, + kms_key=kms_key, + keep_files=True, + ctas_temp_table_name=temp_table, + s3_output=s3_output, + ) + assert wr.catalog.does_table_exist(database=database, table=temp_table) is False + assert len(wr.s3.list_objects(path=s3_output)) > 2 + print(df2) + assert len(df2.index) == 2 + assert len(df2.columns) == 2 + wr.catalog.delete_table_if_exists(database=database, table=table) + wr.s3.delete_objects(path=paths) + + +def test_athena_nested(bucket, database): + table = "test_athena_nested" + path = f"s3://{bucket}/{table}/" + df = pd.DataFrame( + { + "c0": [[1, 2, 3], [4, 5, 6]], + "c1": [[[1, 2], [3, 4]], [[5, 6], [7, 8]]], + "c2": [[["a", "b"], ["c", "d"]], [["e", "f"], ["g", "h"]]], + "c3": [[], [[[[[[[[1]]]]]]]]], + "c4": [{"a": 1}, {"a": 1}], + "c5": [{"a": {"b": {"c": [1, 2]}}}, {"a": {"b": {"c": [3, 4]}}}], + } + ) + paths = wr.s3.to_parquet( + df=df, path=path, index=False, use_threads=True, dataset=True, mode="overwrite", database=database, table=table + )["paths"] + wr.s3.wait_objects_exist(paths=paths) + df2 = wr.athena.read_sql_query(sql=f"SELECT c0, c1, c2, c4 FROM {table}", database=database) + assert len(df2.index) == 2 + assert len(df2.columns) == 4 + + +def test_catalog_versioning(bucket, database): + table = "test_catalog_versioning" + wr.catalog.delete_table_if_exists(database=database, table=table) + path = f"s3://{bucket}/{table}/" + wr.s3.delete_objects(path=path) + + # Version 0 + df = pd.DataFrame({"c0": [1, 2]}) + paths = wr.s3.to_parquet(df=df, path=path, dataset=True, database=database, table=table, mode="overwrite")["paths"] + wr.s3.wait_objects_exist(paths=paths, use_threads=False) + df = wr.athena.read_sql_table(table=table, database=database) + assert len(df.index) == 2 + assert len(df.columns) == 1 + assert str(df.c0.dtype).startswith("Int") + + # Version 1 + df = pd.DataFrame({"c1": ["foo", "boo"]}) + paths = wr.s3.to_parquet( + df=df, path=path, dataset=True, database=database, table=table, mode="overwrite", catalog_versioning=True + )["paths"] + wr.s3.wait_objects_exist(paths=paths, use_threads=False) + df = wr.athena.read_sql_table(table=table, database=database) + assert len(df.index) == 2 + assert len(df.columns) == 1 + assert str(df.c1.dtype) == "string" + + # Version 2 + df = pd.DataFrame({"c1": [1.0, 2.0]}) + paths = wr.s3.to_csv( + df=df, + path=path, + dataset=True, + database=database, + table=table, + mode="overwrite", + catalog_versioning=True, + index=False, + )["paths"] + wr.s3.wait_objects_exist(paths=paths, use_threads=False) + df = wr.athena.read_sql_table(table=table, database=database) + assert len(df.index) == 2 + assert len(df.columns) == 1 + assert str(df.c1.dtype).startswith("float") + + # Version 3 (removing version 2) + df = pd.DataFrame({"c1": [True, False]}) + paths = wr.s3.to_csv( + df=df, + path=path, + dataset=True, + database=database, + table=table, + mode="overwrite", + catalog_versioning=False, + index=False, + )["paths"] + wr.s3.wait_objects_exist(paths=paths, use_threads=False) + df = wr.athena.read_sql_table(table=table, database=database) + assert len(df.index) == 2 + assert len(df.columns) == 1 + assert str(df.c1.dtype).startswith("boolean") + + # Cleaning Up + wr.catalog.delete_table_if_exists(database=database, table=table) + wr.s3.delete_objects(path=path) + + +def test_copy_replacing_filename(bucket): + path = f"s3://{bucket}/test_copy_replacing_filename/" + wr.s3.delete_objects(path=path) + df = pd.DataFrame({"c0": [1, 2]}) + file_path = f"{path}myfile.parquet" + wr.s3.to_parquet(df=df, path=file_path) + wr.s3.wait_objects_exist(paths=[file_path], use_threads=False) + path2 = f"s3://{bucket}/test_copy_replacing_filename2/" + wr.s3.copy_objects( + paths=[file_path], source_path=path, target_path=path2, replace_filenames={"myfile.parquet": "myfile2.parquet"} + ) + expected_file = f"{path2}myfile2.parquet" + wr.s3.wait_objects_exist(paths=[expected_file], use_threads=False) + objs = wr.s3.list_objects(path=path2) + assert objs[0] == expected_file + wr.s3.delete_objects(path=path) + wr.s3.delete_objects(path=path2) diff --git a/testing/test_awswrangler/test_db.py b/testing/test_awswrangler/test_db.py index adcacf4a4..86a57a74d 100644 --- a/testing/test_awswrangler/test_db.py +++ b/testing/test_awswrangler/test_db.py @@ -76,6 +76,11 @@ def external_schema(cloudformation_outputs, parameters, glue_database): yield "aws_data_wrangler_external" +@pytest.fixture(scope="module") +def kms_key_id(cloudformation_outputs): + yield cloudformation_outputs["KmsKeyArn"].split("/", 1)[1] + + @pytest.mark.parametrize("db_type", ["mysql", "redshift", "postgresql"]) def test_sql(parameters, db_type): df = get_df() @@ -386,3 +391,138 @@ def test_redshift_category(bucket, parameters): for df2 in dfs: ensure_data_types_category(df2) wr.s3.delete_objects(path=path) + + +def test_redshift_unload_extras(bucket, parameters, kms_key_id): + table = "test_redshift_unload_extras" + schema = parameters["redshift"]["schema"] + path = f"s3://{bucket}/{table}/" + wr.s3.delete_objects(path=path) + engine = wr.catalog.get_engine(connection=f"aws-data-wrangler-redshift") + df = pd.DataFrame({"id": [1, 2], "name": ["foo", "boo"]}) + wr.db.to_sql(df=df, con=engine, name=table, schema=schema, if_exists="replace", index=False) + paths = wr.db.unload_redshift_to_files( + sql=f"SELECT * FROM {schema}.{table}", + path=path, + con=engine, + iam_role=parameters["redshift"]["role"], + region=wr.s3.get_bucket_region(bucket), + max_file_size=5.0, + kms_key_id=kms_key_id, + partition_cols=["name"], + ) + wr.s3.wait_objects_exist(paths=paths) + df = wr.s3.read_parquet(path=path, dataset=True) + assert len(df.index) == 2 + assert len(df.columns) == 2 + wr.s3.delete_objects(path=path) + df = wr.db.unload_redshift( + sql=f"SELECT * FROM {schema}.{table}", + con=engine, + iam_role=parameters["redshift"]["role"], + path=path, + keep_files=False, + region=wr.s3.get_bucket_region(bucket), + max_file_size=5.0, + kms_key_id=kms_key_id, + ) + assert len(df.index) == 2 + assert len(df.columns) == 2 + wr.s3.delete_objects(path=path) + + +@pytest.mark.parametrize("db_type", ["mysql", "redshift", "postgresql"]) +def test_to_sql_cast(parameters, db_type): + table = "test_to_sql_cast" + schema = parameters[db_type]["schema"] + df = pd.DataFrame( + { + "col": [ + "".join([str(i)[-1] for i in range(1_024)]), + "".join([str(i)[-1] for i in range(1_024)]), + "".join([str(i)[-1] for i in range(1_024)]), + ] + }, + dtype="string", + ) + engine = wr.catalog.get_engine(connection=f"aws-data-wrangler-{db_type}") + wr.db.to_sql( + df=df, + con=engine, + name=table, + schema=schema, + if_exists="replace", + index=False, + index_label=None, + chunksize=None, + method=None, + dtype={"col": sqlalchemy.types.VARCHAR(length=1_024)}, + ) + df2 = wr.db.read_sql_query(sql=f"SELECT * FROM {schema}.{table}", con=engine) + assert df.equals(df2) + + +def test_uuid(parameters): + table = "test_uuid" + schema = parameters["postgresql"]["schema"] + engine = wr.catalog.get_engine(connection=f"aws-data-wrangler-postgresql") + df = pd.DataFrame( + { + "id": [1, 2, 3], + "uuid": [ + "ec0f0482-8d3b-11ea-8b27-8c859043dd95", + "f56ff7c0-8d3b-11ea-be94-8c859043dd95", + "fa043e90-8d3b-11ea-b7e7-8c859043dd95", + ], + } + ) + wr.db.to_sql( + df=df, + con=engine, + name=table, + schema=schema, + if_exists="replace", + index=False, + index_label=None, + chunksize=None, + method=None, + dtype={"uuid": sqlalchemy.dialects.postgresql.UUID}, + ) + df2 = wr.db.read_sql_table(table=table, schema=schema, con=engine) + df["id"] = df["id"].astype("Int64") + df["uuid"] = df["uuid"].astype("string") + assert df.equals(df2) + + +@pytest.mark.parametrize("db_type", ["mysql", "redshift", "postgresql"]) +def test_null(parameters, db_type): + table = "test_null" + schema = parameters[db_type]["schema"] + engine = wr.catalog.get_engine(connection=f"aws-data-wrangler-{db_type}") + df = pd.DataFrame({"id": [1, 2, 3], "nothing": [None, None, None]}) + wr.db.to_sql( + df=df, + con=engine, + name=table, + schema=schema, + if_exists="replace", + index=False, + index_label=None, + chunksize=None, + method=None, + dtype={"nothing": sqlalchemy.types.Integer}, + ) + wr.db.to_sql( + df=df, + con=engine, + name=table, + schema=schema, + if_exists="append", + index=False, + index_label=None, + chunksize=None, + method=None, + ) + df2 = wr.db.read_sql_table(table=table, schema=schema, con=engine) + df["id"] = df["id"].astype("Int64") + assert pd.concat(objs=[df, df], ignore_index=True).equals(df2) diff --git a/testing/test_awswrangler/test_emr.py b/testing/test_awswrangler/test_emr.py index e64329b33..0c0112bf8 100644 --- a/testing/test_awswrangler/test_emr.py +++ b/testing/test_awswrangler/test_emr.py @@ -146,3 +146,41 @@ def test_cluster_single_node(bucket, cloudformation_outputs): wr.emr.submit_steps(cluster_id=cluster_id, steps=steps) wr.emr.terminate_cluster(cluster_id=cluster_id) wr.s3.delete_objects(f"s3://{bucket}/emr-logs/") + + +def test_default_logging_path(cloudformation_outputs): + path = wr.emr._get_default_logging_path(subnet_id=cloudformation_outputs["SubnetId"]) + assert path.startswith("s3://aws-logs-") + assert path.endswith("/elasticmapreduce/") + with pytest.raises(wr.exceptions.InvalidArgumentCombination): + wr.emr._get_default_logging_path() + + +def test_docker(bucket, cloudformation_outputs): + cluster_id = wr.emr.create_cluster( + subnet_id=cloudformation_outputs["SubnetId"], + docker=True, + custom_classifications=[ + { + "Classification": "livy-conf", + "Properties": { + "livy.spark.master": "yarn", + "livy.spark.deploy-mode": "cluster", + "livy.server.session.timeout": "16h", + }, + } + ], + steps=[wr.emr.build_step("spark-submit --deploy-mode cluster s3://bucket/emr.py")], + ) + wr.emr.submit_ecr_credentials_refresh(cluster_id, path=f"s3://{bucket}/emr/") + wr.emr.submit_steps( + cluster_id=cluster_id, + steps=[ + wr.emr.build_spark_step( + path=f"s3://{bucket}/emr/test_docker.py", + docker_image="123456789123.dkr.ecr.us-east-1.amazonaws.com/docker-emr:docker-emr", + ) + ], + ) + wr.emr.submit_spark_step(cluster_id=cluster_id, path=f"s3://{bucket}/emr/test_docker.py") + wr.emr.terminate_cluster(cluster_id=cluster_id) diff --git a/testing/test_awswrangler/test_metadata.py b/testing/test_awswrangler/test_metadata.py index d076c0d94..88e71c5a3 100644 --- a/testing/test_awswrangler/test_metadata.py +++ b/testing/test_awswrangler/test_metadata.py @@ -2,7 +2,7 @@ def test_metadata(): - assert wr.__version__ == "1.0.4" + assert wr.__version__ == "1.1.0" assert wr.__title__ == "awswrangler" assert wr.__description__ == "Pandas on AWS." assert wr.__license__ == "Apache License 2.0" diff --git a/testing/test_awswrangler/test_moto.py b/testing/test_awswrangler/test_moto.py index db12dbe1a..2adc7aec8 100644 --- a/testing/test_awswrangler/test_moto.py +++ b/testing/test_awswrangler/test_moto.py @@ -20,6 +20,21 @@ def emr(): yield True +@pytest.fixture(scope="module") +def sts(): + with moto.mock_sts(): + yield True + + +@pytest.fixture(scope="module") +def subnet(): + with moto.mock_ec2(): + ec2 = boto3.resource("ec2", region_name="us-west-1") + vpc = ec2.create_vpc(CidrBlock="10.0.0.0/16") + subnet = ec2.create_subnet(VpcId=vpc.id, CidrBlock="10.0.0.0/24", AvailabilityZone="us-west-1a") + yield subnet.id + + def test_csv(s3): path = "s3://bucket/test.csv" wr.s3.to_csv(df=get_df_csv(), path=path, index=False) @@ -37,12 +52,13 @@ def test_parquet(s3): assert len(df.columns) == 18 -def test_emr(s3, emr): +def test_emr(s3, emr, sts, subnet): + session = boto3.Session(region_name="us-west-1") cluster_id = wr.emr.create_cluster( cluster_name="wrangler_cluster", logging_s3_path="s3://bucket/emr-logs/", emr_release="emr-5.29.0", - subnet_id="foo", + subnet_id=subnet, emr_ec2_role="EMR_EC2_DefaultRole", emr_role="EMR_DefaultRole", instance_type_master="m5.xlarge", @@ -87,11 +103,12 @@ def test_emr(s3, emr): termination_protected=False, spark_pyarrow=False, tags={"foo": "boo", "bar": "xoo"}, + boto3_session=session, ) - wr.emr.get_cluster_state(cluster_id=cluster_id) + wr.emr.get_cluster_state(cluster_id=cluster_id, boto3_session=session) steps = [] for cmd in ['echo "Hello"', "ls -la"]: steps.append(wr.emr.build_step(name=cmd, command=cmd)) - wr.emr.submit_steps(cluster_id=cluster_id, steps=steps) - wr.emr.terminate_cluster(cluster_id=cluster_id) + wr.emr.submit_steps(cluster_id=cluster_id, steps=steps, boto3_session=session) + wr.emr.terminate_cluster(cluster_id=cluster_id, boto3_session=session) wr.s3.delete_objects("s3://bucket/emr-logs/") diff --git a/testing/test_awswrangler/test_torch.py b/testing/test_awswrangler/test_torch.py new file mode 100644 index 000000000..a19dd64b5 --- /dev/null +++ b/testing/test_awswrangler/test_torch.py @@ -0,0 +1,282 @@ +import io +import logging +import re + +import boto3 +import numpy as np +import pandas as pd +import pytest +import torch +import torchaudio +from PIL import Image +from torch.utils.data import DataLoader +from torchvision.transforms.functional import to_tensor + +import awswrangler as wr + +logging.basicConfig(level=logging.INFO, format="[%(asctime)s][%(levelname)s][%(name)s][%(funcName)s] %(message)s") +logging.getLogger("awswrangler").setLevel(logging.DEBUG) +logging.getLogger("botocore.credentials").setLevel(logging.CRITICAL) + + +@pytest.fixture(scope="module") +def cloudformation_outputs(): + response = boto3.client("cloudformation").describe_stacks(StackName="aws-data-wrangler-test") + outputs = {} + for output in response.get("Stacks")[0].get("Outputs"): + outputs[output.get("OutputKey")] = output.get("OutputValue") + yield outputs + + +@pytest.fixture(scope="module") +def bucket(cloudformation_outputs): + if "BucketName" in cloudformation_outputs: + bucket = cloudformation_outputs["BucketName"] + else: + raise Exception("You must deploy/update the test infrastructure (CloudFormation)") + yield bucket + + +@pytest.fixture(scope="module") +def parameters(cloudformation_outputs): + parameters = dict(postgresql={}, mysql={}, redshift={}) + parameters["postgresql"]["host"] = cloudformation_outputs["PostgresqlAddress"] + parameters["postgresql"]["port"] = 3306 + parameters["postgresql"]["schema"] = "public" + parameters["postgresql"]["database"] = "postgres" + parameters["mysql"]["host"] = cloudformation_outputs["MysqlAddress"] + parameters["mysql"]["port"] = 3306 + parameters["mysql"]["schema"] = "test" + parameters["mysql"]["database"] = "test" + parameters["redshift"]["host"] = cloudformation_outputs["RedshiftAddress"] + parameters["redshift"]["port"] = cloudformation_outputs["RedshiftPort"] + parameters["redshift"]["identifier"] = cloudformation_outputs["RedshiftIdentifier"] + parameters["redshift"]["schema"] = "public" + parameters["redshift"]["database"] = "test" + parameters["redshift"]["role"] = cloudformation_outputs["RedshiftRole"] + parameters["password"] = cloudformation_outputs["DatabasesPassword"] + parameters["user"] = "test" + yield parameters + + +@pytest.mark.parametrize("chunksize", [None, 1, 10]) +@pytest.mark.parametrize("db_type", ["mysql", "redshift", "postgresql"]) +def test_torch_sql(parameters, db_type, chunksize): + schema = parameters[db_type]["schema"] + table = f"test_torch_sql_{db_type}_{str(chunksize).lower()}" + engine = wr.catalog.get_engine(connection=f"aws-data-wrangler-{db_type}") + wr.db.to_sql( + df=pd.DataFrame({"a": [1.0, 2.0, 3.0], "b": [4.0, 5.0, 6.0]}), + con=engine, + name=table, + schema=schema, + if_exists="replace", + index=False, + index_label=None, + chunksize=None, + method=None, + ) + ds = list(wr.torch.SQLDataset(f"SELECT * FROM {schema}.{table}", con=engine, chunksize=chunksize)) + assert torch.all(ds[0].eq(torch.tensor([1.0, 4.0]))) + assert torch.all(ds[1].eq(torch.tensor([2.0, 5.0]))) + assert torch.all(ds[2].eq(torch.tensor([3.0, 6.0]))) + + +@pytest.mark.parametrize("chunksize", [None, 1, 10]) +@pytest.mark.parametrize("db_type", ["mysql", "redshift", "postgresql"]) +@pytest.mark.parametrize("label_col", [2, "c"]) +def test_torch_sql_label(parameters, db_type, chunksize, label_col): + schema = parameters[db_type]["schema"] + table = f"test_torch_sql_label_{db_type}_{str(chunksize).lower()}_{label_col}" + engine = wr.catalog.get_engine(connection=f"aws-data-wrangler-{db_type}") + wr.db.to_sql( + df=pd.DataFrame({"a": [1.0, 2.0, 3.0], "b": [4.0, 5.0, 6.0], "c": [7, 8, 9]}), + con=engine, + name=table, + schema=schema, + if_exists="replace", + index=False, + index_label=None, + chunksize=None, + method=None, + ) + ts = list( + wr.torch.SQLDataset(f"SELECT * FROM {schema}.{table}", con=engine, chunksize=chunksize, label_col=label_col) + ) + assert torch.all(ts[0][0].eq(torch.tensor([1.0, 4.0]))) + assert torch.all(ts[0][1].eq(torch.tensor([7], dtype=torch.long))) + assert torch.all(ts[1][0].eq(torch.tensor([2.0, 5.0]))) + assert torch.all(ts[1][1].eq(torch.tensor([8], dtype=torch.long))) + assert torch.all(ts[2][0].eq(torch.tensor([3.0, 6.0]))) + assert torch.all(ts[2][1].eq(torch.tensor([9], dtype=torch.long))) + + +def test_torch_image_s3(bucket): + folder = "test_torch_image_s3" + path = f"s3://{bucket}/{folder}/" + wr.s3.delete_objects(path=path, boto3_session=boto3.Session()) + s3 = boto3.client("s3") + ref_label = 0 + s3.put_object( + Body=open("docs/source/_static/logo.png", "rb").read(), + Bucket=bucket, + Key=f"{folder}/class={ref_label}/logo.png", + ContentType="image/png", + ) + wr.s3.wait_objects_exist(paths=[f"s3://{bucket}/{folder}/class={ref_label}/logo.png"]) + ds = wr.torch.ImageS3Dataset(path=path, suffix="png", boto3_session=boto3.Session()) + image, label = ds[0] + assert image.shape == torch.Size([4, 494, 1636]) + assert label == torch.tensor(ref_label, dtype=torch.int) + wr.s3.delete_objects(path=path) + + +@pytest.mark.parametrize("drop_last", [True, False]) +def test_torch_image_s3_loader(bucket, drop_last): + folder = f"test_torch_image_s3_loader_{str(drop_last).lower()}" + path = f"s3://{bucket}/{folder}/" + wr.s3.delete_objects(path=path) + client_s3 = boto3.client("s3") + labels = np.random.randint(0, 4, size=(8,)) + for i, label in enumerate(labels): + client_s3.put_object( + Body=open("./docs/source/_static/logo.png", "rb").read(), + Bucket=bucket, + Key=f"{folder}/class={label}/logo{i}.png", + ContentType="image/png", + ) + wr.s3.wait_objects_exist(paths=[f"s3://{bucket}/{folder}/class={label}/logo{i}.png"]) + ds = wr.torch.ImageS3Dataset(path=path, suffix="png", boto3_session=boto3.Session()) + batch_size = 2 + num_train = len(ds) + indices = list(range(num_train)) + loader = DataLoader( + ds, + batch_size=batch_size, + num_workers=4, + sampler=torch.utils.data.sampler.RandomSampler(indices), + drop_last=drop_last, + ) + for i, (image, label) in enumerate(loader): + assert image.shape == torch.Size([batch_size, 4, 494, 1636]) + assert label.dtype == torch.int64 + wr.s3.delete_objects(path=path) + + +def test_torch_lambda_s3(bucket): + path = f"s3://{bucket}/test_torch_lambda_s3/" + wr.s3.delete_objects(path=path) + s3 = boto3.client("s3") + ref_label = 0 + s3.put_object( + Body=open("./docs/source/_static/logo.png", "rb").read(), + Bucket=bucket, + Key=f"test_torch_lambda_s3/class={ref_label}/logo.png", + ContentType="image/png", + ) + wr.s3.wait_objects_exist(paths=[f"s3://{bucket}/test_torch_lambda_s3/class={ref_label}/logo.png"]) + ds = wr.torch.LambdaS3Dataset( + path=path, + suffix="png", + boto3_session=boto3.Session(), + data_fn=lambda x: to_tensor(Image.open(x)), + label_fn=lambda x: int(re.findall(r"/class=(.*?)/", x)[-1]), + ) + image, label = ds[0] + assert image.shape == torch.Size([4, 494, 1636]) + assert label == torch.tensor(ref_label, dtype=torch.int) + wr.s3.delete_objects(path=path) + + +def test_torch_audio_s3(bucket): + size = (1, 8_000 * 5) + audio = torch.randint(low=-25, high=25, size=size) / 100.0 + audio_file = "/tmp/amazing_sound.wav" + torchaudio.save(audio_file, audio, 8_000) + folder = "test_torch_audio_s3" + path = f"s3://{bucket}/{folder}/" + wr.s3.delete_objects(path=path) + s3 = boto3.client("s3") + ref_label = 0 + s3.put_object( + Body=open(audio_file, "rb").read(), + Bucket=bucket, + Key=f"{folder}/class={ref_label}/amazing_sound.wav", + ContentType="audio/wav", + ) + wr.s3.wait_objects_exist(paths=[f"s3://{bucket}/{folder}/class={ref_label}/amazing_sound.wav"]) + s3_audio_file = f"{bucket}/test_torch_audio_s3/class={ref_label}/amazing_sound.wav" + ds = wr.torch.AudioS3Dataset(path=s3_audio_file, suffix="wav") + loader = DataLoader(ds, batch_size=1) + for (audio, rate), label in loader: + assert audio.shape == torch.Size((1, *size)) + wr.s3.delete_objects(path=path) + + +# def test_torch_s3_file_dataset(bucket): +# cifar10 = "s3://fast-ai-imageclas/cifar10.tgz" +# batch_size = 64 +# for image, label in DataLoader( +# wr.torch.S3FilesDataset(cifar10), +# batch_size=batch_size, +# ): +# assert image.shape == torch.Size([batch_size, 3, 32, 32]) +# assert label.dtype == torch.int64 +# break + + +@pytest.mark.parametrize("drop_last", [True, False]) +def test_torch_s3_iterable(bucket, drop_last): + folder = f"test_torch_s3_iterable_{str(drop_last).lower()}" + path = f"s3://{bucket}/{folder}/" + wr.s3.delete_objects(path=path) + batch_size = 32 + client_s3 = boto3.client("s3") + for i in range(3): + batch = torch.randn(100, 3, 32, 32) + buff = io.BytesIO() + torch.save(batch, buff) + buff.seek(0) + client_s3.put_object(Body=buff.read(), Bucket=bucket, Key=f"{folder}/file{i}.pt") + wr.s3.wait_objects_exist(paths=[f"s3://{bucket}/{folder}/file{i}.pt"]) + + for image in DataLoader( + wr.torch.S3IterableDataset(path=f"s3://{bucket}/{folder}/file"), batch_size=batch_size, drop_last=drop_last + ): + if drop_last: + assert image.shape == torch.Size([batch_size, 3, 32, 32]) + else: + assert image[0].shape == torch.Size([3, 32, 32]) + + wr.s3.delete_objects(path=path) + + +@pytest.mark.parametrize("drop_last", [True, False]) +def test_torch_s3_iterable_with_labels(bucket, drop_last): + folder = f"test_torch_s3_iterable_with_labels_{str(drop_last).lower()}" + path = f"s3://{bucket}/{folder}/" + wr.s3.delete_objects(path=path) + batch_size = 32 + client_s3 = boto3.client("s3") + for i in range(3): + batch = (torch.randn(100, 3, 32, 32), torch.randint(2, size=(100,))) + buff = io.BytesIO() + torch.save(batch, buff) + buff.seek(0) + client_s3.put_object(Body=buff.read(), Bucket=bucket, Key=f"{folder}/file{i}.pt") + wr.s3.wait_objects_exist(paths=[f"s3://{bucket}/{folder}/file{i}.pt"]) + + for images, labels in DataLoader( + wr.torch.S3IterableDataset(path=f"s3://{bucket}/{folder}/file"), batch_size=batch_size, drop_last=drop_last + ): + if drop_last: + assert images.shape == torch.Size([batch_size, 3, 32, 32]) + assert labels.dtype == torch.int64 + assert labels.shape == torch.Size([batch_size]) + + else: + assert images[0].shape == torch.Size([3, 32, 32]) + assert labels[0].dtype == torch.int64 + assert labels[0].shape == torch.Size([]) + + wr.s3.delete_objects(path=path) diff --git a/tox.ini b/tox.ini index 9768fd204..f2bb572c2 100644 --- a/tox.ini +++ b/tox.ini @@ -6,10 +6,13 @@ deps = pytest pytest-xdist moto -commands = pytest -n 8 testing/test_awswrangler + -rrequirements-torch.txt +commands = + pytest -n 8 testing/test_awswrangler [testenv:py36] deps = {[testenv]deps} pytest-cov -commands = pytest --cov=awswrangler -n 8 testing/test_awswrangler +commands = + pytest --cov=awswrangler -n 8 testing/test_awswrangler diff --git a/tutorials/14 - PyTorch.ipynb b/tutorials/14 - PyTorch.ipynb new file mode 100644 index 000000000..b7af04627 --- /dev/null +++ b/tutorials/14 - PyTorch.ipynb @@ -0,0 +1,330 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "[![AWS Data Wrangler](_static/logo.png \"AWS Data Wrangler\")](https://github.com/awslabs/aws-data-wrangler)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# PyTorch" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Table of Contents\n", + "* [1.Defining Training Function](#1.-Defining-Training-Function)\n", + "* [2.Training From Amazon S3](#2.-Traoning-From-Amazon-S3)\n", + "\t* [2.1 Writing PyTorch Dataset to S3](#2.1-Writing-PyTorch-Dataset-to-S3)\n", + "\t* [2.2 Training Network](#2.2-Training-Network)\n", + "* [3. Training From SQL Query](#3.-Training-From-SQL-Query)\n", + "\t* [3.1 Writing Data to SQL Database](#3.1-Writing-Data-to-SQL-Database)\n", + "\t* [3.3 Training Network From SQL](#3.3-Reading-single-JSON-file)\n", + "* [4. Creating Custom S3 Dataset](#4.-Creating-Custom-S3-Dataset)\n", + "\t* [4.1 Creating Custom PyTorch Dataset](#4.1-Creating-Custom-PyTorch-Dataset)\n", + "\t* [4.2 Writing Data to S3](#4.2-Writing-Data-to-S3)\n", + "\t* [4.3 Training Network](#4.4-Training-Network)\n", + "* [5. Delete objects](#5.-Delete-objects)" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "import io\n", + "\n", + "import boto3\n", + "import torch\n", + "import torchvision\n", + "import pandas as pd\n", + "import awswrangler as wr\n", + "\n", + "from torch.optim import SGD\n", + "from torch.nn import CrossEntropyLoss\n", + "from torch.utils.data import DataLoader" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "········\n" + ] + } + ], + "source": [ + "import getpass\n", + "bucket = getpass.getpass()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 1. Defining Training Function" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "def train(model, dataset, batch_size=64, epochs=2, device='cpu', num_workers=1):\n", + "\n", + " criterion = CrossEntropyLoss().to(device)\n", + " opt = SGD(model.parameters(), 0.025)\n", + " loader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)\n", + "\n", + " for epoch in range(epochs):\n", + "\n", + " correct = 0 \n", + " model.train()\n", + " for i, (inputs, labels) in enumerate(loader):\n", + "\n", + " # Forward Pass\n", + " outputs = model(inputs)\n", + " \n", + " # Backward Pass\n", + " loss = criterion(outputs, labels)\n", + " loss.backward()\n", + " opt.step()\n", + " opt.zero_grad()\n", + " \n", + " # Accuracy\n", + " _, predicted = torch.max(outputs.data, 1)\n", + " correct += (predicted == labels).sum().item()\n", + " accuracy = 100 * correct / ((i+1) * batch_size)\n", + "\n", + " print(f'batch: {i} loss: {loss.mean().item():.4f} acc: {accuracy:.2f}') " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 2. Training From Amazon S3" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2.1 Writing PyTorch Dataset to S3" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "client_s3 = boto3.client(\"s3\")\n", + "folder = \"tutorial_torch_dataset\"\n", + "\n", + "wr.s3.delete_objects(f\"s3://{bucket}/{folder}\")\n", + "for i in range(3):\n", + " batch = (\n", + " torch.randn(100, 3, 32, 32),\n", + " torch.randint(2, size=(100,)),\n", + " )\n", + " buff = io.BytesIO()\n", + " torch.save(batch, buff)\n", + " buff.seek(0)\n", + " client_s3.put_object(\n", + " Body=buff.read(),\n", + " Bucket=bucket,\n", + " Key=f\"{folder}/file{i}.pt\",\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2.2 Training Network" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "batch: 0 loss: 7.0132 acc: 0.00\n", + "batch: 1 loss: 2.8764 acc: 21.09\n", + "batch: 2 loss: 0.9600 acc: 32.29\n", + "batch: 3 loss: 0.8676 acc: 36.33\n", + "batch: 4 loss: 1.1386 acc: 36.88\n", + "batch: 0 loss: 1.0754 acc: 51.56\n", + "batch: 1 loss: 1.4241 acc: 51.56\n", + "batch: 2 loss: 1.3019 acc: 51.04\n", + "batch: 3 loss: 0.8631 acc: 53.52\n", + "batch: 4 loss: 0.4252 acc: 54.38\n" + ] + } + ], + "source": [ + "train(\n", + " torchvision.models.resnet18(),\n", + " wr.torch.S3IterableDataset(path=f\"{bucket}/{folder}\")\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 2. Training Directly From SQL Query" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2.1 Writing Data to SQL Database" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "eng = wr.catalog.get_engine(\"aws-data-wrangler-redshift\")\n", + "df = pd.DataFrame({\n", + " \"height\": [2, 1.4, 1.7, 1.8, 1.9, 2.2],\n", + " \"weight\": [100.0, 50.0, 70.0, 80.0, 90.0, 160.0],\n", + " \"target\": [1, 0, 0, 1, 1, 1]\n", + "})\n", + "\n", + "wr.db.to_sql(\n", + " df,\n", + " eng,\n", + " schema=\"public\",\n", + " name=\"torch\",\n", + " if_exists=\"replace\",\n", + " index=False\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2.2 Training Network From SQL" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "batch: 0 loss: 8.8708 acc: 50.00\n", + "batch: 1 loss: 88.7789 acc: 50.00\n", + "batch: 2 loss: 0.8655 acc: 33.33\n", + "batch: 0 loss: 0.7036 acc: 50.00\n", + "batch: 1 loss: 0.7034 acc: 50.00\n", + "batch: 2 loss: 0.8447 acc: 33.33\n", + "batch: 0 loss: 0.7012 acc: 50.00\n", + "batch: 1 loss: 0.7010 acc: 50.00\n", + "batch: 2 loss: 0.8250 acc: 33.33\n", + "batch: 0 loss: 0.6992 acc: 50.00\n", + "batch: 1 loss: 0.6991 acc: 50.00\n", + "batch: 2 loss: 0.8063 acc: 33.33\n", + "batch: 0 loss: 0.6975 acc: 50.00\n", + "batch: 1 loss: 0.6974 acc: 50.00\n", + "batch: 2 loss: 0.7886 acc: 33.33\n" + ] + } + ], + "source": [ + "train(\n", + " torch.nn.Sequential(\n", + " torch.nn.Linear(2, 10),\n", + " torch.nn.ReLU(),\n", + " torch.nn.Linear(10, 2), \n", + " ),\n", + " wr.torch.SQLDataset(\n", + " sql=\"SELECT * FROM public.torch\",\n", + " con=eng,\n", + " label_col=\"target\",\n", + " chunksize=2\n", + " ),\n", + " num_workers=0,\n", + " batch_size=2,\n", + " epochs=5\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 3. Delete Objects" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "wr.s3.delete_objects(f\"s3://{bucket}/{folder}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.7.5" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/tutorials/15 - EMR.ipynb b/tutorials/15 - EMR.ipynb new file mode 100644 index 000000000..4e1c627e6 --- /dev/null +++ b/tutorials/15 - EMR.ipynb @@ -0,0 +1,193 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "[![AWS Data Wrangler](_static/logo.png \"AWS Data Wrangler\")](https://github.com/awslabs/aws-data-wrangler)\n", + "\n", + "# 15 - EMR" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import awswrangler as wr\n", + "import boto3" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Enter your bucket name:" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stdin", + "output_type": "stream", + "text": [ + " ··········································\n" + ] + } + ], + "source": [ + "import getpass\n", + "bucket = getpass.getpass()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Enter your Subnet ID:" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "name": "stdin", + "output_type": "stream", + "text": [ + " ························\n" + ] + } + ], + "source": [ + "subnet = getpass.getpass()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Creating EMR Cluster" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "cluster_id = wr.emr.create_cluster(subnet)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Uploading our PySpark script to Amazon S3" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [], + "source": [ + "script = \"\"\"\n", + "from pyspark.sql import SparkSession\n", + "spark = SparkSession.builder.appName(\"docker-awswrangler\").getOrCreate()\n", + "sc = spark.sparkContext\n", + "\n", + "print(\"Spark Initialized\")\n", + "\"\"\"\n", + "\n", + "_ = boto3.client(\"s3\").put_object(\n", + " Body=script,\n", + " Bucket=bucket,\n", + " Key=\"test.py\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Submit PySpark step" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [], + "source": [ + "step_id = wr.emr.submit_step(cluster_id, command=f\"spark-submit s3://{bucket}/test.py\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Wait Step" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [], + "source": [ + "while wr.emr.get_step_state(cluster_id, step_id) != \"COMPLETED\":\n", + " pass" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Terminate Cluster" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [], + "source": [ + "wr.emr.terminate_cluster(cluster_id)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.10" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/tutorials/16 - EMR & Docker.ipynb b/tutorials/16 - EMR & Docker.ipynb new file mode 100644 index 000000000..4ffb2be2b --- /dev/null +++ b/tutorials/16 - EMR & Docker.ipynb @@ -0,0 +1,362 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "[![AWS Data Wrangler](_static/logo.png \"AWS Data Wrangler\")](https://github.com/awslabs/aws-data-wrangler)\n", + "\n", + "# 16 - EMR & Docker" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import awswrangler as wr\n", + "import boto3\n", + "import getpass" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Enter your bucket name:" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stdin", + "output_type": "stream", + "text": [ + " ··········································\n" + ] + } + ], + "source": [ + "bucket = getpass.getpass()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Enter your Subnet ID:" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stdin", + "output_type": "stream", + "text": [ + " ························\n" + ] + } + ], + "source": [ + "subnet = getpass.getpass()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "## Build and Upload Docker Image to ECR repository\n", + "\n", + "Replace the `{ACCOUNT_ID}` placeholder." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false, + "jupyter": { + "outputs_hidden": false + }, + "pycharm": { + "name": "#%%writefile Dockerfile\n" + } + }, + "outputs": [], + "source": [ + "%%writefile Dockerfile\n", + "\n", + "FROM amazoncorretto:8\n", + "\n", + "RUN yum -y update\n", + "RUN yum -y install yum-utils\n", + "RUN yum -y groupinstall development\n", + "\n", + "RUN yum list python3*\n", + "RUN yum -y install python3 python3-dev python3-pip python3-virtualenv\n", + "\n", + "RUN python -V\n", + "RUN python3 -V\n", + "\n", + "ENV PYSPARK_DRIVER_PYTHON python3\n", + "ENV PYSPARK_PYTHON python3\n", + "\n", + "RUN pip3 install --upgrade pip\n", + "RUN pip3 install awswrangler\n", + "\n", + "RUN python3 -c \"import awswrangler as wr\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "\n", + "docker build -t 'local/emr-wrangler' .\n", + "aws ecr create-repository --repository-name emr-wrangler\n", + "docker tag local/emr-wrangler {ACCOUNT_ID}.dkr.ecr.us-east-1.amazonaws.com/emr-wrangler:emr-wrangler\n", + "eval $(aws ecr get-login --region us-east-1 --no-include-email)\n", + "docker push {ACCOUNT_ID}.dkr.ecr.us-east-1.amazonaws.com/emr-wrangler:emr-wrangler" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Creating EMR Cluster" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "cluster_id = wr.emr.create_cluster(subnet, docker=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Refresh ECR credentials in the cluster (expiration time: 12h )" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'s-1B0O45RWJL8CL'" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "wr.emr.submit_ecr_credentials_refresh(cluster_id, path=f\"s3://{bucket}/\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Uploading application script to Amazon S3 (PySpark)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "script = \"\"\"\n", + "from pyspark.sql import SparkSession\n", + "spark = SparkSession.builder.appName(\"docker-awswrangler\").getOrCreate()\n", + "sc = spark.sparkContext\n", + "\n", + "print(\"Spark Initialized\")\n", + "\n", + "import awswrangler as wr\n", + "\n", + "print(f\"Wrangler version: {wr.__version__}\")\n", + "\"\"\"\n", + "\n", + "boto3.client(\"s3\").put_object(Body=script, Bucket=bucket, Key=\"test_docker.py\");" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Submit PySpark step" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [], + "source": [ + "DOCKER_IMAGE = f\"{wr.get_account_id()}.dkr.ecr.us-east-1.amazonaws.com/emr-wrangler:emr-wrangler\"\n", + "\n", + "step_id = wr.emr.submit_spark_step(\n", + " cluster_id,\n", + " f\"s3://{bucket}/test_docker.py\",\n", + " docker_image=DOCKER_IMAGE\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Wait Step" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "while wr.emr.get_step_state(cluster_id, step_id) != \"COMPLETED\":\n", + " pass" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Terminate Cluster" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "wr.emr.terminate_cluster(cluster_id)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Another example with custom configurations" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "cluster_id = wr.emr.create_cluster(\n", + " cluster_name=\"my-demo-cluster-v2\",\n", + " logging_s3_path=f\"s3://{bucket}/emr-logs/\",\n", + " emr_release=\"emr-6.0.0\",\n", + " subnet_id=subnet,\n", + " emr_ec2_role=\"EMR_EC2_DefaultRole\",\n", + " emr_role=\"EMR_DefaultRole\",\n", + " instance_type_master=\"m5.2xlarge\",\n", + " instance_type_core=\"m5.2xlarge\",\n", + " instance_ebs_size_master=50,\n", + " instance_ebs_size_core=50,\n", + " instance_num_on_demand_master=0,\n", + " instance_num_on_demand_core=0,\n", + " instance_num_spot_master=1,\n", + " instance_num_spot_core=2,\n", + " spot_bid_percentage_of_on_demand_master=100,\n", + " spot_bid_percentage_of_on_demand_core=100,\n", + " spot_provisioning_timeout_master=5,\n", + " spot_provisioning_timeout_core=5,\n", + " spot_timeout_to_on_demand_master=False,\n", + " spot_timeout_to_on_demand_core=False,\n", + " python3=True,\n", + " docker=True,\n", + " spark_glue_catalog=True,\n", + " hive_glue_catalog=True,\n", + " presto_glue_catalog=True,\n", + " debugging=True,\n", + " applications=[\"Hadoop\", \"Spark\", \"Hive\", \"Zeppelin\", \"Livy\"],\n", + " visible_to_all_users=True,\n", + " maximize_resource_allocation=True,\n", + " keep_cluster_alive_when_no_steps=True,\n", + " termination_protected=False,\n", + " spark_pyarrow=True\n", + ")\n", + "\n", + "wr.emr.submit_ecr_credentials_refresh(cluster_id, path=f\"s3://{bucket}/emr/\")\n", + "\n", + "DOCKER_IMAGE = f\"{wr.get_account_id()}.dkr.ecr.us-east-1.amazonaws.com/emr-wrangler:emr-wrangler\"\n", + "\n", + "step_id = wr.emr.submit_spark_step(\n", + " cluster_id,\n", + " f\"s3://{bucket}/test_docker.py\",\n", + " docker_image=DOCKER_IMAGE\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "while wr.emr.get_step_state(cluster_id, step_id) != \"COMPLETED\":\n", + " pass\n", + "\n", + "wr.emr.terminate_cluster(cluster_id)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.6.10" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +}