From 06b9d014237b5709cf2d218d879599117be56281 Mon Sep 17 00:00:00 2001 From: Chris Lu Date: Sun, 1 Nov 2026 18:44:49 -0800 Subject: [PATCH] s3vectors: match a filter value against the elements of an array (#12321) * vector: compare json_array_contains elements as typed JSON values Lance's json_array_contains compares each element's text with the value, or with the value wrapped in quotes, so 1 matches "1", true matches "true", and 1.0 misses a stored 1. The filter swap now replaces it like json_extract and json_exists: the value is parsed as JSON once per batch and compared with jsonb's typed equality, and only literal arguments are swapped. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK * s3vectors: match a filter value against the elements of an array A filter never looked inside an array: {"tags": "space"} and {"tags": {"$in": ["space"]}} both missed "tags": ["space", "classic"]. AWS documents that $eq, and so a bare value, matches an array metadata value when any element equals the input. $eq, a bare value and each $in member now also match through json_array_contains; $ne and $nin remain their negations. Range operators, $exists and nested objects are unchanged. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK --------- Co-authored-by: Claude Opus 5.5 --- seaweed-vector/crates/core/src/filter.rs | 127 +++++++++++++++++++---- test/s3vectors/s3vectors_test.py | 19 ++++ weed/s3api/s3vectors/handlers.go | 52 ++++++++-- weed/s3api/s3vectors/handlers_test.go | 31 ++++-- 4 files changed, 190 insertions(+), 39 deletions(-) diff --git a/seaweed-vector/crates/core/src/filter.rs b/seaweed-vector/crates/core/src/filter.rs index da1fabf458..9613e9cafa 100644 --- a/seaweed-vector/crates/core/src/filter.rs +++ b/seaweed-vector/crates/core/src/filter.rs @@ -1,14 +1,16 @@ //! 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. +//! batch and return the same values. json_array_contains is swapped too, and +//! compares the elements as typed JSON values: Lance's compares their text, so +//! 1 would match "1". 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 jsonb::{OwnedJsonb, RawJsonb}; use lance::dataset::scanner::Scanner; use lance::deps::datafusion::common::tree_node::{Transformed, TreeNode}; use lance::deps::datafusion::common::{DataFusionError, Result, ScalarValue}; @@ -34,12 +36,16 @@ fn swap(expr: Expr) -> Result { let op = match f.func.name() { "json_extract" => Op::Extract, "json_exists" => Op::Exists, + "json_array_contains" => Op::Contains, _ => return Ok(Transformed::no(Expr::ScalarFunction(f))), }; - if !matches!( - f.args.get(1), - Some(Expr::Literal(ScalarValue::Utf8(Some(_)), _)) - ) { + let literal = |i: usize| { + matches!( + f.args.get(i), + Some(Expr::Literal(ScalarValue::Utf8(Some(_)), _)) + ) + }; + if !literal(1) || (op == Op::Contains && !literal(2)) { return Ok(Transformed::no(Expr::ScalarFunction(f))); } let udf = Arc::new(ScalarUDF::new_from_impl(JsonPath::new(op))); @@ -54,6 +60,7 @@ fn swap(expr: Expr) -> Result { enum Op { Extract, Exists, + Contains, } #[derive(Debug, PartialEq, Eq, Hash)] @@ -64,12 +71,13 @@ struct JsonPath { impl JsonPath { fn new(op: Op) -> Self { + let mut args = vec![DataType::LargeBinary, DataType::Utf8]; + if op == Op::Contains { + args.push(DataType::Utf8); + } Self { op, - signature: Signature::exact( - vec![DataType::LargeBinary, DataType::Utf8], - Volatility::Immutable, - ), + signature: Signature::exact(args, Volatility::Immutable), } } } @@ -83,6 +91,7 @@ impl ScalarUDFImpl for JsonPath { match self.op { Op::Extract => "json_extract", Op::Exists => "json_exists", + Op::Contains => "json_array_contains", } } @@ -93,17 +102,24 @@ impl ScalarUDFImpl for JsonPath { fn return_type(&self, _: &[DataType]) -> Result { Ok(match self.op { Op::Extract => DataType::Utf8, - Op::Exists => DataType::Boolean, + Op::Exists | Op::Contains => 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, path, value) = match args.args.as_slice() { + [data, ColumnarValue::Scalar(ScalarValue::Utf8(Some(path)))] => (data, path, None), + [ + data, + ColumnarValue::Scalar(ScalarValue::Utf8(Some(path))), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(value))), + ] => (data, path, Some(value)), + _ => { + return Err(exec_err(format!( + "{} expects literal arguments", + self.name() + ))); + } }; let data = data.to_array(args.number_rows)?; let data = data @@ -160,6 +176,37 @@ impl ScalarUDFImpl for JsonPath { } Arc::new(builder.finish()) } + Op::Contains => { + let value = value.map(String::as_str).unwrap_or_default(); + let value: OwnedJsonb = value + .parse() + .map_err(|e| exec_err(format!("Invalid JSON value '{value}': {e}")))?; + 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))); + let found = match selector.select_value(&json_path) { + Ok(Some(array)) => array + .as_raw() + .array_values() + .map_err(|e| { + exec_err(format!("Failed to read array at '{path}': {e}")) + })? + .is_some_and(|items| items.contains(&value)), + Ok(None) => false, + Err(e) => { + return Err(exec_err(format!( + "Failed to check array contains at path '{path}': {e}" + ))); + } + }; + builder.append_value(found); + } + Arc::new(builder.finish()) + } }; Ok(ColumnarValue::Array(out)) } @@ -298,6 +345,48 @@ mod tests { assert!(keys(&dataset, bad, true).await.is_err()); } + #[tokio::test] + async fn array_contains_compares_typed_values() { + let dir = TempDir::new().unwrap(); + let metas = [ + Some(r#"{"k": ["space", "classic"]}"#), + Some(r#"{"k": [1, 2.5]}"#), + Some(r#"{"k": ["1", "2.5"]}"#), + Some(r#"{"k": [true]}"#), + Some(r#"{"k": ["true"]}"#), + Some(r#"{"k": "space"}"#), + Some(r#"{"k": [["space"]]}"#), + Some(r#"{"k": {"a": "space"}}"#), + Some(r#"{"k": ["", "spa\"ce"]}"#), + Some(r#"{"k": []}"#), + Some(r#"{"x": 1}"#), + None, + ]; + let dataset = dataset(&dir, &metas).await; + for (value, want) in [ + (r#""space""#, vec!["k0"]), + ("1", vec!["k1"]), + ("1.0", vec!["k1"]), + ("2.5", vec!["k1"]), + (r#""1""#, vec!["k2"]), + ("true", vec!["k3"]), + (r#""true""#, vec!["k4"]), + (r#""""#, vec!["k8"]), + (r#""spa\"ce""#, vec!["k8"]), + (r#""classic""#, vec!["k0"]), + ("3", vec![]), + ] { + let filter = format!(r#"json_array_contains(`metadata`, '$["k"]', '{value}')"#); + assert_eq!( + keys(&dataset, &filter, true).await.unwrap(), + want, + "{filter}" + ); + } + let bad = r#"json_array_contains(`metadata`, '$["k"]', 'not json')"#; + assert!(keys(&dataset, bad, true).await.is_err()); + } + #[tokio::test] async fn apply_swaps_literal_paths_only() { let dir = TempDir::new().unwrap(); @@ -305,7 +394,7 @@ mod tests { 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'"#, + r#"json_extract(`metadata`, '$["k"]') = '"v"' AND json_exists(`metadata`, '$["k"]') AND json_get_string(`metadata`, 'k') = 'v' AND json_array_contains(`metadata`, '$["k"]', '"v"') AND NOT json_array_contains(`metadata`, '$["k"]', key)"#, ) .unwrap(); let mut swapped = Vec::new(); @@ -325,6 +414,8 @@ mod tests { assert_eq!( swapped, [ + ("json_array_contains".to_string(), false), + ("json_array_contains".to_string(), true), ("json_exists".to_string(), true), ("json_extract".to_string(), true), ("json_get_string".to_string(), false), diff --git a/test/s3vectors/s3vectors_test.py b/test/s3vectors/s3vectors_test.py index 887d7e3191..e8138e7e35 100644 --- a/test/s3vectors/s3vectors_test.py +++ b/test/s3vectors/s3vectors_test.py @@ -204,6 +204,25 @@ def test_filter_types(vectors, index): assert sorted(hit["key"] for hit in hits) == ["0", "3"] +def test_filter_arrays(vectors, index): + rows = {"space": ["space", "classic"], "nums": [1, 2.5], "strs": ["1", "2.5"], "flag": [True], "scalar": "space"} + vectors.put_vectors(**index, vectors=[{"key": key, "data": {"float32": [1, 0, 0, 0]}, "metadata": {"tags": tags}} + for key, tags in rows.items()]) + + def matches(condition): + hits = vectors.query_vectors(**index, queryVector={"float32": [1, 0, 0, 0]}, topK=10, + filter={"tags": condition})["vectors"] + return sorted(hit["key"] for hit in hits) + + assert matches("space") == ["scalar", "space"] + assert matches({"$eq": 1}) == ["nums"] + assert matches({"$eq": "1"}) == ["strs"] + assert matches({"$eq": True}) == ["flag"] + assert matches({"$in": ["classic", 2.5]}) == ["nums", "space"] + assert matches({"$ne": "space"}) == ["flag", "nums", "strs"] + assert matches({"$nin": [1, "2.5"]}) == ["flag", "scalar", "space"] + + def test_cosine_metric(vectors): bucket = "cosine-" + uuid.uuid4().hex[:16] vectors.create_vector_bucket(vectorBucketName=bucket) diff --git a/weed/s3api/s3vectors/handlers.go b/weed/s3api/s3vectors/handlers.go index d647668b18..77213ba2f9 100644 --- a/weed/s3api/s3vectors/handlers.go +++ b/weed/s3api/s3vectors/handlers.go @@ -1712,16 +1712,12 @@ func fieldClause(field string, cond interface{}) (string, error) { return "", fmt.Errorf("$exists expects a boolean") } case "$eq", "$ne": - lit, err := sqlLiteral(operand) - if err != nil { - return "", err - } - match, err := fieldMatch(field, operand, "= "+lit) + match, err := equalMatch(field, operand) if err != nil { return "", err } if op == "$ne" { - match = "NOT COALESCE(" + match + ", false)" + match = "NOT " + match } parts = append(parts, match) case "$in", "$nin": @@ -1761,6 +1757,13 @@ func fieldClause(field string, cond interface{}) (string, error) { for _, col := range cols { groups = append(groups, fmt.Sprintf("COALESCE(%s in (%s), false)", col, strings.Join(members[col], ", "))) } + for _, item := range list { + element, err := elementMatch(field, item) + if err != nil { + return "", err + } + groups = append(groups, element) + } expr := "(" + strings.Join(groups, " OR ") + ")" if op == "$nin" { expr = "NOT " + expr @@ -1789,14 +1792,41 @@ func fieldClause(field string, cond interface{}) (string, error) { } return strings.Join(parts, " AND "), nil default: - lit, err := sqlLiteral(cond) - if err != nil { - return "", err - } - return fieldMatch(field, cond, "= "+lit) + return equalMatch(field, cond) } } +// equalMatch renders "field equals operand": the stored value itself or, as +// AWS documents for $eq, any element of a stored array. +func equalMatch(field string, operand interface{}) (string, error) { + lit, err := sqlLiteral(operand) + if err != nil { + return "", err + } + match, err := fieldMatch(field, operand, "= "+lit) + if err != nil { + return "", err + } + element, err := elementMatch(field, operand) + if err != nil { + return "", err + } + return "(COALESCE(" + match + ", false) OR " + element + ")", nil +} + +// elementMatch renders "the array stored at field holds operand". The engine +// compares the elements as typed JSON values, so 1 does not match "1". +func elementMatch(field string, operand interface{}) (string, error) { + var b strings.Builder + enc := json.NewEncoder(&b) + enc.SetEscapeHTML(false) + if err := enc.Encode(operand); err != nil { + return "", err + } + return fmt.Sprintf("COALESCE(json_array_contains(`metadata`, '%s', '%s'), false)", + sqlEscape(existsPath(field)), sqlEscape(strings.TrimSuffix(b.String(), "\n"))), nil +} + // jsonbString renders s as the JSON text json_extract returns for a stored // string: quoted, with the same escapes the jsonb serializer emits. func jsonbString(s string) string { diff --git a/weed/s3api/s3vectors/handlers_test.go b/weed/s3api/s3vectors/handlers_test.go index 32dbef2d70..76849850f6 100644 --- a/weed/s3api/s3vectors/handlers_test.go +++ b/weed/s3api/s3vectors/handlers_test.go @@ -353,19 +353,26 @@ func TestResolveIndexARNs(t *testing.T) { } func TestTranslateFilter(t *testing.T) { + str := "json_extract(`metadata`, '$[\"k\"]')" + num := "TRY_CAST(json_extract(`metadata`, '$[\"k\"]') AS DOUBLE)" + elem := func(value string) string { + return "COALESCE(json_array_contains(`metadata`, '$[\"k\"]', '" + value + "'), false)" + } cases := []struct { name string filter map[string]interface{} want string }{ {"bare string equality", map[string]interface{}{"k": "v"}, - "json_extract(`metadata`, '$[\"k\"]') = '\"v\"'"}, + "(COALESCE(" + str + " = '\"v\"', false) OR " + elem(`"v"`) + ")"}, {"bare number equality", map[string]interface{}{"k": float64(3)}, - "TRY_CAST(json_extract(`metadata`, '$[\"k\"]') AS DOUBLE) = 3"}, + "(COALESCE(" + num + " = 3, false) OR " + elem("3") + ")"}, {"bare bool equality", map[string]interface{}{"k": true}, - "json_extract(`metadata`, '$[\"k\"]') = 'true'"}, + "(COALESCE(" + str + " = 'true', false) OR " + elem("true") + ")"}, + {"eq fraction", map[string]interface{}{"k": map[string]interface{}{"$eq": float64(2.5)}}, + "(COALESCE(" + num + " = 2.5, false) OR " + elem("2.5") + ")"}, {"ne", map[string]interface{}{"k": map[string]interface{}{"$ne": "v"}}, - "NOT COALESCE(json_extract(`metadata`, '$[\"k\"]') = '\"v\"', false)"}, + "NOT (COALESCE(" + str + " = '\"v\"', false) OR " + elem(`"v"`) + ")"}, {"exists true", map[string]interface{}{"k": map[string]interface{}{"$exists": true}}, "json_exists(`metadata`, '$[\"k\"]')"}, {"exists false", map[string]interface{}{"k": map[string]interface{}{"$exists": false}}, @@ -373,19 +380,23 @@ func TestTranslateFilter(t *testing.T) { {"exists special key", map[string]interface{}{`a.b"c`: map[string]interface{}{"$exists": true}}, "json_exists(`metadata`, '$[\"a.b\\\"c\"]')"}, {"range", map[string]interface{}{"k": map[string]interface{}{"$gt": float64(1)}}, - "TRY_CAST(json_extract(`metadata`, '$[\"k\"]') AS DOUBLE) > 1"}, + num + " > 1"}, {"in strings", map[string]interface{}{"k": map[string]interface{}{"$in": []interface{}{"a", "b"}}}, - "(COALESCE(json_extract(`metadata`, '$[\"k\"]') in ('\"a\"', '\"b\"'), false))"}, + "(COALESCE(" + str + " in ('\"a\"', '\"b\"'), false) OR " + elem(`"a"`) + " OR " + elem(`"b"`) + ")"}, {"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))"}, + "NOT (COALESCE(" + num + " in (1), false) OR " + elem("1") + ")"}, {"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))"}, + "(COALESCE(" + str + " in ('\"a\"', 'true'), false) OR COALESCE(" + num + " in (1, 2), false) OR " + + elem(`"a"`) + " OR " + elem("1") + " OR " + elem("true") + " OR " + elem("2") + ")"}, {"and", map[string]interface{}{"$and": []interface{}{ map[string]interface{}{"a": "1"}, map[string]interface{}{"b": float64(2)}, }}, - "(json_extract(`metadata`, '$[\"a\"]') = '\"1\"' AND TRY_CAST(json_extract(`metadata`, '$[\"b\"]') AS DOUBLE) = 2)"}, + "((COALESCE(json_extract(`metadata`, '$[\"a\"]') = '\"1\"', false) OR COALESCE(json_array_contains(`metadata`, '$[\"a\"]', '\"1\"'), false)) AND " + + "(COALESCE(TRY_CAST(json_extract(`metadata`, '$[\"b\"]') AS DOUBLE) = 2, false) OR COALESCE(json_array_contains(`metadata`, '$[\"b\"]', '2'), false)))"}, {"sql injection", map[string]interface{}{"a'b": "v'1"}, - "json_extract(`metadata`, '$[\"a''b\"]') = '\"v''1\"'"}, + "(COALESCE(json_extract(`metadata`, '$[\"a''b\"]') = '\"v''1\"', false) OR COALESCE(json_array_contains(`metadata`, '$[\"a''b\"]', '\"v''1\"'), false))"}, + {"array value escapes", map[string]interface{}{"k": ``}, + "(COALESCE(" + str + ` = '""', false) OR ` + elem(`""`) + ")"}, } for _, tc := range cases { got, err := translateFilter(tc.filter)