Skip to content

fix: isolate object-store registration by backend and configuration - #5503

Merged
sunchao merged 3 commits into
apache:mainfrom
sunchao:dev/chao/codex/comet-isolation-followup
Sep 10, 2026
Merged

fix: isolate object-store registration by backend and configuration#5503
sunchao merged 3 commits into
apache:mainfrom
sunchao:dev/chao/codex/comet-isolation-followup

Conversation

@sunchao

@sunchao sunchao commented Aug 27, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #5502.

Rationale for this change

A configuration can route s3 through Hadoop while leaving s3a on the native backend. Normalizing native s3a to s3 currently gives those different backends the same process-cache key. Separately, DataFusion's registry identifies stores by scheme and authority, so registering a second configuration can overwrite the first mapping.

This extracts the object-store correctness change from #5453. It contains no new scan counters or producer-lifecycle changes and does not depend on that metrics PR.

What changes are included in this PR?

Include backend identity in the cache key and register each non-local backend/configuration under a deterministic internal URL. The first registration receives the same identity it would receive after another store, so planning order cannot affect routing. Native s3 and s3a share canonical s3 identity; Hadoop-selected schemes retain their physical spelling. Native local files keep file:// behavior.

Reconstruct the physical URI before configuring encryption lookup, stripping only a complete internal identity suffix. This preserves existing s3/s3a key normalization and custom Hadoop schemes. The change isolates object-store lookup, not file-metadata caches.

How are these changes tested?

The focused tests use two distinct in-memory stores and verify actual returned bytes after registration in both orders. Cases cover a normalized S3 alias/backend collision, different native configurations, custom Hadoop routing, native S3 alias reuse, and native versus Hadoop-routed file URLs. Encryption-option tests compare ordinary and isolated physical URIs, including aliases, ports, custom schemes, and local files.

September 10 rebase validation

Rebased onto Apache main at 2d3eca2100d8d8684d31c496b376388a5fd79f18, which includes #5825. The backend decision returned by NormalizedObjectStoreUrl is preserved in the cache key, store selection, and Comet's synthetic registry URL. The two upstream native S3 routing tests now expect the configuration-specific registry URLs while retaining their bucket, path, and explicit Hadoop-routing assertions. All five isolation tests and the encryption URI regression remain enabled.

cargo fmt --all --check and git diff --check passed. An independent source review found no rebase integration issues; the diff remains limited to the two Parquet files, with dependency manifests and lockfiles unchanged.

Attempted from native/ with JDK 21:

cargo test --locked -p datafusion-comet --lib --no-default-features parquet::parquet_support::tests

Dependency resolution stopped before compilation because the configured registry mirror lacks aws-smithy-runtime-api 1.16.0, required by the upstream lockfile. Native compilation and the isolation/encryption tests therefore remain unverified locally for this rebase. Fresh hosted CI is queued. Its default-feature native suite covers the five isolation tests and encryption URI regression; the two upstream store-preparation routing tests require --no-default-features and are excluded from that hosted suite.

Historical validation before this rebase

The full native crate compiled with the two disjoint extracted changes combined at their shared base. All 90 Parquet tests passed, including the six new isolation/encryption regressions, with a confirmed zero process exit status. The four companion producer-lifecycle tests also passed. cargo fmt --all --check and git diff --check pass. Default native features and JDK 21 were used; tests use deterministic in-memory stores rather than a live Hadoop/cloud deployment. Full Spark integration and this branch's hosted CI have not run locally.

@comphead

Copy link
Copy Markdown
Contributor

would be related to #5314

// Registration URLs use a reserved suffix to distinguish backend/configuration
// identities. Encryption must use the physical URI that Spark registered. Match
// the complete suffix, from the right, so custom Hadoop schemes are preserved.
fn physical_object_store_scheme(object_store_url: &ObjectStoreUrl) -> &str {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we already have a similar method is_hdfs_scheme maybe its time to unify them

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I kept these separate because they have different roles: is_hdfs_scheme uses the original URL and configuration to select the Hadoop backend; physical_object_store_scheme removes the internal registration suffix to recover the scheme encryption expects. The latter handles registration URLs for both native and Hadoop stores, so it needs to remain independent of backend selection.

}

type ObjectStoreCache = RwLock<HashMap<(String, u64), Arc<dyn ObjectStore>>>;
type ObjectStoreCacheKey = (String, u64, bool);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please document what is (String, u64, bool) represents

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added documentation directly above ObjectStoreCacheKey in deaf061. The fields are the scheme/host/port, configuration hash, and whether the Hadoop backend is selected (true for Hadoop, including custom schemes routed through it; false for native). Native s3a is normalized to s3, while Hadoop-selected schemes retain their spelling.

} else {
let backend = if is_hdfs_scheme { "hdfs" } else { "native" };
ObjectStoreUrl::parse(format!(
"{scheme}+comet-{config_hash:016x}-{backend}://{}",

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why would comet be a hardcoded part of url?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+comet- is a reserved marker for our internal registration URLs. DataFusion keys stores only by scheme and authority, so we put the configuration hash and backend identity in the scheme while preserving the physical authority. This prevents a later scan of the same bucket with different credentials or backend selection from replacing an earlier scan's store. Encryption lookup recognizes and strips the complete suffix to recover the physical URI. Added a comment explaining this in deaf061.

@comphead comphead left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @sunchao makes sense to me, some small nits

@andygrove andygrove added bug Something isn't working area:scan Parquet scan / data reading labels Sep 6, 2026

@andygrove andygrove left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I read through prepare_object_store_with_configs and the registration path in parquet_support.rs. Two things in here matter together. Adding is_hdfs_scheme to the cache key stops a Hadoop-routed s3 and a native s3a from landing in the same process-cache entry when fs.comet.libhdfs.schemes=s3 is set. Baking that same identity into the scheme of the URL passed to register_object_store stops a second scan in the same RuntimeEnv from overwriting the first scan's registry mapping. That second part is the one that could have handed an existing scan the wrong store or the wrong credentials before this fix. The check_isolated_stores tests cover this well: they prepare two distinct in-memory stores in both orders and read actual bytes back through the returned ObjectStoreUrl, which is a more convincing check than comparing URL strings.

The uri_base fix in parquet_exec.rs is needed for the same reason. Once object_store_url carries the synthetic scheme+comet-hash-backend marker, CometEncryptionConfig would otherwise hand that synthetic URI to the JVM key retriever instead of the physical one Spark registered keys under. physical_object_store_scheme's use of rsplit_once against the 16-hex-digit-plus-backend shape correctly leaves a physical scheme alone when it already contains +comet- but doesn't match that exact grammar, and the custom+comet-existing case in the test covers that directly.

On comphead's question about why comet ends up hardcoded into the URL, at parquet_support.rs:600, I think the answer is that register_object_store only keys on scheme and authority, so the scheme is the only part of the URL available to carry backend and config identity without changing how DataFusion resolves stores more broadly. That reasoning isn't in the code right now. Would a short comment there be worth adding so the next reader doesn't have to rediscover it?

@andygrove

Copy link
Copy Markdown
Member

@sunchao could you respond to @comphead's questions?

@sunchao

sunchao commented Sep 9, 2026

Copy link
Copy Markdown
Member Author

@comphead Thanks for the review. I added the cache-key and registration-URL documentation in deaf061 and replied to all three inline comments.

Agreed that #5314 is related through scheme normalization; this PR keeps backend/configuration identities separate in the cache and registry after routing and normalization.

This follow-up adds eight documentation lines. cargo fmt --all --check and git diff --check passed. Runtime tests were not rerun for this documentation-only update; CI is queued.

@comphead comphead left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @sunchao

@andygrove

Copy link
Copy Markdown
Member

@sunchao could you take a look at the test failure?

      failures:
          parquet::parquet_support::tests::isolates_backends_even_when_s3_alias_and_configs_match
  
      test result: FAILED. 0 passed; 1 failed; 0 ignored; 0 measured; 343 filtered out; finished in 0.34s
  
    stderr ───
  
      thread 'parquet::parquet_support::tests::isolates_backends_even_when_s3_alias_and_configs_match' (13793) panicked at core/src/parquet/parquet_support.rs:818:13:
      assertion `left != right` failed
        left: ObjectStoreUrl { url: Url { scheme: "s3+comet-ba6b8898acdeb3c1-hdfs", cannot_be_a_base: false, username: "", password: None, host: Some(Domain("comet-isolation-backend-alias")), port: None, path: "/", query: None, fragment: None } }
       right: ObjectStoreUrl { url: Url { scheme: "s3+comet-ba6b8898acdeb3c1-hdfs", cannot_be_a_base: false, username: "", password: None, host: Some(Domain("comet-isolation-backend-alias")), port: None, path: "/", query: None, fragment: None } }

@andygrove

Copy link
Copy Markdown
Member

I chased the test failure I pointed you at rather than leaving it with you, and it is not your branch. Sorry for the noise.

isolates_backends_even_when_s3_alias_and_configs_match passes on deaf061bb as it stands. I checked out the head and ran the five isolation tests locally:

test parquet::parquet_support::tests::native_s3_aliases_share_cache_and_registration_identity ... ok
test parquet::parquet_support::tests::keeps_native_file_url_separate_from_explicit_hadoop_file_routing ... ok
test parquet::parquet_support::tests::preserves_custom_hadoop_scheme_when_routing_changes ... ok
test parquet::parquet_support::tests::isolates_backends_even_when_s3_alias_and_configs_match ... ok
test parquet::parquet_support::tests::isolates_native_stores_with_different_configurations ... ok

CI tests the merge with main, and the branch point here is ec0f7975, before #5314 landed. That is what @comphead was pointing at with the #5314 reference. #5314 replaced

let is_hdfs_scheme = is_hdfs_scheme(&url, object_store_configs);
let mut scheme = url.scheme();
if !is_hdfs_scheme && scheme == "s3a" { scheme = "s3"; url.set_scheme("s3")?; }

with

let url = normalize_object_store_url(url.as_str(), object_store_configs)?;
let is_hdfs_scheme = is_hdfs_scheme(&url, object_store_configs);

so the libhdfs classification now runs against the rewritten scheme. Your test's first case is s3a input with fs.comet.libhdfs.schemes=s3, which normalizes to s3:// and then matches the list, so it comes back as the Hadoop backend just like the second case. Both URLs end up s3+comet-ba6b8898acdeb3c1-hdfs, assert_ne! fires, and the values in the failure I quoted are exactly that.

I confirmed the mechanism on 424c31aa7 directly rather than inferring it from the diff:

PROBE typed scheme      = s3a
PROBE typed  is_hdfs    = false
PROBE normalized scheme = s3
PROBE final  is_hdfs    = true

So this is a live behaviour bug on main, not a test problem: with fs.comet.libhdfs.schemes=s3 and no s3a entry, an s3a:// scan is now served by create_hdfs_object_store instead of the native S3 store. I filed it as #5816 with the reproducer. The converse is fine, since listing s3a makes normalize_object_store_url return early and the spelling survives.

Which leaves the question of ordering. Your test is the right regression guard for #5816, so I would rather it landed here and #5816 was fixed on top of it than have the test weakened to accommodate the current main. Could you rebase onto main and confirm whether you would prefer to fold the one-line ordering fix into this PR, given that it is in the function you are already changing, or keep #5816 separate and mark the test #[ignore] with a link until it lands? I have a mild preference for folding it in: taking the classification once from the typed scheme and threading it into normalize_object_store_url is a couple of lines, and it makes the backend component of your new cache key mean what it says.

The documentation follow-up in deaf061bb reads well, and the cache-key and registration-URL notes cover the question I raised about why comet ends up in the scheme.

@sunchao
sunchao force-pushed the dev/chao/codex/comet-isolation-followup branch from deaf061 to c965dd5 Compare September 10, 2026 16:10
@sunchao
sunchao merged commit bd1aa66 into apache:main Sep 10, 2026
79 checks passed
@sunchao

sunchao commented Sep 10, 2026

Copy link
Copy Markdown
Member Author

Merged, thanks!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:scan Parquet scan / data reading bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Avoid object-store cache and registry collisions across backends and configurations

3 participants