fix(merge): make demand folding a per-dimension join

crates/lance-context-core/src/rollout_append.rs

@@ -273,6 +273,15 @@ pub enum StageEvent {
     Failed(String),
 }
 
+/// Merged generation watermark per shard, read from the base table's MemWAL
+/// index at its latest version. Public so executors can report what they
+/// merged alongside releasing ownership. Empty when the table has no WAL
+/// index yet.
+pub async fn merged_watermarks(uri: &str) -> Result<HashMap<Uuid, u64>> {
+    let dataset = Dataset::open(uri).await?;
+    watermarks(&dataset).await
+}
+
 pub(crate) async fn watermarks(dataset: &Dataset) -> Result<HashMap<Uuid, u64>> {
     let indices = dataset.load_indices().await?;
     let Some(index) = indices.iter().find(|i| i.name == MEM_WAL_INDEX_NAME) else {

crates/lance-context-master/src/maintenance_execution.rs

@@ -261,7 +261,7 @@ where
     let scope = MergeWriteScope::with_pinned_authorizer(Arc::new(Guard {
         coordinator: coordinator.clone(),
         execution: running.clone(),
-        uri,
+        uri: uri.clone(),
     }));
     let executor_id = state.admission.status().executor_id;
     let outcome = std::panic::AssertUnwindSafe(async {
@@ -321,9 +321,40 @@ where
     {
         return Err("maintenance ownership changed before completion".into());
     }
-    if !coordinator.release(&proof, &terminal).await? {
+    // A successful WAL catch-up reclaimed generations: publish the merged
+    // watermark per shard in the release transaction so demand and ownership
+    // change together. Other maintenance kinds do not move WAL watermarks. A
+    // watermark read failure must not strand a completed merge; release
+    // without demand and let the next scan correct it.
+    let merged = if maintenance == MaintenanceKind::Catchup && outcome.is_ok() {
+        match lance_context_core::rollout_append::merged_watermarks(&uri).await {
+            Ok(marks) => marks
+                .into_iter()
+                .map(
+                    |(shard, generation)| lance_context_merge::demand::ShardMerged {
+                        shard: shard.to_string(),
+                        merged_through: generation,
+                    },
+                )
+                .collect(),
+            Err(error) => {
+                tracing::warn!(target = %running.target, %error,
+                    "merged watermarks unavailable; releasing without demand update");
+                Vec::new()
+            }
+        }
+    } else {
+        Vec::new()
+    };
+    if !coordinator
+        .release_with_demand(&proof, &terminal, &merged)
+        .await?
+    {
         return Err("maintenance completion claim lost".into());
     }
+    if !merged.is_empty() {
+        metrics::counter!("master_demand_merged_watermarks_total").increment(merged.len() as u64);
+    }
     outcome
 }
 

crates/lance-context-master/src/rollout_append.rs

@@ -626,6 +626,44 @@ mod tests {
                     format!("{target}-{id}")
                 );
             }
+            // The release published merged watermarks for both shards in
+            // the same transaction that gave up ownership, so a fold over
+            // the events shows no pending generations.
+            let marks = lance_context_core::rollout_append::merged_watermarks(&uri)
+                .await
+                .unwrap();
+            assert_eq!(marks.len(), 2, "{target}: {marks:?}");
+            let encoded: String = target.bytes().map(|b| format!("{b:02x}")).collect();
+            let events = first
+                .task_store
+                .etcd_client()
+                .clone()
+                .get(
+                    format!(
+                        "{}/demand-events/{encoded}/",
+                        first.config.etcd.etcd_prefix.trim_end_matches('/')
+                    ),
+                    Some(etcd_client::GetOptions::new().with_prefix()),
+                )
+                .await
+                .unwrap();
+            assert_eq!(events.kvs().len(), 2, "{target}");
+            let mut table = lance_context_merge::demand::TableDemand::default();
+            for kv in events.kvs() {
+                let event: lance_context_merge::demand::DemandEvent =
+                    serde_json::from_slice(kv.value()).unwrap();
+                assert_eq!(
+                    event.source,
+                    lance_context_merge::demand::EventSource::Executor
+                );
+                let expected = marks
+                    .iter()
+                    .find(|(shard, _)| shard.to_string() == event.shard)
+                    .map(|(_, generation)| *generation);
+                assert_eq!(event.merged_through, expected);
+                table.fold(&event, kv.mod_revision(), 0).unwrap();
+            }
+            assert_eq!(table.pending_generations(), 0, "{target}");
         }
     }