fix(rowids): union overlapping segments in RowAddrTreeMap::from

rust/lance-table/src/rowids.rs

@@ -530,30 +530,32 @@ impl From<&RowIdSequence> for RowAddrTreeMap {
     fn from(row_ids: &RowIdSequence) -> Self {
         let mut tree_map = Self::new();
         for segment in &row_ids.0 {
+            let mut seg = Self::new();
             match segment {
                 U64Segment::Range(range) => {
-                    tree_map.insert_range(range.clone());
+                    seg.insert_range(range.clone());
                 }
                 U64Segment::RangeWithBitmap { range, bitmap } => {
-                    tree_map.insert_range(range.clone());
+                    seg.insert_range(range.clone());
                     for (i, val) in range.clone().enumerate() {
                         if !bitmap.get(i) {
-                            tree_map.remove(val);
+                            seg.remove(val);
                         }
                     }
                 }
                 U64Segment::RangeWithHoles { range, holes } => {
-                    tree_map.insert_range(range.clone());
+                    seg.insert_range(range.clone());
                     for hole in holes.iter() {
-                        tree_map.remove(hole);
+                        seg.remove(hole);
                     }
                 }
                 U64Segment::SortedArray(array) | U64Segment::Array(array) => {
                     for val in array.iter() {
-                        tree_map.insert(val);
+                        seg.insert(val);
                     }
                 }
             }
+            tree_map |= seg;
         }
         tree_map
     }
@@ -1040,6 +1042,27 @@ mod test {
         assert_eq!(tree_map, expected);
     }
 
+    #[test]
+    fn test_row_id_sequence_to_treemap_overlapping_segments() {
+        // Compaction can concatenate segments whose ranges overlap but whose
+        // selected ids are disjoint (here: even ids, then odd ids over 0..6).
+        // The tree map must contain every id the sequence yields.
+        let sequence = RowIdSequence(vec![
+            U64Segment::RangeWithBitmap {
+                range: 0..6,
+                bitmap: [true, false, true, false, true, false].as_slice().into(),
+            },
+            U64Segment::RangeWithBitmap {
+                range: 0..6,
+                bitmap: [false, true, false, true, false, true].as_slice().into(),
+            },
+        ]);
+
+        let expected = sequence.iter().collect::<RowAddrTreeMap>();
+        assert_eq!(expected, (0..6).collect::<RowAddrTreeMap>());
+        assert_eq!(RowAddrTreeMap::from(&sequence), expected);
+    }
+
     #[test]
     fn test_row_addr_mask() {
         // 0, 1, 2, 3, 4

rust/lance/src/dataset/tests/dataset_aggregate.rs

@@ -1361,6 +1361,85 @@ async fn test_scanner_count_rows_with_indexed_filter_stable_row_ids() {
     );
 }
 
+#[tokio::test]
+async fn test_scanner_count_rows_indexed_filter_stable_row_ids_after_compaction() {
+    // Update rewrites a scattered subset of rows under stable row ids; the
+    // rewritten copies keep their stable ids, so compaction folds the surviving
+    // and rewritten halves of a fragment into row-id segments whose ranges
+    // overlap. The indexed-filter count must still see every live row.
+    let tmp = tempdir().unwrap();
+    let uri = tmp.path().to_str().unwrap();
+    let ds = gen_batch()
+        .col("x", array::step::<Int64Type>())
+        .col("category", array::cycle::<Int64Type>(vec![1, 2, 3]))
+        .into_dataset_with_params(
+            uri,
+            FragmentCount::from(2),
+            FragmentRowCount::from(50),
+            Some(crate::dataset::WriteParams {
+                max_rows_per_file: 50,
+                enable_stable_row_ids: true,
+                ..Default::default()
+            }),
+        )
+        .await
+        .unwrap();
+    // Update every third row (category == 1), scattered across stable ids.
+    let res = crate::dataset::UpdateBuilder::new(Arc::new(ds))
+        .update_where("category = 1")
+        .unwrap()
+        .set("category", "0")
+        .unwrap()
+        .build()
+        .unwrap()
+        .execute()
+        .await
+        .unwrap();
+    let mut ds = res.new_dataset.as_ref().clone();
+    // Compaction merges the surviving and rewritten fragments, producing a
+    // fragment whose row-id sequence has overlapping segments.
+    crate::dataset::optimize::compact_files(&mut ds, Default::default(), None)
+        .await
+        .unwrap();
+    // Index after compaction: it covers every fragment and no deletions remain,
+    // so the count is answered from the stable-id universe (the path the
+    // overlapping segments corrupt).
+    ds.create_index(
+        &["x"],
+        IndexType::BTree,
+        None,
+        &ScalarIndexParams::default(),
+        true,
+    )
+    .await
+    .unwrap();
+
+    let mut scanner = ds.scan();
+    scanner.filter("x < 100").unwrap();
+    scanner
+        .aggregate(AggregateExpr::builder().count_star().build())
+        .unwrap();
+    let plan = scanner.create_plan().await.unwrap();
+
+    assert_plan_node_equals(
+        plan.clone(),
+        "AggregateExec: mode=Final, gby=[], aggr=[count(Int32(1))]
+  CountFromMask
+    ScalarIndexQuery: query=[x < 100]@x_idx(BTree)",
+    )
+    .await
+    .unwrap();
+
+    let stream = execute_plan(plan, LanceExecutionOptions::default()).unwrap();
+    let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
+    assert_eq!(batches.len(), 1);
+    // No deletions remain after compaction; all 100 rows match `x < 100`.
+    assert_eq!(
+        batches[0].column(0).as_primitive::<Int64Type>().value(0),
+        100,
+    );
+}
+
 #[tokio::test]
 async fn test_scanner_count_rows_with_partial_index_coverage() {
     // Index covers the first two fragments, then a third fragment is