feat: optimize indices

python/src/dataset.rs

@@ -23,7 +23,7 @@ use arrow_schema::{DataType, Schema as ArrowSchema};
 use async_trait::async_trait;
 use chrono::Duration;
 
-use futures::StreamExt;
+use futures::{StreamExt, TryFutureExt};
 use lance::dataset::builder::DatasetBuilder;
 use lance::dataset::UpdateBuilder;
 use lance::dataset::{
@@ -37,6 +37,7 @@ use lance::index::{
 };
 use lance_arrow::as_fixed_size_list_array;
 use lance_core::{datatypes::Schema, format::Fragment, io::object_store::ObjectStoreParams};
+use lance_index::optimize::OptimizeOptions;
 use lance_index::{
     vector::{ivf::IvfBuildParams, pq::PQBuildParams},
     DatasetIndexExt, IndexParams, IndexType,
@@ -639,10 +640,21 @@ impl Dataset {
         })
     }
 
-    fn optimize_indices(&mut self, _kwargs: Option<&PyDict>) -> PyResult<()> {
+    fn optimize_indices(&mut self, kwargs: Option<&PyDict>) -> PyResult<()> {
         let mut new_self = self.ds.as_ref().clone();
-        RT.block_on(None, new_self.optimize_indices())?
-            .map_err(|err| PyIOError::new_err(err.to_string()))?;
+        let mut options: OptimizeOptions = Default::default();
+        if let Some(kwargs) = kwargs {
+            if let Some(num_indices_to_merge) = kwargs.get_item("num_indices_to_merge")? {
+                options.num_indices_to_merge = num_indices_to_merge.extract()?;
+            }
+        }
+        RT.block_on(
+            None,
+            new_self
+                .optimize_indices(&options)
+                .map_err(|err| PyIOError::new_err(err.to_string())),
+        )?
+        .map_err(|err| PyIOError::new_err(err.to_string()))?;
         self.ds = Arc::new(new_self);
         Ok(())
     }

rust/lance-index/src/lib.rs

@@ -22,6 +22,7 @@ use async_trait::async_trait;
 use lance_core::Result;
 use roaring::RoaringBitmap;
 
+pub mod optimize;
 pub mod scalar;
 pub mod traits;
 pub mod vector;
@@ -58,6 +59,7 @@ pub trait Index: Send + Sync {
 }
 
 /// Index Type
+#[derive(Debug, PartialEq)]
 pub enum IndexType {
     // Preserve 0-100 for simple indices.
     Scalar = 0,

rust/lance-index/src/optimize.rs

@@ -0,0 +1,33 @@
+// Copyright 2024 Lance Developers.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/// Options for optimizing all indices.
+#[derive(Debug)]
+pub struct OptimizeOptions {
+    /// Number of delta indices to merge for one column. Default: 1.
+    ///
+    /// If `num_indices_to_merge` is 0, a new delta index will be created.
+    /// If `num_indices_to_merge` is 1, the delta updates will be merged into the latest index.
+    /// If `num_indices_to_merge` is more than 1, the delta updates and latest N indices
+    /// will be merged into one single index.
+    pub num_indices_to_merge: usize,
+}
+
+impl Default for OptimizeOptions {
+    fn default() -> Self {
+        Self {
+            num_indices_to_merge: 1,
+        }
+    }
+}

rust/lance-index/src/traits.rs

@@ -17,7 +17,7 @@ use std::sync::Arc;
 use async_trait::async_trait;
 use lance_core::{format::Index, Result};
 
-use crate::{IndexParams, IndexType};
+use crate::{optimize::OptimizeOptions, IndexParams, IndexType};
 
 // Extends Lance Dataset with secondary index.
 ///
@@ -85,7 +85,7 @@ pub trait DatasetIndexExt {
     async fn load_scalar_index_for_column(&self, col: &str) -> Result<Option<Index>>;
 
     /// Optimize indices.
-    async fn optimize_indices(&mut self) -> Result<()>;
+    async fn optimize_indices(&mut self, options: &OptimizeOptions) -> Result<()>;
 
     /// Find index with a given index_name and return its serialized statistics.
     ///

rust/lance/src/dataset/scanner.rs

@@ -2757,7 +2757,7 @@ mod test {
 
             // UPDATE
 
-            dataset.optimize_indices().await.unwrap();
+            dataset.optimize_indices(&Default::default()).await.unwrap();
             let updated_version = dataset.version().version;
 
             // APPEND -> DELETE

rust/lance/src/index.rs

@@ -26,6 +26,7 @@ use lance_core::io::{
     read_message, read_message_from_buf, read_metadata_offset, reader::read_manifest_indexes,
     Reader,
 };
+use lance_index::optimize::OptimizeOptions;
 use lance_index::pb::index::Implementation;
 use lance_index::scalar::expression::IndexInformationProvider;
 use lance_index::scalar::lance_format::LanceIndexStore;
@@ -296,29 +297,47 @@ impl DatasetIndexExt for Dataset {
     }
 
     #[instrument(skip_all)]
-    async fn optimize_indices(&mut self) -> Result<()> {
+    async fn optimize_indices(&mut self, options: &OptimizeOptions) -> Result<()> {
         let dataset = Arc::new(self.clone());
         // Append index
         let indices = self.load_indices().await?;
 
+        let mut name_to_indices = HashMap::new();
+        for idx in indices.iter() {
+            name_to_indices
+                .entry(idx.name.clone())
+                .or_insert_with(Vec::new)
+                .push(idx);
+        }
+
         let mut new_indices = vec![];
         let mut removed_indices = vec![];
-        for idx in indices.as_slice() {
-            if idx.dataset_version == self.manifest.version {
-                continue;
-            }
-            let Some((new_id, new_frag_ids)) = append_index(dataset.clone(), idx).await? else {
+        for deltas in name_to_indices.values() {
+            let Some((new_id, removed, mut new_frag_ids)) =
+                append_index(dataset.clone(), deltas.as_slice(), options).await?
+            else {
                 continue;
             };
+            for removed_idx in removed.iter() {
+                new_frag_ids |= removed_idx.fragment_bitmap.as_ref().unwrap();
+            }
 
+            let last_idx = deltas.last().expect("Delte indices should not be empty");
             let new_idx = IndexMetadata {
                 uuid: new_id,
-                name: idx.name.clone(),
-                fields: idx.fields.clone(),
+                name: last_idx.name.clone(), // Keep the same name
+                fields: last_idx.fields.clone(),
                 dataset_version: self.manifest.version,
-                fragment_bitmap: new_frag_ids,
+                fragment_bitmap: Some(new_frag_ids),
             };
-            removed_indices.push(idx.clone());
+            removed_indices.extend(removed.iter().map(|&idx| idx.clone()));
+            if deltas.len() > removed.len() {
+                new_indices.extend(
+                    deltas[0..(deltas.len() - removed.len())]
+                        .iter()
+                        .map(|&idx| idx.clone()),
+                );
+            }
             new_indices.push(new_idx);
         }
 
@@ -533,6 +552,8 @@ impl DatasetIndexInternalExt for Dataset {
 
 #[cfg(test)]
 mod tests {
+    use crate::dataset::builder::DatasetBuilder;
+
     use super::*;
 
     use arrow_array::{FixedSizeListArray, RecordBatch, RecordBatchIterator};
@@ -666,4 +687,95 @@ mod tests {
         assert_eq!(stats["num_unindexed_rows"], 512);
         assert_eq!(stats["num_indexed_rows"], 512);
     }
+
+    #[tokio::test]
+    async fn test_optimize_delta_indices() {
+        let test_dir = tempdir().unwrap();
+        let dimensions = 16;
+        let column_name = "vec";
+        let field = Field::new(
+            column_name,
+            DataType::FixedSizeList(
+                Arc::new(Field::new("item", DataType::Float32, true)),
+                dimensions,
+            ),
+            false,
+        );
+        let schema = Arc::new(Schema::new(vec![field]));
+
+        let float_arr = generate_random_array(512 * dimensions as usize);
+
+        let vectors =
+            arrow_array::FixedSizeListArray::try_new_from_values(float_arr, dimensions).unwrap();
+
+        let record_batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vectors)]).unwrap();
+
+        let reader = RecordBatchIterator::new(
+            vec![record_batch.clone()].into_iter().map(Ok),
+            schema.clone(),
+        );
+
+        let test_uri = test_dir.path().to_str().unwrap();
+        let mut dataset = Dataset::write(reader, test_uri, None).await.unwrap();
+        let params = VectorIndexParams::ivf_pq(10, 8, 2, false, MetricType::L2, 10);
+        dataset
+            .create_index(
+                &[column_name],
+                IndexType::Vector,
+                Some("vec_idx".into()),
+                &params,
+                true,
+            )
+            .await
+            .unwrap();
+
+        let stats: serde_json::Value =
+            serde_json::from_str(&dataset.index_statistics("vec_idx").await.unwrap()).unwrap();
+        assert_eq!(stats["num_unindexed_rows"], 0);
+        assert_eq!(stats["num_indexed_rows"], 512);
+        assert_eq!(stats["num_indexed_fragments"], 1);
+        assert_eq!(stats["num_indices"], 1);
+
+        let reader =
+            RecordBatchIterator::new(vec![record_batch].into_iter().map(Ok), schema.clone());
+        dataset.append(reader, None).await.unwrap();
+        let mut dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
+        let stats: serde_json::Value =
+            serde_json::from_str(&dataset.index_statistics("vec_idx").await.unwrap()).unwrap();
+        assert_eq!(stats["num_unindexed_rows"], 512);
+        assert_eq!(stats["num_indexed_rows"], 512);
+        assert_eq!(stats["num_indexed_fragments"], 1);
+        assert_eq!(stats["num_unindexed_fragments"], 1);
+        assert_eq!(stats["num_indices"], 1);
+
+        dataset
+            .optimize_indices(&OptimizeOptions {
+                num_indices_to_merge: 0, // Just create index for delta
+            })
+            .await
+            .unwrap();
+        let mut dataset = DatasetBuilder::from_uri(test_uri).load().await.unwrap();
+
+        let stats: serde_json::Value =
+            serde_json::from_str(&dataset.index_statistics("vec_idx").await.unwrap()).unwrap();
+        assert_eq!(stats["num_unindexed_rows"], 0);
+        assert_eq!(stats["num_indexed_rows"], 1024);
+        assert_eq!(stats["num_indexed_fragments"], 2);
+        assert_eq!(stats["num_unindexed_fragments"], 0);
+        assert_eq!(stats["num_indices"], 2);
+
+        dataset
+            .optimize_indices(&OptimizeOptions {
+                num_indices_to_merge: 2,
+            })
+            .await
+            .unwrap();
+        let stats: serde_json::Value =
+            serde_json::from_str(&dataset.index_statistics("vec_idx").await.unwrap()).unwrap();
+        assert_eq!(stats["num_unindexed_rows"], 0);
+        assert_eq!(stats["num_indexed_rows"], 1024);
+        assert_eq!(stats["num_indexed_fragments"], 2);
+        assert_eq!(stats["num_unindexed_fragments"], 0);
+        assert_eq!(stats["num_indices"], 1);
+    }
 }