feat: support segmented RTree indexes

java/lance-jni/src/blocking_dataset.rs

@@ -1228,8 +1228,10 @@ fn index_metadata_to_segment(metadata: &IndexMetadata) -> Result<IndexSegment> {
     Ok(IndexSegment::new(
         metadata.uuid,
         fragment_bitmap,
+        metadata.fields.iter().copied(),
         index_details,
         metadata.index_version,
+        metadata.dataset_version,
     ))
 }
 

python/python/lance/dataset.py

@@ -3308,6 +3308,7 @@ def _is_segment_native_scalar_index_type(
             "BITMAP",
             "INVERTED",
             "FTS",
+            "RTREE",
             "ZONEMAP",
             "BLOOMFILTER",
         }
@@ -3320,6 +3321,7 @@ def _requires_uncommitted_scalar_index(
         return cls._normalized_index_type(index_type) in {
             "BTREE",
             "BITMAP",
+            "RTREE",
             "ZONEMAP",
             "BLOOMFILTER",
         }
@@ -4277,8 +4279,8 @@ def create_index_uncommitted(
         Create one segment without publishing it and return its metadata.
 
         This is the public distributed-build API for vector, BTREE scalar,
-        canonical bitmap scalar, INVERTED scalar, ZONEMAP scalar, and
-        BLOOMFILTER scalar index construction. Unlike
+        canonical bitmap scalar, INVERTED scalar, RTREE scalar, ZONEMAP scalar,
+        and BLOOMFILTER scalar index construction. Unlike
         :meth:`create_index`, this method does not publish the index into the
         dataset manifest. Instead, it writes one segment under
         ``_indices/<segment_uuid>/`` and returns the resulting
@@ -4294,7 +4296,7 @@ def create_index_uncommitted(
         4. commit the final segment list with
            :meth:`commit_existing_index_segments`
 
-        BTREE, BITMAP, INVERTED, ZONEMAP, and BLOOMFILTER segments may
+        BTREE, BITMAP, INVERTED, RTREE, ZONEMAP, and BLOOMFILTER segments may
         be merged with :meth:`merge_existing_index_segments` before commit.
         Parameters are the same as :meth:`create_index`, with one additional
         requirement:

python/python/tests/test_geo.py

@@ -19,6 +19,19 @@
 )
 
 
+def _query_point_ids(dataset: lance.LanceDataset, wkt: str) -> list[int]:
+    sql = f"""
+          SELECT id, point
+          FROM dataset
+          WHERE St_Intersects(point, ST_GeomFromText('{wkt}'))
+          """
+    return [
+        value
+        for batch in dataset.sql(sql).build().to_batch_records()
+        for value in batch.column("id").to_pylist()
+    ]
+
+
 def test_geo_types(tmp_path: Path):
     uri = str(tmp_path / "test_geo_types.lance")
     # Points
@@ -153,3 +166,275 @@ def query(ds: lance.LanceDataset, has_index=False):
     table_with_index = query(ds, has_index=True)
 
     assert table_with_index == table_without_index
+
+
+def test_rtree_segment_merge_and_commit(tmp_path: Path):
+    num_points = 120
+    points_2d = points(
+        [
+            np.arange(num_points, dtype=np.float64),
+            np.arange(num_points, dtype=np.float64),
+        ]
+    )
+    schema = pa.schema(
+        [
+            pa.field("id", pa.int64()),
+            pa.field(point("xy")).with_name("point"),
+        ]
+    )
+    table = pa.Table.from_arrays(
+        [np.arange(num_points, dtype=np.int64), points_2d], schema=schema
+    )
+    ds = lance.write_dataset(
+        table,
+        str(tmp_path / "segmented_rtree.lance"),
+        max_rows_per_file=40,
+    )
+    fragments = ds.get_fragments()
+    assert len(fragments) == 3
+    segments = [
+        ds.create_index_uncommitted(
+            column="point",
+            index_type="RTREE",
+            name="point_rtree",
+            fragment_ids=[fragment.fragment_id],
+        )
+        for fragment in fragments
+    ]
+
+    merged = ds.merge_existing_index_segments(segments)
+    assert set(merged.fragment_ids) == {fragment.fragment_id for fragment in fragments}
+    ds = ds.commit_existing_index_segments("point_rtree", "point", [merged])
+
+    sql = """
+          SELECT id, point
+          FROM dataset
+          WHERE St_Intersects(point, ST_GeomFromText('LINESTRING (10 10, 110 110)'))
+          """
+    indexed = pa.Table.from_batches(ds.sql(sql).build().to_batch_records())
+    assert indexed["id"].to_pylist() == list(range(10, 111))
+    explain = (
+        pa.Table.from_batches(
+            ds.sql("EXPLAIN ANALYZE " + sql).build().to_batch_records()
+        )
+        .to_pandas()
+        .to_string()
+    )
+    assert "ScalarIndexQuery" in explain
+
+
+def test_staged_rtree_after_rewrite_columns(tmp_path: Path):
+    uri = str(tmp_path / "stale_rtree.lance")
+    point_type = point("xy")
+    schema = pa.schema(
+        [
+            pa.field("id", pa.int64()),
+            pa.field(point_type).with_name("point"),
+        ]
+    )
+    dataset = lance.write_dataset(
+        pa.Table.from_arrays(
+            [
+                pa.array([0, 1], type=pa.int64()),
+                points([np.array([0.0, 1.0]), np.array([0.0, 1.0])]),
+            ],
+            schema=schema,
+        ),
+        uri,
+    )
+    segment = dataset.create_index_uncommitted(
+        column="point",
+        index_type="RTREE",
+        name="point_rtree",
+        fragment_ids=[0],
+    )
+
+    update_schema = pa.schema(
+        [
+            pa.field("_rowid", pa.uint64()),
+            pa.field(point_type).with_name("point"),
+        ]
+    )
+    update = pa.Table.from_arrays(
+        [
+            pa.array([0], type=pa.uint64()),
+            points([np.array([10.0]), np.array([10.0])]),
+        ],
+        schema=update_schema,
+    )
+    fragment, fields = dataset.get_fragment(0).update_columns(update)
+    updated = lance.LanceDataset.commit(
+        uri,
+        lance.LanceOperation.Update(
+            updated_fragments=[fragment],
+            fields_modified=fields,
+        ),
+        read_version=dataset.version,
+    )
+
+    assert _query_point_ids(updated, "POINT (10 10)") == [0]
+    committed = updated.commit_existing_index_segments(
+        "point_rtree",
+        "point",
+        [segment],
+    )
+    assert _query_point_ids(committed, "POINT (10 10)") == [0]
+
+
+def test_rtree_rejects_distributed_uuid_reuse(tmp_path: Path):
+    uri = str(tmp_path / "uuid_reuse.lance")
+    num_points = 120
+    point_type = point("xy")
+    schema = pa.schema(
+        [
+            pa.field("id", pa.int64()),
+            pa.field(point_type).with_name("point"),
+        ]
+    )
+    dataset = lance.write_dataset(
+        pa.Table.from_arrays(
+            [
+                pa.array(range(num_points), type=pa.int64()),
+                points(
+                    [
+                        np.arange(num_points, dtype=np.float64),
+                        np.arange(num_points, dtype=np.float64),
+                    ]
+                ),
+            ],
+            schema=schema,
+        ),
+        uri,
+        max_rows_per_file=40,
+    )
+    dataset.create_scalar_index("point", "RTREE")
+    index_uuid = dataset.describe_indices()[0].segments[0].uuid
+
+    with pytest.raises(
+        ValueError,
+        match="index_uuid is no longer accepted for RTree distributed index builds",
+    ):
+        dataset.create_index_uncommitted(
+            column="point",
+            index_type="RTREE",
+            name="point_rtree_reuse",
+            fragment_ids=[0],
+            index_uuid=index_uuid,
+        )
+
+    assert _query_point_ids(
+        lance.dataset(uri),
+        "LINESTRING (100 100, 110 110)",
+    ) == list(range(100, 111))
+
+
+def test_rtree_merge_all_deleted_stable_row_ids(tmp_path: Path):
+    uri = str(tmp_path / "all_deleted.lance")
+    point_type = point("xy")
+    schema = pa.schema(
+        [
+            pa.field("id", pa.int64()),
+            pa.field(point_type).with_name("point"),
+        ]
+    )
+    dataset = lance.write_dataset(
+        pa.Table.from_arrays(
+            [
+                pa.array([0, 1], type=pa.int64()),
+                points([np.array([0.0, 1.0]), np.array([0.0, 1.0])]),
+            ],
+            schema=schema,
+        ),
+        uri,
+        enable_stable_row_ids=True,
+    )
+    segment = dataset.create_index_uncommitted(
+        column="point",
+        index_type="RTREE",
+        name="point_rtree",
+        fragment_ids=[0],
+    )
+
+    dataset.delete("true")
+    merged = dataset.merge_existing_index_segments([segment])
+    assert merged.fragment_ids == set()
+    committed = dataset.commit_existing_index_segments(
+        "point_rtree",
+        "point",
+        [merged],
+    )
+    assert _query_point_ids(committed, "POINT (0 0)") == []
+
+
+def test_rtree_merge_preserves_newer_fragment_coverage(tmp_path: Path):
+    uri = str(tmp_path / "mixed_versions.lance")
+    point_type = point("xy")
+    schema = pa.schema(
+        [
+            pa.field("id", pa.int64()),
+            pa.field(point_type).with_name("point"),
+        ]
+    )
+
+    def batch(start: int, stop: int) -> pa.Table:
+        values = np.arange(start, stop, dtype=np.float64)
+        return pa.Table.from_arrays(
+            [
+                pa.array(range(start, stop), type=pa.int64()),
+                points([values, values]),
+            ],
+            schema=schema,
+        )
+
+    dataset = lance.write_dataset(batch(0, 40), uri)
+    first = dataset.create_index_uncommitted(
+        column="point",
+        index_type="RTREE",
+        name="point_rtree",
+        fragment_ids=[0],
+    )
+    dataset = lance.write_dataset(batch(40, 80), uri, mode="append")
+    second = dataset.create_index_uncommitted(
+        column="point",
+        index_type="RTREE",
+        name="point_rtree",
+        fragment_ids=[1],
+    )
+
+    merged = dataset.merge_existing_index_segments([first, second])
+    assert merged.fragment_ids == {0, 1}
+    assert merged.dataset_version == dataset.version
+
+    update_schema = pa.schema(
+        [
+            pa.field("_rowid", pa.uint64()),
+            pa.field(point_type).with_name("point"),
+        ]
+    )
+    update = pa.Table.from_arrays(
+        [
+            pa.array([1 << 32], type=pa.uint64()),
+            points([np.array([100.0]), np.array([100.0])]),
+        ],
+        schema=update_schema,
+    )
+    fragment, fields = dataset.get_fragment(1).update_columns(update)
+    dataset = lance.LanceDataset.commit(
+        uri,
+        lance.LanceOperation.Update(
+            updated_fragments=[fragment],
+            fields_modified=fields,
+        ),
+        read_version=dataset.version,
+    )
+
+    committed = dataset.commit_existing_index_segments(
+        "point_rtree",
+        "point",
+        [merged],
+    )
+    assert committed.describe_indices()[0].segments[0].fragment_ids == {0}
+    assert _query_point_ids(
+        committed,
+        "POINT (100 100)",
+    ) == [40]

python/src/dataset.rs

@@ -608,8 +608,10 @@ fn index_metadata_to_segment(metadata: IndexMetadata) -> PyResult<IndexSegment>
     Ok(IndexSegment::new(
         metadata.uuid,
         fragment_bitmap.iter(),
+        metadata.fields,
         index_details,
         metadata.index_version,
+        metadata.dataset_version,
     ))
 }
 

python/src/indices.rs

@@ -550,12 +550,14 @@ async fn do_load_shuffled_vectors(
             .as_ref()
             .expect("vector metadata should include fragment coverage")
             .iter(),
+        metadata.fields.iter().copied(),
         metadata
             .index_details
             .as_ref()
             .expect("vector metadata should include index details")
             .clone(),
         metadata.index_version,
+        metadata.dataset_version,
     );
     ds.commit_existing_index_segments(index_name, column, vec![segment])
         .await