Repository navigation
Support per-SessionContext object store credentials without os.environ (thread-safety) #1623
Description
Is your feature request related to a problem or challenge? Please describe what you are trying to do.
When using datafusion-python as a compute backend from multi-threaded Python applications (e.g., Apache Iceberg's PyIceberg), we need to read Parquet files from cloud storage (S3/GCS/ADLS) in parallel across multiple threads. Each thread creates its own SessionContext and calls register_parquet().
Currently, cloud credentials must be set via os.environ before calling register_parquet() (the Rust object_store crate reads AWS_ACCESS_KEY_ID, etc. from the environment). Since os.environ is process-global, concurrent threads mutating it causes credential cross-contamination.
The workaround is a global threading.RLock() that serializes all file-based DataFusion operations — effectively negating the benefit of the thread pool. A scan over 200 S3 Parquet files runs ~4-8× slower than it should because only one thread can read at a time.
Describe the solution you'd like
Add an optional object_store parameter to register_parquet() (and other register_* / read_* methods) that accepts a pre-configured store instance:
from datafusion import SessionContext
from datafusion.object_store import AmazonS3
# Thread-safe: no os.environ mutation needed
ctx = SessionContext()
store = AmazonS3('my-bucket', region='us-east-1', access_key_id=key, secret_access_key=secret)
ctx.register_object_store('s3://', store, host='my-bucket')
ctx.register_parquet('my_table', 's3://my-bucket/data.parquet')I know register_object_store() already exists (and works!), but the two-step dance of register_object_store + register_parquet requires the caller to:
- Parse the bucket name from the file path
- Determine the correct scheme (s3://, gs://, az://)
- Call register_object_store before register_parquet
A singleobject_store=keyword onregister_parquetwould handle all of this internally.
Current workaround (in a pending PyIceberg PR)
import threading, os
_ENV_LOCK = threading.RLock()
def read_parquet_with_credentials(ctx, path, io_properties):
env_vars = translate_properties_to_env(io_properties)
with _ENV_LOCK: # Serializes ALL concurrent reads
for k, v in env_vars.items():
os.environ[k] = v
try:
ctx.register_parquet('source', path)
result = ctx.sql('SELECT * FROM source').to_arrow_table()
finally:
for k in env_vars:
os.environ.pop(k, None)
return resultDescribe alternatives you've considered
- Use register_object_store directly — This works for S3 (AmazonS3 accepts inline creds), but GoogleCloud only accepts service_account_path (no inline token/JSON), making it insufficient for GCS with ephemeral credentials.
- Wait for PR feat: auto-register ObjectStore and accept it in read/register methods #1476 — The draft PR integrates pyo3-object_store which fully solves this, but it's a large change (361 lines) and has been in draft for 3 months.
- Minimal fix — Add object_store= parameter to register_parquet that calls register_object_store internally after parsing the URL. This is ~20 lines of Python.
Additional context
- Related: Remove the need for registering an ObjectStore for remote files #899 (general ObjectStore ergonomics), PR feat: auto-register ObjectStore and accept it in read/register methods #1476 (full pyo3-object_store integration)
- Downstream consumer: apache/iceberg-python pluggable backend
- DataFusion version: 54.0.0
- The
AmazonS3,GoogleCloud,MicrosoftAzureclasses already exist indatafusion.object_store— the gap is just the ergonomics of passing them toregister_parquetand filling in GoogleCloud's credential options
Activity
qzyu999 commentedon Jul 8, 2026
Accidentally closed and not allowed to reopen, recreating issue at #1624, apologies for the duplicate.