mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-21 12:35:42 +00:00
feat: index builders for NGram, BloomFilter, RTree (#4205)
Add index builders in the lancedb library for NGram, BloomFilter, and RTree indexes.
This commit is contained in:
Colin Patrick McCabe
Generated
+297
@@ -3151,6 +3151,16 @@ version = "0.1.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e1d926b4d407d372f141f93bb444696142c29d32962ccbd3531117cf3aa0bfa9"
|
||||
|
||||
[[package]]
|
||||
name = "earcutr"
|
||||
version = "0.4.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "79127ed59a85d7687c409e9978547cffb7dc79675355ed22da6b66fd5f6ead01"
|
||||
dependencies = [
|
||||
"itertools 0.11.0",
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ecdsa"
|
||||
version = "0.14.8"
|
||||
@@ -3450,6 +3460,12 @@ dependencies = [
|
||||
"rand_distr 0.5.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "float_next_after"
|
||||
version = "1.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8bf7cc16383c4b8d58b9905a8509f02926ce3058053c056376248d958c9df1e8"
|
||||
|
||||
[[package]]
|
||||
name = "fnv"
|
||||
version = "1.0.7"
|
||||
@@ -3776,6 +3792,129 @@ dependencies = [
|
||||
"version_check",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "geo"
|
||||
version = "0.31.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2fc1a1678e54befc9b4bcab6cd43b8e7f834ae8ea121118b0fd8c42747675b4a"
|
||||
dependencies = [
|
||||
"earcutr",
|
||||
"float_next_after",
|
||||
"geo-types",
|
||||
"geographiclib-rs",
|
||||
"i_overlay",
|
||||
"log",
|
||||
"num-traits",
|
||||
"robust",
|
||||
"rstar",
|
||||
"spade",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "geo-traits"
|
||||
version = "0.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2e7c353d12a704ccfab1ba8bfb1a7fe6cb18b665bf89d37f4f7890edcd260206"
|
||||
dependencies = [
|
||||
"geo-types",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "geo-types"
|
||||
version = "0.7.19"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "94776032c45f950d30a13af6113c2ad5625316c9abfbccee4dd5a6695f8fe0f5"
|
||||
dependencies = [
|
||||
"approx",
|
||||
"num-traits",
|
||||
"rayon",
|
||||
"rstar",
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "geoarrow-array"
|
||||
version = "0.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dafe7b7de3fab1a8b7099fd6a6434ca955fa65065f9c19f0f8a133693f3c2b0e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
"arrow-schema",
|
||||
"geo-traits",
|
||||
"geoarrow-schema",
|
||||
"num-traits",
|
||||
"wkb",
|
||||
"wkt",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "geoarrow-expr-geo"
|
||||
version = "0.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8e4a62ac19c86827c6ec81ea584594b3ee96db5a8119b9774d3466c6b373c434"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
"geo",
|
||||
"geo-traits",
|
||||
"geoarrow-array",
|
||||
"geoarrow-schema",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "geoarrow-schema"
|
||||
version = "0.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4d4a7edb2a1d87024a93805332a9c8184a0354836271d42c0d18cf628a5e3cd0"
|
||||
dependencies = [
|
||||
"arrow-schema",
|
||||
"geo-traits",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 1.0.69",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "geodatafusion"
|
||||
version = "0.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fecbdd00d0fff2b04635c1b1e4129c217908f0c2d17539e0a2275308afce2552"
|
||||
dependencies = [
|
||||
"arrow-arith",
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
"arrow-schema",
|
||||
"datafusion",
|
||||
"geo",
|
||||
"geo-traits",
|
||||
"geoarrow-array",
|
||||
"geoarrow-expr-geo",
|
||||
"geoarrow-schema",
|
||||
"geohash",
|
||||
"thiserror 1.0.69",
|
||||
"wkt",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "geographiclib-rs"
|
||||
version = "0.2.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c5a7f08910fd98737a6eda7568e7c5e645093e073328eeef49758cfe8b0489c7"
|
||||
dependencies = [
|
||||
"libm",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "geohash"
|
||||
version = "0.13.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7f58890382f70caccc5fa388981f7ac80c913795042afce9f3e065695d8f7464"
|
||||
dependencies = [
|
||||
"geo-types",
|
||||
"libm",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "getrandom"
|
||||
version = "0.2.17"
|
||||
@@ -3954,6 +4093,15 @@ dependencies = [
|
||||
"zerocopy",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hash32"
|
||||
version = "0.3.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "47d60b12902ba28e2730cd37e95b8c9223af2808df9e902d4df49588d1470606"
|
||||
dependencies = [
|
||||
"byteorder",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hashbrown"
|
||||
version = "0.12.3"
|
||||
@@ -4010,6 +4158,16 @@ version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0049b265b7f201ca9ab25475b22b47fe444060126a51abe00f77d986fc5cc52e"
|
||||
|
||||
[[package]]
|
||||
name = "heapless"
|
||||
version = "0.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad"
|
||||
dependencies = [
|
||||
"hash32",
|
||||
"stable_deref_trait",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "heck"
|
||||
version = "0.4.1"
|
||||
@@ -4327,6 +4485,49 @@ dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "i_float"
|
||||
version = "1.15.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "010025c2c532c8d82e42d0b8bb5184afa449fa6f06c709ea9adcb16c49ae405b"
|
||||
dependencies = [
|
||||
"libm",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "i_key_sort"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9190f86706ca38ac8add223b2aed8b1330002b5cdbbce28fb58b10914d38fc27"
|
||||
|
||||
[[package]]
|
||||
name = "i_overlay"
|
||||
version = "4.0.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "413183068e6e0289e18d7d0a1f661b81546e6918d5453a44570b9ab30cbed1b3"
|
||||
dependencies = [
|
||||
"i_float",
|
||||
"i_key_sort",
|
||||
"i_shape",
|
||||
"i_tree",
|
||||
"rayon",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "i_shape"
|
||||
version = "1.14.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1ea154b742f7d43dae2897fcd5ead86bc7b5eefcedd305a7ebf9f69d44d61082"
|
||||
dependencies = [
|
||||
"i_float",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "i_tree"
|
||||
version = "0.16.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "35e6d558e6d4c7b82bc51d9c771e7a927862a161a7d87bf2b0541450e0e20915"
|
||||
|
||||
[[package]]
|
||||
name = "iana-time-zone"
|
||||
version = "0.1.65"
|
||||
@@ -5061,6 +5262,7 @@ dependencies = [
|
||||
"jsonb",
|
||||
"lance-arrow",
|
||||
"lance-core",
|
||||
"lance-geo",
|
||||
"log",
|
||||
"pin-project",
|
||||
"prost",
|
||||
@@ -5164,6 +5366,21 @@ dependencies = [
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lance-geo"
|
||||
version = "13.0.0-beta.3"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.3#565808c992c3d2adf1909b9753b50a28b6bd8a19"
|
||||
dependencies = [
|
||||
"datafusion",
|
||||
"geo-traits",
|
||||
"geo-types",
|
||||
"geoarrow-array",
|
||||
"geoarrow-schema",
|
||||
"geodatafusion",
|
||||
"lance-core",
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lance-index"
|
||||
version = "13.0.0-beta.3"
|
||||
@@ -5190,6 +5407,9 @@ dependencies = [
|
||||
"dirs",
|
||||
"fst",
|
||||
"futures",
|
||||
"geo-types",
|
||||
"geoarrow-array",
|
||||
"geoarrow-schema",
|
||||
"half",
|
||||
"itertools 0.14.0",
|
||||
"jieba-rs",
|
||||
@@ -5201,6 +5421,7 @@ dependencies = [
|
||||
"lance-datafusion",
|
||||
"lance-encoding",
|
||||
"lance-file",
|
||||
"lance-geo",
|
||||
"lance-index-core",
|
||||
"lance-io",
|
||||
"lance-linalg",
|
||||
@@ -6414,6 +6635,28 @@ dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num_enum"
|
||||
version = "0.7.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5d0bca838442ec211fa11de3a8b0e0e8f3a4522575b5c4c06ed722e005036f26"
|
||||
dependencies = [
|
||||
"num_enum_derive",
|
||||
"rustversion",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num_enum_derive"
|
||||
version = "0.7.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "680998035259dcfcafe653688bf2aa6d3e2dc05e98be6ab46afb089dc84f1df8"
|
||||
dependencies = [
|
||||
"proc-macro-crate",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "number_prefix"
|
||||
version = "0.4.0"
|
||||
@@ -8771,6 +9014,12 @@ dependencies = [
|
||||
"byteorder",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "robust"
|
||||
version = "1.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4e27ee8bb91ca0adcf0ecb116293afa12d393f9c2b9b9cd54d33e8078fe19839"
|
||||
|
||||
[[package]]
|
||||
name = "rsa"
|
||||
version = "0.9.10"
|
||||
@@ -8792,6 +9041,17 @@ dependencies = [
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rstar"
|
||||
version = "0.12.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "421400d13ccfd26dfa5858199c30a5d76f9c54e0dba7575273025b43c5175dbb"
|
||||
dependencies = [
|
||||
"heapless",
|
||||
"num-traits",
|
||||
"smallvec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rstest"
|
||||
version = "0.23.0"
|
||||
@@ -9530,6 +9790,18 @@ dependencies = [
|
||||
"winapi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "spade"
|
||||
version = "2.15.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9699399fd9349b00b184f5635b074f9ec93afffef30c853f8c875b32c0f8c7fa"
|
||||
dependencies = [
|
||||
"hashbrown 0.16.1",
|
||||
"num-traits",
|
||||
"robust",
|
||||
"smallvec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "spin"
|
||||
version = "0.9.9"
|
||||
@@ -11353,6 +11625,31 @@ dependencies = [
|
||||
"wasmparser",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wkb"
|
||||
version = "0.9.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a120b336c7ad17749026d50427c23d838ecb50cd64aaea6254b5030152f890a9"
|
||||
dependencies = [
|
||||
"byteorder",
|
||||
"geo-traits",
|
||||
"num_enum",
|
||||
"thiserror 1.0.69",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wkt"
|
||||
version = "0.14.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "efb2b923ccc882312e559ffaa832a055ba9d1ac0cc8e86b3e25453247e4b81d7"
|
||||
dependencies = [
|
||||
"geo-traits",
|
||||
"geo-types",
|
||||
"log",
|
||||
"num-traits",
|
||||
"thiserror 1.0.69",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "writeable"
|
||||
version = "0.6.3"
|
||||
|
||||
@@ -128,6 +128,8 @@ pprof = { version = "0.14", features = ["flamegraph"] }
|
||||
|
||||
[features]
|
||||
default = []
|
||||
# Native spatial indexes and queries. Remote RTree requests do not require this feature.
|
||||
geo = ["lance/geo"]
|
||||
aws = [
|
||||
"dep:aws-smithy-types",
|
||||
"lance/aws",
|
||||
|
||||
@@ -14,8 +14,8 @@ use crate::{DistanceType, Error, Result, job::Job, table::BaseTable};
|
||||
|
||||
use self::{
|
||||
scalar::{
|
||||
BTreeIndexBuilder, BitmapIndexBuilder, FmIndexBuilder, LabelListIndexBuilder,
|
||||
ZoneMapIndexBuilder,
|
||||
BTreeIndexBuilder, BitmapIndexBuilder, BloomFilterIndexBuilder, FmIndexBuilder,
|
||||
LabelListIndexBuilder, NGramIndexBuilder, RTreeIndexBuilder, ZoneMapIndexBuilder,
|
||||
},
|
||||
vector::{
|
||||
IvfHnswFlatIndexBuilder, IvfHnswPqIndexBuilder, IvfHnswSqIndexBuilder, IvfPqIndexBuilder,
|
||||
@@ -63,6 +63,16 @@ pub enum Index {
|
||||
/// prove they cannot match the predicate.
|
||||
ZoneMap(ZoneMapIndexBuilder),
|
||||
|
||||
/// An NGram index accelerates substring and pattern filters on UTF-8 strings.
|
||||
NGram(NGramIndexBuilder),
|
||||
|
||||
/// A Bloom filter index skips groups of scalar values that cannot match a filter.
|
||||
BloomFilter(BloomFilterIndexBuilder),
|
||||
|
||||
/// An R-tree index accelerates spatial intersection filters on GeoArrow geometries.
|
||||
/// Native creation requires the `geo` feature.
|
||||
RTree(RTreeIndexBuilder),
|
||||
|
||||
/// Full text search index using BM25.
|
||||
///
|
||||
/// The posting block size defaults to 128. Supported values are 128 and 256;
|
||||
@@ -352,6 +362,12 @@ pub enum IndexType {
|
||||
Fm,
|
||||
#[serde(alias = "ZONEMAP", alias = "ZONE_MAP")]
|
||||
ZoneMap,
|
||||
#[serde(alias = "NGRAM", alias = "N_GRAM")]
|
||||
NGram,
|
||||
#[serde(alias = "BLOOM_FILTER", alias = "BLOOMFILTER")]
|
||||
BloomFilter,
|
||||
#[serde(alias = "RTREE", alias = "R_TREE")]
|
||||
RTree,
|
||||
// FTS
|
||||
#[serde(alias = "INVERTED", alias = "Inverted")]
|
||||
FTS,
|
||||
@@ -374,6 +390,9 @@ impl std::fmt::Display for IndexType {
|
||||
Self::LabelList => write!(f, "LABEL_LIST"),
|
||||
Self::Fm => write!(f, "FM"),
|
||||
Self::ZoneMap => write!(f, "ZONEMAP"),
|
||||
Self::NGram => write!(f, "NGRAM"),
|
||||
Self::BloomFilter => write!(f, "BLOOM_FILTER"),
|
||||
Self::RTree => write!(f, "RTREE"),
|
||||
Self::FTS => write!(f, "FTS"),
|
||||
Self::Unknown => write!(f, "UNKNOWN"),
|
||||
}
|
||||
@@ -390,6 +409,9 @@ impl std::str::FromStr for IndexType {
|
||||
"LABEL_LIST" | "LABELLIST" => Ok(Self::LabelList),
|
||||
"FM" | "FMINDEX" => Ok(Self::Fm),
|
||||
"ZONEMAP" | "ZONE_MAP" => Ok(Self::ZoneMap),
|
||||
"NGRAM" | "N_GRAM" => Ok(Self::NGram),
|
||||
"BLOOM_FILTER" | "BLOOMFILTER" => Ok(Self::BloomFilter),
|
||||
"RTREE" | "R_TREE" => Ok(Self::RTree),
|
||||
"FTS" | "INVERTED" => Ok(Self::FTS),
|
||||
"IVF_FLAT" => Ok(Self::IvfFlat),
|
||||
"IVF_SQ" => Ok(Self::IvfSq),
|
||||
@@ -495,3 +517,34 @@ pub struct IndexStatistics {
|
||||
/// The number of parts this index is split into.
|
||||
pub num_indices: Option<u32>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::IndexType;
|
||||
|
||||
#[test]
|
||||
fn builtin_scalar_index_type_names() {
|
||||
for (index_type, canonical, aliases) in [
|
||||
(IndexType::NGram, "NGRAM", ["NGram", "NGRAM", "N_GRAM"]),
|
||||
(
|
||||
IndexType::BloomFilter,
|
||||
"BLOOM_FILTER",
|
||||
["BloomFilter", "BLOOMFILTER", "BLOOM_FILTER"],
|
||||
),
|
||||
(IndexType::RTree, "RTREE", ["RTree", "RTREE", "R_TREE"]),
|
||||
] {
|
||||
assert_eq!(index_type.to_string(), canonical);
|
||||
for alias in aliases {
|
||||
assert_eq!(alias.parse::<IndexType>().unwrap(), index_type);
|
||||
assert_eq!(
|
||||
alias.to_lowercase().parse::<IndexType>().unwrap(),
|
||||
index_type
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::from_value::<IndexType>(serde_json::json!(alias)).unwrap(),
|
||||
index_type
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -82,8 +82,178 @@ pub struct FmIndexBuilder {}
|
||||
#[derive(Debug, Clone, Default, serde::Serialize)]
|
||||
pub struct ZoneMapIndexBuilder {}
|
||||
|
||||
/// Builder for an NGram index over UTF-8 strings.
|
||||
///
|
||||
/// This index accelerates certain substring, `LIKE`, and regular-expression filters.
|
||||
/// It uses Lance's default trigram parameters.
|
||||
///
|
||||
/// ```
|
||||
/// use lancedb::index::{Index, scalar::NGramIndexBuilder};
|
||||
/// # async fn example(table: &lancedb::Table) -> lancedb::Result<()> {
|
||||
/// table.create_index(&["text"], Index::NGram(NGramIndexBuilder::default()))
|
||||
/// .execute().await?;
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
#[derive(Debug, Clone, Default, serde::Serialize)]
|
||||
pub struct NGramIndexBuilder {}
|
||||
|
||||
/// Builder for a Bloom filter index on scalar values.
|
||||
///
|
||||
/// Bloom filters accelerate equality and membership filters by skipping groups of rows that cannot
|
||||
/// match. Candidate rows are checked to remove false positives. A bloom filter is much smaller than
|
||||
/// a btree or bitmap index, but not as precise. It is also limited to equality queries.
|
||||
///
|
||||
/// Unset parameters use Lance's defaults.
|
||||
///
|
||||
/// ```
|
||||
/// use lancedb::index::{Index, scalar::BloomFilterIndexBuilder};
|
||||
/// # async fn example(table: &lancedb::Table) -> lancedb::Result<()> {
|
||||
/// let params = BloomFilterIndexBuilder::default()
|
||||
/// .number_of_items(4096)?
|
||||
/// .probability(0.01)?;
|
||||
/// table.create_index(&["id"], Index::BloomFilter(params))
|
||||
/// .execute().await?;
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
#[derive(Debug, Clone, Default, serde::Serialize)]
|
||||
pub struct BloomFilterIndexBuilder {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
number_of_items: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
probability: Option<f64>,
|
||||
}
|
||||
|
||||
impl BloomFilterIndexBuilder {
|
||||
/// Set the number of rows covered by each Bloom filter.
|
||||
///
|
||||
/// Must be greater than zero. Defaults to 8192, unless overridden by Lance's
|
||||
/// `LANCE_BLOOMFILTER_DEFAULT_NUMBER_OF_ITEMS` environment variable.
|
||||
pub fn number_of_items(mut self, number_of_items: u64) -> crate::Result<Self> {
|
||||
if number_of_items == 0 {
|
||||
return Err(crate::Error::InvalidInput {
|
||||
message: "BloomFilter number_of_items must be greater than zero".into(),
|
||||
});
|
||||
}
|
||||
self.number_of_items = Some(number_of_items);
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
/// Set the desired false-positive probability for each Bloom filter.
|
||||
///
|
||||
/// Must be finite and strictly between zero and one. Lower values use more
|
||||
/// space. Defaults to 0.00057, unless overridden by Lance's
|
||||
/// `LANCE_BLOOMFILTER_DEFAULT_PROBABILITY` environment variable.
|
||||
pub fn probability(mut self, probability: f64) -> crate::Result<Self> {
|
||||
if !probability.is_finite() || probability <= 0.0 || probability >= 1.0 {
|
||||
return Err(crate::Error::InvalidInput {
|
||||
message: "BloomFilter probability must be finite and strictly between zero and one"
|
||||
.into(),
|
||||
});
|
||||
}
|
||||
self.probability = Some(probability);
|
||||
Ok(self)
|
||||
}
|
||||
}
|
||||
|
||||
/// Builder for an R-tree index on GeoArrow geometry columns.
|
||||
///
|
||||
/// This index accelerates spatial intersection filters using geometry bounding
|
||||
/// boxes. Unset parameters use Lance's defaults. Native creation requires
|
||||
/// the `geo` feature; remote creation requires server support.
|
||||
///
|
||||
/// ```
|
||||
/// use lancedb::index::{Index, scalar::RTreeIndexBuilder};
|
||||
/// # async fn example(table: &lancedb::Table) -> lancedb::Result<()> {
|
||||
/// let params = RTreeIndexBuilder::default().page_size(1024)?;
|
||||
/// table.create_index(&["geometry"], Index::RTree(params))
|
||||
/// .execute().await?;
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
#[derive(Debug, Clone, Default, serde::Serialize)]
|
||||
pub struct RTreeIndexBuilder {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
page_size: Option<u32>,
|
||||
}
|
||||
|
||||
impl RTreeIndexBuilder {
|
||||
/// Set the maximum number of entries in each R-tree page.
|
||||
///
|
||||
/// Must be at least 2. Defaults to 4096.
|
||||
pub fn page_size(mut self, page_size: u32) -> crate::Result<Self> {
|
||||
if page_size < 2 {
|
||||
return Err(crate::Error::InvalidInput {
|
||||
message: "RTree page_size must be at least 2".into(),
|
||||
});
|
||||
}
|
||||
self.page_size = Some(page_size);
|
||||
Ok(self)
|
||||
}
|
||||
}
|
||||
|
||||
pub use lance_index::scalar::FullTextSearchQuery;
|
||||
pub use lance_index::scalar::InvertedIndexParams as FtsIndexBuilder;
|
||||
pub use lance_index::scalar::InvertedIndexParams;
|
||||
pub use lance_index::scalar::inverted::DocumentGranularity;
|
||||
pub use lance_index::scalar::inverted::query::*;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn scalar_index_parameters() {
|
||||
assert_eq!(
|
||||
serde_json::to_value(BloomFilterIndexBuilder::default()).unwrap(),
|
||||
serde_json::json!({})
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::to_value(RTreeIndexBuilder::default()).unwrap(),
|
||||
serde_json::json!({})
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::to_value(
|
||||
BloomFilterIndexBuilder::default()
|
||||
.number_of_items(1)
|
||||
.unwrap()
|
||||
)
|
||||
.unwrap(),
|
||||
serde_json::json!({"number_of_items": 1})
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::to_value(
|
||||
BloomFilterIndexBuilder::default()
|
||||
.probability(0.01)
|
||||
.unwrap()
|
||||
)
|
||||
.unwrap(),
|
||||
serde_json::json!({"probability": 0.01})
|
||||
);
|
||||
assert!(
|
||||
BloomFilterIndexBuilder::default()
|
||||
.number_of_items(0)
|
||||
.is_err()
|
||||
);
|
||||
for probability in [
|
||||
f64::NAN,
|
||||
f64::INFINITY,
|
||||
f64::NEG_INFINITY,
|
||||
-0.1,
|
||||
0.0,
|
||||
1.0,
|
||||
1.1,
|
||||
] {
|
||||
assert!(
|
||||
BloomFilterIndexBuilder::default()
|
||||
.probability(probability)
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
for page_size in [0, 1] {
|
||||
assert!(RTreeIndexBuilder::default().page_size(page_size).is_err());
|
||||
}
|
||||
assert!(RTreeIndexBuilder::default().page_size(2).is_ok());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -588,6 +588,9 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
Index::LabelList(p) => ("LABEL_LIST", Some(to_json(p)?)),
|
||||
Index::Fm(p) => ("FM", Some(to_json(p)?)),
|
||||
Index::ZoneMap(p) => ("ZONEMAP", Some(to_json(p)?)),
|
||||
Index::NGram(p) => ("NGRAM", Some(to_json(p)?)),
|
||||
Index::BloomFilter(p) => ("BLOOM_FILTER", Some(to_json(p)?)),
|
||||
Index::RTree(p) => ("RTREE", Some(to_json(p)?)),
|
||||
Index::FTS(p) => {
|
||||
let mut params = to_json(p)?;
|
||||
if p.get_document_granularity().is_list_element() {
|
||||
@@ -6378,6 +6381,33 @@ mod tests {
|
||||
("BTREE", json!({}), Index::BTree(Default::default())),
|
||||
("BITMAP", json!({}), Index::Bitmap(Default::default())),
|
||||
("ZONEMAP", json!({}), Index::ZoneMap(Default::default())),
|
||||
("NGRAM", json!({}), Index::NGram(Default::default())),
|
||||
(
|
||||
"BLOOM_FILTER",
|
||||
json!({}),
|
||||
Index::BloomFilter(Default::default()),
|
||||
),
|
||||
("RTREE", json!({}), Index::RTree(Default::default())),
|
||||
(
|
||||
"BLOOM_FILTER",
|
||||
json!({"number_of_items": 4096, "probability": 0.01}),
|
||||
Index::BloomFilter(
|
||||
crate::index::scalar::BloomFilterIndexBuilder::default()
|
||||
.number_of_items(4096)
|
||||
.unwrap()
|
||||
.probability(0.01)
|
||||
.unwrap(),
|
||||
),
|
||||
),
|
||||
(
|
||||
"RTREE",
|
||||
json!({"page_size": 1024}),
|
||||
Index::RTree(
|
||||
crate::index::scalar::RTreeIndexBuilder::default()
|
||||
.page_size(1024)
|
||||
.unwrap(),
|
||||
),
|
||||
),
|
||||
(
|
||||
"LABEL_LIST",
|
||||
json!({}),
|
||||
|
||||
@@ -265,6 +265,26 @@ impl NativeTable {
|
||||
BuiltinIndexType::ZoneMap,
|
||||
)))
|
||||
}
|
||||
Index::NGram(_) => Ok(Box::new(ScalarIndexParams::for_builtin(
|
||||
BuiltinIndexType::NGram,
|
||||
))),
|
||||
Index::BloomFilter(params) => {
|
||||
let params = serde_json::to_value(params).map_err(|e| Error::InvalidInput {
|
||||
message: format!("failed to serialize index params: {e}"),
|
||||
})?;
|
||||
Ok(Box::new(
|
||||
ScalarIndexParams::for_builtin(BuiltinIndexType::BloomFilter)
|
||||
.with_params(¶ms),
|
||||
))
|
||||
}
|
||||
Index::RTree(params) => {
|
||||
let params = serde_json::to_value(params).map_err(|e| Error::InvalidInput {
|
||||
message: format!("failed to serialize index params: {e}"),
|
||||
})?;
|
||||
Ok(Box::new(
|
||||
ScalarIndexParams::for_builtin(BuiltinIndexType::RTree).with_params(¶ms),
|
||||
))
|
||||
}
|
||||
Index::FTS(fts_opts) => {
|
||||
Self::validate_index_type(field, "FTS", supported_fts_data_type)?;
|
||||
Ok(Box::new(fts_opts))
|
||||
@@ -425,6 +445,9 @@ impl NativeTable {
|
||||
Index::LabelList(_) => IndexType::LabelList,
|
||||
Index::Fm(_) => IndexType::Fm,
|
||||
Index::ZoneMap(_) => IndexType::ZoneMap,
|
||||
Index::NGram(_) => IndexType::NGram,
|
||||
Index::BloomFilter(_) => IndexType::BloomFilter,
|
||||
Index::RTree(_) => IndexType::RTree,
|
||||
Index::FTS(_) => IndexType::Inverted,
|
||||
Index::IvfFlat(_)
|
||||
| Index::IvfSq(_)
|
||||
@@ -439,6 +462,7 @@ impl NativeTable {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use lance::index::DatasetIndexExt;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -1321,6 +1345,191 @@ mod tests {
|
||||
assert_eq!(null_count, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_builtin_scalar_index() {
|
||||
for (index, index_type, predicate, invalid_column) in [
|
||||
(
|
||||
Index::NGram(Default::default()),
|
||||
crate::index::IndexType::NGram,
|
||||
"contains(text, 'abc')",
|
||||
"id",
|
||||
),
|
||||
(
|
||||
Index::BloomFilter(Default::default()),
|
||||
crate::index::IndexType::BloomFilter,
|
||||
"text = 'abc'",
|
||||
"flag",
|
||||
),
|
||||
(
|
||||
Index::BloomFilter(
|
||||
crate::index::scalar::BloomFilterIndexBuilder::default()
|
||||
.number_of_items(2)
|
||||
.unwrap()
|
||||
.probability(0.01)
|
||||
.unwrap(),
|
||||
),
|
||||
crate::index::IndexType::BloomFilter,
|
||||
"text = 'abc'",
|
||||
"flag",
|
||||
),
|
||||
] {
|
||||
let conn = connect("memory://").execute().await.unwrap();
|
||||
let batch = record_batch!(
|
||||
("id", Int32, [1, 2, 3, 4]),
|
||||
("text", Utf8, [Some("abc"), Some("xyz"), None, Some("")]),
|
||||
("flag", Boolean, [true, false, true, false])
|
||||
)
|
||||
.unwrap();
|
||||
let table = conn
|
||||
.create_table("scalar", batch.clone())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
table.add(batch).execute().await.unwrap();
|
||||
|
||||
table
|
||||
.create_index(&["text"], index.clone())
|
||||
.name("text_search".into())
|
||||
.train(false)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
table.list_indices().await.unwrap()[0].index_type,
|
||||
index_type
|
||||
);
|
||||
assert!(
|
||||
table
|
||||
.create_index(&["text"], index.clone())
|
||||
.name("text_search".into())
|
||||
.replace(false)
|
||||
.execute()
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
table
|
||||
.create_index(&["text"], index.clone())
|
||||
.name("text_search".into())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let stats = table.index_stats("text_search").await.unwrap().unwrap();
|
||||
assert_eq!(stats.index_type, index_type);
|
||||
assert_eq!(stats.num_indexed_rows, 8);
|
||||
if let Index::BloomFilter(params) = &index {
|
||||
let expected = serde_json::to_value(params).unwrap();
|
||||
let dataset = table.as_native().unwrap().dataset.get().await.unwrap();
|
||||
let stats: serde_json::Value =
|
||||
serde_json::from_str(&dataset.index_statistics("text_search").await.unwrap())
|
||||
.unwrap();
|
||||
for (key, value) in expected.as_object().unwrap() {
|
||||
assert_eq!(&stats["indices"][0][key], value);
|
||||
}
|
||||
}
|
||||
for (filter, expected_rows) in [(predicate, 2), ("text IS NULL", 2), ("text = ''", 2)] {
|
||||
let query = table.query().only_if(filter);
|
||||
let plan = query.explain_plan(false).await.unwrap();
|
||||
if filter == predicate {
|
||||
assert!(plan.contains("ScalarIndexQuery"), "{plan}");
|
||||
}
|
||||
let batches = query
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
|
||||
expected_rows
|
||||
);
|
||||
}
|
||||
let error = table
|
||||
.create_index(&[invalid_column], index)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
error.to_string().to_lowercase().contains(
|
||||
if index_type == crate::index::IndexType::NGram {
|
||||
"ngram"
|
||||
} else {
|
||||
"bloom"
|
||||
}
|
||||
),
|
||||
"{error}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "geo")]
|
||||
#[tokio::test]
|
||||
async fn test_create_rtree_index() {
|
||||
let points = StructArray::new(
|
||||
vec![
|
||||
Field::new("x", DataType::Float64, false),
|
||||
Field::new("y", DataType::Float64, false),
|
||||
]
|
||||
.into(),
|
||||
vec![
|
||||
Arc::new(arrow_array::Float64Array::from(vec![0.0, 10.0, 0.0])),
|
||||
Arc::new(arrow_array::Float64Array::from(vec![0.0, 10.0, 0.0])),
|
||||
],
|
||||
Some(arrow_buffer::NullBuffer::from(vec![true, true, false])),
|
||||
);
|
||||
let field = Field::new("geometry", points.data_type().clone(), true).with_metadata(
|
||||
std::collections::HashMap::from([
|
||||
("ARROW:extension:name".into(), "geoarrow.point".into()),
|
||||
("ARROW:extension:metadata".into(), "{}".into()),
|
||||
]),
|
||||
);
|
||||
let batch =
|
||||
RecordBatch::try_new(Arc::new(Schema::new(vec![field])), vec![Arc::new(points)])
|
||||
.unwrap();
|
||||
let conn = connect("memory://").execute().await.unwrap();
|
||||
let table = conn.create_table("spatial", batch).execute().await.unwrap();
|
||||
table
|
||||
.create_index(
|
||||
&["geometry"],
|
||||
Index::RTree(
|
||||
crate::index::scalar::RTreeIndexBuilder::default()
|
||||
.page_size(2)
|
||||
.unwrap(),
|
||||
),
|
||||
)
|
||||
.name("spatial_idx".into())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let indices = table.list_indices().await.unwrap();
|
||||
assert_eq!(indices.len(), 1);
|
||||
assert_eq!(indices[0].index_type, crate::index::IndexType::RTree);
|
||||
let stats = table.index_stats("spatial_idx").await.unwrap().unwrap();
|
||||
assert_eq!(stats.index_type, crate::index::IndexType::RTree);
|
||||
assert_eq!(stats.num_indexed_rows, 3);
|
||||
let dataset = table.as_native().unwrap().dataset.get().await.unwrap();
|
||||
let stats: serde_json::Value =
|
||||
serde_json::from_str(&dataset.index_statistics("spatial_idx").await.unwrap()).unwrap();
|
||||
assert_eq!(stats["indices"][0]["page_size"], 2);
|
||||
for predicate in [
|
||||
"ST_Intersects(geometry, ST_GeomFromText('POLYGON ((-1 -1, 1 -1, 1 1, -1 1, -1 -1))'))",
|
||||
"geometry IS NULL",
|
||||
] {
|
||||
let query = table.query().only_if(predicate);
|
||||
let plan = query.explain_plan(false).await.unwrap();
|
||||
assert!(plan.contains("ScalarIndexQuery"), "{plan}");
|
||||
let batches = query
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_index_nested_field_paths() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
|
||||
Reference in New Issue
Block a user