feat: Add Timestream UNLOAD (#2284)

* feat: Add Timestream UNLOAD

* Add unload & unload_to_files

* [skip ci] Update api.rst

* Read column names from CSV metadata

* Add missing parameter to s3.list_objects type stubs

* [skip ci] Docstrings

* [skip ci] Update scale docs

* [skip ci] Update docstrings

* [skip ci] Update tutorial

* Error handling

* [skip ci] Docs formatting fix

---------

Co-authored-by: Lucas Hanson <[email protected]>

Author: kukushking

awswrangler/s3/_list.pyi

@@ -69,6 +69,7 @@ def list_directories(
 def list_objects(
     path: str,
     chunked: Literal[False],
+    suffix: Union[str, List[str], None] = ...,
     ignore_suffix: Union[str, List[str], None] = ...,
     last_modified_begin: Optional[datetime.datetime] = ...,
     last_modified_end: Optional[datetime.datetime] = ...,
@@ -79,6 +80,7 @@ def list_objects(
 @overload
 def list_objects(
     path: str,
+    suffix: Union[str, List[str], None] = ...,
     ignore_suffix: Union[str, List[str], None] = ...,
     last_modified_begin: Optional[datetime.datetime] = ...,
     last_modified_end: Optional[datetime.datetime] = ...,
@@ -90,6 +92,7 @@ def list_objects(
 def list_objects(
     path: str,
     chunked: Literal[True],
+    suffix: Union[str, List[str], None] = ...,
     ignore_suffix: Union[str, List[str], None] = ...,
     last_modified_begin: Optional[datetime.datetime] = ...,
     last_modified_end: Optional[datetime.datetime] = ...,
@@ -101,6 +104,7 @@ def list_objects(
 def list_objects(
     path: str,
     chunked: bool,
+    suffix: Union[str, List[str], None] = ...,
     ignore_suffix: Union[str, List[str], None] = ...,
     last_modified_begin: Optional[datetime.datetime] = ...,
     last_modified_end: Optional[datetime.datetime] = ...,

awswrangler/timestream/__init__.py

@@ -4,8 +4,13 @@
 from awswrangler.timestream._create import create_database, create_table
 from awswrangler.timestream._delete import delete_database, delete_table
 from awswrangler.timestream._list import list_databases, list_tables
-from awswrangler.timestream._read import query
-from awswrangler.timestream._write import batch_load, batch_load_from_files, wait_batch_load_task, write
+from awswrangler.timestream._read import query, unload, unload_to_files
+from awswrangler.timestream._write import (
+    batch_load,
+    batch_load_from_files,
+    wait_batch_load_task,
+    write,
+)
 
 __all__ = [
     "create_database",
@@ -19,4 +24,6 @@
     "batch_load",
     "batch_load_from_files",
     "wait_batch_load_task",
+    "unload_to_files",
+    "unload",
 ]

awswrangler/timestream/_read.py

@@ -1,14 +1,16 @@
 """Amazon Timestream Read Module."""
 
+import json
 import logging
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Literal, Optional, Union, cast
 
 import boto3
 import pandas as pd
 from botocore.config import Config
 
-from awswrangler import _utils
+from awswrangler import _utils, exceptions, s3
+from awswrangler._config import apply_configs
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
@@ -113,6 +115,16 @@ def _paginate_query(
         rows = []
 
 
+def _get_column_names_from_metadata(unload_path: str, boto3_session: Optional[boto3.Session] = None) -> List[str]:
+    client_s3 = _utils.client(service_name="s3", session=boto3_session)
+    metadata_path = s3.list_objects(path=unload_path, suffix="_metadata.json", boto3_session=boto3_session)[0]
+    bucket, key = _utils.parse_path(metadata_path)
+    metadata_content = json.loads(client_s3.get_object(Bucket=bucket, Key=key)["Body"].read().decode("utf-8"))
+    columns = [column["Name"] for column in metadata_content["ColumnInfo"]]
+    _logger.debug("Read %d columns from metadata file in: %s", len(columns), metadata_path)
+    return columns
+
+
 def query(
     sql: str,
     chunked: bool = False,
@@ -155,3 +167,272 @@ def query(
         # Modin's concat() can not concatenate empty data frames
         return pd.concat(results, ignore_index=True)
     return pd.DataFrame()
+
+
+@_utils.validate_distributed_kwargs(
+    unsupported_kwargs=["boto3_session", "s3_additional_kwargs"],
+)
+@apply_configs
+def unload(
+    sql: str,
+    path: str,
+    unload_format: Optional[Literal["CSV", "PARQUET"]] = None,
+    compression: Optional[Literal["GZIP", "NONE"]] = None,
+    partition_cols: Optional[List[str]] = None,
+    encryption: Optional[Literal["SSE_KMS", "SSE_S3"]] = None,
+    kms_key_id: Optional[str] = None,
+    field_delimiter: Optional[str] = ",",
+    escaped_by: Optional[str] = "\\",
+    chunked: Union[bool, int] = False,
+    keep_files: bool = False,
+    use_threads: Union[bool, int] = True,
+    boto3_session: Optional[boto3.Session] = None,
+    s3_additional_kwargs: Optional[Dict[str, str]] = None,
+    pyarrow_additional_kwargs: Optional[Dict[str, Any]] = None,
+) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
+    """
+    Unload query results to Amazon S3 and read the results as Pandas Data Frame.
+
+    https://docs.aws.amazon.com/timestream/latest/developerguide/export-unload.html
+
+    Parameters
+    ----------
+    sql : str
+        SQL query
+    path : str
+        S3 path to write stage files (e.g. s3://bucket_name/any_name/)
+    unload_format : str, optional
+        Format of the unloaded S3 objects from the query.
+        Valid values: "CSV", "PARQUET". Case sensitive. Defaults to "PARQUET"
+    compression : str, optional
+        Compression of the unloaded S3 objects from the query.
+        Valid values: "GZIP", "NONE". Defaults to "GZIP"
+    partition_cols : List[str], optional
+        Specifies the partition keys for the unload operation
+    encryption : str, optional
+        Encryption of the unloaded S3 objects from the query.
+        Valid values: "SSE_KMS", "SSE_S3". Defaults to "SSE_S3"
+    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
+    field_delimiter : str, optional
+        A single ASCII character that is used to separate fields in the output file,
+        such as pipe character (|), a comma (,), or tab (/t). Only used with CSV format
+    escaped_by : str, optional
+        The character that should be treated as an escape character in the data file
+        written to S3 bucket. Only used with CSV format
+    chunked : Union[int, bool]
+        If passed will split the data in a Iterable of DataFrames (Memory friendly).
+        If `True` awswrangler iterates on the data by files in the most efficient way without guarantee of chunksize.
+        If an `INTEGER` is passed awswrangler will iterate on the data by number of rows equal the received INTEGER.
+    keep_files : bool
+        Should keep stage files?
+    use_threads : bool, int
+        True to enable concurrent requests, False to disable multiple threads.
+        If enabled os.cpu_count() will be used as the max number of threads.
+        If integer is provided, specified number is used.
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 session is used if None
+    s3_additional_kwargs : Dict[str, str], optional
+        Forward to botocore requests.
+    pyarrow_additional_kwargs : Dict[str, Any], optional
+        Forwarded to `to_pandas` method converting from PyArrow tables to Pandas DataFrame.
+        Valid values include "split_blocks", "self_destruct", "ignore_metadata".
+        e.g. pyarrow_additional_kwargs={'split_blocks': True}.
+
+    Returns
+    -------
+    Union[pandas.DataFrame, Iterator[pandas.DataFrame]]
+        Result as Pandas DataFrame(s).
+
+    Examples
+    --------
+    Unload and read as Parquet (default).
+
+    >>> import awswrangler as wr
+    >>> df = wr.timestream.unload(
+    ...     sql="SELECT time, measure, dimension FROM database.mytable",
+    ...     path="s3://bucket/extracted_parquet_files/",
+    ... )
+
+    Unload and read partitioned Parquet. Note: partition columns must be at the end of the table.
+
+    >>> import awswrangler as wr
+    >>> df = wr.timestream.unload(
+    ...     sql="SELECT time, measure, dim1, dim2 FROM database.mytable",
+    ...     path="s3://bucket/extracted_parquet_files/",
+    ...     partition_cols=["dim2"],
+    ... )
+
+    Unload and read as CSV.
+
+    >>> import awswrangler as wr
+    >>> df = wr.timestream.unload(
+    ...     sql="SELECT time, measure, dimension FROM database.mytable",
+    ...     path="s3://bucket/extracted_parquet_files/",
+    ...     unload_format="CSV",
+    ... )
+
+    """
+    path = path if path.endswith("/") else f"{path}/"
+
+    if unload_format not in [None, "CSV", "PARQUET"]:
+        raise exceptions.InvalidArgumentValue("<unload_format> argument must be 'CSV' or 'PARQUET'")
+
+    unload_to_files(
+        sql=sql,
+        path=path,
+        unload_format=unload_format,
+        compression=compression,
+        partition_cols=partition_cols,
+        encryption=encryption,
+        kms_key_id=kms_key_id,
+        field_delimiter=field_delimiter,
+        escaped_by=escaped_by,
+        boto3_session=boto3_session,
+    )
+    results_path = f"{path}results/"
+    try:
+        if unload_format == "CSV":
+            column_names: List[str] = _get_column_names_from_metadata(path, boto3_session)
+            return s3.read_csv(
+                path=results_path,
+                header=None,
+                names=[column for column in column_names if column not in set(partition_cols)]
+                if partition_cols is not None
+                else column_names,
+                dataset=True,
+                use_threads=use_threads,
+                boto3_session=boto3_session,
+                s3_additional_kwargs=s3_additional_kwargs,
+            )
+        else:
+            return s3.read_parquet(
+                path=results_path,
+                chunked=chunked,
+                dataset=True,
+                use_threads=use_threads,
+                boto3_session=boto3_session,
+                s3_additional_kwargs=s3_additional_kwargs,
+                pyarrow_additional_kwargs=pyarrow_additional_kwargs,
+            )
+    finally:
+        if keep_files is False:
+            _logger.debug("Deleting objects in S3 path: %s", path)
+            s3.delete_objects(
+                path=path,
+                use_threads=use_threads,
+                boto3_session=boto3_session,
+                s3_additional_kwargs=s3_additional_kwargs,
+            )
+
+
+@_utils.validate_distributed_kwargs(
+    unsupported_kwargs=["boto3_session"],
+)
+@apply_configs
+def unload_to_files(
+    sql: str,
+    path: str,
+    unload_format: Optional[Literal["CSV", "PARQUET"]] = None,
+    compression: Optional[Literal["GZIP", "NONE"]] = None,
+    partition_cols: Optional[List[str]] = None,
+    encryption: Optional[Literal["SSE_KMS", "SSE_S3"]] = None,
+    kms_key_id: Optional[str] = None,
+    field_delimiter: Optional[str] = ",",
+    escaped_by: Optional[str] = "\\",
+    boto3_session: Optional[boto3.Session] = None,
+) -> None:
+    """
+    Unload query results to Amazon S3.
+
+    https://docs.aws.amazon.com/timestream/latest/developerguide/export-unload.html
+
+    Parameters
+    ----------
+    sql : str
+        SQL query
+    path : str
+        S3 path to write stage files (e.g. s3://bucket_name/any_name/)
+    unload_format : str, optional
+        Format of the unloaded S3 objects from the query.
+        Valid values: "CSV", "PARQUET". Case sensitive. Defaults to "PARQUET"
+    compression : str, optional
+        Compression of the unloaded S3 objects from the query.
+        Valid values: "GZIP", "NONE". Defaults to "GZIP"
+    partition_cols : List[str], optional
+        Specifies the partition keys for the unload operation
+    encryption : str, optional
+        Encryption of the unloaded S3 objects from the query.
+        Valid values: "SSE_KMS", "SSE_S3". Defaults to "SSE_S3"
+    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
+    field_delimiter : str, optional
+        A single ASCII character that is used to separate fields in the output file,
+        such as pipe character (|), a comma (,), or tab (/t). Only used with CSV format
+    escaped_by : str, optional
+        The character that should be treated as an escape character in the data file
+        written to S3 bucket. Only used with CSV format
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 session is used if None
+
+    Returns
+    -------
+    None
+
+    Examples
+    --------
+    Unload and read as Parquet (default).
+
+    >>> import awswrangler as wr
+    >>> wr.timestream.unload_to_files(
+    ...     sql="SELECT time, measure, dimension FROM database.mytable",
+    ...     path="s3://bucket/extracted_parquet_files/",
+    ... )
+
+    Unload and read partitioned Parquet. Note: partition columns must be at the end of the table.
+
+    >>> import awswrangler as wr
+    >>> wr.timestream.unload_to_files(
+    ...     sql="SELECT time, measure, dim1, dim2 FROM database.mytable",
+    ...     path="s3://bucket/extracted_parquet_files/",
+    ...     partition_cols=["dim2"],
+    ... )
+
+    Unload and read as CSV.
+
+    >>> import awswrangler as wr
+    >>> wr.timestream.unload_to_files(
+    ...     sql="SELECT time, measure, dimension FROM database.mytable",
+    ...     path="s3://bucket/extracted_parquet_files/",
+    ...     unload_format="CSV",
+    ... )
+
+    """
+    timestream_client = _utils.client(service_name="timestream-query", session=boto3_session)
+
+    partitioned_by_str: str = (
+        f"""partitioned_by = ARRAY [{','.join([f"'{col}'" for col in partition_cols])}],\n"""
+        if partition_cols is not None
+        else ""
+    )
+    kms_key_id_str: str = f"kms_key = '{kms_key_id}',\n" if kms_key_id is not None else ""
+    field_delimiter_str: str = f"field_delimiter = '{field_delimiter}',\n" if unload_format == "CSV" else ""
+    escaped_by_str: str = f"escaped_by = '{escaped_by}',\n" if unload_format == "CSV" else ""
+
+    sql = (
+        f"UNLOAD ({sql})\n"
+        f"TO '{path}'\n"
+        f"WITH (\n"
+        f"{partitioned_by_str}"
+        f"format='{unload_format or 'PARQUET'}',\n"
+        f"compression='{compression or 'GZIP'}',\n"
+        f"{field_delimiter_str}"
+        f"{escaped_by_str}"
+        f"{kms_key_id_str}"
+        f"encryption='{encryption or 'SSE_S3'}'\n"
+        f")"
+    )
+
+    timestream_client.query(QueryString=sql)