lance_index/scalar/
lance_format.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Utilities for serializing and deserializing scalar indices in the lance format
5
6use super::{IndexFile, IndexReader, IndexStore, IndexWriter};
7use arrow_array::RecordBatch;
8use arrow_schema::Schema;
9use async_trait::async_trait;
10use bytes::Bytes;
11use futures::TryStreamExt;
12use lance_core::deepsize::DeepSizeOf;
13use lance_core::{Error, Result, cache::LanceCache};
14use lance_encoding::decoder::{DecoderPlugins, FilterExpression};
15use lance_encoding::version::LanceFileVersion;
16use lance_file::reader::{FileReader as CurrentFileReader, FileReaderOptions, ReaderProjection};
17use lance_file::versions::v1::reader::FileReader as V1FileReader;
18use lance_file::writer as current_writer;
19use lance_io::scheduler::{ScanScheduler, SchedulerConfig};
20use lance_io::utils::CachedFileSize;
21use lance_io::{ReadBatchParams, object_store::ObjectStore};
22use lance_table::format::SelfDescribingFileReader;
23use lance_table::format::list_index_files_with_sizes;
24use object_store::path::Path;
25use std::cmp::min;
26use std::collections::HashMap;
27use std::pin::Pin;
28use std::{any::Any, sync::Arc};
29
30/// An index store that serializes scalar indices using the lance format
31///
32/// Scalar indices are made up of named collections of record batches.  This
33/// struct relies on there being a dedicated directory for the index and stores
34/// each collection in a file in the lance format.
35#[derive(Debug, Clone)]
36pub struct LanceIndexStore {
37    object_store: Arc<ObjectStore>,
38    index_dir: Path,
39    metadata_cache: Arc<LanceCache>,
40    scheduler: Arc<ScanScheduler>,
41    /// Cached file sizes (filename -> size in bytes)
42    /// When set, used to avoid HEAD calls when opening files
43    file_sizes: HashMap<String, u64>,
44    format_version: LanceFileVersion,
45    /// Base I/O priority for all requests this store submits to `scheduler`.
46    io_priority: u64,
47}
48
49impl DeepSizeOf for LanceIndexStore {
50    fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize {
51        // Exclude the shared, session-scoped `metadata_cache` (accounted once by
52        // `Session::deep_size_of_children`): it is not this store's own footprint, and
53        // sizing it here makes every opened index that holds a store report the whole
54        // cache as its own bytes — inflating and N-times double-counting it.
55        self.object_store.deep_size_of_children(context)
56            + self.index_dir.as_ref().deep_size_of_children(context)
57    }
58}
59
60impl LanceIndexStore {
61    /// Create a new index store at the given directory
62    pub fn new(
63        object_store: Arc<ObjectStore>,
64        index_dir: Path,
65        metadata_cache: Arc<LanceCache>,
66    ) -> Self {
67        Self::with_format_version(
68            object_store,
69            index_dir,
70            metadata_cache,
71            LanceFileVersion::V2_0,
72        )
73    }
74
75    /// Create a new index store at the given directory with a specific format version
76    pub fn with_format_version(
77        object_store: Arc<ObjectStore>,
78        index_dir: Path,
79        metadata_cache: Arc<LanceCache>,
80        format_version: LanceFileVersion,
81    ) -> Self {
82        let scheduler = ScanScheduler::new(
83            object_store.clone(),
84            SchedulerConfig::max_bandwidth(&object_store),
85        );
86        Self {
87            object_store,
88            index_dir,
89            metadata_cache,
90            scheduler,
91            file_sizes: HashMap::new(),
92            format_version,
93            io_priority: 0,
94        }
95    }
96
97    /// Set cached file sizes to avoid HEAD calls when opening files.
98    ///
99    /// The map should contain relative paths (e.g., "index.idx") as keys
100    /// and file sizes in bytes as values.
101    pub fn with_file_sizes(mut self, file_sizes: HashMap<String, u64>) -> Self {
102        self.file_sizes = file_sizes;
103        self
104    }
105
106    /// The base I/O priority all this store's requests are submitted at.
107    pub fn io_priority(&self) -> u64 {
108        self.io_priority
109    }
110
111    fn index_file_path(&self, name: &str) -> Result<Path> {
112        let relative_path = Path::parse(name).map_err(|err| {
113            Error::invalid_input(format!("invalid index file path {name:?}: {err}"))
114        })?;
115        if self.index_dir.is_root() {
116            return Ok(relative_path);
117        }
118        if relative_path.is_root() {
119            return Ok(self.index_dir.clone());
120        }
121        Path::parse(format!(
122            "{}/{}",
123            self.index_dir.as_ref(),
124            relative_path.as_ref()
125        ))
126        .map_err(|err| Error::invalid_input(format!("invalid index file path {name:?}: {err}")))
127    }
128}
129
130struct LanceIndexWriter {
131    path: String,
132    inner: current_writer::FileWriter,
133}
134
135#[async_trait]
136impl IndexWriter for LanceIndexWriter {
137    async fn write_record_batch(&mut self, batch: RecordBatch) -> Result<u64> {
138        let offset = self.inner.tell().await?;
139        self.inner.write_batch(&batch).await?;
140        Ok(offset)
141    }
142
143    async fn add_global_buffer(&mut self, data: Bytes) -> Result<u32> {
144        self.inner.add_global_buffer(data).await
145    }
146
147    async fn finish(&mut self) -> Result<IndexFile> {
148        let summary = self.inner.finish().await?;
149        Ok(IndexFile {
150            path: self.path.clone(),
151            size_bytes: summary.size_bytes,
152        })
153    }
154
155    async fn finish_with_metadata(
156        &mut self,
157        metadata: HashMap<String, String>,
158    ) -> Result<IndexFile> {
159        metadata.into_iter().for_each(|(k, v)| {
160            self.inner.add_schema_metadata(k, v);
161        });
162        let summary = self.inner.finish().await?;
163        Ok(IndexFile {
164            path: self.path.clone(),
165            size_bytes: summary.size_bytes,
166        })
167    }
168}
169
170/// Newtype wrapper to allow implementing IndexReader for V1FileReader (a foreign type)
171struct V1IndexReader(V1FileReader);
172
173#[async_trait]
174impl IndexReader for V1IndexReader {
175    async fn read_record_batch(&self, offset: u64, _batch_size: u64) -> Result<RecordBatch> {
176        self.0
177            .read_batch(offset as i32, ReadBatchParams::RangeFull, self.0.schema())
178            .await
179    }
180
181    async fn read_range(
182        &self,
183        range: std::ops::Range<usize>,
184        projection: Option<&[&str]>,
185    ) -> Result<RecordBatch> {
186        let projection = match projection {
187            Some(projection) => self.0.schema().project(projection)?,
188            None => self.0.schema().clone(),
189        };
190        self.0.read_range(range, &projection).await
191    }
192
193    async fn num_batches(&self, _batch_size: u64) -> u32 {
194        self.0.num_batches() as u32
195    }
196
197    fn num_rows(&self) -> usize {
198        self.0.len()
199    }
200
201    fn schema(&self) -> &lance_core::datatypes::Schema {
202        V1FileReader::schema(&self.0)
203    }
204}
205
206/// Newtype wrapper to allow implementing IndexReader for CurrentFileReader (a foreign type)
207struct CurrentIndexReader(CurrentFileReader);
208
209#[async_trait]
210impl IndexReader for CurrentIndexReader {
211    async fn read_record_batch(&self, offset: u64, batch_size: u64) -> Result<RecordBatch> {
212        let start = offset * batch_size;
213        let end = start + batch_size;
214        let end = end.min(self.0.num_rows());
215        self.read_range(start as usize..end as usize, None).await
216    }
217
218    async fn read_global_buffer(&self, n: u32) -> Result<Bytes> {
219        CurrentFileReader::read_global_buffer(&self.0, n).await
220    }
221
222    async fn read_range(
223        &self,
224        range: std::ops::Range<usize>,
225        projection: Option<&[&str]>,
226    ) -> Result<RecordBatch> {
227        if range.is_empty() {
228            return Ok(RecordBatch::new_empty(Arc::new(
229                self.0.schema().as_ref().into(),
230            )));
231        }
232        let projection = if let Some(projection) = projection {
233            ReaderProjection::from_column_names(
234                self.0.metadata().version(),
235                self.0.schema(),
236                projection,
237            )?
238        } else {
239            ReaderProjection::from_whole_schema(self.0.schema(), self.0.metadata().version())
240        };
241        let batches = self
242            .0
243            .read_stream_projected(
244                ReadBatchParams::Range(range),
245                u32::MAX,
246                u32::MAX,
247                projection,
248                FilterExpression::no_filter(),
249            )
250            .await?
251            .try_collect::<Vec<_>>()
252            .await?;
253        assert_eq!(batches.len(), 1);
254        Ok(batches[0].clone())
255    }
256
257    async fn read_ranges(
258        &self,
259        ranges: &[std::ops::Range<usize>],
260        projection: Option<&[&str]>,
261    ) -> Result<RecordBatch> {
262        let empty_batch = || {
263            Ok(RecordBatch::new_empty(Arc::new(
264                self.0.schema().as_ref().into(),
265            )))
266        };
267        if ranges.is_empty() {
268            return empty_batch();
269        }
270        let projection = if let Some(projection) = projection {
271            ReaderProjection::from_column_names(
272                self.0.metadata().version(),
273                self.0.schema(),
274                projection,
275            )?
276        } else {
277            ReaderProjection::from_whole_schema(self.0.schema(), self.0.metadata().version())
278        };
279        // `DecodeBatchScheduler::schedule_ranges` requires sorted,
280        // non-overlapping ranges; sort internally and permute the
281        // result back to caller order so callers don't have to know.
282        let mut order: Vec<usize> = (0..ranges.len()).collect();
283        order.sort_by_key(|&i| ranges[i].start);
284        let already_sorted = order.iter().enumerate().all(|(i, &j)| i == j);
285        let sorted_ranges: Arc<[std::ops::Range<u64>]> = order
286            .iter()
287            .map(|&i| ranges[i].start as u64..ranges[i].end as u64)
288            .collect();
289        let total_rows: u64 = sorted_ranges.iter().map(|r| r.end - r.start).sum();
290        let batches = self
291            .0
292            .read_stream_projected(
293                ReadBatchParams::Ranges(sorted_ranges),
294                (total_rows as u32).max(1),
295                16,
296                projection,
297                FilterExpression::no_filter(),
298            )
299            .await?
300            .try_collect::<Vec<_>>()
301            .await?;
302        let merged = match batches.len() {
303            0 => return empty_batch(),
304            1 => batches.into_iter().next().unwrap(),
305            _ => {
306                let schema = batches[0].schema();
307                arrow_select::concat::concat_batches(&schema, &batches)?
308            }
309        };
310        if already_sorted {
311            return Ok(merged);
312        }
313        let sorted_sizes: Vec<u32> = order
314            .iter()
315            .map(|&i| (ranges[i].end - ranges[i].start) as u32)
316            .collect();
317        let mut sorted_offsets = Vec::with_capacity(sorted_sizes.len());
318        let mut acc = 0u32;
319        for &s in &sorted_sizes {
320            sorted_offsets.push(acc);
321            acc += s;
322        }
323        let mut sorted_pos = vec![0usize; ranges.len()];
324        for (sp, &oi) in order.iter().enumerate() {
325            sorted_pos[oi] = sp;
326        }
327        let mut take_indices = Vec::with_capacity(total_rows as usize);
328        for &sp in &sorted_pos {
329            for k in 0..sorted_sizes[sp] {
330                take_indices.push(sorted_offsets[sp] + k);
331            }
332        }
333        let take_arr = arrow_array::UInt32Array::from(take_indices);
334        Ok(arrow_select::take::take_record_batch(&merged, &take_arr)?)
335    }
336
337    async fn read_range_stream(
338        &self,
339        range: std::ops::Range<usize>,
340        projection: Option<&[&str]>,
341    ) -> Result<Pin<Box<dyn lance_io::stream::RecordBatchStream>>> {
342        if range.is_empty() {
343            return Ok(Box::pin(lance_io::stream::RecordBatchStreamAdapter::new(
344                Arc::new(self.0.schema().as_ref().into()),
345                futures::stream::empty(),
346            )));
347        }
348        let projection = if let Some(projection) = projection {
349            ReaderProjection::from_column_names(
350                self.0.metadata().version(),
351                self.0.schema(),
352                projection,
353            )?
354        } else {
355            ReaderProjection::from_whole_schema(self.0.schema(), self.0.metadata().version())
356        };
357        self.0
358            .read_stream_projected(
359                ReadBatchParams::Range(range),
360                4096,
361                2,
362                projection,
363                FilterExpression::no_filter(),
364            )
365            .await
366    }
367
368    // V2 format has removed the row group concept,
369    // so here we assume each batch is with 4096 rows.
370    async fn num_batches(&self, batch_size: u64) -> u32 {
371        CurrentFileReader::num_rows(&self.0).div_ceil(batch_size) as u32
372    }
373
374    fn num_rows(&self) -> usize {
375        CurrentFileReader::num_rows(&self.0) as usize
376    }
377
378    fn schema(&self) -> &lance_core::datatypes::Schema {
379        CurrentFileReader::schema(&self.0)
380    }
381
382    fn file_size_bytes(&self) -> Option<u64> {
383        // The manifest records each index file's size and passes it to the reader
384        // at open, so it's already in metadata here (no extra I/O).
385        Some(self.0.metadata().file_size())
386    }
387}
388
389#[async_trait]
390impl IndexStore for LanceIndexStore {
391    fn as_any(&self) -> &dyn Any {
392        self
393    }
394
395    fn clone_arc(&self) -> Arc<dyn IndexStore> {
396        Arc::new(self.clone())
397    }
398
399    fn io_parallelism(&self) -> usize {
400        self.object_store.io_parallelism()
401    }
402
403    async fn new_index_file(
404        &self,
405        name: &str,
406        schema: Arc<Schema>,
407    ) -> Result<Box<dyn IndexWriter>> {
408        let path = self.index_file_path(name)?;
409        let schema = schema.as_ref().try_into()?;
410        let writer = self.object_store.create(&path).await?;
411        let writer = current_writer::FileWriter::try_new(
412            writer,
413            schema,
414            current_writer::FileWriterOptions {
415                format_version: Some(self.format_version),
416                ..Default::default()
417            },
418        )?;
419        Ok(Box::new(LanceIndexWriter {
420            path: name.to_string(),
421            inner: writer,
422        }))
423    }
424
425    fn with_io_priority(&self, io_priority: u64) -> Arc<dyn IndexStore> {
426        // The `scheduler` is shared (`Arc`), so this clone is cheap and the new
427        // priority only affects requests this clone submits.
428        Arc::new(Self {
429            io_priority,
430            ..self.clone()
431        })
432    }
433
434    async fn open_index_file(&self, name: &str) -> Result<Arc<dyn IndexReader>> {
435        let path = self.index_file_path(name)?;
436        // Use cached file size if available, otherwise unknown (requires HEAD call)
437        let cached_size = self
438            .file_sizes
439            .get(name)
440            .map(|&size| CachedFileSize::new(size))
441            .unwrap_or_else(CachedFileSize::unknown);
442        let file_scheduler = self
443            .scheduler
444            .open_file_with_priority(&path, self.io_priority, &cached_size)
445            .await?;
446        match CurrentFileReader::try_open(
447            file_scheduler,
448            None,
449            Arc::<DecoderPlugins>::default(),
450            &self.metadata_cache,
451            FileReaderOptions::default(),
452        )
453        .await
454        {
455            Ok(reader) => Ok(Arc::new(CurrentIndexReader(reader))),
456            Err(e) => {
457                // If the error is a version conflict we can try to read the file with v1 reader
458                if let Error::VersionConflict { .. } = e {
459                    let path = self.index_file_path(name)?;
460                    let file_reader = V1FileReader::try_new_self_described(
461                        &self.object_store,
462                        &path,
463                        Some(&self.metadata_cache),
464                    )
465                    .await?;
466                    Ok(Arc::new(V1IndexReader(file_reader)))
467                } else {
468                    Err(e)
469                }
470            }
471        }
472    }
473
474    async fn copy_index_file(&self, name: &str, dest_store: &dyn IndexStore) -> Result<IndexFile> {
475        self.copy_index_file_to(name, name, dest_store).await
476    }
477
478    async fn copy_index_file_to(
479        &self,
480        name: &str,
481        new_name: &str,
482        dest_store: &dyn IndexStore,
483    ) -> Result<IndexFile> {
484        let path = self.index_file_path(name)?;
485
486        let other_store = dest_store.as_any().downcast_ref::<Self>();
487        match other_store {
488            Some(dest_store) if dest_store.object_store.scheme() == self.object_store.scheme() => {
489                // If both this store and the destination are lance stores we can use object_store's copy
490                // This does blindly assume that both stores are using the same underlying object_store
491                // but there is no easy way to verify this and it happens to always be true at the moment
492                let dest_path = dest_store.index_file_path(new_name)?;
493                self.object_store.copy(&path, &dest_path).await?;
494                let size_bytes = match self.file_sizes.get(name) {
495                    Some(size_bytes) => *size_bytes,
496                    None => self.object_store.size(&path).await?,
497                };
498                Ok(IndexFile {
499                    path: new_name.to_string(),
500                    size_bytes,
501                })
502            }
503            _ => {
504                let reader = self.open_index_file(name).await?;
505                let mut writer = dest_store
506                    .new_index_file(new_name, Arc::new(reader.schema().into()))
507                    .await?;
508
509                for offset in (0..reader.num_rows()).step_by(4096) {
510                    let next_offset = min(offset + 4096, reader.num_rows());
511                    let batch = reader.read_range(offset..next_offset, None).await?;
512                    writer.write_record_batch(batch).await?;
513                }
514                writer.finish().await
515            }
516        }
517    }
518
519    async fn rename_index_file(&self, name: &str, new_name: &str) -> Result<IndexFile> {
520        let path = self.index_file_path(name)?;
521        let new_path = self.index_file_path(new_name)?;
522        self.object_store.copy(&path, &new_path).await?;
523        self.object_store.delete(&path).await?;
524        let size_bytes = match self.file_sizes.get(name) {
525            Some(size_bytes) => *size_bytes,
526            None => self.object_store.size(&new_path).await?,
527        };
528        Ok(IndexFile {
529            path: new_name.to_string(),
530            size_bytes,
531        })
532    }
533
534    async fn delete_index_file(&self, name: &str) -> Result<()> {
535        let path = self.index_file_path(name)?;
536        self.object_store.delete(&path).await
537    }
538
539    async fn list_files_with_sizes(&self) -> Result<Vec<IndexFile>> {
540        let files = list_index_files_with_sizes(&self.object_store, &self.index_dir).await?;
541        Ok(files
542            .into_iter()
543            .map(|f| IndexFile {
544                path: f.path,
545                size_bytes: f.size_bytes,
546            })
547            .collect())
548    }
549}
550
551#[cfg(test)]
552mod tests {
553
554    use std::{collections::HashMap, ops::Bound};
555
556    use crate::metrics::NoOpMetricsCollector;
557    use crate::pbold;
558    use crate::scalar::bitmap::BitmapIndexPlugin;
559    use crate::scalar::btree::{BTreeIndexPlugin, BTreeParameters};
560    use crate::scalar::label_list::LabelListIndexPlugin;
561    use crate::scalar::registry::{BasicTrainer, ScalarIndexPlugin, VALUE_COLUMN_NAME};
562    use crate::scalar::{
563        LabelListQuery, SargableQuery, ScalarIndex, SearchResult,
564        bitmap::BitmapIndex,
565        btree::{DEFAULT_BTREE_BATCH_SIZE, train_btree_index},
566    };
567
568    use super::*;
569    use arrow::{buffer::ScalarBuffer, datatypes::UInt8Type};
570    use arrow_array::{
571        ListArray, RecordBatchIterator, RecordBatchReader, StringArray, UInt64Array,
572        cast::AsArray,
573        types::{Int32Type, UInt64Type},
574    };
575    use arrow_schema::Schema as ArrowSchema;
576    use arrow_schema::{DataType, Field, TimeUnit};
577    use arrow_select::take::TakeOptions;
578    use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
579    use datafusion_common::ScalarValue;
580    use futures::FutureExt;
581    use lance_core::ROW_ID;
582    use lance_core::utils::row_addr_remap::RowAddrRemap;
583    use lance_core::utils::tempfile::TempDir;
584    use lance_datagen::{ArrayGeneratorExt, BatchCount, ByteCount, RowCount, array, gen_batch};
585    use lance_select::{RowAddrTreeMap, RowSetOps};
586
587    fn test_store(tempdir: &TempDir) -> Arc<dyn IndexStore> {
588        let test_path = tempdir.obj_path();
589        let (object_store, test_path) = ObjectStore::from_uri(test_path.as_ref())
590            .now_or_never()
591            .unwrap()
592            .unwrap();
593        let cache = Arc::new(lance_core::cache::LanceCache::with_capacity(
594            128 * 1024 * 1024,
595        ));
596        Arc::new(LanceIndexStore::new(object_store, test_path, cache))
597    }
598
599    #[tokio::test]
600    async fn test_store_deep_size_excludes_metadata_cache() {
601        // The metadata cache is session-scoped and shared; a store must not count
602        // it as its own footprint, so growing the cache must not change store size.
603        struct BlobKey;
604        impl lance_core::cache::CacheKey for BlobKey {
605            type ValueType = Vec<u8>;
606            fn key(&self) -> std::borrow::Cow<'_, str> {
607                std::borrow::Cow::Borrowed("blob")
608            }
609            fn type_name() -> &'static str {
610                "Vec<u8>"
611            }
612            fn stable_type_id() -> &'static str {
613                "lance.scalar.lance-format.Blob"
614            }
615            fn schema() -> lance_core::cache::CacheKeySchema {
616                lance_core::cache::CacheKeySchema::new("lance.scalar.lance-format.blob-key", 1)
617            }
618            fn write_key(&self, builder: &mut lance_core::cache::KeyBuilder) {
619                builder.write_variant(0);
620            }
621        }
622
623        let index_dir = TempDir::default();
624        let test_path = index_dir.obj_path();
625        let (object_store, test_path) = ObjectStore::from_uri(test_path.as_ref())
626            .now_or_never()
627            .unwrap()
628            .unwrap();
629        let cache = Arc::new(lance_core::cache::LanceCache::with_capacity(
630            128 * 1024 * 1024,
631        ));
632        let store = LanceIndexStore::new(object_store, test_path, cache.clone());
633
634        let before = store.deep_size_of();
635        cache
636            .insert_with_key(&BlobKey, Arc::new(vec![0u8; 4 * 1024 * 1024]))
637            .await;
638        // Force moka to commit the write so a cache-inclusive size would grow.
639        let _ = cache.size_bytes().await;
640        let after = store.deep_size_of();
641
642        assert_eq!(
643            before, after,
644            "store deep size must exclude the shared metadata cache"
645        );
646    }
647
648    async fn train_index(
649        index_store: &Arc<dyn IndexStore>,
650        data: impl RecordBatchReader + Send + Sync + 'static,
651        custom_batch_size: Option<u64>,
652    ) {
653        let batch_size = custom_batch_size.unwrap_or(DEFAULT_BTREE_BATCH_SIZE);
654        let params = BTreeParameters {
655            zone_size: Some(batch_size),
656            range_id: None,
657        };
658        let params = serde_json::to_string(&params).unwrap();
659        let btree_plugin = BTreeIndexPlugin;
660        let data = lance_datafusion::utils::reader_to_stream(Box::new(data));
661        let request = btree_plugin
662            .new_training_request(
663                &params,
664                &Field::new(VALUE_COLUMN_NAME, DataType::Int32, false),
665            )
666            .unwrap();
667        btree_plugin
668            .train_index(
669                data,
670                index_store.as_ref(),
671                request,
672                None,
673                crate::progress::noop_progress(),
674            )
675            .await
676            .unwrap();
677    }
678
679    fn default_details<T: prost::Message + prost::Name + std::default::Default>() -> prost_types::Any
680    {
681        prost_types::Any::from_msg(&T::default()).unwrap()
682    }
683
684    #[tokio::test]
685    async fn test_global_buffer_round_trip() {
686        let tempdir = TempDir::default();
687        let index_store = test_store(&tempdir);
688
689        let mut writer = index_store
690            .new_index_file("global-buffer.lance", Arc::new(Schema::empty()))
691            .await
692            .unwrap();
693        let expected = bytes::Bytes::from_static(b"scalar-global-buffer");
694        let buffer_idx = writer.add_global_buffer(expected.clone()).await.unwrap();
695        let write_summary = writer.finish().await.unwrap();
696        let files = index_store.list_files_with_sizes().await.unwrap();
697        assert_eq!(files.len(), 1);
698        assert_eq!(files[0].path, "global-buffer.lance");
699        assert_eq!(write_summary.size_bytes, files[0].size_bytes);
700
701        let reader = index_store
702            .open_index_file("global-buffer.lance")
703            .await
704            .unwrap();
705        let actual = reader.read_global_buffer(buffer_idx).await.unwrap();
706
707        assert_eq!(actual, expected);
708    }
709
710    #[tokio::test]
711    async fn test_basic_btree() {
712        let tempdir = TempDir::default();
713        let index_store = test_store(&tempdir);
714        let data = gen_batch()
715            .col(VALUE_COLUMN_NAME, array::step::<Int32Type>())
716            .col(ROW_ID, array::step::<UInt64Type>())
717            .into_reader_rows(RowCount::from(4096), BatchCount::from(100));
718        train_index(&index_store, data, None).await;
719        let index = BTreeIndexPlugin
720            .load_index(
721                index_store,
722                &default_details::<pbold::BTreeIndexDetails>(),
723                None,
724                &LanceCache::no_cache(),
725            )
726            .await
727            .unwrap();
728
729        let result = index
730            .search(
731                &SargableQuery::Equals(ScalarValue::Int32(Some(10000))),
732                &NoOpMetricsCollector,
733            )
734            .await
735            .unwrap();
736
737        assert!(result.is_exact());
738        let row_ids = result.row_addrs().true_rows();
739        assert_eq!(Some(1), row_ids.len());
740        assert!(row_ids.contains(10000));
741
742        let result = index
743            .search(
744                &SargableQuery::Range(
745                    Bound::Unbounded,
746                    Bound::Excluded(ScalarValue::Int32(Some(-100))),
747                ),
748                &NoOpMetricsCollector,
749            )
750            .await
751            .unwrap();
752
753        assert!(result.is_exact());
754        let row_addrs = result.row_addrs().true_rows();
755
756        assert_eq!(Some(0), row_addrs.len());
757
758        let result = index
759            .search(
760                &SargableQuery::Range(
761                    Bound::Unbounded,
762                    Bound::Excluded(ScalarValue::Int32(Some(100))),
763                ),
764                &NoOpMetricsCollector,
765            )
766            .await
767            .unwrap();
768
769        assert!(result.is_exact());
770        let row_addrs = result.row_addrs().true_rows();
771
772        assert_eq!(Some(100), row_addrs.len());
773    }
774
775    #[tokio::test]
776    async fn test_btree_update() {
777        let index_dir = TempDir::default();
778        let index_store = test_store(&index_dir);
779        let data = gen_batch()
780            .col(VALUE_COLUMN_NAME, array::step::<Int32Type>())
781            .col(ROW_ID, array::step::<UInt64Type>())
782            .into_reader_rows(RowCount::from(4096), BatchCount::from(100));
783        train_index(&index_store, data, None).await;
784        let index = BTreeIndexPlugin
785            .load_index(
786                index_store,
787                &default_details::<pbold::BTreeIndexDetails>(),
788                None,
789                &LanceCache::no_cache(),
790            )
791            .await
792            .unwrap();
793
794        let data = gen_batch()
795            .col(
796                VALUE_COLUMN_NAME,
797                array::step_custom::<Int32Type>(4096 * 100, 1),
798            )
799            .col(ROW_ID, array::step_custom::<UInt64Type>(4096 * 100, 1))
800            .into_reader_rows(RowCount::from(4096), BatchCount::from(100));
801
802        let updated_index_dir = TempDir::default();
803        let updated_index_store = test_store(&updated_index_dir);
804        index
805            .update(
806                lance_datafusion::utils::reader_to_stream(Box::new(data)),
807                updated_index_store.as_ref(),
808                None,
809            )
810            .await
811            .unwrap();
812        let updated_index = BTreeIndexPlugin
813            .load_index(
814                updated_index_store,
815                &default_details::<pbold::BTreeIndexDetails>(),
816                None,
817                &LanceCache::no_cache(),
818            )
819            .await
820            .unwrap();
821
822        let result = updated_index
823            .search(
824                &SargableQuery::Equals(ScalarValue::Int32(Some(10000))),
825                &NoOpMetricsCollector,
826            )
827            .await
828            .unwrap();
829
830        assert!(result.is_exact());
831        let row_addrs = result.row_addrs().true_rows();
832
833        assert_eq!(Some(1), row_addrs.len());
834        assert!(row_addrs.contains(10000));
835
836        let result = updated_index
837            .search(
838                &SargableQuery::Equals(ScalarValue::Int32(Some(500_000))),
839                &NoOpMetricsCollector,
840            )
841            .await
842            .unwrap();
843
844        assert!(result.is_exact());
845        let row_addrs = result.row_addrs().true_rows();
846
847        assert_eq!(Some(1), row_addrs.len());
848        assert!(row_addrs.contains(500_000));
849    }
850
851    async fn check(index: &Arc<dyn ScalarIndex>, query: SargableQuery, expected: &[u64]) {
852        let results = index.search(&query, &NoOpMetricsCollector).await.unwrap();
853        assert!(results.is_exact());
854        let expected_arr = RowAddrTreeMap::from_iter(expected);
855        assert_eq!(&results.row_addrs().true_rows(), &expected_arr);
856    }
857
858    #[tokio::test]
859    async fn test_btree_with_gaps() {
860        let tempdir = TempDir::default();
861        let index_store = test_store(&tempdir);
862        let batch_one = gen_batch()
863            .col(
864                VALUE_COLUMN_NAME,
865                array::cycle::<Int32Type>(vec![0, 1, 4, 5]),
866            )
867            .col(ROW_ID, array::cycle::<UInt64Type>(vec![0, 1, 2, 3]))
868            .into_batch_rows(RowCount::from(4));
869        let batch_two = gen_batch()
870            .col(
871                VALUE_COLUMN_NAME,
872                array::cycle::<Int32Type>(vec![10, 11, 11, 15]),
873            )
874            .col(ROW_ID, array::cycle::<UInt64Type>(vec![40, 50, 60, 70]))
875            .into_batch_rows(RowCount::from(4));
876        let batch_three = gen_batch()
877            .col(
878                VALUE_COLUMN_NAME,
879                array::cycle::<Int32Type>(vec![15, 15, 15, 15]),
880            )
881            .col(ROW_ID, array::cycle::<UInt64Type>(vec![400, 500, 600, 700]))
882            .into_batch_rows(RowCount::from(4));
883        let batch_four = gen_batch()
884            .col(
885                VALUE_COLUMN_NAME,
886                array::cycle::<Int32Type>(vec![15, 16, 20, 20]),
887            )
888            .col(
889                ROW_ID,
890                array::cycle::<UInt64Type>(vec![4000, 5000, 6000, 7000]),
891            )
892            .into_batch_rows(RowCount::from(4));
893        let batches = vec![batch_one, batch_two, batch_three, batch_four];
894        let schema = Arc::new(Schema::new(vec![
895            Field::new(VALUE_COLUMN_NAME, DataType::Int32, false),
896            Field::new(ROW_ID, DataType::UInt64, false),
897        ]));
898        let data = RecordBatchIterator::new(batches, schema);
899        train_index(&index_store, data, Some(4)).await;
900        let index = BTreeIndexPlugin
901            .load_index(
902                index_store,
903                &default_details::<pbold::BTreeIndexDetails>(),
904                None,
905                &LanceCache::no_cache(),
906            )
907            .await
908            .unwrap();
909
910        // The above should create four pages
911        //
912        // 0 - 5
913        // 10 - 15
914        // 15 - 15
915        // 15 - 20
916        //
917        // This will help us test various indexing corner cases
918
919        // No results (off the left side)
920        check(
921            &index,
922            SargableQuery::Equals(ScalarValue::Int32(Some(-3))),
923            &[],
924        )
925        .await;
926
927        check(
928            &index,
929            SargableQuery::Range(
930                Bound::Unbounded,
931                Bound::Included(ScalarValue::Int32(Some(-3))),
932            ),
933            &[],
934        )
935        .await;
936
937        check(
938            &index,
939            SargableQuery::Range(
940                Bound::Included(ScalarValue::Int32(Some(-10))),
941                Bound::Included(ScalarValue::Int32(Some(-3))),
942            ),
943            &[],
944        )
945        .await;
946
947        // Hitting the middle of a bucket
948        check(
949            &index,
950            SargableQuery::Equals(ScalarValue::Int32(Some(4))),
951            &[2],
952        )
953        .await;
954
955        // Hitting a gap between two buckets
956        check(
957            &index,
958            SargableQuery::Equals(ScalarValue::Int32(Some(7))),
959            &[],
960        )
961        .await;
962
963        // Hitting the lowest of the overlapping buckets
964        check(
965            &index,
966            SargableQuery::Equals(ScalarValue::Int32(Some(11))),
967            &[50, 60],
968        )
969        .await;
970
971        // Hitting the 15 shared on all three buckets
972        check(
973            &index,
974            SargableQuery::Equals(ScalarValue::Int32(Some(15))),
975            &[70, 400, 500, 600, 700, 4000],
976        )
977        .await;
978
979        // Hitting the upper part of the three overlapping buckets
980        check(
981            &index,
982            SargableQuery::Equals(ScalarValue::Int32(Some(20))),
983            &[6000, 7000],
984        )
985        .await;
986
987        // Ranges that capture multiple buckets
988        check(
989            &index,
990            SargableQuery::Range(
991                Bound::Unbounded,
992                Bound::Included(ScalarValue::Int32(Some(11))),
993            ),
994            &[0, 1, 2, 3, 40, 50, 60],
995        )
996        .await;
997
998        check(
999            &index,
1000            SargableQuery::Range(
1001                Bound::Unbounded,
1002                Bound::Excluded(ScalarValue::Int32(Some(11))),
1003            ),
1004            &[0, 1, 2, 3, 40],
1005        )
1006        .await;
1007
1008        check(
1009            &index,
1010            SargableQuery::Range(
1011                Bound::Included(ScalarValue::Int32(Some(4))),
1012                Bound::Unbounded,
1013            ),
1014            &[
1015                2, 3, 40, 50, 60, 70, 400, 500, 600, 700, 4000, 5000, 6000, 7000,
1016            ],
1017        )
1018        .await;
1019
1020        check(
1021            &index,
1022            SargableQuery::Range(
1023                Bound::Included(ScalarValue::Int32(Some(4))),
1024                Bound::Included(ScalarValue::Int32(Some(11))),
1025            ),
1026            &[2, 3, 40, 50, 60],
1027        )
1028        .await;
1029
1030        check(
1031            &index,
1032            SargableQuery::Range(
1033                Bound::Included(ScalarValue::Int32(Some(4))),
1034                Bound::Excluded(ScalarValue::Int32(Some(11))),
1035            ),
1036            &[2, 3, 40],
1037        )
1038        .await;
1039
1040        check(
1041            &index,
1042            SargableQuery::Range(
1043                Bound::Excluded(ScalarValue::Int32(Some(4))),
1044                Bound::Unbounded,
1045            ),
1046            &[
1047                3, 40, 50, 60, 70, 400, 500, 600, 700, 4000, 5000, 6000, 7000,
1048            ],
1049        )
1050        .await;
1051
1052        check(
1053            &index,
1054            SargableQuery::Range(
1055                Bound::Excluded(ScalarValue::Int32(Some(4))),
1056                Bound::Included(ScalarValue::Int32(Some(11))),
1057            ),
1058            &[3, 40, 50, 60],
1059        )
1060        .await;
1061
1062        check(
1063            &index,
1064            SargableQuery::Range(
1065                Bound::Excluded(ScalarValue::Int32(Some(4))),
1066                Bound::Excluded(ScalarValue::Int32(Some(11))),
1067            ),
1068            &[3, 40],
1069        )
1070        .await;
1071
1072        check(
1073            &index,
1074            SargableQuery::Range(
1075                Bound::Excluded(ScalarValue::Int32(Some(-50))),
1076                Bound::Excluded(ScalarValue::Int32(Some(1000))),
1077            ),
1078            &[
1079                0, 1, 2, 3, 40, 50, 60, 70, 400, 500, 600, 700, 4000, 5000, 6000, 7000,
1080            ],
1081        )
1082        .await;
1083    }
1084
1085    #[tokio::test]
1086    async fn test_btree_types() {
1087        for data_type in &[
1088            DataType::Boolean,
1089            DataType::Int32,
1090            DataType::Utf8,
1091            DataType::Float32,
1092            DataType::Date32,
1093            DataType::Timestamp(TimeUnit::Nanosecond, None),
1094            DataType::Date64,
1095            DataType::Date32,
1096            DataType::Time64(TimeUnit::Nanosecond),
1097            DataType::Time32(TimeUnit::Second),
1098            DataType::FixedSizeBinary(16),
1099            // Not supported today, error from datafusion:
1100            // Min/max accumulator not implemented for Duration(Nanosecond)
1101            // DataType::Duration(TimeUnit::Nanosecond),
1102        ] {
1103            let tempdir = TempDir::default();
1104            let index_store = test_store(&tempdir);
1105            let data: RecordBatch = gen_batch()
1106                .col(VALUE_COLUMN_NAME, array::rand_type(data_type))
1107                .col(ROW_ID, array::step::<UInt64Type>())
1108                .into_batch_rows(RowCount::from(4096 * 3))
1109                .unwrap();
1110
1111            let sample_value = ScalarValue::try_from_array(data.column(0), 0).unwrap();
1112            let sample_row_id = data.column(1).as_primitive::<UInt64Type>().value(0);
1113
1114            let sort_indices = arrow::compute::sort_to_indices(data.column(0), None, None).unwrap();
1115            let sorted_values = arrow_select::take::take(
1116                data.column(0),
1117                &sort_indices,
1118                Some(TakeOptions {
1119                    check_bounds: false,
1120                }),
1121            )
1122            .unwrap();
1123            let sorted_row_ids = arrow_select::take::take(
1124                data.column(1),
1125                &sort_indices,
1126                Some(TakeOptions {
1127                    check_bounds: false,
1128                }),
1129            )
1130            .unwrap();
1131            let sorted_batch =
1132                RecordBatch::try_new(data.schema().clone(), vec![sorted_values, sorted_row_ids])
1133                    .unwrap();
1134
1135            let batch_one = sorted_batch.slice(0, 4096);
1136            let batch_two = sorted_batch.slice(4096, 4096);
1137            let batch_three = sorted_batch.slice(8192, 4096);
1138            let training_data = RecordBatchIterator::new(
1139                vec![batch_one, batch_two, batch_three].into_iter().map(Ok),
1140                data.schema().clone(),
1141            );
1142
1143            train_index(&index_store, training_data, None).await;
1144            let index = BTreeIndexPlugin
1145                .load_index(
1146                    index_store,
1147                    &default_details::<pbold::BTreeIndexDetails>(),
1148                    None,
1149                    &LanceCache::no_cache(),
1150                )
1151                .await
1152                .unwrap();
1153
1154            let result = index
1155                .search(&SargableQuery::Equals(sample_value), &NoOpMetricsCollector)
1156                .await
1157                .unwrap();
1158
1159            assert!(result.is_exact());
1160            let row_addrs = result.row_addrs().true_rows();
1161
1162            // The random data may have had duplicates so there might be more than 1 result
1163            // but even for boolean we shouldn't match the entire thing
1164            assert!(!row_addrs.is_empty());
1165            assert!(row_addrs.len().unwrap() < data.num_rows() as u64);
1166            assert!(row_addrs.contains(sample_row_id));
1167        }
1168    }
1169
1170    #[tokio::test]
1171    async fn btree_entire_null_page() {
1172        let tempdir = TempDir::default();
1173        let index_store = test_store(&tempdir);
1174        let batch = gen_batch()
1175            .col(
1176                VALUE_COLUMN_NAME,
1177                array::rand_utf8(ByteCount::from(0), false).with_nulls(&[true]),
1178            )
1179            .col(ROW_ID, array::step::<UInt64Type>())
1180            .into_batch_rows(RowCount::from(4096));
1181        assert_eq!(
1182            batch.as_ref().unwrap()[VALUE_COLUMN_NAME].null_count(),
1183            4096
1184        );
1185        let batches = vec![batch];
1186        let schema = Arc::new(Schema::new(vec![
1187            Field::new(VALUE_COLUMN_NAME, DataType::Utf8, true),
1188            Field::new(ROW_ID, DataType::UInt64, false),
1189        ]));
1190        let data = RecordBatchIterator::new(batches, schema);
1191        let data = lance_datafusion::utils::reader_to_stream(Box::new(data));
1192
1193        train_btree_index(
1194            data,
1195            index_store.as_ref(),
1196            DEFAULT_BTREE_BATCH_SIZE,
1197            None,
1198            None,
1199        )
1200        .await
1201        .unwrap();
1202
1203        let index = BTreeIndexPlugin
1204            .load_index(
1205                index_store,
1206                &default_details::<pbold::BTreeIndexDetails>(),
1207                None,
1208                &LanceCache::no_cache(),
1209            )
1210            .await
1211            .unwrap();
1212
1213        let result = index
1214            .search(
1215                &SargableQuery::Equals(ScalarValue::Utf8(Some("foo".to_string()))),
1216                &NoOpMetricsCollector,
1217            )
1218            .await
1219            .unwrap();
1220
1221        assert!(result.is_exact());
1222        let row_addrs = result.row_addrs().true_rows();
1223
1224        assert!(row_addrs.is_empty());
1225
1226        let result = index
1227            .search(&SargableQuery::IsNull(), &NoOpMetricsCollector)
1228            .await
1229            .unwrap();
1230        assert!(result.is_exact());
1231        let row_addrs = result.row_addrs().true_rows();
1232        assert_eq!(row_addrs.len(), Some(4096));
1233    }
1234
1235    async fn train_bitmap(
1236        index_store: &Arc<dyn IndexStore>,
1237        data: impl RecordBatchReader + Send + Sync + 'static,
1238    ) {
1239        // Sort the data by value column (nulls first) to match the production
1240        // scanner behavior (TrainingOrdering::Values).
1241        let schema = data.schema();
1242        let batches: Vec<_> = data
1243            .into_iter()
1244            .collect::<std::result::Result<Vec<_>, _>>()
1245            .unwrap();
1246        let combined = arrow::compute::concat_batches(&schema, &batches).unwrap();
1247        let options = arrow::compute::SortOptions {
1248            descending: false,
1249            nulls_first: true,
1250        };
1251        let indices =
1252            arrow::compute::sort_to_indices(combined.column(0), Some(options), None).unwrap();
1253        let sorted_columns: Vec<_> = combined
1254            .columns()
1255            .iter()
1256            .map(|col| arrow::compute::take(col.as_ref(), &indices, None).unwrap())
1257            .collect();
1258        let sorted_batch = RecordBatch::try_new(schema.clone(), sorted_columns).unwrap();
1259        let stream = Box::pin(RecordBatchStreamAdapter::new(
1260            schema,
1261            futures::stream::once(async move { Ok(sorted_batch) }),
1262        ));
1263
1264        let request = BitmapIndexPlugin
1265            .new_training_request("{}", &Field::new(VALUE_COLUMN_NAME, DataType::Int32, false))
1266            .unwrap();
1267        BitmapIndexPlugin
1268            .train_index(
1269                stream,
1270                index_store.as_ref(),
1271                request,
1272                None,
1273                crate::progress::noop_progress(),
1274            )
1275            .await
1276            .unwrap();
1277    }
1278
1279    #[tokio::test]
1280    async fn test_bitmap_working() {
1281        let tempdir = TempDir::default();
1282        let index_store = test_store(&tempdir);
1283
1284        let schema = Arc::new(ArrowSchema::new(vec![
1285            Field::new(VALUE_COLUMN_NAME, DataType::Utf8, true),
1286            Field::new(ROW_ID, DataType::UInt64, false),
1287        ]));
1288
1289        let batch1 = RecordBatch::try_new(
1290            schema.clone(),
1291            vec![
1292                Arc::new(StringArray::from(vec![Some("abcd"), None, Some("abcd")])),
1293                Arc::new(UInt64Array::from(vec![1, 2, 3])),
1294            ],
1295        )
1296        .unwrap();
1297
1298        let batch2 = RecordBatch::try_new(
1299            schema.clone(),
1300            vec![
1301                Arc::new(StringArray::from(vec![
1302                    Some("apple"),
1303                    Some("hello"),
1304                    Some("abcd"),
1305                ])),
1306                Arc::new(UInt64Array::from(vec![4, 5, 6])),
1307            ],
1308        )
1309        .unwrap();
1310
1311        let batches = vec![batch1, batch2];
1312        let data = RecordBatchIterator::new(batches.into_iter().map(Ok), schema);
1313        train_bitmap(&index_store, data).await;
1314
1315        let index = BitmapIndex::load(index_store, None, &LanceCache::no_cache())
1316            .await
1317            .unwrap();
1318
1319        let result = index
1320            .search(
1321                &SargableQuery::Equals(ScalarValue::Utf8(None)),
1322                &NoOpMetricsCollector,
1323            )
1324            .await
1325            .unwrap();
1326
1327        assert!(result.is_exact());
1328        let row_addrs = result.row_addrs().true_rows();
1329        assert_eq!(Some(1), row_addrs.len());
1330        assert!(row_addrs.contains(2));
1331
1332        let result = index
1333            .search(
1334                &SargableQuery::Equals(ScalarValue::Utf8(Some("abcd".to_string()))),
1335                &NoOpMetricsCollector,
1336            )
1337            .await
1338            .unwrap();
1339
1340        assert!(result.is_exact());
1341        let row_addrs = result.row_addrs().true_rows();
1342        assert_eq!(Some(3), row_addrs.len());
1343        assert!(row_addrs.contains(1));
1344        assert!(row_addrs.contains(3));
1345        assert!(row_addrs.contains(6));
1346    }
1347
1348    #[tokio::test]
1349    async fn test_basic_bitmap() {
1350        let tempdir = TempDir::default();
1351        let index_store = test_store(&tempdir);
1352        let data = gen_batch()
1353            .col(VALUE_COLUMN_NAME, array::step::<Int32Type>())
1354            .col(ROW_ID, array::step::<UInt64Type>())
1355            .into_reader_rows(RowCount::from(4096), BatchCount::from(100));
1356        train_bitmap(&index_store, data).await;
1357        let index = BitmapIndex::load(index_store, None, &LanceCache::no_cache())
1358            .await
1359            .unwrap();
1360
1361        let result = index
1362            .search(
1363                &SargableQuery::Equals(ScalarValue::Int32(Some(10000))),
1364                &NoOpMetricsCollector,
1365            )
1366            .await
1367            .unwrap();
1368
1369        assert!(result.is_exact());
1370        let row_addrs = result.row_addrs().true_rows();
1371        assert_eq!(Some(1), row_addrs.len());
1372        assert!(row_addrs.contains(10000));
1373
1374        let result = index
1375            .search(
1376                &SargableQuery::Range(
1377                    Bound::Unbounded,
1378                    Bound::Excluded(ScalarValue::Int32(Some(-100))),
1379                ),
1380                &NoOpMetricsCollector,
1381            )
1382            .await
1383            .unwrap();
1384
1385        assert!(result.is_exact());
1386        let row_addrs = result.row_addrs().true_rows();
1387        assert!(row_addrs.is_empty());
1388
1389        let result = index
1390            .search(
1391                &SargableQuery::Range(
1392                    Bound::Unbounded,
1393                    Bound::Excluded(ScalarValue::Int32(Some(100))),
1394                ),
1395                &NoOpMetricsCollector,
1396            )
1397            .await
1398            .unwrap();
1399
1400        assert!(result.is_exact());
1401        let row_addrs = result.row_addrs().true_rows();
1402        assert_eq!(Some(100), row_addrs.len());
1403    }
1404
1405    async fn check_bitmap(index: &BitmapIndex, query: SargableQuery, expected: &[u64]) {
1406        let results = index.search(&query, &NoOpMetricsCollector).await.unwrap();
1407        assert!(results.is_exact());
1408        let expected_arr = RowAddrTreeMap::from_iter(expected);
1409        assert_eq!(&results.row_addrs().true_rows(), &expected_arr);
1410    }
1411
1412    #[tokio::test]
1413    async fn test_bitmap_with_gaps() {
1414        let tempdir = TempDir::default();
1415        let index_store = test_store(&tempdir);
1416        let batch_one = gen_batch()
1417            .col(
1418                VALUE_COLUMN_NAME,
1419                array::cycle::<Int32Type>(vec![0, 1, 4, 5]),
1420            )
1421            .col(ROW_ID, array::cycle::<UInt64Type>(vec![0, 1, 2, 3]))
1422            .into_batch_rows(RowCount::from(4));
1423        let batch_two = gen_batch()
1424            .col(
1425                VALUE_COLUMN_NAME,
1426                array::cycle::<Int32Type>(vec![10, 11, 11, 15]),
1427            )
1428            .col(ROW_ID, array::cycle::<UInt64Type>(vec![40, 50, 60, 70]))
1429            .into_batch_rows(RowCount::from(4));
1430        let batch_three = gen_batch()
1431            .col(
1432                VALUE_COLUMN_NAME,
1433                array::cycle::<Int32Type>(vec![15, 15, 15, 15]),
1434            )
1435            .col(ROW_ID, array::cycle::<UInt64Type>(vec![400, 500, 600, 700]))
1436            .into_batch_rows(RowCount::from(4));
1437        let batch_four = gen_batch()
1438            .col(
1439                VALUE_COLUMN_NAME,
1440                array::cycle::<Int32Type>(vec![15, 16, 20, 20]),
1441            )
1442            .col(
1443                ROW_ID,
1444                array::cycle::<UInt64Type>(vec![4000, 5000, 6000, 7000]),
1445            )
1446            .into_batch_rows(RowCount::from(4));
1447        let batches = vec![batch_one, batch_two, batch_three, batch_four];
1448        let schema = Arc::new(Schema::new(vec![
1449            Field::new(VALUE_COLUMN_NAME, DataType::Int32, false),
1450            Field::new(ROW_ID, DataType::UInt64, false),
1451        ]));
1452        let data = RecordBatchIterator::new(batches, schema);
1453        train_bitmap(&index_store, data).await;
1454        let index = BitmapIndex::load(index_store, None, &LanceCache::no_cache())
1455            .await
1456            .unwrap();
1457
1458        // The above should create four pages
1459        //
1460        // 0 - 5
1461        // 10 - 15
1462        // 15 - 15
1463        // 15 - 20
1464        //
1465        // This will help us test various indexing corner cases
1466
1467        // No results (off the left side)
1468        check_bitmap(
1469            &index,
1470            SargableQuery::Equals(ScalarValue::Int32(Some(-3))),
1471            &[],
1472        )
1473        .await;
1474
1475        check_bitmap(
1476            &index,
1477            SargableQuery::Range(
1478                Bound::Unbounded,
1479                Bound::Included(ScalarValue::Int32(Some(-3))),
1480            ),
1481            &[],
1482        )
1483        .await;
1484
1485        check_bitmap(
1486            &index,
1487            SargableQuery::Range(
1488                Bound::Included(ScalarValue::Int32(Some(-10))),
1489                Bound::Included(ScalarValue::Int32(Some(-3))),
1490            ),
1491            &[],
1492        )
1493        .await;
1494
1495        // Hitting the middle of a bucket
1496        check_bitmap(
1497            &index,
1498            SargableQuery::Equals(ScalarValue::Int32(Some(4))),
1499            &[2],
1500        )
1501        .await;
1502
1503        // Hitting a gap between two buckets
1504        check_bitmap(
1505            &index,
1506            SargableQuery::Equals(ScalarValue::Int32(Some(7))),
1507            &[],
1508        )
1509        .await;
1510
1511        // Hitting the lowest of the overlapping buckets
1512        check_bitmap(
1513            &index,
1514            SargableQuery::Equals(ScalarValue::Int32(Some(11))),
1515            &[50, 60],
1516        )
1517        .await;
1518
1519        // Hitting the 15 shared on all three buckets
1520        check_bitmap(
1521            &index,
1522            SargableQuery::Equals(ScalarValue::Int32(Some(15))),
1523            &[70, 400, 500, 600, 700, 4000],
1524        )
1525        .await;
1526
1527        // Hitting the upper part of the three overlapping buckets
1528        check_bitmap(
1529            &index,
1530            SargableQuery::Equals(ScalarValue::Int32(Some(20))),
1531            &[6000, 7000],
1532        )
1533        .await;
1534
1535        // Ranges that capture multiple buckets
1536        check_bitmap(
1537            &index,
1538            SargableQuery::Range(
1539                Bound::Unbounded,
1540                Bound::Included(ScalarValue::Int32(Some(11))),
1541            ),
1542            &[0, 1, 2, 3, 40, 50, 60],
1543        )
1544        .await;
1545
1546        check_bitmap(
1547            &index,
1548            SargableQuery::Range(
1549                Bound::Unbounded,
1550                Bound::Excluded(ScalarValue::Int32(Some(11))),
1551            ),
1552            &[0, 1, 2, 3, 40],
1553        )
1554        .await;
1555
1556        check_bitmap(
1557            &index,
1558            SargableQuery::Range(
1559                Bound::Included(ScalarValue::Int32(Some(4))),
1560                Bound::Unbounded,
1561            ),
1562            &[
1563                2, 3, 40, 50, 60, 70, 400, 500, 600, 700, 4000, 5000, 6000, 7000,
1564            ],
1565        )
1566        .await;
1567
1568        check_bitmap(
1569            &index,
1570            SargableQuery::Range(
1571                Bound::Included(ScalarValue::Int32(Some(4))),
1572                Bound::Included(ScalarValue::Int32(Some(11))),
1573            ),
1574            &[2, 3, 40, 50, 60],
1575        )
1576        .await;
1577
1578        check_bitmap(
1579            &index,
1580            SargableQuery::Range(
1581                Bound::Included(ScalarValue::Int32(Some(4))),
1582                Bound::Excluded(ScalarValue::Int32(Some(11))),
1583            ),
1584            &[2, 3, 40],
1585        )
1586        .await;
1587
1588        check_bitmap(
1589            &index,
1590            SargableQuery::Range(
1591                Bound::Excluded(ScalarValue::Int32(Some(4))),
1592                Bound::Unbounded,
1593            ),
1594            &[
1595                3, 40, 50, 60, 70, 400, 500, 600, 700, 4000, 5000, 6000, 7000,
1596            ],
1597        )
1598        .await;
1599
1600        check_bitmap(
1601            &index,
1602            SargableQuery::Range(
1603                Bound::Excluded(ScalarValue::Int32(Some(4))),
1604                Bound::Included(ScalarValue::Int32(Some(11))),
1605            ),
1606            &[3, 40, 50, 60],
1607        )
1608        .await;
1609
1610        check_bitmap(
1611            &index,
1612            SargableQuery::Range(
1613                Bound::Excluded(ScalarValue::Int32(Some(4))),
1614                Bound::Excluded(ScalarValue::Int32(Some(11))),
1615            ),
1616            &[3, 40],
1617        )
1618        .await;
1619
1620        check_bitmap(
1621            &index,
1622            SargableQuery::Range(
1623                Bound::Excluded(ScalarValue::Int32(Some(-50))),
1624                Bound::Excluded(ScalarValue::Int32(Some(1000))),
1625            ),
1626            &[
1627                0, 1, 2, 3, 40, 50, 60, 70, 400, 500, 600, 700, 4000, 5000, 6000, 7000,
1628            ],
1629        )
1630        .await;
1631    }
1632
1633    #[tokio::test]
1634    async fn test_bitmap_update() {
1635        let index_dir = TempDir::default();
1636        let index_store = test_store(&index_dir);
1637        let data = gen_batch()
1638            .col(VALUE_COLUMN_NAME, array::step::<Int32Type>())
1639            .col(ROW_ID, array::step::<UInt64Type>())
1640            .into_reader_rows(RowCount::from(4096), BatchCount::from(1));
1641        train_bitmap(&index_store, data).await;
1642        let index = BitmapIndex::load(index_store, None, &LanceCache::no_cache())
1643            .await
1644            .unwrap();
1645
1646        let data = gen_batch()
1647            .col(VALUE_COLUMN_NAME, array::step_custom::<Int32Type>(4096, 1))
1648            .col(ROW_ID, array::step_custom::<UInt64Type>(4096, 1))
1649            .into_reader_rows(RowCount::from(4096), BatchCount::from(1));
1650
1651        let updated_index_dir = TempDir::default();
1652        let updated_index_store = test_store(&updated_index_dir);
1653        index
1654            .update(
1655                lance_datafusion::utils::reader_to_stream(Box::new(data)),
1656                updated_index_store.as_ref(),
1657                None,
1658            )
1659            .await
1660            .unwrap();
1661        let updated_index = BitmapIndex::load(updated_index_store, None, &LanceCache::no_cache())
1662            .await
1663            .unwrap();
1664
1665        let result = updated_index
1666            .search(
1667                &SargableQuery::Equals(ScalarValue::Int32(Some(5000))),
1668                &NoOpMetricsCollector,
1669            )
1670            .await
1671            .unwrap();
1672
1673        assert!(result.is_exact());
1674        let row_addrs = result.row_addrs().true_rows();
1675        assert_eq!(Some(1), row_addrs.len());
1676        assert!(row_addrs.contains(5000));
1677    }
1678
1679    #[tokio::test]
1680    async fn test_bitmap_remap() {
1681        let index_dir = TempDir::default();
1682        let index_store = test_store(&index_dir);
1683        let data = gen_batch()
1684            .col(VALUE_COLUMN_NAME, array::step::<Int32Type>())
1685            .col(ROW_ID, array::step::<UInt64Type>())
1686            .into_reader_rows(RowCount::from(50), BatchCount::from(1));
1687        train_bitmap(&index_store, data).await;
1688        let index = BitmapIndex::load(index_store, None, &LanceCache::no_cache())
1689            .await
1690            .unwrap();
1691
1692        let mapping = (0..50)
1693            .map(|i| {
1694                let map_result = if i == 5 {
1695                    Some(65)
1696                } else if i == 7 {
1697                    None
1698                } else {
1699                    Some(i)
1700                };
1701                (i, map_result)
1702            })
1703            .collect::<HashMap<_, _>>();
1704
1705        let remapped_dir = TempDir::default();
1706        let remapped_store = test_store(&remapped_dir);
1707        index
1708            .remap(&RowAddrRemap::direct(mapping), remapped_store.as_ref())
1709            .await
1710            .unwrap();
1711        let remapped_index = BitmapIndex::load(remapped_store, None, &LanceCache::no_cache())
1712            .await
1713            .unwrap();
1714
1715        // Remapped to new value
1716        assert!(
1717            remapped_index
1718                .search(
1719                    &SargableQuery::Equals(ScalarValue::Int32(Some(5))),
1720                    &NoOpMetricsCollector
1721                )
1722                .await
1723                .unwrap()
1724                .row_addrs()
1725                .selected(65)
1726        );
1727        // Deleted
1728        assert!(
1729            remapped_index
1730                .search(
1731                    &SargableQuery::Equals(ScalarValue::Int32(Some(7))),
1732                    &NoOpMetricsCollector
1733                )
1734                .await
1735                .unwrap()
1736                .row_addrs()
1737                .is_empty()
1738        );
1739        // Not remapped
1740        assert!(
1741            remapped_index
1742                .search(
1743                    &SargableQuery::Equals(ScalarValue::Int32(Some(3))),
1744                    &NoOpMetricsCollector
1745                )
1746                .await
1747                .unwrap()
1748                .row_addrs()
1749                .selected(3)
1750        );
1751    }
1752
1753    async fn train_tag(
1754        index_store: &Arc<dyn IndexStore>,
1755        data: impl RecordBatchReader + Send + Sync + 'static,
1756    ) {
1757        let data = lance_datafusion::utils::reader_to_stream(Box::new(data));
1758        let request = LabelListIndexPlugin
1759            .new_training_request(
1760                "{}",
1761                &Field::new(
1762                    VALUE_COLUMN_NAME,
1763                    DataType::List(Arc::new(Field::new("item", DataType::UInt8, false))),
1764                    false,
1765                ),
1766            )
1767            .unwrap();
1768        LabelListIndexPlugin
1769            .train_index(
1770                data,
1771                index_store.as_ref(),
1772                request,
1773                None,
1774                crate::progress::noop_progress(),
1775            )
1776            .await
1777            .unwrap();
1778    }
1779
1780    #[tokio::test]
1781    async fn test_label_list_index() {
1782        let tempdir = TempDir::default();
1783        let index_store = test_store(&tempdir);
1784        let data = gen_batch()
1785            .col(
1786                VALUE_COLUMN_NAME,
1787                array::rand_type(&DataType::List(Arc::new(Field::new(
1788                    "item",
1789                    DataType::UInt8,
1790                    false,
1791                )))),
1792            )
1793            .col(ROW_ID, array::step::<UInt64Type>())
1794            .into_batch_rows(RowCount::from(40960))
1795            .unwrap();
1796
1797        let batch_reader = RecordBatchIterator::new(vec![Ok(data.clone())], data.schema());
1798
1799        // This is probably enough data that we can be assured each tag is used at least once
1800        train_tag(&index_store, batch_reader).await;
1801
1802        // We scan through each list, if it was a match we run match_fn to check
1803        // if the match was correct if it was not a match we run no_match_fn to check
1804        // if the no-match was correct
1805        type MatchFn = Box<dyn Fn(&ScalarBuffer<u8>) -> bool>;
1806        let check = |query: LabelListQuery, match_fn: MatchFn, no_match_fn: MatchFn| {
1807            let index_store = index_store.clone();
1808            let data = data.clone();
1809            async move {
1810                let index = LabelListIndexPlugin
1811                    .load_index(
1812                        index_store,
1813                        &default_details::<pbold::LabelListIndexDetails>(),
1814                        None,
1815                        &LanceCache::no_cache(),
1816                    )
1817                    .await
1818                    .unwrap();
1819                let result = index.search(&query, &NoOpMetricsCollector).await.unwrap();
1820                assert!(result.is_exact());
1821                let row_addrs = result.row_addrs().true_rows();
1822
1823                let row_addrs_set = row_addrs
1824                    .row_addrs()
1825                    .unwrap()
1826                    .map(u64::from)
1827                    .collect::<std::collections::HashSet<_>>();
1828
1829                for (list, row_id) in data
1830                    .column(0)
1831                    .as_list::<i32>()
1832                    .iter()
1833                    .zip(data.column(1).as_primitive::<UInt64Type>())
1834                {
1835                    let list = list.unwrap();
1836                    let row_id = row_id.unwrap();
1837                    let vals = list.as_primitive::<UInt8Type>().values();
1838                    if row_addrs_set.contains(&row_id) {
1839                        assert!(match_fn(vals));
1840                    } else {
1841                        assert!(no_match_fn(vals));
1842                    }
1843                }
1844            }
1845        };
1846
1847        // Simple check for 1 value (doesn't matter intersection vs union)
1848        check(
1849            LabelListQuery::HasAnyLabel(vec![ScalarValue::UInt8(Some(1))]),
1850            Box::new(|vals| vals.contains(&1)),
1851            Box::new(|vals| !vals.contains(&1)),
1852        )
1853        .await;
1854        check(
1855            LabelListQuery::HasAllLabels(vec![ScalarValue::UInt8(Some(1))]),
1856            Box::new(|vals| vals.contains(&1)),
1857            Box::new(|vals| !vals.contains(&1)),
1858        )
1859        .await;
1860        // Set intersection
1861        check(
1862            LabelListQuery::HasAllLabels(vec![
1863                ScalarValue::UInt8(Some(1)),
1864                ScalarValue::UInt8(Some(2)),
1865            ]),
1866            // Match must have 1 and 2
1867            Box::new(|vals| vals.contains(&1) && vals.contains(&2)),
1868            // No-match must either not have 1 or not have 2
1869            Box::new(|vals| !vals.contains(&1) || !vals.contains(&2)),
1870        )
1871        .await;
1872        // Set union
1873        check(
1874            LabelListQuery::HasAnyLabel(vec![
1875                ScalarValue::UInt8(Some(1)),
1876                ScalarValue::UInt8(Some(2)),
1877            ]),
1878            // Match either have 1 or have 2
1879            Box::new(|vals| vals.contains(&1) || vals.contains(&2)),
1880            // No-match must not have 1 and not have 2
1881            Box::new(|vals| !vals.contains(&1) && !vals.contains(&2)),
1882        )
1883        .await;
1884    }
1885
1886    #[tokio::test]
1887    async fn test_label_list_null_handling() {
1888        let tempdir = TempDir::default();
1889        let index_store = test_store(&tempdir);
1890
1891        // Create test data with null items within lists:
1892        // Row 0: [1, 2] - no nulls
1893        // Row 1: [3, null] - has a null item
1894        // Row 2: [4] - no nulls
1895        let list_array = ListArray::from_iter_primitive::<UInt8Type, _, _>(vec![
1896            Some(vec![Some(1), Some(2)]),
1897            Some(vec![Some(3), None]),
1898            Some(vec![Some(4)]),
1899        ]);
1900        let row_ids = UInt64Array::from_iter_values(0..3);
1901        // Create schema with nullable list items to match the ListArray
1902        let schema = Arc::new(Schema::new(vec![
1903            Field::new(
1904                VALUE_COLUMN_NAME,
1905                DataType::List(Arc::new(Field::new("item", DataType::UInt8, true))),
1906                true,
1907            ),
1908            Field::new(ROW_ID, DataType::UInt64, false),
1909        ]));
1910        let batch = RecordBatch::try_new(
1911            schema.clone(),
1912            vec![Arc::new(list_array), Arc::new(row_ids)],
1913        )
1914        .unwrap();
1915
1916        let batch_reader = RecordBatchIterator::new(vec![Ok(batch)], schema);
1917        train_tag(&index_store, batch_reader).await;
1918
1919        let index = LabelListIndexPlugin
1920            .load_index(
1921                index_store,
1922                &default_details::<pbold::LabelListIndexDetails>(),
1923                None,
1924                &LanceCache::no_cache(),
1925            )
1926            .await
1927            .unwrap();
1928
1929        // Test: Search for lists containing value 1
1930        // Row 0: [1, 2] - contains 1 → TRUE
1931        // Row 1: [3, null] - null elements are ignored → FALSE
1932        // Row 2: [4] - doesn't contain 1 → FALSE
1933        let query = LabelListQuery::HasAnyLabel(vec![ScalarValue::UInt8(Some(1))]);
1934        let result = index.search(&query, &NoOpMetricsCollector).await.unwrap();
1935
1936        match result {
1937            SearchResult::Exact(row_ids) => {
1938                let actual_rows: Vec<u64> = row_ids
1939                    .true_rows()
1940                    .row_addrs()
1941                    .unwrap()
1942                    .map(u64::from)
1943                    .collect();
1944                assert_eq!(
1945                    actual_rows,
1946                    vec![0],
1947                    "Should find row 0 where list contains 1"
1948                );
1949
1950                assert!(
1951                    row_ids.null_rows().is_empty(),
1952                    "null_row_ids should be empty when null elements are ignored"
1953                );
1954            }
1955            _ => panic!("Expected Exact search result"),
1956        }
1957    }
1958
1959    #[tokio::test]
1960    async fn test_label_list_bitmap_only_layout_is_compatible() {
1961        let tempdir = TempDir::default();
1962        let index_store = test_store(&tempdir);
1963
1964        // Simulate an older released layout that only had the bitmap lookup file.
1965        let values = arrow_array::UInt8Array::from(vec![1, 2]);
1966        let row_ids = UInt64Array::from(vec![0, 2]);
1967        let schema = Arc::new(Schema::new(vec![
1968            Field::new(VALUE_COLUMN_NAME, DataType::UInt8, true),
1969            Field::new(ROW_ID, DataType::UInt64, false),
1970        ]));
1971        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(values), Arc::new(row_ids)])
1972            .unwrap();
1973
1974        BitmapIndexPlugin::train_bitmap_index(
1975            lance_datafusion::utils::reader_to_stream(Box::new(RecordBatchIterator::new(
1976                vec![Ok(batch)],
1977                schema,
1978            ))),
1979            index_store.as_ref(),
1980        )
1981        .await
1982        .unwrap();
1983
1984        let index = LabelListIndexPlugin
1985            .load_index(
1986                index_store,
1987                &default_details::<pbold::LabelListIndexDetails>(),
1988                None,
1989                &LanceCache::no_cache(),
1990            )
1991            .await
1992            .unwrap();
1993
1994        let query = LabelListQuery::HasAnyLabel(vec![ScalarValue::UInt8(Some(1))]);
1995        let result = index.search(&query, &NoOpMetricsCollector).await.unwrap();
1996
1997        match result {
1998            SearchResult::Exact(row_ids) => {
1999                assert!(row_ids.null_rows().is_empty());
2000                let actual_rows: Vec<u64> = row_ids
2001                    .true_rows()
2002                    .row_addrs()
2003                    .unwrap()
2004                    .map(u64::from)
2005                    .collect();
2006                assert_eq!(actual_rows, vec![0]);
2007            }
2008            _ => panic!("Expected Exact search result"),
2009        }
2010    }
2011}
lance_index/scalar/lance_format.rs

lance_index/scalar/
lance_format.rs
