diff --git a/.bumpversion.toml b/.bumpversion.toml index 32f3260d6c1..43b5bc02da0 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -80,6 +80,11 @@ filename = "Cargo.toml" search = 'lance-index = {{ version = "={current_version}"' replace = 'lance-index = {{ version = "={new_version}"' +[[tool.bumpversion.files]] +filename = "Cargo.toml" +search = 'lance-index-core = {{ version = "={current_version}"' +replace = 'lance-index-core = {{ version = "={new_version}"' + [[tool.bumpversion.files]] filename = "Cargo.toml" search = 'lance-io = {{ version = "={current_version}"' diff --git a/Cargo.lock b/Cargo.lock index 13df305c896..9c8ac672bc9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4832,6 +4832,7 @@ dependencies = [ "lance-encoding", "lance-file", "lance-geo", + "lance-index-core", "lance-io", "lance-linalg", "lance-select", @@ -4864,6 +4865,28 @@ dependencies = [ "uuid", ] +[[package]] +name = "lance-index-core" +version = "9.0.0-beta.24" +dependencies = [ + "arrow-array", + "arrow-schema", + "arrow-select", + "async-trait", + "bytes", + "datafusion", + "datafusion-common", + "datafusion-expr", + "futures", + "lance-core", + "lance-io", + "lance-select", + "prost-types", + "roaring", + "serde", + "serde_json", +] + [[package]] name = "lance-io" version = "9.0.0-beta.24" diff --git a/Cargo.toml b/Cargo.toml index 047813738d1..9f2ab2163a5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ members = [ "rust/lance-file", "rust/lance-geo", "rust/lance-index", + "rust/lance-index-core", "rust/lance-io", "rust/lance-linalg", "rust/lance-namespace", @@ -67,6 +68,7 @@ lance-encoding = { version = "=9.0.0-beta.24", path = "./rust/lance-encoding" } lance-file = { version = "=9.0.0-beta.24", path = "./rust/lance-file" } lance-geo = { version = "=9.0.0-beta.24", path = "./rust/lance-geo" } lance-index = { version = "=9.0.0-beta.24", path = "./rust/lance-index" } +lance-index-core = { version = "=9.0.0-beta.24", path = "./rust/lance-index-core" } lance-io = { version = "=9.0.0-beta.24", path = "./rust/lance-io", default-features = false } lance-linalg = { version = "=9.0.0-beta.24", path = "./rust/lance-linalg" } lance-namespace = { version = "=9.0.0-beta.24", path = "./rust/lance-namespace" } diff --git a/java/lance-jni/Cargo.lock b/java/lance-jni/Cargo.lock index a3330674937..5fb4d1cf750 100644 --- a/java/lance-jni/Cargo.lock +++ b/java/lance-jni/Cargo.lock @@ -4001,6 +4001,7 @@ dependencies = [ "lance-encoding", "lance-file", "lance-geo", + "lance-index-core", "lance-io", "lance-linalg", "lance-select", @@ -4029,6 +4030,28 @@ dependencies = [ "uuid", ] +[[package]] +name = "lance-index-core" +version = "9.0.0-beta.24" +dependencies = [ + "arrow-array", + "arrow-schema", + "arrow-select", + "async-trait", + "bytes", + "datafusion", + "datafusion-common", + "datafusion-expr", + "futures", + "lance-core", + "lance-io", + "lance-select", + "prost-types", + "roaring", + "serde", + "serde_json", +] + [[package]] name = "lance-io" version = "9.0.0-beta.24" diff --git a/python/Cargo.lock b/python/Cargo.lock index 6f771efd251..cabbe2be7c3 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -4411,6 +4411,7 @@ dependencies = [ "lance-encoding", "lance-file", "lance-geo", + "lance-index-core", "lance-io", "lance-linalg", "lance-select", @@ -4439,6 +4440,28 @@ dependencies = [ "uuid", ] +[[package]] +name = "lance-index-core" +version = "9.0.0-beta.24" +dependencies = [ + "arrow-array", + "arrow-schema", + "arrow-select", + "async-trait", + "bytes", + "datafusion", + "datafusion-common", + "datafusion-expr", + "futures", + "lance-core", + "lance-io", + "lance-select", + "prost-types", + "roaring", + "serde", + "serde_json", +] + [[package]] name = "lance-io" version = "9.0.0-beta.24" diff --git a/python/src/utils.rs b/python/src/utils.rs index 4f7d6d7dde2..edc376d35c8 100644 --- a/python/src/utils.rs +++ b/python/src/utils.rs @@ -26,7 +26,6 @@ use lance::Result; use lance::datatypes::Schema; use lance_arrow::FixedSizeListArrayExt; use lance_file::previous::writer::FileWriter as PreviousFileWriter; -use lance_index::scalar::IndexWriter; use lance_index::vector::hnsw::{HNSW, builder::HnswBuildParams}; use lance_index::vector::kmeans::{ KMeans as LanceKMeans, KMeansAlgoFloat, KMeansParams, compute_partitions, @@ -255,7 +254,7 @@ impl Hnsw { rt().block_on(Some(py), async { let batch = self.hnsw.to_batch()?; let metadata = batch.schema_ref().metadata().clone(); - writer.write_record_batch(batch).await?; + writer.write(&[batch]).await?; writer.finish_with_metadata(&metadata).await?; Result::Ok(()) })? diff --git a/rust/lance-index-core/Cargo.toml b/rust/lance-index-core/Cargo.toml new file mode 100644 index 00000000000..5024f2eee0b --- /dev/null +++ b/rust/lance-index-core/Cargo.toml @@ -0,0 +1,33 @@ +[package] +name = "lance-index-core" +version.workspace = true +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +readme = "README.md" +description = "Core traits and types for Lance index plugins" +keywords.workspace = true +categories.workspace = true +rust-version.workspace = true + +[dependencies] +async-trait.workspace = true +arrow-array.workspace = true +arrow-schema.workspace = true +arrow-select.workspace = true +bytes.workspace = true +datafusion-common.workspace = true +datafusion-expr.workspace = true +datafusion.workspace = true +futures.workspace = true +lance-core.workspace = true +lance-io.workspace = true +lance-select.workspace = true +prost-types.workspace = true +roaring.workspace = true +serde.workspace = true +serde_json.workspace = true + +[lints] +workspace = true diff --git a/rust/lance-index-core/src/lib.rs b/rust/lance-index-core/src/lib.rs new file mode 100644 index 00000000000..393064014e4 --- /dev/null +++ b/rust/lance-index-core/src/lib.rs @@ -0,0 +1,288 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{any::Any, sync::Arc}; + +use async_trait::async_trait; +use lance_core::deepsize::DeepSizeOf; +use lance_core::{Error, Result}; +use roaring::RoaringBitmap; +use serde::{Deserialize, Serialize}; +use std::convert::TryFrom; + +pub mod metrics; +pub mod scalar; + +/// Generic methods common across all types of secondary indices +/// +#[async_trait] +pub trait Index: Send + Sync + DeepSizeOf { + /// Cast to [Any]. + fn as_any(&self) -> &dyn Any; + + /// Cast to [Index] + fn as_index(self: Arc) -> Arc; + + /// Retrieve index statistics as a JSON Value + fn statistics(&self) -> Result; + + /// Prewarm the index. + /// + /// This will load the index into memory and cache it. + async fn prewarm(&self) -> Result<()>; + + /// Get the type of the index + fn index_type(&self) -> IndexType; + + /// Read through the index and determine which fragment ids are covered by the index + /// + /// This is a kind of slow operation. It's better to use the fragment_bitmap. This + /// only exists for cases where the fragment_bitmap has become corrupted or missing. + async fn calculate_included_frags(&self) -> Result; +} + +/// Index Type +#[derive(Debug, PartialEq, Eq, Copy, Hash, Clone, DeepSizeOf, Serialize, Deserialize)] +pub enum IndexType { + // Preserve 0-100 for simple indices. + Scalar = 0, // Legacy scalar index, alias to BTree + + BTree = 1, // BTree + + Bitmap = 2, // Bitmap + + LabelList = 3, // LabelList + + Inverted = 4, // Inverted + + NGram = 5, // NGram + + FragmentReuse = 6, + + MemWal = 7, + + ZoneMap = 8, // ZoneMap + + BloomFilter = 9, // Bloom filter + + RTree = 10, // RTree + + Fm = 11, // FM-Index + + // 100+ and up for vector index. + /// Flat vector index. + Vector = 100, // Legacy vector index, alias to IvfPq + IvfFlat = 101, + IvfSq = 102, + IvfPq = 103, + IvfHnswSq = 104, + IvfHnswPq = 105, + IvfHnswFlat = 106, + IvfRq = 107, +} + +impl std::fmt::Display for IndexType { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match self { + Self::Scalar | Self::BTree => write!(f, "BTree"), + Self::Bitmap => write!(f, "Bitmap"), + Self::LabelList => write!(f, "LabelList"), + Self::Inverted => write!(f, "Inverted"), + Self::NGram => write!(f, "NGram"), + Self::FragmentReuse => write!(f, "FragmentReuse"), + Self::MemWal => write!(f, "MemWal"), + Self::ZoneMap => write!(f, "ZoneMap"), + Self::BloomFilter => write!(f, "BloomFilter"), + Self::RTree => write!(f, "RTree"), + Self::Fm => write!(f, "Fm"), + Self::Vector | Self::IvfPq => write!(f, "IVF_PQ"), + Self::IvfFlat => write!(f, "IVF_FLAT"), + Self::IvfSq => write!(f, "IVF_SQ"), + Self::IvfHnswSq => write!(f, "IVF_HNSW_SQ"), + Self::IvfHnswPq => write!(f, "IVF_HNSW_PQ"), + Self::IvfHnswFlat => write!(f, "IVF_HNSW_FLAT"), + Self::IvfRq => write!(f, "IVF_RQ"), + } + } +} + +impl TryFrom for IndexType { + type Error = Error; + + fn try_from(value: i32) -> Result { + match value { + v if v == Self::Scalar as i32 => Ok(Self::Scalar), + v if v == Self::BTree as i32 => Ok(Self::BTree), + v if v == Self::Bitmap as i32 => Ok(Self::Bitmap), + v if v == Self::LabelList as i32 => Ok(Self::LabelList), + v if v == Self::NGram as i32 => Ok(Self::NGram), + v if v == Self::Inverted as i32 => Ok(Self::Inverted), + v if v == Self::FragmentReuse as i32 => Ok(Self::FragmentReuse), + v if v == Self::MemWal as i32 => Ok(Self::MemWal), + v if v == Self::ZoneMap as i32 => Ok(Self::ZoneMap), + v if v == Self::BloomFilter as i32 => Ok(Self::BloomFilter), + v if v == Self::RTree as i32 => Ok(Self::RTree), + v if v == Self::Fm as i32 => Ok(Self::Fm), + v if v == Self::Vector as i32 => Ok(Self::Vector), + v if v == Self::IvfFlat as i32 => Ok(Self::IvfFlat), + v if v == Self::IvfSq as i32 => Ok(Self::IvfSq), + v if v == Self::IvfPq as i32 => Ok(Self::IvfPq), + v if v == Self::IvfHnswSq as i32 => Ok(Self::IvfHnswSq), + v if v == Self::IvfHnswPq as i32 => Ok(Self::IvfHnswPq), + v if v == Self::IvfHnswFlat as i32 => Ok(Self::IvfHnswFlat), + v if v == Self::IvfRq as i32 => Ok(Self::IvfRq), + _ => Err(Error::invalid_input_source( + format!("the input value {} is not a valid IndexType", value).into(), + )), + } + } +} + +impl TryFrom<&str> for IndexType { + type Error = Error; + + fn try_from(value: &str) -> Result { + match value { + "BTree" | "BTREE" => Ok(Self::BTree), + "Bitmap" | "BITMAP" => Ok(Self::Bitmap), + "LabelList" | "LABELLIST" => Ok(Self::LabelList), + "Inverted" | "INVERTED" => Ok(Self::Inverted), + "NGram" | "NGRAM" => Ok(Self::NGram), + "ZoneMap" | "ZONEMAP" => Ok(Self::ZoneMap), + "BloomFilter" | "BLOOMFILTER" | "BLOOM_FILTER" => Ok(Self::BloomFilter), + "RTree" | "RTREE" | "R_TREE" => Ok(Self::RTree), + "Fm" | "FM" => Ok(Self::Fm), + "Vector" | "VECTOR" => Ok(Self::Vector), + "IVF_FLAT" => Ok(Self::IvfFlat), + "IVF_SQ" => Ok(Self::IvfSq), + "IVF_PQ" => Ok(Self::IvfPq), + "IVF_RQ" => Ok(Self::IvfRq), + "IVF_HNSW_FLAT" => Ok(Self::IvfHnswFlat), + "IVF_HNSW_SQ" => Ok(Self::IvfHnswSq), + "IVF_HNSW_PQ" => Ok(Self::IvfHnswPq), + "FragmentReuse" => Ok(Self::FragmentReuse), + "MemWal" => Ok(Self::MemWal), + _ => Err(Error::invalid_input(format!( + "invalid index type: {}", + value + ))), + } + } +} + +impl IndexType { + pub fn is_scalar(&self) -> bool { + matches!( + self, + Self::Scalar + | Self::BTree + | Self::Bitmap + | Self::LabelList + | Self::Inverted + | Self::NGram + | Self::ZoneMap + | Self::BloomFilter + | Self::RTree + | Self::Fm, + ) + } + + pub fn is_vector(&self) -> bool { + matches!( + self, + Self::Vector + | Self::IvfPq + | Self::IvfHnswSq + | Self::IvfHnswPq + | Self::IvfHnswFlat + | Self::IvfFlat + | Self::IvfSq + | Self::IvfRq + ) + } + + pub fn is_system(&self) -> bool { + matches!(self, Self::FragmentReuse | Self::MemWal) + } + + /// Returns the current format version of the index type, + /// bump this when the index format changes. + /// Indices which higher version than these will be ignored for compatibility, + /// This would happen when creating index in a newer version of Lance, + /// but then opening the index in older version of Lance + pub fn version(&self) -> i32 { + match self { + Self::Scalar => 0, + Self::BTree => 0, + Self::Bitmap => 0, + Self::LabelList => 0, + Self::Inverted => 0, + Self::NGram => 0, + Self::FragmentReuse => 0, + Self::MemWal => 0, + Self::ZoneMap => 0, + Self::BloomFilter => 0, + Self::RTree => 0, + Self::Fm => 0, + + // IMPORTANT: if any vector index subtype needs a format bump that is + // not backward compatible, its new version must be set to + // (current max vector index version + 1), even if only one subtype + // changed. Compatibility filtering currently cannot distinguish vector + // subtypes from details-only metadata, so vector versions effectively + // share one global monotonic compatibility level. + Self::Vector + | Self::IvfFlat + | Self::IvfSq + | Self::IvfPq + | Self::IvfHnswSq + | Self::IvfHnswPq + | Self::IvfHnswFlat => 1, + Self::IvfRq => 2, + } + } + + /// Returns the target partition size for the index type. + /// + /// This is used to compute the number of partitions for the index. + /// The partition size is optimized for the best performance of the index. + /// + /// This is for vector indices only. + pub fn target_partition_size(&self) -> usize { + match self { + Self::Vector => 8192, + Self::IvfFlat => 4096, + Self::IvfSq => 8192, + Self::IvfPq => 8192, + Self::IvfRq => 4096, + Self::IvfHnswFlat => 1 << 20, + Self::IvfHnswSq => 1 << 20, + Self::IvfHnswPq => 1 << 20, + _ => 8192, + } + } + + /// Returns the highest supported vector index version in this Lance build. + pub fn max_vector_version() -> u32 { + [ + Self::Vector, + Self::IvfFlat, + Self::IvfSq, + Self::IvfPq, + Self::IvfHnswSq, + Self::IvfHnswPq, + Self::IvfHnswFlat, + Self::IvfRq, + ] + .into_iter() + .map(|index_type| index_type.version() as u32) + .max() + .unwrap_or(1) + } +} + +pub trait IndexParams: Send + Sync { + fn as_any(&self) -> &dyn Any; + + fn index_name(&self) -> &str; +} diff --git a/rust/lance-index-core/src/metrics.rs b/rust/lance-index-core/src/metrics.rs new file mode 100644 index 00000000000..8c0c119a3c3 --- /dev/null +++ b/rust/lance-index-core/src/metrics.rs @@ -0,0 +1,117 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::sync::atomic::{AtomicUsize, Ordering}; + +pub const AND_CANDIDATES_SEEN_METRIC: &str = "and_candidates_seen"; +pub const AND_CANDIDATES_PRUNED_BEFORE_RETURN_METRIC: &str = "and_candidates_pruned_before_return"; +pub const AND_FULL_SCORES_METRIC: &str = "and_full_scores"; +pub const FREQS_COLLECTED_METRIC: &str = "freqs_collected"; + +/// A trait used by the index to report metrics +/// +/// Callers can implement this trait to collect metrics +pub trait MetricsCollector: Send + Sync { + /// Record partition loads + /// + /// Many indices consist of partitions that may need to be loaded + /// into cache. For example, an inverted index or ngram index has a + /// posting list for each token. + /// + /// In the ideal case, these shards are in the cache and will not need + /// to be loaded from disk. This method should not be called if the + /// shard is in the cache. + fn record_parts_loaded(&self, num_parts: usize); + + /// Record a shard load + fn record_part_load(&self) { + self.record_parts_loaded(1); + } + + /// Record an index load + /// + /// This should be called when a scalar index is loaded from storage. + /// It should not be called if the index is already in memory. + fn record_index_loads(&self, num_indexes: usize); + + /// Record an index load + fn record_index_load(&self) { + self.record_index_loads(1); + } + + /// Record the number of "comparisons" made by the index + /// + /// What exactly constitutes a comparison depends on the index type. + /// For example, a B-tree index may make comparisons while searching for a value. + /// On the other hand, a bitmap index makes comparisons when computing the intersection + /// of two bitmaps. + /// + /// The goal is to provide some visibility into the compute cost of the search + fn record_comparisons(&self, num_comparisons: usize); + + /// Record AND candidates returned from WAND alignment to the scoring loop. + /// + /// This excludes candidates pruned before `next()` returns. Use this with + /// `record_and_candidates_pruned_before_return` to recover total aligned + /// AND candidates. + fn record_and_candidates_seen(&self, _num_candidates: usize) {} + + /// Record AND candidates pruned during WAND alignment before `next()` returns. + fn record_and_candidates_pruned_before_return(&self, _num_candidates: usize) {} + + fn record_and_full_scores(&self, _num_scores: usize) {} + + fn record_freqs_collected(&self, _num_collections: usize) {} + + /// Returns an optional sink for recording exact I/O statistics (bytes read, + /// IOPS, and requests) performed on behalf of this collector. + /// + /// Index implementations that read from a + /// [`lance_io::scheduler::ScanScheduler`] can attach the returned handle to + /// their file readers so the I/O performed for a single query is measured + /// and attributed here. The default returns `None`, meaning the caller does + /// not want I/O measured (and index implementations should then take their + /// normal, uninstrumented read path). + fn io_stats(&self) -> Option { + None + } +} + +/// A no-op metrics collector that does nothing +pub struct NoOpMetricsCollector; + +impl MetricsCollector for NoOpMetricsCollector { + fn record_parts_loaded(&self, _num_parts: usize) {} + fn record_index_loads(&self, _num_indexes: usize) {} + fn record_comparisons(&self, _num_comparisons: usize) {} +} + +#[derive(Default)] +pub struct LocalMetricsCollector { + pub parts_loaded: AtomicUsize, + pub index_loads: AtomicUsize, + pub comparisons: AtomicUsize, +} + +impl LocalMetricsCollector { + pub fn dump_into(self, other: &dyn MetricsCollector) { + other.record_parts_loaded(self.parts_loaded.load(Ordering::Relaxed)); + other.record_index_loads(self.index_loads.load(Ordering::Relaxed)); + other.record_comparisons(self.comparisons.load(Ordering::Relaxed)); + } +} + +impl MetricsCollector for LocalMetricsCollector { + fn record_parts_loaded(&self, num_parts: usize) { + self.parts_loaded.fetch_add(num_parts, Ordering::Relaxed); + } + + fn record_index_loads(&self, num_indexes: usize) { + self.index_loads.fetch_add(num_indexes, Ordering::Relaxed); + } + + fn record_comparisons(&self, num_comparisons: usize) { + self.comparisons + .fetch_add(num_comparisons, Ordering::Relaxed); + } +} diff --git a/rust/lance-index-core/src/scalar.rs b/rust/lance-index-core/src/scalar.rs new file mode 100644 index 00000000000..e209a32cdb8 --- /dev/null +++ b/rust/lance-index-core/src/scalar.rs @@ -0,0 +1,585 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Abstract scalar index traits and types for Lance index plugins + +use arrow_array::{BooleanArray, RecordBatch, UInt64Array}; +use arrow_schema::Schema; +use async_trait::async_trait; +use bytes::Bytes; +use datafusion::physical_plan::SendableRecordBatchStream; +use datafusion_common::scalar::ScalarValue; +use datafusion_expr::Expr; +use lance_core::deepsize::DeepSizeOf; +use lance_core::utils::row_addr_remap::RowAddrRemap; +use lance_core::{Error, Result}; +use lance_io::stream::{RecordBatchStream, RecordBatchStreamAdapter}; +use lance_select::{NullableRowAddrSet, RowAddrTreeMap, RowSetOps}; +use roaring::{RoaringBitmap, RoaringTreemap}; +use serde::Serialize; +use std::collections::HashMap; +use std::fmt::Debug; +use std::pin::Pin; +use std::{any::Any, sync::Arc}; + +use crate::metrics::MetricsCollector; +use crate::{Index, IndexParams, IndexType}; + +/// Metadata about a single file within an index segment. +#[derive(Debug, Clone, PartialEq, DeepSizeOf)] +pub struct IndexFile { + /// Path relative to the index directory (e.g., "index.idx", "auxiliary.idx") + pub path: String, + /// Size of the file in bytes + pub size_bytes: u64, +} + +pub const LANCE_SCALAR_INDEX: &str = "__lance_scalar_index"; + +/// Builtin index types supported by the Lance library +/// +/// This is primarily for convenience to avoid a bunch of string +/// constants and provide some auto-complete. This type should not +/// be used in the manifest as plugins cannot add new entries. +#[derive(Debug, Clone, PartialEq, Eq, DeepSizeOf)] +pub enum BuiltinIndexType { + BTree, + Bitmap, + LabelList, + NGram, + ZoneMap, + BloomFilter, + RTree, + Inverted, + Fm, +} + +impl BuiltinIndexType { + pub fn as_str(&self) -> &str { + match self { + Self::BTree => "btree", + Self::Bitmap => "bitmap", + Self::LabelList => "labellist", + Self::NGram => "ngram", + Self::ZoneMap => "zonemap", + Self::Inverted => "inverted", + Self::BloomFilter => "bloomfilter", + Self::RTree => "rtree", + Self::Fm => "fm", + } + } +} + +impl TryFrom for BuiltinIndexType { + type Error = Error; + + fn try_from(value: IndexType) -> Result { + match value { + IndexType::BTree => Ok(Self::BTree), + IndexType::Bitmap => Ok(Self::Bitmap), + IndexType::LabelList => Ok(Self::LabelList), + IndexType::NGram => Ok(Self::NGram), + IndexType::ZoneMap => Ok(Self::ZoneMap), + IndexType::Inverted => Ok(Self::Inverted), + IndexType::BloomFilter => Ok(Self::BloomFilter), + IndexType::RTree => Ok(Self::RTree), + IndexType::Fm => Ok(Self::Fm), + _ => Err(Error::index("Invalid index type".to_string())), + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub struct ScalarIndexParams { + /// The type of index to create + /// + /// Plugins may add additional index types. Index type lookup is case-insensitive. + pub index_type: String, + /// The parameters to train the index + /// + /// This should be a JSON string. The contents of the JSON string will be specific to the + /// index type. If not set, then default parameters will be used for the index type. + pub params: Option, +} + +impl Default for ScalarIndexParams { + fn default() -> Self { + Self { + index_type: BuiltinIndexType::BTree.as_str().to_string(), + params: None, + } + } +} + +impl ScalarIndexParams { + /// Creates a new ScalarIndexParams from one of the builtin index types + pub fn for_builtin(index_type: BuiltinIndexType) -> Self { + Self { + index_type: index_type.as_str().to_string(), + params: None, + } + } + + /// Create a new ScalarIndexParams with the given index type + pub fn new(index_type: String) -> Self { + Self { + index_type, + params: None, + } + } + + /// Set the parameters for the index + pub fn with_params(mut self, params: &ParamsType) -> Self { + self.params = Some(serde_json::to_string(params).unwrap()); + self + } +} + +impl IndexParams for ScalarIndexParams { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn index_name(&self) -> &str { + LANCE_SCALAR_INDEX + } +} + +/// Trait for storing an index (or parts of an index) into storage +#[async_trait] +pub trait IndexWriter: Send { + /// Writes a record batch into the file, returning the 0-based index of the batch in the file + /// + /// E.g. if this is the third time this is called this method will return 2 + async fn write_record_batch(&mut self, batch: RecordBatch) -> Result; + /// Adds a global buffer and returns its index. + async fn add_global_buffer(&mut self, _data: Bytes) -> Result { + Err(Error::not_supported( + "global buffers are not supported by this index writer", + )) + } + /// Finishes writing the file and closes the file + async fn finish(&mut self) -> Result; + /// Finishes writing the file and closes the file with additional metadata + async fn finish_with_metadata( + &mut self, + metadata: HashMap, + ) -> Result; +} + +/// Trait for reading an index (or parts of an index) from storage +#[async_trait] +pub trait IndexReader: Send + Sync { + /// Read the n-th record batch from the file + async fn read_record_batch(&self, n: u64, batch_size: u64) -> Result; + /// Reads a global buffer by index. + async fn read_global_buffer(&self, _index: u32) -> Result { + Err(Error::not_supported( + "global buffers are not supported by this index reader", + )) + } + /// Read the range of rows from the file. + /// If projection is Some, only return the columns in the projection, + /// nested columns like Some(&["x.y"]) are not supported. + /// If projection is None, return all columns. + async fn read_range( + &self, + range: std::ops::Range, + projection: Option<&[&str]>, + ) -> Result; + /// Read multiple ranges and concatenate into a single batch. + /// Default impl runs `read_range`s in parallel via `try_join_all`. + async fn read_ranges( + &self, + ranges: &[std::ops::Range], + projection: Option<&[&str]>, + ) -> Result { + if ranges.is_empty() { + return self.read_range(0..0, projection).await; + } + let futures = ranges + .iter() + .map(|r| self.read_range(r.clone(), projection)); + let batches = futures::future::try_join_all(futures).await?; + let schema = batches[0].schema(); + Ok(arrow_select::concat::concat_batches(&schema, &batches)?) + } + /// Read a range of rows as a stream of record batches. + /// + /// This allows the caller to process rows incrementally without loading the + /// entire range into memory at once. + /// + /// The default implementation falls back to [`Self::read_range`] and wraps + /// the result in a single-item stream. + async fn read_range_stream( + &self, + range: std::ops::Range, + projection: Option<&[&str]>, + ) -> Result>> { + let batch = self.read_range(range, projection).await?; + let schema = batch.schema(); + Ok(Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::once(async move { Ok(batch) }), + ))) + } + /// Return the number of batches in the file + async fn num_batches(&self, batch_size: u64) -> u32; + /// Return the number of rows in the file + fn num_rows(&self) -> usize; + /// Return the metadata of the file + fn schema(&self) -> &lance_core::datatypes::Schema; + /// Best-effort on-disk byte size of the file when the reader already knows it + /// without extra I/O, else `None`. Used to size prewarm chunks. + fn file_size_bytes(&self) -> Option { + None + } +} + +/// Trait abstracting I/O away from index logic +/// +/// Scalar indices are currently serialized as indexable arrow record batches stored in +/// named "files". The index store is responsible for serializing and deserializing +/// these batches into file data (e.g. as .lance files or .parquet files, etc.) +#[async_trait] +pub trait IndexStore: std::fmt::Debug + Send + Sync + DeepSizeOf { + fn as_any(&self) -> &dyn Any; + fn clone_arc(&self) -> Arc; + + /// Suggested I/O parallelism for the store + fn io_parallelism(&self) -> usize; + + /// Create a new file and return a writer to store data in the file + async fn new_index_file(&self, name: &str, schema: Arc) + -> Result>; + + /// Open an existing file for retrieval + async fn open_index_file(&self, name: &str) -> Result>; + + /// Return a store that submits its I/O at the given base priority. + fn with_io_priority(&self, io_priority: u64) -> Arc; + + /// Copy a range of batches from an index file from this store to another + /// + /// This is often useful when remapping or updating + async fn copy_index_file(&self, name: &str, dest_store: &dyn IndexStore) -> Result; + + /// Copy an index file from this store to a new name in another store, leaving the source intact + async fn copy_index_file_to( + &self, + name: &str, + new_name: &str, + dest_store: &dyn IndexStore, + ) -> Result { + if name == new_name { + self.copy_index_file(name, dest_store).await + } else { + Err(Error::not_supported(format!( + "copying index file {name} to {new_name} is not supported by this index store" + ))) + } + } + + /// Rename an index file + async fn rename_index_file(&self, name: &str, new_name: &str) -> Result; + + /// Delete an index file (used in the tmp spill store to keep tmp size down) + async fn delete_index_file(&self, name: &str) -> Result<()>; + + /// List all files in the index directory with their sizes. + /// + /// Returns a list of (relative_path, size_bytes) tuples. + /// Used to capture file metadata after index creation/modification. + async fn list_files_with_sizes(&self) -> Result>; +} + +/// Different scalar indices may support different kinds of queries +/// +/// For example, a btree index can support a wide range of queries (e.g. x > 7) +/// while an index based on FTS only supports queries like "x LIKE 'foo'" +/// +/// This trait is used when we need an object that can represent any kind of query +/// +/// Note: if you are implementing this trait for a query type then you probably also +/// need to implement the scalar query parser trait to create instances of your query at parse time. +pub trait AnyQuery: std::fmt::Debug + Any + Send + Sync { + /// Cast the query as Any to allow for downcasting + fn as_any(&self) -> &dyn Any; + /// Format the query as a string for display purposes + fn format(&self, col: &str) -> String; + /// Convert the query to a datafusion expression + fn to_expr(&self, col: String) -> Expr; + /// Compare this query to another query + fn dyn_eq(&self, other: &dyn AnyQuery) -> bool; +} + +impl PartialEq for dyn AnyQuery { + fn eq(&self, other: &Self) -> bool { + self.dyn_eq(other) + } +} + +/// The result of a search operation against a scalar index +#[derive(Debug, PartialEq)] +pub enum SearchResult { + /// The exact row ids that satisfy the query + Exact(NullableRowAddrSet), + /// Any row id satisfying the query will be in this set but not every + /// row id in this set will satisfy the query, a further recheck step + /// is needed + AtMost(NullableRowAddrSet), + /// All of the given row ids satisfy the query but there may be more + /// + /// No scalar index actually returns this today but it can arise from + /// boolean operations (e.g. NOT(AtMost(x)) == AtLeast(NOT(x))) + AtLeast(NullableRowAddrSet), +} + +impl SearchResult { + pub fn exact(row_ids: impl Into) -> Self { + Self::Exact(NullableRowAddrSet::new(row_ids.into(), Default::default())) + } + + pub fn at_most(row_ids: impl Into) -> Self { + Self::AtMost(NullableRowAddrSet::new(row_ids.into(), Default::default())) + } + + pub fn at_least(row_ids: impl Into) -> Self { + Self::AtLeast(NullableRowAddrSet::new(row_ids.into(), Default::default())) + } + + pub fn with_nulls(self, nulls: impl Into) -> Self { + match self { + Self::Exact(row_ids) => Self::Exact(row_ids.with_nulls(nulls.into())), + Self::AtMost(row_ids) => Self::AtMost(row_ids.with_nulls(nulls.into())), + Self::AtLeast(row_ids) => Self::AtLeast(row_ids.with_nulls(nulls.into())), + } + } + + pub fn row_addrs(&self) -> &NullableRowAddrSet { + match self { + Self::Exact(row_addrs) => row_addrs, + Self::AtMost(row_addrs) => row_addrs, + Self::AtLeast(row_addrs) => row_addrs, + } + } + + pub fn is_exact(&self) -> bool { + matches!(self, Self::Exact(_)) + } +} + +/// Brief information about an index that was created +pub struct CreatedIndex { + /// The details of the index that was created + /// + /// These should be stored somewhere as they will be needed to + /// load the index later. + pub index_details: prost_types::Any, + /// The version of the index that was created + /// + /// This can be used to determine if a reader is able to load the index. + pub index_version: u32, + /// List of files and their sizes for this index + /// + /// This enables skipping HEAD calls when opening indices and provides + /// visibility into index storage size via describe_indices(). + pub files: Vec, +} + +/// The ordering that training data must satisfy +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TrainingOrdering { + /// The input will arrive sorted by the value column in ascending order + Values, + /// The input will arrive sorted by the address column in ascending order + Addresses, + /// The input will arrive in an arbitrary order + None, +} + +#[derive(Debug, Clone)] +pub struct TrainingCriteria { + pub ordering: TrainingOrdering, + pub needs_row_ids: bool, + pub needs_row_addrs: bool, +} + +impl TrainingCriteria { + pub fn new(ordering: TrainingOrdering) -> Self { + Self { + ordering, + needs_row_ids: false, + needs_row_addrs: false, + } + } + + pub fn with_row_id(mut self) -> Self { + self.needs_row_ids = true; + self + } + + pub fn with_row_addr(mut self) -> Self { + self.needs_row_addrs = true; + self + } +} + +/// The criteria that specifies how to update an index +pub struct UpdateCriteria { + /// If true, then we need to read the old data to update the index + /// + /// This should be avoided if possible but is left in for some legacy paths + pub requires_old_data: bool, + /// The criteria required for data (both old and new) + pub data_criteria: TrainingCriteria, +} + +/// Filter used when merging existing scalar-index rows during update. +/// +/// The caller must pick a filter mode that matches the row-id semantics of the +/// dataset: +/// - address-style row IDs: fragment filtering is valid +/// - stable row IDs: use exact row-id membership instead +#[derive(Debug, Clone)] +pub enum OldIndexDataFilter { + /// Keeps track of which fragments are still valid and which are no longer valid. + /// + /// This is valid for address-style row IDs. + Fragments { + to_keep: RoaringBitmap, + to_remove: RoaringBitmap, + }, + /// Keep old rows whose row IDs are in this exact allow-list. + /// + /// This is required for stable row IDs, where row IDs are opaque and + /// should not be interpreted as encoded row addresses. + RowIds(RowAddrTreeMap), +} + +impl OldIndexDataFilter { + /// Build a boolean mask that keeps only row IDs selected by this filter. + pub fn filter_row_ids(&self, row_ids: &UInt64Array) -> BooleanArray { + match self { + Self::Fragments { to_keep, .. } => row_ids + .iter() + .map(|id| id.map(|id| to_keep.contains((id >> 32) as u32))) + .collect(), + Self::RowIds(valid_row_ids) => row_ids + .iter() + .map(|id| id.map(|id| valid_row_ids.contains(id))) + .collect(), + } + } + + /// Apply this filter in place to a set of existing (old) row ids/addresses, + /// retaining only the rows the filter selects to keep. Used by index types + /// that merge old postings directly (e.g. bitmap) instead of re-scanning a + /// row-id array through [`Self::filter_row_ids`]. + pub fn retain_old_rows(&self, rows: &mut RowAddrTreeMap) { + match self { + Self::Fragments { to_keep, .. } => rows.retain_fragments(to_keep.iter()), + Self::RowIds(valid_row_ids) => *rows &= valid_row_ids, + } + } +} + +impl UpdateCriteria { + pub fn requires_old_data(data_criteria: TrainingCriteria) -> Self { + Self { + requires_old_data: true, + data_criteria, + } + } + + pub fn only_new_data(data_criteria: TrainingCriteria) -> Self { + Self { + requires_old_data: false, + data_criteria, + } + } +} + +/// A trait for a scalar index, a structure that can determine row ids that satisfy scalar queries +#[async_trait] +pub trait ScalarIndex: Send + Sync + std::fmt::Debug + Index + DeepSizeOf { + /// Search the scalar index + /// + /// Returns all row ids that satisfy the query, these row ids are not necessarily ordered + async fn search( + &self, + query: &dyn AnyQuery, + metrics: &dyn MetricsCollector, + ) -> Result; + + /// Returns true if this index reports matches as physical row addresses + /// (`fragment_id << 32 | offset`) rather than row ids + /// + /// Address-domain indices (e.g. zone map, bloom filter) are built over the + /// `_rowaddr` column. On a dataset with stable row ids the address and + /// row-id domains diverge, so these results must be translated back to row + /// ids (via the per-fragment row-id sequences, known only at the dataset + /// layer) before they are combined with row-id results or handed to the + /// scan. The default (row-id domain) needs no translation. + fn results_are_row_addresses(&self) -> bool { + false + } + + /// Returns true if the remap operation is supported + fn can_remap(&self) -> bool; + + /// Remap the row ids, creating a new remapped version of this index in `dest_store` + async fn remap( + &self, + mapping: &RowAddrRemap, + dest_store: &dyn IndexStore, + ) -> Result; + + /// Add the new data into the index, creating an updated version of the index in `dest_store` + /// + /// If `old_data_filter` is provided, old index data will be filtered before + /// merge according to the chosen filter mode. + async fn update( + &self, + new_data: SendableRecordBatchStream, + dest_store: &dyn IndexStore, + old_data_filter: Option, + ) -> Result; + + /// Returns the criteria that will be used to update the index + fn update_criteria(&self) -> UpdateCriteria; + + /// Derive the index parameters from the current index + /// + /// This returns a ScalarIndexParams that can be used to recreate an index + /// with the same configuration on another dataset. + fn derive_index_params(&self) -> Result; + + /// Global `[min, max]` of the indexed column from index metadata, without a + /// scan, or `None` if this index type cannot supply a sound bound. When + /// `Some`, the range is a superset of live values (conservative under + /// deletes): safe to prune with, not guaranteed tight. + fn value_range(&self) -> Option<(ScalarValue, ScalarValue)> { + None + } +} + +/// Abstraction over any type that can remap row IDs during index loading. +/// +/// This decouples scalar index plugins from the table-level frag reuse index type. +/// The frag reuse index implements this trait, but callers may also supply custom +/// implementations for testing or other remapping strategies. +pub trait RowIdRemapper: Send + Sync + std::fmt::Debug { + /// Remap a single row id. Returns `None` if the row was deleted. + fn remap_row_id(&self, row_id: u64) -> Option; + /// Remap all addresses in a [`RowAddrTreeMap`], dropping deleted rows. + fn remap_row_addrs_tree_map(&self, row_addrs: &RowAddrTreeMap) -> RowAddrTreeMap; + /// Remap all row ids in a [`RoaringTreemap`], dropping deleted rows. + fn remap_row_ids_roaring_tree_map(&self, row_ids: &RoaringTreemap) -> RoaringTreemap; + /// Remap the row-id column at `row_id_idx` inside `batch`, dropping deleted rows. + fn remap_row_ids_record_batch( + &self, + batch: RecordBatch, + row_id_idx: usize, + ) -> Result; +} diff --git a/rust/lance-index/Cargo.toml b/rust/lance-index/Cargo.toml index b9d1a5fde29..652bf73cc26 100644 --- a/rust/lance-index/Cargo.toml +++ b/rust/lance-index/Cargo.toml @@ -39,6 +39,7 @@ jsonb.workspace = true lance-arrow.workspace = true lance-arrow-stats.workspace = true lance-core.workspace = true +lance-index-core.workspace = true lance-datafusion.workspace = true lance-encoding.workspace = true lance-file.workspace = true diff --git a/rust/lance-index/src/frag_reuse.rs b/rust/lance-index/src/frag_reuse.rs index 4c70db44094..7072cdbe446 100644 --- a/rust/lance-index/src/frag_reuse.rs +++ b/rust/lance-index/src/frag_reuse.rs @@ -5,14 +5,16 @@ //! //! The data structures and table-format logic live in //! [`lance_table::system_index::frag_reuse`]; this module re-exports them and -//! implements the local [`Index`] trait for [`FragReuseIndex`]. +//! provides newtype wrappers that implement the [`Index`] and [`RowIdRemapper`] +//! traits. use std::any::Any; use std::sync::Arc; use arrow_array::RecordBatch; use async_trait::async_trait; -use lance_core::{Error, Result}; +use lance_core::Result; +use lance_core::deepsize::DeepSizeOf; use lance_select::RowAddrTreeMap; use roaring::{RoaringBitmap, RoaringTreemap}; use serde::Serialize; @@ -22,25 +24,22 @@ pub use lance_table::system_index::frag_reuse::*; use crate::scalar::RowIdRemapper; use crate::{Index, IndexType}; -impl RowIdRemapper for FragReuseIndex { - fn remap_row_id(&self, row_id: u64) -> Option { - self.remap_row_id(row_id) - } +/// Newtype wrapping [`FragReuseIndex`] so that `lance-index` can implement +/// the `Index` and `RowIdRemapper` traits (orphan rules prevent implementing +/// them directly in `lance-table`). +pub struct FragReuseIndexHandle(pub Arc); - fn remap_row_addrs_tree_map(&self, row_addrs: &RowAddrTreeMap) -> RowAddrTreeMap { - self.remap_row_addrs_tree_map(row_addrs) +impl std::fmt::Debug for FragReuseIndexHandle { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_tuple("FragReuseIndexHandle") + .field(&self.0) + .finish() } +} - fn remap_row_ids_roaring_tree_map(&self, row_ids: &RoaringTreemap) -> RoaringTreemap { - self.remap_row_ids_roaring_tree_map(row_ids) - } - - fn remap_row_ids_record_batch( - &self, - batch: RecordBatch, - row_id_idx: usize, - ) -> Result { - self.remap_row_ids_record_batch(batch, row_id_idx) +impl DeepSizeOf for FragReuseIndexHandle { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + self.0.deep_size_of_children(context) } } @@ -50,7 +49,7 @@ struct FragReuseStatistics { } #[async_trait] -impl Index for FragReuseIndex { +impl Index for FragReuseIndexHandle { fn as_any(&self) -> &dyn Any { self } @@ -61,10 +60,10 @@ impl Index for FragReuseIndex { fn statistics(&self) -> Result { let stats = FragReuseStatistics { - num_versions: self.details.versions.len(), + num_versions: self.0.details.versions.len(), }; serde_json::to_value(stats).map_err(|e| { - Error::internal(format!( + lance_core::Error::internal(format!( "failed to serialize fragment reuse index statistics: {}", e )) @@ -83,3 +82,25 @@ impl Index for FragReuseIndex { unimplemented!() } } + +impl RowIdRemapper for FragReuseIndexHandle { + fn remap_row_id(&self, row_id: u64) -> Option { + self.0.remap_row_id(row_id) + } + + fn remap_row_addrs_tree_map(&self, row_addrs: &RowAddrTreeMap) -> RowAddrTreeMap { + self.0.remap_row_addrs_tree_map(row_addrs) + } + + fn remap_row_ids_roaring_tree_map(&self, row_ids: &RoaringTreemap) -> RoaringTreemap { + self.0.remap_row_ids_roaring_tree_map(row_ids) + } + + fn remap_row_ids_record_batch( + &self, + batch: RecordBatch, + row_id_idx: usize, + ) -> Result { + self.0.remap_row_ids_record_batch(batch, row_id_idx) + } +} diff --git a/rust/lance-index/src/lib.rs b/rust/lance-index/src/lib.rs index 61b45550367..cf60ece2fcf 100644 --- a/rust/lance-index/src/lib.rs +++ b/rust/lance-index/src/lib.rs @@ -9,16 +9,9 @@ //! API stability is not guaranteed. //! -use std::{any::Any, sync::Arc}; - use crate::frag_reuse::FRAG_REUSE_INDEX_NAME; use crate::mem_wal::MEM_WAL_INDEX_NAME; -use async_trait::async_trait; -use lance_core::deepsize::DeepSizeOf; -use lance_core::{Error, Result}; -use roaring::RoaringBitmap; use serde::{Deserialize, Serialize}; -use std::convert::TryFrom; pub mod frag_reuse; pub mod mem_wal; @@ -33,6 +26,9 @@ pub mod vector; pub use crate::traits::*; +// Re-export core traits from lance-index-core +pub use lance_index_core::{Index, IndexParams, IndexType}; + pub const INDEX_FILE_NAME: &str = "index.idx"; /// The name of the auxiliary index file. /// @@ -75,280 +71,6 @@ pub mod cache_pb { include!(concat!(env!("OUT_DIR"), "/lance.index.cache.rs")); } -/// Generic methods common across all types of secondary indices -/// -#[async_trait] -pub trait Index: Send + Sync + DeepSizeOf { - /// Cast to [Any]. - fn as_any(&self) -> &dyn Any; - - /// Cast to [Index] - fn as_index(self: Arc) -> Arc; - - /// Retrieve index statistics as a JSON Value - fn statistics(&self) -> Result; - - /// Prewarm the index. - /// - /// This will load the index into memory and cache it. - async fn prewarm(&self) -> Result<()>; - - /// Get the type of the index - fn index_type(&self) -> IndexType; - - /// Read through the index and determine which fragment ids are covered by the index - /// - /// This is a kind of slow operation. It's better to use the fragment_bitmap. This - /// only exists for cases where the fragment_bitmap has become corrupted or missing. - async fn calculate_included_frags(&self) -> Result; -} - -/// Index Type -#[derive(Debug, PartialEq, Eq, Copy, Hash, Clone, DeepSizeOf)] -pub enum IndexType { - // Preserve 0-100 for simple indices. - Scalar = 0, // Legacy scalar index, alias to BTree - - BTree = 1, // BTree - - Bitmap = 2, // Bitmap - - LabelList = 3, // LabelList - - Inverted = 4, // Inverted - - NGram = 5, // NGram - - FragmentReuse = 6, - - MemWal = 7, - - ZoneMap = 8, // ZoneMap - - BloomFilter = 9, // Bloom filter - - RTree = 10, // RTree - - Fm = 11, // FM-Index - - // 100+ and up for vector index. - /// Flat vector index. - Vector = 100, // Legacy vector index, alias to IvfPq - IvfFlat = 101, - IvfSq = 102, - IvfPq = 103, - IvfHnswSq = 104, - IvfHnswPq = 105, - IvfHnswFlat = 106, - IvfRq = 107, -} - -impl std::fmt::Display for IndexType { - fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - match self { - Self::Scalar | Self::BTree => write!(f, "BTree"), - Self::Bitmap => write!(f, "Bitmap"), - Self::LabelList => write!(f, "LabelList"), - Self::Inverted => write!(f, "Inverted"), - Self::NGram => write!(f, "NGram"), - Self::FragmentReuse => write!(f, "FragmentReuse"), - Self::MemWal => write!(f, "MemWal"), - Self::ZoneMap => write!(f, "ZoneMap"), - Self::BloomFilter => write!(f, "BloomFilter"), - Self::RTree => write!(f, "RTree"), - Self::Fm => write!(f, "Fm"), - Self::Vector | Self::IvfPq => write!(f, "IVF_PQ"), - Self::IvfFlat => write!(f, "IVF_FLAT"), - Self::IvfSq => write!(f, "IVF_SQ"), - Self::IvfHnswSq => write!(f, "IVF_HNSW_SQ"), - Self::IvfHnswPq => write!(f, "IVF_HNSW_PQ"), - Self::IvfHnswFlat => write!(f, "IVF_HNSW_FLAT"), - Self::IvfRq => write!(f, "IVF_RQ"), - } - } -} - -impl TryFrom for IndexType { - type Error = Error; - - fn try_from(value: i32) -> Result { - match value { - v if v == Self::Scalar as i32 => Ok(Self::Scalar), - v if v == Self::BTree as i32 => Ok(Self::BTree), - v if v == Self::Bitmap as i32 => Ok(Self::Bitmap), - v if v == Self::LabelList as i32 => Ok(Self::LabelList), - v if v == Self::NGram as i32 => Ok(Self::NGram), - v if v == Self::Inverted as i32 => Ok(Self::Inverted), - v if v == Self::FragmentReuse as i32 => Ok(Self::FragmentReuse), - v if v == Self::MemWal as i32 => Ok(Self::MemWal), - v if v == Self::ZoneMap as i32 => Ok(Self::ZoneMap), - v if v == Self::BloomFilter as i32 => Ok(Self::BloomFilter), - v if v == Self::RTree as i32 => Ok(Self::RTree), - v if v == Self::Fm as i32 => Ok(Self::Fm), - v if v == Self::Vector as i32 => Ok(Self::Vector), - v if v == Self::IvfFlat as i32 => Ok(Self::IvfFlat), - v if v == Self::IvfSq as i32 => Ok(Self::IvfSq), - v if v == Self::IvfPq as i32 => Ok(Self::IvfPq), - v if v == Self::IvfHnswSq as i32 => Ok(Self::IvfHnswSq), - v if v == Self::IvfHnswPq as i32 => Ok(Self::IvfHnswPq), - v if v == Self::IvfHnswFlat as i32 => Ok(Self::IvfHnswFlat), - v if v == Self::IvfRq as i32 => Ok(Self::IvfRq), - _ => Err(Error::invalid_input_source( - format!("the input value {} is not a valid IndexType", value).into(), - )), - } - } -} - -impl TryFrom<&str> for IndexType { - type Error = Error; - - fn try_from(value: &str) -> Result { - match value { - "BTree" | "BTREE" => Ok(Self::BTree), - "Bitmap" | "BITMAP" => Ok(Self::Bitmap), - "LabelList" | "LABELLIST" => Ok(Self::LabelList), - "Inverted" | "INVERTED" => Ok(Self::Inverted), - "NGram" | "NGRAM" => Ok(Self::NGram), - "ZoneMap" | "ZONEMAP" => Ok(Self::ZoneMap), - "BloomFilter" | "BLOOMFILTER" | "BLOOM_FILTER" => Ok(Self::BloomFilter), - "RTree" | "RTREE" | "R_TREE" => Ok(Self::RTree), - "Fm" | "FM" => Ok(Self::Fm), - "Vector" | "VECTOR" => Ok(Self::Vector), - "IVF_FLAT" => Ok(Self::IvfFlat), - "IVF_SQ" => Ok(Self::IvfSq), - "IVF_PQ" => Ok(Self::IvfPq), - "IVF_RQ" => Ok(Self::IvfRq), - "IVF_HNSW_FLAT" => Ok(Self::IvfHnswFlat), - "IVF_HNSW_SQ" => Ok(Self::IvfHnswSq), - "IVF_HNSW_PQ" => Ok(Self::IvfHnswPq), - "FragmentReuse" => Ok(Self::FragmentReuse), - "MemWal" => Ok(Self::MemWal), - _ => Err(Error::invalid_input(format!( - "invalid index type: {}", - value - ))), - } - } -} - -impl IndexType { - pub fn is_scalar(&self) -> bool { - matches!( - self, - Self::Scalar - | Self::BTree - | Self::Bitmap - | Self::LabelList - | Self::Inverted - | Self::NGram - | Self::ZoneMap - | Self::BloomFilter - | Self::RTree - | Self::Fm, - ) - } - - pub fn is_vector(&self) -> bool { - matches!( - self, - Self::Vector - | Self::IvfPq - | Self::IvfHnswSq - | Self::IvfHnswPq - | Self::IvfHnswFlat - | Self::IvfFlat - | Self::IvfSq - | Self::IvfRq - ) - } - - pub fn is_system(&self) -> bool { - matches!(self, Self::FragmentReuse | Self::MemWal) - } - - /// Returns the current format version of the index type, - /// bump this when the index format changes. - /// Indices which higher version than these will be ignored for compatibility, - /// This would happen when creating index in a newer version of Lance, - /// but then opening the index in older version of Lance - pub fn version(&self) -> i32 { - match self { - Self::Scalar => 0, - Self::BTree => 0, - Self::Bitmap => 0, - Self::LabelList => 0, - Self::Inverted => 0, - Self::NGram => 0, - Self::FragmentReuse => 0, - Self::MemWal => 0, - Self::ZoneMap => 0, - Self::BloomFilter => 0, - Self::RTree => 0, - Self::Fm => 0, - - // IMPORTANT: if any vector index subtype needs a format bump that is - // not backward compatible, its new version must be set to - // (current max vector index version + 1), even if only one subtype - // changed. Compatibility filtering currently cannot distinguish vector - // subtypes from details-only metadata, so vector versions effectively - // share one global monotonic compatibility level. - Self::Vector - | Self::IvfFlat - | Self::IvfSq - | Self::IvfPq - | Self::IvfHnswSq - | Self::IvfHnswPq - | Self::IvfHnswFlat => VECTOR_INDEX_VERSION as i32, - Self::IvfRq => IVF_RQ_INDEX_VERSION as i32, - } - } - - /// Returns the target partition size for the index type. - /// - /// This is used to compute the number of partitions for the index. - /// The partition size is optimized for the best performance of the index. - /// - /// This is for vector indices only. - pub fn target_partition_size(&self) -> usize { - match self { - Self::Vector => 8192, - Self::IvfFlat => 4096, - Self::IvfSq => 8192, - Self::IvfPq => 8192, - Self::IvfRq => 4096, - Self::IvfHnswFlat => 1 << 20, - Self::IvfHnswSq => 1 << 20, - Self::IvfHnswPq => 1 << 20, - _ => 8192, - } - } - - /// Returns the highest supported vector index version in this Lance build. - pub fn max_vector_version() -> u32 { - [ - Self::Vector, - Self::IvfFlat, - Self::IvfSq, - Self::IvfPq, - Self::IvfHnswSq, - Self::IvfHnswPq, - Self::IvfHnswFlat, - Self::IvfRq, - ] - .into_iter() - .map(|index_type| index_type.version() as u32) - .max() - .unwrap_or(VECTOR_INDEX_VERSION) - } -} - -pub trait IndexParams: Send + Sync { - fn as_any(&self) -> &dyn Any; - - fn index_name(&self) -> &str; -} - #[derive(Serialize, Deserialize, Debug)] pub struct IndexMetadata { #[serde(rename = "type")] diff --git a/rust/lance-index/src/mem_wal.rs b/rust/lance-index/src/mem_wal.rs index 9bd72ff7866..b169d8c6994 100644 --- a/rust/lance-index/src/mem_wal.rs +++ b/rust/lance-index/src/mem_wal.rs @@ -5,13 +5,14 @@ //! //! The data structures and table-format logic live in //! [`lance_table::system_index::mem_wal`]; this module re-exports them and -//! implements the local [`Index`] trait for [`MemWalIndex`]. +//! provides a newtype wrapper that implements the [`Index`] trait. use std::any::Any; use std::sync::Arc; use async_trait::async_trait; -use lance_core::Error; +use lance_core::Result; +use lance_core::deepsize::DeepSizeOf; use roaring::RoaringBitmap; use serde::Serialize; @@ -19,6 +20,17 @@ pub use lance_table::system_index::mem_wal::*; use crate::{Index, IndexType}; +/// Newtype wrapping [`MemWalIndex`] so that `lance-index` can implement +/// the `Index` trait (orphan rules prevent implementing it directly in +/// `lance-table`). +pub struct MemWalIndexHandle(pub Arc); + +impl DeepSizeOf for MemWalIndexHandle { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + self.0.deep_size_of_children(context) + } +} + #[derive(Serialize)] struct MemWalStatistics { num_shards: u32, @@ -29,7 +41,7 @@ struct MemWalStatistics { } #[async_trait] -impl Index for MemWalIndex { +impl Index for MemWalIndexHandle { fn as_any(&self) -> &dyn Any { self } @@ -38,23 +50,23 @@ impl Index for MemWalIndex { self } - fn statistics(&self) -> lance_core::Result { + fn statistics(&self) -> Result { let stats = MemWalStatistics { - num_shards: self.details.num_shards, - num_merged_generations: self.details.merged_generations.len(), - num_shard_specs: self.details.sharding_specs.len(), - num_maintained_indexes: self.details.maintained_indexes.len(), - num_index_catchup_entries: self.details.index_catchup.len(), + num_shards: self.0.details.num_shards, + num_merged_generations: self.0.details.merged_generations.len(), + num_shard_specs: self.0.details.sharding_specs.len(), + num_maintained_indexes: self.0.details.maintained_indexes.len(), + num_index_catchup_entries: self.0.details.index_catchup.len(), }; serde_json::to_value(stats).map_err(|e| { - Error::internal(format!( + lance_core::Error::internal(format!( "failed to serialize MemWAL index statistics: {}", e )) }) } - async fn prewarm(&self) -> lance_core::Result<()> { + async fn prewarm(&self) -> Result<()> { Ok(()) } @@ -62,7 +74,7 @@ impl Index for MemWalIndex { IndexType::MemWal } - async fn calculate_included_frags(&self) -> lance_core::Result { + async fn calculate_included_frags(&self) -> Result { Ok(RoaringBitmap::new()) } } diff --git a/rust/lance-index/src/metrics.rs b/rust/lance-index/src/metrics.rs index 8c0c119a3c3..7d1ea2964c6 100644 --- a/rust/lance-index/src/metrics.rs +++ b/rust/lance-index/src/metrics.rs @@ -1,117 +1,4 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use std::sync::atomic::{AtomicUsize, Ordering}; - -pub const AND_CANDIDATES_SEEN_METRIC: &str = "and_candidates_seen"; -pub const AND_CANDIDATES_PRUNED_BEFORE_RETURN_METRIC: &str = "and_candidates_pruned_before_return"; -pub const AND_FULL_SCORES_METRIC: &str = "and_full_scores"; -pub const FREQS_COLLECTED_METRIC: &str = "freqs_collected"; - -/// A trait used by the index to report metrics -/// -/// Callers can implement this trait to collect metrics -pub trait MetricsCollector: Send + Sync { - /// Record partition loads - /// - /// Many indices consist of partitions that may need to be loaded - /// into cache. For example, an inverted index or ngram index has a - /// posting list for each token. - /// - /// In the ideal case, these shards are in the cache and will not need - /// to be loaded from disk. This method should not be called if the - /// shard is in the cache. - fn record_parts_loaded(&self, num_parts: usize); - - /// Record a shard load - fn record_part_load(&self) { - self.record_parts_loaded(1); - } - - /// Record an index load - /// - /// This should be called when a scalar index is loaded from storage. - /// It should not be called if the index is already in memory. - fn record_index_loads(&self, num_indexes: usize); - - /// Record an index load - fn record_index_load(&self) { - self.record_index_loads(1); - } - - /// Record the number of "comparisons" made by the index - /// - /// What exactly constitutes a comparison depends on the index type. - /// For example, a B-tree index may make comparisons while searching for a value. - /// On the other hand, a bitmap index makes comparisons when computing the intersection - /// of two bitmaps. - /// - /// The goal is to provide some visibility into the compute cost of the search - fn record_comparisons(&self, num_comparisons: usize); - - /// Record AND candidates returned from WAND alignment to the scoring loop. - /// - /// This excludes candidates pruned before `next()` returns. Use this with - /// `record_and_candidates_pruned_before_return` to recover total aligned - /// AND candidates. - fn record_and_candidates_seen(&self, _num_candidates: usize) {} - - /// Record AND candidates pruned during WAND alignment before `next()` returns. - fn record_and_candidates_pruned_before_return(&self, _num_candidates: usize) {} - - fn record_and_full_scores(&self, _num_scores: usize) {} - - fn record_freqs_collected(&self, _num_collections: usize) {} - - /// Returns an optional sink for recording exact I/O statistics (bytes read, - /// IOPS, and requests) performed on behalf of this collector. - /// - /// Index implementations that read from a - /// [`lance_io::scheduler::ScanScheduler`] can attach the returned handle to - /// their file readers so the I/O performed for a single query is measured - /// and attributed here. The default returns `None`, meaning the caller does - /// not want I/O measured (and index implementations should then take their - /// normal, uninstrumented read path). - fn io_stats(&self) -> Option { - None - } -} - -/// A no-op metrics collector that does nothing -pub struct NoOpMetricsCollector; - -impl MetricsCollector for NoOpMetricsCollector { - fn record_parts_loaded(&self, _num_parts: usize) {} - fn record_index_loads(&self, _num_indexes: usize) {} - fn record_comparisons(&self, _num_comparisons: usize) {} -} - -#[derive(Default)] -pub struct LocalMetricsCollector { - pub parts_loaded: AtomicUsize, - pub index_loads: AtomicUsize, - pub comparisons: AtomicUsize, -} - -impl LocalMetricsCollector { - pub fn dump_into(self, other: &dyn MetricsCollector) { - other.record_parts_loaded(self.parts_loaded.load(Ordering::Relaxed)); - other.record_index_loads(self.index_loads.load(Ordering::Relaxed)); - other.record_comparisons(self.comparisons.load(Ordering::Relaxed)); - } -} - -impl MetricsCollector for LocalMetricsCollector { - fn record_parts_loaded(&self, num_parts: usize) { - self.parts_loaded.fetch_add(num_parts, Ordering::Relaxed); - } - - fn record_index_loads(&self, num_indexes: usize) { - self.index_loads.fetch_add(num_indexes, Ordering::Relaxed); - } - - fn record_comparisons(&self, num_comparisons: usize) { - self.comparisons - .fetch_add(num_comparisons, Ordering::Relaxed); - } -} +pub use lance_index_core::metrics::*; diff --git a/rust/lance-index/src/scalar.rs b/rust/lance-index/src/scalar.rs index 21ae7aac71c..33e4082e114 100644 --- a/rust/lance-index/src/scalar.rs +++ b/rust/lance-index/src/scalar.rs @@ -4,19 +4,13 @@ //! Scalar indices for metadata search & filtering use arrow::buffer::{OffsetBuffer, ScalarBuffer}; -use arrow_array::{BooleanArray, ListArray, RecordBatch, UInt64Array}; -use arrow_schema::{Field, Schema}; -use async_trait::async_trait; -use bytes::Bytes; +use arrow_array::ListArray; +use arrow_schema::Field; use datafusion::functions::regex::regexplike::RegexpLikeFunc; use datafusion::functions::string::contains::ContainsFunc; use datafusion::functions_nested::array_has; -use datafusion::physical_plan::SendableRecordBatchStream; use datafusion_common::{Column, scalar::ScalarValue}; -use lance_core::utils::row_addr_remap::RowAddrRemap; -use std::collections::{HashMap, HashSet}; -use std::fmt::Debug; -use std::pin::Pin; +use std::collections::HashSet; use std::{any::Any, ops::Bound, sync::Arc}; use datafusion_expr::{ @@ -24,17 +18,17 @@ use datafusion_expr::{ expr::{Like, ScalarFunction}, }; use inverted::query::{FtsQuery, FtsQueryNode, FtsSearchParams, MatchQuery, fill_fts_query_column}; -use lance_core::deepsize::DeepSizeOf; -use lance_core::{Error, Result}; -use lance_io::stream::{RecordBatchStream, RecordBatchStreamAdapter}; -use lance_select::{NullableRowAddrSet, RowAddrTreeMap, RowSetOps}; -use roaring::{RoaringBitmap, RoaringTreemap}; -use serde::Serialize; - -use crate::metrics::MetricsCollector; -use crate::scalar::registry::TrainingCriteria; -use crate::{Index, IndexParams, IndexType}; -pub use lance_table::format::IndexFile; +use lance_core::Result; + +use lance_datafusion::udf::CONTAINS_TOKENS_UDF; + +use crate::IndexParams; +pub use crate::metrics::MetricsCollector; +pub use lance_index_core::scalar::{ + AnyQuery, BuiltinIndexType, CreatedIndex, IndexFile, IndexReader, IndexStore, IndexWriter, + LANCE_SCALAR_INDEX, OldIndexDataFilter, RowIdRemapper, ScalarIndex, ScalarIndexParams, + SearchResult, TrainingCriteria, TrainingOrdering, UpdateCriteria, +}; pub mod bitmap; pub mod bloomfilter; @@ -53,117 +47,39 @@ pub mod zoned; pub mod zonemap; pub use inverted::tokenizer::InvertedIndexParams; -use lance_datafusion::udf::CONTAINS_TOKENS_UDF; -pub const LANCE_SCALAR_INDEX: &str = "__lance_scalar_index"; - -/// Builtin index types supported by the Lance library +/// Convert a `Vec<`[`lance_index_core::scalar::IndexFile`]`>` to a +/// `Vec<`[`lance_table::format::IndexFile`]`>`. /// -/// This is primarily for convenience to avoid a bunch of string -/// constants and provide some auto-complete. This type should not -/// be used in the manifest as plugins cannot add new entries. -#[derive(Debug, Clone, PartialEq, Eq, DeepSizeOf)] -pub enum BuiltinIndexType { - BTree, - Bitmap, - LabelList, - NGram, - ZoneMap, - BloomFilter, - RTree, - Inverted, - Fm, -} - -impl BuiltinIndexType { - pub fn as_str(&self) -> &str { - match self { - Self::BTree => "btree", - Self::Bitmap => "bitmap", - Self::LabelList => "labellist", - Self::NGram => "ngram", - Self::ZoneMap => "zonemap", - Self::Inverted => "inverted", - Self::BloomFilter => "bloomfilter", - Self::RTree => "rtree", - Self::Fm => "fm", - } - } -} - -impl TryFrom for BuiltinIndexType { - type Error = Error; - - fn try_from(value: IndexType) -> Result { - match value { - IndexType::BTree => Ok(Self::BTree), - IndexType::Bitmap => Ok(Self::Bitmap), - IndexType::LabelList => Ok(Self::LabelList), - IndexType::NGram => Ok(Self::NGram), - IndexType::ZoneMap => Ok(Self::ZoneMap), - IndexType::Inverted => Ok(Self::Inverted), - IndexType::BloomFilter => Ok(Self::BloomFilter), - IndexType::RTree => Ok(Self::RTree), - IndexType::Fm => Ok(Self::Fm), - _ => Err(Error::index("Invalid index type".to_string())), - } - } -} - -#[derive(Debug, Clone, PartialEq)] -pub struct ScalarIndexParams { - /// The type of index to create - /// - /// Plugins may add additional index types. Index type lookup is case-insensitive. - pub index_type: String, - /// The parameters to train the index - /// - /// This should be a JSON string. The contents of the JSON string will be specific to the - /// index type. If not set, then default parameters will be used for the index type. - pub params: Option, -} - -impl Default for ScalarIndexParams { - fn default() -> Self { - Self { - index_type: BuiltinIndexType::BTree.as_str().to_string(), - params: None, - } - } -} - -impl ScalarIndexParams { - /// Creates a new ScalarIndexParams from one of the builtin index types - pub fn for_builtin(index_type: BuiltinIndexType) -> Self { - Self { - index_type: index_type.as_str().to_string(), - params: None, - } - } - - /// Create a new ScalarIndexParams with the given index type - pub fn new(index_type: String) -> Self { - Self { - index_type, - params: None, - } - } - - /// Set the parameters for the index - pub fn with_params(mut self, params: &ParamsType) -> Self { - self.params = Some(serde_json::to_string(params).unwrap()); - self - } -} - -impl IndexParams for ScalarIndexParams { - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn index_name(&self) -> &str { - LANCE_SCALAR_INDEX - } +/// These two structs have identical fields; this helper bridges the crate +/// boundary without relying on orphan-rule–violating `From` impls. +pub fn index_files_to_table( + files: Vec, +) -> Vec { + files + .into_iter() + .map(|f| lance_table::format::IndexFile { + path: f.path, + size_bytes: f.size_bytes, + }) + .collect() +} + +/// Convert a `Vec<`[`lance_table::format::IndexFile`]`>` to a +/// `Vec<`[`lance_index_core::scalar::IndexFile`]`>`. +/// +/// These two structs have identical fields; this helper bridges the crate +/// boundary without relying on orphan-rule–violating `From` impls. +pub fn table_files_to_index( + files: Vec, +) -> Vec { + files + .into_iter() + .map(|f| lance_index_core::scalar::IndexFile { + path: f.path, + size_bytes: f.size_bytes, + }) + .collect() } impl IndexParams for InvertedIndexParams { @@ -176,180 +92,6 @@ impl IndexParams for InvertedIndexParams { } } -/// Trait for storing an index (or parts of an index) into storage -#[async_trait] -pub trait IndexWriter: Send { - /// Writes a record batch into the file, returning the 0-based index of the batch in the file - /// - /// E.g. if this is the third time this is called this method will return 2 - async fn write_record_batch(&mut self, batch: RecordBatch) -> Result; - /// Adds a global buffer and returns its index. - async fn add_global_buffer(&mut self, _data: Bytes) -> Result { - Err(Error::not_supported( - "global buffers are not supported by this index writer", - )) - } - /// Finishes writing the file and closes the file - async fn finish(&mut self) -> Result; - /// Finishes writing the file and closes the file with additional metadata - async fn finish_with_metadata( - &mut self, - metadata: HashMap, - ) -> Result; -} - -/// Trait for reading an index (or parts of an index) from storage -#[async_trait] -pub trait IndexReader: Send + Sync { - /// Read the n-th record batch from the file - async fn read_record_batch(&self, n: u64, batch_size: u64) -> Result; - /// Reads a global buffer by index. - async fn read_global_buffer(&self, _index: u32) -> Result { - Err(Error::not_supported( - "global buffers are not supported by this index reader", - )) - } - /// Read the range of rows from the file. - /// If projection is Some, only return the columns in the projection, - /// nested columns like Some(&["x.y"]) are not supported. - /// If projection is None, return all columns. - async fn read_range( - &self, - range: std::ops::Range, - projection: Option<&[&str]>, - ) -> Result; - /// Read multiple ranges and concatenate into a single batch. - /// Default impl runs `read_range`s in parallel via `try_join_all`. - async fn read_ranges( - &self, - ranges: &[std::ops::Range], - projection: Option<&[&str]>, - ) -> Result { - if ranges.is_empty() { - return self.read_range(0..0, projection).await; - } - let futures = ranges - .iter() - .map(|r| self.read_range(r.clone(), projection)); - let batches = futures::future::try_join_all(futures).await?; - let schema = batches[0].schema(); - Ok(arrow_select::concat::concat_batches(&schema, &batches)?) - } - /// Read a range of rows as a stream of record batches. - /// - /// This allows the caller to process rows incrementally without loading the - /// entire range into memory at once. - /// - /// The default implementation falls back to [`Self::read_range`] and wraps - /// the result in a single-item stream. - async fn read_range_stream( - &self, - range: std::ops::Range, - projection: Option<&[&str]>, - ) -> Result>> { - let batch = self.read_range(range, projection).await?; - let schema = batch.schema(); - Ok(Box::pin(RecordBatchStreamAdapter::new( - schema, - futures::stream::once(async move { Ok(batch) }), - ))) - } - /// Return the number of batches in the file - async fn num_batches(&self, batch_size: u64) -> u32; - /// Return the number of rows in the file - fn num_rows(&self) -> usize; - /// Return the metadata of the file - fn schema(&self) -> &lance_core::datatypes::Schema; - /// Best-effort on-disk byte size of the file when the reader already knows it - /// without extra I/O, else `None`. Used to size prewarm chunks. - fn file_size_bytes(&self) -> Option { - None - } -} - -/// Trait abstracting I/O away from index logic -/// -/// Scalar indices are currently serialized as indexable arrow record batches stored in -/// named "files". The index store is responsible for serializing and deserializing -/// these batches into file data (e.g. as .lance files or .parquet files, etc.) -#[async_trait] -pub trait IndexStore: std::fmt::Debug + Send + Sync + DeepSizeOf { - fn as_any(&self) -> &dyn Any; - fn clone_arc(&self) -> Arc; - - /// Suggested I/O parallelism for the store - fn io_parallelism(&self) -> usize; - - /// Create a new file and return a writer to store data in the file - async fn new_index_file(&self, name: &str, schema: Arc) - -> Result>; - - /// Open an existing file for retrieval - async fn open_index_file(&self, name: &str) -> Result>; - - /// Return a store that submits its I/O at the given base priority. - fn with_io_priority(&self, io_priority: u64) -> Arc; - - /// Copy a range of batches from an index file from this store to another - /// - /// This is often useful when remapping or updating - async fn copy_index_file(&self, name: &str, dest_store: &dyn IndexStore) -> Result; - - /// Copy an index file from this store to a new name in another store, leaving the source intact - async fn copy_index_file_to( - &self, - name: &str, - new_name: &str, - dest_store: &dyn IndexStore, - ) -> Result { - if name == new_name { - self.copy_index_file(name, dest_store).await - } else { - Err(Error::not_supported(format!( - "copying index file {name} to {new_name} is not supported by this index store" - ))) - } - } - - /// Rename an index file - async fn rename_index_file(&self, name: &str, new_name: &str) -> Result; - - /// Delete an index file (used in the tmp spill store to keep tmp size down) - async fn delete_index_file(&self, name: &str) -> Result<()>; - - /// List all files in the index directory with their sizes. - /// - /// Returns a list of (relative_path, size_bytes) tuples. - /// Used to capture file metadata after index creation/modification. - async fn list_files_with_sizes(&self) -> Result>; -} - -/// Different scalar indices may support different kinds of queries -/// -/// For example, a btree index can support a wide range of queries (e.g. x > 7) -/// while an index based on FTS only supports queries like "x LIKE 'foo'" -/// -/// This trait is used when we need an object that can represent any kind of query -/// -/// Note: if you are implementing this trait for a query type then you probably also -/// need to implement the [crate::scalar::expression::ScalarQueryParser] trait to -/// create instances of your query at parse time. -pub trait AnyQuery: std::fmt::Debug + Any + Send + Sync { - /// Cast the query as Any to allow for downcasting - fn as_any(&self) -> &dyn Any; - /// Format the query as a string for display purposes - fn format(&self, col: &str) -> String; - /// Convert the query to a datafusion expression - fn to_expr(&self, col: String) -> Expr; - /// Compare this query to another query - fn dyn_eq(&self, other: &dyn AnyQuery) -> bool; -} - -impl PartialEq for dyn AnyQuery { - fn eq(&self, other: &Self) -> bool { - self.dyn_eq(other) - } -} /// A full text search query #[derive(Debug, Clone, PartialEq)] pub struct FullTextSearchQuery { @@ -892,149 +634,6 @@ impl AnyQuery for GeoQuery { } } -/// The result of a search operation against a scalar index -#[derive(Debug, PartialEq)] -pub enum SearchResult { - /// The exact row ids that satisfy the query - Exact(NullableRowAddrSet), - /// Any row id satisfying the query will be in this set but not every - /// row id in this set will satisfy the query, a further recheck step - /// is needed - AtMost(NullableRowAddrSet), - /// All of the given row ids satisfy the query but there may be more - /// - /// No scalar index actually returns this today but it can arise from - /// boolean operations (e.g. NOT(AtMost(x)) == AtLeast(NOT(x))) - AtLeast(NullableRowAddrSet), -} - -impl SearchResult { - pub fn exact(row_ids: impl Into) -> Self { - Self::Exact(NullableRowAddrSet::new(row_ids.into(), Default::default())) - } - - pub fn at_most(row_ids: impl Into) -> Self { - Self::AtMost(NullableRowAddrSet::new(row_ids.into(), Default::default())) - } - - pub fn at_least(row_ids: impl Into) -> Self { - Self::AtLeast(NullableRowAddrSet::new(row_ids.into(), Default::default())) - } - - pub fn with_nulls(self, nulls: impl Into) -> Self { - match self { - Self::Exact(row_ids) => Self::Exact(row_ids.with_nulls(nulls.into())), - Self::AtMost(row_ids) => Self::AtMost(row_ids.with_nulls(nulls.into())), - Self::AtLeast(row_ids) => Self::AtLeast(row_ids.with_nulls(nulls.into())), - } - } - - pub fn row_addrs(&self) -> &NullableRowAddrSet { - match self { - Self::Exact(row_addrs) => row_addrs, - Self::AtMost(row_addrs) => row_addrs, - Self::AtLeast(row_addrs) => row_addrs, - } - } - - pub fn is_exact(&self) -> bool { - matches!(self, Self::Exact(_)) - } -} - -/// Brief information about an index that was created -pub struct CreatedIndex { - /// The details of the index that was created - /// - /// These should be stored somewhere as they will be needed to - /// load the index later. - pub index_details: prost_types::Any, - /// The version of the index that was created - /// - /// This can be used to determine if a reader is able to load the index. - pub index_version: u32, - /// List of files and their sizes for this index - /// - /// This enables skipping HEAD calls when opening indices and provides - /// visibility into index storage size via describe_indices(). - pub files: Vec, -} - -/// The criteria that specifies how to update an index -pub struct UpdateCriteria { - /// If true, then we need to read the old data to update the index - /// - /// This should be avoided if possible but is left in for some legacy paths - pub requires_old_data: bool, - /// The criteria required for data (both old and new) - pub data_criteria: TrainingCriteria, -} - -/// Filter used when merging existing scalar-index rows during update. -/// -/// The caller must pick a filter mode that matches the row-id semantics of the -/// dataset: -/// - address-style row IDs: fragment filtering is valid -/// - stable row IDs: use exact row-id membership instead -#[derive(Debug, Clone)] -pub enum OldIndexDataFilter { - /// Keeps track of which fragments are still valid and which are no longer valid. - /// - /// This is valid for address-style row IDs. - Fragments { - to_keep: RoaringBitmap, - to_remove: RoaringBitmap, - }, - /// Keep old rows whose row IDs are in this exact allow-list. - /// - /// This is required for stable row IDs, where row IDs are opaque and - /// should not be interpreted as encoded row addresses. - RowIds(RowAddrTreeMap), -} - -impl OldIndexDataFilter { - /// Build a boolean mask that keeps only row IDs selected by this filter. - pub fn filter_row_ids(&self, row_ids: &UInt64Array) -> BooleanArray { - match self { - Self::Fragments { to_keep, .. } => row_ids - .iter() - .map(|id| id.map(|id| to_keep.contains((id >> 32) as u32))) - .collect(), - Self::RowIds(valid_row_ids) => row_ids - .iter() - .map(|id| id.map(|id| valid_row_ids.contains(id))) - .collect(), - } - } - - /// Apply this filter in place to a set of existing (old) row ids/addresses, - /// retaining only the rows the filter selects to keep. Used by index types - /// that merge old postings directly (e.g. bitmap) instead of re-scanning a - /// row-id array through [`Self::filter_row_ids`]. - pub fn retain_old_rows(&self, rows: &mut RowAddrTreeMap) { - match self { - Self::Fragments { to_keep, .. } => rows.retain_fragments(to_keep.iter()), - Self::RowIds(valid_row_ids) => *rows &= valid_row_ids, - } - } -} - -impl UpdateCriteria { - pub fn requires_old_data(data_criteria: TrainingCriteria) -> Self { - Self { - requires_old_data: true, - data_criteria, - } - } - - pub fn only_new_data(data_criteria: TrainingCriteria) -> Self { - Self { - requires_old_data: false, - data_criteria, - } - } -} - /// Compute the lexicographically next prefix by incrementing the last character's code point. /// Returns None if no valid upper bound exists. /// @@ -1091,90 +690,6 @@ fn next_unicode_char(c: char) -> Option { char::from_u32(next_cp) } -/// A trait for a scalar index, a structure that can determine row ids that satisfy scalar queries -#[async_trait] -pub trait ScalarIndex: Send + Sync + std::fmt::Debug + Index + DeepSizeOf { - /// Search the scalar index - /// - /// Returns all row ids that satisfy the query, these row ids are not necessarily ordered - async fn search( - &self, - query: &dyn AnyQuery, - metrics: &dyn MetricsCollector, - ) -> Result; - - /// Returns true if this index reports matches as physical row addresses - /// (`fragment_id << 32 | offset`) rather than row ids - /// - /// Address-domain indices (e.g. zone map, bloom filter) are built over the - /// `_rowaddr` column. On a dataset with stable row ids the address and - /// row-id domains diverge, so these results must be translated back to row - /// ids (via the per-fragment row-id sequences, known only at the dataset - /// layer) before they are combined with row-id results or handed to the - /// scan. The default (row-id domain) needs no translation. - fn results_are_row_addresses(&self) -> bool { - false - } - - /// Returns true if the remap operation is supported - fn can_remap(&self) -> bool; - - /// Remap the row ids, creating a new remapped version of this index in `dest_store` - async fn remap( - &self, - mapping: &RowAddrRemap, - dest_store: &dyn IndexStore, - ) -> Result; - - /// Add the new data into the index, creating an updated version of the index in `dest_store` - /// - /// If `old_data_filter` is provided, old index data will be filtered before - /// merge according to the chosen filter mode. - async fn update( - &self, - new_data: SendableRecordBatchStream, - dest_store: &dyn IndexStore, - old_data_filter: Option, - ) -> Result; - - /// Returns the criteria that will be used to update the index - fn update_criteria(&self) -> UpdateCriteria; - - /// Derive the index parameters from the current index - /// - /// This returns a ScalarIndexParams that can be used to recreate an index - /// with the same configuration on another dataset. - fn derive_index_params(&self) -> Result; - - /// Global `[min, max]` of the indexed column from index metadata, without a - /// scan, or `None` if this index type cannot supply a sound bound. When - /// `Some`, the range is a superset of live values (conservative under - /// deletes): safe to prune with, not guaranteed tight. - fn value_range(&self) -> Option<(ScalarValue, ScalarValue)> { - None - } -} - -/// Abstraction over any type that can remap row IDs during index loading. -/// -/// This decouples scalar index plugins from the table-level [`crate::frag_reuse::FragReuseIndex`] -/// type. [`crate::frag_reuse::FragReuseIndex`] implements this trait, but callers may also -/// supply custom implementations for testing or other remapping strategies. -pub trait RowIdRemapper: Send + Sync + std::fmt::Debug { - /// Remap a single row id. Returns `None` if the row was deleted. - fn remap_row_id(&self, row_id: u64) -> Option; - /// Remap all addresses in a [`RowAddrTreeMap`], dropping deleted rows. - fn remap_row_addrs_tree_map(&self, row_addrs: &RowAddrTreeMap) -> RowAddrTreeMap; - /// Remap all row ids in a [`RoaringTreemap`], dropping deleted rows. - fn remap_row_ids_roaring_tree_map(&self, row_ids: &RoaringTreemap) -> RoaringTreemap; - /// Remap the row-id column at `row_id_idx` inside `batch`, dropping deleted rows. - fn remap_row_ids_record_batch( - &self, - batch: RecordBatch, - row_id_idx: usize, - ) -> Result; -} - #[cfg(test)] mod tests { use super::*; diff --git a/rust/lance-index/src/scalar/bloomfilter.rs b/rust/lance-index/src/scalar/bloomfilter.rs index 4a05cc8fc69..1f1ffea7ebd 100644 --- a/rust/lance-index/src/scalar/bloomfilter.rs +++ b/rust/lance-index/src/scalar/bloomfilter.rs @@ -7,6 +7,7 @@ //! It is a space-efficient data structure that can be used to test whether an element is a member of a set. //! It's an inexact filter - they may include false positives that require rechecking. +use crate::pb; use crate::scalar::expression::{BloomFilterQueryParser, ScalarQueryParser}; use crate::scalar::registry::{ BasicTrainer, ScalarIndexPlugin, TrainingCriteria, TrainingOrdering, TrainingRequest, @@ -14,7 +15,6 @@ use crate::scalar::registry::{ use crate::scalar::{ BloomFilterQuery, BuiltinIndexType, CreatedIndex, IndexFile, ScalarIndexParams, UpdateCriteria, }; -use crate::{Any, pb}; use arrow_array::{Array, UInt64Array}; use arrow_schema::{DataType, Field}; use lance_arrow_stats::StatisticsAccumulator; @@ -23,6 +23,7 @@ use lance_core::utils::bloomfilter::sbbf::{Sbbf, SbbfBuilder}; use lance_core::utils::row_addr_remap::RowAddrRemap; use lance_select::RowAddrTreeMap; use serde::{Deserialize, Serialize}; +use std::any::Any; use std::collections::HashMap; use std::sync::LazyLock; diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index bd63f05edd5..8e9bd1d2ba9 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -2562,9 +2562,7 @@ pub async fn train_btree_index( Ok(vec![pages_file, lookup_file]) } -fn find_single_partition_files( - files: &[lance_table::format::IndexFile], -) -> Result> { +fn find_single_partition_files(files: &[super::IndexFile]) -> Result> { let lookup_files = files .iter() .filter_map(|file| { @@ -6285,7 +6283,7 @@ mod tests { #[tokio::test] async fn test_btree_index_state_reconstruct_applies_frag_reuse_index() { - use crate::frag_reuse::{FragReuseIndex, FragReuseIndexDetails}; + use crate::frag_reuse::{FragReuseIndex, FragReuseIndexDetails, FragReuseIndexHandle}; use std::collections::HashMap; use uuid::Uuid; @@ -6318,11 +6316,11 @@ mod tests { // Querying for value == 0 should now return row 5000, confirming reconstruct threaded // the FragReuseIndex through to the rebuilt BTreeIndex. let frag_reuse_index: Arc = - Arc::new(FragReuseIndex::new( + Arc::new(FragReuseIndexHandle(Arc::new(FragReuseIndex::new( Uuid::new_v4(), vec![HashMap::from([(0u64, Some(5000u64))])], FragReuseIndexDetails { versions: vec![] }, - )); + )))); let reconstructed = state .reconstruct( test_store.clone(), diff --git a/rust/lance-index/src/scalar/expression.rs b/rust/lance-index/src/scalar/expression.rs index cd76b68ac66..4c7fe2b59ff 100644 --- a/rust/lance-index/src/scalar/expression.rs +++ b/rust/lance-index/src/scalar/expression.rs @@ -1662,12 +1662,16 @@ impl std::fmt::Display for ScalarIndexExpr { } } -impl From for NullableIndexExprResult { - fn from(result: SearchResult) -> Self { - match result { - SearchResult::Exact(mask) => Self::exact(NullableRowAddrMask::AllowList(mask)), - SearchResult::AtMost(mask) => Self::at_most(NullableRowAddrMask::AllowList(mask)), - SearchResult::AtLeast(mask) => Self::at_least(NullableRowAddrMask::AllowList(mask)), +fn search_result_to_nullable(result: SearchResult) -> NullableIndexExprResult { + match result { + SearchResult::Exact(mask) => { + NullableIndexExprResult::exact(NullableRowAddrMask::AllowList(mask)) + } + SearchResult::AtMost(mask) => { + NullableIndexExprResult::at_most(NullableRowAddrMask::AllowList(mask)) + } + SearchResult::AtLeast(mask) => { + NullableIndexExprResult::at_least(NullableRowAddrMask::AllowList(mask)) } } } @@ -1707,7 +1711,7 @@ impl ScalarIndexExpr { .load_index(&search.column, &search.index_name, metrics) .await?; let search_result = index.search(search.query.as_ref(), metrics).await?; - let result: NullableIndexExprResult = search_result.into(); + let result = search_result_to_nullable(search_result); if index.results_are_row_addresses() { // Translate address-domain results to the row-id domain // before combining or scanning; otherwise stable-row-id diff --git a/rust/lance-index/src/scalar/lance_format.rs b/rust/lance-index/src/scalar/lance_format.rs index 0be5707d859..76ebd0cf0da 100644 --- a/rust/lance-index/src/scalar/lance_format.rs +++ b/rust/lance-index/src/scalar/lance_format.rs @@ -13,11 +13,8 @@ use lance_core::deepsize::DeepSizeOf; use lance_core::{Error, Result, cache::LanceCache}; use lance_encoding::decoder::{DecoderPlugins, FilterExpression}; use lance_encoding::version::LanceFileVersion; -use lance_file::previous::{ - reader::FileReader as PreviousFileReader, - writer::{FileWriter as PreviousFileWriter, ManifestProvider as PreviousManifestProvider}, -}; -use lance_file::reader::{self as current_reader, FileReaderOptions, ReaderProjection}; +use lance_file::previous::reader::FileReader as PreviousFileReader; +use lance_file::reader::{FileReader as CurrentFileReader, FileReaderOptions, ReaderProjection}; use lance_file::writer as current_writer; use lance_io::scheduler::{ScanScheduler, SchedulerConfig}; use lance_io::utils::CachedFileSize; @@ -130,34 +127,6 @@ impl LanceIndexStore { } } -#[async_trait] -impl IndexWriter for PreviousFileWriter { - async fn write_record_batch(&mut self, batch: RecordBatch) -> Result { - let offset = self.tell().await?; - self.write(&[batch]).await?; - Ok(offset as u64) - } - - async fn finish(&mut self) -> Result { - let summary = Self::finish(self).await?; - Ok(IndexFile { - path: String::new(), - size_bytes: summary.size_bytes, - }) - } - - async fn finish_with_metadata( - &mut self, - metadata: HashMap, - ) -> Result { - let summary = Self::finish_with_metadata(self, &metadata).await?; - Ok(IndexFile { - path: String::new(), - size_bytes: summary.size_bytes, - }) - } -} - struct LanceIndexWriter { path: String, inner: current_writer::FileWriter, @@ -198,10 +167,14 @@ impl IndexWriter for LanceIndexWriter { } } +/// Newtype wrapper to allow implementing IndexReader for PreviousFileReader (a foreign type) +struct PreviousIndexReader(PreviousFileReader); + #[async_trait] -impl IndexReader for PreviousFileReader { +impl IndexReader for PreviousIndexReader { async fn read_record_batch(&self, offset: u64, _batch_size: u64) -> Result { - self.read_batch(offset as i32, ReadBatchParams::RangeFull, self.schema()) + self.0 + .read_batch(offset as i32, ReadBatchParams::RangeFull, self.0.schema()) .await } @@ -211,36 +184,39 @@ impl IndexReader for PreviousFileReader { projection: Option<&[&str]>, ) -> Result { let projection = match projection { - Some(projection) => self.schema().project(projection)?, - None => self.schema().clone(), + Some(projection) => self.0.schema().project(projection)?, + None => self.0.schema().clone(), }; - self.read_range(range, &projection).await + self.0.read_range(range, &projection).await } async fn num_batches(&self, _batch_size: u64) -> u32 { - self.num_batches() as u32 + self.0.num_batches() as u32 } fn num_rows(&self) -> usize { - self.len() + self.0.len() } fn schema(&self) -> &lance_core::datatypes::Schema { - Self::schema(self) + PreviousFileReader::schema(&self.0) } } +/// Newtype wrapper to allow implementing IndexReader for CurrentFileReader (a foreign type) +struct CurrentIndexReader(CurrentFileReader); + #[async_trait] -impl IndexReader for current_reader::FileReader { +impl IndexReader for CurrentIndexReader { async fn read_record_batch(&self, offset: u64, batch_size: u64) -> Result { let start = offset * batch_size; let end = start + batch_size; - let end = end.min(self.num_rows()); + let end = end.min(self.0.num_rows()); self.read_range(start as usize..end as usize, None).await } async fn read_global_buffer(&self, n: u32) -> Result { - Self::read_global_buffer(self, n).await + CurrentFileReader::read_global_buffer(&self.0, n).await } async fn read_range( @@ -250,19 +226,20 @@ impl IndexReader for current_reader::FileReader { ) -> Result { if range.is_empty() { return Ok(RecordBatch::new_empty(Arc::new( - self.schema().as_ref().into(), + self.0.schema().as_ref().into(), ))); } let projection = if let Some(projection) = projection { ReaderProjection::from_column_names( - self.metadata().version(), - self.schema(), + self.0.metadata().version(), + self.0.schema(), projection, )? } else { - ReaderProjection::from_whole_schema(self.schema(), self.metadata().version()) + ReaderProjection::from_whole_schema(self.0.schema(), self.0.metadata().version()) }; let batches = self + .0 .read_stream_projected( ReadBatchParams::Range(range), u32::MAX, @@ -284,7 +261,7 @@ impl IndexReader for current_reader::FileReader { ) -> Result { let empty_batch = || { Ok(RecordBatch::new_empty(Arc::new( - self.schema().as_ref().into(), + self.0.schema().as_ref().into(), ))) }; if ranges.is_empty() { @@ -292,12 +269,12 @@ impl IndexReader for current_reader::FileReader { } let projection = if let Some(projection) = projection { ReaderProjection::from_column_names( - self.metadata().version(), - self.schema(), + self.0.metadata().version(), + self.0.schema(), projection, )? } else { - ReaderProjection::from_whole_schema(self.schema(), self.metadata().version()) + ReaderProjection::from_whole_schema(self.0.schema(), self.0.metadata().version()) }; // `DecodeBatchScheduler::schedule_ranges` requires sorted, // non-overlapping ranges; sort internally and permute the @@ -311,6 +288,7 @@ impl IndexReader for current_reader::FileReader { .collect(); let total_rows: u64 = sorted_ranges.iter().map(|r| r.end - r.start).sum(); let batches = self + .0 .read_stream_projected( ReadBatchParams::Ranges(sorted_ranges), (total_rows as u32).max(1), @@ -363,47 +341,48 @@ impl IndexReader for current_reader::FileReader { ) -> Result>> { if range.is_empty() { return Ok(Box::pin(lance_io::stream::RecordBatchStreamAdapter::new( - Arc::new(self.schema().as_ref().into()), + Arc::new(self.0.schema().as_ref().into()), futures::stream::empty(), ))); } let projection = if let Some(projection) = projection { ReaderProjection::from_column_names( - self.metadata().version(), - self.schema(), + self.0.metadata().version(), + self.0.schema(), projection, )? } else { - ReaderProjection::from_whole_schema(self.schema(), self.metadata().version()) + ReaderProjection::from_whole_schema(self.0.schema(), self.0.metadata().version()) }; - self.read_stream_projected( - ReadBatchParams::Range(range), - 4096, - 2, - projection, - FilterExpression::no_filter(), - ) - .await + self.0 + .read_stream_projected( + ReadBatchParams::Range(range), + 4096, + 2, + projection, + FilterExpression::no_filter(), + ) + .await } // V2 format has removed the row group concept, // so here we assume each batch is with 4096 rows. async fn num_batches(&self, batch_size: u64) -> u32 { - Self::num_rows(self).div_ceil(batch_size) as u32 + CurrentFileReader::num_rows(&self.0).div_ceil(batch_size) as u32 } fn num_rows(&self) -> usize { - Self::num_rows(self) as usize + CurrentFileReader::num_rows(&self.0) as usize } fn schema(&self) -> &lance_core::datatypes::Schema { - Self::schema(self) + CurrentFileReader::schema(&self.0) } fn file_size_bytes(&self) -> Option { // The manifest records each index file's size and passes it to the reader // at open, so it's already in metadata here (no extra I/O). - Some(self.metadata().file_size()) + Some(self.0.metadata().file_size()) } } @@ -464,7 +443,7 @@ impl IndexStore for LanceIndexStore { .scheduler .open_file_with_priority(&path, self.io_priority, &cached_size) .await?; - match current_reader::FileReader::try_open( + match CurrentFileReader::try_open( file_scheduler, None, Arc::::default(), @@ -473,7 +452,7 @@ impl IndexStore for LanceIndexStore { ) .await { - Ok(reader) => Ok(Arc::new(reader)), + Ok(reader) => Ok(Arc::new(CurrentIndexReader(reader))), Err(e) => { // If the error is a version conflict we can try to read the file with v1 reader if let Error::VersionConflict { .. } = e { @@ -484,7 +463,7 @@ impl IndexStore for LanceIndexStore { Some(&self.metadata_cache), ) .await?; - Ok(Arc::new(file_reader)) + Ok(Arc::new(PreviousIndexReader(file_reader))) } else { Err(e) } @@ -558,7 +537,14 @@ impl IndexStore for LanceIndexStore { } async fn list_files_with_sizes(&self) -> Result> { - list_index_files_with_sizes(&self.object_store, &self.index_dir).await + let files = list_index_files_with_sizes(&self.object_store, &self.index_dir).await?; + Ok(files + .into_iter() + .map(|f| IndexFile { + path: f.path, + size_bytes: f.size_bytes, + }) + .collect()) } } diff --git a/rust/lance-index/src/scalar/registry.rs b/rust/lance-index/src/scalar/registry.rs index 39dfff2c2c8..fb7a460e9ab 100644 --- a/rust/lance-index/src/scalar/registry.rs +++ b/rust/lance-index/src/scalar/registry.rs @@ -18,46 +18,11 @@ use crate::progress::IndexBuildProgress; use crate::registry::IndexPluginRegistry; use crate::scalar::RowIdRemapper; use crate::scalar::{CreatedIndex, IndexStore, ScalarIndex, expression::ScalarQueryParser}; +// Re-export training types that were previously defined here +pub use crate::scalar::{TrainingCriteria, TrainingOrdering}; pub const VALUE_COLUMN_NAME: &str = "value"; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum TrainingOrdering { - /// The input will arrive sorted by the value column in ascending order - Values, - /// The input will arrive sorted by the address column in ascending order - Addresses, - /// The input will arrive in an arbitrary order - None, -} - -#[derive(Debug, Clone)] -pub struct TrainingCriteria { - pub ordering: TrainingOrdering, - pub needs_row_ids: bool, - pub needs_row_addrs: bool, -} - -impl TrainingCriteria { - pub fn new(ordering: TrainingOrdering) -> Self { - Self { - ordering, - needs_row_ids: false, - needs_row_addrs: false, - } - } - - pub fn with_row_id(mut self) -> Self { - self.needs_row_ids = true; - self - } - - pub fn with_row_addr(mut self) -> Self { - self.needs_row_addrs = true; - self - } -} - /// A trait object for plugin-specific training parameters and data requirements. /// /// Returned by [`BasicTrainer::new_training_request`]. The caller uses diff --git a/rust/lance-index/src/scalar/rtree/sort/hilbert_sort.rs b/rust/lance-index/src/scalar/rtree/sort/hilbert_sort.rs index e6c10a20575..a8256c659c2 100644 --- a/rust/lance-index/src/scalar/rtree/sort/hilbert_sort.rs +++ b/rust/lance-index/src/scalar/rtree/sort/hilbert_sort.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use crate::Result; use crate::scalar::rtree::sort::Sorter; use arrow_array::{ArrayRef, UInt32Array}; use arrow_schema::{ArrowError, DataType as ArrowDataType, Field as ArrowField, Field}; @@ -19,6 +18,7 @@ use datafusion_physical_expr::expressions::Column as DFColumn; use datafusion_physical_expr::{PhysicalExpr, ScalarFunctionExpr}; use geoarrow_array::array::from_arrow_array; use geoarrow_array::{GeoArrowArray, GeoArrowArrayAccessor}; +use lance_core::Result; use lance_datafusion::exec::{LanceExecutionOptions, OneShotExec, execute_plan}; use lance_geo::bbox::{BoundingBox, bounding_box}; use std::any::Any; diff --git a/rust/lance-index/src/scalar/zonemap.rs b/rust/lance-index/src/scalar/zonemap.rs index 766a563593d..65d58b49f81 100644 --- a/rust/lance-index/src/scalar/zonemap.rs +++ b/rust/lance-index/src/scalar/zonemap.rs @@ -12,7 +12,6 @@ //! false positives that require rechecking. //! //! -use crate::Any; use crate::pbold; use crate::scalar::expression::{SargableQueryParser, ScalarQueryParser}; use crate::scalar::registry::{ @@ -26,6 +25,7 @@ use lance_arrow_stats::StatisticsAccumulator; use lance_core::cache::{LanceCache, WeakLanceCache}; use lance_core::utils::row_addr_remap::RowAddrRemap; use serde::{Deserialize, Serialize}; +use std::any::Any; use std::sync::LazyLock; use arrow_array::{ diff --git a/rust/lance-index/src/vector/flat/transform.rs b/rust/lance-index/src/vector/flat/transform.rs index 75a465ce262..f9fdca0819c 100644 --- a/rust/lance-index/src/vector/flat/transform.rs +++ b/rust/lance-index/src/vector/flat/transform.rs @@ -26,7 +26,7 @@ impl FlatTransformer { impl Transformer for FlatTransformer { #[instrument(name = "FlatTransformer::transform", level = "debug", skip_all)] - fn transform(&self, batch: &RecordBatch) -> crate::Result { + fn transform(&self, batch: &RecordBatch) -> lance_core::Result { let input_arr = batch .column_by_name(&self.input_column) .ok_or(Error::index(format!( diff --git a/rust/lance-index/src/vector/hnsw/builder.rs b/rust/lance-index/src/vector/hnsw/builder.rs index cc1ac1abf81..a020d011fba 100644 --- a/rust/lance-index/src/vector/hnsw/builder.rs +++ b/rust/lance-index/src/vector/hnsw/builder.rs @@ -1330,7 +1330,6 @@ mod tests { use rstest::rstest; use super::HnswGraph; - use crate::scalar::IndexWriter; use crate::vector::storage::{DistCalculator, VectorStore}; use crate::vector::v3::subindex::IvfSubIndex; use crate::vector::{ @@ -1375,7 +1374,7 @@ mod tests { .unwrap(); let batch = builder.to_batch().unwrap(); let metadata = batch.schema_ref().metadata().clone(); - writer.write_record_batch(batch).await.unwrap(); + writer.write(&[batch]).await.unwrap(); writer.finish_with_metadata(&metadata).await.unwrap(); let reader = PreviousFileReader::try_new_self_described(&object_store, &path, None) @@ -1436,7 +1435,7 @@ mod tests { .unwrap(); let batch = builder.to_batch().unwrap(); let metadata = batch.schema_ref().metadata().clone(); - writer.write_record_batch(batch).await.unwrap(); + writer.write(&[batch]).await.unwrap(); writer.finish_with_metadata(&metadata).await.unwrap(); let reader = PreviousFileReader::try_new_self_described(&object_store, &path, None) diff --git a/rust/lance-index/src/vector/kmeans.rs b/rust/lance-index/src/vector/kmeans.rs index b11fb70bed0..07dc067b263 100644 --- a/rust/lance-index/src/vector/kmeans.rs +++ b/rust/lance-index/src/vector/kmeans.rs @@ -45,7 +45,7 @@ use { }; use crate::vector::utils::SimpleIndex; -use crate::{Error, Result}; +use lance_core::{Error, Result}; /// KMean initialization method. #[derive(Debug, PartialEq)] diff --git a/rust/lance-namespace-impls/src/dir/manifest.rs b/rust/lance-namespace-impls/src/dir/manifest.rs index ab916821d29..1e52fc85bfe 100644 --- a/rust/lance-namespace-impls/src/dir/manifest.rs +++ b/rust/lance-namespace-impls/src/dir/manifest.rs @@ -36,7 +36,9 @@ use lance_index::progress::noop_progress; use lance_index::registry::IndexPluginRegistry; use lance_index::scalar::lance_format::LanceIndexStore; use lance_index::scalar::registry::VALUE_COLUMN_NAME; -use lance_index::scalar::{BuiltinIndexType, CreatedIndex, ScalarIndexParams}; +use lance_index::scalar::{ + BuiltinIndexType, CreatedIndex, ScalarIndexParams, index_files_to_table, +}; use lance_io::object_store::{ObjectStore, ObjectStoreParams}; use lance_io::stream::RecordBatchStream as LanceRecordBatchStream; use lance_namespace::LanceNamespace; @@ -1279,7 +1281,7 @@ impl ManifestNamespace { index_version: trained_index.created_index.index_version as i32, created_at: None, base_id: None, - files: Some(trained_index.created_index.files), + files: Some(index_files_to_table(trained_index.created_index.files)), }) } diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index a258acd8985..f8f06a59fc9 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -2236,6 +2236,7 @@ mod tests { use lance_datagen::Dimension; use lance_file::version::LanceFileVersion; use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME; + use lance_index::frag_reuse::FragReuseIndexHandle; use lance_index::scalar::{ BuiltinIndexType, FullTextSearchQuery, InvertedIndexParams, ScalarIndexParams, }; @@ -3432,7 +3433,9 @@ mod tests { open_frag_reuse_index(frag_reuse_index_meta.uuid, frag_reuse_details.as_ref()) .await .unwrap(); - let stats = frag_reuse_index.statistics().unwrap(); + let stats = FragReuseIndexHandle(Arc::new(frag_reuse_index.clone())) + .statistics() + .unwrap(); assert_eq!( serde_json::to_string(&stats).unwrap(), dataset diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index b960405b3ee..a1d38f0e4b1 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -25,8 +25,8 @@ use lance_file::previous::reader::FileReader as PreviousFileReader; use lance_file::reader::FileReaderOptions; use lance_index::INDEX_METADATA_SCHEMA_KEY; pub use lance_index::IndexParams; -use lance_index::frag_reuse::{FRAG_REUSE_INDEX_NAME, FragReuseIndex}; -use lance_index::mem_wal::{MEM_WAL_INDEX_NAME, MemWalIndex}; +use lance_index::frag_reuse::{FRAG_REUSE_INDEX_NAME, FragReuseIndex, FragReuseIndexHandle}; +use lance_index::mem_wal::{MEM_WAL_INDEX_NAME, MemWalIndex, MemWalIndexHandle}; use lance_index::optimize::OptimizeOptions; use lance_index::pb::index::Implementation; pub use lance_index::progress::{IndexBuildProgress, NoopIndexBuildProgress}; @@ -34,7 +34,7 @@ use lance_index::scalar::expression::{IndexInformationProvider, MultiQueryParser use lance_index::scalar::inverted::{InvertedIndex, InvertedIndexPlugin}; use lance_index::scalar::lance_format::LanceIndexStore; use lance_index::scalar::registry::{TrainingCriteria, TrainingOrdering}; -use lance_index::scalar::{CreatedIndex, ScalarIndex}; +use lance_index::scalar::{CreatedIndex, ScalarIndex, index_files_to_table, table_files_to_index}; use lance_index::vector::bq::builder::RabitQuantizer; use lance_index::vector::flat::index::{FlatBinQuantizer, FlatIndex, FlatQuantizer}; use lance_index::vector::hnsw::HNSW; @@ -827,7 +827,7 @@ pub(crate) async fn remap_index( ) .unwrap(), index_version, - files, + files: table_files_to_index(files), } } _ => { @@ -843,7 +843,7 @@ pub(crate) async fn remap_index( new_id, index_details: created_index.index_details, index_version: created_index.index_version, - files: Some(created_index.files), + files: Some(index_files_to_table(created_index.files)), })) } @@ -1814,7 +1814,7 @@ async fn index_statistics_frag_reuse(ds: &Dataset) -> Result { .open_frag_reuse_index(&NoOpMetricsCollector) .await? .expect("FragmentReuse index does not exist"); - serialize_index_statistics(&index.statistics()?) + serialize_index_statistics(&FragReuseIndexHandle(index).statistics()?) } async fn index_statistics_mem_wal(ds: &Dataset) -> Result { @@ -1822,7 +1822,7 @@ async fn index_statistics_mem_wal(ds: &Dataset) -> Result { .open_mem_wal_index(&NoOpMetricsCollector) .await? .expect("MemWal index does not exist"); - serialize_index_statistics(&index.statistics()?) + serialize_index_statistics(&MemWalIndexHandle(index).statistics()?) } async fn index_statistics_scalar( @@ -2082,7 +2082,7 @@ impl DatasetIndexInternalExt for Dataset { let frag_reuse_cache_key = FragReuseIndexCacheKey::new(uuid, frag_reuse_uuid.as_ref()); if let Some(index) = self.index_cache.get_with_key(&frag_reuse_cache_key).await { - return Ok(index.as_index()); + return Ok(Arc::new(FragReuseIndexHandle(index)).as_index()); } // Sometimes we want to open an index and we don't care if it is a scalar or vector index. diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index edcd4357ae4..3a6b0e42a44 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -11,8 +11,8 @@ use lance_index::{ optimize::OptimizeOptions, progress::NoopIndexBuildProgress, scalar::{ - CreatedIndex, OldIndexDataFilter, ScalarIndex, inverted::InvertedIndex, - lance_format::LanceIndexStore, + CreatedIndex, OldIndexDataFilter, ScalarIndex, index_files_to_table, + inverted::InvertedIndex, lance_format::LanceIndexStore, table_files_to_index, }, }; use lance_select::{RowAddrTreeMap, RowSetOps}; @@ -592,7 +592,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( CreatedIndex { index_details: vector_index_details_default(), index_version: lance_index::IndexType::Vector.version() as u32, - files, + files: table_files_to_index(files), }, )) } else { @@ -655,7 +655,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( // index_version <= our max supported version, so we can safely // write the current library's version for this index type. index_version: lance_index::IndexType::Vector.version() as u32, - files, + files: table_files_to_index(files), }, )) } @@ -734,7 +734,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( new_fragment_bitmap: dataset.fragment_bitmap.as_ref().clone(), new_index_version: created_index.index_version as i32, new_index_details: created_index.index_details, - files: created_index.files, + files: index_files_to_table(created_index.files), })); } @@ -850,7 +850,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( new_fragment_bitmap, new_index_version: created_index.index_version as i32, new_index_details: created_index.index_details, - files: created_index.files, + files: index_files_to_table(created_index.files), })) } diff --git a/rust/lance/src/index/create.rs b/rust/lance/src/index/create.rs index c50477bb88c..7e5fe98bbf0 100644 --- a/rust/lance/src/index/create.rs +++ b/rust/lance/src/index/create.rs @@ -24,7 +24,10 @@ use lance_index::progress::{IndexBuildProgress, NoopIndexBuildProgress}; use lance_index::{IndexParams, IndexType, scalar::CreatedIndex}; use lance_index::{ metrics::NoOpMetricsCollector, - scalar::{LANCE_SCALAR_INDEX, ScalarIndexParams, inverted::tokenizer::InvertedIndexParams}, + scalar::{ + LANCE_SCALAR_INDEX, ScalarIndexParams, index_files_to_table, + inverted::tokenizer::InvertedIndexParams, table_files_to_index, + }, }; use lance_table::format::{IndexMetadata, list_index_files_with_sizes}; use std::{collections::HashMap, future::IntoFuture, sync::Arc}; @@ -432,7 +435,7 @@ impl<'a> CreateIndexBuilder<'a> { CreatedIndex { index_details: vector_index_details(vec_params), index_version, - files, + files: table_files_to_index(files), } } // Can't use if let Some(...) here because it's not stable yet. @@ -471,7 +474,7 @@ impl<'a> CreateIndexBuilder<'a> { CreatedIndex { index_details: vector_index_details_default(), index_version: self.index_type.version() as u32, - files, + files: table_files_to_index(files), } } (IndexType::FragmentReuse, _) => { @@ -504,7 +507,7 @@ impl<'a> CreateIndexBuilder<'a> { index_version: created_index.index_version as i32, created_at: Some(chrono::Utc::now()), base_id: None, - files: Some(created_index.files), + files: Some(index_files_to_table(created_index.files)), }) } .boxed() @@ -671,7 +674,7 @@ impl<'a> CreateIndexBuilder<'a> { index_version: created_index.index_version as i32, created_at: Some(chrono::Utc::now()), base_id: None, - files: Some(created_index.files), + files: Some(index_files_to_table(created_index.files)), }; let segments = vec![metadata.into_index_segment()?]; let new_indices = @@ -740,7 +743,7 @@ impl<'a> CreateIndexBuilder<'a> { index_version: created_index.index_version as i32, created_at: Some(chrono::Utc::now()), base_id: None, - files: Some(created_index.files), + files: Some(index_files_to_table(created_index.files)), }); } @@ -2892,7 +2895,7 @@ mod tests { let mut legacy_segment = segment.clone(); legacy_segment.uuid = legacy_uuid; legacy_segment.index_version = LABEL_LIST_NULLS_MIN_VERSION; - legacy_segment.files = Some(vec![legacy_file]); + legacy_segment.files = Some(index_files_to_table(vec![legacy_file])); let err = dataset .merge_existing_index_segments(vec![legacy_segment]) diff --git a/rust/lance/src/index/scalar.rs b/rust/lance/src/index/scalar.rs index 79e6eed44c9..d31b96c9202 100644 --- a/rust/lance/src/index/scalar.rs +++ b/rust/lance/src/index/scalar.rs @@ -35,6 +35,7 @@ use lance_core::datatypes::Field; use lance_core::utils::tracing::{IO_TYPE_OPEN_SCALAR, TRACE_IO_EVENTS}; use lance_core::{Error, ROW_ADDR, ROW_ID, Result}; use lance_datafusion::exec::LanceExecutionOptions; +use lance_index::frag_reuse::FragReuseIndexHandle; use lance_index::metrics::{MetricsCollector, NoOpMetricsCollector}; use lance_index::pbold::{ BTreeIndexDetails, BitmapIndexDetails, InvertedIndexDetails, LabelListIndexDetails, @@ -460,7 +461,7 @@ pub async fn open_scalar_index( .for_index(&index.uuid, frag_reuse_index.as_ref().map(|f| &f.uuid)); let frag_reuse_index: Option> = - frag_reuse_index.map(|f| f as Arc); + frag_reuse_index.map(|f| Arc::new(FragReuseIndexHandle(f)) as Arc); // Runs only on a cold miss, and at most once even under concurrent opens // (the plugin coalesces). The compat check lives here because a warm hit was diff --git a/rust/lance/src/index/scalar/bitmap.rs b/rust/lance/src/index/scalar/bitmap.rs index 2eb5702ee28..84c8b6a8b91 100644 --- a/rust/lance/src/index/scalar/bitmap.rs +++ b/rust/lance/src/index/scalar/bitmap.rs @@ -3,6 +3,7 @@ use lance_index::metrics::NoOpMetricsCollector; use lance_index::scalar::bitmap::BitmapIndex; +use lance_index::scalar::index_files_to_table; use lance_index::scalar::lance_format::LanceIndexStore; use lance_table::format::IndexMetadata; use roaring::RoaringBitmap; @@ -70,7 +71,7 @@ pub(in crate::index) async fn merge_segments( index_version: created_index.index_version as i32, created_at: Some(chrono::Utc::now()), base_id: None, - files: Some(created_index.files), + files: Some(index_files_to_table(created_index.files)), ..segments[0].clone() }) } diff --git a/rust/lance/src/index/scalar/btree.rs b/rust/lance/src/index/scalar/btree.rs index 4339b8c183b..7089997721c 100644 --- a/rust/lance/src/index/scalar/btree.rs +++ b/rust/lance/src/index/scalar/btree.rs @@ -15,7 +15,7 @@ use lance_index::pbold::BTreeIndexDetails; use lance_index::scalar::btree::BTreeIndex; use lance_index::scalar::lance_format::LanceIndexStore; use lance_index::scalar::registry::VALUE_COLUMN_NAME; -use lance_index::scalar::{CreatedIndex, OldIndexDataFilter}; +use lance_index::scalar::{CreatedIndex, OldIndexDataFilter, index_files_to_table}; use lance_table::format::IndexMetadata; use uuid::Uuid; @@ -161,6 +161,6 @@ pub(crate) async fn merge_segments( index_version: created_index.index_version as i32, created_at: Some(chrono::Utc::now()), base_id: None, - files: Some(created_index.files), + files: Some(index_files_to_table(created_index.files)), }) } diff --git a/rust/lance/src/index/scalar/fmindex.rs b/rust/lance/src/index/scalar/fmindex.rs index 32684ebf9ab..9c1baadbbb6 100644 --- a/rust/lance/src/index/scalar/fmindex.rs +++ b/rust/lance/src/index/scalar/fmindex.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use lance_index::scalar::index_files_to_table; use lance_table::format::IndexMetadata; use roaring::RoaringBitmap; use std::sync::Arc; @@ -76,7 +77,7 @@ pub(in crate::index) async fn merge_segments( index_version: created_index.index_version as i32, created_at: Some(chrono::Utc::now()), base_id: None, - files: Some(created_index.files), + files: Some(index_files_to_table(created_index.files)), ..segments[0].clone() }); } @@ -111,7 +112,7 @@ pub(in crate::index) async fn merge_segments( index_version: created_index.index_version as i32, created_at: Some(chrono::Utc::now()), base_id: None, - files: Some(created_index.files), + files: Some(index_files_to_table(created_index.files)), ..segments[0].clone() }) } diff --git a/rust/lance/src/index/scalar/inverted.rs b/rust/lance/src/index/scalar/inverted.rs index 000d2c3139c..b41dfa562b9 100644 --- a/rust/lance/src/index/scalar/inverted.rs +++ b/rust/lance/src/index/scalar/inverted.rs @@ -11,6 +11,7 @@ use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use lance_core::ROW_ID; use lance_index::metrics::NoOpMetricsCollector; use lance_index::pbold::InvertedIndexDetails; +use lance_index::scalar::index_files_to_table; use lance_index::scalar::inverted::InvertedIndex; use lance_index::scalar::lance_format::LanceIndexStore; use lance_index::scalar::registry::VALUE_COLUMN_NAME; @@ -137,7 +138,7 @@ pub(crate) async fn merge_segments( index_version: created_index.index_version as i32, created_at: Some(chrono::Utc::now()), base_id: None, - files: Some(created_index.files), + files: Some(index_files_to_table(created_index.files)), ..segments[0].clone() }) } diff --git a/rust/lance/src/index/scalar/label_list.rs b/rust/lance/src/index/scalar/label_list.rs index 27bc49643bb..884346f0d6b 100644 --- a/rust/lance/src/index/scalar/label_list.rs +++ b/rust/lance/src/index/scalar/label_list.rs @@ -3,6 +3,7 @@ use lance_index::metrics::NoOpMetricsCollector; use lance_index::scalar::IndexStore; +use lance_index::scalar::index_files_to_table; use lance_index::scalar::label_list::{ BITMAP_LOOKUP_NAME, LABEL_LIST_NULLS_METADATA_KEY, LABEL_LIST_NULLS_MIN_VERSION, LabelListIndex, }; @@ -142,7 +143,7 @@ pub(in crate::index) async fn merge_segments( index_version: created_index.index_version as i32, created_at: Some(chrono::Utc::now()), base_id: None, - files: Some(created_index.files), + files: Some(index_files_to_table(created_index.files)), ..segments[0].clone() }) } diff --git a/rust/lance/src/index/scalar/zonemap.rs b/rust/lance/src/index/scalar/zonemap.rs index 0cbd98f2c40..a3524b4955e 100644 --- a/rust/lance/src/index/scalar/zonemap.rs +++ b/rust/lance/src/index/scalar/zonemap.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use lance_index::metrics::NoOpMetricsCollector; +use lance_index::scalar::index_files_to_table; use lance_index::scalar::lance_format::LanceIndexStore; use lance_index::scalar::zonemap::ZoneMapIndex; use lance_table::format::IndexMetadata; @@ -80,7 +81,7 @@ pub(in crate::index) async fn merge_segments( index_version: created_index.index_version as i32, created_at: Some(chrono::Utc::now()), base_id: None, - files: Some(created_index.files), + files: Some(index_files_to_table(created_index.files)), ..segments[0].clone() }) } diff --git a/rust/lance/src/index/vector/ivf/io.rs b/rust/lance/src/index/vector/ivf/io.rs index f8f713a07af..a49398a0750 100644 --- a/rust/lance/src/index/vector/ivf/io.rs +++ b/rust/lance/src/index/vector/ivf/io.rs @@ -25,7 +25,6 @@ use lance_core::utils::tokio::{get_num_compute_intensive_cpus, spawn_cpu}; use lance_file::previous::reader::FileReader as PreviousFileReader; use lance_file::previous::writer::FileWriter as PreviousFileWriter; use lance_index::metrics::NoOpMetricsCollector; -use lance_index::scalar::IndexWriter; use lance_index::vector::hnsw::HNSW; use lance_index::vector::hnsw::{HnswMetadata, builder::HnswBuildParams}; use lance_index::vector::ivf::storage::IvfModel; @@ -584,7 +583,7 @@ async fn build_and_write_hnsw( ) -> Result { let batch = params.build(vectors, distance_type).await?.to_batch()?; let metadata = batch.schema_ref().metadata().clone(); - writer.write_record_batch(batch).await?; + writer.write(&[batch]).await?; Ok(writer.finish_with_metadata(&metadata).await?.num_rows as usize) } @@ -597,7 +596,7 @@ async fn build_and_write_pq_storage( ) -> Result<()> { let storage = spawn_cpu(move || build_pq_storage(metric_type, row_ids, code_array, pq)).await?; - writer.write_record_batch(storage.batch().clone()).await?; + writer.write(&[storage.batch().clone()]).await?; writer.finish().await?; Ok(()) } diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index 1301871b560..732e6befd90 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -2072,6 +2072,7 @@ mod tests { use lance_file::reader::{FileReader, FileReaderOptions}; use lance_file::writer::FileWriter; use lance_index::IndexType; + use lance_index::optimize::OptimizeOptions; use lance_index::progress::IndexBuildProgress; use lance_index::vector::DIST_COL; use lance_index::vector::hnsw::builder::HnswBuildParams; @@ -2086,7 +2087,6 @@ mod tests { storage::STORAGE_METADATA_KEY, }; use lance_index::{INDEX_AUXILIARY_FILE_NAME, metrics::NoOpMetricsCollector}; - use lance_index::{optimize::OptimizeOptions, scalar::IndexReader}; use lance_io::{ object_store::ObjectStore, scheduler::{ScanScheduler, SchedulerConfig}, @@ -5284,9 +5284,22 @@ mod tests { // Rewrite auxiliary file with PQ codebook inlined into schema metadata. let mut metadata = reader.schema().metadata.clone(); - let batch = reader - .read_range(0..reader.num_rows() as usize, None) + let projection = lance_file::reader::ReaderProjection::from_whole_schema( + reader.schema(), + reader.metadata().version(), + ); + let batches = reader + .read_stream_projected( + lance_io::ReadBatchParams::RangeFull, + u32::MAX, + u32::MAX, + projection, + lance_encoding::decoder::FilterExpression::no_filter(), + ) .await?; + use futures::TryStreamExt as _; + let batches = batches.try_collect::>().await?; + let batch = arrow::compute::concat_batches(&batches[0].schema(), &batches)?; let new_aux_path = new_dir.clone().join(INDEX_AUXILIARY_FILE_NAME); let mut writer = FileWriter::try_new( obj_store.create(&new_aux_path).await?,