Skip to content

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

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

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions