Skip to content

Support per-SessionContext object store credentials without os.environ (thread-safety) #1624

Description

@qzyu999

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:

  1. Parse the bucket name from the file path
  2. Determine the correct scheme (s3://, gs://, az://)
  3. Call register_object_store before register_parquet
    A single object_store= keyword on register_parquet would 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 result

Describe alternatives you've considered

  1. 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.
  2. 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.
  3. 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

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions