feat: Add s3.read_orc and s3.to_orc (#2312)

Author: LeonLuttenberger

awswrangler/catalog/__init__.py

@@ -1,13 +1,22 @@
 """Amazon Glue Catalog Module."""
 
-from awswrangler.catalog._add import add_column, add_csv_partitions, add_json_partitions, add_parquet_partitions  # noqa
+from awswrangler.catalog._add import (
+    add_column,
+    add_csv_partitions,
+    add_json_partitions,
+    add_orc_partitions,
+    add_parquet_partitions,
+)
+
+# noqa
 from awswrangler.catalog._create import (  # noqa
     _create_csv_table,
     _create_json_table,
     _create_parquet_table,
     create_csv_table,
     create_database,
     create_json_table,
+    create_orc_table,
     create_parquet_table,
     overwrite_table_parameters,
     upsert_table_parameters,
@@ -54,6 +63,7 @@
     "add_csv_partitions",
     "add_json_partitions",
     "add_parquet_partitions",
+    "add_orc_partitions",
     "does_table_exist",
     "delete_column",
     "drop_duplicated_columns",
@@ -68,6 +78,7 @@
     "create_csv_table",
     "create_database",
     "create_parquet_table",
+    "create_orc_table",
     "create_json_table",
     "overwrite_table_parameters",
     "upsert_table_parameters",

awswrangler/catalog/_add.py

@@ -11,6 +11,7 @@
     _check_column_type,
     _csv_partition_definition,
     _json_partition_definition,
+    _orc_partition_definition,
     _parquet_partition_definition,
     _update_table_definition,
 )
@@ -294,6 +295,84 @@ def add_parquet_partitions(
         )
 
 
+@apply_configs
+def add_orc_partitions(
+    database: str,
+    table: str,
+    partitions_values: Dict[str, List[str]],
+    bucketing_info: Optional[typing.BucketingInfoTuple] = None,
+    catalog_id: Optional[str] = None,
+    compression: Optional[str] = None,
+    boto3_session: Optional[boto3.Session] = None,
+    columns_types: Optional[Dict[str, str]] = None,
+    partitions_parameters: Optional[Dict[str, str]] = None,
+) -> None:
+    """Add partitions (metadata) to a ORC Table in the AWS Glue Catalog.
+
+    Parameters
+    ----------
+    database : str
+        Database name.
+    table : str
+        Table name.
+    partitions_values: Dict[str, List[str]]
+        Dictionary with keys as S3 path locations and values as a list of partitions values as str
+        (e.g. {'s3://bucket/prefix/y=2020/m=10/': ['2020', '10']}).
+    bucketing_info: Tuple[List[str], int], optional
+        Tuple consisting of the column names used for bucketing as the first element and the number of buckets as the
+        second element.
+        Only `str`, `int` and `bool` are supported as column data types for bucketing.
+    catalog_id : str, optional
+        The ID of the Data Catalog from which to retrieve Databases.
+        If none is provided, the AWS account ID is used by default.
+    compression: str, optional
+        Compression style (``None``, ``snappy``, ``zlib``, etc).
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 session will be used if boto3_session receive None.
+    columns_types: Optional[Dict[str, str]]
+        Only required for Hive compability.
+        Dictionary with keys as column names and values as data types (e.g. {'col0': 'bigint', 'col1': 'double'}).
+        P.S. Only materialized columns please, not partition columns.
+    partitions_parameters: Optional[Dict[str, str]]
+        Dictionary with key-value pairs defining partition parameters.
+
+    Returns
+    -------
+    None
+        None.
+
+    Examples
+    --------
+    >>> import awswrangler as wr
+    >>> wr.catalog.add_orc_partitions(
+    ...     database='default',
+    ...     table='my_table',
+    ...     partitions_values={
+    ...         's3://bucket/prefix/y=2020/m=10/': ['2020', '10'],
+    ...         's3://bucket/prefix/y=2020/m=11/': ['2020', '11'],
+    ...         's3://bucket/prefix/y=2020/m=12/': ['2020', '12']
+    ...     }
+    ... )
+
+    """
+    table = sanitize_table_name(table=table)
+    if partitions_values:
+        inputs: List[Dict[str, Any]] = [
+            _orc_partition_definition(
+                location=k,
+                values=v,
+                bucketing_info=bucketing_info,
+                compression=compression,
+                columns_types=columns_types,
+                partitions_parameters=partitions_parameters,
+            )
+            for k, v in partitions_values.items()
+        ]
+        _add_partitions(
+            database=database, table=table, boto3_session=boto3_session, inputs=inputs, catalog_id=catalog_id
+        )
+
+
 @apply_configs
 def add_column(
     database: str,

awswrangler/catalog/_create.py

@@ -7,7 +7,12 @@
 
 from awswrangler import _utils, exceptions, typing
 from awswrangler._config import apply_configs
-from awswrangler.catalog._definitions import _csv_table_definition, _json_table_definition, _parquet_table_definition
+from awswrangler.catalog._definitions import (
+    _csv_table_definition,
+    _json_table_definition,
+    _orc_table_definition,
+    _parquet_table_definition,
+)
 from awswrangler.catalog._delete import delete_all_partitions, delete_table_if_exists
 from awswrangler.catalog._get import _get_table_input
 from awswrangler.catalog._utils import _catalog_id, _transaction_id, sanitize_column_name, sanitize_table_name
@@ -335,6 +340,68 @@ def _create_parquet_table(
     )
 
 
+def _create_orc_table(
+    database: str,
+    table: str,
+    path: str,
+    columns_types: Dict[str, str],
+    table_type: Optional[str],
+    partitions_types: Optional[Dict[str, str]],
+    bucketing_info: Optional[typing.BucketingInfoTuple],
+    catalog_id: Optional[str],
+    compression: Optional[str],
+    description: Optional[str],
+    parameters: Optional[Dict[str, str]],
+    columns_comments: Optional[Dict[str, str]],
+    mode: str,
+    catalog_versioning: bool,
+    transaction_id: Optional[str],
+    athena_partition_projection_settings: Optional[typing.AthenaPartitionProjectionSettings],
+    boto3_session: Optional[boto3.Session],
+    catalog_table_input: Optional[Dict[str, Any]],
+) -> None:
+    table = sanitize_table_name(table=table)
+    partitions_types = {} if partitions_types is None else partitions_types
+    _logger.debug("catalog_table_input: %s", catalog_table_input)
+
+    table_input: Dict[str, Any]
+    if (catalog_table_input is not None) and (mode in ("append", "overwrite_partitions")):
+        table_input = catalog_table_input
+
+        is_table_updated = _update_table_input(table_input, columns_types)
+        if is_table_updated:
+            mode = "update"
+    else:
+        table_input = _orc_table_definition(
+            table=table,
+            path=path,
+            columns_types=columns_types,
+            table_type=table_type,
+            partitions_types=partitions_types,
+            bucketing_info=bucketing_info,
+            compression=compression,
+        )
+    table_exist: bool = catalog_table_input is not None
+    _logger.debug("table_exist: %s", table_exist)
+    _create_table(
+        database=database,
+        table=table,
+        description=description,
+        parameters=parameters,
+        columns_comments=columns_comments,
+        mode=mode,
+        catalog_versioning=catalog_versioning,
+        boto3_session=boto3_session,
+        table_input=table_input,
+        table_type=table_type,
+        table_exist=table_exist,
+        partitions_types=partitions_types,
+        transaction_id=transaction_id,
+        athena_partition_projection_settings=athena_partition_projection_settings,
+        catalog_id=catalog_id,
+    )
+
+
 def _create_csv_table(  # pylint: disable=too-many-arguments,too-many-locals
     database: str,
     table: str,
@@ -827,6 +894,168 @@ def create_parquet_table(
     )
 
 
+@apply_configs
+def create_orc_table(
+    database: str,
+    table: str,
+    path: str,
+    columns_types: Dict[str, str],
+    table_type: Optional[str] = None,
+    partitions_types: Optional[Dict[str, str]] = None,
+    bucketing_info: Optional[typing.BucketingInfoTuple] = None,
+    catalog_id: Optional[str] = None,
+    compression: Optional[str] = None,
+    description: Optional[str] = None,
+    parameters: Optional[Dict[str, str]] = None,
+    columns_comments: Optional[Dict[str, str]] = None,
+    mode: Literal["overwrite", "append"] = "overwrite",
+    catalog_versioning: bool = False,
+    transaction_id: Optional[str] = None,
+    athena_partition_projection_settings: Optional[typing.AthenaPartitionProjectionSettings] = None,
+    boto3_session: Optional[boto3.Session] = None,
+) -> None:
+    """Create a ORC Table (Metadata Only) in the AWS Glue Catalog.
+
+    'https://docs.aws.amazon.com/athena/latest/ug/data-types.html'
+
+    Parameters
+    ----------
+    database : str
+        Database name.
+    table : str
+        Table name.
+    path : str
+        Amazon S3 path (e.g. s3://bucket/prefix/).
+    columns_types: Dict[str, str]
+        Dictionary with keys as column names and values as data types (e.g. {'col0': 'bigint', 'col1': 'double'}).
+    table_type: str, optional
+        The type of the Glue Table (EXTERNAL_TABLE, GOVERNED...). Set to EXTERNAL_TABLE if None
+    partitions_types: Dict[str, str], optional
+        Dictionary with keys as partition names and values as data types (e.g. {'col2': 'date'}).
+    bucketing_info: Tuple[List[str], int], optional
+        Tuple consisting of the column names used for bucketing as the first element and the number of buckets as the
+        second element.
+        Only `str`, `int` and `bool` are supported as column data types for bucketing.
+    catalog_id : str, optional
+        The ID of the Data Catalog from which to retrieve Databases.
+        If none is provided, the AWS account ID is used by default.
+    compression: str, optional
+        Compression style (``None``, ``snappy``, ``gzip``, etc).
+    description: str, optional
+        Table description
+    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
+        '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.
+    transaction_id: str, optional
+        The ID of the transaction (i.e. used with GOVERNED tables).
+    athena_partition_projection_settings: typing.AthenaPartitionProjectionSettings, optional
+        Parameters of the Athena Partition Projection (https://docs.aws.amazon.com/athena/latest/ug/partition-projection.html).
+        AthenaPartitionProjectionSettings is a `TypedDict`, meaning the passed parameter can be instantiated either as an
+        instance of AthenaPartitionProjectionSettings or as a regular Python dict.
+
+        Following projection parameters are supported:
+
+        .. list-table:: Projection Parameters
+           :header-rows: 1
+
+           * - Name
+             - Type
+             - Description
+           * - projection_types
+             - Optional[Dict[str, str]]
+             - Dictionary of partitions names and Athena projections types.
+               Valid types: "enum", "integer", "date", "injected"
+               https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html
+               (e.g. {'col_name': 'enum', 'col2_name': 'integer'})
+           * - projection_ranges
+             - Optional[Dict[str, str]]
+             - Dictionary of partitions names and Athena projections ranges.
+               https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html
+               (e.g. {'col_name': '0,10', 'col2_name': '-1,8675309'})
+           * - projection_values
+             - Optional[Dict[str, str]]
+             - Dictionary of partitions names and Athena projections values.
+               https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html
+               (e.g. {'col_name': 'A,B,Unknown', 'col2_name': 'foo,boo,bar'})
+           * - projection_intervals
+             - Optional[Dict[str, str]]
+             - Dictionary of partitions names and Athena projections intervals.
+               https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html
+               (e.g. {'col_name': '1', 'col2_name': '5'})
+           * - projection_digits
+             - Optional[Dict[str, str]]
+             - Dictionary of partitions names and Athena projections digits.
+               https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html
+               (e.g. {'col_name': '1', 'col2_name': '2'})
+           * - projection_formats
+             - Optional[Dict[str, str]]
+             - Dictionary of partitions names and Athena projections formats.
+               https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html
+               (e.g. {'col_date': 'yyyy-MM-dd', 'col2_timestamp': 'yyyy-MM-dd HH:mm:ss'})
+           * - projection_storage_location_template
+             - Optional[str]
+             - Value which is allows Athena to properly map partition values if the S3 file locations do not follow
+               a typical `.../column=value/...` pattern.
+               https://docs.aws.amazon.com/athena/latest/ug/partition-projection-setting-up.html
+               (e.g. s3://bucket/table_root/a=${a}/${b}/some_static_subdirectory/${c}/)
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 session will be used if boto3_session receive None.
+
+    Returns
+    -------
+    None
+        None.
+
+    Examples
+    --------
+    >>> import awswrangler as wr
+    >>> wr.catalog.create_parquet_table(
+    ...     database='default',
+    ...     table='my_table',
+    ...     path='s3://bucket/prefix/',
+    ...     columns_types={'col0': 'bigint', 'col1': 'double'},
+    ...     partitions_types={'col2': 'date'},
+    ...     compression='snappy',
+    ...     description='My own table!',
+    ...     parameters={'source': 'postgresql'},
+    ...     columns_comments={'col0': 'Column 0.', 'col1': 'Column 1.', 'col2': 'Partition.'}
+    ... )
+
+    """
+    catalog_table_input: Optional[Dict[str, Any]] = _get_table_input(
+        database=database,
+        table=table,
+        boto3_session=boto3_session,
+        transaction_id=transaction_id,
+        catalog_id=catalog_id,
+    )
+    _create_orc_table(
+        database=database,
+        table=table,
+        path=path,
+        columns_types=columns_types,
+        table_type=table_type,
+        partitions_types=partitions_types,
+        bucketing_info=bucketing_info,
+        catalog_id=catalog_id,
+        compression=compression,
+        description=description,
+        parameters=parameters,
+        columns_comments=columns_comments,
+        mode=mode,
+        catalog_versioning=catalog_versioning,
+        transaction_id=transaction_id,
+        athena_partition_projection_settings=athena_partition_projection_settings,
+        boto3_session=boto3_session,
+        catalog_table_input=catalog_table_input,
+    )
+
+
 @apply_configs
 def create_csv_table(  # pylint: disable=too-many-arguments,too-many-locals
     database: str,