Skip to content

Add CTAS bucketing to wr.athena.read_sql_query #802

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jul 9, 2021
Merged

Add CTAS bucketing to wr.athena.read_sql_query #802

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a724fa1
Add CTAS bucketing to wr.athena_read_sql_query
kukushking Jul 7, 2021
08d43a7
Add CTAS bucketing to read_sql_table
kukushking Jul 7, 2021
f452c67
Formatting
kukushking Jul 7, 2021
992ee45
Merge branch 'main' into feat-read-sql-ctas-bucketing
jaidisido Jul 9, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 22 additions & 1 deletion awswrangler/athena/_read.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import re
import sys
import uuid
from typing import Any, Dict, Iterator, List, Match, NamedTuple, Optional, Union
from typing import Any, Dict, Iterator, List, Match, NamedTuple, Optional, Tuple, Union

import boto3
import botocore.exceptions
Expand Down Expand Up @@ -385,18 +385,25 @@ def _resolve_query_without_cache_ctas(
wg_config: _WorkGroupConfig,
alt_database: Optional[str],
name: Optional[str],
ctas_bucketing_info: Optional[Tuple[List[str], int]],
use_threads: bool,
s3_additional_kwargs: Optional[Dict[str, Any]],
boto3_session: boto3.Session,
) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
path: str = f"{s3_output}/{name}"
ext_location: str = "\n" if wg_config.enforced is True else f",\n external_location = '{path}'\n"
fully_qualified_name: str = f'"{alt_database}"."{name}"' if alt_database else f'"{database}"."{name}"'
bucketing_str = (
(f",\n" f" bucketed_by = ARRAY{ctas_bucketing_info[0]},\n" f" bucket_count = {ctas_bucketing_info[1]}")
if ctas_bucketing_info
else ""
)
sql = (
f"CREATE TABLE {fully_qualified_name}\n"
f"WITH(\n"
f" format = 'Parquet',\n"
f" parquet_compression = 'SNAPPY'"
f"{bucketing_str}"
f"{ext_location}"
f") AS\n"
f"{sql}"
Expand Down Expand Up @@ -521,6 +528,7 @@ def _resolve_query_without_cache(
keep_files: bool,
ctas_database_name: Optional[str],
ctas_temp_table_name: Optional[str],
ctas_bucketing_info: Optional[Tuple[List[str], int]],
use_threads: bool,
s3_additional_kwargs: Optional[Dict[str, Any]],
boto3_session: boto3.Session,
Expand Down Expand Up @@ -553,6 +561,7 @@ def _resolve_query_without_cache(
wg_config=wg_config,
alt_database=ctas_database_name,
name=name,
ctas_bucketing_info=ctas_bucketing_info,
use_threads=use_threads,
s3_additional_kwargs=s3_additional_kwargs,
boto3_session=boto3_session,
Expand Down Expand Up @@ -593,6 +602,7 @@ def read_sql_query(
keep_files: bool = True,
ctas_database_name: Optional[str] = None,
ctas_temp_table_name: Optional[str] = None,
ctas_bucketing_info: Optional[Tuple[List[str], int]] = None,
use_threads: bool = True,
boto3_session: Optional[boto3.Session] = None,
max_cache_seconds: int = 0,
Expand Down Expand Up @@ -733,6 +743,10 @@ def read_sql_query(
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_{uuid.uuid4().hex()}"`.
On S3 this directory will be under under the pattern: `f"{s3_output}/{ctas_temp_table_name}/"`.
ctas_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.
Copy link
Contributor

@jaidisido jaidisido Jul 9, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I never know when we should verify the user input or just let it fail...

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.
Expand Down Expand Up @@ -841,6 +855,7 @@ def read_sql_query(
keep_files=keep_files,
ctas_database_name=ctas_database_name,
ctas_temp_table_name=ctas_temp_table_name,
ctas_bucketing_info=ctas_bucketing_info,
use_threads=use_threads,
s3_additional_kwargs=s3_additional_kwargs,
boto3_session=session,
Expand All @@ -861,6 +876,7 @@ def read_sql_table(
keep_files: bool = True,
ctas_database_name: Optional[str] = None,
ctas_temp_table_name: Optional[str] = None,
ctas_bucketing_info: Optional[Tuple[List[str], int]] = None,
use_threads: bool = True,
boto3_session: Optional[boto3.Session] = None,
max_cache_seconds: int = 0,
Expand Down Expand Up @@ -995,6 +1011,10 @@ def read_sql_table(
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_{uuid.uuid4().hex}"`.
On S3 this directory will be under under the pattern: `f"{s3_output}/{ctas_temp_table_name}/"`.
ctas_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.
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.
Expand Down Expand Up @@ -1053,6 +1073,7 @@ def read_sql_table(
keep_files=keep_files,
ctas_database_name=ctas_database_name,
ctas_temp_table_name=ctas_temp_table_name,
ctas_bucketing_info=ctas_bucketing_info,
use_threads=use_threads,
boto3_session=boto3_session,
max_cache_seconds=max_cache_seconds,
Expand Down
27 changes: 27 additions & 0 deletions tests/test_athena.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,33 @@ def test_athena_ctas(path, path2, path3, glue_table, glue_table2, glue_database,
assert len(wr.s3.list_objects(path=path3)) == 0


def test_athena_read_sql_ctas_bucketing(path, path2, glue_table, glue_table2, glue_database, glue_ctas_database):
df = pd.DataFrame({"c0": [0, 1], "c1": ["foo", "bar"]})
wr.s3.to_parquet(
df=df,
path=path,
dataset=True,
database=glue_database,
table=glue_table,
)
df_ctas = wr.athena.read_sql_query(
sql=f"SELECT * FROM {glue_table}",
ctas_approach=True,
database=glue_database,
ctas_database_name=glue_ctas_database,
ctas_temp_table_name=glue_table2,
ctas_bucketing_info=(["c0"], 1),
s3_output=path2,
)
df_no_ctas = wr.athena.read_sql_query(
sql=f"SELECT * FROM {glue_table}",
ctas_approach=False,
database=glue_database,
s3_output=path2,
)
assert df_ctas.equals(df_no_ctas)


def test_athena(path, glue_database, glue_table, kms_key, workgroup0, workgroup1):
wr.catalog.delete_table_if_exists(database=glue_database, table=glue_table)
wr.s3.to_parquet(
Expand Down