diff --git a/java/lance-jni/src/mem_wal.rs b/java/lance-jni/src/mem_wal.rs index c4f56d7b97d..9047c6288b0 100644 --- a/java/lance-jni/src/mem_wal.rs +++ b/java/lance-jni/src/mem_wal.rs @@ -1265,9 +1265,6 @@ fn build_writer_config(env: &mut JNIEnv, config: &JObject) -> Result Result( fn memtable_stats_to_java<'a>(env: &mut JNIEnv<'a>, stats: &MemTableStats) -> Result> { let max_buffered = box_u64_opt(env, stats.max_buffered_batch_position)?; - let max_flushed = box_u64_opt(env, stats.max_flushed_batch_position)?; let pending_start = box_u64_opt(env, stats.pending_wal_start_batch_position)?; let pending_end = box_u64_opt(env, stats.pending_wal_end_batch_position)?; Ok(env.new_object( "org/lance/memwal/MemTableStats", - "(JJJJLjava/lang/Long;Ljava/lang/Long;Ljava/lang/Long;Ljava/lang/Long;JJJ)V", + "(JJJJLjava/lang/Long;JJLjava/lang/Long;Ljava/lang/Long;JJJ)V", &[ JValueGen::Long(stats.row_count as i64), JValueGen::Long(stats.batch_count as i64), JValueGen::Long(stats.estimated_size as i64), JValueGen::Long(stats.generation as i64), JValueGen::Object(&max_buffered), - JValueGen::Object(&max_flushed), + JValueGen::Long(stats.durable_batch_count as i64), + JValueGen::Long(stats.global_offset as i64), JValueGen::Object(&pending_start), JValueGen::Object(&pending_end), JValueGen::Long(stats.pending_wal_batch_count as i64), diff --git a/java/src/main/java/org/lance/memwal/MemTableStats.java b/java/src/main/java/org/lance/memwal/MemTableStats.java index a993d771b42..58533e48194 100644 --- a/java/src/main/java/org/lance/memwal/MemTableStats.java +++ b/java/src/main/java/org/lance/memwal/MemTableStats.java @@ -24,7 +24,8 @@ public class MemTableStats { private final long estimatedSizeBytes; private final long generation; private final Optional maxBufferedBatchPosition; - private final Optional maxFlushedBatchPosition; + private final long durableBatchCount; + private final long globalOffset; private final Optional pendingWalStartBatchPosition; private final Optional pendingWalEndBatchPosition; private final long pendingWalBatchCount; @@ -37,7 +38,8 @@ public MemTableStats( long estimatedSizeBytes, long generation, Long maxBufferedBatchPosition, - Long maxFlushedBatchPosition, + long durableBatchCount, + long globalOffset, Long pendingWalStartBatchPosition, Long pendingWalEndBatchPosition, long pendingWalBatchCount, @@ -48,7 +50,8 @@ public MemTableStats( this.estimatedSizeBytes = estimatedSizeBytes; this.generation = generation; this.maxBufferedBatchPosition = Optional.ofNullable(maxBufferedBatchPosition); - this.maxFlushedBatchPosition = Optional.ofNullable(maxFlushedBatchPosition); + this.durableBatchCount = durableBatchCount; + this.globalOffset = globalOffset; this.pendingWalStartBatchPosition = Optional.ofNullable(pendingWalStartBatchPosition); this.pendingWalEndBatchPosition = Optional.ofNullable(pendingWalEndBatchPosition); this.pendingWalBatchCount = pendingWalBatchCount; @@ -81,9 +84,14 @@ public Optional maxBufferedBatchPosition() { return maxBufferedBatchPosition; } - /** Highest WAL batch position flushed from the MemTable, if any. */ - public Optional maxFlushedBatchPosition() { - return maxFlushedBatchPosition; + /** Writer-global count of WAL-durable batches (exclusive). */ + public long durableBatchCount() { + return durableBatchCount; + } + + /** Writer-global coordinate of this MemTable's batch 0. */ + public long globalOffset() { + return globalOffset; } /** First WAL batch position pending flush, if any. */ @@ -119,7 +127,8 @@ public String toString() { .add("estimatedSizeBytes", estimatedSizeBytes) .add("generation", generation) .add("maxBufferedBatchPosition", maxBufferedBatchPosition.orElse(null)) - .add("maxFlushedBatchPosition", maxFlushedBatchPosition.orElse(null)) + .add("durableBatchCount", durableBatchCount) + .add("globalOffset", globalOffset) .add("pendingWalStartBatchPosition", pendingWalStartBatchPosition.orElse(null)) .add("pendingWalEndBatchPosition", pendingWalEndBatchPosition.orElse(null)) .add("pendingWalBatchCount", pendingWalBatchCount) diff --git a/java/src/main/java/org/lance/memwal/ShardWriterConfig.java b/java/src/main/java/org/lance/memwal/ShardWriterConfig.java index 40310907d3e..d6bf2aeb383 100644 --- a/java/src/main/java/org/lance/memwal/ShardWriterConfig.java +++ b/java/src/main/java/org/lance/memwal/ShardWriterConfig.java @@ -28,7 +28,6 @@ */ public class ShardWriterConfig { private Optional durableWrite = Optional.empty(); - private Optional syncIndexedWrite = Optional.empty(); private Optional maxWalBufferSize = Optional.empty(); private Optional maxWalFlushIntervalMs = Optional.empty(); private Optional maxMemtableSize = Optional.empty(); @@ -36,8 +35,6 @@ public class ShardWriterConfig { private Optional maxMemtableBatches = Optional.empty(); private Optional maxUnflushedMemtableBytes = Optional.empty(); private Optional manifestScanBatchSize = Optional.empty(); - private Optional asyncIndexBufferRows = Optional.empty(); - private Optional asyncIndexIntervalMs = Optional.empty(); private Optional backpressureLogIntervalMs = Optional.empty(); private Optional statsLogIntervalMs = Optional.empty(); private List hnswParams = Collections.emptyList(); @@ -48,12 +45,6 @@ public ShardWriterConfig withDurableWrite(boolean durableWrite) { return this; } - /** Whether indexed writes are applied synchronously. */ - public ShardWriterConfig withSyncIndexedWrite(boolean syncIndexedWrite) { - this.syncIndexedWrite = Optional.of(syncIndexedWrite); - return this; - } - /** Maximum size of the in-memory WAL buffer, in bytes. */ public ShardWriterConfig withMaxWalBufferSize(long maxWalBufferSize) { Preconditions.checkArgument( @@ -118,26 +109,6 @@ public ShardWriterConfig withManifestScanBatchSize(long manifestScanBatchSize) { return this; } - /** Number of rows buffered before an asynchronous index update is triggered. */ - public ShardWriterConfig withAsyncIndexBufferRows(long asyncIndexBufferRows) { - Preconditions.checkArgument( - asyncIndexBufferRows >= 0, - "asyncIndexBufferRows must not be negative, got %s", - asyncIndexBufferRows); - this.asyncIndexBufferRows = Optional.of(asyncIndexBufferRows); - return this; - } - - /** Interval between asynchronous index updates, in milliseconds. */ - public ShardWriterConfig withAsyncIndexIntervalMs(long asyncIndexIntervalMs) { - Preconditions.checkArgument( - asyncIndexIntervalMs >= 0, - "asyncIndexIntervalMs must not be negative, got %s", - asyncIndexIntervalMs); - this.asyncIndexIntervalMs = Optional.of(asyncIndexIntervalMs); - return this; - } - /** Interval between backpressure log messages, in milliseconds. */ public ShardWriterConfig withBackpressureLogIntervalMs(long backpressureLogIntervalMs) { Preconditions.checkArgument( @@ -172,10 +143,6 @@ public Optional durableWrite() { return durableWrite; } - public Optional syncIndexedWrite() { - return syncIndexedWrite; - } - public Optional maxWalBufferSize() { return maxWalBufferSize; } @@ -204,14 +171,6 @@ public Optional manifestScanBatchSize() { return manifestScanBatchSize; } - public Optional asyncIndexBufferRows() { - return asyncIndexBufferRows; - } - - public Optional asyncIndexIntervalMs() { - return asyncIndexIntervalMs; - } - public Optional backpressureLogIntervalMs() { return backpressureLogIntervalMs; } diff --git a/java/src/test/java/org/lance/memwal/MemWalTest.java b/java/src/test/java/org/lance/memwal/MemWalTest.java index 57d0b5b81ad..3ee998733ea 100644 --- a/java/src/test/java/org/lance/memwal/MemWalTest.java +++ b/java/src/test/java/org/lance/memwal/MemWalTest.java @@ -409,7 +409,6 @@ void testShardWriterDeleteMasksBaseRow(@TempDir Path tempDir) throws Exception { ShardWriterConfig config = new ShardWriterConfig() .withDurableWrite(true) - .withSyncIndexedWrite(true) .withMaxWalBufferSize(1) .withMaxWalFlushIntervalMs(10); diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index def635d58a6..0aad8f5806b 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -5086,7 +5086,6 @@ def initialize_mem_wal( identity_column: Optional[str] = None, unsharded: bool = False, durable_write: Optional[bool] = None, - sync_indexed_write: Optional[bool] = None, max_wal_buffer_size: Optional[int] = None, max_wal_flush_interval_ms: Optional[int] = None, max_memtable_size: Optional[int] = None, @@ -5094,8 +5093,6 @@ def initialize_mem_wal( max_memtable_batches: Optional[int] = None, max_unflushed_memtable_bytes: Optional[int] = None, manifest_scan_batch_size: Optional[int] = None, - async_index_buffer_rows: Optional[int] = None, - async_index_interval_ms: Optional[int] = None, backpressure_log_interval_ms: Optional[int] = None, stats_log_interval_ms: Optional[int] = None, hnsw_params: Optional[Dict[str, Dict[str, int]]] = None, @@ -5156,7 +5153,6 @@ def initialize_mem_wal( identity_column=identity_column, unsharded=unsharded, durable_write=durable_write, - sync_indexed_write=sync_indexed_write, max_wal_buffer_size=max_wal_buffer_size, max_wal_flush_interval_ms=max_wal_flush_interval_ms, max_memtable_size=max_memtable_size, @@ -5164,8 +5160,6 @@ def initialize_mem_wal( max_memtable_batches=max_memtable_batches, max_unflushed_memtable_bytes=max_unflushed_memtable_bytes, manifest_scan_batch_size=manifest_scan_batch_size, - async_index_buffer_rows=async_index_buffer_rows, - async_index_interval_ms=async_index_interval_ms, backpressure_log_interval_ms=backpressure_log_interval_ms, stats_log_interval_ms=stats_log_interval_ms, hnsw_params=hnsw_params, @@ -5188,7 +5182,6 @@ def mem_wal_writer( shard_id: str, *, durable_write: Optional[bool] = None, - sync_indexed_write: Optional[bool] = None, max_wal_buffer_size: Optional[int] = None, max_wal_flush_interval_ms: Optional[int] = None, max_memtable_size: Optional[int] = None, @@ -5196,8 +5189,6 @@ def mem_wal_writer( max_memtable_batches: Optional[int] = None, max_unflushed_memtable_bytes: Optional[int] = None, manifest_scan_batch_size: Optional[int] = None, - async_index_buffer_rows: Optional[int] = None, - async_index_interval_ms: Optional[int] = None, backpressure_log_interval_ms: Optional[int] = None, stats_log_interval_ms: Optional[int] = None, hnsw_params: Optional[Dict[str, Dict[str, int]]] = None, @@ -5215,8 +5206,6 @@ def mem_wal_writer( ``str(uuid.uuid4())``). durable_write : bool, optional Whether to fsync WAL writes (default: ``True``). - sync_indexed_write : bool, optional - Whether index updates are synchronous (default: ``True``). max_wal_buffer_size : int, optional Maximum WAL buffer size in bytes (default: 10 MB). max_wal_flush_interval_ms : int, optional @@ -5231,10 +5220,6 @@ def mem_wal_writer( Maximum unflushed bytes before backpressure (default: 1 GB). manifest_scan_batch_size : int, optional Batch size for manifest scans (default: 2). - async_index_buffer_rows : int, optional - Buffer rows for async index updates (default: 10 000). - async_index_interval_ms : int, optional - Interval for async index updates in milliseconds (default: 1000). backpressure_log_interval_ms : int, optional Interval for backpressure log messages in milliseconds (default: 30 000). @@ -5282,7 +5267,6 @@ def mem_wal_writer( name: val for name, val in [ ("durable_write", durable_write), - ("sync_indexed_write", sync_indexed_write), ("max_wal_buffer_size", max_wal_buffer_size), ("max_wal_flush_interval_ms", max_wal_flush_interval_ms), ("max_memtable_size", max_memtable_size), @@ -5290,8 +5274,6 @@ def mem_wal_writer( ("max_memtable_batches", max_memtable_batches), ("max_unflushed_memtable_bytes", max_unflushed_memtable_bytes), ("manifest_scan_batch_size", manifest_scan_batch_size), - ("async_index_buffer_rows", async_index_buffer_rows), - ("async_index_interval_ms", async_index_interval_ms), ("backpressure_log_interval_ms", backpressure_log_interval_ms), ("stats_log_interval_ms", stats_log_interval_ms), ("hnsw_params", hnsw_params), diff --git a/python/python/tests/test_mem_wal.py b/python/python/tests/test_mem_wal.py index 2871baad986..32904ae6038 100644 --- a/python/python/tests/test_mem_wal.py +++ b/python/python/tests/test_mem_wal.py @@ -209,7 +209,6 @@ def test_shard_writer_delete_binding_masks_base_row(tmp_path): with ds.mem_wal_writer( shard_id, durable_write=True, - sync_indexed_write=True, max_wal_buffer_size=1, max_wal_flush_interval_ms=10, ) as writer: @@ -355,7 +354,6 @@ def test_shard_writer_e2e_correctness(tmp_path): writer = ds.mem_wal_writer( shard_id, durable_write=True, - sync_indexed_write=True, max_wal_buffer_size=10 * 1024, # 10 KB max_wal_flush_interval_ms=50, max_memtable_size=80, # flush after ~80 rows @@ -405,9 +403,7 @@ def test_shard_writer_e2e_correctness(tmp_path): # === New writer: write and read back via active MemTable scanner === ds2 = lance.dataset(ds_path) shard_id2 = str(uuid.uuid4()) - with ds2.mem_wal_writer( - shard_id2, durable_write=False, sync_indexed_write=True - ) as writer2: + with ds2.mem_wal_writer(shard_id2, durable_write=False) as writer2: verify_batch = _e2e_batch(schema, start_id=10000, num_rows=10) writer2.put(pa.Table.from_batches([verify_batch])) result = writer2.lsm_scanner().to_table() @@ -522,7 +518,6 @@ def test_initialize_mem_wal_writer_config_defaults(tmp_path): # Duration knobs are recorded in milliseconds with a `_ms` suffix. assert defaults["max_wal_flush_interval_ms"] == "250" # Every ShardWriterConfig tunable is recorded once any default is set. - assert "sync_indexed_write" in defaults assert "enable_memtable" in defaults diff --git a/python/src/dataset.rs b/python/src/dataset.rs index a361e0b63a8..54e8b9340bc 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -164,7 +164,6 @@ fn stats_log_interval_from_millis(ms: u64) -> Option { #[allow(clippy::too_many_arguments)] fn writer_config_from_kwargs( durable_write: Option, - sync_indexed_write: Option, max_wal_buffer_size: Option, max_wal_flush_interval_ms: Option, max_memtable_size: Option, @@ -172,8 +171,6 @@ fn writer_config_from_kwargs( max_memtable_batches: Option, max_unflushed_memtable_bytes: Option, manifest_scan_batch_size: Option, - async_index_buffer_rows: Option, - async_index_interval_ms: Option, backpressure_log_interval_ms: Option, stats_log_interval_ms: Option, hnsw_params: Option>>, @@ -186,10 +183,6 @@ fn writer_config_from_kwargs( config = config.with_durable_write(v); any = true; } - if let Some(v) = sync_indexed_write { - config = config.with_sync_indexed_write(v); - any = true; - } if let Some(v) = max_wal_buffer_size { config = config.with_max_wal_buffer_size(v); any = true; @@ -218,14 +211,6 @@ fn writer_config_from_kwargs( config = config.with_manifest_scan_batch_size(v); any = true; } - if let Some(v) = async_index_buffer_rows { - config = config.with_async_index_buffer_rows(v); - any = true; - } - if let Some(v) = async_index_interval_ms { - config = config.with_async_index_interval(Duration::from_millis(v)); - any = true; - } if let Some(v) = backpressure_log_interval_ms { config = config.with_backpressure_log_interval(Duration::from_millis(v)); any = true; @@ -3414,7 +3399,6 @@ impl Dataset { identity_column=None, unsharded=false, durable_write=None, - sync_indexed_write=None, max_wal_buffer_size=None, max_wal_flush_interval_ms=None, max_memtable_size=None, @@ -3422,8 +3406,6 @@ impl Dataset { max_memtable_batches=None, max_unflushed_memtable_bytes=None, manifest_scan_batch_size=None, - async_index_buffer_rows=None, - async_index_interval_ms=None, backpressure_log_interval_ms=None, stats_log_interval_ms=None, hnsw_params=None, @@ -3437,7 +3419,6 @@ impl Dataset { identity_column: Option, unsharded: bool, durable_write: Option, - sync_indexed_write: Option, max_wal_buffer_size: Option, max_wal_flush_interval_ms: Option, max_memtable_size: Option, @@ -3445,8 +3426,6 @@ impl Dataset { max_memtable_batches: Option, max_unflushed_memtable_bytes: Option, manifest_scan_batch_size: Option, - async_index_buffer_rows: Option, - async_index_interval_ms: Option, backpressure_log_interval_ms: Option, stats_log_interval_ms: Option, hnsw_params: Option>>, @@ -3474,7 +3453,6 @@ impl Dataset { let writer_config = writer_config_from_kwargs( durable_write, - sync_indexed_write, max_wal_buffer_size, max_wal_flush_interval_ms, max_memtable_size, @@ -3482,8 +3460,6 @@ impl Dataset { max_memtable_batches, max_unflushed_memtable_bytes, manifest_scan_batch_size, - async_index_buffer_rows, - async_index_interval_ms, backpressure_log_interval_ms, stats_log_interval_ms, hnsw_params, @@ -3565,7 +3541,6 @@ impl Dataset { shard_id, *, durable_write=None, - sync_indexed_write=None, max_wal_buffer_size=None, max_wal_flush_interval_ms=None, max_memtable_size=None, @@ -3573,8 +3548,6 @@ impl Dataset { max_memtable_batches=None, max_unflushed_memtable_bytes=None, manifest_scan_batch_size=None, - async_index_buffer_rows=None, - async_index_interval_ms=None, backpressure_log_interval_ms=None, stats_log_interval_ms=None, hnsw_params=None, @@ -3584,7 +3557,6 @@ impl Dataset { py: Python<'_>, shard_id: String, durable_write: Option, - sync_indexed_write: Option, max_wal_buffer_size: Option, max_wal_flush_interval_ms: Option, max_memtable_size: Option, @@ -3592,8 +3564,6 @@ impl Dataset { max_memtable_batches: Option, max_unflushed_memtable_bytes: Option, manifest_scan_batch_size: Option, - async_index_buffer_rows: Option, - async_index_interval_ms: Option, backpressure_log_interval_ms: Option, stats_log_interval_ms: Option, hnsw_params: Option>>, @@ -3605,7 +3575,6 @@ impl Dataset { let config = writer_config_from_kwargs( durable_write, - sync_indexed_write, max_wal_buffer_size, max_wal_flush_interval_ms, max_memtable_size, @@ -3613,8 +3582,6 @@ impl Dataset { max_memtable_batches, max_unflushed_memtable_bytes, manifest_scan_batch_size, - async_index_buffer_rows, - async_index_interval_ms, backpressure_log_interval_ms, stats_log_interval_ms, hnsw_params, diff --git a/python/src/mem_wal.rs b/python/src/mem_wal.rs index e415acfc38d..bbac185e625 100644 --- a/python/src/mem_wal.rs +++ b/python/src/mem_wal.rs @@ -949,10 +949,8 @@ fn memtable_stats_to_pydict(py: Python<'_>, stats: &MemTableStats) -> PyResult

MemTableStats { estimated_size: 0, generation: stats_before_close.generation.saturating_add(1), max_buffered_batch_position: None, - max_flushed_batch_position: None, + durable_batch_count: 0, + global_offset: 0, pending_wal_start_batch_position: None, pending_wal_end_batch_position: None, pending_wal_batch_count: 0, diff --git a/rust/lance/benches/mem_wal/fts/mem_wal_fineweb_fts.rs b/rust/lance/benches/mem_wal/fts/mem_wal_fineweb_fts.rs index 63f9368c3c5..188e2c73ce2 100644 --- a/rust/lance/benches/mem_wal/fts/mem_wal_fineweb_fts.rs +++ b/rust/lance/benches/mem_wal/fts/mem_wal_fineweb_fts.rs @@ -111,11 +111,6 @@ impl Mode { fn durable_write(self) -> bool { matches!(self, Self::SyncNoIndex | Self::SyncIndexed) } - - /// Index update happens inline in `put` (only meaningful when indexed). - fn sync_indexed_write(self) -> bool { - matches!(self, Self::SyncIndexed) - } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -428,7 +423,6 @@ fn shard_writer_config(args: &Args, shard_id: Uuid, disable_auto_flush: bool) -> }; let mut config = ShardWriterConfig::new(shard_id) .with_durable_write(args.mode.durable_write()) - .with_sync_indexed_write(args.mode.sync_indexed_write()) .with_max_memtable_size(args.max_memtable_size) .with_max_unflushed_memtable_bytes(args.max_unflushed_memtable_bytes) .with_max_memtable_rows(max_rows) @@ -638,7 +632,7 @@ async fn run_read(args: &Args, uri: &str, corpus: &[String]) -> Result Result Result { wal_persist_retry_base_delay: std::time::Duration::from_millis(50), shard_spec_id: 0, durable_write: false, - sync_indexed_write: false, max_memtable_size: args.max_memtable_rows * row_bytes * 100, max_memtable_rows: args.max_memtable_rows, max_memtable_batches: batches_per_gen, diff --git a/rust/lance/benches/mem_wal/kv/mem_wal_kv_point_lookup.rs b/rust/lance/benches/mem_wal/kv/mem_wal_kv_point_lookup.rs index 757c2539642..e5c6cf234e1 100644 --- a/rust/lance/benches/mem_wal/kv/mem_wal_kv_point_lookup.rs +++ b/rust/lance/benches/mem_wal/kv/mem_wal_kv_point_lookup.rs @@ -792,7 +792,6 @@ async fn run_lance( wal_persist_retry_base_delay: std::time::Duration::from_millis(50), shard_spec_id: 0, durable_write: true, - sync_indexed_write: true, max_memtable_size: big, max_memtable_rows: args.rows * 4 + 1_000_000, max_memtable_batches: args.rows / args.batch_rows + 1_000_000, diff --git a/rust/lance/benches/mem_wal/point_lookup/mem_wal_point_lookup_bench.rs b/rust/lance/benches/mem_wal/point_lookup/mem_wal_point_lookup_bench.rs index cb9e6413ac1..d723cfde9c8 100644 --- a/rust/lance/benches/mem_wal/point_lookup/mem_wal_point_lookup_bench.rs +++ b/rust/lance/benches/mem_wal/point_lookup/mem_wal_point_lookup_bench.rs @@ -338,7 +338,6 @@ async fn run_lookup(args: &Args) -> Result { wal_persist_retry_base_delay: std::time::Duration::from_millis(50), shard_spec_id: 0, durable_write: false, - sync_indexed_write: false, max_memtable_size: max_memtable_rows * 200, max_memtable_rows, max_wal_flush_interval: Some(Duration::from_secs(60)), diff --git a/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs b/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs index 7b9f35b73c9..a4968161647 100644 --- a/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs +++ b/rust/lance/benches/mem_wal/vector/hnsw/mem_wal_recall_hnsw.rs @@ -389,7 +389,6 @@ async fn run_checkpoint( wal_persist_retry_base_delay: std::time::Duration::from_millis(50), shard_spec_id: 0, durable_write: false, - sync_indexed_write: true, max_memtable_size: cp.saturating_mul(row_size_estimate).saturating_mul(4), max_memtable_rows: cp.saturating_mul(2), max_memtable_batches: total_batches_max.saturating_mul(2).max(8_000), @@ -416,11 +415,11 @@ async fn run_checkpoint( writer.put(vec![batch]).await?; } - let target_batch_pos = total_batches.saturating_sub(1); + let target_indexed_count = total_batches; let mut spins = 0u64; loop { let active = writer.active_memtable_ref().await?; - if active.index_store.max_visible_batch_position() >= target_batch_pos { + if active.index_store.indexed_count() >= target_indexed_count { break; } drop(active); diff --git a/rust/lance/benches/mem_wal/vector/mem_wal_index_micro.rs b/rust/lance/benches/mem_wal/vector/mem_wal_index_micro.rs index 8bff64dfcca..c6fb42834a3 100644 --- a/rust/lance/benches/mem_wal/vector/mem_wal_index_micro.rs +++ b/rust/lance/benches/mem_wal/vector/mem_wal_index_micro.rs @@ -10,7 +10,9 @@ //! index maintained, opens `mem_wal_writer` with a `ShardWriterConfig` //! sized to hold the largest checkpoint without flushing, and times //! `writer.put(batch)` calls. Index updates therefore go through the -//! parallel `IndexStore::insert_batches_parallel` path. At each +//! `IndexStore::insert_batches` path, which indexes inline at or below +//! `PARALLEL_INDEX_MIN_ROWS` rows and spawns a thread per index above it — +//! so the checkpoint sizes here straddle that crossover. At each //! checkpoint, queries are issued against `active_memtable_ref()` via //! `MemTableScanner::nearest`. //! @@ -212,7 +214,6 @@ async fn main() -> lance_core::Result<()> { wal_persist_retry_base_delay: std::time::Duration::from_millis(50), shard_spec_id: 0, durable_write, - sync_indexed_write: true, max_memtable_size: max_rows.saturating_mul(row_size_estimate).saturating_mul(4), max_memtable_rows: max_rows.saturating_mul(2), max_memtable_batches: total_batches_max.saturating_mul(2).max(8_000), @@ -253,7 +254,7 @@ async fn main() -> lance_core::Result<()> { while next_cp_idx < checkpoints.len() && total_inserted >= checkpoints[next_cp_idx] { let cp = checkpoints[next_cp_idx]; - let target_batch_pos = (cp / batch_size).saturating_sub(1); + let target_indexed_count = cp / batch_size; // The WAL flush handler only updates the index watermark when a // flush is triggered, and the time-based trigger inside the // writer runs only when `put()` is called. After the final put @@ -265,7 +266,7 @@ async fn main() -> lance_core::Result<()> { let mut spins = 0u64; loop { let active = writer.active_memtable_ref().await?; - if active.index_store.max_visible_batch_position() >= target_batch_pos { + if active.index_store.indexed_count() >= target_indexed_count { break; } drop(active); @@ -375,12 +376,11 @@ async fn measure_flush( memtable.set_indexes(registry); let total_batches = cp.div_ceil(batch_size); - for (wal_pos, i) in (0_u64..).zip(0..total_batches) { + for i in 0..total_batches { let start = (i * batch_size) as i64; let rows = batch_size.min(cp - i * batch_size); let batch = make_batch(start, rows, dim); - let frag_id = memtable.insert(batch).await?; - memtable.mark_wal_flushed(&[frag_id], wal_pos + 1, &[i]); + let _frag_id = memtable.insert(batch).await?; } let temp_dir = @@ -406,11 +406,19 @@ async fn measure_flush( let flusher = MemTableFlusher::new(store, base_path, uri, shard_id, manifest_store); // total_batches WAL entries were stamped at positions 1..=total_batches - // by the mark_wal_flushed loop above (1-based positions). + // by the loop above (1-based positions). let covered_wal_entry_position = total_batches as u64; + // Every batch was appended above, so the writer-global durable count is all of them. + let durable = total_batches; let t = Instant::now(); let _result = flusher - .flush_with_indexes(&memtable, epoch, index_configs, covered_wal_entry_position) + .flush_with_indexes( + &memtable, + epoch, + index_configs, + covered_wal_entry_position, + durable, + ) .await?; let elapsed = t.elapsed(); diff --git a/rust/lance/benches/mem_wal/vector/mem_wal_vector_bench.rs b/rust/lance/benches/mem_wal/vector/mem_wal_vector_bench.rs index 632bd37626c..a8fe8e81f29 100644 --- a/rust/lance/benches/mem_wal/vector/mem_wal_vector_bench.rs +++ b/rust/lance/benches/mem_wal/vector/mem_wal_vector_bench.rs @@ -520,7 +520,6 @@ async fn run_search(args: &Args) -> Result { wal_persist_retry_base_delay: std::time::Duration::from_millis(50), shard_spec_id: 0, durable_write: false, - sync_indexed_write: false, max_memtable_size: args.max_memtable_rows * row_bytes * 2, max_memtable_rows: args.max_memtable_rows, max_unflushed_memtable_bytes: args.max_memtable_rows * row_bytes * 6, diff --git a/rust/lance/benches/mem_wal/write/mem_wal_replay.rs b/rust/lance/benches/mem_wal/write/mem_wal_replay.rs index 14ec406e5e5..6fda0748006 100644 --- a/rust/lance/benches/mem_wal/write/mem_wal_replay.rs +++ b/rust/lance/benches/mem_wal/write/mem_wal_replay.rs @@ -179,7 +179,6 @@ async fn populate_shard_wal( let dataset = Dataset::open(dataset_uri).await.unwrap(); let mut config = ShardWriterConfig::new(shard_id); config.durable_write = true; - config.sync_indexed_write = false; let writer = dataset.mem_wal_writer(shard_id, config).await.unwrap(); for i in 0..num_entries { diff --git a/rust/lance/benches/mem_wal/write/mem_wal_shard_writer_backpressure.rs b/rust/lance/benches/mem_wal/write/mem_wal_shard_writer_backpressure.rs index f4e4fb47395..79298daa872 100644 --- a/rust/lance/benches/mem_wal/write/mem_wal_shard_writer_backpressure.rs +++ b/rust/lance/benches/mem_wal/write/mem_wal_shard_writer_backpressure.rs @@ -99,10 +99,6 @@ impl Mode { fn durable_write(self) -> bool { matches!(self, Self::SyncNoIndex | Self::SyncIndexed) } - - fn sync_indexed_write(self) -> bool { - matches!(self, Self::SyncIndexed) - } } /// Which index the MemTable maintains in the indexed (`*_idx`) modes. @@ -192,7 +188,6 @@ struct Args { max_memtable_batches: Option, max_wal_buffer_size: usize, max_wal_flush_interval_ms: u64, - async_index_buffer_rows: usize, sample_interval_ms: u64, target_rows_per_sec: Option, num_partitions: usize, @@ -223,7 +218,6 @@ impl Default for Args { max_memtable_batches: None, max_wal_buffer_size: 10 * 1024 * 1024, max_wal_flush_interval_ms: 100, - async_index_buffer_rows: 10_000, sample_interval_ms: 500, target_rows_per_sec: None, num_partitions: 1, @@ -343,11 +337,9 @@ async fn run(args: Args) -> Result<()> { let shard_id = Uuid::new_v4(); let mut config = ShardWriterConfig::new(shard_id) .with_durable_write(args.mode.durable_write()) - .with_sync_indexed_write(args.mode.sync_indexed_write()) .with_max_memtable_size(args.max_memtable_size) .with_max_unflushed_memtable_bytes(args.max_unflushed_memtable_bytes) .with_max_wal_buffer_size(args.max_wal_buffer_size) - .with_async_index_buffer_rows(args.async_index_buffer_rows) .with_max_memtable_rows(memtable_limits.rows) .with_max_memtable_batches(memtable_limits.batches); if args.max_wal_flush_interval_ms == 0 { @@ -544,7 +536,6 @@ async fn run(args: Args) -> Result<()> { "max_memtable_batches": memtable_limits.batches, "max_wal_buffer_size": args.max_wal_buffer_size, "max_wal_flush_interval_ms": args.max_wal_flush_interval_ms, - "async_index_buffer_rows": args.async_index_buffer_rows, "sample_interval_ms": args.sample_interval_ms, "skip_close": args.skip_close, "setup_seconds": setup_s, @@ -650,7 +641,7 @@ fn push_sample( "active_memtable_bytes": memtable.as_ref().map(|stats| stats.estimated_size), "active_memtable_generation": memtable.as_ref().map(|stats| stats.generation), "active_memtable_max_buffered_batch_position": memtable.as_ref().and_then(|stats| stats.max_buffered_batch_position), - "active_memtable_max_flushed_batch_position": memtable.as_ref().and_then(|stats| stats.max_flushed_batch_position), + "active_memtable_durable_batch_count": memtable.as_ref().map(|stats| stats.durable_batch_count), "wal_queue_pending_batches": memtable.as_ref().map(|stats| stats.pending_wal_batch_count), "wal_queue_pending_rows": memtable.as_ref().map(|stats| stats.pending_wal_row_count), "wal_queue_pending_bytes": memtable.as_ref().map(|stats| stats.pending_wal_estimated_bytes), @@ -667,7 +658,7 @@ fn memtable_stats_json(memtable: Option<&MemTableStats>) -> serde_json::Value { "estimated_size": stats.estimated_size, "generation": stats.generation, "max_buffered_batch_position": stats.max_buffered_batch_position, - "max_flushed_batch_position": stats.max_flushed_batch_position, + "durable_batch_count": stats.durable_batch_count, "wal_queue_pending_start_batch_position": stats.pending_wal_start_batch_position, "wal_queue_pending_end_batch_position": stats.pending_wal_end_batch_position, "wal_queue_pending_batches": stats.pending_wal_batch_count, @@ -980,7 +971,6 @@ fn parse_args() -> Result { "--max-wal-flush-interval-ms" => { args.max_wal_flush_interval_ms = parse(&flag, &value)?; } - "--async-index-buffer-rows" => args.async_index_buffer_rows = parse(&flag, &value)?, "--sample-interval-ms" => args.sample_interval_ms = parse(&flag, &value)?, "--target-rows-per-sec" => args.target_rows_per_sec = Some(parse(&flag, &value)?), "--num-partitions" => args.num_partitions = parse(&flag, &value)?, diff --git a/rust/lance/benches/mem_wal/write/mem_wal_write.rs b/rust/lance/benches/mem_wal/write/mem_wal_write.rs index db28c5b33ab..b5018fe09f1 100644 --- a/rust/lance/benches/mem_wal/write/mem_wal_write.rs +++ b/rust/lance/benches/mem_wal/write/mem_wal_write.rs @@ -30,7 +30,6 @@ //! - `BATCH_SIZE`: Number of rows per write batch (default: 20) //! - `NUM_BATCHES`: Total number of batches to write (default: 1000) //! - `DURABLE_WRITE`: yes/no/both (default: no) - whether writes wait for WAL flush -//! - `INDEXED_WRITE`: yes/no/both (default: no) - whether writes update indexes synchronously //! - `MAX_WAL_BUFFER_SIZE`: WAL buffer size in bytes (default: 1MB from ShardWriterConfig) //! - `MAX_FLUSH_INTERVAL_MS`: WAL flush interval in milliseconds, 0 to disable (default: 1000ms) //! - `MAX_MEMTABLE_SIZE`: MemTable size threshold in bytes (default: 64MB from ShardWriterConfig) @@ -44,8 +43,7 @@ //! `no` uses WAL-only mode (no MemTable, no indexes, no Lance flushes; pure WAL throughput). //! `both` runs each combination twice, once per mode, side-by-side. //! When `no` or `both`, the WAL-only branch always runs with -//! `MEMWAL_MAINTAINED_INDEXES=none` and skips `INDEXED_WRITE=yes` -//! (sync-indexed writes require a MemTable). +//! `MEMWAL_MAINTAINED_INDEXES=none`. //! - `SAMPLE_SIZE`: Number of benchmark iterations (default: 10, minimum: 10) #![allow(clippy::print_stdout, clippy::print_stderr)] @@ -119,11 +117,6 @@ fn get_durable_write_options() -> Vec { parse_yes_no_both("DURABLE_WRITE", "no") } -/// Get indexed write settings from environment. -fn get_indexed_write_options() -> Vec { - parse_yes_no_both("INDEXED_WRITE", "no") -} - /// Get enable_memtable settings from environment. Default `yes` keeps /// existing benchmark behavior; `no` runs WAL-only mode; `both` runs both /// modes side-by-side for comparison. @@ -442,31 +435,26 @@ fn build_label( num_batches: usize, batch_size: usize, durable: bool, - indexed: bool, enable_memtable: bool, storage: &str, ) -> String { let durable_str = if durable { "durable" } else { "nondurable" }; - // sync_indexed_write controls sync vs async index updates - let indexed_str = if indexed { "sync_idx" } else { "async_idx" }; let mode_str = if enable_memtable { "memtable" } else { "wal_only" }; format!( - "{}x{} {} {} {} ({})", - num_batches, batch_size, mode_str, durable_str, indexed_str, storage + "{}x{} {} {} ({})", + num_batches, batch_size, mode_str, durable_str, storage ) } /// Build dataset name prefix from config options. -fn build_name_prefix(durable: bool, indexed: bool, enable_memtable: bool) -> String { +fn build_name_prefix(durable: bool, enable_memtable: bool) -> String { let d = if durable { "d" } else { "nd" }; - // sync_indexed_write: sync (si) vs async (ai) - let i = if indexed { "si" } else { "ai" }; let m = if enable_memtable { "mt" } else { "wo" }; - format!("{}_{}_{}", m, d, i) + format!("{}_{}", m, d) } /// Benchmark Lance MemWAL write throughput. @@ -490,7 +478,6 @@ fn bench_lance_memwal_write(c: &mut Criterion) { let maintained_indexes = get_maintained_indexes(); let durable_options = get_durable_write_options(); - let indexed_options = get_indexed_write_options(); let enable_memtable_options = get_enable_memtable_options(); let max_wal_buffer_size = get_max_wal_buffer_size(); let max_flush_interval = get_max_flush_interval(); @@ -550,70 +537,56 @@ fn bench_lance_memwal_write(c: &mut Criterion) { // Generate benchmarks for all combinations for &enable_memtable in &enable_memtable_options { for &durable in &durable_options { - for &indexed in &indexed_options { - if !enable_memtable && indexed { - eprintln!( - "Skipping wal_only + sync_idx (sync_indexed_write requires a MemTable)" - ); - continue; - } - - let label = build_label( - num_batches, - batch_size, - durable, - indexed, - enable_memtable, - storage_label, - ); - let name_prefix = build_name_prefix(durable, indexed, enable_memtable); - - // WAL-only mode never uses indexes; force the dataset - // setup to skip the MemWAL index list. - let effective_indexes: Vec = if enable_memtable { - maintained_indexes.clone() - } else { - Vec::new() - }; - - // Create dataset ONCE before benchmark iterations - // Each iteration will use a different shard on the same dataset - let dataset = rt.block_on(create_dataset( - &schema, - &name_prefix, - vector_dim, - &effective_indexes, - &dataset_prefix, - )); - let dataset_uri = dataset.uri().to_string(); - - // Pre-generate all batches before timing (outside iter_custom) - let batches: Arc> = Arc::new( - (0..num_batches) - .map(|i| { - create_test_batch( - &schema, - (i * batch_size) as i64, - batch_size, - vector_dim, - ) - }) - .collect(), - ); - - println!("Running: {}", label); - - // Track if we've printed stats (only print once across all samples) - let stats_printed = Arc::new(AtomicBool::new(false)); - - group.bench_with_input( - BenchmarkId::new("Lance MemWAL", &label), - &(batch_size, num_batches, durable, indexed, row_size_bytes), - |b, &(_batch_size, _num_batches, durable, indexed, row_size_bytes)| { - let dataset_uri = dataset_uri.clone(); - let batches = batches.clone(); - let stats_printed = stats_printed.clone(); - b.to_async(&rt).iter_custom(|iters| { + let label = build_label( + num_batches, + batch_size, + durable, + enable_memtable, + storage_label, + ); + let name_prefix = build_name_prefix(durable, enable_memtable); + + // WAL-only mode never uses indexes; force the dataset + // setup to skip the MemWAL index list. + let effective_indexes: Vec = if enable_memtable { + maintained_indexes.clone() + } else { + Vec::new() + }; + + // Create dataset ONCE before benchmark iterations + // Each iteration will use a different shard on the same dataset + let dataset = rt.block_on(create_dataset( + &schema, + &name_prefix, + vector_dim, + &effective_indexes, + &dataset_prefix, + )); + let dataset_uri = dataset.uri().to_string(); + + // Pre-generate all batches before timing (outside iter_custom) + let batches: Arc> = Arc::new( + (0..num_batches) + .map(|i| { + create_test_batch(&schema, (i * batch_size) as i64, batch_size, vector_dim) + }) + .collect(), + ); + + println!("Running: {}", label); + + // Track if we've printed stats (only print once across all samples) + let stats_printed = Arc::new(AtomicBool::new(false)); + + group.bench_with_input( + BenchmarkId::new("Lance MemWAL", &label), + &(batch_size, num_batches, durable, row_size_bytes), + |b, &(_batch_size, _num_batches, durable, row_size_bytes)| { + let dataset_uri = dataset_uri.clone(); + let batches = batches.clone(); + let stats_printed = stats_printed.clone(); + b.to_async(&rt).iter_custom(|iters| { let dataset_uri = dataset_uri.clone(); let batches = batches.clone(); let stats_printed = stats_printed.clone(); @@ -631,10 +604,10 @@ fn bench_lance_memwal_write(c: &mut Criterion) { shard_id, shard_spec_id: 0, max_wal_persist_retries: 3, - wal_persist_retry_base_delay: - std::time::Duration::from_millis(50), + wal_persist_retry_base_delay: std::time::Duration::from_millis( + 50, + ), durable_write: durable, - sync_indexed_write: indexed, max_wal_buffer_size: max_wal_buffer_size .unwrap_or(default_config.max_wal_buffer_size), max_wal_flush_interval: max_flush_interval @@ -643,8 +616,6 @@ fn bench_lance_memwal_write(c: &mut Criterion) { .unwrap_or(default_config.max_memtable_size), max_memtable_rows: default_config.max_memtable_rows, max_memtable_batches: default_config.max_memtable_batches, - async_index_buffer_rows: default_config.async_index_buffer_rows, - async_index_interval: default_config.async_index_interval, manifest_scan_batch_size: default_config .manifest_scan_batch_size, max_unflushed_memtable_bytes: default_config @@ -706,9 +677,8 @@ fn bench_lance_memwal_write(c: &mut Criterion) { total_duration } }) - }, - ); - } + }, + ); } } diff --git a/rust/lance/src/dataset/mem_wal/api.rs b/rust/lance/src/dataset/mem_wal/api.rs index 1bd95092e7c..4351ecc95b7 100644 --- a/rust/lance/src/dataset/mem_wal/api.rs +++ b/rust/lance/src/dataset/mem_wal/api.rs @@ -392,10 +392,6 @@ fn writer_config_to_defaults(config: &ShardWriterConfig) -> HashMap HashMap Result<()> { + // Not caller input: the graph was sized below what the memtable holds. + // See the matching note in `storage.rs::append_batch`. if needed_len > self.nodes.len() { - return Err(Error::invalid_input(format!( - "graph capacity {} exhausted: need {needed_len}", + return Err(Error::internal(format!( + "HNSW graph capacity {} exhausted: need {needed_len}; \ + the graph is sized below the memtable's row capacity", self.nodes.len() ))); } diff --git a/rust/lance/src/dataset/mem_wal/hnsw/storage.rs b/rust/lance/src/dataset/mem_wal/hnsw/storage.rs index bbeb57a5fe2..543c1172500 100644 --- a/rust/lance/src/dataset/mem_wal/hnsw/storage.rs +++ b/rust/lance/src/dataset/mem_wal/hnsw/storage.rs @@ -269,24 +269,31 @@ impl ArrowFixedSizeListVectorStore { ))); }; + // Exhaustion is not a caller-input problem — the batch is well-formed and + // has already passed the memtable's schema gate. It means the HNSW store + // was sized below what the memtable will hold before it flushes, which is + // a shard-construction bug. `invalid_input` mislabelled it as the writer's + // fault and hid that. let start = self.committed_len.load(Ordering::Relaxed); let end = start.checked_add(num_rows).ok_or_else(|| { - Error::invalid_input(format!( + Error::internal(format!( "vector count overflow: start={}, batch_len={}", start, num_rows )) })?; if end > self.capacity { - return Err(Error::invalid_input(format!( - "capacity {} exhausted: inserting rows [{}..{})", + return Err(Error::internal(format!( + "HNSW vector store capacity {} exhausted: inserting rows [{}..{}); \ + the store is sized below the memtable's row capacity", self.capacity, start, end ))); } let batch_idx = self.committed_batches.load(Ordering::Relaxed); if batch_idx >= self.max_batches { - return Err(Error::invalid_input(format!( - "max_batches {} exhausted", + return Err(Error::internal(format!( + "HNSW vector store max_batches {} exhausted; \ + the store is sized below the memtable's batch capacity", self.max_batches ))); } diff --git a/rust/lance/src/dataset/mem_wal/index.rs b/rust/lance/src/dataset/mem_wal/index.rs index 8077f06149e..fc2ae180c81 100644 --- a/rust/lance/src/dataset/mem_wal/index.rs +++ b/rust/lance/src/dataset/mem_wal/index.rs @@ -27,7 +27,9 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use datafusion::common::ScalarValue; use super::memtable::batch_store::StoredBatch; +use super::wal::WriterCursors; use arrow_array::RecordBatch; +use arrow_schema::{DataType, Schema as ArrowSchema}; use lance_core::datatypes::Schema as LanceSchema; use lance_core::{Error, Result}; use lance_index::pbold; @@ -58,6 +60,19 @@ use pk_key::encode_pk_batch; /// [`BTreeMemIndex`]'s byte backend indexes it directly. const PK_KEY_COLUMN: &str = "__pk_key__"; +/// Row count at or below which [`IndexStore::insert_batches`] indexes inline +/// rather than spawning a thread per index. +/// +/// The spawn is one OS thread *per index* — tens of microseconds each, and a table can +/// carry several BTrees alongside its HNSW and FTS — so for a small batch it costs more +/// than the indexing it parallelizes. Small batches are not the exceptional case: a +/// durable put triggers a WAL flush covering only the batch it just inserted, so this +/// path is routinely called with a single short batch. +/// +/// The crossover depends on per-row HNSW cost, which varies with dimension and +/// `ef_construction`; tune against `benches/mem_wal/vector/mem_wal_index_micro.rs`. +const PARALLEL_INDEX_MIN_ROWS: usize = 64; + /// The memtable's primary-key index, used to answer "newest visible version of /// this key" for dedup. Single-column PKs reuse the column's compact typed /// [`BTreeMemIndex`] (no second copy); composite PKs key a `BTreeMemIndex` on @@ -80,6 +95,169 @@ enum PkIndex { // Index Store // ============================================================================ +/// Validate every configured in-memory index, and the composite primary key, +/// against the shard schema. Call once at shard open, before any write can land. +/// +/// This is what makes poison-and-replay *terminating*. An index insert that +/// fails deterministically on a row that is already WAL-durable cannot be +/// recovered from: the writer poisons, the operator reopens, replay re-reads the +/// same WAL rows, the same insert fails again, and `open()` propagates it — a +/// shard that never comes back. Every such failure is an index *config* +/// disagreeing with the schema, never a property of the data, so one pass here +/// closes the whole class before a single row is accepted. +/// +/// The data-dependent errors inside the index layer are already unreachable +/// through `put`: `MemTable::insert_batches_only` does a full `Arc` +/// equality check, so a batch that would trip one is rejected before it reaches +/// the batch store, let alone the WAL. +/// +/// It also rejects a config whose `field_id` names a different column than its +/// `column`. Index *selection* keys off `field_id` — a single-column PK reuses +/// the BTree whose `field_id` matches its key — so a config resolved only by name +/// could be bound under the wrong identity, serving stale reads and flushing the +/// wrong column into the durable PK sidecar. `lance_schema` supplies the +/// authoritative name→id mapping. +pub fn validate_index_configs( + configs: &[MemIndexConfig], + schema: &ArrowSchema, + lance_schema: &LanceSchema, + pk_columns: &[String], +) -> Result<()> { + for config in configs { + let column = config.column(); + let field = schema.field_with_name(column).map_err(|_| { + Error::invalid_input(format!( + "index '{}' is configured on column '{}', which is not in the shard schema; \ + available columns: [{}]", + config.name(), + column, + schema + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>() + .join(", ") + )) + })?; + + match config { + // BTree falls back to per-row `ScalarValue` extraction, so it + // accepts any column type the schema can hold. Existence is the + // only precondition. + MemIndexConfig::BTree(_) => {} + MemIndexConfig::Fts(_) => { + if !matches!( + field.data_type(), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) { + return Err(Error::invalid_input(format!( + "FTS index '{}' requires a Utf8, LargeUtf8, or Utf8View column; \ + column '{}' is {:?}", + config.name(), + column, + field.data_type() + ))); + } + } + MemIndexConfig::Hnsw(_) => match field.data_type() { + DataType::FixedSizeList(item, dim) => { + if item.data_type() != &DataType::Float32 { + return Err(Error::invalid_input(format!( + "HNSW index '{}' requires a FixedSizeList column; \ + column '{}' has item type {:?}", + config.name(), + column, + item.data_type() + ))); + } + // `HnswMemIndex.dim` is a placeholder until the first batch + // pins it (`hnsw.rs`), so a zero-width vector would only + // surface at insert time — i.e. on already-durable data. + if *dim <= 0 { + return Err(Error::invalid_input(format!( + "HNSW index '{}' requires a vector dimension > 0; column '{}' has \ + dimension {dim}", + config.name(), + column, + ))); + } + } + other => { + return Err(Error::invalid_input(format!( + "HNSW index '{}' requires a FixedSizeList column; \ + column '{}' is {:?}", + config.name(), + column, + other + ))); + } + }, + } + + // The column resolves, but index selection keys off `field_id`, not name. + // A config whose `field_id` identifies a *different* column would be bound + // under the wrong identity (e.g. reused as the single-column PK index), so + // reject any `field_id` that does not name the resolved column. + let resolved_field_id = lance_schema + .field(column) + .expect("column resolved in the Arrow schema is present in the Lance schema") + .id; + if resolved_field_id != config.field_id() { + return Err(Error::invalid_input(format!( + "index '{}' is configured with field_id {} but its column '{}' has field_id {} \ + in the shard schema", + config.name(), + config.field_id(), + column, + resolved_field_id, + ))); + } + } + + // Every PK column must exist in the schema. A single-column PK aliases a + // BTree entry (any type); only a *composite* PK builds an order-preserving + // encoded key, and only some types encode. + for column in pk_columns { + let field = schema.field_with_name(column).map_err(|_| { + Error::invalid_input(format!( + "primary-key column '{column}' is not in the shard schema" + )) + })?; + if pk_columns.len() > 1 && !is_encodable_pk_type(field.data_type()) { + return Err(Error::invalid_input(format!( + "composite primary-key column '{column}' has type {:?}, which has no \ + order-preserving key encoding", + field.data_type() + ))); + } + } + + Ok(()) +} + +/// Types `pk_key::encode_value` can encode into an order-preserving composite key. +fn is_encodable_pk_type(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::Date32 + | DataType::Date64 + | DataType::Boolean + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Binary + | DataType::LargeBinary + | DataType::FixedSizeBinary(_) + ) +} + /// Configuration for an index in MemWAL. /// /// Each variant contains all the configuration needed for that index type. @@ -238,9 +416,9 @@ impl MemIndexConfig { /// Indexes are keyed by index name. Each index stores its field_id for /// stable column-to-index resolution (column name → field_id → index). /// -/// The store also carries the MemTable's `max_visible_batch_position` -/// watermark — the highest batch position that is durable in the WAL and -/// therefore safe for scanners to read. Scanners snapshot this at plan +/// The store also carries the MemTable's two cursors: `indexed_count` (what the +/// index layer has ingested) and `visible_count` (what is indexed *and* durable, +/// and therefore safe for scanners to read). Scanners snapshot the latter at plan /// construction time so every plan keys on a stable MVCC cursor. pub struct IndexStore { /// BTree indexes keyed by index name. `Arc` so the primary-key BTrees can be @@ -254,10 +432,23 @@ pub struct IndexStore { /// primary key. Queried via [`Self::pk_newest_visible`] (see /// [`Self::enable_pk_index`]). pk_index: Option, - /// Maximum batch position that is durable in the WAL and therefore - /// visible to scanners. Advanced unconditionally after a WAL append - /// succeeds; not gated on whether any indexes are configured. - max_visible_batch_position: AtomicUsize, + /// How many batches of this memtable have been fully indexed. An exclusive + /// count: 0 means none. + /// + /// This has only ever been an *indexed* cursor — it is advanced at the end of + /// `insert_batches`, once every index insert for the batch has completed, and + /// never before. It was named `max_visible_batch_position` and treated as a + /// visibility cursor by five read sites, which is how rows became readable + /// before they were durable. Publishing is a separate step, and it is the + /// writer's to make — see `visible_count`. + indexed_count: AtomicUsize, + + /// The writer's cursors, and this memtable's coordinate within them. `None` + /// for a bare `IndexStore` (tests, benches), where visibility is just the + /// indexed prefix. + /// + /// Visibility is **derived, never stored**: see `visible_count`. + durability: Option<(Arc, usize)>, /// Conservative flag set once this memtable has observed any primary-key /// rewrite while maintaining a search index. Search planners can push top-k /// into HNSW/FTS for append-only PK data, but must switch to @@ -272,7 +463,8 @@ impl Default for IndexStore { hnsw_indexes: HashMap::new(), fts_indexes: HashMap::new(), pk_index: None, - max_visible_batch_position: AtomicUsize::new(0), + indexed_count: AtomicUsize::new(0), + durability: None, pk_has_overrides: AtomicBool::new(false), } } @@ -300,10 +492,7 @@ impl std::fmt::Debug for IndexStore { } }, ) - .field( - "max_visible_batch_position", - &self.max_visible_batch_position.load(Ordering::Acquire), - ) + .field("indexed_count", &self.indexed_count.load(Ordering::Acquire)) .field( "pk_has_overrides", &self.pk_has_overrides.load(Ordering::Acquire), @@ -695,26 +884,25 @@ impl IndexStore { let had_existing = self.insert_composite_pk(batch, row_offset, track_pk_overrides)?; self.mark_pk_overrides_if_needed(had_existing); - // Update global watermark after all indexes have been updated + // Update the indexed prefix after every index has been updated. if let Some(bp) = batch_position { - self.advance_max_visible_batch_position(bp); + self.advance_indexed_count(bp + 1); } Ok(()) } - /// Advance the visibility watermark to at least `batch_pos`. + /// Advance the indexed prefix to at least `count` batches. /// - /// The watermark only ever moves forward (idempotent max). The vector - /// planner relies on the insert paths setting `pk_has_overrides` before - /// calling this method, so any snapshot that can see a PK rewrite also - /// observes `pk_has_overrides == true`. - pub(crate) fn advance_max_visible_batch_position(&self, batch_pos: usize) { - let mut current = self.max_visible_batch_position.load(Ordering::Acquire); - while batch_pos > current { - match self.max_visible_batch_position.compare_exchange_weak( + /// Only ever moves forward (idempotent max). The vector planner relies on the + /// insert paths setting `pk_has_overrides` before this is called, so any + /// snapshot that can see a PK rewrite also observes `pk_has_overrides == true`. + pub(crate) fn advance_indexed_count(&self, count: usize) { + let mut current = self.indexed_count.load(Ordering::Acquire); + while count > current { + match self.indexed_count.compare_exchange_weak( current, - batch_pos, + count, Ordering::Release, Ordering::Acquire, ) { @@ -724,69 +912,16 @@ impl IndexStore { } } - /// Insert multiple batches into all indexes with cross-batch optimization. - #[instrument(name = "idx_insert_batches", level = "debug", skip_all, fields(batch_count = batches.len()))] - pub fn insert_batches(&self, batches: &[StoredBatch]) -> Result<()> { - if batches.is_empty() { - return Ok(()); - } - - let track_pk_overrides = self.should_track_pk_overrides(); - // BTree indexes: iterate batches (no cross-batch optimization benefit) - for index in self.btree_indexes.values() { - let track_this_index = track_pk_overrides && self.is_single_pk_btree(index); - let mut had_existing = false; - for stored in batches { - if track_this_index { - had_existing |= - index.insert_and_report_existing(&stored.data, stored.row_offset)?; - } else { - index.insert(&stored.data, stored.row_offset)?; - } - } - self.mark_pk_overrides_if_needed(had_existing); - } - - // HNSW indexes: use batched insert - for index in self.hnsw_indexes.values() { - index.insert_batches(batches)?; - } - - // FTS indexes: iterate batches (potential future optimization) - for index in self.fts_indexes.values() { - for stored in batches { - index.insert(&stored.data, stored.row_offset)?; - } - } - - // Single-column PK aliases a `btree_indexes` entry (maintained above); - // a composite PK has its own index, maintained here. - let mut had_existing = false; - for stored in batches { - had_existing |= - self.insert_composite_pk(&stored.data, stored.row_offset, track_pk_overrides)?; - } - self.mark_pk_overrides_if_needed(had_existing); - - // Update global watermark to the max batch position - let max_bp = batches.iter().map(|b| b.batch_position).max().unwrap(); - self.advance_max_visible_batch_position(max_bp); - - Ok(()) - } - - /// Insert multiple batches into all indexes in parallel. + /// Insert multiple batches into every index. /// - /// Each individual index runs in its own thread, regardless of type. - /// This maximizes parallelism when multiple indexes are maintained. - /// - /// This is used during WAL flush to parallelize index updates with WAL I/O. - /// Insert batches into all indexes in parallel. + /// Above `PARALLEL_INDEX_MIN_ROWS` rows each index runs on its own thread, which + /// maximizes parallelism when several indexes are maintained. At or below it they run + /// inline on the calling thread: the spawn is one OS thread *per index*, and for a + /// handful of rows that costs more than the indexing itself. /// /// Returns a map of index names to their update durations for performance tracking. - #[allow(clippy::print_stderr)] - #[instrument(name = "idx_insert_batches_parallel", level = "debug", skip_all, fields(batch_count = batches.len()))] - pub fn insert_batches_parallel( + #[instrument(name = "idx_insert_batches", level = "debug", skip_all, fields(batch_count = batches.len()))] + pub fn insert_batches( &self, batches: &[StoredBatch], ) -> Result> { @@ -797,117 +932,129 @@ impl IndexStore { } let track_pk_overrides = self.should_track_pk_overrides(); - // Use std::thread::scope for parallel CPU-bound work - std::thread::scope(|scope| { - // Each handle returns (index_name, index_type, duration, Result) - let mut handles: Vec<( - &str, - &str, - std::thread::ScopedJoinHandle<'_, (std::time::Duration, Result)>, - )> = Vec::new(); - - // Spawn a thread for each BTree index - for (name, index) in &self.btree_indexes { - let track_this_index = track_pk_overrides && self.is_single_pk_btree(index); - let handle = scope.spawn(move || -> (std::time::Duration, Result) { - let start = Instant::now(); - let result = (|| { - let mut had_existing = false; - for stored in batches { - if track_this_index { - had_existing |= index - .insert_and_report_existing(&stored.data, stored.row_offset)?; - } else { - index.insert(&stored.data, stored.row_offset)?; - } - } - Ok(had_existing) - })(); - (start.elapsed(), result) - }); - handles.push((name.as_str(), "btree", handle)); - } - // Spawn a thread for each HNSW index - for (name, index) in &self.hnsw_indexes { - let handle = scope.spawn(move || -> (std::time::Duration, Result) { - let start = Instant::now(); - let result = index.insert_batches(batches).map(|_| false); - (start.elapsed(), result) - }); - handles.push((name.as_str(), "hnsw", handle)); - } + // One task per index, boxed so the inline and the threaded path drive the very + // same closures. Each reports whether it saw an already-present PK. + type IndexTask<'a> = Box Result + Send + Sync + 'a>; + let mut tasks: Vec<(&str, IndexTask<'_>)> = Vec::new(); - // Spawn a thread for each FTS index - for (name, index) in &self.fts_indexes { - let handle = scope.spawn(move || -> (std::time::Duration, Result) { - let start = Instant::now(); - let result = (|| { - for stored in batches { + for (name, index) in &self.btree_indexes { + let track_this_index = track_pk_overrides && self.is_single_pk_btree(index); + tasks.push(( + name.as_str(), + Box::new(move || { + let mut had_existing = false; + for stored in batches { + if track_this_index { + had_existing |= index + .insert_and_report_existing(&stored.data, stored.row_offset)?; + } else { index.insert(&stored.data, stored.row_offset)?; } - Ok(false) - })(); - (start.elapsed(), result) - }); - handles.push((name.as_str(), "fts", handle)); - } - - // Collect results, log timing, and check for errors. Keep the raw - // `Duration` so sub-millisecond timings (the steady-state case for - // BTree updates) are preserved instead of getting truncated to 0. - let mut first_error: Option = None; - let mut timings: Vec<(&str, &str, std::time::Duration)> = Vec::new(); - let mut had_existing_pk = false; - - for (name, idx_type, handle) in handles { - match handle.join() { - Ok((duration, Ok(had_existing))) => { - timings.push((name, idx_type, duration)); - had_existing_pk |= had_existing; - } - Ok((duration, Err(e))) => { - timings.push((name, idx_type, duration)); - if first_error.is_none() { - first_error = Some(e); - } } - Err(_) => { - if first_error.is_none() { - first_error = - Some(Error::internal(format!("Index '{}' thread panicked", name))); - } + Ok(had_existing) + }), + )); + } + + for (name, index) in &self.hnsw_indexes { + tasks.push(( + name.as_str(), + Box::new(move || index.insert_batches(batches).map(|_| false)), + )); + } + + for (name, index) in &self.fts_indexes { + tasks.push(( + name.as_str(), + Box::new(move || { + for stored in batches { + index.insert(&stored.data, stored.row_offset)?; } - } - } + Ok(false) + }), + )); + } - if let Some(e) = first_error { - return Err(e); - } - self.mark_pk_overrides_if_needed(had_existing_pk); - - let duration_map: std::collections::HashMap = timings - .into_iter() - .map(|(name, _idx_type, duration)| (name.to_string(), duration)) - .collect(); - - // Single-column PK aliases a `btree_indexes` entry — its thread above - // already maintained it (and joined). A composite PK has its own - // index; maintain it here before the watermark advances so the - // visible prefix is fully indexed. - let mut had_existing = false; - for stored in batches { - had_existing |= - self.insert_composite_pk(&stored.data, stored.row_offset, track_pk_overrides)?; + // Keep the raw `Duration` so sub-millisecond timings (the steady state for BTree + // updates) survive instead of truncating to 0. + let total_rows: usize = batches.iter().map(|b| b.num_rows).sum(); + let results: Vec<(&str, std::time::Duration, Result)> = + if tasks.len() < 2 || total_rows <= PARALLEL_INDEX_MIN_ROWS { + tasks + .iter() + .map(|(name, task)| { + let start = Instant::now(); + let result = task(); + (*name, start.elapsed(), result) + }) + .collect() + } else { + std::thread::scope(|scope| { + let handles: Vec<_> = tasks + .iter() + .map(|(name, task)| { + let handle = scope.spawn(move || { + let start = Instant::now(); + let result = task(); + (start.elapsed(), result) + }); + (*name, handle) + }) + .collect(); + + handles + .into_iter() + .map(|(name, handle)| match handle.join() { + Ok((duration, result)) => (name, duration, result), + Err(_) => ( + name, + std::time::Duration::ZERO, + Err(Error::internal(format!("Index '{}' thread panicked", name))), + ), + }) + .collect() + }) + }; + + // Every task ran to completion whether or not a peer failed (the threaded path + // joins all handles unconditionally). Keep the first error; there is no rollback, + // so a failure here is terminal for the writer. + let mut first_error: Option = None; + let mut had_existing_pk = false; + let mut duration_map = + std::collections::HashMap::::with_capacity(results.len()); + + for (name, duration, result) in results { + duration_map.insert(name.to_string(), duration); + match result { + Ok(had_existing) => had_existing_pk |= had_existing, + Err(e) if first_error.is_none() => first_error = Some(e), + Err(_) => {} } - self.mark_pk_overrides_if_needed(had_existing); + } + + if let Some(e) = first_error { + return Err(e); + } + self.mark_pk_overrides_if_needed(had_existing_pk); + + // Single-column PK aliases a `btree_indexes` entry — its task above already + // maintained it. A composite PK has its own index; maintain it here before the + // watermark advances so the visible prefix is fully indexed. + let mut had_existing = false; + for stored in batches { + had_existing |= + self.insert_composite_pk(&stored.data, stored.row_offset, track_pk_overrides)?; + } + self.mark_pk_overrides_if_needed(had_existing); - // Update global watermark to the max batch position - let max_bp = batches.iter().map(|b| b.batch_position).max().unwrap(); - self.advance_max_visible_batch_position(max_bp); + // The indexed prefix now covers every batch up to and including the + // highest position in this call, so the count is that position plus one. + let max_bp = batches.iter().map(|b| b.batch_position).max().unwrap(); + self.advance_indexed_count(max_bp + 1); - Ok(duration_map) - }) + Ok(duration_map) } /// Get a BTree index by name. @@ -985,15 +1132,36 @@ impl IndexStore { self.btree_indexes.len() + self.hnsw_indexes.len() + self.fts_indexes.len() } - /// Get the visibility watermark (max batch position safe to read). + /// How many batches of this memtable have been fully indexed (exclusive + /// count; 0 before any batch is indexed). /// - /// Returns the highest batch position whose data is durable in the WAL - /// and therefore visible to scanners. Scanners snapshot this at plan - /// construction time so every plan runs against a stable cursor. + /// This is the *indexed* cursor, not the visibility watermark: it advances + /// once every index insert for a batch completes, regardless of WAL + /// durability. Readers must snapshot [`Self::visible_count`], which derives + /// what is safe to read from this cursor and the writer's durability cursor. + pub fn indexed_count(&self) -> usize { + self.indexed_count.load(Ordering::Acquire) + } + + /// The prefix of this memtable that readers may see. Snapshot this, never + /// `indexed_count`. /// - /// Returns 0 before any WAL flush has advanced the watermark. - pub fn max_visible_batch_position(&self) -> usize { - self.max_visible_batch_position.load(Ordering::Acquire) + /// Derived on every call from the two cursors rather than cached, so there is + /// no published value that can be left stale by a race between the two tasks + /// that advance them. A bare `IndexStore` has no writer, so its visible + /// prefix is simply what has been indexed. + pub fn visible_count(&self) -> usize { + let indexed = self.indexed_count(); + match &self.durability { + Some((cursors, global_offset)) => cursors.visible_count(indexed, *global_offset), + None => indexed, + } + } + + /// Bind this memtable's indexes to the writer's cursors. Called once at + /// construction, before the memtable is published. + pub(crate) fn set_durability(&mut self, cursors: Arc, global_offset: usize) { + self.durability = Some((cursors, global_offset)); } } @@ -1003,6 +1171,7 @@ mod tests { use arrow_array::{Int32Array, StringArray}; use arrow_schema::{DataType, Field, Schema as ArrowSchema}; use log::warn; + use rstest::rstest; use std::sync::Arc; use uuid::Uuid; @@ -1047,6 +1216,21 @@ mod tests { .unwrap() } + fn create_sized_batch(schema: &ArrowSchema, start_id: i32, num_rows: usize) -> RecordBatch { + let ids: Vec = (0..num_rows as i32).map(|i| start_id + i).collect(); + let names: Vec = ids.iter().map(|id| format!("name-{id}")).collect(); + let descriptions: Vec = ids.iter().map(|id| format!("hello world {id}")).collect(); + RecordBatch::try_new( + Arc::new(schema.clone()), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(names)), + Arc::new(StringArray::from(descriptions)), + ], + ) + .unwrap() + } + fn fts_index_metadata(index_version: i32) -> IndexMetadata { let details = pbold::InvertedIndexDetails::try_from(&InvertedIndexParams::default()).unwrap(); @@ -1492,8 +1676,124 @@ mod tests { assert!(registry.get_fts_by_field_id(2).is_some()); } + fn vector_schema() -> Arc { + Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("description", DataType::Utf8, true), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4), + true, + ), + Field::new( + "f64_vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float64, true)), 4), + true, + ), + ])) + } + + /// Every index config that would fail *deterministically* on insert must be + /// rejected at open instead. Such a config also fails on WAL replay, so once + /// a row is durable the shard could never reopen — poison-and-replay would + /// not terminate. + #[rstest] + #[case::btree_ok(MemIndexConfig::BTree(BTreeIndexConfig { + name: "idx".into(), field_id: 0, column: "id".into(), + }), None)] + #[case::btree_missing_column(MemIndexConfig::BTree(BTreeIndexConfig { + name: "idx".into(), field_id: 9, column: "nope".into(), + }), Some("not in the shard schema"))] + // Column exists, but its field_id names a *different* column ("id" is 0, not 1). + #[case::btree_field_id_column_mismatch(MemIndexConfig::BTree(BTreeIndexConfig { + name: "idx".into(), field_id: 1, column: "id".into(), + }), Some("has field_id 0"))] + #[case::fts_ok(MemIndexConfig::Fts(FtsIndexConfig::new( + "idx".into(), 1, "description".into(), + )), None)] + #[case::fts_non_utf8(MemIndexConfig::Fts(FtsIndexConfig::new( + "idx".into(), 0, "id".into(), + )), Some("requires a Utf8, LargeUtf8, or Utf8View column"))] + #[case::fts_missing_column(MemIndexConfig::Fts(FtsIndexConfig::new( + "idx".into(), 9, "nope".into(), + )), Some("not in the shard schema"))] + #[case::hnsw_ok(MemIndexConfig::Hnsw(Box::new(HnswIndexConfig::new( + "idx".into(), 2, "vector".into(), DistanceType::L2, + ))), None)] + #[case::hnsw_not_a_vector(MemIndexConfig::Hnsw(Box::new(HnswIndexConfig::new( + "idx".into(), 0, "id".into(), DistanceType::L2, + ))), Some("requires a FixedSizeList column"))] + #[case::hnsw_wrong_item_type(MemIndexConfig::Hnsw(Box::new(HnswIndexConfig::new( + "idx".into(), 3, "f64_vector".into(), DistanceType::L2, + ))), Some("item type Float64"))] + #[case::hnsw_missing_column(MemIndexConfig::Hnsw(Box::new(HnswIndexConfig::new( + "idx".into(), 9, "nope".into(), DistanceType::L2, + ))), Some("not in the shard schema"))] + fn test_validate_index_configs( + #[case] config: MemIndexConfig, + #[case] expected_error: Option<&str>, + ) { + let schema = vector_schema(); + let lance_schema = LanceSchema::try_from(schema.as_ref()).unwrap(); + let result = validate_index_configs(&[config], &schema, &lance_schema, &[]); + match expected_error { + None => result.expect("valid config must pass validation"), + Some(fragment) => { + let message = result + .expect_err("invalid config must be rejected") + .to_string(); + assert!( + message.contains(fragment), + "error must explain the mismatch; wanted {fragment:?}, got {message:?}" + ); + } + } + } + + /// A composite PK builds an order-preserving encoded key, so its columns must + /// be encodable. A single-column PK aliases a BTree entry, which accepts any + /// type — so it must *not* be rejected here. + #[test] + fn test_validate_composite_pk_column_types() { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, false), + Field::new( + "coords", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2), + true, + ), + ])); + let lance_schema = LanceSchema::try_from(schema.as_ref()).unwrap(); + + validate_index_configs(&[], &schema, &lance_schema, &["id".into(), "name".into()]) + .expect("Int32 + Utf8 composite PK must be encodable"); + + let err = + validate_index_configs(&[], &schema, &lance_schema, &["id".into(), "coords".into()]) + .expect_err("a FixedSizeList PK column has no order-preserving encoding"); + assert!( + err.to_string().contains("order-preserving key encoding"), + "error must name the reason, got {err}" + ); + + // A single-column PK of the same type is fine: it aliases a BTree. + validate_index_configs(&[], &schema, &lance_schema, &["coords".into()]) + .expect("single-column PK aliases a BTree and accepts any type"); + + // But every PK column must exist. A single-column PK naming an absent + // column is rejected here, not left to fail deterministically on every + // later index build and WAL replay. + let err = validate_index_configs(&[], &schema, &lance_schema, &["missing".into()]) + .expect_err("a single-column PK on an absent column must be rejected"); + assert!( + err.to_string().contains("not in the shard schema"), + "error must name the missing column, got {err}" + ); + } + #[test] - fn test_index_store_max_visible_batch_position() { + fn test_index_store_indexed_count() { let schema = create_test_schema(); let mut registry = IndexStore::new(); @@ -1502,7 +1802,7 @@ mod tests { registry.add_fts("desc_idx".to_string(), 2, "description".to_string()); // Initial watermark should be 0 (no data indexed yet) - assert_eq!(registry.max_visible_batch_position(), 0); + assert_eq!(registry.indexed_count(), 0); // Insert with batch position tracking let batch = create_test_batch(&schema, 0); @@ -1510,20 +1810,54 @@ mod tests { .insert_with_batch_position(&batch, 0, Some(5)) .unwrap(); - // Now watermark should be 5 - assert_eq!(registry.max_visible_batch_position(), 5); + // Indexing batch position 5 means the prefix [0, 6) is indexed. + assert_eq!(registry.indexed_count(), 6); // Insert with higher batch position registry .insert_with_batch_position(&batch, 3, Some(10)) .unwrap(); - // Watermark should advance to 10 - assert_eq!(registry.max_visible_batch_position(), 10); + // Advances to cover batch position 10. + assert_eq!(registry.indexed_count(), 11); - // Insert without batch position shouldn't change watermark + // Insert without batch position shouldn't change the cursor registry.insert(&batch, 6).unwrap(); - assert_eq!(registry.max_visible_batch_position(), 10); + assert_eq!(registry.indexed_count(), 11); + } + + /// `insert_batches` picks the inline or the threaded path by row count, so + /// exercise both and assert they leave the same index state: every row indexed + /// exactly once, in every index, with a timing reported for each. + #[rstest] + #[case::inline(8)] + #[case::threaded(PARALLEL_INDEX_MIN_ROWS + 64)] + fn test_insert_batches_indexes_every_row_once(#[case] num_rows: usize) { + let schema = create_test_schema(); + let mut registry = IndexStore::new(); + registry.add_btree("id_idx".to_string(), 0, "id".to_string()); + registry.add_fts("desc_idx".to_string(), 2, "description".to_string()); + + let batch = create_sized_batch(&schema, 0, num_rows); + let durations = registry + .insert_batches(&[StoredBatch::new(batch, 0, 2)]) + .unwrap(); + + assert_eq!(durations.len(), 2, "expected one timing per index"); + assert!(durations.contains_key("id_idx")); + assert!(durations.contains_key("desc_idx")); + + let btree = registry.get_btree("id_idx").unwrap(); + for id in 0..num_rows as i32 { + let positions = btree.get(&ScalarValue::Int32(Some(id))); + assert_eq!( + positions.len(), + 1, + "id={id} should be indexed exactly once, got {positions:?}" + ); + } + assert_eq!(registry.get_fts("desc_idx").unwrap().doc_count(), num_rows); + assert_eq!(registry.indexed_count(), 3); } #[test] diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index 6e0d12a276e..04b07fcbda2 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -1098,23 +1098,20 @@ impl FtsMemIndex { let st = self.state.load_full(); let batch_position = st.tail.next_position(); + // A missing column is a config error, not an empty document: silently + // appending an empty batch would leave the index empty forever while the + // shard reported healthy. BTree and HNSW both reject this; so do we. + // `validate_index_configs` catches it at open, so reaching here means a + // batch got past the memtable's schema-equality gate. let Some(col_idx) = batch .schema() .column_with_name(&self.column_name) .map(|(idx, _)| idx) else { - // Column missing: nothing to index, but publish an empty batch so - // the tail's visibility counters keep up with the writer. - st.tail.append_batch( - batch_position, - row_offset, - batch.num_rows() as u32, - vec![0; batch.num_rows()], - 0, - FxHashMap::default(), - self.params.has_positions(), - ); - return Ok(()); + return Err(Error::invalid_input(format!( + "FTS index column '{}' is not in the inserted batch schema", + self.column_name + ))); }; let column = batch.column(col_idx); diff --git a/rust/lance/src/dataset/mem_wal/memtable.rs b/rust/lance/src/dataset/mem_wal/memtable.rs index 77611fd7c47..dec498f0ef6 100644 --- a/rust/lance/src/dataset/mem_wal/memtable.rs +++ b/rust/lance/src/dataset/mem_wal/memtable.rs @@ -7,7 +7,6 @@ pub mod batch_store; pub mod flush; pub mod scanner; -use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -79,13 +78,6 @@ pub struct MemTable { /// Generation number (incremented on flush). generation: u64, - /// WAL batch mapping: batch_position -> (wal_entry_position, position within WAL entry). - wal_batch_mapping: HashMap, - /// Last WAL entry position that has been flushed. - last_flushed_wal_entry_position: u64, - /// Set of batch IDs that have been flushed to WAL. - flushed_batch_positions: HashSet, - /// Primary key bloom filter for staleness detection. pk_bloom_filter: Sbbf, /// Primary key field IDs (for bloom filter updates). @@ -187,6 +179,29 @@ impl MemTable { pk_field_ids: Vec, cache_config: CacheConfig, batch_capacity: usize, + ) -> Result { + Self::with_capacity_at( + schema, + generation, + pk_field_ids, + cache_config, + batch_capacity, + 0, + ) + } + + /// Create a memtable whose batch 0 sits at `global_offset` in the writer's + /// batch sequence. Every memtable after the writer's first is rotated in by + /// `freeze_memtable`, which stamps the outgoing memtable's `global_end()` + /// here so writer-global cursors stay mappable onto local batch positions. + #[allow(clippy::too_many_arguments)] + pub fn with_capacity_at( + schema: Arc, + generation: u64, + pk_field_ids: Vec, + cache_config: CacheConfig, + batch_capacity: usize, + global_offset: usize, ) -> Result { let lance_schema = Schema::try_from(schema.as_ref())?; @@ -205,7 +220,7 @@ impl MemTable { let dataset_uri = format!("memory://{}", Uuid::new_v4()); // Create lock-free batch store - let batch_store = Arc::new(BatchStore::with_capacity(batch_capacity)); + let batch_store = Arc::new(BatchStore::with_capacity_at(batch_capacity, global_offset)); // Create memtable_flush_completion cell immediately so backpressure can // wait on it even before the memtable is frozen. Every memtable will @@ -220,9 +235,6 @@ impl MemTable { cache_config, cached_dataset: RwLock::new(None), generation, - wal_batch_mapping: HashMap::new(), - last_flushed_wal_entry_position: 0, - flushed_batch_positions: HashSet::new(), pk_bloom_filter, pk_field_ids, // Initialize with an empty IndexStore so the visibility cursor has @@ -483,51 +495,39 @@ impl MemTable { /// Get batches visible up to a specific batch position (inclusive). /// - /// A batch at position `i` is visible if `i <= max_visible_batch_position`. + /// A batch at position `i` is visible if `i <= visible_count`. /// /// # Arguments /// - /// * `max_visible_batch_position` - The maximum batch position to include (inclusive) + /// * `visible_count` - The maximum batch position to include (inclusive) /// /// # Returns /// /// Vector of visible batches. - pub async fn get_visible_batches(&self, max_visible_batch_position: usize) -> Vec { - self.batch_store - .visible_record_batches(max_visible_batch_position) + pub async fn get_visible_batches(&self, visible_count: usize) -> Vec { + self.batch_store.visible_record_batches(visible_count) } /// Get batch positions visible up to a specific batch position (inclusive). /// /// This is useful for filtering index results by visibility. - pub async fn get_max_visible_batch_positions( - &self, - max_visible_batch_position: usize, - ) -> Vec { - self.batch_store - .max_visible_batch_positions(max_visible_batch_position) + pub async fn get_visible_batch_positions(&self, visible_count: usize) -> Vec { + self.batch_store.visible_batch_positions(visible_count) } /// Check if a specific batch is visible at a given visibility position. /// /// Returns true if the batch is visible, false if not visible or doesn't exist. - pub async fn is_batch_visible( - &self, - batch_position: usize, - max_visible_batch_position: usize, - ) -> bool { + pub async fn is_batch_visible(&self, batch_position: usize, visible_count: usize) -> bool { self.batch_store - .is_batch_visible(batch_position, max_visible_batch_position) + .is_batch_visible(batch_position, visible_count) } /// Scan batches visible up to a specific batch position. /// /// This combines `get_visible_batches` with the scan interface. - pub async fn scan_batches_at_position( - &self, - max_visible_batch_position: usize, - ) -> Result> { - Ok(self.get_visible_batches(max_visible_batch_position).await) + pub async fn scan_batches_at_position(&self, visible_count: usize) -> Result> { + Ok(self.get_visible_batches(visible_count).await) } /// Update the bloom filter with primary keys from a batch. @@ -561,30 +561,6 @@ impl MemTable { Ok(()) } - /// Mark batches as flushed to WAL. - /// - /// Updates the WAL batch mapping for use during MemTable flush. - /// Also updates the batch_store's watermark to the highest flushed batch_position. - pub fn mark_wal_flushed( - &mut self, - batch_positions: &[usize], - wal_entry_position: u64, - positions: &[usize], - ) { - for (idx, &batch_position) in batch_positions.iter().enumerate() { - self.wal_batch_mapping - .insert(batch_position, (wal_entry_position, positions[idx])); - self.flushed_batch_positions.insert(batch_position); - } - self.last_flushed_wal_entry_position = wal_entry_position; - - // Update batch_store watermark to the highest batch_position flushed (inclusive) - if let Some(&max_batch_position) = batch_positions.iter().max() { - self.batch_store - .set_max_flushed_batch_position(max_batch_position); - } - } - /// Get or create a Dataset for reading. /// /// Uses caching based on the configured eventual consistency strategy: @@ -726,16 +702,6 @@ impl MemTable { self.batch_store.estimated_bytes() + self.pk_bloom_filter.estimated_memory_size() } - /// Get the WAL batch mapping. - pub fn wal_batch_mapping(&self) -> &HashMap { - &self.wal_batch_mapping - } - - /// Get the last flushed WAL entry position. - pub fn last_flushed_wal_entry_position(&self) -> u64 { - self.last_flushed_wal_entry_position - } - /// Get the bloom filter for serialization. pub fn bloom_filter(&self) -> &Sbbf { &self.pk_bloom_filter @@ -757,17 +723,15 @@ impl MemTable { self.indexes.take() } - /// Check if all batches have been flushed to WAL. - pub fn all_flushed_to_wal(&self) -> bool { - self.batch_store.pending_wal_flush_count() == 0 + /// Whether every committed batch in this memtable is WAL-durable, given the + /// writer-global durability cursor. The L0 flush's precondition. + pub fn all_flushed_to_wal(&self, durable: usize) -> bool { + self.batch_store.pending_wal_flush_count(durable) == 0 } - /// Get unflushed batch IDs. - pub fn unflushed_batch_positions(&self) -> Vec { - let batch_count = self.batch_count(); - (0..batch_count) - .filter(|id| !self.flushed_batch_positions.contains(id)) - .collect() + /// Writer-global coordinate one past this memtable's last committed batch. + pub fn global_end(&self) -> usize { + self.batch_store.global_end() } /// Get cache configuration. @@ -785,18 +749,13 @@ impl MemTable { self.batch_store.remaining_capacity() } - /// Check if batch store is full. - pub fn is_batch_store_full(&self) -> bool { - self.batch_store.is_full() - } - /// Create a scanner for querying this MemTable. /// /// # Arguments /// - /// * `max_visible_batch_position` - Maximum batch position visible (inclusive) + /// * `visible_count` - Maximum batch position visible (inclusive) /// - /// The scanner captures the current `max_visible_batch_position` from the + /// The scanner captures the current `visible_count` from the /// `IndexStore` at construction time to ensure consistent visibility. /// /// # Panics @@ -931,29 +890,11 @@ mod tests { assert_eq!(total_rows, 15); } + /// `all_flushed_to_wal(durable)` is the L0 flush's precondition (`flush.rs:171`): + /// false while any committed batch is still un-appended, true once the + /// durability watermark covers every one of them. #[tokio::test] - async fn test_memtable_wal_mapping() { - let schema = create_test_schema(); - let mut memtable = MemTable::new(schema.clone(), 1, vec![]).unwrap(); - - let batch_position = memtable - .insert(create_test_batch(&schema, 10)) - .await - .unwrap(); - assert!(!memtable.all_flushed_to_wal()); - - memtable.mark_wal_flushed(&[batch_position], 5, &[0]); - - assert!(memtable.all_flushed_to_wal()); - assert_eq!( - memtable.wal_batch_mapping().get(&batch_position), - Some(&(5, 0)) - ); - assert_eq!(memtable.last_flushed_wal_entry_position(), 5); - } - - #[tokio::test] - async fn test_memtable_unflushed_batches() { + async fn test_all_flushed_to_wal_tracks_the_durability_watermark() { let schema = create_test_schema(); let mut memtable = MemTable::new(schema.clone(), 1, vec![]).unwrap(); @@ -965,12 +906,16 @@ mod tests { .insert(create_test_batch(&schema, 5)) .await .unwrap(); + assert!(!memtable.all_flushed_to_wal(0), "nothing is durable yet"); - assert_eq!(memtable.unflushed_batch_positions(), vec![batch1, batch2]); - - memtable.mark_wal_flushed(&[batch1], 1, &[0]); + let durable = batch1 + 1; + assert!( + !memtable.all_flushed_to_wal(durable), + "batch2 is still waiting on its WAL append" + ); - assert_eq!(memtable.unflushed_batch_positions(), vec![batch2]); + let durable = batch2 + 1; + assert!(memtable.all_flushed_to_wal(durable)); } #[tokio::test] @@ -992,23 +937,21 @@ mod tests { .await .unwrap(); - // max_visible_batch_position=1 means positions 0 and 1 are visible - let visible = memtable.get_visible_batches(1).await; + // A count of N exposes the prefix [0, N). + let visible = memtable.get_visible_batches(2).await; assert_eq!(visible.len(), 2); let total_rows: usize = visible.iter().map(|b| b.num_rows()).sum(); assert_eq!(total_rows, 15); // 10 + 5 - // max_visible_batch_position=2 means all batches are visible - let visible = memtable.get_visible_batches(2).await; + let visible = memtable.get_visible_batches(3).await; assert_eq!(visible.len(), 3); - // max_visible_batch_position=0 means only position 0 is visible - let visible = memtable.get_visible_batches(0).await; - assert_eq!(visible.len(), 1); + // A count of 0 exposes nothing — not "batch 0". + assert!(memtable.get_visible_batches(0).await.is_empty()); } #[tokio::test] - async fn test_memtable_get_max_visible_batch_positions() { + async fn test_memtable_get_visible_batch_positions() { let schema = create_test_schema(); let mut memtable = MemTable::new(schema.clone(), 1, vec![]).unwrap(); @@ -1026,17 +969,15 @@ mod tests { .await .unwrap(); - // max_visible_batch_position=1 means positions 0 and 1 visible - let visible_ids = memtable.get_max_visible_batch_positions(1).await; + // A count of N exposes the prefix [0, N). + let visible_ids = memtable.get_visible_batch_positions(2).await; assert_eq!(visible_ids, vec![0, 1]); - // max_visible_batch_position=2 means all positions visible - let visible_ids = memtable.get_max_visible_batch_positions(2).await; + let visible_ids = memtable.get_visible_batch_positions(3).await; assert_eq!(visible_ids, vec![0, 1, 2]); - // max_visible_batch_position=0 means only position 0 visible - let visible_ids = memtable.get_max_visible_batch_positions(0).await; - assert_eq!(visible_ids, vec![0]); + // A count of 0 exposes nothing. + assert!(memtable.get_visible_batch_positions(0).await.is_empty()); } #[tokio::test] @@ -1057,14 +998,14 @@ mod tests { .await .unwrap(); // position 2 - // batch_position 0 is visible when max_visible_batch_position >= 0 - assert!(memtable.is_batch_visible(0, 0).await); + // A count of 0 means nothing is visible, batch 0 included. + assert!(!memtable.is_batch_visible(0, 0).await); + + // Batch i is visible once the count exceeds i. assert!(memtable.is_batch_visible(0, 1).await); assert!(memtable.is_batch_visible(0, 2).await); - - // batch_position 2 is only visible when max_visible_batch_position >= 2 assert!(!memtable.is_batch_visible(2, 1).await); - assert!(memtable.is_batch_visible(2, 2).await); + assert!(!memtable.is_batch_visible(2, 2).await); assert!(memtable.is_batch_visible(2, 3).await); // Non-existent batch @@ -1085,12 +1026,21 @@ mod tests { .await .unwrap(); // position 1 - let batches = memtable.scan_batches_at_position(0).await.unwrap(); + let batches = memtable.scan_batches_at_position(1).await.unwrap(); assert_eq!(batches.len(), 1); assert_eq!(batches[0].num_rows(), 10); - let batches = memtable.scan_batches_at_position(1).await.unwrap(); + let batches = memtable.scan_batches_at_position(2).await.unwrap(); assert_eq!(batches.len(), 2); + + // Nothing indexed yet => nothing scannable. + assert!( + memtable + .scan_batches_at_position(0) + .await + .unwrap() + .is_empty() + ); } #[tokio::test] @@ -1101,7 +1051,7 @@ mod tests { assert_eq!(memtable.batch_capacity(), 3); assert_eq!(memtable.remaining_batch_capacity(), 3); - assert!(!memtable.is_batch_store_full()); + assert!(!memtable.batch_store().is_full()); // Fill up the store memtable @@ -1117,7 +1067,7 @@ mod tests { .await .unwrap(); - assert!(memtable.is_batch_store_full()); + assert!(memtable.batch_store().is_full()); assert_eq!(memtable.remaining_batch_capacity(), 0); // Next insert should fail diff --git a/rust/lance/src/dataset/mem_wal/memtable/batch_store.rs b/rust/lance/src/dataset/mem_wal/memtable/batch_store.rs index 30e47e1a9bb..f16f9bea1fd 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/batch_store.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/batch_store.rs @@ -195,10 +195,14 @@ pub struct BatchStore { /// Estimated size in bytes (for flush threshold). estimated_bytes: AtomicUsize, - /// WAL flush watermark: the last batch ID that has been flushed to WAL (inclusive). - /// Uses usize::MAX as sentinel for "nothing flushed yet". - /// This is per-memtable tracking, not global. - max_flushed_batch_position: AtomicUsize, + /// Writer-global coordinate of this store's batch 0. + /// + /// A *coordinate*, not a cursor: stamped once at construction and never + /// moved. `global_position = global_offset + local_position`. Batch + /// positions restart at 0 in every memtable, so this is the only thing that + /// lets a writer-global cursor (the WAL durability count) be mapped onto a + /// particular store. + global_offset: usize, } // SAFETY: Safe to share across threads because: @@ -221,6 +225,13 @@ impl BatchStore { /// /// Panics if capacity is 0. pub fn with_capacity(capacity: usize) -> Self { + Self::with_capacity_at(capacity, 0) + } + + /// Create a store whose batch 0 sits at `global_offset` in the writer's + /// batch sequence. Used by `freeze_memtable` for every memtable after the + /// first; the first starts at 0. + pub fn with_capacity_at(capacity: usize, global_offset: usize) -> Self { assert!(capacity > 0, "capacity must be > 0"); // Allocate uninitialized storage @@ -235,7 +246,7 @@ impl BatchStore { capacity, total_rows: AtomicUsize::new(0), estimated_bytes: AtomicUsize::new(0), - max_flushed_batch_position: AtomicUsize::new(usize::MAX), // Nothing flushed yet + global_offset, } } @@ -437,74 +448,54 @@ impl BatchStore { // WAL Flush Tracking API // ========================================================================= - /// Get the WAL flush watermark (the last batch ID that was flushed, inclusive). - /// Returns None if nothing has been flushed yet. + /// Writer-global coordinate one past this store's last committed batch. #[inline] - pub fn max_flushed_batch_position(&self) -> Option { - let watermark = self.max_flushed_batch_position.load(Ordering::Acquire); - if watermark == usize::MAX { - None - } else { - Some(watermark) - } + pub fn global_end(&self) -> usize { + self.global_offset + self.committed_len.load(Ordering::Acquire) } - /// Update the WAL flush watermark after successful WAL flush. - /// - /// # Arguments - /// - /// * `batch_position` - The last batch ID that was flushed (inclusive) + /// This store's writer-global coordinate for batch 0. #[inline] - pub fn set_max_flushed_batch_position(&self, batch_position: usize) { - debug_assert!( - batch_position != usize::MAX, - "batch_position cannot be usize::MAX (reserved as sentinel)" - ); - self.max_flushed_batch_position - .store(batch_position, Ordering::Release); + pub fn global_offset(&self) -> usize { + self.global_offset } - /// Get the number of batches pending WAL flush. + /// The local exclusive end of this store covered by a writer-global cursor. + /// + /// Saturating in both directions, and both directions are reachable in + /// normal operation: a cursor *below* this store's offset means "nothing + /// here yet" (the store was rotated in after the cursor last advanced — + /// the ordinary state of a fresh memtable), and a cursor beyond its end + /// clamps to what is committed. + /// + /// This is the **only** place the global-to-local subtraction is written. + /// Open-coding it underflows on every memtable rotation, which in release + /// wraps to a huge end and makes the whole new memtable instantly visible. #[inline] - pub fn pending_wal_flush_count(&self) -> usize { - let committed = self.committed_len.load(Ordering::Acquire); - let watermark = self.max_flushed_batch_position.load(Ordering::Acquire); - if watermark == usize::MAX { - // Nothing flushed yet, all committed batches are pending - committed - } else { - // Batches [0, watermark] are flushed, so pending = committed - (watermark + 1) - committed.saturating_sub(watermark + 1) - } + pub fn local_end(&self, global_cursor: usize) -> usize { + global_cursor + .saturating_sub(self.global_offset) + .min(self.committed_len.load(Ordering::Acquire)) } - /// Check if all committed batches have been WAL-flushed. + /// Batches in this store still waiting on their WAL append. #[inline] - pub fn is_wal_flush_complete(&self) -> bool { - self.pending_wal_flush_count() == 0 + pub fn pending_wal_flush_count(&self, durable: usize) -> usize { + self.committed_len.load(Ordering::Acquire) - self.local_end(durable) } - /// Get the range of batch IDs pending WAL flush: [start, end). - /// Returns None if nothing pending. + /// Local range `[start, end)` of batches still waiting on their WAL append, + /// or `None` when the store is fully durable. #[inline] - pub fn pending_wal_flush_range(&self) -> Option<(usize, usize)> { - let committed = self.committed_len.load(Ordering::Acquire); - let watermark = self.max_flushed_batch_position.load(Ordering::Acquire); - let start = if watermark == usize::MAX { - 0 - } else { - watermark + 1 - }; - if committed > start { - Some((start, committed)) - } else { - None - } + pub fn pending_wal_flush_range(&self, durable: usize) -> Option<(usize, usize)> { + let start = self.local_end(durable); + let end = self.committed_len.load(Ordering::Acquire); + (end > start).then_some((start, end)) } /// Get a point-in-time summary of batches pending WAL flush. - pub fn pending_wal_flush_stats(&self) -> PendingWalFlushStats { - let Some((start, end)) = self.pending_wal_flush_range() else { + pub fn pending_wal_flush_stats(&self, durable: usize) -> PendingWalFlushStats { + let Some((start, end)) = self.pending_wal_flush_range(durable) else { return PendingWalFlushStats::default(); }; @@ -643,66 +634,53 @@ impl BatchStore { // Visibility API // ========================================================================= - /// Get batches visible up to a specific batch position (inclusive). + /// Batches in the visible prefix `[0, visible_count)`. /// - /// A batch at position `i` is visible if `i <= max_visible_batch_position`. - pub fn visible_batches(&self, max_visible_batch_position: usize) -> Vec<&StoredBatch> { - let len = self.committed_len.load(Ordering::Acquire); - let end = (max_visible_batch_position + 1).min(len); + /// `visible_count` is an **exclusive count**, not an inclusive position: 0 + /// means nothing is visible. As an inclusive position, 0 meant *both* + /// "nothing visible" and "batch 0 is visible", so a batch that was committed + /// to the store but not yet indexed or WAL-durable was readable for a full + /// PUT round-trip. The count makes that off-by-one inexpressible. + pub fn visible_batches(&self, visible_count: usize) -> Vec<&StoredBatch> { + let end = visible_count.min(self.committed_len.load(Ordering::Acquire)); (0..end).filter_map(|i| self.get(i)).collect() } - /// Get batch positions visible up to a specific batch position (inclusive). - pub fn max_visible_batch_positions(&self, max_visible_batch_position: usize) -> Vec { - let len = self.committed_len.load(Ordering::Acquire); - let end = (max_visible_batch_position + 1).min(len); + /// Positions of the batches in the visible prefix. + pub fn visible_batch_positions(&self, visible_count: usize) -> Vec { + let end = visible_count.min(self.committed_len.load(Ordering::Acquire)); (0..end).collect() } - /// The inclusive maximum visible *row* position at `max_visible_batch_position`, - /// or `None` when no rows are visible. The visible batches are the committed - /// prefix `[0, last_visible_idx]`; each batch carries its cumulative - /// `row_offset`, so this is the end of the last visible batch minus one. - /// Used to bound MVCC seeks against the maintained PK-position index. - pub fn max_visible_row(&self, max_visible_batch_position: usize) -> Option { - let len = self.committed_len.load(Ordering::Acquire); - if len == 0 { - return None; - } - let last_visible_idx = max_visible_batch_position.min(len - 1); - let last = self.get(last_visible_idx)?; + /// The inclusive maximum visible *row* position, or `None` when no rows are + /// visible. Each batch carries its cumulative `row_offset`, so this is the + /// end of the last visible batch minus one. Bounds MVCC seeks against the + /// maintained PK-position index. + pub fn max_visible_row(&self, visible_count: usize) -> Option { + let end = visible_count.min(self.committed_len.load(Ordering::Acquire)); + let last = self.get(end.checked_sub(1)?)?; let visible_end = last.row_offset + last.num_rows as u64; // exclusive visible_end.checked_sub(1) } - /// Check if a specific batch is visible at a given visibility position. + /// Whether a batch falls inside the visible prefix. #[inline] - pub fn is_batch_visible( - &self, - batch_position: usize, - max_visible_batch_position: usize, - ) -> bool { + pub fn is_batch_visible(&self, batch_position: usize, visible_count: usize) -> bool { let len = self.committed_len.load(Ordering::Acquire); - batch_position < len && batch_position <= max_visible_batch_position + batch_position < len && batch_position < visible_count } - /// Get visible RecordBatches (clones the data). - pub fn visible_record_batches(&self, max_visible_batch_position: usize) -> Vec { - self.visible_batches(max_visible_batch_position) + /// Visible RecordBatches (clones the data). + pub fn visible_record_batches(&self, visible_count: usize) -> Vec { + self.visible_batches(visible_count) .into_iter() .map(|b| b.data.clone()) .collect() } - /// Get visible RecordBatches with their row offsets. - /// - /// Returns tuples of (batch, row_offset) for each visible batch. - /// The row_offset is the starting row position for that batch. - pub fn visible_batches_with_offsets( - &self, - max_visible_batch_position: usize, - ) -> Vec<(RecordBatch, u64)> { - self.visible_batches(max_visible_batch_position) + /// Visible RecordBatches paired with the row position each one starts at. + pub fn visible_batches_with_offsets(&self, visible_count: usize) -> Vec<(RecordBatch, u64)> { + self.visible_batches(visible_count) .into_iter() .map(|b| (b.data.clone(), b.row_offset)) .collect() @@ -935,17 +913,52 @@ mod tests { store.append(create_test_batch(10)).unwrap(); // position 3 store.append(create_test_batch(10)).unwrap(); // position 4 - // max_visible_batch_position=2 means positions 0, 1, 2 are visible - let visible = store.max_visible_batch_positions(2); - assert_eq!(visible, vec![0, 1, 2]); + // A count of N exposes the prefix [0, N). + assert_eq!(store.visible_batch_positions(3), vec![0, 1, 2]); + assert_eq!(store.visible_batch_positions(5), vec![0, 1, 2, 3, 4]); + + // A count of 0 exposes nothing. Under the old inclusive cursor this + // case was indistinguishable from "batch 0 is visible", so every + // memtable leaked its first batch before it was indexed or durable. + assert!(store.visible_batch_positions(0).is_empty()); + + // Beyond the committed range, clamp. + assert_eq!(store.visible_batch_positions(99), vec![0, 1, 2, 3, 4]); + } + + /// The zero of the visibility cursor must be unambiguous. + /// + /// `BatchStore::append` publishes `committed_len` on the put path, under the + /// state lock, *before* the WAL flush that indexes the batch is even + /// triggered — and that flush is a ~100ms S3 PUT on another task. So batch 0 + /// sits committed and readable for a full round-trip before it is indexed or + /// durable. As an inclusive position, a cursor of 0 meant both "nothing is + /// visible" and "batch 0 is visible", so every read arm backed by the batch + /// store served that batch while the index-backed arms did not — the tiers + /// actively disagreed. An exclusive count makes the state inexpressible. + #[test] + fn test_zero_cursor_hides_the_committed_but_unindexed_prefix() { + let store = BatchStore::with_capacity(4); + store.append(create_test_batch(10)).unwrap(); + store.append(create_test_batch(10)).unwrap(); + + // Committed, but nothing indexed yet: every visibility query must agree + // that there is nothing to read. + assert!(store.visible_batches(0).is_empty()); + assert!(store.visible_batch_positions(0).is_empty()); + assert!(store.visible_record_batches(0).is_empty()); + assert!(store.visible_batches_with_offsets(0).is_empty()); + assert!(!store.is_batch_visible(0, 0)); + assert_eq!(store.max_visible_row(0), None); - // max_visible_batch_position=4 means all visible - let visible = store.max_visible_batch_positions(4); - assert_eq!(visible, vec![0, 1, 2, 3, 4]); + // The batches are there — they are simply not yet published. + assert_eq!(store.len(), 2); - // max_visible_batch_position=0 means only position 0 visible - let visible = store.max_visible_batch_positions(0); - assert_eq!(visible, vec![0]); + // Indexing batch 0 publishes exactly batch 0. + assert_eq!(store.visible_batches(1).len(), 1); + assert!(store.is_batch_visible(0, 1)); + assert!(!store.is_batch_visible(1, 1)); + assert_eq!(store.max_visible_row(1), Some(9)); } #[test] @@ -956,14 +969,17 @@ mod tests { store.append(create_test_batch(10)).unwrap(); // position 1 store.append(create_test_batch(10)).unwrap(); // position 2 - // Batch at position 0 is visible when max_visible_batch_position >= 0 - assert!(store.is_batch_visible(0, 0)); + // A count of 0 means *nothing* is visible — including batch 0. As an + // inclusive position this case was indistinguishable from "batch 0 is + // visible", so a batch that was committed to the store but not yet + // indexed or WAL-durable was readable for a full PUT round-trip. + assert!(!store.is_batch_visible(0, 0)); + + // Batch i is visible once the count exceeds i. assert!(store.is_batch_visible(0, 1)); assert!(store.is_batch_visible(0, 2)); - - // Batch at position 2 is only visible when max_visible_batch_position >= 2 assert!(!store.is_batch_visible(2, 1)); - assert!(store.is_batch_visible(2, 2)); + assert!(!store.is_batch_visible(2, 2)); assert!(store.is_batch_visible(2, 3)); // Batch 3 doesn't exist @@ -972,7 +988,7 @@ mod tests { #[test] fn test_max_visible_row() { - // (1) Empty store: no rows are visible at any position. + // (1) Empty store: no rows are visible at any count. let store = BatchStore::with_capacity(10); assert_eq!(store.max_visible_row(0), None); assert_eq!(store.max_visible_row(100), None); @@ -982,23 +998,24 @@ mod tests { store.append(create_test_batch(20)).unwrap(); // position 1 store.append(create_test_batch(30)).unwrap(); // position 2 - // (2) A position within range yields the inclusive end of that prefix. - assert_eq!(store.max_visible_row(0), Some(9)); // batch 0: 0..10 - assert_eq!(store.max_visible_row(1), Some(29)); // batch 1: 10..30 - assert_eq!(store.max_visible_row(2), Some(59)); // batch 2: 30..60 + // (2) A count of 0 means nothing is visible — not "batch 0 is visible". + assert_eq!(store.max_visible_row(0), None); + + // (3) A count of N yields the inclusive last row of the prefix [0, N). + assert_eq!(store.max_visible_row(1), Some(9)); // batch 0: 0..10 + assert_eq!(store.max_visible_row(2), Some(29)); // + batch 1: 10..30 + assert_eq!(store.max_visible_row(3), Some(59)); // + batch 2: 30..60 - // (3) A position beyond the committed range clamps to the last batch, - // i.e. the inclusive max over all rows. + // (4) A count beyond the committed range clamps to the last batch. assert_eq!(store.max_visible_row(100), Some(59)); - // (4) An empty leading batch contributes no rows: at its own position - // the inclusive end underflows to None, while a later non-empty batch - // is reported correctly. + // (5) An empty leading batch contributes no rows, so a prefix covering + // only it still yields None, while a later non-empty batch is reported. let store = BatchStore::with_capacity(10); store.append(create_test_batch(0)).unwrap(); // position 0: rows [0,0) store.append(create_test_batch(5)).unwrap(); // position 1: rows [0,5) - assert_eq!(store.max_visible_row(0), None); // empty prefix → no rows - assert_eq!(store.max_visible_row(1), Some(4)); // through batch 1 + assert_eq!(store.max_visible_row(1), None); // empty prefix → no rows + assert_eq!(store.max_visible_row(2), Some(4)); // through batch 1 } #[test] diff --git a/rust/lance/src/dataset/mem_wal/memtable/flush.rs b/rust/lance/src/dataset/mem_wal/memtable/flush.rs index e88be539e30..b843c3f1823 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/flush.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/flush.rs @@ -219,6 +219,7 @@ impl MemTableFlusher { memtable: &MemTable, epoch: u64, covered_wal_entry_position: u64, + durable: usize, ) -> Result { self.manifest_store.check_fenced(epoch).await?; @@ -226,7 +227,7 @@ impl MemTableFlusher { return Err(Error::invalid_input("Cannot flush empty MemTable")); } - if !memtable.all_flushed_to_wal() { + if !memtable.all_flushed_to_wal(durable) { return Err(Error::invalid_input( "MemTable has unflushed fragments - WAL flush required first", )); @@ -452,6 +453,7 @@ impl MemTableFlusher { epoch: u64, index_configs: &[MemIndexConfig], covered_wal_entry_position: u64, + durable: usize, ) -> Result { self.manifest_store.check_fenced(epoch).await?; @@ -459,7 +461,7 @@ impl MemTableFlusher { return Err(Error::invalid_input("Cannot flush empty MemTable")); } - if !memtable.all_flushed_to_wal() { + if !memtable.all_flushed_to_wal(durable) { return Err(Error::invalid_input( "MemTable has unflushed fragments - WAL flush required first", )); @@ -1258,11 +1260,12 @@ mod tests { .await .unwrap(); - // Not flushed to WAL yet - assert!(!memtable.all_flushed_to_wal()); + // Nothing is durable yet, so the L0 flush must refuse. + let durable = 0; + assert!(!memtable.all_flushed_to_wal(durable)); let flusher = MemTableFlusher::new(store, base_path, base_uri, shard_id, manifest_store); - let result = flusher.flush(&memtable, epoch, 0).await; + let result = flusher.flush(&memtable, epoch, 0, 0).await; assert!(result.is_err()); assert!( @@ -1291,7 +1294,7 @@ mod tests { let memtable = MemTable::new(schema, 1, vec![]).unwrap(); let flusher = MemTableFlusher::new(store, base_path, base_uri, shard_id, manifest_store); - let result = flusher.flush(&memtable, epoch, 0).await; + let result = flusher.flush(&memtable, epoch, 0, 0).await; assert!(result.is_err()); assert!(result.unwrap_err().to_string().contains("empty MemTable")); @@ -1319,8 +1322,8 @@ mod tests { .unwrap(); // Simulate WAL flush - memtable.mark_wal_flushed(&[frag_id], 1, &[0]); - assert!(memtable.all_flushed_to_wal()); + let durable = frag_id + 1; + assert!(memtable.all_flushed_to_wal(durable)); let flusher = MemTableFlusher::new( store.clone(), @@ -1329,7 +1332,7 @@ mod tests { shard_id, manifest_store.clone(), ); - let result = flusher.flush(&memtable, epoch, 1).await.unwrap(); + let result = flusher.flush(&memtable, epoch, 1, durable).await.unwrap(); assert_eq!(result.generation.generation, 1); assert_eq!(result.rows_flushed, 10); @@ -1383,7 +1386,7 @@ mod tests { .insert(create_test_batch(&schema, 10)) .await .unwrap(); - memtable.mark_wal_flushed(&[frag_id], 1, &[0]); + let durable = frag_id + 1; let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); let warmer: Arc = Arc::new(CountingWarmer { @@ -1400,7 +1403,7 @@ mod tests { ) .with_warmer(Some(warmer)); // Flush must succeed despite the warmer erroring. - let result = flusher.flush(&memtable, epoch, 1).await.unwrap(); + let result = flusher.flush(&memtable, epoch, 1, durable).await.unwrap(); assert_eq!(result.generation.generation, 1); assert_eq!( @@ -1444,7 +1447,7 @@ mod tests { ) .unwrap(); let frag_id = memtable.insert(batch).await.unwrap(); - memtable.mark_wal_flushed(&[frag_id], 1, &[0]); + let durable = frag_id + 1; let flusher = MemTableFlusher::new( store.clone(), @@ -1453,7 +1456,7 @@ mod tests { shard_id, manifest_store, ); - let result = flusher.flush(&memtable, epoch, 1).await.unwrap(); + let result = flusher.flush(&memtable, epoch, 1, durable).await.unwrap(); assert_eq!(result.rows_flushed, 5, "all physical rows are written"); // Scanning the flushed generation must honor the deletion vector and @@ -1547,7 +1550,7 @@ mod tests { ) .unwrap(); let frag_id = memtable.insert(batch).await.unwrap(); - memtable.mark_wal_flushed(&[frag_id], 1, &[0]); + let durable = frag_id + 1; let flusher = MemTableFlusher::new( store.clone(), @@ -1557,7 +1560,7 @@ mod tests { manifest_store.clone(), ); let result = flusher - .flush_with_indexes(&memtable, epoch, &[], 1) + .flush_with_indexes(&memtable, epoch, &[], 1, durable) .await .unwrap(); @@ -1647,7 +1650,7 @@ mod tests { ) .unwrap(); let frag_id = memtable.insert(batch).await.unwrap(); - memtable.mark_wal_flushed(&[frag_id], 1, &[0]); + let durable = frag_id + 1; let flusher = MemTableFlusher::new( store.clone(), @@ -1657,7 +1660,7 @@ mod tests { manifest_store.clone(), ); // The plain-flush path — what the writer dispatches to with no indexes. - let result = flusher.flush(&memtable, epoch, 1).await.unwrap(); + let result = flusher.flush(&memtable, epoch, 1, durable).await.unwrap(); let gen_path = base_path .clone() @@ -1741,7 +1744,7 @@ mod tests { ) .unwrap(); let frag_id = memtable.insert(batch).await.unwrap(); - memtable.mark_wal_flushed(&[frag_id], 1, &[0]); + let durable = frag_id + 1; let flusher = MemTableFlusher::new( store.clone(), @@ -1751,7 +1754,7 @@ mod tests { manifest_store.clone(), ); let result = flusher - .flush_with_indexes(&memtable, epoch, &index_configs, 1) + .flush_with_indexes(&memtable, epoch, &index_configs, 1, durable) .await .unwrap(); assert_eq!(result.rows_flushed, 5, "all physical rows are written"); @@ -1873,7 +1876,7 @@ mod tests { .unwrap(); // Simulate WAL flush - memtable.mark_wal_flushed(&[frag_id], 1, &[0]); + let durable = frag_id + 1; let flusher = MemTableFlusher::new( store.clone(), @@ -1883,7 +1886,7 @@ mod tests { manifest_store.clone(), ); let result = flusher - .flush_with_indexes(&memtable, epoch, &index_configs, 1) + .flush_with_indexes(&memtable, epoch, &index_configs, 1, durable) .await .unwrap(); @@ -2010,7 +2013,7 @@ mod tests { let frag_id = memtable.insert(batch).await.unwrap(); // Simulate WAL flush - memtable.mark_wal_flushed(&[frag_id], 1, &[0]); + let durable = frag_id + 1; let flusher = MemTableFlusher::new( store.clone(), @@ -2020,7 +2023,7 @@ mod tests { manifest_store.clone(), ); let result = flusher - .flush_with_indexes(&memtable, epoch, &index_configs, 1) + .flush_with_indexes(&memtable, epoch, &index_configs, 1, durable) .await .unwrap(); @@ -2160,7 +2163,7 @@ mod tests { let frag_id = memtable.insert(batch).await.unwrap(); // Simulate WAL flush - memtable.mark_wal_flushed(&[frag_id], 1, &[0]); + let durable = frag_id + 1; let flusher = MemTableFlusher::new( store.clone(), @@ -2170,7 +2173,7 @@ mod tests { manifest_store.clone(), ); let result = flusher - .flush_with_indexes(&memtable, epoch, &index_configs, 1) + .flush_with_indexes(&memtable, epoch, &index_configs, 1, durable) .await .unwrap(); diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs index e1f0d6e689c..e0a06ff85b9 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs @@ -378,7 +378,7 @@ impl ScalarPredicate { /// /// # Index Visibility Model /// -/// The scanner captures `max_visible_batch_position` from the `IndexStore` at +/// The scanner captures `visible_count` from the `IndexStore` at /// construction time. This frozen visibility ensures queries only see data /// that has been indexed, providing consistent results. /// @@ -401,8 +401,8 @@ pub struct MemTableScanner { indexes: Arc, schema: SchemaRef, /// Frozen visibility captured at scanner construction time. - /// This is the `max_visible_batch_position` from the IndexStore. - max_visible_batch_position: usize, + /// This is the `visible_count` from the IndexStore. + visible_count: usize, projection: Option>, filter: Option, limit: Option, @@ -427,7 +427,7 @@ pub struct MemTableScanner { impl MemTableScanner { /// Create a new scanner. /// - /// Captures `max_visible_batch_position` from the `IndexStore` at construction + /// Captures `visible_count` from the `IndexStore` at construction /// time to ensure consistent query visibility. /// /// # Arguments @@ -439,13 +439,13 @@ impl MemTableScanner { // Snapshot the visibility cursor at construction time. The cursor is // advanced by `flush_from_batch_store` after the WAL append succeeds, // so this snapshot reflects WAL-durable data. - let max_visible_batch_position = indexes.max_visible_batch_position(); + let visible_count = indexes.visible_count(); Self { batch_store, indexes, schema, - max_visible_batch_position, + visible_count, projection: None, filter: None, limit: None, @@ -504,12 +504,12 @@ impl MemTableScanner { self } - /// The `max_visible_batch_position` snapshot this scanner latched at + /// The `visible_count` snapshot this scanner latched at /// construction. A downstream recency filter must key on this same snapshot /// (not a fresh read of the IndexStore watermark, which a concurrent append /// could have advanced) so it stays consistent with the rows the search saw. - pub fn max_visible_batch_position(&self) -> usize { - self.max_visible_batch_position + pub fn visible_count(&self) -> usize { + self.visible_count } /// Include the _rowaddr column in output. @@ -971,7 +971,7 @@ impl MemTableScanner { let scan = MemTableScanExec::with_filter( self.batch_store.clone(), - self.max_visible_batch_position, + self.visible_count, projection_indices, self.output_schema(), self.schema.clone(), @@ -1033,7 +1033,7 @@ impl MemTableScanner { Ok(Arc::new(MemTableDedupScanExec::new( self.batch_store.clone(), - self.max_visible_batch_position, + self.visible_count, projection_indices, self.output_schema(), pk_indices, @@ -1056,7 +1056,7 @@ impl MemTableScanner { return self.plan_full_scan().await; } - let max_visible = self.max_visible_batch_position; + let max_visible = self.visible_count; let projection_indices = self.compute_projection_indices()?; let index_exec = BTreeIndexExec::new( @@ -1095,7 +1095,7 @@ impl MemTableScanner { } async fn plan_vector_search(&self, query: &VectorQuery) -> Result> { - let max_visible = self.max_visible_batch_position; + let max_visible = self.visible_count; let projection_indices = self.compute_projection_indices()?; let base_schema = self.base_output_schema(); let filter_predicate = self.filter_predicate()?; @@ -1154,7 +1154,7 @@ impl MemTableScanner { return self.empty_fts_plan(); } - let max_visible = self.max_visible_batch_position; + let max_visible = self.visible_count; let projection_indices = self.compute_projection_indices()?; let filter_predicate = self.filter_predicate()?; if let Some(pk_columns) = &self.pk_columns { @@ -1426,7 +1426,7 @@ mod tests { let indexes = Arc::new(index_store); let scanner = MemTableScanner::new(batch_store, indexes, schema.clone()); let result = scanner.try_into_batch().await.unwrap(); - // max_visible_batch_position is 1, so we see batches 0 and 1 (20 rows) + // visible_count is 1, so we see batches 0 and 1 (20 rows) assert_eq!(result.num_rows(), 20); } diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec.rs index 2a593a0f215..ac54e4c5617 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec.rs @@ -36,7 +36,7 @@ pub use vector::VectorIndexExec; pub(super) fn newest_pk_positions( batch_store: &BatchStore, pk_columns: &[String], - max_visible_batch_position: usize, + visible_count: usize, max_visible_row: u64, ) -> DataFusionResult> { let mut newest: HashMap, u64> = HashMap::new(); @@ -46,7 +46,7 @@ pub(super) fn newest_pk_positions( if n == 0 { continue; } - if batch_position > max_visible_batch_position { + if batch_position > visible_count { current_row += n as u64; continue; } diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/brute_force_vector.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/brute_force_vector.rs index e47a715ceaa..f6c870d2c7a 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/brute_force_vector.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/brute_force_vector.rs @@ -47,7 +47,7 @@ const DEFAULT_DISTANCE_TYPE: DistanceType = DistanceType::L2; pub struct MemTableBruteForceVectorExec { batch_store: Arc, query: VectorQuery, - max_visible_batch_position: usize, + visible_count: usize, projection: Option>, output_schema: SchemaRef, properties: Arc, @@ -68,10 +68,7 @@ impl Debug for MemTableBruteForceVectorExec { f.debug_struct("MemTableBruteForceVectorExec") .field("column", &self.query.column) .field("k", &self.query.k) - .field( - "max_visible_batch_position", - &self.max_visible_batch_position, - ) + .field("visible_count", &self.visible_count) .field("with_row_id", &self.with_row_id) .finish() } @@ -84,7 +81,7 @@ impl MemTableBruteForceVectorExec { pub fn new( batch_store: Arc, query: VectorQuery, - max_visible_batch_position: usize, + visible_count: usize, projection: Option>, base_schema: SchemaRef, with_row_id: bool, @@ -110,7 +107,7 @@ impl MemTableBruteForceVectorExec { Ok(Self { batch_store, query, - max_visible_batch_position, + visible_count, projection, output_schema, properties, @@ -159,7 +156,7 @@ impl MemTableBruteForceVectorExec { Ok(Some(mask)) } - /// Last row position visible under `max_visible_batch_position`, or `None` + /// Last row position visible under `visible_count`, or `None` /// if no batches are visible. Identical to `VectorIndexExec`'s helper so /// both arms cut at the same MVCC boundary. fn compute_max_visible_row(&self) -> Option { @@ -168,7 +165,7 @@ impl MemTableBruteForceVectorExec { for (batch_position, stored_batch) in self.batch_store.iter().enumerate() { let batch_end = current_row + stored_batch.num_rows as u64; - if batch_position <= self.max_visible_batch_position { + if batch_position < self.visible_count { max_visible_row_exclusive = batch_end; } current_row = batch_end; @@ -224,7 +221,7 @@ impl MemTableBruteForceVectorExec { newest_pk_positions( &self.batch_store, pk_columns, - self.max_visible_batch_position, + self.visible_count, max_visible_row, ) .map_err(|e| Error::invalid_input(e.to_string()))?, @@ -244,7 +241,7 @@ impl MemTableBruteForceVectorExec { if n == 0 { continue; } - if batch_position > self.max_visible_batch_position { + if batch_position > self.visible_count { current_row += n as u64; continue; } @@ -610,7 +607,7 @@ mod tests { MemTableBruteForceVectorExec::new( store, query, - /* max_visible_batch_position = */ usize::MAX, + /* visible_count = */ usize::MAX, None, schema, false, @@ -663,7 +660,7 @@ mod tests { } #[tokio::test] - async fn respects_max_visible_batch_position() { + async fn respects_indexed_count() { // Two batches of two rows. Freeze at batch 0 — only ids 0,1 are // visible candidates; the (closer) ids 2,3 in batch 1 are excluded. let schema = make_schema(); @@ -673,7 +670,7 @@ mod tests { let query = query_for([0.0, 0.0], 4); let exec = Arc::new( MemTableBruteForceVectorExec::new( - store, query, /* max_visible_batch_position = */ 0, None, schema, false, + store, query, /* visible_count = */ 1, None, schema, false, ) .expect("ctor"), ); diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/btree.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/btree.rs index b9e1e3fa469..01781d07ad8 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/btree.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/btree.rs @@ -30,7 +30,7 @@ pub struct BTreeIndexExec { batch_store: Arc, indexes: Arc, predicate: ScalarPredicate, - max_visible_batch_position: usize, + visible_count: usize, projection: Option>, output_schema: SchemaRef, properties: Arc, @@ -47,10 +47,7 @@ impl Debug for BTreeIndexExec { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("BTreeIndexExec") .field("predicate", &self.predicate) - .field( - "max_visible_batch_position", - &self.max_visible_batch_position, - ) + .field("visible_count", &self.visible_count) .field("with_row_id", &self.with_row_id) .field("with_row_address", &self.with_row_address) .field("column", &self.column) @@ -66,7 +63,7 @@ impl BTreeIndexExec { /// * `batch_store` - Lock-free batch store containing data /// * `indexes` - Index registry with BTree indexes /// * `predicate` - Scalar predicate to apply - /// * `max_visible_batch_position` - MVCC visibility sequence number + /// * `visible_count` - MVCC visibility sequence number /// * `projection` - Optional column indices to project /// * `output_schema` - Schema after projection (should include _rowid/_rowaddr if requested) /// * `with_row_id` - Whether to include _rowid column (row position) @@ -76,7 +73,7 @@ impl BTreeIndexExec { batch_store: Arc, indexes: Arc, predicate: ScalarPredicate, - max_visible_batch_position: usize, + visible_count: usize, projection: Option>, output_schema: SchemaRef, with_row_id: bool, @@ -102,7 +99,7 @@ impl BTreeIndexExec { batch_store, indexes, predicate, - max_visible_batch_position, + visible_count, projection, output_schema, properties, @@ -113,7 +110,7 @@ impl BTreeIndexExec { }) } - /// Compute the maximum visible row position based on max_visible_batch_position. + /// Compute the maximum visible row position based on visible_count. /// Returns None if no batches are visible. fn compute_max_visible_row(&self) -> Option { let mut max_visible_row_exclusive: u64 = 0; @@ -121,7 +118,7 @@ impl BTreeIndexExec { for (batch_position, stored_batch) in self.batch_store.iter().enumerate() { let batch_end = current_row + stored_batch.num_rows as u64; - if batch_position <= self.max_visible_batch_position { + if batch_position < self.visible_count { max_visible_row_exclusive = batch_end; } current_row = batch_end; @@ -429,7 +426,7 @@ mod tests { batch_store, indexes, predicate, - 0, // max_visible_batch_position (batch at position 0) + 1, // visible_count (batch at position 0) None, schema, false, @@ -473,7 +470,7 @@ mod tests { batch_store, indexes, predicate, - 0, + 1, None, schema, false, @@ -518,7 +515,7 @@ mod tests { batch_store.clone(), indexes.clone(), predicate.clone(), - 0, + 1, None, schema.clone(), false, @@ -538,7 +535,7 @@ mod tests { batch_store, indexes, predicate, - 1, + 2, None, schema, false, @@ -587,7 +584,7 @@ mod tests { batch_store, indexes, predicate, - 0, + 1, None, schema_with_rowid.clone(), true, @@ -651,7 +648,7 @@ mod tests { batch_store.clone(), indexes.clone(), predicate.clone(), - 0, + 1, None, schema.clone(), false, @@ -679,7 +676,7 @@ mod tests { batch_store, indexes, predicate, - 0, + 1, None, schema_with_rowid, true, diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/dedup_scan.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/dedup_scan.rs index 071581aa065..66d73becb69 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/dedup_scan.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/dedup_scan.rs @@ -43,7 +43,7 @@ use crate::dataset::mem_wal::write::BatchStore; /// that satisfy the (optional) predicate. See the module doc. pub struct MemTableDedupScanExec { batch_store: Arc, - max_visible_batch_position: usize, + visible_count: usize, /// Column indices to project (into the source schema). projection: Option>, output_schema: SchemaRef, @@ -61,10 +61,7 @@ pub struct MemTableDedupScanExec { impl Debug for MemTableDedupScanExec { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("MemTableDedupScanExec") - .field( - "max_visible_batch_position", - &self.max_visible_batch_position, - ) + .field("visible_count", &self.visible_count) .field("projection", &self.projection) .field("pk_indices", &self.pk_indices) .field("with_row_address", &self.with_row_address) @@ -78,7 +75,7 @@ impl MemTableDedupScanExec { #[allow(clippy::too_many_arguments)] pub fn new( batch_store: Arc, - max_visible_batch_position: usize, + visible_count: usize, projection: Option>, output_schema: SchemaRef, pk_indices: Vec, @@ -96,7 +93,7 @@ impl MemTableDedupScanExec { Self { batch_store, - max_visible_batch_position, + visible_count, projection, output_schema, pk_indices, @@ -183,7 +180,7 @@ impl ExecutionPlan for MemTableDedupScanExec { // back-to-front below. let mut batches = self .batch_store - .visible_batches_with_offsets(self.max_visible_batch_position); + .visible_batches_with_offsets(self.visible_count); batches.reverse(); let projection = self.projection.clone(); @@ -342,7 +339,7 @@ mod tests { /// Run the exec and collect (id -> (value, rowaddr)). async fn run( store: Arc, - max_visible: usize, + visible_count: usize, filter: Option, ) -> HashMap, u64)> { let filter_predicate = filter.map(|expr| { @@ -353,7 +350,7 @@ mod tests { let filter_expr = None; let exec = MemTableDedupScanExec::new( store, - max_visible, + visible_count, None, output_schema(), vec![0], @@ -387,11 +384,11 @@ mod tests { // id=10 inserted (100) then updated to NULL, all in one batch. store.append(batch(&[(10, Some(100)), (10, None)])).unwrap(); - let no_filter = run(store.clone(), 0, None).await; + let no_filter = run(store.clone(), 1, None).await; assert_eq!(no_filter.len(), 1); assert_eq!(no_filter[&10].0, None, "newest version of id=10 is NULL"); - let not_null = run(store, 0, Some(col("value").is_not_null())).await; + let not_null = run(store, 1, Some(col("value").is_not_null())).await; assert!( !not_null.contains_key(&10), "id=10 newest is NULL; the stale value=100 must not leak under value IS NOT NULL" @@ -408,7 +405,7 @@ mod tests { store.append(batch(&[(20, Some(999)), (30, None)])).unwrap(); // No filter: newest per PK = {10:NULL@2, 20:999@3, 30:NULL@4}. - let all = run(store.clone(), 1, None).await; + let all = run(store.clone(), 2, None).await; assert_eq!(all.len(), 3); assert_eq!(all[&10], (None, 2)); assert_eq!(all[&20], (Some(999), 3)); @@ -416,7 +413,7 @@ mod tests { // value IS NOT NULL: only id=20 (newest 999) survives; 10 and 30 are // newest-NULL so they must be absent (no stale leak). - let not_null = run(store, 1, Some(col("value").is_not_null())).await; + let not_null = run(store, 2, Some(col("value").is_not_null())).await; assert_eq!(not_null.len(), 1); assert_eq!(not_null[&20], (Some(999), 3)); } @@ -429,7 +426,7 @@ mod tests { // id=40 inserted NULL then updated to 400 (newest non-NULL). store.append(batch(&[(40, None), (40, Some(400))])).unwrap(); - let is_null = run(store, 0, Some(col("value").is_null())).await; + let is_null = run(store, 1, Some(col("value").is_null())).await; assert!( !is_null.contains_key(&40), "id=40 newest is 400; the stale NULL must not leak under value IS NULL" diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/fts.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/fts.rs index 77c1a27b980..e35ee6d1b6a 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/fts.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/fts.rs @@ -47,14 +47,14 @@ pub struct FtsIndexExec { batch_store: Arc, indexes: Arc, query: FtsQuery, - max_visible_batch_position: usize, + visible_count: usize, projection: Option>, output_schema: SchemaRef, properties: Arc, metrics: ExecutionPlanMetricsSet, /// Pre-computed batch ranges for O(log n) lookup. batch_ranges: Vec, - /// Maximum visible row position based on max_visible_batch_position (None if nothing visible). + /// Maximum visible row position based on visible_count (None if nothing visible). max_visible_row: Option, /// Whether to include _rowid column (row position) in output. with_row_id: bool, @@ -72,10 +72,7 @@ impl Debug for FtsIndexExec { f.debug_struct("FtsIndexExec") .field("column", &self.query.column) .field("query_type", &self.query.query_type) - .field( - "max_visible_batch_position", - &self.max_visible_batch_position, - ) + .field("visible_count", &self.visible_count) .field("with_row_id", &self.with_row_id) .finish() } @@ -89,7 +86,7 @@ impl FtsIndexExec { /// * `batch_store` - Lock-free batch store containing data /// * `indexes` - Index registry with FTS indexes /// * `query` - FTS query parameters - /// * `max_visible_batch_position` - MVCC visibility sequence number + /// * `visible_count` - MVCC visibility sequence number /// * `projection` - Optional column indices to project /// * `base_schema` - Schema before adding score column (and _rowid if with_row_id) /// * `with_row_id` - Whether to include _rowid column (row position) @@ -97,7 +94,7 @@ impl FtsIndexExec { batch_store: Arc, indexes: Arc, query: FtsQuery, - max_visible_batch_position: usize, + visible_count: usize, projection: Option>, base_schema: SchemaRef, with_row_id: bool, @@ -148,7 +145,7 @@ impl FtsIndexExec { end: batch_end, batch_id, }); - if batch_id <= max_visible_batch_position { + if batch_id < visible_count { max_visible_row_exclusive = batch_end as u64; } current_row = batch_end; @@ -165,7 +162,7 @@ impl FtsIndexExec { batch_store, indexes, query, - max_visible_batch_position, + visible_count, projection, output_schema, properties, @@ -458,7 +455,7 @@ impl FtsIndexExec { Some(newest_pk_positions( &self.batch_store, pk_columns, - self.max_visible_batch_position, + self.visible_count, max_visible_row, )?) }; @@ -641,7 +638,7 @@ mod tests { let query = FtsQuery::match_query("text", "hello"); - let exec = FtsIndexExec::new(batch_store, indexes, query, 0, None, schema, false).unwrap(); + let exec = FtsIndexExec::new(batch_store, indexes, query, 1, None, schema, false).unwrap(); let ctx = Arc::new(TaskContext::default()); let stream = exec.execute(0, ctx).unwrap(); @@ -682,7 +679,7 @@ mod tests { batch_store.clone(), indexes.clone(), query.clone(), - 0, + 1, None, schema.clone(), false, @@ -697,7 +694,7 @@ mod tests { assert_eq!(total_rows, 2); // "hello" in batch1 docs 0 and 2 // Query with max_visible=1 should see both batches - let exec = FtsIndexExec::new(batch_store, indexes, query, 1, None, schema, false).unwrap(); + let exec = FtsIndexExec::new(batch_store, indexes, query, 2, None, schema, false).unwrap(); let ctx = Arc::new(TaskContext::default()); let stream = exec.execute(0, ctx).unwrap(); diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/scan.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/scan.rs index de2f2f05d67..6126f392f1f 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/scan.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/scan.rs @@ -30,12 +30,12 @@ pub const ROW_ADDRESS_COLUMN: &str = "_rowaddr"; /// ExecutionPlan node that scans all visible batches from a MemTable. /// /// This node implements visibility filtering, returning only batches -/// where `batch_position <= max_visible_batch_position`. +/// where `batch_position <= visible_count`. /// /// Supports filter pushdown for efficient predicate evaluation during scan. pub struct MemTableScanExec { batch_store: Arc, - max_visible_batch_position: usize, + visible_count: usize, projection: Option>, output_schema: SchemaRef, /// Schema of the source data (before projection), used for filter evaluation. @@ -55,10 +55,7 @@ pub struct MemTableScanExec { impl Debug for MemTableScanExec { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("MemTableScanExec") - .field( - "max_visible_batch_position", - &self.max_visible_batch_position, - ) + .field("visible_count", &self.visible_count) .field("projection", &self.projection) .field("with_row_id", &self.with_row_id) .field("with_row_address", &self.with_row_address) @@ -73,20 +70,20 @@ impl MemTableScanExec { /// # Arguments /// /// * `batch_store` - Lock-free batch store containing data - /// * `max_visible_batch_position` - Maximum batch position visible (inclusive) + /// * `visible_count` - Maximum batch position visible (inclusive) /// * `projection` - Optional column indices to project /// * `output_schema` - Schema after projection (should include _rowid/_rowaddr if requested) /// * `with_row_id` - Whether to include _rowid column (row position) pub fn new( batch_store: Arc, - max_visible_batch_position: usize, + visible_count: usize, projection: Option>, output_schema: SchemaRef, with_row_id: bool, ) -> Self { Self::with_filter( batch_store, - max_visible_batch_position, + visible_count, projection, output_schema.clone(), output_schema, @@ -102,7 +99,7 @@ impl MemTableScanExec { /// # Arguments /// /// * `batch_store` - Lock-free batch store containing data - /// * `max_visible_batch_position` - Maximum batch position visible (inclusive) + /// * `visible_count` - Maximum batch position visible (inclusive) /// * `projection` - Optional column indices to project /// * `output_schema` - Schema after projection (should include _rowid/_rowaddr if requested) /// * `source_schema` - Schema of source data (before projection), used for filter evaluation @@ -113,7 +110,7 @@ impl MemTableScanExec { #[allow(clippy::too_many_arguments)] pub fn with_filter( batch_store: Arc, - max_visible_batch_position: usize, + visible_count: usize, projection: Option>, output_schema: SchemaRef, source_schema: SchemaRef, @@ -131,7 +128,7 @@ impl MemTableScanExec { Self { batch_store, - max_visible_batch_position, + visible_count, projection, output_schema, source_schema, @@ -221,7 +218,7 @@ impl ExecutionPlan for MemTableScanExec { // Get visible batches with their row offsets let batches_with_offsets = self .batch_store - .visible_batches_with_offsets(self.max_visible_batch_position); + .visible_batches_with_offsets(self.visible_count); let projection = self.projection.clone(); let schema = self.output_schema.clone(); @@ -394,7 +391,7 @@ mod tests { batch_store.append(batch).unwrap(); // Batch is at position 0, max_visible=0 means position 0 is visible - let exec = MemTableScanExec::new(batch_store, 0, None, schema, false); + let exec = MemTableScanExec::new(batch_store, 1, None, schema, false); let ctx = Arc::new(TaskContext::default()); let stream = exec.execute(0, ctx).unwrap(); @@ -420,8 +417,8 @@ mod tests { .append(create_test_batch(&schema, 20, 10)) .unwrap(); - // max_visible_batch_position=1 means positions 0 and 1 are visible (2 batches) - let exec = MemTableScanExec::new(batch_store.clone(), 1, None, schema.clone(), false); + // visible_count=1 means positions 0 and 1 are visible (2 batches) + let exec = MemTableScanExec::new(batch_store.clone(), 2, None, schema.clone(), false); let ctx = Arc::new(TaskContext::default()); let stream = exec.execute(0, ctx).unwrap(); let batches: Vec = stream.try_collect().await.unwrap(); @@ -442,7 +439,7 @@ mod tests { // Project only "id" column (index 0) let projected_schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); - let exec = MemTableScanExec::new(batch_store, 0, Some(vec![0]), projected_schema, false); + let exec = MemTableScanExec::new(batch_store, 1, Some(vec![0]), projected_schema, false); let ctx = Arc::new(TaskContext::default()); let stream = exec.execute(0, ctx).unwrap(); @@ -459,7 +456,7 @@ mod tests { let batch_store = Arc::new(BatchStore::with_capacity(100)); // Empty store with max_visible=0 should return no batches - let exec = MemTableScanExec::new(batch_store, 0, None, schema, false); + let exec = MemTableScanExec::new(batch_store, 1, None, schema, false); let ctx = Arc::new(TaskContext::default()); let stream = exec.execute(0, ctx).unwrap(); @@ -481,7 +478,7 @@ mod tests { .unwrap(); // max_visible=1 means positions 0 and 1 are visible - let exec = MemTableScanExec::new(batch_store, 1, None, schema, false); + let exec = MemTableScanExec::new(batch_store, 2, None, schema, false); let stats = exec.partition_statistics(None).unwrap(); // Statistics are Absent to avoid DataFusion analysis bugs @@ -508,7 +505,7 @@ mod tests { Field::new("_rowid", DataType::UInt64, true), ])); - let exec = MemTableScanExec::new(batch_store, 1, None, schema_with_rowid, true); + let exec = MemTableScanExec::new(batch_store, 2, None, schema_with_rowid, true); let ctx = Arc::new(TaskContext::default()); let stream = exec.execute(0, ctx).unwrap(); diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/vector.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/vector.rs index 58d4250f0c2..0c8f4c27ff9 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/vector.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/vector.rs @@ -34,7 +34,7 @@ pub struct VectorIndexExec { batch_store: Arc, indexes: Arc, query: VectorQuery, - max_visible_batch_position: usize, + visible_count: usize, projection: Option>, output_schema: SchemaRef, properties: Arc, @@ -58,10 +58,7 @@ impl Debug for VectorIndexExec { if let Some(metric) = &self.query.distance_type { debug.field("distance_type", metric); } - debug.field( - "max_visible_batch_position", - &self.max_visible_batch_position, - ); + debug.field("visible_count", &self.visible_count); debug.field("with_row_id", &self.with_row_id); debug.finish() } @@ -75,7 +72,7 @@ impl VectorIndexExec { /// * `batch_store` - Lock-free batch store containing data /// * `indexes` - Index registry with HNSW vector indexes /// * `query` - Vector query parameters - /// * `max_visible_batch_position` - MVCC visibility sequence number + /// * `visible_count` - MVCC visibility sequence number /// * `projection` - Optional column indices to project /// * `base_schema` - Schema after projection (will add _distance column, and _rowid if with_row_id) /// * `with_row_id` - Whether to include _rowid column (row position) @@ -83,7 +80,7 @@ impl VectorIndexExec { batch_store: Arc, indexes: Arc, query: VectorQuery, - max_visible_batch_position: usize, + visible_count: usize, projection: Option>, base_schema: SchemaRef, with_row_id: bool, @@ -119,7 +116,7 @@ impl VectorIndexExec { batch_store, indexes, query, - max_visible_batch_position, + visible_count, projection, output_schema, properties, @@ -128,9 +125,9 @@ impl VectorIndexExec { }) } - /// Compute the maximum visible row position based on max_visible_batch_position. + /// Compute the maximum visible row position based on visible_count. /// - /// Returns the last row position that is visible at the given max_visible_batch_position, + /// Returns the last row position that is visible at the given visible_count, /// or None if no batches are visible. fn compute_max_visible_row(&self) -> Option { let mut max_visible_row_exclusive: u64 = 0; @@ -138,7 +135,7 @@ impl VectorIndexExec { for (batch_position, stored_batch) in self.batch_store.iter().enumerate() { let batch_end = current_row + stored_batch.num_rows as u64; - if batch_position <= self.max_visible_batch_position { + if batch_position < self.visible_count { max_visible_row_exclusive = batch_end; } current_row = batch_end; diff --git a/rust/lance/src/dataset/mem_wal/scanner/block_list.rs b/rust/lance/src/dataset/mem_wal/scanner/block_list.rs index 30a19650005..e74ebda0f63 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/block_list.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/block_list.rs @@ -328,7 +328,7 @@ fn in_memory_membership( batch_store: &Arc, index_store: &Arc, ) -> GenMembership { - let max_visible_row = batch_store.max_visible_row(index_store.max_visible_batch_position()); + let max_visible_row = batch_store.max_visible_row(index_store.visible_count()); GenMembership::InMemory { index_store: index_store.clone(), max_visible_row, @@ -346,9 +346,10 @@ fn bounded_in_memory_membership( index_store: &Arc, batch_count: u64, ) -> GenMembership { - let max_visible_row = batch_count - .checked_sub(1) - .and_then(|last_batch| batch_store.max_visible_row(last_batch as usize)); + // `batch_count` is already an exclusive count, and so is what + // `max_visible_row` takes, so it passes straight through: no + // count-to-inclusive-position conversion, and `0` yields `None` on its own. + let max_visible_row = batch_store.max_visible_row(batch_count as usize); GenMembership::InMemory { index_store: index_store.clone(), max_visible_row, diff --git a/rust/lance/src/dataset/mem_wal/scanner/exec/pk_block_filter.rs b/rust/lance/src/dataset/mem_wal/scanner/exec/pk_block_filter.rs index 4903f0de3a8..b81a470a882 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/exec/pk_block_filter.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/exec/pk_block_filter.rs @@ -317,7 +317,7 @@ mod tests { let (bp, off, _) = store.append(b.clone()).unwrap(); index.insert_with_batch_position(&b, off, Some(bp)).unwrap(); } - let max_visible_row = store.max_visible_row(index.max_visible_batch_position()); + let max_visible_row = store.max_visible_row(index.visible_count()); GenMembership::InMemory { index_store: Arc::new(index), max_visible_row, diff --git a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs index 60000666f52..346cee851f4 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/fts_search.rs @@ -98,7 +98,7 @@ fn active_source_can_execute_fts(source: &LsmDataSource, column: &str) -> bool { .get_fts_by_column(column) .is_some_and(|index| !index.is_empty()) && batch_store - .max_visible_row(index_store.max_visible_batch_position()) + .max_visible_row(index_store.visible_count()) .is_some() } _ => false, diff --git a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs index d80e0e26795..6270b84c4a0 100644 --- a/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs +++ b/rust/lance/src/dataset/mem_wal/scanner/point_lookup.rs @@ -891,7 +891,12 @@ fn probe_position( if len == 0 { return Ok(ProbePos::Miss); } - let last_visible_idx = index_store.max_visible_batch_position().min(len - 1); + // The cursor is an exclusive count, so the last visible batch sits at + // `count - 1`. A count of 0 means nothing is visible yet — not "batch 0". + let visible_count = index_store.visible_count().min(len); + let Some(last_visible_idx) = visible_count.checked_sub(1) else { + return Ok(ProbePos::Miss); + }; let last = batch_store.get(last_visible_idx).ok_or_else(|| { lance_core::Error::internal("point-lookup: visible batch index out of range") })?; @@ -1388,8 +1393,9 @@ mod tests { let batch_store = Arc::new(BatchStore::with_capacity(16)); let mut index_store = IndexStore::new(); - // BTree on the PK so that `max_visible_batch_position` advances as - // we insert, otherwise the scanner sees no batches at all. + // BTree on the PK: the point lookup resolves keys through the indexed PK + // path, which this exercises. (`indexed_count`/`visible_count` advance + // from the batch position regardless of whether any index is configured.) index_store.add_btree("id_idx".to_string(), 0, "id".to_string()); // Two writes to pk=1, then an unrelated pk=2. The "new" row goes diff --git a/rust/lance/src/dataset/mem_wal/wal.rs b/rust/lance/src/dataset/mem_wal/wal.rs index b4caff180e1..4082cd7ece9 100644 --- a/rust/lance/src/dataset/mem_wal/wal.rs +++ b/rust/lance/src/dataset/mem_wal/wal.rs @@ -9,7 +9,8 @@ use std::io::Cursor; use std::sync::Arc; use std::sync::Mutex as StdMutex; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::MutexGuard as StdMutexGuard; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::time::Instant; use std::time::Duration; @@ -87,68 +88,200 @@ impl WalFlushFailure { } } -/// Watcher for batch durability using watermark-based tracking. +/// The writer's cursors, shared by the WAL-append task, the index-apply task, +/// every memtable's `IndexStore`, and every `put` waiting to become visible. /// -/// Uses a shared watch channel that broadcasts the durable watermark. -/// The watcher waits until the watermark reaches or exceeds its target batch ID. -#[derive(Clone)] -pub struct BatchDurableWatcher { - /// Watch receiver for the durable watermark. - rx: watch::Receiver, - /// Target batch ID to wait for. - target_batch_position: usize, - /// Terminal flush failure shared with the flusher. When set, the watermark - /// can never reach the target, so `wait` returns this typed error instead of - /// blocking forever. +/// Two cursors are stored, one view is derived: +/// +/// - `durable` — writer-global count of WAL-durable batches, advanced by the +/// WAL-append task. Exclusive; 0 means none. +/// - `indexed` — per-memtable count of indexed batches, advanced by the +/// index-apply task. It lives on the memtable's own `IndexStore`, not here, +/// because each memtable has its own indexes. +/// - `visible` — **derived, never stored**: a batch is visible once it is +/// indexed and, under `durable_write`, also durable. +/// +/// Deriving `visible` rather than caching it is deliberate. A cached +/// `min(indexed, durable)` recomputed by two independent tasks is the classic +/// store-buffer race: with `Release`/`Acquire` each task can read the other's +/// *pre-store* value, so both compute a minimum below the true one, and a +/// max-clamped publish then leaves the cached value permanently short. A `put` +/// blocked on it would hang until some unrelated write happened to move a cursor +/// again. With nothing cached there is nothing to leave stale — `notify` is a +/// bare wake-up and every waiter recomputes from the cursors themselves. +pub struct WriterCursors { + durable: AtomicUsize, + /// Bumped whenever a cursor advances or the writer poisons. Carries no + /// value; it exists only to wake waiters, which then recompute. + notify_tx: watch::Sender, + notify_rx: watch::Receiver, + /// First terminal failure. Shared so a poisoned writer wakes every waiter + /// with the typed error, instead of leaving it blocked on a cursor that can + /// never advance again. terminal_error: Arc>>, + /// Whether durability is part of visibility. Per-writer, not per-write: + /// `visible` is a writer-wide definition, so a mix would need two visibility + /// views over one memtable. + durable_write: bool, +} + +impl WriterCursors { + pub fn new(durable_write: bool) -> Self { + let (notify_tx, notify_rx) = watch::channel(0); + Self { + durable: AtomicUsize::new(0), + notify_tx, + notify_rx, + terminal_error: Arc::new(StdMutex::new(None)), + durable_write, + } + } + + /// Writer-global count of WAL-durable batches. + pub fn durable(&self) -> usize { + self.durable.load(Ordering::Acquire) + } + + pub fn durable_write(&self) -> bool { + self.durable_write + } + + /// Advance the durability cursor. Monotonic, so an out-of-order completion + /// can never walk it backwards. + pub(crate) fn advance_durable(&self, global_count: usize) { + self.durable.fetch_max(global_count, Ordering::AcqRel); + self.wake(); + } + + /// Wake every waiter so it recomputes. Called after any cursor advances, and + /// after the writer poisons. + pub(crate) fn wake(&self) { + self.notify_tx.send_modify(|version| *version += 1); + } + + /// The visible prefix of one memtable, given its indexed prefix and its + /// writer-global coordinate. + pub fn visible_count(&self, indexed_count: usize, global_offset: usize) -> usize { + if !self.durable_write { + return indexed_count; + } + indexed_count.min(self.durable().saturating_sub(global_offset)) + } + + /// Lock `terminal_error`, ignoring mutex poisoning. + /// + /// A panic under this lock cannot tear the `Option` it guards: the sole + /// writer builds the value first and assigns it whole, so a panic mid-section + /// leaves the slot exactly as it was. Surfacing poison as an error instead + /// would mask the latched `FenceReason` behind an unrelated "mutex poisoned" + /// precisely when a caller needs the real reason — and would strand + /// `mark_terminal_failure`, which has no error to return. + fn lock_terminal_error(&self) -> StdMutexGuard<'_, Option> { + self.terminal_error + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + pub(crate) fn check_poisoned(&self) -> Result<()> { + if let Some(failure) = self.lock_terminal_error().clone() { + return Err(failure.into_error()); + } + Ok(()) + } + + pub(crate) fn mark_terminal_failure(&self, error: &Error) { + { + let mut slot = self.lock_terminal_error(); + if slot.is_none() { + *slot = Some(WalFlushFailure::from_error(error)); + } + } + // Wake waiters without advancing anything: each re-checks `terminal_error` + // and returns the error rather than blocking on a cursor that can no + // longer move. + self.wake(); + } +} + +/// Blocks a `put` until its batches are visible: indexed, and — in durable mode +/// — WAL-durable too. +/// +/// Recomputes the condition on every wake instead of comparing against a cached +/// watermark. `notify` only says "something moved"; this decides what that means. +pub struct BatchDurableWatcher { + cursors: Arc, + rx: watch::Receiver, + /// The memtable the batches landed in; its `indexed_count` is the apply + /// cursor being waited on. `None` in WAL-only mode, which has no indexes. + indexes: Option>, + /// Local exclusive count this write needs indexed. + target_indexed: usize, + /// Writer-global exclusive count this write needs durable. Batch positions + /// restart at 0 in every memtable while the durability cursor spans the + /// writer's whole life, so the caller must globalize this. + target_durable: usize, } impl BatchDurableWatcher { - /// Create a new watcher for a specific batch ID. pub fn new( - rx: watch::Receiver, - target_batch_position: usize, - terminal_error: Arc>>, + cursors: Arc, + indexes: Option>, + target_indexed: usize, + target_durable: usize, ) -> Self { + let rx = cursors.notify_rx.clone(); Self { + cursors, rx, - target_batch_position, - terminal_error, + indexes, + target_indexed, + target_durable, } } - /// Wait until the batch is durable. - /// - /// Returns Ok(()) when `durable_watermark >= target_batch_position`, or - /// Err if a terminal flush failure (e.g. a fence) means the watermark can - /// never reach the target. + /// Whether the write is readable yet. + fn is_visible(&self) -> bool { + // WAL-only mode has no indexes, so there is nothing to index-wait on. + let indexed = match &self.indexes { + Some(indexes) => indexes.indexed_count(), + None => self.target_indexed, + }; + if indexed < self.target_indexed { + return false; + } + !self.cursors.durable_write() || self.cursors.durable() >= self.target_durable + } + + /// Wait until the write is visible, or until the writer poisons — in which + /// case no cursor will ever reach the target, so surface the typed error + /// rather than blocking forever. pub async fn wait(&mut self) -> Result<()> { loop { - if let Some(failure) = self.terminal_error.lock().unwrap().clone() { - return Err(failure.into_error()); - } - let current = *self.rx.borrow(); - if current >= self.target_batch_position { + // Mark the current version seen *before* testing, so a wake-up landing + // between the test and `changed()` below is not lost. + self.rx.borrow_and_update(); + self.cursors.check_poisoned()?; + if self.is_visible() { return Ok(()); } self.rx .changed() .await - .map_err(|_| Error::io("Durable watermark channel closed"))?; + .map_err(|_| Error::io("Writer cursor channel closed"))?; } } - /// Check if the batch is already durable (non-blocking). + /// Non-blocking check. pub fn is_durable(&self) -> bool { - *self.rx.borrow() >= self.target_batch_position + self.is_visible() } } impl std::fmt::Debug for BatchDurableWatcher { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("BatchDurableWatcher") - .field("target_batch_position", &self.target_batch_position) - .field("current_watermark", &*self.rx.borrow()) + .field("target_indexed", &self.target_indexed) + .field("target_durable", &self.target_durable) .finish() } } @@ -164,20 +297,15 @@ pub struct WalEntry { pub num_batches: usize, } -/// Result of a parallel WAL flush with index update. +/// Result of a WAL flush. Append-only: index application runs on its own task +/// (see `apply_index_range`), which records its own stats, so this no longer +/// carries index-update timing or row counts. #[derive(Debug, Clone)] pub struct WalFlushResult { /// WAL entry that was written (if any). pub entry: Option, /// Duration of WAL I/O operation. pub wal_io_duration: std::time::Duration, - /// Overall wall-clock duration of the index update operation. - /// This includes any overhead from thread scheduling and context switching. - pub index_update_duration: std::time::Duration, - /// Per-index update durations. Key is index name, value is duration. - pub index_update_duration_breakdown: std::collections::HashMap, - /// Number of rows indexed. - pub rows_indexed: usize, /// Size of WAL data written in bytes. pub wal_bytes: usize, } @@ -185,32 +313,142 @@ pub struct WalFlushResult { /// Source for a WAL flush — either a `BatchStore` range (MemTable mode) or /// a drainable in-memory pending queue (WAL-only mode). pub enum WalFlushSource { - /// MemTable mode: read a `[max_flushed+1, end_batch_position)` range - /// from a `BatchStore`. Indexes are updated in parallel with the WAL - /// append. - BatchStore { - batch_store: Arc, - indexes: Option>, - }, + /// MemTable mode: append the `[durable, end_batch_position)` range of a + /// `BatchStore` to the WAL. Append-only — the index apply runs on its own + /// task, so a failed append can no longer publish rows through it. + BatchStore { batch_store: Arc }, + /// Timer-driven: append whichever store still owes the WAL an append, + /// resolved when the message is *handled*, not when it is enqueued. + /// + /// Carries no store on purpose. `MessageFactory` is synchronous and cannot + /// take the async state lock, and capturing an `Arc` once at + /// handler construction would pin the first memtable forever. + /// + /// Resolution walks the live stores **oldest first** and takes the first with + /// `global_end() > durable`. It must not simply take "the active memtable": + /// a tick enqueued before a freeze is handled *after* it, would resolve to + /// the new memtable, and would append its batches ahead of the outgoing + /// memtable's tail. WAL entry positions are assigned in append-call order, + /// replay walks them ascending, row positions follow, and primary-key recency + /// is "newest visible row position wins" — so an out-of-order append silently + /// inverts dedup after a crash: the stale row wins. A full scan looks fine. + NextPending, /// WAL-only mode: drain all pending batches from the shared /// `WalOnlyState`. There are no in-memory indexes to update. WalOnly { state: Arc }, } impl WalFlushSource { - fn pending_count(&self) -> usize { + fn kind(&self) -> &'static str { match self { - Self::BatchStore { batch_store, .. } => batch_store.pending_wal_flush_count(), - Self::WalOnly { state } => state - .pending - .lock() - .ok() - .map(|p| p.batches.len()) - .unwrap_or(0), + Self::BatchStore { .. } => "BatchStore", + Self::WalOnly { .. } => "WalOnly", + Self::NextPending => "NextPending", } } } +/// Message to trigger an index apply. +/// +/// Carries the store the batches actually landed in, captured on the put path +/// *before* any freeze can rotate the memtable — pairing a new store with an old +/// store's end position would index the wrong range. +#[derive(Clone)] +pub struct TriggerIndexApply { + pub batch_store: Arc, + pub indexes: Arc, + /// Local exclusive end of the range to cover. + pub end_batch_position: usize, +} + +impl std::fmt::Debug for TriggerIndexApply { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TriggerIndexApply") + .field("end_batch_position", &self.end_batch_position) + .finish() + } +} + +/// What an [`apply_index_range`] call actually indexed. `rows_indexed == 0` +/// marks a routine coalesced no-op (the range was already covered) that must +/// not be recorded as an index update. +#[derive(Debug, Default, Clone, Copy)] +pub struct IndexApplyStats { + pub rows_indexed: usize, + /// Wall-clock time spent applying the range, including the thread-scheduling + /// overhead of the blocking hand-off. + pub duration: std::time::Duration, +} + +/// Apply a contiguous range of batches to a memtable's in-memory indexes. +/// +/// Runs on its own task, as the single sequential consumer of its own channel. +/// Being a single consumer is what makes it safe: it guarantees in-order, +/// contiguous ranges, which is exactly what `HnswGraph::insert_batch` requires — +/// it hard-rejects any range whose start is not `indexed_len`. Ordering comes +/// from the task, not from the flush interval, so triggering per-put is exactly +/// as safe as triggering on a timer. +/// +/// It has its own channel rather than sharing the WAL flusher's, because +/// `TaskDispatcher::run` awaits `handle()` inline: a shared channel would put a +/// ~100ms S3 PUT in front of every latency-sensitive index apply. +pub async fn apply_index_range( + cursors: &Arc, + message: TriggerIndexApply, +) -> Result { + let TriggerIndexApply { + batch_store, + indexes, + end_batch_position, + } = message; + + // Self-batching: a message handled while more puts queue behind it covers + // everything committed so far, so the ones behind it find their range already + // applied. Redundant messages are therefore *routine*, not exceptional, and + // must be a clean no-op — never a call into `insert_batches`, whose HNSW arm + // hard-rejects a non-contiguous start, and which is now terminal for the + // writer. + let start = indexes.indexed_count(); + if end_batch_position <= start { + return Ok(IndexApplyStats::default()); + } + + // Every position in the range must exist. `get` returns `None` only for a + // position past `committed_len`, so a hole means the caller asked to index a + // batch the store never committed. Silently skipping it would let + // `insert_batches` advance `indexed_count` past a never-indexed batch — + // rows counted visible but absent from every index. Fail loudly; the handler + // poisons and reopen rebuilds the indexes from the WAL. + let stored: Vec = (start..end_batch_position) + .map(|position| { + batch_store.get(position).cloned().ok_or_else(|| { + Error::internal(format!( + "index apply range [{start}, {end_batch_position}) is missing batch \ + position {position}; batch_store committed_len is {}", + batch_store.len() + )) + }) + }) + .collect::>()?; + + // `insert_batches` advances `indexed_count` itself, once every index has + // taken the batch. Time the whole hand-off so the recorded latency reflects + // the blocking-pool scheduling too, matching the old inline-flush measurement. + let rows_indexed: usize = stored.iter().map(|b| b.num_rows).sum(); + let apply_start = Instant::now(); + tokio::task::spawn_blocking(move || indexes.insert_batches(&stored)) + .await + .map_err(|e| Error::internal(format!("Index apply task panicked: {e}")))??; + let duration = apply_start.elapsed(); + + // Wake anything waiting to become visible. + cursors.wake(); + Ok(IndexApplyStats { + rows_indexed, + duration, + }) +} + /// Message to trigger a WAL flush. /// /// Carries a `source` describing where to read batches from (BatchStore range @@ -232,7 +470,7 @@ pub struct TriggerWalFlush { impl std::fmt::Debug for TriggerWalFlush { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("TriggerWalFlush") - .field("pending_batches", &self.source.pending_count()) + .field("source", &self.source.kind()) .field("end_batch_position", &self.end_batch_position) .finish() } @@ -361,11 +599,9 @@ impl WalOnlyState { /// shared `WalAppender`. The flusher delegates the actual WAL write to the /// appender, optionally running a parallel index update in MemTable mode. pub struct WalFlusher { - /// Watch channel sender for durable watermark. - /// Broadcasts the highest batch_position that is now durable. - durable_watermark_tx: watch::Sender, - /// Watch channel receiver for creating new watchers. - durable_watermark_rx: watch::Receiver, + /// The writer's cursors. Shared with the index-apply task and every + /// memtable's `IndexStore`, so all three agree on what is visible. + cursors: Arc, /// Underlying WAL append primitive — owns object store, epoch, and /// position discovery. wal_appender: Arc, @@ -377,35 +613,37 @@ pub struct WalFlusher { /// Created at construction and recreated after each flush. /// Used by backpressure to wait for WAL flushes. wal_flush_cell: std::sync::Mutex>>, - /// First terminal flush failure, shared with every `BatchDurableWatcher`. It - /// wakes durability waiters (the watermark never advances) and is read by - /// `check_poisoned` so the write path fails fast. - terminal_error: Arc>>, } impl WalFlusher { /// Create a new WAL flusher backed by an existing `WalAppender`. /// /// The appender owns object store, epoch, and position state. The - /// flusher adds the durability watermark, trigger channel, and - /// completion cell on top. + /// flusher adds the trigger channel and completion cell on top, and shares + /// the writer's cursors. pub fn new(wal_appender: Arc) -> Self { + // Defaults to durable visibility; the writer replaces this with cursors + // built from its own config. + Self::with_cursors(wal_appender, Arc::new(WriterCursors::new(true))) + } + + pub fn with_cursors(wal_appender: Arc, cursors: Arc) -> Self { let shard_id = wal_appender.shard_id(); - // Initialize durable watermark at 0 (no batches durable yet) - let (durable_watermark_tx, durable_watermark_rx) = watch::channel(0); - // Create initial WAL flush cell for backpressure let wal_flush_cell = WatchableOnceCell::new(); Self { - durable_watermark_tx, - durable_watermark_rx, + cursors, wal_appender, shard_id, flush_tx: None, wal_flush_cell: std::sync::Mutex::new(Some(wal_flush_cell)), - terminal_error: Arc::new(StdMutex::new(None)), } } + /// The writer's cursors, for the index-apply task and the memtables. + pub fn cursors(&self) -> &Arc { + &self.cursors + } + /// Set the flush channel for background flush handler. pub fn set_flush_channel(&mut self, tx: mpsc::UnboundedSender) { self.flush_tx = Some(tx); @@ -420,45 +658,65 @@ impl WalFlusher { /// /// Returns a `BatchDurableWatcher` that can be awaited for durability. /// The actual batch data is stored in the BatchStore. - pub fn track_batch(&self, batch_position: usize) -> BatchDurableWatcher { - // Return a watcher that waits for this batch to become durable - // batch_position is 0-indexed, so we wait for watermark > batch_position (i.e., >= batch_position + 1) + /// Watch a write until it becomes visible. + /// + /// `target_indexed` is a **memtable-local** exclusive count; `target_durable` + /// is a **writer-global** one. They are different coordinate spaces on + /// purpose: batch positions restart at 0 in every memtable, while the + /// durability cursor spans the writer's whole life. A local durable target + /// would already be satisfied by a *previous* memtable's appends, and would + /// ack a write that never reached the WAL. Callers globalize via + /// `BatchStore::global_offset`. + pub fn track_batch( + &self, + indexes: Option>, + target_indexed: usize, + target_durable: usize, + ) -> BatchDurableWatcher { BatchDurableWatcher::new( - self.durable_watermark_rx.clone(), - batch_position + 1, - Arc::clone(&self.terminal_error), + Arc::clone(&self.cursors), + indexes, + target_indexed, + target_durable, ) } - /// Latch a terminal flush failure and wake every durability waiter (the - /// watermark never advances, so they must observe the error, not block). + /// The writer-global WAL durability cursor: how many batches of this + /// writer's batch sequence are durable. Exclusive count; 0 means none. + pub fn durable(&self) -> usize { + self.cursors.durable() + } + + /// Advance the durability cursor and wake waiters. + pub(crate) fn advance_durable(&self, global_count: usize) { + self.cursors.advance_durable(global_count); + } + + /// Latch a terminal flush failure and wake every waiter (no cursor will + /// advance again, so they must observe the error rather than block). /// Idempotent: only the first failure is retained. fn mark_terminal_failure(&self, error: &Error) { - { - let mut slot = self.terminal_error.lock().unwrap(); - if slot.is_none() { - *slot = Some(WalFlushFailure::from_error(error)); - } - } - // Wake `wait`ers without advancing the watermark; each re-checks - // `terminal_error` and returns the error. - self.durable_watermark_tx.send_modify(|_| {}); + self.cursors.mark_terminal_failure(error); + } + + /// Latch a terminal failure from outside the flush path (the index-apply + /// task). Same effect: reads and writes fail fast, waiters wake with the + /// typed error, and recovery is reopen -> replay. + pub(crate) fn poison(&self, error: &Error) { + self.cursors.mark_terminal_failure(error); } /// Fail fast with the typed error if this writer has been fenced (by a peer - /// or its own persistence failure). The write path calls this before touching - /// the memtable so a poisoned writer can't diverge further. Recovery is to - /// reopen the shard (replay the WAL). + /// or its own persistence failure). Both the read and write paths call this + /// so a poisoned writer can neither diverge further nor serve a snapshot + /// that replay will not reproduce. Recovery is to reopen and replay. pub fn check_poisoned(&self) -> Result<()> { - if let Some(failure) = self.terminal_error.lock().unwrap().clone() { - return Err(failure.into_error()); - } - Ok(()) + self.cursors.check_poisoned() } /// Get the current durable watermark. pub fn durable_watermark(&self) -> usize { - *self.durable_watermark_rx.borrow() + self.cursors.durable() } /// Get a watcher for WAL flush completion. @@ -527,14 +785,16 @@ impl WalFlusher { end_batch_position: usize, ) -> Result { let result = match source { - WalFlushSource::BatchStore { - batch_store, - indexes, - } => { - self.flush_from_batch_store(batch_store, indexes.clone(), end_batch_position) + WalFlushSource::BatchStore { batch_store } => { + self.flush_from_batch_store(batch_store, end_batch_position) .await } WalFlushSource::WalOnly { state } => self.flush_from_wal_only(state).await, + // The handler resolves a tick to a concrete store before it gets + // here, because only it can take the async state lock. + WalFlushSource::NextPending => Err(Error::internal( + "WalFlushSource::NextPending must be resolved by the flush handler", + )), }; // A terminal failure means the watermark can never advance; latch the // poison so waiters wake with the typed error and later writes fail fast. @@ -546,80 +806,64 @@ impl WalFlusher { result } + /// Append this store's un-appended suffix to the WAL. **Append-only.** + /// + /// The index apply used to run here, concurrently, under a `tokio::join!`. + /// That was the source of the dirty read: `join!` runs both arms to + /// completion and does not cancel the index arm when the append fails, so a + /// failed append still advanced the cursor readers keyed off. The two + /// operations have nothing in common — one is an in-memory microsecond write, + /// the other a ~100ms S3 PUT billed per call — so they now run on separate + /// tasks with separate cursors, and this one only ever touches the WAL. async fn flush_from_batch_store( &self, batch_store: &BatchStore, - indexes: Option>, end_batch_position: usize, ) -> Result { - // Get current flush position from per-memtable watermark (inclusive) - // start_batch_position is the first batch to flush - let start_batch_position = batch_store - .max_flushed_batch_position() - .map(|w| w + 1) - .unwrap_or(0); - - // If we've already flushed past this end, nothing to do + // Where this store's un-appended suffix begins, derived from the + // writer-global durability cursor. `local_end` clamps a cursor that + // predates this memtable to 0 and one past its end to `committed_len`. + let start_batch_position = batch_store.local_end(self.durable()); + + // Already appended past this end: nothing to do. Redundant triggers are + // routine (a put and the freeze can both target the same range), so this + // is the common case, not an error. if start_batch_position >= end_batch_position { return Ok(empty_flush_result()); } - // Collect batches in range [start_batch_position, end_batch_position) - let mut stored_batches: Vec = - Vec::with_capacity(end_batch_position - start_batch_position); - - for batch_position in start_batch_position..end_batch_position { - if let Some(stored) = batch_store.get(batch_position) { - stored_batches.push(stored.clone()); - } - } - - if stored_batches.is_empty() { - return Ok(empty_flush_result()); - } + // Every position in the range must exist. `get` returns `None` only for a + // position past `committed_len`, so a hole means we were asked to append a + // batch the store never committed. Silently skipping it while still + // advancing durability to `end_batch_position` (below) would mark an + // un-appended batch durable and lose it on replay — a divergence that + // survives the crash and hides from a full scan. Return a terminal error + // so `flush` poisons the writer; reopen replays the WAL. + let stored_batches: Vec = (start_batch_position..end_batch_position) + .map(|batch_position| { + batch_store.get(batch_position).cloned().ok_or_else(|| { + Error::writer_poisoned(format!( + "WAL flush range [{start_batch_position}, {end_batch_position}) is \ + missing batch position {batch_position}; batch_store committed_len is {}", + batch_store.len() + )) + }) + }) + .collect::>()?; - let rows_to_index: usize = stored_batches.iter().map(|b| b.num_rows).sum(); let record_batches: Vec = stored_batches.iter().map(|s| s.data.clone()).collect(); - let appender = self.wal_appender.clone(); - let (append_result, index_result) = if let Some(idx_registry) = indexes { - let wal_future = async move { - let start = Instant::now(); - let r = appender.append(record_batches).await?; - Ok::<_, Error>((r, start.elapsed())) - }; - let index_future = async { - let start = Instant::now(); - let per_index = tokio::task::spawn_blocking(move || { - idx_registry.insert_batches_parallel(&stored_batches) - }) - .await - .map_err(|e| Error::internal(format!("Index update task panicked: {}", e)))??; - Ok::<_, Error>((start.elapsed(), per_index)) - }; - tokio::join!(wal_future, index_future) - } else { - let wal_future = async move { - let start = Instant::now(); - let r = appender.append(record_batches).await?; - Ok::<_, Error>((r, start.elapsed())) - }; - ( - wal_future.await, - Ok((std::time::Duration::ZERO, std::collections::HashMap::new())), - ) - }; - - let (append_result, wal_io_duration) = append_result?; - let (index_update_duration, index_update_duration_breakdown) = index_result?; - - // Update per-memtable watermark (inclusive: last batch ID that was flushed) - batch_store.set_max_flushed_batch_position(end_batch_position - 1); + let start = Instant::now(); + let append_result = self.wal_appender.append(record_batches).await?; + let wal_io_duration = start.elapsed(); - // Notify durability waiters (global channel) - let _ = self.durable_watermark_tx.send(end_batch_position); - // Signal WAL flush completion for backpressure waiters + // Advance the writer-global durability cursor and wake waiters. The range + // just appended is `[start, end)` *local to this store*, so it must be + // lifted into the writer's coordinate space before it is published — + // otherwise a fresh memtable's small local end would be compared against a + // cursor carrying a previous memtable's larger one. + self.advance_durable(batch_store.global_offset() + end_batch_position); self.signal_wal_flush_complete(); Ok(WalFlushResult { @@ -629,9 +873,6 @@ impl WalFlusher { num_batches: append_result.num_batches, }), wal_io_duration, - index_update_duration, - index_update_duration_breakdown, - rows_indexed: rows_to_index, wal_bytes: append_result.wal_bytes, }) } @@ -664,9 +905,6 @@ impl WalFlusher { num_batches: append_result.num_batches, }), wal_io_duration, - index_update_duration: std::time::Duration::ZERO, - index_update_duration_breakdown: std::collections::HashMap::new(), - rows_indexed: 0, wal_bytes: append_result.wal_bytes, }) } @@ -699,9 +937,6 @@ pub fn empty_flush_result() -> WalFlushResult { WalFlushResult { entry: None, wal_io_duration: std::time::Duration::ZERO, - index_update_duration: std::time::Duration::ZERO, - index_update_duration_breakdown: std::collections::HashMap::new(), - rows_indexed: 0, wal_bytes: 0, } } @@ -1555,18 +1790,22 @@ mod tests { fn batch_store_source(batch_store: &Arc) -> WalFlushSource { WalFlushSource::BatchStore { batch_store: batch_store.clone(), - indexes: None, } } - fn batch_store_source_with_indexes( - batch_store: &Arc, - indexes: &Arc, - ) -> WalFlushSource { - WalFlushSource::BatchStore { - batch_store: batch_store.clone(), - indexes: Some(indexes.clone()), - } + /// Run the index-apply task's body, as the writer's index task would. + async fn apply_all(batch_store: &Arc, indexes: &Arc) -> Result<()> { + let cursors = Arc::new(WriterCursors::new(true)); + apply_index_range( + &cursors, + TriggerIndexApply { + batch_store: batch_store.clone(), + indexes: indexes.clone(), + end_batch_position: batch_store.len(), + }, + ) + .await + .map(|_| ()) } #[tokio::test] @@ -1576,7 +1815,7 @@ mod tests { let buffer = build_test_flusher(store, &base_path, shard_id, 1); // Track a batch - let watcher = buffer.track_batch(0); + let watcher = buffer.track_batch(None, 0, 1); // Watcher should not be durable yet assert!(!watcher.is_durable()); @@ -1595,7 +1834,7 @@ mod tests { let batch_store = Arc::new(BatchStore::with_capacity(10)); batch_store.append(create_test_batch(&schema, 10)).unwrap(); - let mut watcher = flusher.track_batch(0); + let mut watcher = flusher.track_batch(None, 0, 1); // wait() must NOT resolve before the flush happens let result = @@ -1628,13 +1867,13 @@ mod tests { batch_store.append(batch2).unwrap(); // Track batch IDs in WAL flusher - let mut watcher1 = buffer.track_batch(0); - let mut watcher2 = buffer.track_batch(1); + let mut watcher1 = buffer.track_batch(None, 0, 1); + let mut watcher2 = buffer.track_batch(None, 0, 2); // Verify initial state assert!(!watcher1.is_durable()); assert!(!watcher2.is_durable()); - assert!(batch_store.max_flushed_batch_position().is_none()); + assert_eq!(buffer.durable(), 0); // Flush all pending batches let source = batch_store_source(&batch_store); @@ -1646,8 +1885,8 @@ mod tests { assert_eq!(entry.position, FIRST_WAL_ENTRY_POSITION); assert_eq!(entry.writer_epoch, 1); assert_eq!(entry.num_batches, 2); - // After flushing 2 batches (positions 0 and 1), max flushed position is 1 (inclusive) - assert_eq!(batch_store.max_flushed_batch_position(), Some(1)); + // Two batches appended => the writer-global durable count is 2 (exclusive). + assert_eq!(buffer.durable(), 2); // Watchers should be notified watcher1.wait().await.unwrap(); @@ -1658,12 +1897,19 @@ mod tests { // Regression test for the visibility-cursor bug: with an empty IndexStore // (the common case for WAL-managed tables that mirror an index-less base - // dataset), a WAL flush must still advance `max_visible_batch_position` so - // scanners can see every batch up to the durable position — not just - // batch 0. Before the fix, the cursor stayed at 0 for the lifetime of the - // memtable and scanners returned only the first row. + /// The index apply and the WAL append are separate tasks with separate + /// cursors. Appending makes a range durable; it does not index it. Indexing + /// makes it indexed; it does not make it durable. Only both together make it + /// visible. + /// + /// This also covers the empty-registry case (a memtable with no configured + /// indexes), which used to be skipped entirely by the flush's index arm and + /// so left the cursor stuck at 0 for the memtable's whole life. + #[rstest::rstest] + #[case::no_indexes(false)] + #[case::btree_index(true)] #[tokio::test] - async fn test_wal_flush_advances_visibility_with_empty_indexes() { + async fn test_append_and_index_advance_separate_cursors(#[case] with_btree: bool) { let (store, base_path, _temp_dir) = create_local_store().await; let shard_id = Uuid::new_v4(); let flusher = build_test_flusher(store, &base_path, shard_id, 1); @@ -1674,43 +1920,88 @@ mod tests { batch_store.append(create_test_batch(&schema, 5)).unwrap(); } - // Empty registry, mimicking a memtable with `index_configs = []`. - let indexes = Arc::new(IndexStore::new()); - assert_eq!(indexes.max_visible_batch_position(), 0); + let mut idx = IndexStore::new(); + if with_btree { + idx.add_btree("id_idx".to_string(), 0, "id".to_string()); + } + let indexes = Arc::new(idx); - let source = batch_store_source_with_indexes(&batch_store, &indexes); - flusher.flush(&source, batch_store.len()).await.unwrap(); + // The append alone makes the range durable and indexes nothing. + flusher + .flush(&batch_store_source(&batch_store), batch_store.len()) + .await + .unwrap(); + assert_eq!(flusher.durable(), 3); + assert_eq!(indexes.indexed_count(), 0); - // Cursor must advance to the highest flushed batch position (2), - // making all three batches visible to scanners. - assert_eq!(indexes.max_visible_batch_position(), 2); - assert_eq!(batch_store.max_flushed_batch_position(), Some(2)); + // The index apply alone advances the index cursor. + apply_all(&batch_store, &indexes).await.unwrap(); + assert_eq!(indexes.indexed_count(), 3); } - // Regression guard for the indexed path: with at least one BTree index - // configured, the cursor advance still fires (this was already working - // before the fix — keeping the test to lock in the behavior). + /// A WAL flush asked to cover a range past the store's committed length must + /// fail terminally, not silently short-append. The store is append-only, so + /// `get` returns `None` only past `committed_len`; skipping that position + /// while still advancing durability to `end_batch_position` would mark an + /// un-appended batch durable and lose it on replay. The flush must poison + /// instead, and durability must not move. #[tokio::test] - async fn test_wal_flush_advances_visibility_with_btree_index() { + async fn test_flush_rejects_range_past_committed_len() { let (store, base_path, _temp_dir) = create_local_store().await; let shard_id = Uuid::new_v4(); let flusher = build_test_flusher(store, &base_path, shard_id, 1); let schema = create_test_schema(); let batch_store = Arc::new(BatchStore::with_capacity(10)); - for _ in 0..3 { - batch_store.append(create_test_batch(&schema, 5)).unwrap(); - } + batch_store.append(create_test_batch(&schema, 5)).unwrap(); + batch_store.append(create_test_batch(&schema, 5)).unwrap(); - let mut idx = IndexStore::new(); - idx.add_btree("id_idx".to_string(), 0, "id".to_string()); - let indexes = Arc::new(idx); + // Two batches committed (positions 0, 1); ask to flush through position 2. + let err = flusher + .flush(&batch_store_source(&batch_store), batch_store.len() + 1) + .await + .unwrap_err(); + assert_eq!(err.fence_reason(), Some(FenceReason::PersistenceFailure)); + assert!( + err.to_string().contains("missing batch position 2"), + "unexpected error: {err}" + ); - let source = batch_store_source_with_indexes(&batch_store, &indexes); - flusher.flush(&source, batch_store.len()).await.unwrap(); + // Terminal: the flusher is poisoned and durability never advanced past + // what was actually appended. + assert!(flusher.check_poisoned().is_err()); + assert_eq!(flusher.durable(), 0); + } + + /// The index-apply path has the same invariant: a range past the committed + /// length must error rather than silently under-index and advance the cursor + /// past a batch that was never inserted. + #[tokio::test] + async fn test_index_apply_rejects_range_past_committed_len() { + let schema = create_test_schema(); + let batch_store = Arc::new(BatchStore::with_capacity(10)); + batch_store.append(create_test_batch(&schema, 5)).unwrap(); - assert_eq!(indexes.max_visible_batch_position(), 2); - assert_eq!(batch_store.max_flushed_batch_position(), Some(2)); + let indexes = Arc::new(IndexStore::new()); + let cursors = Arc::new(WriterCursors::new(true)); + + // One batch committed (position 0); ask to index through position 1. + let err = apply_index_range( + &cursors, + TriggerIndexApply { + batch_store: batch_store.clone(), + indexes: indexes.clone(), + end_batch_position: batch_store.len() + 1, + }, + ) + .await + .unwrap_err(); + assert!( + err.to_string().contains("missing batch position 1"), + "unexpected error: {err}" + ); + // The cursor did not advance past the hole. + assert_eq!(indexes.indexed_count(), 0); } #[tokio::test] @@ -1726,8 +2017,8 @@ mod tests { batch_store.append(create_test_batch(&schema, 5)).unwrap(); // Track batch IDs and flush all pending batches - let _watcher1 = buffer.track_batch(0); - let _watcher2 = buffer.track_batch(1); + let _watcher1 = buffer.track_batch(None, 0, 1); + let _watcher2 = buffer.track_batch(None, 0, 2); let source = batch_store_source(&batch_store); let result = buffer.flush(&source, batch_store.len()).await.unwrap(); let entry = result.entry.unwrap(); @@ -1905,7 +2196,7 @@ mod tests { // A durable put on the predecessor: stage a batch and track it. let batch_store = Arc::new(BatchStore::with_capacity(10)); batch_store.append(create_test_batch(&schema, 1)).unwrap(); - let mut watcher = flusher.track_batch(0); + let mut watcher = flusher.track_batch(None, 0, 1); // Flushing collides with the sentinel and fences. Both the flush result // and the watcher must report the fence — and the watcher must resolve @@ -2067,6 +2358,132 @@ mod tests { ); } + /// A failed WAL append must not make rows visible, even when the index apply + /// has already taken them. + /// + /// The two now run on separate tasks, so an index apply that lands while the + /// append is failing advances `indexed_count` — which is fine, and + /// unavoidable: indexes are derived state and replay rebuilds them. What must + /// not happen is for that to make the rows *readable*, because they are not in + /// the WAL and replay will not reproduce them. + /// + /// Visibility is derived, not published, so this holds by construction: with + /// `durable_write`, `visible = min(indexed, durable)`, and a failed append + /// leaves `durable` at 0. + #[tokio::test] + async fn test_failed_append_indexes_but_stays_invisible() { + let (store, base, controls) = failing_memory_store().await; + let shard_id = Uuid::new_v4(); + controls.fail_wal_puts(usize::MAX); + let manifest_store = Arc::new(ShardManifestStore::new(store.clone(), &base, shard_id, 2)); + let (epoch, _) = manifest_store.claim_epoch(0).await.unwrap(); + let appender = Arc::new(WalAppender::with_claimed_epoch( + store, + base, + shard_id, + manifest_store, + epoch, + 0, + WalRetryConfig { + max_retries: 1, + base_delay: Duration::from_millis(1), + }, + )); + let cursors = Arc::new(WriterCursors::new(true)); + let flusher = WalFlusher::with_cursors(appender, Arc::clone(&cursors)); + + let schema = create_test_schema(); + let batch_store = Arc::new(BatchStore::with_capacity(10)); + batch_store.append(create_test_batch(&schema, 1)).unwrap(); + + let mut idx = IndexStore::new(); + idx.add_btree("id_idx".to_string(), 0, "id".to_string()); + idx.set_durability(Arc::clone(&cursors), 0); + let indexes = Arc::new(idx); + + // The index apply succeeds. + apply_index_range( + &cursors, + TriggerIndexApply { + batch_store: batch_store.clone(), + indexes: indexes.clone(), + end_batch_position: 1, + }, + ) + .await + .unwrap(); + assert_eq!(indexes.indexed_count(), 1); + + // The append does not. + let err = flusher + .flush(&batch_store_source(&batch_store), batch_store.len()) + .await + .expect_err("the WAL PUT is failing, so the append must fail"); + assert_eq!(err.fence_reason(), Some(FenceReason::PersistenceFailure)); + + // Indexed, but not durable — so not visible. + assert_eq!(flusher.durable(), 0); + assert_eq!( + indexes.visible_count(), + 0, + "a row whose WAL append failed must never become readable" + ); + } + + /// An index-apply failure poisons the writer. + /// + /// A partial apply cannot be rolled back — `insert_batches` joins every index + /// thread unconditionally, so a failure leaves the others fully applied, and + /// none of HNSW, FTS or BTree has a delete. Continuing would re-cover the + /// range on the next attempt and corrupt the indexes that *did* succeed. So + /// the failure is terminal: reads and writes fail fast, and recovery is + /// reopen -> replay, which rebuilds the indexes from the WAL. + #[tokio::test] + async fn test_index_failure_poisons_the_writer() { + let (store, base_path, _temp_dir) = create_local_store().await; + let shard_id = Uuid::new_v4(); + let flusher = build_test_flusher(store, &base_path, shard_id, 1); + let cursors = Arc::clone(flusher.cursors()); + + let schema = create_test_schema(); + let batch_store = Arc::new(BatchStore::with_capacity(10)); + batch_store.append(create_test_batch(&schema, 1)).unwrap(); + + // An HNSW index on `id`, which is an Int32 and not a vector, so every + // insert of this batch fails deterministically. `validate_index_configs` + // rejects this at shard open — that is what makes poison-and-replay + // terminating — so the store has to be built by hand to reach it at all. + let mut idx = IndexStore::new(); + idx.add_hnsw( + "bad_hnsw".to_string(), + 0, + "id".to_string(), + lance_linalg::distance::DistanceType::L2, + 128, + 8, + ); + let indexes = Arc::new(idx); + + let err = apply_index_range( + &cursors, + TriggerIndexApply { + batch_store: batch_store.clone(), + indexes: indexes.clone(), + end_batch_position: 1, + }, + ) + .await + .expect_err("indexing an Int32 column as a vector must fail"); + + // The index task latches it, exactly as `IndexApplyHandler` does. + flusher.poison(&err); + assert!( + flusher.check_poisoned().is_err(), + "the writer must be poisoned rather than limp on with a corrupt index" + ); + assert_eq!(indexes.visible_count(), 0); + } + // A persistence failure during flush latches the poison: the flush result, // `check_poisoned`, and the durability watcher all report the typed error // (rather than the watcher hanging on a watermark that never advances). @@ -2094,7 +2511,7 @@ mod tests { let schema = create_test_schema(); let batch_store = Arc::new(BatchStore::with_capacity(10)); batch_store.append(create_test_batch(&schema, 1)).unwrap(); - let mut watcher = flusher.track_batch(0); + let mut watcher = flusher.track_batch(None, 0, 1); let source = batch_store_source(&batch_store); let flush_err = flusher.flush(&source, batch_store.len()).await.unwrap_err(); @@ -2114,4 +2531,33 @@ mod tests { .expect_err("watcher must surface the poison"); assert_eq!(waited.fence_reason(), Some(FenceReason::PersistenceFailure)); } + + // A panic under the `terminal_error` lock must not cost the writer its + // latched failure. Recovery is reopen -> replay, which is driven by the real + // `FenceReason`; reporting the mutex poisoning instead would bury it. + #[tokio::test] + async fn test_poisoned_terminal_error_mutex_still_reports_typed_failure() { + let cursors = Arc::new(WriterCursors::new(true)); + cursors.mark_terminal_failure(&Error::writer_poisoned("injected persistence failure")); + + let terminal_error = Arc::clone(&cursors.terminal_error); + let panicked = std::thread::spawn(move || { + let _guard = terminal_error.lock().unwrap(); + panic!("poison the terminal error mutex"); + }) + .join(); + assert!(panicked.is_err()); + assert!(cursors.terminal_error.is_poisoned()); + + let error = cursors.check_poisoned().unwrap_err(); + assert_eq!(error.fence_reason(), Some(FenceReason::PersistenceFailure)); + assert!(error.to_string().contains("injected persistence failure")); + + // The latch still takes writes, and still keeps the first failure. + cursors.mark_terminal_failure(&Error::fenced_by_peer("later peer fence")); + assert_eq!( + cursors.check_poisoned().unwrap_err().fence_reason(), + Some(FenceReason::PersistenceFailure) + ); + } } diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 3cb4586ebb4..ce1acffbc78 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -37,6 +37,7 @@ use uuid::Uuid; pub use super::index::{ BTreeIndexConfig, BTreeMemIndex, FtsIndexConfig, HnswIndexConfig, IndexStore, MemIndexConfig, + validate_index_configs, }; pub use super::memtable::CacheConfig; pub use super::memtable::MemTable; @@ -49,8 +50,8 @@ pub use super::wal::{WalEntry, WalEntryData, WalFlushFailure, WalFlushResult, Wa use super::memtable::flush::TriggerMemTableFlush; use super::scanner::GenerationWarmer; use super::wal::{ - BatchDurableWatcher, TriggerWalFlush, WalAppender, WalFlushSource, WalOnlyState, - WalRetryConfig, WalTailer, empty_flush_result, + BatchDurableWatcher, TriggerIndexApply, TriggerWalFlush, WalAppender, WalFlushSource, + WalOnlyState, WalRetryConfig, WalTailer, WriterCursors, apply_index_range, empty_flush_result, }; use super::{TOMBSTONE, schema_with_tombstone}; use crate::session::Session; @@ -84,17 +85,6 @@ pub struct ShardWriterConfig { /// - Lower latency, batched S3 operations pub durable_write: bool, - /// Whether to update indexes synchronously on each write. - /// - /// When true: - /// - Newly written data is immediately searchable via indexes - /// - Higher latency due to index update overhead - /// - /// When false: - /// - Index updates are deferred - /// - New data may not appear in index-accelerated queries immediately - pub sync_indexed_write: bool, - /// Maximum WAL buffer size in bytes before triggering a flush. /// /// This is a soft threshold - write batches are atomic and won't be split. @@ -166,23 +156,6 @@ pub struct ShardWriterConfig { /// Default: 30 seconds pub backpressure_log_interval: Duration, - /// Maximum rows to buffer before flushing to async indexes. - /// - /// Only applies when `sync_indexed_write` is false. Larger values enable - /// better vectorization but increase memory usage and latency before data - /// becomes searchable. - /// - /// Default: 10,000 rows - pub async_index_buffer_rows: usize, - - /// Maximum time to buffer before flushing to async indexes. - /// - /// Only applies when `sync_indexed_write` is false. Ensures bounded latency - /// for data to become searchable even during low write throughput. - /// - /// Default: 1 second - pub async_index_interval: Duration, - /// Interval for periodic stats logging. /// /// Stats (write throughput, backpressure events, memtable size) are logged @@ -224,8 +197,7 @@ pub struct ShardWriterConfig { /// `durable_write` settings as MemTable mode. /// /// MemTable-tied tunables (`max_memtable_size`, `max_memtable_rows`, - /// `max_memtable_batches`, `sync_indexed_write`, `async_index_buffer_rows`, - /// `async_index_interval`) are ignored when `enable_memtable == false`. + /// `max_memtable_batches`) are ignored when `enable_memtable == false`. /// /// For raw single-entry synchronous atomic appends with no buffering and /// no background tasks, use `WalAppender` directly — it is a strictly @@ -269,7 +241,6 @@ impl Default for ShardWriterConfig { shard_id: Uuid::new_v4(), shard_spec_id: 0, durable_write: true, - sync_indexed_write: true, max_wal_buffer_size: 10 * 1024 * 1024, // 10MB max_wal_flush_interval: Some(Duration::from_millis(100)), // 100ms max_wal_persist_retries: 3, @@ -280,8 +251,6 @@ impl Default for ShardWriterConfig { manifest_scan_batch_size: 2, max_unflushed_memtable_bytes: 1024 * 1024 * 1024, // 1GB backpressure_log_interval: Duration::from_secs(30), - async_index_buffer_rows: 10_000, - async_index_interval: Duration::from_secs(1), stats_log_interval: Some(Duration::from_secs(60)), // 1 minute frozen_memtable_grace: Duration::ZERO, enable_memtable: true, @@ -314,12 +283,6 @@ impl ShardWriterConfig { self } - /// Set indexed writes mode. - pub fn with_sync_indexed_write(mut self, indexed: bool) -> Self { - self.sync_indexed_write = indexed; - self - } - /// Set maximum WAL buffer size. pub fn with_max_wal_buffer_size(mut self, size: usize) -> Self { self.max_wal_buffer_size = size; @@ -382,18 +345,6 @@ impl ShardWriterConfig { self } - /// Set async index buffer rows. - pub fn with_async_index_buffer_rows(mut self, rows: usize) -> Self { - self.async_index_buffer_rows = rows; - self - } - - /// Set async index interval. - pub fn with_async_index_interval(mut self, interval: Duration) -> Self { - self.async_index_interval = interval; - self - } - /// Set stats logging interval. Use None to disable periodic stats logging. pub fn with_stats_log_interval(mut self, interval: Option) -> Self { self.stats_log_interval = interval; @@ -467,7 +418,13 @@ impl TaskDispatcher { let mut ticker_intervals: Vec<(Interval, MessageFactory)> = tickers .into_iter() .map(|(duration, factory)| { - let interval = interval_at(tokio::time::Instant::now() + duration, duration); + let mut interval = interval_at(tokio::time::Instant::now() + duration, duration); + // `Burst` (the default) replays every tick missed while `handle()` + // was running. A WAL append can easily outlast its own interval, so + // the missed ticks pile up, the ticker arm below is always ready, + // and — being `biased` — it starves `rx` indefinitely: freeze + // completion cells and `close()`'s final append never get handled. + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); (interval, factory) }) .collect(); @@ -510,12 +467,12 @@ impl TaskDispatcher { debug!("Task '{}' received cancellation", self.name); break Ok(()); } - _ = first_interval.tick() => { - let message = (ticker_intervals[0].1)(); - if let Err(e) = self.handler.handle(message).await { - error!("Task '{}' error handling ticker message: {}", self.name, e); - } - } + // Explicit messages outrank the ticker. A tick is a backstop + // — it only ever *adds* an append that a real trigger would + // have made anyway — whereas a message may be a freeze's + // completion cell or `close()`'s final append, which nothing + // else will deliver. Polling the ticker first (as this did) + // lets a ready tick starve them. msg = self.rx.recv() => { match msg { Some(message) => { @@ -529,6 +486,12 @@ impl TaskDispatcher { } } } + _ = first_interval.tick() => { + let message = (ticker_intervals[0].1)(); + if let Err(e) = self.handler.handle(message).await { + error!("Task '{}' error handling ticker message: {}", self.name, e); + } + } } } }; @@ -850,14 +813,46 @@ fn now_millis() -> u64 { /// Aborts with an error if any replayed entry's `writer_epoch` is strictly /// greater than `our_epoch` — that indicates a successor writer claimed the /// shard between our `claim_epoch` and this replay, fencing us. +/// Outcome of replaying a shard's WAL into memory. +struct ReplayResult { + /// The active memtable — the final, partial one replay left unsealed. A fresh + /// shard yields an empty one; every sealed memtable was flushed to a Lance + /// generation during replay and is not returned. + active: MemTable, + /// One past the highest WAL entry position observed — the next write position. + next_wal_position: u64, +} + +/// Replay a shard's WAL into memory, flushing sealed memtables as the batch store +/// fills. +/// +/// A single memtable holds at most `max_memtable_batches` batches, but a WAL is +/// unbounded — so replay must rotate exactly as the live write path does. It +/// seals a full memtable and, because the data is already durable, flushes it to +/// a Lance generation right here (the same `MemTableFlusher::flush` the live path +/// uses), rather than holding every sealed memtable in memory until open +/// finishes. That bounds resident memory to ~two memtables and truncates the WAL +/// as it goes, so a later reopen replays only the unflushed tail. Only the final +/// partial memtable is returned, as the active one. +/// +/// `make_memtable(generation, global_offset)` builds a fresh, cursor-bound +/// memtable. Rotation happens at WAL-entry boundaries, never mid-entry, so each +/// sealed memtable covers a clean range of complete entries and stamps the last +/// one as its flushed generation's `replay_after_wal_entry_position`. +#[allow(clippy::too_many_arguments)] async fn replay_memtable_from_wal( object_store: Arc, base_path: Path, shard_id: Uuid, our_epoch: u64, manifest: &ShardManifest, - memtable: &mut MemTable, -) -> Result { + base_generation: u64, + mut make_memtable: impl FnMut(u64, usize) -> Result, + flusher: &MemTableFlusher, + wal_flusher: &WalFlusher, + index_configs: &[MemIndexConfig], + max_memtable_size: usize, +) -> Result { // WAL positions are 1-based (see `FIRST_WAL_ENTRY_POSITION`), so a // cursor of 0 means "no flush has ever stamped this shard" and replay // starts at position 1. After flushing position N the cursor holds N @@ -868,14 +863,11 @@ async fn replay_memtable_from_wal( // contents into the base table. let start_position = manifest.replay_after_wal_entry_position.saturating_add(1); - // The MemTable is always freshly built before this function runs, so - // any existing BatchStore entries can only have come from this replay - // pass. We index everything in `[0, batch_count)` at the end. - debug_assert_eq!(memtable.batch_count(), 0); - let tailer = WalTailer::new(object_store, base_path, shard_id); let mut position = start_position; + let mut active = make_memtable(base_generation, 0)?; + loop { match tailer.read_entry(position).await? { // The first NotFound proves the WAL tip is at `position`, which @@ -894,13 +886,60 @@ async fn replay_memtable_from_wal( // Entries written before deletes existed lack `_tombstone`; // inject `false` so they match the extended memtable schema. // Normal entries already carry it and pass through unchanged. - let target_schema = memtable.schema().clone(); + let target_schema = active.schema().clone(); let batches = entry .batches .into_iter() .map(|b| ensure_tombstone_column(b, &target_schema)) .collect::>>()?; - memtable.insert_batches_only(batches).await?; + + // Seal + flush at the entry boundary on the *same* criteria the + // live path uses (`maybe_trigger_memtable_flush`): the memtable + // is at or over `max_memtable_size` bytes, or this whole entry + // won't fit the batch store. The byte trigger is the one that + // matters beyond avoiding overflow — it is what keeps a memtable + // under `max_memtable_rows`, and therefore keeps the in-memory + // HNSW index (sized to `max_memtable_rows`) from exhausting its + // capacity when the final active memtable is indexed. + // + // Rotate at the entry boundary so no entry is split across two + // memtables and each sealed one covers a clean range of complete + // entries. Never rotate an empty memtable — if a single entry + // has more batches than a memtable can hold, a fresh one would + // overflow too, the same hard limit the live put path has, left + // to the insert below to surface. + if !active.batch_store().is_empty() + && memtable_reached_flush_threshold( + &active, + max_memtable_size, + batches.len(), + ) + { + let store = active.batch_store(); + // The last entry this memtable fully absorbed is the one + // before the entry about to be inserted. + let covered = position.saturating_sub(1); + let generation = active.generation() + 1; + let global_end = store.global_end(); + + // The sealed data is already durable in the WAL — mark it + // so the flush's `all_flushed_to_wal` precondition holds and + // no WAL re-append is attempted. + wal_flusher.advance_durable(global_end); + flush_replayed_memtable( + flusher, + &active, + our_epoch, + covered, + global_end, + index_configs, + ) + .await?; + + active = make_memtable(generation, global_end)?; + } + + active.insert_batches_only(batches).await?; } position = position.checked_add(1).ok_or_else(|| { Error::io(format!( @@ -912,21 +951,18 @@ async fn replay_memtable_from_wal( } } - // Update in-memory indexes with the replayed batches so readers see them - // through the index path (matching what would have happened on the - // pre-crash writer's WAL flush). Indexes from the previous writer don't - // persist; this rebuilds them from the WAL. - if let Some(indexes) = memtable.indexes_arc() { - let batches_after = memtable.batch_count(); - if batches_after > 0 { - let store = memtable.batch_store(); - let mut stored: Vec = Vec::with_capacity(batches_after); - for pos in 0..batches_after { - if let Some(s) = store.get(pos) { - stored.push(s.clone()); - } - } - tokio::task::spawn_blocking(move || indexes.insert_batches_parallel(&stored)) + // Rebuild the active memtable's in-memory indexes from the batches just + // replayed, so readers see them through the index path — matching what the + // pre-crash writer's flush would have done. Sealed memtables needed no + // in-memory index build: they were flushed straight to disk and are gone. + if let Some(indexes) = active.indexes_arc() { + let batch_count = active.batch_count(); + if batch_count > 0 { + let store = active.batch_store(); + let stored: Vec = (0..batch_count) + .filter_map(|pos| store.get(pos).cloned()) + .collect(); + tokio::task::spawn_blocking(move || indexes.insert_batches(&stored)) .await .map_err(|e| { Error::internal(format!("WAL replay index update task panicked: {}", e)) @@ -934,7 +970,50 @@ async fn replay_memtable_from_wal( } } - Ok(position) + Ok(ReplayResult { + active, + next_wal_position: position, + }) +} + +/// Whether a memtable has reached the threshold at which it should be sealed and +/// flushed: at or over `max_memtable_size` bytes, or without room in its batch +/// store for `incoming_batches` more. +/// +/// The single source of truth for the flush trigger, shared by the live put path +/// (`maybe_trigger_memtable_flush`, checking post-insert with `incoming_batches = +/// 1` — "is there room for the next batch") and by replay (checking pre-insert +/// with the next WAL entry's batch count). Keeping one predicate is what stops the +/// two from drifting — e.g. someone adding a third criterion to one and not the +/// other, which is the exact class of bug this whole change set is about. +fn memtable_reached_flush_threshold( + memtable: &MemTable, + max_memtable_size: usize, + incoming_batches: usize, +) -> bool { + memtable.estimated_size() >= max_memtable_size + || memtable.batch_store().remaining_capacity() < incoming_batches +} + +/// Flush a sealed replay memtable to a Lance generation, choosing the indexed +/// path when secondary indexes are configured (mirroring the live memtable-flush +/// handler). Commits the manifest, stamping `covered` as the generation's +/// `replay_after_wal_entry_position` so a later reopen skips these entries. +async fn flush_replayed_memtable( + flusher: &MemTableFlusher, + memtable: &MemTable, + epoch: u64, + covered: u64, + durable: usize, + index_configs: &[MemIndexConfig], +) -> Result<()> { + if index_configs.is_empty() { + flusher.flush(memtable, epoch, covered, durable).await?; + } else { + Box::pin(flusher.flush_with_indexes(memtable, epoch, index_configs, covered, durable)) + .await?; + } + Ok(()) } /// Pair each primary-key column name with its field id (both derived from the @@ -1015,6 +1094,10 @@ struct SharedWriterState { state: Arc>, wal_flusher: Arc, wal_flush_tx: mpsc::UnboundedSender, + /// The index-apply task's channel. Separate from the WAL flusher's on + /// purpose: `TaskDispatcher::run` awaits `handle()` inline, so sharing one + /// would put a ~100ms S3 PUT in front of every latency-sensitive index apply. + index_apply_tx: mpsc::UnboundedSender, memtable_flush_tx: mpsc::UnboundedSender, config: ShardWriterConfig, schema: Arc, @@ -1033,6 +1116,7 @@ impl SharedWriterState { state: Arc>, wal_flusher: Arc, wal_flush_tx: mpsc::UnboundedSender, + index_apply_tx: mpsc::UnboundedSender, memtable_flush_tx: mpsc::UnboundedSender, config: ShardWriterConfig, schema: Arc, @@ -1046,6 +1130,7 @@ impl SharedWriterState { state, wal_flusher, wal_flush_tx, + index_apply_tx, memtable_flush_tx, config, schema, @@ -1057,59 +1142,100 @@ impl SharedWriterState { } } + /// Ask the index-apply task to cover `[indexed, end_batch_position)` of this + /// store. Cheap and idempotent: a range already covered is a no-op, which is + /// the common case under load, since one apply coalesces the puts queued + /// behind it. + fn trigger_index_apply( + &self, + batch_store: Arc, + indexes: Arc, + end_batch_position: usize, + ) -> Result<()> { + self.index_apply_tx + .send(TriggerIndexApply { + batch_store, + indexes, + end_batch_position, + }) + .map_err(|_| Error::io("index apply channel closed")) + } + /// Freeze the current memtable and send it to the flush handler. /// /// Takes `&mut WriterState` directly since caller already holds the lock. fn freeze_memtable(&self, state: &mut WriterState) -> Result { - let pending_wal_range = state.memtable.batch_store().pending_wal_flush_range(); + let durable = self.wal_flusher.durable(); + let pending_wal_range = state + .memtable + .batch_store() + .pending_wal_flush_range(durable); let last_wal_entry_position = state.last_flushed_wal_entry_position; let old_batch_store = state.memtable.batch_store(); - let old_indexes = state.memtable.indexes_arc(); let next_generation = state.memtable.generation() + 1; - let mut new_memtable = MemTable::with_capacity( + // The incoming memtable's batch 0 continues the writer's batch sequence + // where the outgoing one ends. Without this coordinate, local positions + // (which restart at 0 every rotation) cannot be mapped onto the + // writer-global durability cursor. + let next_global_offset = old_batch_store.global_end(); + let mut new_memtable = MemTable::with_capacity_at( self.schema.clone(), next_generation, self.pk_field_ids.clone(), CacheConfig::default(), self.max_memtable_batches, + next_global_offset, )?; - // Build an IndexStore when there are user indexes *or* a primary key: - // the PK dedup index (and its flushed on-disk sidecar) is required for - // cross-generation dedup even when no secondary index is configured. - if !self.index_configs.is_empty() || !self.pk_columns.is_empty() { - let mut indexes = IndexStore::from_configs( - &self.index_configs, - self.max_memtable_rows, - self.max_memtable_batches, - )?; + // Always build and bind an IndexStore, even with no user indexes and no + // primary key. It is what carries the memtable's `indexed_count`, and + // binding it to the writer's cursors is what lets a reader derive the + // visible prefix — so an index-less memtable that skipped this would fall + // back to `visible == indexed` and publish rows before they were durable. + // (A PK memtable also needs the PK dedup index and its flushed sidecar.) + let mut indexes = IndexStore::from_configs( + &self.index_configs, + self.max_memtable_rows, + self.max_memtable_batches, + )?; + if !self.pk_columns.is_empty() { indexes.enable_pk_index(&pk_index_columns(&self.pk_columns, &self.pk_field_ids)); - new_memtable.set_indexes_arc(Arc::new(indexes)); } + indexes.set_durability(Arc::clone(self.wal_flusher.cursors()), next_global_offset); + new_memtable.set_indexes_arc(Arc::new(indexes)); let mut old_memtable = std::mem::replace(&mut state.memtable, new_memtable); old_memtable.freeze(last_wal_entry_position); + + // Set up completion tracking on the outgoing table before it is retained + // and before any fallible dispatch, so the retained table already carries + // its cells and a failed send below can poison-and-return without leaving + // partial state to unwind. let _memtable_flush_watcher = old_memtable.create_memtable_flush_completion(); - if pending_wal_range.is_some() { + // The outgoing memtable may still owe an index apply — the puts that + // filled it triggered one, but the task need not have drained yet, and + // this is the last chance to name that store. Its L0 flush is gated on + // the WAL append (below), not on indexing, so without this its tail could + // stay unindexed and invisible for the rest of its life. + let pending_index_apply = match old_memtable.indexes_arc() { + Some(old_indexes) if old_indexes.indexed_count() < old_batch_store.len() => { + Some((old_batch_store.clone(), old_indexes, old_batch_store.len())) + } + _ => None, + }; + + let pending_wal_flush = if pending_wal_range.is_some() { let completion_cell: WatchableOnceCell< std::result::Result, > = WatchableOnceCell::new(); - let completion_reader = completion_cell.reader(); - old_memtable.set_wal_flush_completion(completion_reader); - - let end_batch_position = old_batch_store.len(); - self.wal_flusher.trigger_flush( - WalFlushSource::BatchStore { - batch_store: old_batch_store, - indexes: old_indexes, - }, - end_batch_position, - Some(completion_cell), - )?; - } + old_memtable.set_wal_flush_completion(completion_cell.reader()); + Some((old_batch_store.len(), completion_cell)) + } else { + None + }; let frozen_size = old_memtable.estimated_size(); state.frozen_memtable_bytes += frozen_size; @@ -1123,14 +1249,38 @@ impl SharedWriterState { let frozen_memtable = Arc::new(old_memtable); - // Keep this generation queryable past its manifest commit (swept after - // the grace by `SweepExpired`). Arc refcount, not a copy — the flush - // task holds it alive for the whole drain anyway. + // Retain the outgoing table in the read view *before* the fallible + // dispatches below. `state.memtable` was already replaced, so a failed + // send that returned here without this push would drop the table and its + // accepted rows would silently vanish from every scan. Keep it queryable + // past its manifest commit too (swept after the grace by `SweepExpired`); + // Arc refcount, not a copy — the flush task holds it alive anyway. state.frozen_memtables.push_back(FrozenMemTable { memtable: frozen_memtable.clone(), flushed_at_ms: None, }); + // Dispatch can only fail if a background task's channel is already closed, + // i.e. the writer is being torn down. Poison so the read path fails fast + // with the typed error instead of serving the retained-but-never-durable + // tail, then return — the table stays in the read view. + if let Some((batch_store, indexes, end_batch_position)) = pending_index_apply { + self.trigger_index_apply(batch_store, indexes, end_batch_position) + .inspect_err(|e| self.wal_flusher.poison(e))?; + } + + if let Some((end_batch_position, completion_cell)) = pending_wal_flush { + self.wal_flusher + .trigger_flush( + WalFlushSource::BatchStore { + batch_store: old_batch_store, + }, + end_batch_position, + Some(completion_cell), + ) + .inspect_err(|e| self.wal_flusher.poison(e))?; + } + debug!( "Frozen memtable generation {}, pending_count = {}", next_generation - 1, @@ -1145,9 +1295,19 @@ impl SharedWriterState { Ok(next_generation) } - /// Track batch for WAL durability. - fn track_batch_for_wal(&self, batch_position: usize) -> super::wal::BatchDurableWatcher { - self.wal_flusher.track_batch(batch_position) + /// Watch for a write to become visible: indexed, and — in durable mode — + /// WAL-durable too. + /// + /// `target_indexed` is memtable-local; `target_durable` is writer-global. + /// Different coordinate spaces on purpose — see `WalFlusher::track_batch`. + fn track_batch_for_wal( + &self, + indexes: Option>, + target_indexed: usize, + target_durable: usize, + ) -> super::wal::BatchDurableWatcher { + self.wal_flusher + .track_batch(indexes, target_indexed, target_durable) } /// Check if memtable flush is needed and trigger if so. @@ -1158,8 +1318,10 @@ impl SharedWriterState { return Ok(()); } - let should_flush = state.memtable.estimated_size() >= self.config.max_memtable_size - || state.memtable.is_batch_store_full(); + // Checked post-insert: flush if there is no longer room for even one more + // batch (or the byte threshold is crossed). Same predicate replay uses. + let should_flush = + memtable_reached_flush_threshold(&state.memtable, self.config.max_memtable_size, 1); if should_flush { state.flush_requested = true; @@ -1178,10 +1340,9 @@ impl SharedWriterState { let batch_count = state.memtable.batch_count(); let total_bytes = state.memtable.estimated_size(); let batch_store = state.memtable.batch_store(); - let indexes = state.memtable.indexes_arc(); // Check if there are any unflushed batches - let has_pending = batch_store.pending_wal_flush_count() > 0; + let has_pending = batch_store.pending_wal_flush_count(self.wal_flusher.durable()) > 0; // Check time-based trigger first let time_trigger = if let Some(interval) = self.config.max_wal_flush_interval { @@ -1210,10 +1371,7 @@ impl SharedWriterState { // If time trigger fired, send a flush message if time_trigger.is_some() { let _ = self.wal_flush_tx.send(TriggerWalFlush { - source: WalFlushSource::BatchStore { - batch_store, - indexes, - }, + source: WalFlushSource::BatchStore { batch_store }, end_batch_position: batch_count, done: None, }); @@ -1238,7 +1396,6 @@ impl SharedWriterState { let _ = self.wal_flush_tx.send(TriggerWalFlush { source: WalFlushSource::BatchStore { batch_store: batch_store.clone(), - indexes: indexes.clone(), }, end_batch_position: batch_count, done: None, @@ -1362,6 +1519,47 @@ impl ShardWriter { config.manifest_scan_batch_size, )); + // Derive PK metadata and run every side-effect-free validation *before* + // claiming the epoch. `claim_epoch` durably bumps the stored epoch and, + // for a successor, `write_fence_sentinel` fences the predecessor — so an + // open doomed by purely local input (a bad flush interval, an index config + // that disagrees with the schema) must fail here, before it can knock the + // healthy incumbent off the shard. Memtable-only: WAL-only mode has no + // indexes and no memtable ticker to validate. + let memtable_validation = if config.enable_memtable { + let lance_schema = Schema::try_from(schema.as_ref())?; + let pk_fields = lance_schema.unenforced_primary_key(); + let pk_field_ids: Vec = pk_fields.iter().map(|f| f.id).collect(); + let pk_columns: Vec = pk_fields.iter().map(|f| f.name.clone()).collect(); + + // A durable writer with no flush ticker cannot make progress. The WAL + // append is timer-, size-, and freeze-driven now — the put path no + // longer triggers its own — so with no ticker a durable put below the + // size threshold waits on an append that nothing will ever make. Reject + // the combination rather than deadlock on it. + if config.durable_write + && config + .max_wal_flush_interval + .is_none_or(|interval| interval.is_zero()) + { + return Err(Error::invalid_input( + "durable_write requires a non-zero max_wal_flush_interval: the WAL append is \ + driven by that ticker, so without it a durable put would wait forever for an \ + append that never happens", + )); + } + + // Reject an index config that disagrees with the schema *before* a + // single row is accepted. Such a config fails deterministically on + // every insert, including inserts replayed from the WAL — so once a row + // is durable the shard can never reopen. Fail the open instead. + validate_index_configs(&index_configs, schema.as_ref(), &lance_schema, &pk_columns)?; + + Some((pk_field_ids, pk_columns)) + } else { + None + }; + // Claim the shard (epoch-based fencing) — done once, then shared // with the WalAppender via `with_claimed_epoch`. let (epoch, manifest) = manifest_store.claim_epoch(config.shard_spec_id).await?; @@ -1403,7 +1601,12 @@ impl ShardWriter { } // Create WAL flusher backed by the shared appender. - let mut wal_flusher = WalFlusher::new(wal_appender); + // Build the cursors from *this writer's* config. `durable_write` is what + // decides whether durability is part of visibility, so a flusher that + // defaulted it would leave a non-durable put waiting on a durability + // cursor nothing ever advances. + let cursors = Arc::new(WriterCursors::new(config.durable_write)); + let mut wal_flusher = WalFlusher::with_cursors(wal_appender, cursors); let (wal_flush_tx, wal_flush_rx) = mpsc::unbounded_channel(); wal_flusher.set_flush_channel(wal_flush_tx.clone()); @@ -1413,11 +1616,15 @@ impl ShardWriter { let task_executor = Arc::new(TaskExecutor::new()); let mode = if config.enable_memtable { + let (pk_field_ids, pk_columns) = memtable_validation + .expect("memtable_validation is Some when enable_memtable is true"); Self::open_memtable_mode( &config, &schema, &manifest, &index_configs, + pk_field_ids, + pk_columns, wal_flusher.clone(), wal_flush_tx, wal_flush_rx, @@ -1459,6 +1666,8 @@ impl ShardWriter { schema: &Arc, manifest: &ShardManifest, index_configs: &[MemIndexConfig], + pk_field_ids: Vec, + pk_columns: Vec, wal_flusher: Arc, wal_flush_tx: mpsc::UnboundedSender, wal_flush_rx: mpsc::UnboundedReceiver, @@ -1471,50 +1680,93 @@ impl ShardWriter { stats: SharedWriteStats, task_executor: &Arc, ) -> Result { - // Create MemTable with primary key field IDs from schema - let lance_schema = Schema::try_from(schema.as_ref())?; - let pk_fields = lance_schema.unenforced_primary_key(); - let pk_field_ids: Vec = pk_fields.iter().map(|f| f.id).collect(); - let pk_columns: Vec = pk_fields.iter().map(|f| f.name.clone()).collect(); - let mut memtable = MemTable::with_capacity( - schema.clone(), - manifest.current_generation, - pk_field_ids.clone(), - CacheConfig::default(), - config.max_memtable_batches, - )?; - - // Create indexes if configured and set them on the MemTable. The - // PK-position index is enabled before any WAL replay below so replayed - // rows are recorded in it. A primary key alone (no secondary index) - // still needs the PK index so flush writes its on-disk dedup sidecar. - if !index_configs.is_empty() || !pk_columns.is_empty() { + // PK metadata and index/interval validation were resolved in `open` + // before the epoch was claimed (a doomed open must not fence the + // incumbent first). + + // Build a fresh, cursor-bound memtable at a given generation and + // writer-global coordinate. Replay calls this for the first memtable and + // after every rotation. Always builds and binds an `IndexStore`, even + // with no user indexes and no primary key — see the note in + // `freeze_memtable` for why an index-less memtable still needs one. + let make_bound_memtable = |generation: u64, global_offset: usize| -> Result { + let mut memtable = MemTable::with_capacity_at( + schema.clone(), + generation, + pk_field_ids.clone(), + CacheConfig::default(), + config.max_memtable_batches, + global_offset, + )?; let mut indexes = IndexStore::from_configs( index_configs, config.max_memtable_rows, config.max_memtable_batches, )?; - indexes.enable_pk_index(&pk_index_columns(&pk_columns, &pk_field_ids)); + if !pk_columns.is_empty() { + indexes.enable_pk_index(&pk_index_columns(&pk_columns, &pk_field_ids)); + } + indexes.set_durability(Arc::clone(wal_flusher.cursors()), global_offset); memtable.set_indexes_arc(Arc::new(indexes)); - } + Ok(memtable) + }; + + // The flusher writes sealed memtables to Lance generations — both the + // ones replay seals below and the ones the live path freezes later. + let flusher = Arc::new( + MemTableFlusher::new( + object_store.clone(), + base_path.clone(), + base_uri.clone(), + shard_id, + manifest_store.clone(), + ) + .with_warmer(config.warmer.clone()) + .with_storage_context(config.store_params.clone(), config.session.clone()), + ); // Replay any WAL entries written after the last successfully-flushed - // generation. Each entry's writer_epoch is checked against ours; an - // entry with a strictly greater epoch indicates a successor writer - // claimed the shard between our `claim_epoch` and replay, so we - // abort the open with a fence error. The replay walked the tailer - // up to the WAL tip, so we hand the discovered next-write position - // straight to the appender — its first append skips the - // discover_next_position probe entirely. - let next_wal_position = replay_memtable_from_wal( + // generation, flushing sealed memtables to Lance generations as the batch + // store fills. Each entry's writer_epoch is checked against ours; an entry + // with a strictly greater epoch means a successor claimed the shard + // between our `claim_epoch` and replay, so we abort with a fence error. + // Replay walks the tailer to the WAL tip and returns the discovered + // next-write position, so the appender's first append skips the + // discover_next_position probe. + let ReplayResult { + active: memtable, + next_wal_position, + } = replay_memtable_from_wal( object_store.clone(), base_path.clone(), shard_id, epoch, manifest, - &mut memtable, + manifest.current_generation, + make_bound_memtable, + &flusher, + &wal_flusher, + index_configs, + config.max_memtable_size, ) .await?; + + // Mark the active memtable's replayed batches durable. They came *from* + // the WAL, and replay has already re-derived its indexes over them. + // + // Without this the durability cursor stays at the last sealed generation, + // so the next WAL flush re-covers the active tail: it re-appends the + // already-durable rows *and* re-inserts every replayed row into the + // indexes. None of the three in-memory indexes is idempotent (HNSW mints + // fresh node ids for the same row, FTS increments doc_count/df rather than + // recomputing them, BTree is a multiset), so a full scan keeps looking + // healthy while every index-accelerated query silently returns duplicates + // — and it compounds, because the WAL now holds those rows twice. + // + // `global_end()` is the writer-global batch count through this memtable, + // since its coordinate continues where the last sealed generation ended. + wal_flusher.advance_durable(memtable.batch_store().global_end()); + wal_flusher .wal_appender() .seed_next_position(next_wal_position) @@ -1545,17 +1797,15 @@ impl ShardWriter { let (memtable_flush_tx, memtable_flush_rx) = mpsc::unbounded_channel(); - let flusher = Arc::new( - MemTableFlusher::new(object_store, base_path, base_uri, shard_id, manifest_store) - .with_warmer(config.warmer.clone()) - .with_storage_context(config.store_params.clone(), config.session.clone()), - ); - let backpressure = BackpressureController::new(config.clone()); // Background WAL flush handler — parallel WAL I/O + index updates. - let wal_handler = - WalFlushHandler::new(wal_flusher.clone(), Some(state.clone()), stats.clone()); + let wal_handler = WalFlushHandler::new( + wal_flusher.clone(), + Some(state.clone()), + config.max_wal_flush_interval, + stats.clone(), + ); task_executor.add_handler( "wal_flusher".to_string(), Box::new(wal_handler), @@ -1567,9 +1817,10 @@ impl ShardWriter { let memtable_handler = MemTableFlushHandler::new( state.clone(), flusher, + wal_flusher.clone(), epoch, index_configs.to_vec(), - stats, + stats.clone(), config.frozen_memtable_grace, ); task_executor.add_handler( @@ -1578,11 +1829,27 @@ impl ShardWriter { memtable_flush_rx, )?; + // The index-apply task. Its own channel and its own dispatcher: the + // dispatcher awaits `handle()` inline, so sharing the WAL flusher's + // channel would queue every index apply behind a ~100ms S3 PUT. + let (index_apply_tx, index_apply_rx) = mpsc::unbounded_channel(); + let index_handler = IndexApplyHandler { + cursors: Arc::clone(wal_flusher.cursors()), + wal_flusher: wal_flusher.clone(), + stats, + }; + task_executor.add_handler( + "index_applier".to_string(), + Box::new(index_handler), + index_apply_rx, + )?; + // Shared state used by `put()` to dispatch trigger checks. let writer_state = Arc::new(SharedWriterState::new( state.clone(), wal_flusher, wal_flush_tx, + index_apply_tx, memtable_flush_tx, config.clone(), schema.clone(), @@ -1610,7 +1877,7 @@ impl ShardWriter { ) -> Result { // Background WAL flush handler — no MemTable state to consult, so // pass `None` for the frozen-vs-active detection. - let wal_handler = WalFlushHandler::new(wal_flusher, None, stats); + let wal_handler = WalFlushHandler::new(wal_flusher, None, None, stats); task_executor.add_handler( "wal_flusher".to_string(), Box::new(wal_handler), @@ -1865,46 +2132,74 @@ impl ShardWriter { // 1. Insert all batches into memtable atomically let results = state.memtable.insert_batches_only(batches).await?; - // Get batch position range + // 2. Capture the store the batches actually landed in, *before* step + // 4 below can freeze and swap the active memtable. Reading it + // afterwards hands the flush trigger the **new** store paired with + // the **old** store's end position, so the new store's watermark + // jumps past batches that were never appended. + let batch_store = state.memtable.batch_store(); + let indexes = state.memtable.indexes_arc(); + let start_pos = results.first().map(|(pos, _, _)| *pos).unwrap_or(0); let end_pos = results.last().map(|(pos, _, _)| pos + 1).unwrap_or(0); let batch_positions = start_pos..end_pos; - // 2. Track last batch for WAL durability - let durable_watcher = writer_state.track_batch_for_wal(end_pos.saturating_sub(1)); + // 3. Watch for this write to become *visible*: indexed, and — under + // `durable_write` — WAL-durable too. + // + // The two targets live in different coordinate spaces. `end_pos` + // is memtable-local, which is what the index apply works in. The + // durability cursor is writer-global, because batch positions + // restart at 0 in every memtable while that cursor spans the + // writer's whole life — a local durable target would already be + // satisfied by a *previous* memtable's appends, so the first N + // puts into every post-rotation memtable would ack as durable with + // no WAL append ever happening. + let durable_watcher = writer_state.track_batch_for_wal( + indexes.clone(), + end_pos, + batch_store.global_offset() + end_pos, + ); - // 3. Check if WAL flush should be triggered + // 4. Check if WAL flush should be triggered writer_state.maybe_trigger_wal_flush(&mut state); - // 4. Check if memtable flush is needed + // 5. Check if memtable flush is needed (may freeze and rotate) if let Err(e) = writer_state.maybe_trigger_memtable_flush(&mut state) { warn!("Failed to trigger memtable flush: {}", e); } - // Get batch_store and indexes while we have the lock (for durable_write case) - let batch_store = state.memtable.batch_store(); - let indexes = state.memtable.indexes_arc(); - (batch_positions, durable_watcher, batch_store, indexes) }; // Lock released here self.stats.record_put(start.elapsed()); - // Trigger the flush here (outside the lock) so the watcher can resolve; - // only the `wait()` is the caller's to schedule. - let watcher = if self.config.durable_write { - self.wal_flusher.trigger_flush( - WalFlushSource::BatchStore { - batch_store, - indexes, - }, - batch_positions.end, - None, - )?; - Some(durable_watcher) - } else { - None - }; + // Trigger the index apply, in **both** modes. This is what makes reads + // read-your-writes regardless of `durable_write`: skipping durability now + // costs the caller durability only, not visibility. It is cheap + // (in-memory, ~ms), so there is no reason to batch it onto the WAL's + // schedule — that schedule exists to bound S3 API cost, which an + // in-memory index apply does not incur. + if let Some(indexes) = indexes { + writer_state.trigger_index_apply(batch_store, indexes, batch_positions.end)?; + } + + // The WAL append is *not* triggered here. It happens on the background + // ticker (and on the size trigger, and at freeze/close), which is the only + // way the flush interval can mean anything: while every durable put + // triggered its own append, the interval could add a redundant trigger but + // never delay or batch one. + // + // The cost is real and accepted: a single client's sequential *durable* + // throughput drops from ~10 writes/sec (one PUT round-trip) to roughly one + // per tick. That is a policy choice — the interval should mean what it + // says, and S3 API cost should be bounded. Latency-sensitive callers want + // `durable_write: false`, which now costs them durability only, not + // visibility. + + // The watcher is returned in both modes now. A non-durable put still + // waits — for its index apply (~ms), not for an S3 PUT (~100ms). + let watcher = Some(durable_watcher); Ok((WriteResult { batch_positions }, watcher)) } @@ -2082,18 +2377,24 @@ impl ShardWriter { /// Get current MemTable statistics. Returns an error in WAL-only mode /// (no MemTable exists). + /// + /// Deliberately does *not* `check_poisoned`, unlike the read and write + /// paths: a poisoned writer is exactly when an operator most needs to see + /// its state, and the caller deciding whether to evict reads these stats. pub async fn memtable_stats(&self) -> Result { let state_lock = self.memtable_state_lock()?; let state = state_lock.read().await; let batch_store = state.memtable.batch_store(); - let pending_wal = batch_store.pending_wal_flush_stats(); + let durable = self.wal_flusher.durable(); + let pending_wal = batch_store.pending_wal_flush_stats(durable); Ok(MemTableStats { row_count: state.memtable.row_count(), batch_count: state.memtable.batch_count(), estimated_size: state.memtable.estimated_size(), generation: state.memtable.generation(), max_buffered_batch_position: batch_store.max_buffered_batch_position(), - max_flushed_batch_position: batch_store.max_flushed_batch_position(), + durable_batch_count: durable, + global_offset: batch_store.global_offset(), pending_wal_start_batch_position: pending_wal.start_batch_position, pending_wal_end_batch_position: pending_wal.end_batch_position, pending_wal_batch_count: pending_wal.batch_count, @@ -2107,11 +2408,12 @@ impl ShardWriter { /// The scanner provides read access to all data currently in the MemTable, /// with optional filtering, projection, and index support. /// - /// The scanner captures the current `max_visible_batch_position` from the + /// The scanner captures the current `visible_count` from the /// `IndexStore` at construction time to ensure consistent visibility. /// - /// Returns an error in WAL-only mode. + /// Returns an error in WAL-only mode, or if the writer is poisoned. pub async fn scan(&self) -> Result { + self.wal_flusher.check_poisoned()?; let state_lock = self.memtable_state_lock()?; let state = state_lock.read().await; Ok(state.memtable.scan()) @@ -2121,10 +2423,11 @@ impl ShardWriter { /// Prefer [`Self::in_memory_memtable_refs`] on the read path — it also /// carries frozen-awaiting-flush generations. /// - /// Returns an error in WAL-only mode. + /// Returns an error in WAL-only mode, or if the writer is poisoned. pub async fn active_memtable_ref( &self, ) -> Result { + self.wal_flusher.check_poisoned()?; let state_lock = self.memtable_state_lock()?; let state = state_lock.read().await; Ok(in_memory_ref(&state.memtable)) @@ -2136,10 +2439,11 @@ impl ShardWriter { /// path uses this instead of [`Self::active_memtable_ref`] so a /// concurrent reader sees no hole while a flush drains. /// - /// Returns an error in WAL-only mode. + /// Returns an error in WAL-only mode, or if the writer is poisoned. pub async fn in_memory_memtable_refs( &self, ) -> Result { + self.wal_flusher.check_poisoned()?; let state_lock = self.memtable_state_lock()?; let state = state_lock.read().await; Ok(crate::dataset::mem_wal::scanner::InMemoryMemTables { @@ -2349,20 +2653,41 @@ impl ShardWriter { writer_state, .. } => { - // Send final WAL flush message and wait for completion + // Drain *both* tasks against the active memtable. The index apply + // and the WAL append are independent now, so closing has to + // settle both: the L0 flush below turns this memtable into a + // Lance generation, and a generation whose indexes never saw the + // tail is a generation with a hole in it. let st = state.read().await; let batch_store = st.memtable.batch_store(); let indexes = st.memtable.indexes_arc(); let batch_count = st.memtable.batch_count(); drop(st); + if batch_count > 0 + && let Some(indexes) = indexes + && indexes.indexed_count() < batch_count + { + let mut watcher = self.wal_flusher.track_batch( + Some(Arc::clone(&indexes)), + batch_count, + 0, // durability is settled by the WAL flush below + ); + writer_state.trigger_index_apply( + Arc::clone(&batch_store), + indexes, + batch_count, + )?; + watcher.wait().await?; + } + if batch_count > 0 { + // Append-only source: on this branch the index apply is a + // separate task (drained above), so the final WAL flush carries + // no indexes. #7769's failure propagation still applies. let stage_result = Self::flush_final_wal( &writer_state.wal_flush_tx, - WalFlushSource::BatchStore { - batch_store, - indexes, - }, + WalFlushSource::BatchStore { batch_store }, batch_count, ) .await; @@ -2463,7 +2788,12 @@ pub struct MemTableStats { pub estimated_size: usize, pub generation: u64, pub max_buffered_batch_position: Option, - pub max_flushed_batch_position: Option, + /// Writer-global count of WAL-durable batches. Exclusive: 0 means none. + /// Compare against `global_offset + batch_count` to see what this memtable + /// still owes the WAL. + pub durable_batch_count: usize, + /// Writer-global coordinate of this memtable's batch 0. + pub global_offset: usize, pub pending_wal_start_batch_position: Option, pub pending_wal_end_batch_position: Option, pub pending_wal_batch_count: usize, @@ -2478,16 +2808,76 @@ pub struct WalStats { pub next_wal_entry_position: u64, } -/// Background handler for WAL flush operations. +/// The oldest store that still owes the WAL an append, or `None` when everything +/// is durable. /// -/// This handler does parallel WAL I/O + index updates during flush. -/// Indexes are passed through the TriggerWalFlush message. +/// Ordering is the whole point. WAL entry positions are assigned in append-call +/// order; replay walks them ascending; row positions follow; and primary-key +/// recency is "newest visible row position wins". So appending a newer memtable +/// ahead of an older one's tail silently inverts dedup after a crash. +/// +/// Taking "the active memtable" would do exactly that, because a timer tick +/// enqueued before a freeze is handled after it and would resolve to the incoming +/// memtable. Selecting by cursor instead makes the target a function of what is +/// actually durable, not of when the timer happened to fire. +fn next_pending_store( + frozen: impl Iterator>, + active: Arc, + durable: usize, +) -> Option> { + frozen + .chain(std::iter::once(active)) + .find(|store| store.global_end() > durable) +} + +/// The index-apply task: one sequential consumer of the index-apply channel. +/// +/// Sequential consumption is the safety property. `HnswGraph::insert_batch` hard- +/// rejects any range whose start is not its `indexed_len`, so the apply must see +/// contiguous, in-order ranges — and a single consumer guarantees that +/// regardless of how many putters race behind it. Ordering comes from the task, +/// not from a flush interval, which is why triggering per-put is exactly as safe +/// as triggering on a timer, and lets a put become visible in milliseconds +/// instead of waiting on an S3 round-trip. +struct IndexApplyHandler { + cursors: Arc, + wal_flusher: Arc, + stats: SharedWriteStats, +} + +#[async_trait] +impl MessageHandler for IndexApplyHandler { + async fn handle(&mut self, message: TriggerIndexApply) -> Result<()> { + match apply_index_range(&self.cursors, message).await { + Ok(applied) => { + // A coalesced no-op indexes nothing (`rows_indexed == 0`); + // recording it would inflate the count and skew avg latency. + if applied.rows_indexed > 0 { + self.stats + .record_index_update(applied.duration, applied.rows_indexed); + } + Ok(()) + } + // An index apply cannot be partially rolled back, so a failure is + // terminal: poison, and let reopen rebuild the indexes from the WAL. + // See the note in `WalFlusher::flush_from_batch_store`. + Err(e) => { + self.wal_flusher.poison(&e); + Err(e) + } + } + } +} + struct WalFlushHandler { wal_flusher: Arc, /// MemTable-mode writer state, used to detect "frozen vs active" flushes /// via Arc::ptr_eq on the active batch_store. `None` when running in /// WAL-only mode (no MemTable, no frozen-vs-active distinction). memtable_state: Option>>, + /// How often to append in the background. `None` disables the ticker, leaving + /// the append size-triggered (and freeze/close-triggered) only. + flush_interval: Option, stats: SharedWriteStats, } @@ -2495,11 +2885,13 @@ impl WalFlushHandler { fn new( wal_flusher: Arc, memtable_state: Option>>, + flush_interval: Option, stats: SharedWriteStats, ) -> Self { Self { wal_flusher, memtable_state, + flush_interval, stats, } } @@ -2507,6 +2899,33 @@ impl WalFlushHandler { #[async_trait] impl MessageHandler for WalFlushHandler { + /// Append periodically in the background. + /// + /// This is what the flush interval was always supposed to mean. It routed to + /// a timer that was only ever *evaluated on the write path*, so it could add + /// a redundant trigger but never delay or batch one — with every durable put + /// triggering its own append, tuning the knob did nothing at all. + /// + /// The ticker exists to bound S3 API cost: an append is a PUT, billed per + /// call, and it is the only thing on this schedule. The index apply is not — + /// it is in-memory and free to batch, so it runs per-put on its own task. + fn tickers(&mut self) -> Vec<(Duration, MessageFactory)> { + // No interval => no ticker. A zero interval would panic in tokio. + let Some(interval) = self.flush_interval.filter(|d| !d.is_zero()) else { + return vec![]; + }; + // The tick names no store: `MessageFactory` is synchronous and cannot take + // the async state lock. `handle()` resolves it against the cursor. + vec![( + interval, + Box::new(|| TriggerWalFlush { + source: WalFlushSource::NextPending, + end_batch_position: 0, + done: None, + }), + )] + } + async fn handle(&mut self, message: TriggerWalFlush) -> Result<()> { let TriggerWalFlush { source, @@ -2514,6 +2933,16 @@ impl MessageHandler for WalFlushHandler { done, } = message; + // A timer tick names no store — resolve it now, at handle time. + let (source, end_batch_position) = match source { + WalFlushSource::NextPending => match self.resolve_next_pending().await { + Some(resolved) => resolved, + // Everything is already durable; the tick has nothing to do. + None => return Ok(()), + }, + other => (other, end_batch_position), + }; + let result = self.do_flush(source, end_batch_position).await; // Propagate the just-appended WAL entry position back into the @@ -2544,6 +2973,44 @@ impl MessageHandler for WalFlushHandler { } impl WalFlushHandler { + /// Pick the store the WAL still owes an append, oldest first. + /// + /// **WAL entries must be appended in global batch-position order for the + /// writer's lifetime.** `WalAppender::append` assigns each entry's position + /// from its own counter, in call order; replay walks those positions + /// ascending and assigns row positions in that order; and primary-key recency + /// is "newest visible row position wins". So append order fixes dedup order. + /// Append two memtables out of order and replay silently hands the dedup to + /// the *stale* row — corruption that survives the crash that caused it, and + /// that a full scan cannot see. + /// + /// Resolving to "the active memtable" would break exactly that: a tick + /// enqueued before a freeze is handled after it, resolves to the incoming + /// memtable, and appends its batches ahead of the outgoing memtable's tail. + /// So the target is a function of `durable`, not of when the timer fired. + /// + /// Safe to walk the frozen list because a store that still owes an append + /// cannot be swept: its L0 flush is blocked on the completion cell that only + /// that append fires. + async fn resolve_next_pending(&self) -> Option<(WalFlushSource, usize)> { + let state_lock = self.memtable_state.as_ref()?; + let state = state_lock.read().await; + let durable = self.wal_flusher.durable(); + + next_pending_store( + state + .frozen_memtables + .iter() + .map(|frozen| frozen.memtable.batch_store()), + state.memtable.batch_store(), + durable, + ) + .map(|store| { + let end = store.len(); + (WalFlushSource::BatchStore { batch_store: store }, end) + }) + } + /// Unified flush method for both active and frozen memtables and for /// WAL-only mode. /// @@ -2563,18 +3030,6 @@ impl WalFlushHandler { ) -> Result { let start = Instant::now(); - // Whether this flush actually updates any in-memory indexes — only - // a BatchStore source carrying a non-empty `IndexStore` does. Used - // to gate the `record_index_update` stat so WAL-only flushes don't - // pollute the index-update counters. - let has_indexes = matches!( - &source, - WalFlushSource::BatchStore { - indexes: Some(_), - .. - } - ); - // Early-out for BatchStore sources where the watermark already // covers the requested end position. Detection of "frozen flush" // requires the active memtable's batch_store; WAL-only handlers @@ -2582,8 +3037,7 @@ impl WalFlushHandler { // BatchStore source, so the early-out simplifies to the watermark // comparison. if let WalFlushSource::BatchStore { batch_store, .. } = &source { - let max_flushed = batch_store.max_flushed_batch_position(); - let flushed_up_to = max_flushed.map(|p| p + 1).unwrap_or(0); + let flushed_up_to = batch_store.local_end(self.wal_flusher.durable()); let is_frozen_flush = if let Some(state_lock) = &self.memtable_state { let state = state_lock.read().await; !Arc::ptr_eq(batch_store, &state.memtable.batch_store()) @@ -2608,12 +3062,6 @@ impl WalFlushHandler { self.stats .record_wal_flush(start.elapsed(), flush_result.wal_bytes); self.stats.record_wal_io(flush_result.wal_io_duration); - if has_indexes { - self.stats.record_index_update( - flush_result.index_update_duration, - flush_result.rows_indexed, - ); - } } Ok(flush_result) @@ -2629,6 +3077,9 @@ impl WalFlushHandler { struct MemTableFlushHandler { state: Arc>, flusher: Arc, + /// Source of the writer-global durability cursor, which the L0 flush asserts + /// covers the whole frozen memtable before it writes a generation. + wal_flusher: Arc, epoch: u64, /// Secondary index configs to rebuild on each flushed generation. When /// non-empty the handler flushes via [`MemTableFlusher::flush_with_indexes`] @@ -2643,9 +3094,11 @@ struct MemTableFlushHandler { } impl MemTableFlushHandler { + #[allow(clippy::too_many_arguments)] fn new( state: Arc>, flusher: Arc, + wal_flusher: Arc, epoch: u64, index_configs: Vec, stats: SharedWriteStats, @@ -2654,6 +3107,7 @@ impl MemTableFlushHandler { Self { state, flusher, + wal_flusher, epoch, index_configs, stats, @@ -2765,9 +3219,15 @@ impl MemTableFlushHandler { // dataset open when there are no indexes to build. The indexed // path's future is boxed to keep this async block's nesting // under the type-layout recursion limit. + // Read the durability cursor *after* the WAL-append completion above, + // not before: the append that makes this memtable durable is the very + // thing we just waited on, so a cursor sampled earlier would still be + // short of it and trip the flush precondition. + let durable = self.wal_flusher.durable(); + if self.index_configs.is_empty() { self.flusher - .flush(&memtable, self.epoch, covered_wal_entry_position) + .flush(&memtable, self.epoch, covered_wal_entry_position, durable) .await } else { Box::pin(self.flusher.flush_with_indexes( @@ -2775,6 +3235,7 @@ impl MemTableFlushHandler { self.epoch, &self.index_configs, covered_wal_entry_position, + durable, )) .await } @@ -3411,7 +3872,7 @@ mod tests { /// await), but the tombstone still lands in the in-memory tier. The delete /// analog of `test_put_no_wait_non_durable_returns_no_watcher`. #[tokio::test] - async fn test_shard_writer_delete_no_wait_non_durable_returns_no_watcher() { + async fn test_non_durable_delete_is_read_your_writes() { let (store, base_path, base_uri, _temp) = create_local_store().await; let schema = create_pk_test_schema(); let config = ShardWriterConfig { @@ -3432,7 +3893,10 @@ mod tests { .delete_no_wait(vec![id_only_keys(&[2])]) .await .unwrap(); - assert!(watcher.is_none(), "non-durable delete has nothing to await"); + // As with a put: a non-durable delete awaits its index apply, not an S3 + // round-trip. It is read-your-writes, just not durable. + let mut watcher = watcher.expect("a non-durable delete awaits its index apply"); + watcher.wait().await.unwrap(); // Tombstone landed in the in-memory tier (5 rows + 1 tombstone). assert_eq!(writer.memtable_stats().await.unwrap().row_count, 6); @@ -3495,7 +3959,6 @@ mod tests { ShardWriterConfig { shard_id, durable_write: false, - sync_indexed_write: true, manifest_scan_batch_size: 2, ..Default::default() } @@ -3820,9 +4283,8 @@ mod tests { shard_id: Uuid::new_v4(), shard_spec_id: 0, durable_write: false, - sync_indexed_write: false, max_wal_buffer_size: 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), max_memtable_size: 64 * 1024 * 1024, manifest_scan_batch_size: 2, ..Default::default() @@ -3863,9 +4325,8 @@ mod tests { shard_id: Uuid::new_v4(), shard_spec_id: 0, durable_write: true, - sync_indexed_write: false, max_wal_buffer_size: 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), max_memtable_size: 64 * 1024 * 1024, manifest_scan_batch_size: 2, ..Default::default() @@ -3892,7 +4353,7 @@ mod tests { } #[tokio::test] - async fn test_put_no_wait_non_durable_returns_no_watcher() { + async fn test_non_durable_put_is_read_your_writes() { let (store, base_path, base_uri, _temp_dir) = create_local_store().await; let schema = create_test_schema(); @@ -3900,9 +4361,8 @@ mod tests { shard_id: Uuid::new_v4(), shard_spec_id: 0, durable_write: false, - sync_indexed_write: false, max_wal_buffer_size: 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), max_memtable_size: 64 * 1024 * 1024, manifest_scan_batch_size: 2, ..Default::default() @@ -3915,7 +4375,21 @@ mod tests { let batch = create_test_batch(&schema, 0, 10); let (result, watcher) = writer.put_no_wait(vec![batch]).await.unwrap(); assert_eq!(result.batch_positions, 0..1); - assert!(watcher.is_none(), "non-durable put has nothing to await"); + + // A non-durable put still has something to await: its *index apply*. + // Skipping durability now costs the caller durability only — not + // visibility. Before the index apply was split off the WAL flush, a + // non-durable write was not read-your-writes at all: the row stayed + // invisible until some later flush happened to index it. + let mut watcher = watcher.expect("a non-durable put awaits its index apply"); + watcher.wait().await.unwrap(); + + let scanned = writer.scan().await.unwrap().try_into_batch().await.unwrap(); + assert_eq!( + scanned.num_rows(), + 10, + "a non-durable put must be readable as soon as it returns" + ); let stats = writer.memtable_stats().await.unwrap(); assert_eq!(stats.row_count, 10); @@ -3932,9 +4406,8 @@ mod tests { shard_id: Uuid::new_v4(), shard_spec_id: 0, durable_write: false, - sync_indexed_write: false, max_wal_buffer_size: 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), max_memtable_size: 64 * 1024 * 1024, manifest_scan_batch_size: 2, ..Default::default() @@ -3967,9 +4440,8 @@ mod tests { shard_id: Uuid::new_v4(), shard_spec_id: 0, durable_write: false, - sync_indexed_write: true, max_wal_buffer_size: 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), max_memtable_size: 64 * 1024 * 1024, manifest_scan_batch_size: 2, ..Default::default() @@ -4019,9 +4491,8 @@ mod tests { shard_id, shard_spec_id: 0, durable_write: false, - sync_indexed_write: true, max_wal_buffer_size: 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), max_memtable_size: 64 * 1024 * 1024, manifest_scan_batch_size: 2, ..Default::default() @@ -4108,9 +4579,8 @@ mod tests { shard_id: Uuid::new_v4(), shard_spec_id: 0, durable_write: false, - sync_indexed_write: false, max_wal_buffer_size: 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), max_memtable_size: 1024, // Very small - will trigger flush quickly manifest_scan_batch_size: 2, ..Default::default() @@ -4215,9 +4685,8 @@ mod tests { shard_id: Uuid::new_v4(), shard_spec_id: 0, durable_write: true, - sync_indexed_write: false, max_wal_buffer_size: 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), max_memtable_size: 64, manifest_scan_batch_size: 2, ..Default::default() @@ -4260,9 +4729,8 @@ mod tests { shard_id: Uuid::new_v4(), shard_spec_id: 0, durable_write: true, - sync_indexed_write: false, max_wal_buffer_size: 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), // Tiny size threshold — every batch crosses it. max_memtable_size: 64, manifest_scan_batch_size: 2, @@ -4327,9 +4795,8 @@ mod tests { shard_id, shard_spec_id: 0, durable_write: true, - sync_indexed_write: false, max_wal_buffer_size: 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), max_memtable_size: usize::MAX, max_unflushed_memtable_bytes: usize::MAX, manifest_scan_batch_size: 2, @@ -4407,7 +4874,6 @@ mod tests { shard_spec_id: 0, durable_write: false, enable_memtable, - sync_indexed_write: false, max_wal_buffer_size: usize::MAX, max_wal_flush_interval: None, max_wal_persist_retries: 0, @@ -4458,9 +4924,8 @@ mod tests { shard_id: Uuid::new_v4(), shard_spec_id: 0, durable_write: true, - sync_indexed_write: false, max_wal_buffer_size: 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), // Tiny size threshold so a few batches cross it. max_memtable_size: 1024, manifest_scan_batch_size: 2, @@ -4601,7 +5066,7 @@ mod tests { durable_write: true, enable_memtable: false, max_wal_buffer_size: 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), manifest_scan_batch_size: 2, ..Default::default() } @@ -4892,9 +5357,8 @@ mod tests { shard_id, shard_spec_id: 0, durable_write: true, - sync_indexed_write: false, max_wal_buffer_size: 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), max_memtable_size: 64 * 1024 * 1024, manifest_scan_batch_size: 2, ..Default::default() @@ -4918,8 +5382,9 @@ mod tests { } // A durable write whose WAL PUT keeps failing poisons the writer with a - // typed persistence failure; the next write fails fast with the same reason; - // and once storage heals, reopening replays the WAL and writes resume. + // typed persistence failure; the next write *and every read* fail fast with + // the same reason; and once storage heals, reopening replays the WAL and + // writes resume. #[tokio::test] async fn test_writer_poisons_on_persistence_failure_and_recovers_on_reopen() { let (store, base_path, controls) = failing_memory_store().await; @@ -4958,6 +5423,31 @@ mod tests { .await .unwrap_err(); assert_eq!(err.fence_reason(), Some(FenceReason::PersistenceFailure)); + + // ...and rejects *reads* too. Batch 0 was committed to the BatchStore + // before its WAL PUT failed, so a poisoned writer that still served + // reads would hand out a row that is not durable and that replay will + // not reproduce — a divergent snapshot. Mirrors SlateDB's + // `check_closed()` at the top of every read. + for reason in [ + writer.scan().await.err().and_then(|e| e.fence_reason()), + writer + .active_memtable_ref() + .await + .err() + .and_then(|e| e.fence_reason()), + writer + .in_memory_memtable_refs() + .await + .err() + .and_then(|e| e.fence_reason()), + ] { + assert_eq!(reason, Some(FenceReason::PersistenceFailure)); + } + + // Stats stay readable: this is what an operator (and the eviction path) + // inspects to decide what to do about the poisoned shard. + writer.memtable_stats().await.unwrap(); drop(writer); // Storage heals: reopening replays the WAL and accepts writes again. @@ -4971,32 +5461,249 @@ mod tests { .unwrap(); } - /// Replay-on-open recovers durable WAL entries that were never flushed - /// to a Lance generation. Setup: writer A durably writes batches, drops - /// without close (so MemTable freeze never runs); writer B reopens and - /// must see A's rows in its MemTable scan. + /// A doomed open must fail on local validation *before* it claims the epoch, + /// so it cannot fence the healthy writer already serving the shard. The + /// durable-write/flush-interval and index-config checks used to run *after* + /// `claim_epoch` (and, for a successor, after `write_fence_sentinel`), so a + /// rejected open still bumped the stored epoch and fenced the incumbent. #[tokio::test] - async fn test_memtable_replay_recovers_unflushed_writes() { + async fn test_doomed_open_does_not_fence_incumbent() { let (store, base_path, base_uri, _temp_dir) = create_local_store().await; let schema = schema_with_pk(); let shard_id = Uuid::new_v4(); - // Writer A: write two durable batches, then drop without close. - // The WAL files persist; the in-memory MemTable does not. - { - let writer_a = ShardWriter::open( - store.clone(), - base_path.clone(), - base_uri.clone(), - memtable_config_with_pk(shard_id), - schema.clone(), - vec![], - ) - .await - .unwrap(); - writer_a - .put(vec![create_test_batch(&schema, 0, 5)]) - .await + let writer_a = ShardWriter::open( + store.clone(), + base_path.clone(), + base_uri.clone(), + memtable_config_with_pk(shard_id), + schema.clone(), + vec![], + ) + .await + .unwrap(); + writer_a + .put(vec![create_test_batch(&schema, 0, 1)]) + .await + .unwrap(); + + // A durable writer with no flush interval is rejected — the put path no + // longer triggers its own append, so it would deadlock. On the old path + // this rejection landed only after the epoch had already been claimed. + let doomed_config = ShardWriterConfig { + max_wal_flush_interval: None, + ..memtable_config_with_pk(shard_id) + }; + let err = ShardWriter::open( + store.clone(), + base_path.clone(), + base_uri.clone(), + doomed_config, + schema.clone(), + vec![], + ) + .await + .map(|_| ()) + .expect_err("durable_write with no flush interval must be rejected"); + assert!( + err.to_string().contains("max_wal_flush_interval"), + "unexpected error: {err}" + ); + + // The incumbent is untouched: not fenced, still accepting writes. + writer_a.check_fenced().await.unwrap(); + writer_a + .put(vec![create_test_batch(&schema, 1, 1)]) + .await + .unwrap(); + writer_a.close().await.unwrap(); + } + + /// A failed dispatch during `freeze_memtable` must not drop the outgoing + /// table's rows from the read view. The active memtable is replaced before + /// the WAL-flush and index-apply sends; a send that failed (background tasks + /// gone) used to return before the outgoing table was retained in + /// `frozen_memtables`, so its accepted rows silently vanished — a scan + /// returned 0 rows with no error. The writer must instead retain the table + /// and poison, so reads fail fast rather than serve a divergent snapshot. + #[tokio::test] + async fn test_freeze_dispatch_failure_retains_rows_and_poisons() { + let (store, base_path, base_uri, _temp_dir) = create_local_store().await; + let schema = schema_with_pk(); + let shard_id = Uuid::new_v4(); + + // Non-durable + no ticker: the put is read-your-writes (waits for its + // index apply) but nothing is WAL-flushed, so the freeze below still owes + // a WAL append. + let config = ShardWriterConfig { + durable_write: false, + max_wal_flush_interval: None, + ..memtable_config_with_pk(shard_id) + }; + let writer = ShardWriter::open(store, base_path, base_uri, config, schema.clone(), vec![]) + .await + .unwrap(); + + writer + .put(vec![create_test_batch(&schema, 0, 10)]) + .await + .unwrap(); + assert_eq!(writer.memtable_stats().await.unwrap().row_count, 10); + + // Tear the background tasks down out from under the writer, so the + // freeze's dispatch sends hit closed channels. + writer.abort().await.unwrap(); + + let err = writer + .force_seal_active() + .await + .expect_err("force_seal_active must surface the failed dispatch"); + assert!( + err.to_string().contains("channel closed"), + "unexpected error: {err}" + ); + + // The failure poisoned the writer: reads fail fast instead of returning a + // silent zero-row snapshot of a shard whose rows were dropped. + assert!( + writer.scan().await.is_err(), + "a poisoned writer must reject reads, not serve a divergent snapshot" + ); + assert!(writer.in_memory_memtable_refs().await.is_err()); + } + + /// A WAL holding more batches than one memtable's capacity must reopen. + /// + /// One memtable holds at most `max_memtable_batches` batches, but a WAL is + /// unbounded, so replay has to rotate — seal the full memtable, start a fresh + /// one — exactly as the live write path does. Before, replay stuffed + /// everything into a single memtable and `open()` failed outright with + /// "MemTable batch store is full", leaving the shard permanently unopenable. + #[tokio::test] + async fn test_replay_rotates_when_wal_exceeds_one_memtable() { + let (store, base_path, base_uri, _temp_dir) = create_local_store().await; + let schema = schema_with_pk(); + let shard_id = Uuid::new_v4(); + + const N: i32 = 8; + + // Writer A has a *large* capacity, so its eight one-batch puts all land in + // a single memtable and it never freezes or flushes a generation of its + // own. Dropping it without close leaves an eight-entry WAL and no + // generations — a WAL that no single small memtable could hold. + let writer_a_config = ShardWriterConfig { + max_memtable_batches: 1000, + ..memtable_config_with_pk(shard_id) + }; + // Writer B has a *two-batch* capacity, so replaying that eight-entry WAL is + // exactly what must rotate. Keeping the configs distinct isolates replay + // rotation from the live rotation writer A would otherwise do concurrently. + let config = ShardWriterConfig { + max_memtable_batches: 2, + ..memtable_config_with_pk(shard_id) + }; + + { + let writer_a = ShardWriter::open( + store.clone(), + base_path.clone(), + base_uri.clone(), + writer_a_config, + schema.clone(), + vec![], + ) + .await + .unwrap(); + for id in 0..N { + writer_a + .put(vec![create_test_batch(&schema, id, 1)]) + .await + .unwrap(); + } + // Drop without close: only the WAL survives, and it holds more batches + // than writer B's memtable can. + } + + // Total rows across the active memtable plus every flushed generation. + // Distinct ids, so no cross-generation dedup — a plain sum is exact. + async fn total_rows(writer: &ShardWriter, base_uri: &str, shard_id: Uuid) -> usize { + let mut rows = writer.memtable_stats().await.unwrap().row_count; + let manifest = writer.manifest().await.unwrap().unwrap(); + for fg in &manifest.flushed_generations { + let gen_uri = format!("{}/_mem_wal/{}/{}", base_uri, shard_id, fg.path); + let dataset = crate::Dataset::open(&gen_uri).await.unwrap(); + rows += dataset.count_rows(None).await.unwrap(); + } + rows + } + + // Reopen. This used to fail with a full-batch-store error. + let writer_b = ShardWriter::open( + store.clone(), + base_path.clone(), + base_uri.clone(), + config.clone(), + schema.clone(), + vec![], + ) + .await + .expect("a WAL larger than one memtable must still reopen"); + + // Rotation produced sealed memtables, and replay flushed each to a Lance + // generation rather than holding it in memory or leaving it in the WAL. + let manifest = writer_b.manifest().await.unwrap().unwrap(); + assert!( + !manifest.flushed_generations.is_empty(), + "replay must have sealed and flushed at least one full memtable" + ); + + // Every row survived, split between the flushed generations and the + // active (partial) memtable. + assert_eq!( + total_rows(&writer_b, &base_uri, shard_id).await as i32, + N, + "every replayed row must be durable, across generations and the active memtable" + ); + writer_b.close().await.unwrap(); + + // Because the sealed memtables were flushed, the manifest's replay cursor + // advanced past their WAL entries — so a second reopen replays only the + // tail and still accounts for every row. The WAL truncates across reopens + // rather than growing without bound. + let writer_c = + ShardWriter::open(store, base_path, base_uri.clone(), config, schema, vec![]) + .await + .unwrap(); + assert_eq!(total_rows(&writer_c, &base_uri, shard_id).await as i32, N); + writer_c.close().await.unwrap(); + } + + /// Replay-on-open recovers durable WAL entries that were never flushed + /// to a Lance generation. Setup: writer A durably writes batches, drops + /// without close (so MemTable freeze never runs); writer B reopens and + /// must see A's rows in its MemTable scan. + #[tokio::test] + async fn test_memtable_replay_recovers_unflushed_writes() { + let (store, base_path, base_uri, _temp_dir) = create_local_store().await; + let schema = schema_with_pk(); + let shard_id = Uuid::new_v4(); + + // Writer A: write two durable batches, then drop without close. + // The WAL files persist; the in-memory MemTable does not. + { + let writer_a = ShardWriter::open( + store.clone(), + base_path.clone(), + base_uri.clone(), + memtable_config_with_pk(shard_id), + schema.clone(), + vec![], + ) + .await + .unwrap(); + writer_a + .put(vec![create_test_batch(&schema, 0, 5)]) + .await .unwrap(); writer_a .put(vec![create_test_batch(&schema, 100, 3)]) @@ -5033,6 +5740,426 @@ mod tests { writer_b.close().await.unwrap(); } + /// Replayed batches are already WAL-durable, so the first flush after a + /// reopen must not re-append them to the WAL or re-insert them into the + /// indexes. Before replay stamped the durability cursor it stayed at + /// "nothing flushed", so the next flush re-covered `[0, end)`: it appended + /// the already-durable rows a second time *and* re-indexed them. None of + /// the in-memory indexes is idempotent, so an indexed PK lookup returned + /// the row twice while a full scan still looked healthy. + #[tokio::test] + async fn test_replay_does_not_reappend_or_reindex() { + use crate::dataset::mem_wal::wal::WalTailer; + + let (store, base_path, base_uri, _temp_dir) = create_local_store().await; + let schema = schema_with_pk(); + let shard_id = Uuid::new_v4(); + + // Writer A: two durable batches (5 + 3 = 8 rows), dropped without close. + { + let writer_a = ShardWriter::open( + store.clone(), + base_path.clone(), + base_uri.clone(), + memtable_config_with_pk(shard_id), + schema.clone(), + vec![], + ) + .await + .unwrap(); + writer_a + .put(vec![create_test_batch(&schema, 0, 5)]) + .await + .unwrap(); + writer_a + .put(vec![create_test_batch(&schema, 100, 3)]) + .await + .unwrap(); + } + + // Writer B reopens and replays A's two batches. + let writer_b = ShardWriter::open( + store.clone(), + base_path.clone(), + base_uri, + memtable_config_with_pk(shard_id), + schema.clone(), + vec![], + ) + .await + .unwrap(); + + let stats = writer_b.memtable_stats().await.unwrap(); + assert_eq!(stats.batch_count, 2); + assert_eq!( + stats.durable_batch_count, 2, + "replayed batches came from the WAL, so the durability cursor must already cover them" + ); + + // One more durable put. Its flush must cover only the new batch. + writer_b + .put(vec![create_test_batch(&schema, 200, 2)]) + .await + .unwrap(); + + let tailer = WalTailer::new(store, base_path, shard_id); + let first = tailer.first_position().await.unwrap(); + let next = tailer.next_position().await.unwrap(); + let mut wal_rows = 0; + for position in first..next { + if let Some(entry) = tailer.read_entry(position).await.unwrap() { + wal_rows += entry.batches.iter().map(|b| b.num_rows()).sum::(); + } + } + assert_eq!( + wal_rows, 10, + "WAL must hold 8 replayed + 2 new rows; a re-covering flush re-appends the replayed 8" + ); + + // The indexed arm must not see a replayed row twice. + let mut scanner = writer_b.scan().await.unwrap(); + scanner.filter("id = 0").unwrap(); + let hit = scanner.try_into_batch().await.unwrap(); + assert_eq!( + hit.num_rows(), + 1, + "indexed PK lookup returned the replayed row more than once" + ); + + let all = writer_b + .scan() + .await + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(all.num_rows(), 10); + + writer_b.close().await.unwrap(); + } + + /// A non-durable put is readable through the **index-backed** arms the moment + /// it returns, not just through a full scan. + /// + /// This is what splitting the index apply off the WAL flush buys. Before, the + /// index apply only ran as one arm of the flush, so with `durable_write: + /// false` nothing triggered it on the put path at all: the row sat in the + /// batch store, unindexed, until some later flush happened along. A full scan + /// (which reads the batch store directly) could still find it, while every + /// index-accelerated query could not — the tiers disagreed. Now the apply is + /// triggered per-put in both modes, so `put` returning means "indexed". + #[tokio::test] + async fn test_non_durable_put_is_visible_through_the_index() { + let (store, base_path, base_uri, _temp_dir) = create_local_store().await; + let schema = schema_with_pk(); + let shard_id = Uuid::new_v4(); + + let config = ShardWriterConfig { + durable_write: false, + ..memtable_config_with_pk(shard_id) + }; + let writer = ShardWriter::open(store, base_path, base_uri, config, schema.clone(), vec![]) + .await + .unwrap(); + + writer + .put(vec![create_test_batch(&schema, 0, 5)]) + .await + .unwrap(); + + // The indexed PK lookup must find it — this is the arm that saw nothing + // before, because the index apply had never run. + let mut scanner = writer.scan().await.unwrap(); + scanner.filter("id = 3").unwrap(); + let hit = scanner.try_into_batch().await.unwrap(); + assert_eq!( + hit.num_rows(), + 1, + "an index-backed lookup must see a non-durable put as soon as it returns" + ); + + // ...and so must the unindexed full scan, i.e. the tiers agree. + let all = writer.scan().await.unwrap().try_into_batch().await.unwrap(); + assert_eq!(all.num_rows(), 5); + + writer.close().await.unwrap(); + } + + /// The durability cursor is writer-global, so a put into a *post-rotation* + /// memtable must still wait for its own WAL append. + /// + /// Batch positions restart at 0 in every memtable, but the durability watch + /// channel spans the writer's whole life and is never reset. When the put + /// path targeted a memtable-local position, the first N puts into every + /// memtable after the first were already "satisfied" by the *previous* + /// memtable's N appends: they acked instantly, with no WAL append, and + /// `durable_write: true` silently degraded to non-durable. Worse, the next + /// append then sent a *smaller* value, walking the watermark backwards and + /// hanging any watcher still waiting on the old, higher one. + /// + /// The cursor is now a writer-global exclusive count and every target is + /// lifted through the store's `global_offset`, so it only ever moves forward + /// and a post-rotation put can only be acked by its own append. + #[tokio::test] + async fn test_durable_ack_after_rotation_requires_its_own_wal_append() { + use crate::dataset::mem_wal::wal::WalTailer; + + let (store, base_path, base_uri, _temp_dir) = create_local_store().await; + let schema = schema_with_pk(); + let shard_id = Uuid::new_v4(); + + // A two-batch memtable, so the third put forces a freeze + rotation. + let config = ShardWriterConfig { + max_memtable_batches: 2, + ..memtable_config_with_pk(shard_id) + }; + + let writer = ShardWriter::open( + store.clone(), + base_path.clone(), + base_uri, + config, + schema.clone(), + vec![], + ) + .await + .unwrap(); + + // Fill and rotate the first memtable. + for i in 0..3 { + writer + .put(vec![create_test_batch(&schema, i * 5, 5)]) + .await + .unwrap(); + } + let stats = writer.memtable_stats().await.unwrap(); + assert!( + stats.global_offset > 0, + "expected a rotation; the active memtable is still the writer's first" + ); + + // Every batch acked so far must be genuinely durable, and the cursor must + // cover the active memtable's entire prefix rather than lagging inside it. + let stats = writer.memtable_stats().await.unwrap(); + assert!( + stats.durable_batch_count >= stats.global_offset + stats.batch_count, + "durable_write acked a put the WAL never received: durable={} but the active \ + memtable spans [{}, {})", + stats.durable_batch_count, + stats.global_offset, + stats.global_offset + stats.batch_count + ); + + // And the WAL really holds every row we acked (3 puts x 5 rows). + writer.close().await.unwrap(); + let tailer = WalTailer::new(store, base_path, shard_id); + let first = tailer.first_position().await.unwrap(); + let next = tailer.next_position().await.unwrap(); + let mut wal_rows = 0; + for position in first..next { + if let Some(entry) = tailer.read_entry(position).await.unwrap() { + wal_rows += entry.batches.iter().map(|b| b.num_rows()).sum::(); + } + } + assert_eq!( + wal_rows, 15, + "every acked row must be in the WAL; a post-rotation put that acked without an \ + append would leave rows missing" + ); + } + + /// A background tick must append the **oldest** store that still owes the WAL, + /// never "whatever memtable is active". + /// + /// A tick carries no store: it is enqueued by a timer and resolved when it is + /// handled. So a tick enqueued before a freeze is handled *after* it, and + /// resolving to the active memtable would append the incoming memtable's + /// batches ahead of the outgoing memtable's tail. WAL entry positions are + /// assigned in append-call order, replay walks them ascending, row positions + /// follow, and primary-key recency is "newest visible row position wins" — so + /// that inverts dedup after a crash, handing the key to the stale row. It + /// survives the crash that caused it, and a full scan cannot see it. + #[test] + fn test_next_pending_store_picks_the_oldest_owing_an_append() { + let schema = create_test_schema(); + + // A frozen store of 2 batches at coordinate 0, and the active store that + // rotated in behind it at coordinate 2. + let frozen = Arc::new(BatchStore::with_capacity(4)); + frozen.append(create_test_batch(&schema, 0, 1)).unwrap(); + frozen.append(create_test_batch(&schema, 1, 1)).unwrap(); + let active = Arc::new(BatchStore::with_capacity_at(4, 2)); + active.append(create_test_batch(&schema, 2, 1)).unwrap(); + + let frozen_list = || std::iter::once(Arc::clone(&frozen)); + + // Nothing durable: the frozen store owes the oldest append, so it wins — + // even though the active memtable also has un-appended batches. + let picked = next_pending_store(frozen_list(), Arc::clone(&active), 0).unwrap(); + assert!( + Arc::ptr_eq(&picked, &frozen), + "the outgoing memtable's tail must be appended before the incoming one's head" + ); + + // Still true partway through the frozen store. + let picked = next_pending_store(frozen_list(), Arc::clone(&active), 1).unwrap(); + assert!(Arc::ptr_eq(&picked, &frozen)); + + // Once the frozen store is fully durable, the active one is next. + let picked = next_pending_store(frozen_list(), Arc::clone(&active), 2).unwrap(); + assert!(Arc::ptr_eq(&picked, &active)); + + // Everything durable: nothing to do. + assert!(next_pending_store(frozen_list(), Arc::clone(&active), 3).is_none()); + } + + /// A durable writer with no flush ticker cannot make progress, so `open()` + /// rejects it rather than letting a put block forever. + #[tokio::test] + async fn test_open_rejects_durable_write_without_a_ticker() { + let (store, base_path, base_uri, _temp_dir) = create_local_store().await; + let schema = schema_with_pk(); + let shard_id = Uuid::new_v4(); + + let config = ShardWriterConfig { + durable_write: true, + max_wal_flush_interval: None, + ..memtable_config_with_pk(shard_id) + }; + + let Err(err) = ShardWriter::open(store, base_path, base_uri, config, schema, vec![]).await + else { + panic!("durable_write with no ticker must be rejected"); + }; + assert!( + err.to_string().contains("max_wal_flush_interval"), + "the error must name the knob, got: {err}" + ); + } + + /// WAL entries must be appended in global batch-position order across a + /// memtable rotation, because append order *is* primary-key recency order. + /// + /// `WalAppender::append` assigns each entry's position from its own counter, + /// in call order. Replay walks those positions ascending and assigns row + /// positions in that order. Primary-key recency is "newest visible row + /// position wins". So an out-of-order append silently inverts dedup after a + /// crash — the stale row wins — and a full scan cannot see it. + /// + /// The hazard is the background ticker. If it resolved to "whatever memtable + /// is active" rather than to the oldest store still owing an append, a tick + /// enqueued before a freeze but handled after it would append the *incoming* + /// memtable's batches ahead of the outgoing memtable's tail. So the target is + /// resolved from the durability cursor, not from wall-clock timing. + #[tokio::test] + async fn test_wal_append_order_preserves_pk_recency_across_rotation() { + use crate::dataset::mem_wal::wal::WalTailer; + + let (store, base_path, base_uri, _temp_dir) = create_local_store().await; + let schema = schema_with_pk(); + let shard_id = Uuid::new_v4(); + + // Two batches per memtable, so the second put fills it and rotates. + let config = ShardWriterConfig { + max_memtable_batches: 2, + ..memtable_config_with_pk(shard_id) + }; + + let writer = ShardWriter::open( + store.clone(), + base_path.clone(), + base_uri, + config, + schema.clone(), + vec![], + ) + .await + .unwrap(); + + // Memtable 1: ids 7 then 20 (the second fills it and triggers the freeze). + writer + .put(vec![create_test_batch(&schema, 7, 1)]) + .await + .unwrap(); + writer + .put(vec![create_test_batch(&schema, 20, 1)]) + .await + .unwrap(); + // Memtable 2 overwrites id=7 with a newer row. Its append must land in the + // WAL *after* memtable 1's, or replay would resolve id=7 to the stale copy. + writer + .put(vec![create_test_batch(&schema, 30, 1)]) + .await + .unwrap(); + writer.close().await.unwrap(); + + // Walk the WAL in entry order and collect the ids as replay would see them. + let tailer = WalTailer::new(store, base_path, shard_id); + let first = tailer.first_position().await.unwrap(); + let next = tailer.next_position().await.unwrap(); + let mut ids: Vec = Vec::new(); + for position in first..next { + let Some(entry) = tailer.read_entry(position).await.unwrap() else { + continue; + }; + for batch in &entry.batches { + let column = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + ids.extend((0..column.len()).map(|i| column.value(i))); + } + } + + assert_eq!( + ids, + vec![7, 20, 30], + "WAL entries must follow global batch-position order; memtable 1's rows must \ + precede memtable 2's, or replay inverts primary-key recency" + ); + } + + /// An index config that disagrees with the schema fails `open()` outright. + /// It must not be allowed to accept writes: the insert would fail + /// deterministically on every batch, including batches replayed from the + /// WAL, so once a row was durable the shard could never reopen. Before this + /// check, an FTS index on a non-Utf8 column silently indexed nothing and the + /// shard reported healthy. + #[tokio::test] + async fn test_open_rejects_index_config_that_disagrees_with_schema() { + let (store, base_path, base_uri, _temp_dir) = create_local_store().await; + let schema = schema_with_pk(); + let shard_id = Uuid::new_v4(); + + // `id` is Int32, not a string column. + let bad_fts = MemIndexConfig::Fts(FtsIndexConfig::new( + "bad_fts".to_string(), + 0, + "id".to_string(), + )); + + let Err(err) = ShardWriter::open( + store, + base_path, + base_uri, + memtable_config_with_pk(shard_id), + schema, + vec![bad_fts], + ) + .await + else { + panic!("open must reject an FTS index on a non-Utf8 column"); + }; + + let message = err.to_string(); + assert!( + message.contains("bad_fts") && message.contains("Utf8"), + "error must name the index and the constraint, got: {message}" + ); + } + /// Replay is a no-op on a fresh shard: the MemTable starts empty. #[tokio::test] async fn test_memtable_replay_no_op_on_fresh_shard() { @@ -5481,6 +6608,55 @@ mod tests { ); } + #[tokio::test] + async fn test_memtable_stats_record_index_update() { + // MemTable mode with a BTree index: index application runs on its own + // task and must record an index-update stat. Regression for the stat + // silently reading zero after index apply moved off the WAL-flush path. + let (store, base_path, base_uri, _temp_dir) = create_local_store().await; + let schema = create_pk_test_schema(); + let index_configs = vec![MemIndexConfig::BTree(BTreeIndexConfig { + name: "id_idx".to_string(), + field_id: 0, + column: "id".to_string(), + })]; + + let writer = ShardWriter::open( + store, + base_path, + base_uri, + flush_test_config(Uuid::new_v4()), + schema.clone(), + index_configs, + ) + .await + .unwrap(); + writer + .put(vec![create_test_batch(&schema, 0, 3)]) + .await + .unwrap(); + + let stats_handle = writer.stats_handle(); + // `close()` drains the index-apply task, so the apply is settled here. + writer.close().await.unwrap(); + + let snapshot = stats_handle.snapshot(); + assert!( + snapshot.index_update_count >= 1, + "the index apply must record an index-update stat, got {}", + snapshot.index_update_count + ); + assert_eq!( + snapshot.index_update_rows, 3, + "every indexed row must be counted exactly once, got {}", + snapshot.index_update_rows + ); + assert!( + snapshot.avg_index_update_latency().is_some(), + "a recorded index update must expose an average latency" + ); + } + #[tokio::test] async fn test_force_seal_active_and_wait_for_flush_drain() { let (store, base_path, base_uri, _temp_dir) = create_local_store().await; @@ -5492,9 +6668,8 @@ mod tests { shard_id: Uuid::new_v4(), shard_spec_id: 0, durable_write: false, - sync_indexed_write: false, max_wal_buffer_size: 64 * 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), max_memtable_size: 64 * 1024 * 1024, manifest_scan_batch_size: 2, ..Default::default() @@ -5547,9 +6722,8 @@ mod tests { shard_id: Uuid::new_v4(), shard_spec_id: 0, durable_write: false, - sync_indexed_write: false, max_wal_buffer_size: 64 * 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), max_memtable_size: 64 * 1024 * 1024, manifest_scan_batch_size: 2, ..Default::default() @@ -5602,9 +6776,8 @@ mod tests { shard_id: Uuid::new_v4(), shard_spec_id: 0, durable_write: false, - sync_indexed_write: false, max_wal_buffer_size: 64 * 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), max_memtable_size: 64 * 1024 * 1024, manifest_scan_batch_size: 2, // Short grace so the sweep is observable without a slow test. @@ -5662,9 +6835,8 @@ mod tests { shard_id: Uuid::new_v4(), shard_spec_id: 0, durable_write: false, - sync_indexed_write: false, max_wal_buffer_size: 64 * 1024 * 1024, - max_wal_flush_interval: None, + max_wal_flush_interval: Some(Duration::from_millis(10)), max_memtable_size: 64 * 1024 * 1024, manifest_scan_batch_size: 2, frozen_memtable_grace: Duration::ZERO, @@ -5973,9 +7145,7 @@ mod shard_writer_tests { Some("100") ); // Every tunable field is present. - assert!(defaults.contains_key("sync_indexed_write")); assert!(defaults.contains_key("enable_memtable")); - assert!(defaults.contains_key("async_index_interval_ms")); // add_writer_config_default records arbitrary keys. assert_eq!( defaults.get("custom_knob").map(String::as_str), @@ -6187,9 +7357,7 @@ mod shard_writer_tests { .expect("Failed to initialize MemWAL"); let shard_id = Uuid::new_v4(); - let config = ShardWriterConfig::new(shard_id) - .with_durable_write(true) - .with_sync_indexed_write(true); + let config = ShardWriterConfig::new(shard_id).with_durable_write(true); let writer = dataset .mem_wal_writer(shard_id, config) .await @@ -6527,9 +7695,7 @@ mod shard_writer_tests { // Create shard writer let shard_id = Uuid::new_v4(); - let config = ShardWriterConfig::new(shard_id) - .with_durable_write(false) - .with_sync_indexed_write(false); + let config = ShardWriterConfig::new(shard_id).with_durable_write(false); let writer = dataset .mem_wal_writer(shard_id, config) @@ -6587,9 +7753,7 @@ mod shard_writer_tests { .expect("Failed to initialize MemWAL"); let shard_id = Uuid::new_v4(); - let config = ShardWriterConfig::new(shard_id) - .with_durable_write(true) - .with_sync_indexed_write(true); + let config = ShardWriterConfig::new(shard_id).with_durable_write(true); let writer = dataset .mem_wal_writer(shard_id, config) @@ -6728,9 +7892,7 @@ mod shard_writer_tests { // Create shard writer with default config let shard_id = Uuid::new_v4(); - let config = ShardWriterConfig::new(shard_id) - .with_durable_write(false) - .with_sync_indexed_write(false); + let config = ShardWriterConfig::new(shard_id).with_durable_write(false); let writer = dataset .mem_wal_writer(shard_id, config) @@ -6807,7 +7969,6 @@ mod shard_writer_tests { let shard_id = Uuid::new_v4(); let config = ShardWriterConfig::new(shard_id) .with_durable_write(true) // Ensure WAL files are written - .with_sync_indexed_write(true) .with_max_memtable_size(50 * 1024) // 50KB - triggers flush after ~8 batches .with_max_wal_buffer_size(10 * 1024) // 10KB WAL buffer .with_max_wal_flush_interval(Duration::from_millis(50)); // Fast flush @@ -6931,9 +8092,13 @@ mod shard_writer_tests { // Re-open dataset and create new writer to verify recovery let dataset = Dataset::open(&uri).await.expect("Failed to reopen dataset"); let new_shard_id = Uuid::new_v4(); - let new_config = ShardWriterConfig::new(new_shard_id) - .with_durable_write(false) - .with_sync_indexed_write(true); + // `durable_write(true)` so the put waits for its flush, which is what + // currently publishes the rows. A non-durable put is *not* yet + // read-your-writes: the index apply is welded to the WAL flush, so the + // rows stay invisible until the next flush. This test used to pass with + // `durable_write(false)` only because an un-advanced cursor of 0 was + // misread as "batch 0 is visible" — it was asserting the dirty read. + let new_config = ShardWriterConfig::new(new_shard_id).with_durable_write(true); let new_writer = dataset .mem_wal_writer(new_shard_id, new_config)