Loading polars.LazyFrame from S3 via PolarsParquetIOManager fails with "Generic S3 error: Missing bucket name"

[ilyanekhay-uta]

I'm trying to pass a polars.LazyFrame between two assets with code that roughly looks like so (omitting details in order to get a minimal reproducible example):

@asset(
    io_manager_key="polars_parquet_io_manager",
)
def upstream() -> pl.LazyFrame:
    return pl.LazyFrame(...)

@asset()
def downstream(upstream: pl.LazyFrame) -> ...:
    upstream.collect()

defs = Definitions(
    assets=[
        upstream,
        downstream,
    ],
    resources={
        "polars_parquet_io_manager": PolarsParquetIOManager(
            base_dir=f"s3://{get_dagster_bucket_name()}/path/to/folder"
        ),
    },
)

The upstream asset materializes and stores to S3 fine. However, calling .collect() in downstream fails with:

The above exception was caused by the following exception: polars.exceptions.ComputeError: Generic S3 error: Missing bucket name

I was hoping that adding an IO manager will help abstract away the bucket name / file location, but surprisingly this isn't working. Is there anything I could be missing in this setup?

[ilyanekhay-uta]

For the time being, I've been able to successfully patch this issue by doing the following:

1. Copied the dagster_polars.io_managers.parquet.scan_parquet function into my own code as patched_scan_parquet
2. Changed the last line of it from

return pl.scan_parquet(str(path), storage_options=storage_options, **kwargs)  # type: ignore

to

return pl.scan_parquet(str(path), storage_options=get_storage_options(), **kwargs)  # type: ignore

where get_storage_options() looks like this:

def get_storage_options():
    # Reference: https://docs.rs/object_store/latest/object_store/aws/enum.AmazonS3ConfigKey.html
    boto_session = boto3.session.Session()
    credentials: Credentials = boto_session.get_credentials()
    storage_options = {
        "aws_access_key_id": credentials.access_key,
        "aws_secret_access_key": credentials.secret_key,
        "default_region": boto_session.region_name,
        "aws_bucket_name": get_dagster_bucket_name(),
    }
    if credentials.token is not None:
        storage_options["aws_session_token"] = credentials.token
    return storage_options

3. Patched the implementation:

dagster_polars.io_managers.parquet.scan_parquet = patched_scan_parquet

[aschreiber1]

It looks like there is still a bug here, it looks like the read and write clients expect different arguments in storage_options and there is no way to unify them. The write client expects things like:
storage_options={
"endpoint_url": endpoint_override,
"key": access_key,
"secret": secret_key
}
But the read client expects:
storage_options={
"endpoint_url": endpoint_override,
"aws_secret_access_key_id": secret_key,
"aws_access_key_id": access_key
}