Skip to content
Open
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
31 changes: 27 additions & 4 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,26 +102,47 @@ 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}/"
ctx.register_object_store("s3://", s3, None)
ctx.register_object_store("s3://", s3, host=bucket_name)

ctx.register_parquet("trips", path)

ctx.table("trips").show()
```

### Use S3 in SQL

Use `CREATE EXTERNAL TABLE` to give an S3 path a table name that you can query
with SQL. The statement uses the object store registered for the bucket on the
same {py:class}`~datafusion.context.SessionContext`.

After registering the object store above, create and query an external 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]}