Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
30 changes: 27 additions & 3 deletions docs/source/user-guide/data-sources.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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}/"
Expand All @@ -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
Expand Down
23 changes: 22 additions & 1 deletion python/tests/test_object_store_param.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]}
Loading