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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011UpvafvK8JpagF94MHWqwK

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

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<Expr> {
         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<Expr> {
 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<DataType> {
         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<ColumnarValue> {
-        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),

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)

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 {

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": `<a&"b">`},
+			"(COALESCE(" + str + ` = '"<a&\"b\">"', false) OR ` + elem(`"<a&\"b\">"`) + ")"},
 	}
 	for _, tc := range cases {
 		got, err := translateFilter(tc.filter)