rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-08-18 12:08:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: support ST_DWithin filters

Browse Source
This commit is contained in:
Gatefixer
2026-08-05 19:43:33 +00:00
parent c7ea91f3ea
commit e459892bb5
4 changed files with 493 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+297
View File
@@ -3081,6 +3081,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"
@@ -3380,6 +3390,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"
@@ -3690,6 +3706,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"
@@ -3865,6 +4004,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"
@@ -3921,6 +4069,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"
@@ -4235,6 +4393,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"
@@ -4971,6 +5172,7 @@ dependencies = [
"lance-arrow",
"lance-core",
"lance-datagen",
"lance-geo",
"log",
"pin-project",
"prost",
@@ -5075,6 +5277,21 @@ dependencies = [
"tracing",
]
[[package]]
name = "lance-geo"
version = "10.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v10.1.0-beta.1#68f4d4c1d0c4871b067557c61fc405078f1ab3b7"
dependencies = [
"datafusion",
"geo-traits",
"geo-types",
"geoarrow-array",
"geoarrow-schema",
"geodatafusion",
"lance-core",
"serde",
]
[[package]]
name = "lance-index"
version = "10.1.0-beta.1"
@@ -5102,6 +5319,9 @@ dependencies = [
"dirs",
"fst",
"futures",
"geo-types",
"geoarrow-array",
"geoarrow-schema",
"half",
"itertools 0.14.0",
"jieba-rs",
@@ -5114,6 +5334,7 @@ dependencies = [
"lance-datagen",
"lance-encoding",
"lance-file",
"lance-geo",
"lance-index-core",
"lance-io",
"lance-linalg",
@@ -6299,6 +6520,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"
@@ -8536,6 +8779,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"
@@ -8557,6 +8806,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"
@@ -9326,6 +9586,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"
@@ -11127,6 +11399,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"
Cargo.toml
+1 -1
View File
@@ -13,7 +13,7 @@ categories = ["database-implementations"]
rust-version = "1.91.0"
[workspace.dependencies]
lance = { "version" = "=10.1.0-beta.1", default-features = false, "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance = { "version" = "=10.1.0-beta.1", default-features = false, features = ["geo"], "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-core = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-datagen = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-file = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
rust/lancedb/src/query.rs
+45 -4
View File
@@ -32,6 +32,7 @@ use crate::{
};
mod hybrid;
mod sql;
pub(crate) const DEFAULT_TOP_K: usize = 10;
@@ -561,8 +562,8 @@ impl<T: HasQuery> QueryBase for T {
}
fn only_if(mut self, filter: impl AsRef<str>) -> Self {
self.mut_query()
.add_filter(QueryFilter::Sql(filter.as_ref().to_string()));
let filter = sql::rewrite_st_dwithin(filter.as_ref());
self.mut_query().add_filter(QueryFilter::Sql(filter));
self
}
@@ -1648,8 +1649,8 @@ mod tests {
use super::*;
use arrow::{array::downcast_array, compute::concat_batches, datatypes::Int32Type};
use arrow_array::{
FixedSizeListArray, Float32Array, Int32Array, RecordBatch, StringArray, cast::AsArray,
types::Float32Type,
FixedSizeListArray, Float32Array, Float64Array, Int32Array, RecordBatch, StringArray,
cast::AsArray, types::Float32Type,
};
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
use futures::{StreamExt, TryStreamExt};
@@ -2024,6 +2025,46 @@ mod tests {
assert!(result.is_err());
}
#[tokio::test]
async fn test_st_dwithin_filter() {
let tmp_dir = tempdir().unwrap();
let uri = tmp_dir.path().to_str().unwrap();
let batch = RecordBatch::try_from_iter([
("id", Arc::new(Int32Array::from_iter_values([1, 2, 3])) as _),
(
"x",
Arc::new(Float64Array::from_iter_values([0.0, 1.0, 3.0])) as _,
),
(
"y",
Arc::new(Float64Array::from_iter_values([0.0, 1.0, 4.0])) as _,
),
])
.unwrap();
let table = connect(uri)
.execute()
.await
.unwrap()
.create_table("points", batch)
.execute()
.await
.unwrap();
let batches = table
.query()
.only_if("ST_DWithin(ST_Point(x, y), ST_Point(0.0, 0.0), 2.0)")
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let batch = concat_batches(&batches[0].schema(), &batches).unwrap();
let ids = batch["id"].as_primitive::<Int32Type>();
assert_eq!(ids.values(), &[1, 2]);
}
fn make_non_empty_batches() -> Box<dyn arrow_array::RecordBatchReader + Send> {
let vec = Box::new(RandomVector::new().named("vector".to_string()));
let id = Box::new(IncrementingInt32::new().named("id".to_string()));
rust/lancedb/src/query/sql.rs
+150
View File
@@ -0,0 +1,150 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! SQL compatibility helpers for query filters.
use std::{any::TypeId, ops::ControlFlow};
use datafusion_sql::sqlparser::{
ast::{
BinaryOperator, Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Ident, ObjectName,
SelectItem, SetExpr, Statement, visit_expressions_mut,
},
dialect::{Dialect, GenericDialect},
parser::Parser,
tokenizer::{Token, Tokenizer},
};
#[derive(Debug, Default)]
struct LanceDialect(GenericDialect);
impl Dialect for LanceDialect {
fn dialect(&self) -> TypeId {
self.0.dialect()
}
fn is_identifier_start(&self, ch: char) -> bool {
self.0.is_identifier_start(ch)
}
fn is_identifier_part(&self, ch: char) -> bool {
self.0.is_identifier_part(ch)
}
fn is_delimited_identifier_start(&self, ch: char) -> bool {
ch == '`'
}
}
/// Rewrite `ST_DWithin(a, b, distance)` as `ST_Distance(a, b) <= distance`.
///
/// GeoDataFusion does not currently expose `ST_DWithin`, but its definition is
/// exactly this distance comparison. Invalid SQL and unsupported function
/// shapes are left untouched so the query planner can report the usual error.
pub(super) fn rewrite_st_dwithin(filter: &str) -> String {
if !filter.to_ascii_lowercase().contains("st_dwithin") {
return filter.to_string();
}
let Some(mut expr) = parse_filter(filter) else {
return filter.to_string();
};
let mut rewritten = false;
let _ = visit_expressions_mut(&mut expr, |expr| {
let Expr::Function(function) = expr else {
return ControlFlow::<()>::Continue(());
};
if !function.name.to_string().eq_ignore_ascii_case("st_dwithin") {
return ControlFlow::Continue(());
}
let FunctionArguments::List(arguments) = &function.args else {
return ControlFlow::Continue(());
};
let [
FunctionArg::Unnamed(FunctionArgExpr::Expr(first)),
FunctionArg::Unnamed(FunctionArgExpr::Expr(second)),
FunctionArg::Unnamed(FunctionArgExpr::Expr(distance)),
] = arguments.args.as_slice()
else {
return ControlFlow::Continue(());
};
let mut st_distance = function.clone();
st_distance.name = ObjectName::from(vec![Ident::new("ST_Distance")]);
if let FunctionArguments::List(arguments) = &mut st_distance.args {
arguments.args = vec![
FunctionArg::Unnamed(FunctionArgExpr::Expr(first.clone())),
FunctionArg::Unnamed(FunctionArgExpr::Expr(second.clone())),
];
}
*expr = Expr::BinaryOp {
left: Box::new(Expr::Function(st_distance)),
op: BinaryOperator::LtEq,
right: Box::new(distance.clone()),
};
rewritten = true;
ControlFlow::Continue(())
});
if rewritten {
expr.to_string()
} else {
filter.to_string()
}
}
fn parse_filter(filter: &str) -> Option<Expr> {
let statement = format!("SELECT 1 FROM t WHERE {filter}");
let dialect = LanceDialect::default();
let mut token_iter = Tokenizer::new(&dialect, &statement)
.tokenize()
.ok()?
.into_iter();
let mut previous = token_iter.next()?;
let mut tokens = Vec::new();
// Match Lance's support for `==` as an equality operator.
for next in token_iter {
if let (Token::Eq, Token::Eq) = (&previous, &next) {
continue;
}
tokens.push(std::mem::replace(&mut previous, next));
}
tokens.push(previous);
let statement = Parser::new(&dialect)
.with_tokens(tokens)
.parse_statement()
.ok()?;
if let Statement::Query(query) = statement
&& let SetExpr::Select(select) = *query.body
&& let Some(expr) = select.selection
&& matches!(select.projection.as_slice(), [SelectItem::UnnamedExpr(_)])
{
Some(expr)
} else {
None
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rewrites_nested_and_case_insensitive_dwithin() {
assert_eq!(
rewrite_st_dwithin("id == 1 AND st_dwithin(ST_Point(x, y), ST_Point(0, 0), radius)"),
"id = 1 AND ST_Distance(ST_Point(x, y), ST_Point(0, 0)) <= radius"
);
}
#[test]
fn leaves_other_filters_unchanged() {
assert_eq!(rewrite_st_dwithin("id == 1"), "id == 1");
assert_eq!(
rewrite_st_dwithin("ST_DWithin(point, origin)"),
"ST_DWithin(point, origin)"
);
}
}