From 5a60ce04e3ec2dd515defa7ca10497011e9aeff9 Mon Sep 17 00:00:00 2001 From: Chris Lu Date: Sun, 1 Nov 2026 12:06:20 -0800 Subject: [PATCH] vector: stop re-parsing a filter's JSONPath for every scanned row (#12302) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * vector: parse a filter's JSONPath once per batch instead of per row Lance's json_extract and json_exists parse their JSONPath argument again for every row they see. Every translated S3 Vectors filter goes through them, so a flat scan paid about 1 µs per row before the predicate compared anything, and a 10%-selective {"category": {"$eq": "c3"}} ran about 50% slower than the same query unfiltered. The path is always a literal, so the scanner now parses the filter as before and swaps in copies of the two functions that parse it once per batch. They call the same jsonb selector and serializer, so results and errors do not change. Calls whose path is not a literal keep Lance's functions. The filter SQL is unchanged, so the gateway, existing tables and the distributed workers (which run query::ann) need no change. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * s3vectors: read the field once per extractor in $in and $nin $in and $nin rendered one IN per member, each with its own json_extract, so an N-member list parsed the row's metadata N times. Members that share an extractor (strings and booleans, or numbers) now go into one IN list, so a list reads the field at most twice. Matching is unchanged: each group still only sees its own JSON type and still coalesces NULL before $nin negates it. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * s3vectors: rename the $in literal map so codespell passes Codespell reads "lits" as a misspelling of "list". Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK --------- Co-authored-by: Claude Opus 5.5 --- seaweed-vector/Cargo.lock | 1 + seaweed-vector/crates/core/Cargo.toml | 1 + seaweed-vector/crates/core/src/filter.rs | 334 +++++++++++++++++++++++ seaweed-vector/crates/core/src/lib.rs | 1 + seaweed-vector/crates/core/src/query.rs | 2 +- test/s3vectors/s3vectors_test.py | 4 + weed/s3api/s3vectors/handlers.go | 19 +- weed/s3api/s3vectors/handlers_test.go | 4 +- 8 files changed, 359 insertions(+), 7 deletions(-) create mode 100644 seaweed-vector/crates/core/src/filter.rs diff --git a/seaweed-vector/Cargo.lock b/seaweed-vector/Cargo.lock index df09ba67fe..e1a3225934 100644 --- a/seaweed-vector/Cargo.lock +++ b/seaweed-vector/Cargo.lock @@ -6461,6 +6461,7 @@ dependencies = [ "arrow-schema", "clap", "futures", + "jsonb", "lance", "lance-index", "lance-linalg", diff --git a/seaweed-vector/crates/core/Cargo.toml b/seaweed-vector/crates/core/Cargo.toml index 8730650e0a..eb792af28b 100644 --- a/seaweed-vector/crates/core/Cargo.toml +++ b/seaweed-vector/crates/core/Cargo.toml @@ -26,6 +26,7 @@ arrow-array = "58" arrow-cast = "58" arrow-schema = "58" roaring = "0.11" +jsonb = { version = "0.5", default-features = false } anyhow.workspace = true clap = { version = "4", features = ["derive", "env"] } futures = "0.3" diff --git a/seaweed-vector/crates/core/src/filter.rs b/seaweed-vector/crates/core/src/filter.rs new file mode 100644 index 0000000000..da1fabf458 --- /dev/null +++ b/seaweed-vector/crates/core/src/filter.rs @@ -0,0 +1,334 @@ +//! Lance's json_extract and json_exists parse their JSONPath again for every +//! row. A filter's path is a literal, so swap in copies that parse it once per +//! batch and return the same values. + +use std::sync::Arc; + +use arrow_array::builder::{BooleanBuilder, StringBuilder}; +use arrow_array::{Array, ArrayRef, LargeBinaryArray, new_null_array}; +use arrow_schema::DataType; +use jsonb::RawJsonb; +use jsonb::jsonpath::{Selector, parse_json_path}; +use lance::dataset::scanner::Scanner; +use lance::deps::datafusion::common::tree_node::{Transformed, TreeNode}; +use lance::deps::datafusion::common::{DataFusionError, Result, ScalarValue}; +use lance::deps::datafusion::logical_expr::expr::ScalarFunction; +use lance::deps::datafusion::logical_expr::{ + ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, +}; + +/// Set `filter` on the scanner with the JSONPath functions swapped. +pub(crate) fn apply(scanner: &mut Scanner, filter: &str) -> anyhow::Result<()> { + scanner.filter(filter)?; + if let Some(expr) = scanner.get_expr_filter()? { + scanner.filter_expr(swap(expr)?); + } + Ok(()) +} + +fn swap(expr: Expr) -> Result { + expr.transform(|e| { + let Expr::ScalarFunction(f) = e else { + return Ok(Transformed::no(e)); + }; + let op = match f.func.name() { + "json_extract" => Op::Extract, + "json_exists" => Op::Exists, + _ => return Ok(Transformed::no(Expr::ScalarFunction(f))), + }; + if !matches!( + f.args.get(1), + Some(Expr::Literal(ScalarValue::Utf8(Some(_)), _)) + ) { + return Ok(Transformed::no(Expr::ScalarFunction(f))); + } + let udf = Arc::new(ScalarUDF::new_from_impl(JsonPath::new(op))); + Ok(Transformed::yes(Expr::ScalarFunction( + ScalarFunction::new_udf(udf, f.args), + ))) + }) + .map(|t| t.data) +} + +#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)] +enum Op { + Extract, + Exists, +} + +#[derive(Debug, PartialEq, Eq, Hash)] +struct JsonPath { + op: Op, + signature: Signature, +} + +impl JsonPath { + fn new(op: Op) -> Self { + Self { + op, + signature: Signature::exact( + vec![DataType::LargeBinary, DataType::Utf8], + Volatility::Immutable, + ), + } + } +} + +fn exec_err(msg: String) -> DataFusionError { + DataFusionError::Execution(msg) +} + +impl ScalarUDFImpl for JsonPath { + fn name(&self) -> &str { + match self.op { + Op::Extract => "json_extract", + Op::Exists => "json_exists", + } + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _: &[DataType]) -> Result { + Ok(match self.op { + Op::Extract => DataType::Utf8, + Op::Exists => DataType::Boolean, + }) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let [data, ColumnarValue::Scalar(ScalarValue::Utf8(Some(path)))] = args.args.as_slice() + else { + return Err(exec_err(format!( + "{} expects a literal JSONPath", + self.name() + ))); + }; + let data = data.to_array(args.number_rows)?; + let data = data + .as_any() + .downcast_ref::() + .ok_or_else(|| exec_err("First argument must be LargeBinary".to_string()))?; + // The original parses per non-null row, so an all-null batch never fails. + if data.null_count() == data.len() { + let return_type = self.return_type(&[])?; + return Ok(ColumnarValue::Array(new_null_array( + &return_type, + data.len(), + ))); + } + let json_path = parse_json_path(path.as_bytes()) + .map_err(|e| exec_err(format!("Invalid JSONPath '{path}': {e}")))?; + let out: ArrayRef = match self.op { + Op::Extract => { + let mut builder = StringBuilder::with_capacity(data.len(), 1024); + for i in 0..data.len() { + if data.is_null(i) { + builder.append_null(); + continue; + } + let mut selector = Selector::new(RawJsonb::new(data.value(i))); + match selector.select_value(&json_path) { + Ok(Some(value)) => builder.append_value(value.to_string()), + Ok(None) => builder.append_null(), + Err(e) => { + return Err(exec_err(format!( + "Failed to select value from path '{path}': {e}" + ))); + } + } + } + Arc::new(builder.finish()) + } + Op::Exists => { + let mut builder = BooleanBuilder::with_capacity(data.len()); + for i in 0..data.len() { + if data.is_null(i) { + builder.append_null(); + continue; + } + let mut selector = Selector::new(RawJsonb::new(data.value(i))); + match selector.exists(&json_path) { + Ok(exists) => builder.append_value(exists), + Err(e) => { + return Err(exec_err(format!( + "Failed to check existence of path '{path}': {e}" + ))); + } + } + } + Arc::new(builder.finish()) + } + }; + Ok(ColumnarValue::Array(out)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::testing::DIM; + use arrow_array::{ + FixedSizeListArray, RecordBatch, RecordBatchIterator, StringArray, types::Float32Type, + }; + use arrow_schema::{Field, Schema}; + use futures::TryStreamExt; + use lance::Dataset; + use tempfile::TempDir; + + async fn dataset(dir: &TempDir, metas: &[Option<&str>]) -> Dataset { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIM as i32, + ), + false, + ), + Field::new("metadata", DataType::Utf8, true).with_metadata( + [("ARROW:extension:name".to_string(), "arrow.json".to_string())] + .into_iter() + .collect(), + ), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from( + (0..metas.len()) + .map(|i| format!("k{i}")) + .collect::>(), + )), + Arc::new( + FixedSizeListArray::from_iter_primitive::( + (0..metas.len()).map(|i| Some(vec![Some(i as f32); DIM])), + DIM as i32, + ), + ), + Arc::new(StringArray::from(metas.to_vec())), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema); + Dataset::write(reader, dir.path().to_str().unwrap(), None) + .await + .unwrap() + } + + async fn keys(dataset: &Dataset, filter: &str, swapped: bool) -> anyhow::Result> { + let mut scanner = dataset.scan(); + if swapped { + apply(&mut scanner, filter)?; + } else { + scanner.filter(filter)?; + } + scanner.project(&["key"])?; + let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; + let mut out: Vec = batches + .iter() + .flat_map(|b| { + let keys = b.column(0).as_any().downcast_ref::().unwrap(); + keys.iter() + .map(|k| k.unwrap().to_string()) + .collect::>() + }) + .collect(); + out.sort(); + Ok(out) + } + + #[tokio::test] + async fn swapped_functions_match_lance() { + let dir = TempDir::new().unwrap(); + let metas = [ + Some(r#"{"k": 3}"#), + Some(r#"{"k": "3"}"#), + Some(r#"{"k": true}"#), + Some(r#"{"k": "v"}"#), + Some(r#"{"k": "a"}"#), + Some(r#"{"k": "b"}"#), + Some(r#"{"x": 1}"#), + Some(r#"{"k": [1, 2]}"#), + Some(r#"{"k": null}"#), + Some(r#"{"k": 1}"#), + Some(r#"{"k": 2.5}"#), + Some(r#"{"k": "line\nbreak"}"#), + Some(r#"{"a": "1", "b": 2}"#), + Some(r#"{"a.b\"c": 1}"#), + Some(r#"{"a'b": "v'1"}"#), + Some(r#"{"k": {"x": [1, true, "text"]}}"#), + None, + ]; + let dataset = dataset(&dir, &metas).await; + let k = r#"`metadata`, '$["k"]'"#; + let filters = [ + format!(r#"json_extract({k}) = '"v"'"#), + format!("TRY_CAST(json_extract({k}) AS DOUBLE) = 3"), + format!("json_extract({k}) = 'true'"), + format!(r#"NOT COALESCE(json_extract({k}) = '"v"', false)"#), + format!("json_exists({k})"), + format!("NOT json_exists({k})"), + r#"json_exists(`metadata`, '$["a.b\"c"]')"#.to_string(), + format!("TRY_CAST(json_extract({k}) AS DOUBLE) > 1"), + format!( + r#"(COALESCE(json_extract({k}) in ('"a"'), false) OR COALESCE(json_extract({k}) in ('"b"'), false))"# + ), + format!("NOT (COALESCE(TRY_CAST(json_extract({k}) AS DOUBLE) in (1), false))"), + format!( + r#"(COALESCE(json_extract({k}) in ('"a"', 'true'), false) OR COALESCE(TRY_CAST(json_extract({k}) AS DOUBLE) in (1, 2.5), false))"# + ), + r#"(json_extract(`metadata`, '$["a"]') = '"1"' AND TRY_CAST(json_extract(`metadata`, '$["b"]') AS DOUBLE) = 2)"#.to_string(), + r#"json_extract(`metadata`, '$["a''b"]') = '"v''1"'"#.to_string(), + format!(r#"json_extract({k}) = '"line\nbreak"'"#), + format!(r#"json_extract({k}) = '{{"x":[1,true,"text"]}}'"#), + format!("json_extract({k}) IS NULL"), + "json_extract(`metadata`, '$.*') IS NOT NULL".to_string(), + ]; + for filter in &filters { + let want = keys(&dataset, filter, false).await.unwrap(); + let got = keys(&dataset, filter, true).await.unwrap(); + assert_eq!(got, want, "{filter}"); + } + // A bad path fails the query either way. + let bad = "json_extract(`metadata`, '$[') = 'x'"; + assert!(keys(&dataset, bad, false).await.is_err()); + assert!(keys(&dataset, bad, true).await.is_err()); + } + + #[tokio::test] + async fn apply_swaps_literal_paths_only() { + let dir = TempDir::new().unwrap(); + let dataset = dataset(&dir, &[Some(r#"{"k": "v"}"#)]).await; + let mut scanner = dataset.scan(); + apply( + &mut scanner, + r#"json_extract(`metadata`, '$["k"]') = '"v"' AND json_exists(`metadata`, '$["k"]') AND json_get_string(`metadata`, 'k') = 'v'"#, + ) + .unwrap(); + let mut swapped = Vec::new(); + scanner + .get_expr_filter() + .unwrap() + .unwrap() + .apply(|e| { + if let Expr::ScalarFunction(f) = e { + let inner: &dyn std::any::Any = f.func.inner().as_ref(); + swapped.push((f.func.name().to_string(), inner.is::())); + } + Ok(lance::deps::datafusion::common::tree_node::TreeNodeRecursion::Continue) + }) + .unwrap(); + swapped.sort(); + assert_eq!( + swapped, + [ + ("json_exists".to_string(), true), + ("json_extract".to_string(), true), + ("json_get_string".to_string(), false), + ] + ); + } +} diff --git a/seaweed-vector/crates/core/src/lib.rs b/seaweed-vector/crates/core/src/lib.rs index def0eb215f..5d8013ed24 100644 --- a/seaweed-vector/crates/core/src/lib.rs +++ b/seaweed-vector/crates/core/src/lib.rs @@ -6,6 +6,7 @@ //! it does not, always covering the fragments no index has reached yet. pub mod coverage; +mod filter; pub mod maintenance; pub mod query; pub mod segments; diff --git a/seaweed-vector/crates/core/src/query.rs b/seaweed-vector/crates/core/src/query.rs index ea6cf15e80..9cf7ef9ee7 100644 --- a/seaweed-vector/crates/core/src/query.rs +++ b/seaweed-vector/crates/core/src/query.rs @@ -117,7 +117,7 @@ async fn scanner( .use_index(use_index) .with_row_id(); if let Some(filter) = &query.filter { - scanner.filter(filter)?; + crate::filter::apply(&mut scanner, filter)?; scanner.prefilter(!query.postfilter); } if let Some(nprobes) = query.nprobes { diff --git a/test/s3vectors/s3vectors_test.py b/test/s3vectors/s3vectors_test.py index 365f74060a..887d7e3191 100644 --- a/test/s3vectors/s3vectors_test.py +++ b/test/s3vectors/s3vectors_test.py @@ -198,6 +198,10 @@ def test_filter_types(vectors, index): hits = vectors.query_vectors(**index, queryVector={"float32": [1, 0, 0, 0]}, topK=4, filter={"value": {"$eq": value}})["vectors"] assert [hit["key"] for hit in hits] == [str(i)] + for op, members in (("$in", [1, "true"]), ("$nin", ["1", True])): + hits = vectors.query_vectors(**index, queryVector={"float32": [1, 0, 0, 0]}, topK=4, + filter={"value": {op: members}})["vectors"] + assert sorted(hit["key"] for hit in hits) == ["0", "3"] def test_cosine_metric(vectors): diff --git a/weed/s3api/s3vectors/handlers.go b/weed/s3api/s3vectors/handlers.go index b02b17b8f3..9b3704aa60 100644 --- a/weed/s3api/s3vectors/handlers.go +++ b/weed/s3api/s3vectors/handlers.go @@ -1723,20 +1723,29 @@ func fieldClause(field string, cond interface{}) (string, error) { } continue } - // Members may mix types; each extractor only sees its own. A - // member whose typed extractor returns NULL must not let NOT + // Members may mix types; each extractor only sees its own, so + // members sharing one go into a single IN that reads the field + // once. A group whose extractor returns NULL must not let NOT // invert NULL into NULL under $nin, so every group coalesces. - var groups []string + var cols []string + members := map[string][]string{} for _, item := range list { lit, err := sqlLiteral(item) if err != nil { return "", err } - g, err := fieldMatch(field, item, "in ("+lit+")") + col, err := fieldAccess(field, item) if err != nil { return "", err } - groups = append(groups, fmt.Sprintf("COALESCE(%s, false)", g)) + if _, seen := members[col]; !seen { + cols = append(cols, col) + } + members[col] = append(members[col], lit) + } + groups := make([]string, 0, len(cols)) + for _, col := range cols { + groups = append(groups, fmt.Sprintf("COALESCE(%s in (%s), false)", col, strings.Join(members[col], ", "))) } expr := "(" + strings.Join(groups, " OR ") + ")" if op == "$nin" { diff --git a/weed/s3api/s3vectors/handlers_test.go b/weed/s3api/s3vectors/handlers_test.go index d7ceccf683..7aad7447bb 100644 --- a/weed/s3api/s3vectors/handlers_test.go +++ b/weed/s3api/s3vectors/handlers_test.go @@ -374,9 +374,11 @@ func TestTranslateFilter(t *testing.T) { {"range", map[string]interface{}{"k": map[string]interface{}{"$gt": float64(1)}}, "TRY_CAST(json_extract(`metadata`, '$[\"k\"]') AS DOUBLE) > 1"}, {"in strings", map[string]interface{}{"k": map[string]interface{}{"$in": []interface{}{"a", "b"}}}, - "(COALESCE(json_extract(`metadata`, '$[\"k\"]') in ('\"a\"'), false) OR COALESCE(json_extract(`metadata`, '$[\"k\"]') in ('\"b\"'), false))"}, + "(COALESCE(json_extract(`metadata`, '$[\"k\"]') in ('\"a\"', '\"b\"'), false))"}, {"nin numbers", map[string]interface{}{"k": map[string]interface{}{"$nin": []interface{}{float64(1)}}}, "NOT (COALESCE(TRY_CAST(json_extract(`metadata`, '$[\"k\"]') AS DOUBLE) in (1), false))"}, + {"in mixed types", map[string]interface{}{"k": map[string]interface{}{"$in": []interface{}{"a", float64(1), true, float64(2)}}}, + "(COALESCE(json_extract(`metadata`, '$[\"k\"]') in ('\"a\"', 'true'), false) OR COALESCE(TRY_CAST(json_extract(`metadata`, '$[\"k\"]') AS DOUBLE) in (1, 2), false))"}, {"and", map[string]interface{}{"$and": []interface{}{ map[string]interface{}{"a": "1"}, map[string]interface{}{"b": float64(2)}, }},