diff --git a/docs/source/user-guide/data-sources.md b/docs/source/user-guide/data-sources.md index 22e666837..cd6044c33 100644 --- a/docs/source/user-guide/data-sources.md +++ b/docs/source/user-guide/data-sources.md @@ -89,7 +89,9 @@ ctx.create_dataframe([[batch]]).show() ## Object Store DataFusion has support for multiple storage options in addition to local files. -The example below requires an appropriate S3 account with access credentials. +The example below requires access to the S3 bucket. Set `AWS_ACCESS_KEY_ID` and +`AWS_SECRET_ACCESS_KEY` in the environment before running it. For temporary +credentials, also set `AWS_SESSION_TOKEN`. Supported Object Stores are @@ -100,16 +102,17 @@ Supported Object Stores are - {py:class}`~datafusion.object_store.MicrosoftAzure` ```python +from datafusion import SessionContext from datafusion.object_store import AmazonS3 region = "us-east-1" bucket_name = "yellow-trips" +ctx = SessionContext() + s3 = AmazonS3( bucket_name=bucket_name, region=region, - access_key_id=os.getenv("AWS_ACCESS_KEY_ID"), - secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"), ) path = f"s3://{bucket_name}/" @@ -120,6 +123,27 @@ ctx.register_parquet("trips", path) ctx.table("trips").show() ``` +### Query S3 data with SQL + +To reach the same data from SQL, give the S3 path a table name with +`CREATE EXTERNAL TABLE`. The registered object store carries the credentials and +the region, so the statement itself needs only the location. + +Register the store for the bucket as shown above, then use that same +{py:class}`~datafusion.context.SessionContext` to create and query the table: + +```python +ctx.sql( + f""" + CREATE EXTERNAL TABLE trips_sql + STORED AS PARQUET + LOCATION '{path}' + """ +).collect() + +ctx.sql("SELECT count(passenger_count) FROM trips_sql").show() +``` + ## Other DataFrame Libraries DataFusion can import DataFrames directly from other libraries, such as diff --git a/python/tests/test_object_store_param.py b/python/tests/test_object_store_param.py index 13d51d792..64be19f63 100644 --- a/python/tests/test_object_store_param.py +++ b/python/tests/test_object_store_param.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -"""Tests for the object_store parameter on register/read file methods.""" +"""Tests for object store registration and register/read file methods.""" import contextlib from pathlib import Path @@ -138,3 +138,24 @@ def test_parquet_methods_with_local_object_store(ctx, tmp_path, method_name): dataframe = ctx.read_parquet(path, object_store=LocalFileSystem()) assert dataframe.collect()[0].column("value").to_pylist() == [10, 20, 30] + + +def test_sql_external_table_uses_registered_object_store(ctx, tmp_path): + """Read a remote URL through a registered store without network access.""" + table = pa.table({"passenger_count": [1, None, 3]}) + pq.write_table(table, tmp_path / "trips.parquet") + + # Back the S3 URL with local files to test SQL's registry lookup, not AWS. + store = LocalFileSystem(prefix=str(tmp_path)) + ctx.register_object_store("s3://", store, host="test-bucket") + ctx.sql( + """ + CREATE EXTERNAL TABLE trips_sql + STORED AS PARQUET + LOCATION 's3://test-bucket/' + """ + ).collect() + + assert ctx.table("trips_sql").to_pydict() == table.to_pydict() + result = ctx.sql("SELECT count(passenger_count) AS count FROM trips_sql") + assert result.to_pydict() == {"count": [2]}