Add configurable scalar index on id column (#50)

## Summary
- Add `IdIndexType` enum (None/BTree/ZoneMap) for configurable scalar
index on the `id` column
- BTree indices are maintained by MemWAL during writes; ZoneMap indices
are excluded from MemWAL maintenance (unsupported by MemWAL) and rebuilt
after compaction
- Index is created eagerly on store open and ensured after each
compaction cycle
- Exposed via Python API as `id_index_type` parameter (`"btree"`,
`"zonemap"`, or `None`)

## Test plan
- [x] Rust tests: btree index creation on open, persistence after
compaction
- [x] Rust tests: zonemap index creation on open, persistence after
compaction
- [x] Rust tests: no index created by default (IdIndexType::None)
- [x] Rust tests: idempotent ensure_id_index (no version bump on
re-check)
- [x] Python tests: btree/zonemap creation, invalid type rejection, data
integrity across compaction cycles
- [x] Full existing test suite passes (19/19)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Beinan Wang <beinanwang@microsoft.com>

Author: beinan

crates/lance-context-core/src/lib.rs

@@ -7,7 +7,9 @@ mod store;
 
 pub use context::{Context, ContextEntry, Snapshot};
 pub use record::{ContextRecord, SearchResult, StateMetadata};
-pub use store::{CompactionConfig, CompactionStats, ContextStore, ContextStoreOptions};
+pub use store::{
+    CompactionConfig, CompactionStats, ContextStore, ContextStoreOptions, IdIndexType,
+};
 
 // Re-export CompactionMetrics from lance for Python bindings
 pub use lance::dataset::optimize::CompactionMetrics;

crates/lance-context-core/src/store.rs

@@ -22,6 +22,8 @@ use lance::index::DatasetIndexExt;
 use lance::io::{ObjectStoreParams, StorageOptionsAccessor};
 use lance::{Error as LanceError, Result as LanceResult};
 use lance_index::mem_wal::MEM_WAL_INDEX_NAME;
+use lance_index::scalar::ScalarIndexParams;
+use lance_index::IndexType;
 use tokio::sync::Mutex;
 use tokio::task::JoinHandle;
 use tracing::{error, info, warn};
@@ -32,6 +34,7 @@ use crate::record::{ContextRecord, SearchResult, StateMetadata};
 /// Embedding length used for the semantic index column.
 const DEFAULT_EMBEDDING_DIM: i32 = 1536;
 const DEFAULT_SEARCH_LIMIT: usize = 10;
+const ID_INDEX_NAME: &str = "id_idx";
 
 /// Configuration for background compaction.
 #[derive(Debug, Clone)]
@@ -72,6 +75,18 @@ impl Default for CompactionConfig {
     }
 }
 
+/// Type of scalar index on the `id` column.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
+pub enum IdIndexType {
+    /// No index on the id column.
+    #[default]
+    None,
+    /// Zone-map index (min/max per fragment, lightweight).
+    ZoneMap,
+    /// B-tree index (point lookups, heavier).
+    BTree,
+}
+
 /// Statistics about compaction status and history.
 #[derive(Debug, Clone)]
 pub struct CompactionStats {
@@ -106,6 +121,7 @@ pub struct ContextStore {
     compaction_state: Arc<Mutex<CompactionState>>,
     pub compaction_config: CompactionConfig,
     blob_columns: HashSet<String>,
+    id_index_type: IdIndexType,
 }
 
 /// Additional configuration when opening a [`ContextStore`].
@@ -116,6 +132,8 @@ pub struct ContextStoreOptions {
     /// Column names that should use Lance V1 blob encoding.
     /// Valid values: `"text_payload"`, `"binary_payload"`.
     pub blob_columns: HashSet<String>,
+    /// Type of scalar index to create on the `id` column.
+    pub id_index_type: IdIndexType,
 }
 
 impl ContextStoreOptions {
@@ -164,8 +182,12 @@ impl ContextStore {
             })),
             compaction_config: options.compaction,
             blob_columns,
+            id_index_type: options.id_index_type,
         };
 
+        // Ensure id index if configured
+        store.ensure_id_index().await?;
+
         // Start background compaction if enabled
         store.start_background_compaction().await?;
 
@@ -192,8 +214,14 @@ impl ContextStore {
             let has_mem_wal = indices.iter().any(|i| i.name == MEM_WAL_INDEX_NAME);
 
             if !has_mem_wal {
-                let maintained_indexes: Vec<String> =
-                    indices.iter().map(|i| i.name.clone()).collect();
+                // ZoneMap indices are not supported by MemWAL; exclude them
+                let maintained_indexes: Vec<String> = indices
+                    .iter()
+                    .filter(|i| {
+                        !(self.id_index_type == IdIndexType::ZoneMap && i.name == ID_INDEX_NAME)
+                    })
+                    .map(|i| i.name.clone())
+                    .collect();
                 self.dataset
                     .initialize_mem_wal()
                     .unsharded()
@@ -348,6 +376,7 @@ impl ContextStore {
                 state.last_compaction = Some(Utc::now());
                 state.total_compactions += 1;
                 state.last_error = None;
+                drop(state); // Release lock before ensure_id_index
 
                 info!(
                     "Compaction completed in {:?}: removed {} fragments ({}files), added {} fragments ({} files)",
@@ -361,6 +390,12 @@ impl ContextStore {
                 // Reload dataset to see new version
                 self.dataset = Dataset::open(self.dataset.uri()).await?;
 
+                // Ensure id index exists after compaction
+                // (handles first-time creation on previously empty dataset)
+                if let Err(e) = self.ensure_id_index().await {
+                    warn!("Failed to ensure id index after compaction: {}", e);
+                }
+
                 Ok(metrics)
             }
             Err(e) => {
@@ -408,6 +443,44 @@ impl ContextStore {
         })
     }
 
+    /// Ensure the configured id index exists on the dataset.
+    async fn ensure_id_index(&mut self) -> LanceResult<()> {
+        if self.id_index_type == IdIndexType::None {
+            return Ok(());
+        }
+
+        let indices = self.dataset.load_indices().await?;
+        if indices.iter().any(|i| i.name == ID_INDEX_NAME) {
+            return Ok(());
+        }
+
+        self.create_id_index().await
+    }
+
+    /// Create (or replace) the scalar index on the `id` column.
+    pub async fn create_id_index(&mut self) -> LanceResult<()> {
+        let index_type = match self.id_index_type {
+            IdIndexType::ZoneMap => IndexType::ZoneMap,
+            IdIndexType::BTree => IndexType::BTree,
+            IdIndexType::None => return Ok(()),
+        };
+
+        info!("Creating {:?} index on id column", index_type);
+
+        let params = ScalarIndexParams::default();
+
+        self.dataset
+            .create_index_builder(&["id"], index_type, &params)
+            .name(ID_INDEX_NAME.to_string())
+            .replace(true)
+            .await?;
+
+        // Reload dataset to pick up new index
+        self.dataset = Dataset::open(self.dataset.uri()).await?;
+
+        Ok(())
+    }
+
     /// Start background compaction task if enabled.
     async fn start_background_compaction(&mut self) -> LanceResult<()> {
         if !self.compaction_config.enabled {
@@ -1379,4 +1452,133 @@ mod tests {
             assert_eq!(results_binary[0].text_payload, record.text_payload);
         });
     }
+
+    #[test]
+    fn test_id_index_btree() {
+        let dir = TempDir::new().unwrap();
+        let uri = dir.path().to_string_lossy().to_string();
+        let runtime = tokio::runtime::Runtime::new().unwrap();
+
+        runtime.block_on(async {
+            let options = ContextStoreOptions {
+                id_index_type: IdIndexType::BTree,
+                ..Default::default()
+            };
+            let mut store = ContextStore::open_with_options(&uri, options)
+                .await
+                .unwrap();
+
+            // Index should be created eagerly on open
+            let indices = store.dataset.load_indices().await.unwrap();
+            assert!(
+                indices.iter().any(|i| i.name == ID_INDEX_NAME),
+                "btree index should be created on open"
+            );
+
+            // Add data and verify it still works with the index
+            for i in 0..5 {
+                store
+                    .add(&[text_record(&format!("btree-{i}"), i as f32)])
+                    .await
+                    .unwrap();
+            }
+            store.compact(None).await.unwrap();
+
+            // Index should still exist after compaction
+            let indices = store.dataset.load_indices().await.unwrap();
+            assert!(
+                indices.iter().any(|i| i.name == ID_INDEX_NAME),
+                "btree index should persist after compaction"
+            );
+        });
+    }
+
+    #[test]
+    fn test_id_index_zonemap() {
+        let dir = TempDir::new().unwrap();
+        let uri = dir.path().to_string_lossy().to_string();
+        let runtime = tokio::runtime::Runtime::new().unwrap();
+
+        runtime.block_on(async {
+            let options = ContextStoreOptions {
+                id_index_type: IdIndexType::ZoneMap,
+                ..Default::default()
+            };
+            let mut store = ContextStore::open_with_options(&uri, options)
+                .await
+                .unwrap();
+
+            // Index should be created eagerly on open
+            let indices = store.dataset.load_indices().await.unwrap();
+            assert!(
+                indices.iter().any(|i| i.name == ID_INDEX_NAME),
+                "zonemap index should be created on open"
+            );
+
+            for i in 0..5 {
+                store
+                    .add(&[text_record(&format!("zm-{i}"), i as f32)])
+                    .await
+                    .unwrap();
+            }
+            store.compact(None).await.unwrap();
+
+            let indices = store.dataset.load_indices().await.unwrap();
+            assert!(
+                indices.iter().any(|i| i.name == ID_INDEX_NAME),
+                "zonemap index should persist after compaction"
+            );
+        });
+    }
+
+    #[test]
+    fn test_id_index_none_by_default() {
+        let dir = TempDir::new().unwrap();
+        let uri = dir.path().to_string_lossy().to_string();
+        let runtime = tokio::runtime::Runtime::new().unwrap();
+
+        runtime.block_on(async {
+            let mut store = ContextStore::open(&uri).await.unwrap();
+
+            store.add(&[text_record("no-idx-1", 0.0)]).await.unwrap();
+            store.compact(None).await.unwrap();
+
+            let indices = store.dataset.load_indices().await.unwrap();
+            assert!(
+                !indices.iter().any(|i| i.name == ID_INDEX_NAME),
+                "no id index should be created when IdIndexType::None"
+            );
+        });
+    }
+
+    #[test]
+    fn test_id_index_idempotent() {
+        let dir = TempDir::new().unwrap();
+        let uri = dir.path().to_string_lossy().to_string();
+        let runtime = tokio::runtime::Runtime::new().unwrap();
+
+        runtime.block_on(async {
+            let options = ContextStoreOptions {
+                id_index_type: IdIndexType::BTree,
+                ..Default::default()
+            };
+            let mut store = ContextStore::open_with_options(&uri, options)
+                .await
+                .unwrap();
+
+            for i in 0..5 {
+                store
+                    .add(&[text_record(&format!("idem-{i}"), i as f32)])
+                    .await
+                    .unwrap();
+            }
+
+            // Create index twice -- second call should be a no-op
+            store.create_id_index().await.unwrap();
+            let v1 = store.version();
+            store.ensure_id_index().await.unwrap();
+            let v2 = store.version();
+            assert_eq!(v1, v2, "ensure_id_index should not recreate existing index");
+        });
+    }
 }

python/python/lance_context/api.py

@@ -263,6 +263,7 @@ def __init__(
         compaction_min_fragments: int = 5,
         compaction_target_rows: int = 1_000_000,
         quiet_hours: list[tuple[int, int]] | None = None,
+        id_index_type: str | None = None,
     ) -> None:
         options = _merge_storage_options(
             storage_options,
@@ -282,11 +283,12 @@ def __init__(
             "quiet_hours": quiet_hours or [],
         }
 
-        if options or compaction_config["enabled"]:
+        if options or compaction_config["enabled"] or id_index_type:
             self._inner = _Context.create(
                 uri,
                 storage_options=options or None,
                 compaction_config=compaction_config,
+                id_index_type=id_index_type,
             )
         else:
             self._inner = _Context.create(uri)
@@ -308,6 +310,7 @@ def create(
         compaction_min_fragments: int = 5,
         compaction_target_rows: int = 1_000_000,
         quiet_hours: list[tuple[int, int]] | None = None,
+        id_index_type: str | None = None,
     ) -> Context:
         return cls(
             uri,
@@ -323,6 +326,7 @@ def create(
             compaction_min_fragments=compaction_min_fragments,
             compaction_target_rows=compaction_target_rows,
             quiet_hours=quiet_hours,
+            id_index_type=id_index_type,
         )
 
     def uri(self) -> str:

python/src/lib.rs

@@ -11,7 +11,7 @@ use tokio::runtime::Runtime;
 use lance_context::serde::CONTENT_TYPE_TEXT;
 use lance_context::{
     CompactionConfig, CompactionMetrics, CompactionStats, Context as RustContext, ContextRecord,
-    ContextStore, ContextStoreOptions, SearchResult,
+    ContextStore, ContextStoreOptions, IdIndexType, SearchResult,
 };
 
 const DEFAULT_BINARY_CONTENT_TYPE: &str = "application/octet-stream";
@@ -110,23 +110,37 @@ fn compaction_config_from_dict<'py>(
 #[pymethods]
 impl Context {
     #[classmethod]
-    #[pyo3(signature = (uri, *, storage_options=None, compaction_config=None, blob_columns=None))]
+    #[pyo3(signature = (uri, *, storage_options=None, compaction_config=None, blob_columns=None, id_index_type=None))]
     fn create(
         _cls: &Bound<'_, PyType>,
         py: Python<'_>,
         uri: &str,
         storage_options: Option<&Bound<'_, PyDict>>,
         compaction_config: Option<&Bound<'_, PyDict>>,
         blob_columns: Option<Vec<String>>,
+        id_index_type: Option<String>,
     ) -> PyResult<Self> {
         let runtime = Arc::new(Runtime::new().map_err(to_py_err)?);
 
         let blob_set: HashSet<String> = blob_columns.unwrap_or_default().into_iter().collect();
 
+        let id_idx = match id_index_type.as_deref() {
+            Some("btree") => IdIndexType::BTree,
+            Some("zonemap") => IdIndexType::ZoneMap,
+            Some("none") | None => IdIndexType::None,
+            Some(other) => {
+                return Err(PyRuntimeError::new_err(format!(
+                    "invalid id_index_type '{}': valid values are 'btree', 'zonemap'",
+                    other
+                )))
+            }
+        };
+
         let options = ContextStoreOptions {
             storage_options: storage_options_from_dict(storage_options)?,
             compaction: compaction_config_from_dict(compaction_config)?,
             blob_columns: blob_set,
+            id_index_type: id_idx,
         };
 
         let store_res =