feat(mem_wal): thread store_params + session through WAL write/read paths#7735
Conversation
📝 WalkthroughWalkthroughChangesMemWAL now preserves dataset object-store parameters and sessions across writers, flushers, derived-generation opens, caches, prewarming, block-list computation, and scan planning. Derived generation parameters clear only the path-bound store binding while retaining other settings. MemWAL storage-context propagation
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
…aths Federated (namespace-managed) tables open with a refreshing credential accessor that Dataset now stashes in store_params. Reuse it for every derived-URI open so WAL writes (ShardWriter / MemTableFlusher) and fresh-tier reads (flushed-generation scanner) resolve to the same vended-credential store instead of ambient identity. - Dataset::store_params() accessor - ShardWriterConfig + MemTableFlusher carry store_params + session - LsmScanner + 4 planners + flushed_cache + block_list thread store_params Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
- cargo fmt the mem_wal write/scanner sources - drop redundant clone on owned branch_location().path - add store_params/session fields to mem_wal_write bench config Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Condense the repeated "reuse the base's store / federated refreshing accessor / instead of an ambient one" rationale into succinct one-liners that describe behavior. No functional change. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
1cf18a9 to
7caf63a
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/lance/src/dataset/mem_wal/api.rs`:
- Around line 710-720: Update the configuration assignment in the relevant
flusher setup to use the existing Dataset::store_params() accessor followed by
cloned(), matching the call at the other mem-WAL site; remove the direct
self.store_params.as_deref() access while leaving the session and object-store
handling unchanged.
In `@rust/lance/src/dataset/mem_wal/scanner/flushed_cache.rs`:
- Around line 262-263: Extend the tests around cache.get_or_open, including the
cases at the referenced call sites, to pass a non-None ObjectStoreParams
containing a custom object_store_wrapper. Assert that opening the DatasetBuilder
invokes the wrapper and propagates the supplied store parameters, while
retaining the existing None-path coverage.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: f1fd31a8-17c9-42fb-9b79-a8b7243cebe8
📒 Files selected for processing (12)
rust/lance/benches/mem_wal/write/mem_wal_write.rsrust/lance/src/dataset.rsrust/lance/src/dataset/mem_wal/api.rsrust/lance/src/dataset/mem_wal/memtable/flush.rsrust/lance/src/dataset/mem_wal/scanner/block_list.rsrust/lance/src/dataset/mem_wal/scanner/builder.rsrust/lance/src/dataset/mem_wal/scanner/flushed_cache.rsrust/lance/src/dataset/mem_wal/scanner/fts_search.rsrust/lance/src/dataset/mem_wal/scanner/planner.rsrust/lance/src/dataset/mem_wal/scanner/point_lookup.rsrust/lance/src/dataset/mem_wal/scanner/vector_search.rsrust/lance/src/dataset/mem_wal/write.rs
| let first = cache.get_or_open(&uri, None, None).await.unwrap(); | ||
| let second = cache.get_or_open(&uri, None, None).await.unwrap(); |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
No test exercises store_params actually taking effect.
All updated tests pass None for the new store_params argument, so nothing here verifies that a supplied ObjectStoreParams is actually attached to the opened DatasetBuilder (e.g. via a custom object_store_wrapper and asserting it was invoked). Given credential propagation is the core motivation of this PR, a regression test covering the non-None path would materially increase confidence.
Also applies to: 294-294, 322-323, 344-349, 358-361, 399-399
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance/src/dataset/mem_wal/scanner/flushed_cache.rs` around lines 262 -
263, Extend the tests around cache.get_or_open, including the cases at the
referenced call sites, to pass a non-None ObjectStoreParams containing a custom
object_store_wrapper. Assert that opening the DatasetBuilder invokes the wrapper
and propagates the supplied store parameters, while retaining the existing
None-path coverage.
…URIs The deprecated `ObjectStoreParams::object_store` binding pins a store to one location: `ObjectStore::from_uri_and_params` and `DatasetBuilder::build_object_store` both take the path from its `Url` and ignore the URI they were asked to open. Threading the base dataset's params verbatim onto a `_mem_wal/<shard>/<gen>` URI therefore aimed every generation write and open at the base table itself — the flush failed with "dataset already exists", and derived opens (scanner planners, the standalone PK BTree, prewarm) returned base rows as generation rows. Adapt the params with `derived_store_params` wherever they cross to a derived URI, dropping only the path-bound store and keeping storage options, wrapper, credentials, and block size. The base table's own re-open keeps the params verbatim — that is where the binding still points correctly — so `base_storage_version` still inherits the base's storage version rather than silently falling back to the default. Splits the ambiguous `open_derived` into `open_base` and `open_generation` to make that distinction explicit. Only surfaced for datasets opened with an explicit object store; the `storage_options` path was unaffected.
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rust/lance/src/dataset/mem_wal/scanner/builder.rs (1)
344-357: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winAlign
with_store_paramswithnew()with_store_paramsshould storederived_store_params(&store_params)instead of the raw value.new()already strips the path-boundobject_storebinding so flushed-generation opens resolve against the generation URI; the setter should preserve the same invariant forwithout_base_tableand future callers.rust/lance/src/dataset/mem_wal/scanner/builder.rs:344-357🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance/src/dataset/mem_wal/scanner/builder.rs` around lines 344 - 357, Update Builder::with_store_params to store the result of derived_store_params(&store_params) rather than the raw ObjectStoreParams, matching the invariant established by Builder::new. Leave with_session unchanged and preserve the setter’s fluent return behavior for without_base_table and other callers.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/lance/src/dataset/mem_wal/util.rs`:
- Around line 145-151: Change the visibility of derived_store_params from pub to
pub(crate), keeping its implementation unchanged so it remains accessible
throughout the rust/lance crate without exposing it as part of the public API.
---
Outside diff comments:
In `@rust/lance/src/dataset/mem_wal/scanner/builder.rs`:
- Around line 344-357: Update Builder::with_store_params to store the result of
derived_store_params(&store_params) rather than the raw ObjectStoreParams,
matching the invariant established by Builder::new. Leave with_session unchanged
and preserve the setter’s fluent return behavior for without_base_table and
other callers.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 44cc76e6-50d7-4607-9dc3-0a291e5e7cb1
📒 Files selected for processing (5)
rust/lance/src/dataset/mem_wal/api.rsrust/lance/src/dataset/mem_wal/memtable/flush.rsrust/lance/src/dataset/mem_wal/scanner/builder.rsrust/lance/src/dataset/mem_wal/util.rsrust/lance/src/dataset/mem_wal/write.rs
westonpace
left a comment
westonpace
left a comment
There was a problem hiding this comment.
Review from Claude attached below. The only thing that seems noteworthy is perhaps the cache key issue. The breaking changes are probably fine at this stage. Do we expect users to have multiple object stores that differ only by store params?
Here is a thorough code review of PR #7735.
PR #7735: feat(mem_wal): thread store_params + session through WAL write/read paths
Overview
This PR fixes a correctness bug where MemWAL tables opened with vended credentials (via a namespace client) would fall back to ambient identity when opening derived URIs (WAL log, flushed generations). The core fix threads store_params and session from the base Dataset through the entire write/read pipeline so every derived-URI open uses the same credential store. A key subtlety is handled via derived_store_params: a path-bound object_store binding in ObjectStoreParams (the deprecated object_store field) cannot be reused for generation URIs, as it would silently redirect all opens/writes to the base table path.
Code Quality & Style
Strengths:
- The approach is principled:
Option<T>fields defaulting toNonemeans zero behavioral change for existing callers — no migration required. derived_store_paramsisolates the path-binding hazard in one place with a clear doc comment, preventing future callers from reintroducing the bug.- Comments explain the why (credential binding, path-binding hazard) rather than restating the code — consistent with project standards.
- Repetitive boilerplate (
with_store_paramsacross four planners) is unavoidable given the struct-per-planner design; it's not over-engineered.
Concerns:
-
ShardWriterConfigis a public struct with public fields. Addingstore_paramsandsessionto it is a breaking change for any downstream consumer that constructsShardWriterConfig { ..Default::default() }with field destructuring or exhaustive pattern matching. The PR notes this is consumed by a companion PR, but the struct ispub, so this deserves a comment or a#[non_exhaustive]annotation if it isn't already marked that way. Worth checking. -
with_storage_contexttakes twoOptions together. The method signaturewith_storage_context(store_params: Option<ObjectStoreParams>, session: Option<Arc<Session>>)is slightly unusual — typically builder methods take a concrete value, notOption. A caller might accidentally passNone, Noneand think they've configured storage. Consider two separate setters (with_store_params/with_session) that mirror the planner API, or at least add a note in the doc. -
FlushedMemTableCache::get_or_opencache key is the path only. After this PR, the same path can be opened with differentstore_params(e.g., if credentials rotate). Becausestore_paramsis not part of the cache key, a second call with different params would silently return the first-opened dataset. For static-credential use cases this is fine, but for refreshing credential accessors (the stated motivation) this is a latent correctness issue. The PR description says credential rotation is handled by the "refreshing credential accessor" inside the params — verify that the accessor itself is stateful and doesn't need re-injection on re-open. -
open_flushed_datasetparameter order changed. The new signature is(path, session, store_params, cache, warmer). Thestore_paramsbeing inserted betweensessionandcache(rather than at the end or grouped withsession) is a minor ergonomics issue, but more importantly it increases the risk of callers passing arguments in the wrong order sincesessionandstore_paramsare bothOptions of different types.
Potential Issues / Risks
-
derived_store_paramscalled on every scan/flush — correct behavior, but involves a clone ofObjectStoreParamson each flushed-generation open. For the commonNonecase this is a no-op clone of an empty struct, so it's fine. Worth confirmingObjectStoreParamsisCloneat low cost (it is, based on the diff). -
open_basereusesstore_paramsverbatim whileopen_generationstripsobject_store. The comment explains this well. The asymmetry is correct but easy to misuse in future — a doc clarification or an assertion (debug_assert!(self.store_params.as_ref().map_or(true, |p| ...))) would guard against future regressions here. -
LsmScanner::newcallsderived_store_paramsat construction time (builder.rs:243). This is correct since the scanner only ever opens generation URIs, not the base. But ifLsmScannerever gains a path to open the base itself, this early adaptation would silently break it. A note inLsmScanner::newpointing to this invariant would help. -
Dataset::store_params()returnsOption<&ObjectStoreParams>, exposing the deprecatedobject_storefield publicly. If callers inspect this field they'll hit the deprecated warning. This is unavoidable since the deprecated field is needed by thederived_store_paramslogic, but the accessor's doc comment doesn't warn consumers to usederived_store_paramswhen using params for a derived path — consider adding a note. -
Test coverage: The regression test
test_flush_and_read_with_path_bound_object_storecovers the path-bound store binding case. However, the stated motivation is vended credentials via a refreshing accessor — there's no test that verifies thestorage_options_accessor(theStorageOptionsAccessor) is actually propagated through to a generation open. The unit testtest_derived_store_params_drops_only_the_path_bound_storechecks the struct transformation, but an integration-level test of the accessor actually being invoked would be more robust.
Minor Issues
api.rs:600— inline comment// Every open below targets a generation URI, never the base's own.is true only for theprewarmpath. The comment is slightly misleading given the function header context.- Benchmark file
mem_wal_write.rs:658— passthrough ofstore_params/sessionfromdefault_configis correct but adds no benchmark coverage for the vended-credential path. Not a blocker. write.rsimportsuse crate::session::Session;is added afteruse super::manifest::ShardManifestStore;, breaking the conventional group ordering (external → internal). Minor but inconsistent with the surrounding code style.
Summary
The PR is well-scoped, mechanically correct, and the derived_store_params abstraction cleanly isolates the tricky path-binding concern. The main actionable concerns are:
- Cache key correctness with rotating credentials — confirm the accessor handles this without re-injection.
ShardWriterConfigbreaking change — add#[non_exhaustive]if not already present.- Test gap — a test that the
storage_options_accessoris propagated through to a generation open would close the gap between the stated motivation (refreshing credentials) and what is actually exercised.
westonpace
left a comment
westonpace
left a comment
There was a problem hiding this comment.
(meant to mark that previous review as approve)
…rams `LsmScanner::new` adapts the base's params for generation URIs, but the `with_store_params` setter stored them raw — and `without_base_table` leaves `store_params: None`, so the setter is the *only* way a fresh-tier-only scanner gets credentials. A caller handing it the base's params (the intended use) therefore carried the path-bound `object_store` binding onto every generation open, redirecting it at the base table: the same bug b7a07d0 fixed for `new`, still live on the one path the downstream consumer actually uses. Apply `derived_store_params` in the setter so both entry points hold the invariant. Also narrow `derived_store_params` to `pub(crate)` (no caller outside the crate now needs it), reuse the `Dataset::store_params()` accessor in `mem_wal_writer` instead of re-inlining it, and document why `FlushedMemTableCache` keys on path alone: credential rotation is safe because a vended store holds the live `StorageOptionsAccessor` and re-resolves per request, but one path served under two different `ObjectStoreParams` would reuse the first opener's store. Tests: the store params a base was opened with must reach the generation write (flush) and the generation read (scan), and a fresh-tier scan configured through `with_store_params` must resolve the generation rather than the base. Both were verified to fail when the threading is removed at each site. The observable signal differs per direction: `ObjectStore::create` writes local files straight through `tokio::fs`, so a generation's fragments never reach a wrapping store under `file://` and the write side keys on the manifest instead; the read side keys on the fragments, since the flusher has already warmed the manifest into the shared session cache. Co-Authored-By: Claude Opus 4.8 <[email protected]>
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
rust/lance/src/dataset/mem_wal/util.rs (1)
145-151: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider adding a doc comment to
derived_store_params.The
pub(crate)visibility fix from the prior review is in place. The function itself still has no doc comment, even though its non-obvious behavior (dropping the path-boundobject_storebinding while preserving credentials/storage options/wrapper/block size) is exactly the subject of extensive tests and call-site docs elsewhere (e.g.with_store_paramsinbuilder.rs). Since this helper is shared acrossapi.rs,builder.rs, and the flusher, consolidating that rationale at the definition would help future readers who land here first.♻️ Proposed doc comment
+/// Adapt `ObjectStoreParams` opened for the base table into params safe to +/// reuse for a *derived* URI (a flushed generation). Drops the deprecated, +/// path-bound `object_store` binding — reusing it would redirect every +/// generation open at the base table's own path — while preserving storage +/// options, credentials, and the object-store wrapper. pub(crate) fn derived_store_params(params: &ObjectStoreParams) -> ObjectStoreParams {🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance/src/dataset/mem_wal/util.rs` around lines 145 - 151, Add a concise Rust doc comment directly above derived_store_params explaining that it clones ObjectStoreParams while clearing the path-bound object_store and preserving credentials, storage options, wrapper, and block size for reuse by callers.rust/lance/src/dataset/mem_wal/scanner/flushed_cache.rs (1)
139-144: 🗄️ Data Integrity & Integration | 🟠 Major | ⚖️ Poor tradeoffKeep
DatasetCache::get_or_openbackward-compatible.DatasetCacheis public and re-exported, so adding a required parameter here breaks any out-of-crateimpl DatasetCache; add a new method or deprecate the old signature instead.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance/src/dataset/mem_wal/scanner/flushed_cache.rs` around lines 139 - 144, Preserve the existing public `DatasetCache::get_or_open` signature so external implementations remain source-compatible. Add a separate method for the new `session` and `store_params` inputs, or retain the old method as a wrapper with appropriate defaults, and update internal callers to use the extended behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/lance/src/dataset/mem_wal/write.rs`:
- Around line 6871-6964: Update both new scanner tests, including
test_store_params_reach_generation_write_and_read and the corresponding test
near it, to use LsmScanner::try_into_batch() instead of
try_into_stream().try_collect::<Vec<_>>(). Preserve each test’s existing
row-count and column assertions by deriving them from the single returned
RecordBatch, and remove the now-unused TryStreamExt import if no longer needed.
---
Outside diff comments:
In `@rust/lance/src/dataset/mem_wal/scanner/flushed_cache.rs`:
- Around line 139-144: Preserve the existing public `DatasetCache::get_or_open`
signature so external implementations remain source-compatible. Add a separate
method for the new `session` and `store_params` inputs, or retain the old method
as a wrapper with appropriate defaults, and update internal callers to use the
extended behavior.
In `@rust/lance/src/dataset/mem_wal/util.rs`:
- Around line 145-151: Add a concise Rust doc comment directly above
derived_store_params explaining that it clones ObjectStoreParams while clearing
the path-bound object_store and preserving credentials, storage options,
wrapper, and block size for reuse by callers.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 39b59921-c22a-4e32-8d07-d8ce26899ccf
📒 Files selected for processing (6)
rust/lance/src/dataset/mem_wal/api.rsrust/lance/src/dataset/mem_wal/scanner/builder.rsrust/lance/src/dataset/mem_wal/scanner/flushed_cache.rsrust/lance/src/dataset/mem_wal/test_util.rsrust/lance/src/dataset/mem_wal/util.rsrust/lance/src/dataset/mem_wal/write.rs
| async fn test_store_params_reach_generation_write_and_read() { | ||
| use crate::dataset::builder::DatasetBuilder; | ||
| use crate::dataset::mem_wal::scanner::{LsmScanner, ShardSnapshot}; | ||
| use crate::dataset::mem_wal::test_util::observable_store_params; | ||
| use futures::TryStreamExt; | ||
| use tempfile::TempDir; | ||
|
|
||
| let vector_dim = 8; | ||
| let schema = create_test_schema(vector_dim); | ||
| let temp_dir = TempDir::new().unwrap(); | ||
| let uri = format!("file://{}", temp_dir.path().display()); | ||
|
|
||
| let initial = create_test_batch(&schema, 0, 16, vector_dim); | ||
| let batches = RecordBatchIterator::new([Ok(initial)], schema.clone()); | ||
| Dataset::write(batches, &uri, Some(WriteParams::default())) | ||
| .await | ||
| .expect("Failed to create dataset"); | ||
|
|
||
| // Open the base through an observable store, exactly as a namespace | ||
| // client would hand in a vended-credential store. | ||
| let (store_params, controls) = observable_store_params(); | ||
| let mut dataset = DatasetBuilder::from_uri(&uri) | ||
| .with_store_params(store_params) | ||
| .load() | ||
| .await | ||
| .expect("Failed to open dataset"); | ||
| dataset | ||
| .initialize_mem_wal() | ||
| .execute() | ||
| .await | ||
| .expect("Failed to initialize MemWAL"); | ||
|
|
||
| let shard_id = Uuid::new_v4(); | ||
| let writer = dataset | ||
| .mem_wal_writer(shard_id, ShardWriterConfig::new(shard_id)) | ||
| .await | ||
| .expect("Failed to create writer"); | ||
| writer | ||
| .put(vec![create_test_batch(&schema, 1_000, 8, vector_dim)]) | ||
| .await | ||
| .expect("Failed to write"); | ||
| writer.force_seal_active().await.unwrap(); | ||
| writer.wait_for_flush_drain().await.expect("flush failed"); | ||
|
|
||
| let manifest = writer.manifest().await.unwrap().expect("manifest exists"); | ||
| assert_eq!(manifest.flushed_generations.len(), 1); | ||
| let flushed = manifest.flushed_generations[0].clone(); | ||
|
|
||
| // The generation's own Lance manifest is the signal to key on. Keying on | ||
| // the generation folder alone would pass vacuously: sidecars like | ||
| // `{gen}/bloom_filter.bin` are written through the *base* dataset's | ||
| // store, which is observable no matter what the params do. And the | ||
| // fragments can't be used either — `ObjectStore::create` writes local | ||
| // files through `tokio::fs`, bypassing the object store entirely, so | ||
| // `{gen}/data/` never reaches a wrapper under `file://`. The manifest | ||
| // goes through `put_opts`, and only the flusher's `Dataset::write` / | ||
| // `open_generation` writes it — both of which must carry the params. | ||
| let gen_manifest = format!("{}/_versions", flushed.path); | ||
|
|
||
| assert!( | ||
| controls.wrote_under(&gen_manifest), | ||
| "the flush must write the generation through the base's store params, \ | ||
| not a store resolved from the generation URI alone" | ||
| ); | ||
|
|
||
| // And the read path must resolve the generation through them too. | ||
| let snapshot = ShardSnapshot::new(shard_id) | ||
| .with_current_generation(manifest.current_generation) | ||
| .with_flushed_generation(flushed.generation, flushed.path.clone()); | ||
| let scanner = LsmScanner::new(Arc::new(dataset), vec![snapshot], vec!["id".to_string()]); | ||
| let rows: usize = scanner | ||
| .try_into_stream() | ||
| .await | ||
| .unwrap() | ||
| .try_collect::<Vec<_>>() | ||
| .await | ||
| .expect("scan failed") | ||
| .iter() | ||
| .map(|batch| batch.num_rows()) | ||
| .sum(); | ||
| assert_eq!(rows, 24); | ||
|
|
||
| // Reads key on the data files, not the manifest: the flusher already | ||
| // pulled the generation's manifest into the shared session cache, so the | ||
| // scan's open serves it from memory and never touches the store. The | ||
| // fragments are read through it (reads have no local bypass), as is the | ||
| // generation's standalone PK index. | ||
| assert!( | ||
| controls.read_under(&format!("{}/data/", flushed.path)), | ||
| "the scan must read the generation through the base's store params" | ||
| ); | ||
|
|
||
| writer.close().await.unwrap(); | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Use try_into_batch() instead of try_into_stream().try_collect() in the new tests.
Both new tests manually collect the stream into a Vec<RecordBatch> and then sum/flat-map row counts and columns. LsmScanner::try_into_batch() already exists for exactly this (single concatenated batch), simplifying both assertions.
As per coding guidelines, "Use .try_into_batch() instead of .try_into_stream().try_collect() for scanner results in tests."
♻️ Proposed simplification
- let scanner = LsmScanner::new(Arc::new(dataset), vec![snapshot], vec!["id".to_string()]);
- let rows: usize = scanner
- .try_into_stream()
- .await
- .unwrap()
- .try_collect::<Vec<_>>()
- .await
- .expect("scan failed")
- .iter()
- .map(|batch| batch.num_rows())
- .sum();
- assert_eq!(rows, 24);
+ let scanner = LsmScanner::new(Arc::new(dataset), vec![snapshot], vec!["id".to_string()]);
+ let batch = scanner.try_into_batch().await.expect("scan failed");
+ assert_eq!(batch.num_rows(), 24);- let arrow_schema: Arc<ArrowSchema> = schema.clone();
- let batches = LsmScanner::without_base_table(
- arrow_schema,
- uri.clone(),
- vec![snapshot],
- vec!["id".to_string()],
- )
- .with_session(dataset.session())
- .with_store_params(store_params)
- .try_into_stream()
- .await
- .unwrap()
- .try_collect::<Vec<_>>()
- .await
- .expect("scan must open the generation, not the base");
-
- let rows: usize = batches.iter().map(|batch| batch.num_rows()).sum();
- assert_eq!(
- rows, 8,
+ let arrow_schema: Arc<ArrowSchema> = schema.clone();
+ let batch = LsmScanner::without_base_table(
+ arrow_schema,
+ uri.clone(),
+ vec![snapshot],
+ vec!["id".to_string()],
+ )
+ .with_session(dataset.session())
+ .with_store_params(store_params)
+ .try_into_batch()
+ .await
+ .expect("scan must open the generation, not the base");
+
+ assert_eq!(
+ batch.num_rows(), 8,
"fresh tier holds only the 8 WAL rows; 16 means the generation open \
was redirected at the base table"
);
- let ids: Vec<i64> = batches
- .iter()
- .flat_map(|batch| {
- batch
- .column_by_name("id")
- .unwrap()
- .as_any()
- .downcast_ref::<Int64Array>()
- .unwrap()
- .values()
- .to_vec()
- })
- .collect();
+ let ids: Vec<i64> = batch
+ .column_by_name("id")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap()
+ .values()
+ .to_vec();Also applies to: 6971-7065
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance/src/dataset/mem_wal/write.rs` around lines 6871 - 6964, Update
both new scanner tests, including
test_store_params_reach_generation_write_and_read and the corresponding test
near it, to use LsmScanner::try_into_batch() instead of
try_into_stream().try_collect::<Vec<_>>(). Preserve each test’s existing
row-count and column assertions by deriving them from the single returned
RecordBatch, and remove the now-unused TryStreamExt import if no longer needed.
2 participants
Summary
Threads
store_params+sessionthrough the mem_wal write and fresh-tier read paths so WAL tables opened via a namespace client (with a refreshing credential accessor) resolve every derived-URI open to the same vended-credential store instead of ambient identity.Datasetalready stashes the options it was opened with — including the credential accessor — in itsstore_paramsfield; this exposes that and reuses it everywhere mem_wal opens a derived URI (WAL log, generations, flushed datasets).What changed
Dataset::store_params()accessor.ShardWriterConfig+MemTableFlushercarrystore_params+session; the flusher's generation opens/writes use them (open_derived,WriteParams).LsmScanner+ the four planners (scan / vector / fts / point-lookup) +flushed_cache+block_listthreadstore_paramsinto flushed-generation opens.All new fields default to
None, so existing (ambient) callers and the ~dozen test sites are unaffected in behavior.Context
Enables federated (namespace-managed) WAL tables with vended credentials in the downstream sophon consumer (companion PR there). Without this, WAL writes/reads of derived URIs sign with ambient credentials and fail against per-table vended stores (e.g. MinIO / S3-compatible endpoints supplied per-request).
Validation
503 mem_wal tests green; clippy clean. Exercised end-to-end against MinIO downstream (dir-namespace, vended creds, write → flush → compact → read).
🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
Enhancements