diff --git a/rust/lance-index/src/scalar/rtree.rs b/rust/lance-index/src/scalar/rtree.rs index 0d6b54a02ca..4ece20a4c7e 100644 --- a/rust/lance-index/src/scalar/rtree.rs +++ b/rust/lance-index/src/scalar/rtree.rs @@ -14,8 +14,6 @@ use crate::scalar::{ }; use crate::{Index, IndexType, pb}; use arrow_array::UInt32Array; -use arrow_array::cast::AsArray; -use arrow_array::types::UInt64Type; use arrow_array::{Array, BinaryArray, RecordBatch, UInt64Array}; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; use async_trait::async_trait; @@ -31,7 +29,6 @@ use geoarrow_schema::{Dimension, RectType}; use lance_arrow::RecordBatchExt; use lance_core::cache::{CacheKey, LanceCache, WeakLanceCache}; use lance_core::deepsize::DeepSizeOf; -use lance_core::utils::address::RowAddress; use lance_core::utils::row_addr_remap::RowAddrRemap; use lance_core::utils::tempfile::TempDir; use lance_core::{Error, ROW_ID, Result}; @@ -218,7 +215,7 @@ pub fn extract_bounding_boxes( } struct BboxStreamStats { - null_map: RowAddrTreeMap, + null_row_ids: RowAddrTreeMap, total_bbox: BoundingBox, // Number of non-null items num_items: usize, @@ -314,7 +311,7 @@ impl RTreeIndex { return Ok(RowAddrTreeMap::default()); } - let mut row_addrs = RowAddrTreeMap::new(); + let mut row_ids = RowAddrTreeMap::new(); let mut stack = vec![self.metadata.num_pages - 1]; while let Some(page_idx) = stack.pop() { @@ -332,7 +329,7 @@ impl RTreeIndex { let bbox_array = extract_bounding_boxes(batch.column(0).as_ref(), batch.schema().field(0))?; - let rowaddr_or_pageid_array = batch + let rowid_or_pageid_array = batch .column(1) .as_any() .downcast_ref::() @@ -342,17 +339,17 @@ impl RTreeIndex { let rect = bbox_array.value(i).unwrap(); if bbox.rect_intersects(&rect) { if is_leaf { - let row_addr = rowaddr_or_pageid_array.value(i); - row_addrs.insert(row_addr); + let row_id = rowid_or_pageid_array.value(i); + row_ids.insert(row_id); } else { - let page_id = rowaddr_or_pageid_array.value(i); + let page_id = rowid_or_pageid_array.value(i); stack.push(page_id); } } } } - Ok(row_addrs) + Ok(row_ids) } async fn search_null(&self, metrics: &dyn MetricsCollector) -> Result { @@ -481,22 +478,7 @@ impl Index for RTreeIndex { } async fn calculate_included_frags(&self) -> Result { - let mut frag_ids = RoaringBitmap::default(); - - let mut reader_stream = self.clone().into_data_stream().await?; - while let Some(page) = reader_stream.try_next().await? { - let mut page_frag_ids = page - .column(1) - .as_primitive::() - .iter() - .flatten() - .map(|row_addr| RowAddress::from(row_addr).fragment_id()) - .collect::>(); - page_frag_ids.sort(); - page_frag_ids.dedup(); - frag_ids |= RoaringBitmap::from_sorted_iter(page_frag_ids).unwrap(); - } - Ok(frag_ids) + unimplemented!() } } @@ -513,15 +495,15 @@ impl ScalarIndex for RTreeIndex { let geo_array = extract_bounding_boxes(query.value.to_array()?.as_ref(), &query.field)?; let bbox = total_bounds(&geo_array)?; - let mut rowids = self.search_bbox(bbox, metrics).await?; + let mut row_ids = self.search_bbox(bbox, metrics).await?; let mut null_map = self.search_null(metrics).await?; if let Some(fri) = &self.frag_reuse_index { - rowids = fri.remap_row_addrs_tree_map(&rowids); + row_ids = fri.remap_row_addrs_tree_map(&row_ids); null_map = fri.remap_row_addrs_tree_map(&null_map); } Ok(SearchResult::AtMost(NullableRowAddrSet::new( - rowids, null_map, + row_ids, null_map, ))) } GeoQuery::IsNull => { @@ -581,7 +563,7 @@ impl ScalarIndex for RTreeIndex { new_bbox.add_rect(&self.metadata.bbox); let merge_stats = BboxStreamStats { - null_map: RowAddrTreeMap::union_all(&[&null_map, &stats.null_map]), + null_row_ids: RowAddrTreeMap::union_all(&[&null_map, &stats.null_row_ids]), total_bbox: new_bbox, num_items: self.metadata.num_items + stats.num_items, }; @@ -663,10 +645,10 @@ impl RTreeIndexPlugin { )); } - let row_id_field = schema.field_with_name(ROW_ID)?; - if *row_id_field.data_type() != DataType::UInt64 { + let row_id_field = schema.field(1); + if row_id_field.name() != ROW_ID || *row_id_field.data_type() != DataType::UInt64 { return Err(Error::invalid_input_source( - "Second field in RTree index schema must be of type UInt64".into(), + "Second field in RTree index schema must be _rowid with type UInt64".into(), )); } Ok(()) @@ -705,7 +687,7 @@ impl RTreeIndexPlugin { page_size: u32, spill_store: Arc, ) -> Result<(SendableRecordBatchStream, BboxStreamStats)> { - let mut null_rowaddrs = RowAddrTreeMap::new(); + let mut null_row_ids = RowAddrTreeMap::new(); let mut total_bbox = BoundingBox::new(); let mut num_non_null_rows = 0; @@ -717,7 +699,7 @@ impl RTreeIndexPlugin { while let Some(batch) = data.try_next().await? { let bbox_array = extract_bounding_boxes(&batch.column(0), batch.schema().field(0))?; - let rowaddr_array = batch + let row_id_array = batch .column(1) .as_any() .downcast_ref::() @@ -731,8 +713,8 @@ impl RTreeIndexPlugin { for i in 0..num_rows { if bbox_array.is_null(i) { - let rowaddr = rowaddr_array.value(i); - null_rowaddrs.insert(rowaddr); + let row_id = row_id_array.value(i); + null_row_ids.insert(row_id); } else { non_null_indexes.push(i as u32); } @@ -762,7 +744,7 @@ impl RTreeIndexPlugin { Ok(( Box::pin(new_data), BboxStreamStats { - null_map: null_rowaddrs, + null_row_ids, total_bbox, num_items: num_non_null_rows, }, @@ -901,7 +883,7 @@ impl RTreeIndexPlugin { ) .await?; - let nulls_file = Self::write_nulls(store, stats.null_map).await?; + let nulls_file = Self::write_nulls(store, stats.null_row_ids).await?; Ok(vec![page_file, nulls_file]) } @@ -1150,13 +1132,13 @@ mod tests { let num_points = 10000; let null_probability = 0.001; // 0.1% - let mut expected_nulls = Vec::new(); + let mut expected_nulls = RowAddrTreeMap::new(); let mut point_builder = PointBuilder::new(point_type.clone()); for i in 0..num_points { if rng.random_bool(null_probability) { point_builder.push_null(); - expected_nulls.push(RowAddress::new_from_parts(0, i as u32)); + expected_nulls.insert(i as u64); } else { let x = rng.random_range(-1000.0..1000.0); let y = rng.random_range(-1000.0..1000.0); @@ -1166,21 +1148,21 @@ mod tests { let point_arr = point_builder.finish(); let (rtree_index, _store, _tmpdir) = train_index(&point_arr, None).await; - let row_addrs = rtree_index + let row_ids = rtree_index .search_null(&NoOpMetricsCollector) .await .unwrap(); - let mut actual_nulls = row_addrs.row_addrs().unwrap().collect::>(); - actual_nulls.sort(); - expected_nulls.sort(); - - assert_eq!(actual_nulls, expected_nulls); + assert_eq!(row_ids, expected_nulls); } #[tokio::test] async fn test_update_and_search() { - fn gen_data(num_items: u32, frag_id: u32, nulls_addrs: &mut RowAddrTreeMap) -> RectArray { + fn gen_data( + num_items: u32, + row_id_start: u64, + null_row_ids: &mut RowAddrTreeMap, + ) -> RectArray { let bbox_type = RectType::new(Dimension::XY, Default::default()); let mut rng = rand::rng(); @@ -1190,7 +1172,7 @@ mod tests { for i in 0..num_items { if rng.random_bool(null_probability) { rect_builder.push_null(); - nulls_addrs.insert(RowAddress::new_from_parts(frag_id, i).into()); + null_row_ids.insert(row_id_start + i as u64); } else { let x1 = rng.random_range(-1000.0..1000.0); let y1 = rng.random_range(-1000.0..1000.0); @@ -1206,10 +1188,9 @@ mod tests { rect_builder.finish() } - let mut nulls_addrs = RowAddrTreeMap::default(); + let mut null_row_ids = RowAddrTreeMap::default(); - let frag_id = 0; - let rect_arr = gen_data(10000, frag_id, &mut nulls_addrs); + let rect_arr = gen_data(10000, 0, &mut null_row_ids); let (rtree_index, _store, _tmpdir) = train_index(&rect_arr, Some(16)).await; @@ -1220,14 +1201,14 @@ mod tests { Arc::new(LanceCache::no_cache()), )); - let new_frag_id = 1; - let new_rect_arr = gen_data(10000, 1, &mut nulls_addrs); - let new_rowaddr_arr = (0..new_rect_arr.len()) - .map(|off| RowAddress::new_from_parts(new_frag_id, off as u32).into()) + let new_row_id_start = rect_arr.len() as u64; + let new_rect_arr = gen_data(10000, new_row_id_start, &mut null_row_ids); + let new_row_ids = (0..new_rect_arr.len()) + .map(|off| new_row_id_start + off as u64) .collect::>(); let stream = convert_bbox_rowid_batch_stream( &new_rect_arr, - Arc::new(UInt64Array::from(new_rowaddr_arr.clone())), + Arc::new(UInt64Array::from(new_row_ids.clone())), ); rtree_index .update(stream, new_store.as_ref(), None) @@ -1243,17 +1224,17 @@ mod tests { coord! { x: 10.5, y: 1.5 }, coord! { x: 99.5, y: 200.5 }, )); - let row_addrs = new_rtree_index + let row_ids = new_rtree_index .search_bbox(search_bbox, &NoOpMetricsCollector) .await .unwrap(); - let mut expected_row_addrs = RowAddrTreeMap::new(); + let mut expected_row_ids = RowAddrTreeMap::new(); for i in 0..rect_arr.len() { if !rect_arr.is_null(i) { let bbox = BoundingBox::new_with_rect(&rect_arr.value(i).unwrap()); if search_bbox.rect_intersects(&bbox) { - expected_row_addrs.insert(i as u64); + expected_row_ids.insert(i as u64); } } } @@ -1261,18 +1242,18 @@ mod tests { if !new_rect_arr.is_null(i) { let bbox = BoundingBox::new_with_rect(&new_rect_arr.value(i).unwrap()); if search_bbox.rect_intersects(&bbox) { - expected_row_addrs.insert(new_rowaddr_arr.get(i).copied().unwrap()); + expected_row_ids.insert(new_row_ids.get(i).copied().unwrap()); } } } - assert_eq!(row_addrs, expected_row_addrs); + assert_eq!(row_ids, expected_row_ids); let actual_nulls = new_rtree_index .search_null(&NoOpMetricsCollector) .await .unwrap(); - assert_eq!(actual_nulls, nulls_addrs); + assert_eq!(actual_nulls, null_row_ids); } #[tokio::test] diff --git a/rust/lance/src/dataset/tests/dataset_geo.rs b/rust/lance/src/dataset/tests/dataset_geo.rs index a43718dd7d4..2d1d26bfe97 100644 --- a/rust/lance/src/dataset/tests/dataset_geo.rs +++ b/rust/lance/src/dataset/tests/dataset_geo.rs @@ -6,12 +6,13 @@ use std::vec; use crate::Dataset; use crate::dataset::tests::dataset_transactions::execute_sql; +use crate::dataset::write::WriteParams; use crate::index::DatasetIndexExt; use arrow_array::RecordBatch; use arrow_array::RecordBatchIterator; use arrow_array::cast::AsArray; -use arrow_array::types::Float64Type; +use arrow_array::types::{Float64Type, UInt64Type}; use datafusion::common::{assert_contains, assert_not_contains}; use geo_types::{Rect, coord, line_string}; use geoarrow_array::{ @@ -19,9 +20,11 @@ use geoarrow_array::{ builder::{LineStringBuilder, PointBuilder, PolygonBuilder}, }; use geoarrow_schema::{Dimension, LineStringType, PointType, PolygonType}; +use lance_core::utils::address::RowAddress; use lance_core::utils::tempfile::TempStrDir; use lance_index::IndexType; use lance_index::scalar::ScalarIndexParams; +use rstest::rstest; #[tokio::test] async fn test_geo_types() { @@ -150,8 +153,9 @@ async fn test_geo_sql() { ); } +#[rstest] #[tokio::test] -async fn test_geo_rtree_index() { +async fn test_geo_rtree_index(#[values(true, false)] use_stable_row_ids: bool) { // 1. Creates arrow table linestring spatial data let line_string_type = LineStringType::new(Dimension::XY, Default::default()); @@ -177,15 +181,26 @@ async fn test_geo_rtree_index() { // 2. Write to lance let lance_path = TempStrDir::default(); let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); - let mut dataset = Dataset::write(reader, &lance_path, Some(Default::default())) + let params = WriteParams { + max_rows_per_file: 1000, + enable_stable_row_ids: use_stable_row_ids, + ..Default::default() + }; + let mut dataset = Dataset::write(reader, &lance_path, Some(params)) .await .unwrap(); + assert!(dataset.fragments().len() > 1); - async fn assert_intersects_sql(dataset: &mut Dataset, has_index: bool) { - // Executes a SQL query with St_Distance function - let sql = "SELECT linestring from dataset where St_Intersects(linestring, ST_GeomFromText('LINESTRING ( 2 0, 0 2 )'))"; + async fn assert_intersects_sql( + dataset: &mut Dataset, + has_index: bool, + use_stable_row_ids: bool, + ) -> Vec<(u64, u64)> { + let sql = "SELECT _rowid, _rowaddr from dataset where St_Intersects(linestring, ST_GeomFromText('LINESTRING ( 2002 2000, 2000 2002 )'))"; let batches = dataset .sql(sql) + .with_row_id(true) + .with_row_addr(true) .build() .await .unwrap() @@ -193,14 +208,33 @@ async fn test_geo_rtree_index() { .await .unwrap(); - let mut num_rows = 0; - for b in batches { - num_rows += b.num_rows(); + let mut rows = Vec::new(); + for batch in &batches { + let row_ids = batch + .column_by_name("_rowid") + .unwrap() + .as_primitive::(); + let row_addrs = batch + .column_by_name("_rowaddr") + .unwrap() + .as_primitive::(); + for idx in 0..batch.num_rows() { + rows.push((row_ids.value(idx), row_addrs.value(idx))); + } + } + rows.sort_unstable(); + assert_eq!(2, rows.len()); + for (row_id, row_addr) in &rows { + if use_stable_row_ids { + assert_ne!(row_id, row_addr); + } + assert!(RowAddress::from(*row_addr).fragment_id() > 0); } - assert_eq!(2, num_rows); let batches = dataset .sql(&format!("Explain {}", sql)) + .with_row_id(true) + .with_row_addr(true) .build() .await .unwrap() @@ -213,9 +247,10 @@ async fn test_geo_rtree_index() { } else { assert_not_contains!(&plan, "ScalarIndexQuery"); } + rows } - assert_intersects_sql(&mut dataset, false).await; + let unindexed_rows = assert_intersects_sql(&mut dataset, false, use_stable_row_ids).await; dataset .create_index( @@ -228,5 +263,6 @@ async fn test_geo_rtree_index() { .await .unwrap(); - assert_intersects_sql(&mut dataset, true).await; + let indexed_rows = assert_intersects_sql(&mut dataset, true, use_stable_row_ids).await; + assert_eq!(indexed_rows, unindexed_rows); }