From e459892bb5ee611daf53fe5e61d51aeb70e41dd5 Mon Sep 17 00:00:00 2001 From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 19:43:33 +0000 Subject: [PATCH] fix: support ST_DWithin filters --- Cargo.lock | 297 ++++++++++++++++++++++++++++++++++ Cargo.toml | 2 +- rust/lancedb/src/query.rs | 49 +++++- rust/lancedb/src/query/sql.rs | 150 +++++++++++++++++ 4 files changed, 493 insertions(+), 5 deletions(-) create mode 100644 rust/lancedb/src/query/sql.rs diff --git a/Cargo.lock b/Cargo.lock index 995424b20..0101dcebd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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" diff --git a/Cargo.toml b/Cargo.toml index 35370e58f..ae57c75d6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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" } diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index b76865043..e70dcbf13 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -32,6 +32,7 @@ use crate::{ }; mod hybrid; +mod sql; pub(crate) const DEFAULT_TOP_K: usize = 10; @@ -561,8 +562,8 @@ impl QueryBase for T { } fn only_if(mut self, filter: impl AsRef) -> 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::>() + .await + .unwrap(); + let batch = concat_batches(&batches[0].schema(), &batches).unwrap(); + let ids = batch["id"].as_primitive::(); + + assert_eq!(ids.values(), &[1, 2]); + } + fn make_non_empty_batches() -> Box { let vec = Box::new(RandomVector::new().named("vector".to_string())); let id = Box::new(IncrementingInt32::new().named("id".to_string())); diff --git a/rust/lancedb/src/query/sql.rs b/rust/lancedb/src/query/sql.rs new file mode 100644 index 000000000..80f9ceefe --- /dev/null +++ b/rust/lancedb/src/query/sql.rs @@ -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 { + 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)" + ); + } +}