fix: include ABFS container in object store cache key

[peterxcli]

Which issue does this PR close?

Closes #4993.

Rationale for this change

ABFS URLs encode the container in URL userinfo, for example
abfss://container@account.dfs.core.windows.net/path. The process-wide object store cache
previously keyed stores by scheme, host, and port only, so different containers in the same
storage account shared one cache entry. Since each Azure object store is bound to the container
from its URL, a later read could reuse the first container's store and return data from the wrong
container.

This also prevents containers sharing one Hadoop configuration map from resolving different fs.azure.sas.<container>.<account> tokens into the same cached store.

What changes are included in this PR?

• Preserve URL userinfo when building object store keys for abfs and abfss, while retaining
  the existing host-only behavior for other schemes.
• Update the cache documentation to describe container-scoped keys.
• Add an offline regression test that prepares two containers in the same account with identical
  configuration and verifies that they resolve to distinct object store instances.

How are these changes tested?

• cargo test --manifest-path native/Cargo.toml -p datafusion-comet --lib
  (129 passed; 4 ignored)
• cargo clippy --manifest-path native/Cargo.toml -p datafusion-comet --lib --tests -- -D warnings
• cargo fmt --manifest-path native/Cargo.toml --all -- --check

The regression constructs the Azure stores locally and does not issue network requests.

[andygrove]

Nice targeted fix. I confirmed the analysis by reading native/core/src/parquet/objectstore/azure.rs — MicrosoftAzureBuilder::parse_url binds the container into the store instance and the resource Path is container-relative, so a store genuinely cannot be reused across containers. Position::BeforeUsername..AfterPort correctly includes the userinfo when present and collapses to BeforeHost when it isn't, so non-Azure schemes are untouched.

A few small things worth considering:

1. DataFusion's own registry has the same blind spot. RuntimeEnv::register_object_store uses get_url_key, which also strips userinfo (datafusion-execution/src/object_store.rs). So within a single RuntimeEnv, only one abfss store per account can be registered — a second registration would overwrite the first at the DF layer. Today this is fine because Comet constructs a fresh RuntimeEnv per Parquet file read, as the module doc calls out. But that invariant is now load-bearing for correctness on Azure. A sentence tying the two together in the module doc block would help a future reader who tries to share a RuntimeEnv across container reads.

2. A one-line comment on the is_azure_scheme(scheme) branch explaining that ABFS URLs encode the container in the userinfo would save future readers a trip to Object store cache key drops the ABFS container, so two containers in one storage account share a store instance #4993.

3. Two small extensions to the new test, if you want them:

   • Assert that a second call with the same container returns the same Arc, so a future change that accidentally invalidates the same-URL cache path is also caught.
   • Add an s3://bucket@…-style case to pin down that the non-Azure branch still keys by host only.

None of these are blockers.

[peterxcli]

@andygrove thanks for the review, added comment on the abfs branch and new test coverage per your review comment

[peterxcli]

df PR: apache/datafusion#23935

[andygrove]

Fix looks good and I confirmed it works. Two containers in the same account now resolve to their own stores across tasks, where on main the second one gets the first's data.

One question before I approve. Within a single plan, prepare_object_store_with_configs registers against the shared session_ctx.runtime_env() (planner.rs:1528), and DataFusion's get_url_key strips userinfo, so a second container in the same account would overwrite the first in the registry and both scans would resolve to the same store. Your apache/datafusion#23935 fixes that upstream. Do you know whether a single plan can actually reference two containers in one account today? If it can, could we get a follow-up issue tracking the DF bump?

[peterxcli]

prepare_object_store_with_configs registers against the shared session_ctx.runtime_env() (planner.rs:1528), and DataFusion's get_url_key strips userinfo, so a second container in the same account would overwrite the first in the registry and both scans would resolve to the same store. Your apache/datafusion#23935 fixes that upstream. Do you know whether a single plan can actually reference two containers in one account today? If it can, could we get a follow-up issue tracking the DF bump?

We didn't use the datafusion's get_url_key.

[andygrove]

We didn't use the datafusion's get_url_key.

I think we may be talking past each other here. Comet doesn't call get_url_key directly, but runtime_env.register_object_store() and runtime_env.object_store() both go through DefaultObjectStoreRegistry, which does call it. In datafusion-execution-54.1.0/src/object_store.rs:268 it slices Position::BeforeHost..Position::AfterPort, and the doc comment on it reads "The credential info will be removed."

I put together a quick probe on your branch that registers two containers against one shared RuntimeEnv, the way planner.rs:1587 does:

#[test]
fn test_shared_runtime_env_collides_across_containers() {
    let configs = HashMap::from([("fs.azure.account.key".into(), "c2VjcmV0".into())]);
    // One plan => one shared RuntimeEnv, as in planner.rs
    let runtime_env = Arc::new(RuntimeEnv::default());
    let register = |container: &str| {
        super::prepare_object_store_with_configs(
            Arc::clone(&runtime_env),
            format!("abfss://{container}@shared-acct.dfs.core.windows.net/path/file.parquet"),
            &configs,
        )
        .unwrap()
        .0
    };
    let url_a = register("container-a");
    let url_b = register("container-b");
    let store_a = runtime_env.object_store(&url_a).unwrap();
    let store_b = runtime_env.object_store(&url_b).unwrap();
    println!("url_a={url_a} url_b={url_b}");
    println!("same store: {}", Arc::ptr_eq(&store_a, &store_b));
    println!("store_a={store_a:?}");
}

Output:

url_a=abfss://container-a@shared-acct.dfs.core.windows.net/
url_b=abfss://container-b@shared-acct.dfs.core.windows.net/
same store: true
store_a=MicrosoftAzure { ... container: "container-b" ... }

Both URLs resolve to the same store, and that store is bound to container-b. So a scan of container-a would silently read container-b's data. FileScanConfig resolves the store at execute time (datafusion-datasource/src/file_scan_config/mod.rs:640), after every scan in the plan has already registered, so the last registration wins. A join or union across two containers in one storage account would hit this. The native_iceberg_compat path in parquet/mod.rs:163 is safe because it builds a fresh SessionContext per file, but native_datafusion shares the task-scoped one.

Would you be up for closing this in the same PR? Comet already builds its RuntimeEnv through RuntimeEnvBuilder at jni_api.rs:559, and DF 54.1 exposes RuntimeEnvBuilder::with_object_store_registry, so a small registry impl that keys on the full authority would fix it now without waiting on apache/datafusion#23935 and the version bump. If you'd rather keep this PR scoped to the cache key, that's reasonable, but could we get an issue filed for the registry side so it doesn't get lost?

Two smaller things while I'm here.

The two new tests build a fresh RuntimeEnv per call, which is what makes the assertions hold. That's the right shape for testing the Comet cache in isolation. It might be worth adding one more that shares a single RuntimeEnv across both registrations, since that's what planner.rs actually does and it's the case that still resolves to one store today.

On the module doc at parquet_support.rs:501, since the container is now part of the Comet cache key but not part of the DataFusion registry key, could we add a sentence noting that per-container isolation depends on the registry side too? As written a reader could come away thinking the key change alone is sufficient.

One thing I noticed that works in your favour and isn't mentioned in the description: translate_hadoop_configs resolves fs.azure.sas.<container>.<account>, so before this change two containers sharing one config map would resolve different SAS tokens into the same cached store. This fixes that as well.

[peterxcli]

@andygrove thanks for further explanation! you're really right on this. I've addressed all of your review/suggestion.

runtime_env.register_object_store() and runtime_env.object_store() both go through DefaultObjectStoreRegistry, which does call it. In datafusion-execution-54.1.0/src/object_store.rs:268 it slices Position::BeforeHost..Position::AfterPort, and the doc comment on it reads "The credential info will be removed."

I put together a quick probe on your branch that registers two containers against one shared RuntimeEnv, the way planner.rs:1587 does:
....
Both URLs resolve to the same store, and that store is bound to container-b. So a scan of container-a would silently read container-b's data. FileScanConfig resolves the store at execute time (datafusion-datasource/src/file_scan_config/mod.rs:640), after every scan in the plan has already registered, so the last registration wins. A join or union across two containers in one storage account would hit this. The native_iceberg_compat path in parquet/mod.rs:163 is safe because it builds a fresh SessionContext per file, but native_datafusion shares the task-scoped one.

Thanks for catching this! sorry my previous response just refuse to address this because that's my blindspot.

The two new tests build a fresh RuntimeEnv per call, which is what makes the assertions hold. That's the right shape for testing the Comet cache in isolation. It might be worth adding one more that shares a single RuntimeEnv across both registrations, since that's what planner.rs actually does and it's the case that still resolves to one store today.

added as test_shared_runtime_env_uses_distinct_azure_container_stores in native/core/src/parquet/parquet_support.rs

On the module doc at parquet_support.rs:501, since the container is now part of the Comet cache key but not part of the DataFusion registry key, could we add a sentence noting that per-container isolation depends on the registry side too? As written a reader could come away thinking the key change alone is sufficient.

added.

One thing I noticed that works in your favour and isn't mentioned in the description: translate_hadoop_configs resolves fs.azure.sas.<container>.<account>, so before this change two containers sharing one config map would resolve different SAS tokens into the same cached store. This fixes that as well.

Nice!