Enable AWS S3 remote data source for Presto GPU benchmarks - #376
Enable AWS S3 remote data source for Presto GPU benchmarks#376kingcrimsontianyu wants to merge 11 commits into
Conversation
kingcrimsontianyu
added
improvement
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
simoneves
left a comment
simoneves
left a comment
There was a problem hiding this comment.
LGTM but will test properly
|
Sure! This has been tested on my laptop with S3 TPC-H/TPC-DS SF-1. The planned next step is to test on an EC2 instance with S3 SF-1K. |
| # When external_location_base is set (e.g. s3://bucket/prefix/sf100), point the | ||
| # table at that base. Otherwise fall back to the local bind-mounted file path. | ||
| if external_location_base: | ||
| location = f"{external_location_base}/{table_name}" |
There was a problem hiding this comment.
Do the trailing backslashes need to be stripped here similar to what you have in L38 in generate_table_schemas.py?
paul-aiyedun
left a comment
paul-aiyedun
left a comment
There was a problem hiding this comment.
Changes overall look good to me. However, I had comments about extending existing logic and using more specific options names.
|
|
||
| # AWS S3. | ||
| hive.s3.ssl.enabled=true | ||
| hive.s3.path-style-access=false |
There was a problem hiding this comment.
Was hive.s3.use-instance-credentials=false intentionally not added here?
| def create_duckdb_table(table_name, data_path): | ||
| create_table(table_name, get_abs_file_path(__file__, data_path)) | ||
| # Only resolve local path | ||
| if "://" not in data_path: |
There was a problem hiding this comment.
Consider adding a direct flag parameter e.g. is_s3_location here instead of the :// check.
| "table). Mutually exclusive with --external-location-base.", | ||
| ) | ||
| parser.add_argument( | ||
| "--external-location-base", |
There was a problem hiding this comment.
Consider a more specific option name like s3-data-dir-path.
Same comment applies to other files with a similar flag.
| nargs="+", | ||
| default=None, | ||
| help="Table names to generate schemas for. Required with --external-location-base since a " | ||
| "remote prefix cannot be listed.", |
There was a problem hiding this comment.
I believe S3 supports listing objects with a specified prefix?
|
|
||
| Usage: $0 [OPTIONS] | ||
|
|
||
| Registers benchmark external tables in a Presto Hive schema whose data lives at an |
There was a problem hiding this comment.
I think this script is effectively doing the same thing as setup_benchmark_tables.sh. Can we extend setup_benchmark_tables.sh with S3 options (e.g. --s3/--data-location {file|s3}, --s3-data-dir-path, etc.) instead of adding a new script?
| external_dir = "" | ||
| # Capture everything between two single quotes | ||
| location_match = re.search(r"external_location = '([^']*)'", create_table_text[0]) | ||
| location = location_match.group(1) if location_match else "" |
There was a problem hiding this comment.
Should we assert that location_match and location exist?
| def _extract_scale_factor(metadata: dict): | ||
| """Return the scale factor from parsed metadata, whether it is a top-level field or | ||
| nested under 'options'.""" | ||
| return metadata.get("scale_factor") or metadata.get("options", {}).get("scale_factor") | ||
|
|
||
|
|
||
| def read_scale_factor(metadata_uri: str): | ||
| """Read the scale_factor field from a metadata.json at ``metadata_uri``.""" | ||
| # For local data | ||
| if not str(metadata_uri).startswith("s3://"): | ||
| with open(metadata_uri) as file: | ||
| return _extract_scale_factor(json.load(file)) | ||
| # For remote data | ||
| ensure_remote_access(metadata_uri) | ||
| raw = duckdb.sql(f"SELECT content FROM read_text('{metadata_uri}')").fetchone()[0] | ||
| return _extract_scale_factor(json.loads(raw)) |
There was a problem hiding this comment.
I don't think scale factor access/reading is DuckDB specific logic?
| _s3_configured = False | ||
|
|
||
|
|
||
| def ensure_remote_access(path) -> None: |
There was a problem hiding this comment.
Consider something like configure_data_access, since this also applies to local files.
4 participants
This PR lets Presto GPU TPC-H/TPC-DS benchmarks and integration tests run against data in a remote object store (AWS S3) with no local copy of the dataset. Only the Hive metastore stays local.
The key changes include:
presto/scripts/register_external_tables.shregisters Hive tables at a URI base without generating or copying data.create_hive_tables.pygains--external-location-base. When set, each table'sEXTERNAL_LOCATIONis<base>/<table>. Otherwise it falls back to the localfile:path..sqlfiles now useEXTERNAL_LOCATION = '{location}'instead of the hardcoded'file:{file_path}'.docker-compose.common.ymlfile passesAWS_ACCESS_KEY_ID/SECRET_ACCESS_KEY/SESSION_TOKENandAWS_DEFAULT_REGION/AWS_REGIONto the coordinator and native worker.duckdb_utils.ensure_remote_accessfunction configures httpfs and S3credential_chainsecret.read_scale_factorderives the scale factor frommetadata.json(local via stdlibjson, and S3 via DuckDB). This allows the integration tests to continue working for remote data source.