fix(merge_insert): deduplicate indexed candidates - #8176
Conversation
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
The change restores stream-wide set semantics for indexed candidates and fixes the batch-dependent false ambiguity without weakening genuine duplicate-source handling. The remaining concern is the lifetime of the retained candidate state on large indexed merges.
| Self::map_batch(lookups, dataset, deletion_mask, batch, metrics).await | ||
| } | ||
| }); | ||
| let mut emitted_row_addrs = RoaringTreemap::new(); |
There was a problem hiding this comment.
This treemap retains every unique candidate address until the stream ends, but the allocation is not registered with the DataFusion memory pool. A broad or many-batch indexed merge can therefore add O(unique target candidates) untracked memory on top of the existing per-batch masks and arrays. Please consider memory-accounting this state with a bounded failure under pressure, or use a spillable distinct if large candidate streams are expected.
There was a problem hiding this comment.
This revision resolves the prior lifetime-accounting concern and preserves the verified cross-batch correctness fix. The retained set is now tracked at batch boundaries; if the pool is intended to be a hard peak-memory bound, pre-reservation remains the safer contract.
| .copied() | ||
| .filter(|row_addr| self.emitted.insert(*row_addr)) | ||
| .collect(); | ||
| if let Err(error) = self.reservation.try_resize(self.emitted.deep_size_of()) { |
There was a problem hiding this comment.
The reservation is checked only after the complete candidate batch has been inserted and the unseen array allocated, so a broad batch can exceed the pool before returning ResourcesExhausted. Rollback restores logical membership, but the accounting is not a strict peak-memory bound. If hard boundedness is required, reserve conservatively before mutation and shrink to the measured size, or account insertion in bounded chunks. A finite-pool regression covering successful growth, later rejection, preserved prior candidates, and release on drop would also make the intended contract explicit.
There was a problem hiding this comment.
❌ Gate recommendation: request changes.
Empty index results are valid and must remain successful empty candidate batches. Account the retained map baseline before comparing measured growth so the new memory bound does not reject no-match indexed merges.
| } | ||
|
|
||
| let measured_size = self.emitted.deep_size_of(); | ||
| if measured_size > provisional_size { |
There was a problem hiding this comment.
An initial empty candidate batch always enters this error branch: both reservation.size() and unaccounted_candidates are zero, but RowAddrTreeMap::deep_size_of() includes size_of_val(self). A normal no-match probe therefore returns ResourcesExhausted even when the pool has capacity. Include the empty-map baseline in the provisional reservation (or return an empty array before this comparison) and keep a first-empty regression.
Reproducer
#[test]
fn test_map_index_candidates_accept_empty_first_batch() {
let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1024));
let reservation =
MemoryConsumer::new(MAP_INDEX_CANDIDATES_MEMORY_CONSUMER).register(&pool);
let mut candidates = DistinctRowAddrs::new(reservation);
let empty = candidates
.retain_unseen(&UInt64Array::from(Vec::<u64>::new()))
.unwrap();
assert!(empty.is_empty());
}cargo test -p lance test_map_index_candidates_accept_empty_first_batch -- --test-threads=1 failed on d945a74bca1b7185731d375df830e7da156a0e4f with ResourcesExhausted("MapIndexExecCandidates batch exceeded its 256-byte per-candidate reservation").
There was a problem hiding this comment.
This revision fixes the empty-candidate regression while preserving cross-batch deduplication and finite-pool rollback. The retained-state reservation is suitable as cooperative pool accounting; avoid relying on it as an allocator-level peak cap without representation-threshold evidence.
| /// Conservative upper bound on the [`DeepSizeOf`] growth from inserting one | ||
| /// address: a new B-tree entry, Roaring container metadata, and its payload. | ||
| /// The reservation is shrunk to the measured size after each input batch. | ||
| const ROW_ADDR_INSERT_RESERVATION_BYTES: usize = 256; |
There was a problem hiding this comment.
This constant conservatively bounds the repository’s approximate DeepSizeOf metric, not demonstrated allocator-level peak usage: RowAddrSelection measures a Roaring bitmap with serialized_size(), and the provisional reservation is shrunk while the returned output array remains alive. The current mechanism is reasonable as cooperative pool accounting, but it should not be treated as a hard physical-memory cap unless representation-threshold and output-lifetime measurements establish that contract.
Xuanwo
left a comment
There was a problem hiding this comment.
Please also add an end-to-end python test. I'm not sure if I fully understand the bug you want to fix.
| pub index_name: String, | ||
| } | ||
|
|
||
| const MAP_INDEX_CANDIDATES_MEMORY_CONSUMER: &str = "MapIndexExecCandidates"; |
There was a problem hiding this comment.
What's this for? Can we avoid this?
There was a problem hiding this comment.
Addressed in b268921. Kept the reservation because cross-batch deduplication retains every unique candidate until the stream ends; removing it would leave O(unique candidates) outside the task pool. Added an inline comment distinguishing this stream-lifetime state from per-batch reservations.
There was a problem hiding this comment.
✅ Gate recommendation: approve.
The requested Python end-to-end regression now preserves two source batches, exercises both scalar indexes, and verifies the correct three updates without duplicated target rows. Production behavior and the previously verified accounting contract are unchanged.
|
Addressed in b268921. Added a Python end-to-end regression that passes a two-batch RecordBatchReader through the indexed composite-key merge and verifies the three intended updates complete without a false ambiguity error. |
Summary
Root cause
MapIndexExec produces a set of matching row addresses for each source batch, but its output stream did not de-duplicate addresses between batches. Composite-key IsIn probes can over-match because their per-column values are not tuple-correlated, so two batches could emit the same target address. Reading that target row twice made the exact downstream join appear ambiguous.
The fix retains previously emitted row addresses in a compressed RowAddrTreeMap for the lifetime of the index-mapping stream. Before mutating the retained set or allocating its output for a batch, MapIndexExecCandidates reserves a conservative per-candidate allowance through the DataFusion task pool, then shrinks the retained-state reservation to the approximate DeepSizeOf measurement. This is cooperative task-pool accounting, not an allocator-level peak-memory cap: DeepSizeOf is approximate and the returned output allocation lives beyond the retained-map measurement. A rejected batch preserves earlier candidates and reservations, and the reservation is released with the stream.
Validation
Fixes #8057