Vectorset cmds

docs/supported-commands/Redis/VECTOR_SET.md

@@ -1,51 +1,53 @@
+# Redis `vector_set` commands (9/12 implemented)
 
-## Unsupported vector_set commands
-> To implement support for a command, see [here](/guides/implement-command/)
-
-#### [VADD](https://redis.io/commands/vadd/) <small>(not implemented)</small>
+## [VADD](https://redis.io/commands/vadd/)
 
 Add a new element to a vector set, or update its vector if it already exists.
 
-#### [VCARD](https://redis.io/commands/vcard/) <small>(not implemented)</small>
+## [VCARD](https://redis.io/commands/vcard/)
 
 Return the number of elements in a vector set.
 
-#### [VDIM](https://redis.io/commands/vdim/) <small>(not implemented)</small>
+## [VDIM](https://redis.io/commands/vdim/)
 
 Return the dimension of vectors in the vector set.
 
-#### [VEMB](https://redis.io/commands/vemb/) <small>(not implemented)</small>
+## [VEMB](https://redis.io/commands/vemb/)
 
 Return the vector associated with an element.
 
-#### [VGETATTR](https://redis.io/commands/vgetattr/) <small>(not implemented)</small>
+## [VGETATTR](https://redis.io/commands/vgetattr/)
 
 Retrieve the JSON attributes of elements.
 
-#### [VINFO](https://redis.io/commands/vinfo/) <small>(not implemented)</small>
-
-Return information about a vector set.
-
-#### [VLINKS](https://redis.io/commands/vlinks/) <small>(not implemented)</small>
-
-Return the neighbors of an element at each layer in the HNSW graph.
-
-#### [VRANDMEMBER](https://redis.io/commands/vrandmember/) <small>(not implemented)</small>
+## [VRANDMEMBER](https://redis.io/commands/vrandmember/)
 
 Return one or multiple random members from a vector set.
 
-#### [VRANGE](https://redis.io/commands/vrange/) <small>(not implemented)</small>
+## [VRANGE](https://redis.io/commands/vrange/)
 
 Return elements in a lexicographical range
 
-#### [VREM](https://redis.io/commands/vrem/) <small>(not implemented)</small>
+## [VREM](https://redis.io/commands/vrem/)
 
 Remove an element from a vector set.
 
-#### [VSETATTR](https://redis.io/commands/vsetattr/) <small>(not implemented)</small>
+## [VSETATTR](https://redis.io/commands/vsetattr/)
 
 Associate or remove the JSON attributes of elements.
 
+
+## Unsupported vector_set commands
+> To implement support for a command, see [here](/guides/implement-command/)
+
+#### [VINFO](https://redis.io/commands/vinfo/) <small>(not implemented)</small>
+
+Return information about a vector set.
+
+#### [VLINKS](https://redis.io/commands/vlinks/) <small>(not implemented)</small>
+
+Return the neighbors of an element at each layer in the HNSW graph.
+
 #### [VSIM](https://redis.io/commands/vsim/) <small>(not implemented)</small>
 
 Return elements by vector similarity.

fakeredis/_commands.py

@@ -319,6 +319,10 @@ def __init__(self, value: Union[bytes, BeforeAny, AfterAny], exclusive: bool):
         self.value = value
         self.exclusive = exclusive
 
+    @property
+    def inclusive(self) -> bool:
+        return not self.exclusive
+
     @classmethod
     def decode(cls, value: bytes) -> "StringTest":
         if value == b"-":

fakeredis/_fakesocket.py

@@ -24,6 +24,7 @@
     TopkCommandsMixin,
     TDigestCommandsMixin,
     TimeSeriesCommandsMixin,
+    VectorSetCommandsMixin,
 )
 from ._basefakesocket import BaseFakeSocket
 from ._server import FakeServer
@@ -56,6 +57,7 @@ class FakeSocket(
     TimeSeriesCommandsMixin,
     DragonflyCommandsMixin,
     AclCommandsMixin,
+    VectorSetCommandsMixin,
 ):
     def __init__(
         self,

fakeredis/_typing.py

@@ -1,5 +1,5 @@
 import sys
-from typing import Tuple, Union
+from typing import Tuple, Union, Dict, Any, List
 
 import redis
 
@@ -22,9 +22,10 @@
 lib_version = metadata.version("fakeredis")
 VersionType = Tuple[int, ...]
 ServerType = Literal["redis", "dragonfly", "valkey"]
-
+JsonType = Union[str, int, float, bool, None, Dict[str, Any], List[Any]]
 RaiseErrorTypes = (redis.ResponseError, redis.AuthenticationError)
 ResponseErrorType = redis.ResponseError
+
 try:
     import valkey
 

fakeredis/model/__init__.py

@@ -38,6 +38,14 @@
     "TDigest",
 ]
 
+try:
+    import numpy as np  # noqa: F401
+    from ._vectorset import VectorSet, Vector  # noqa: F401
+
+    __all__.extend(["VectorSet", "Vector"])
+except ImportError:
+    pass
+
 try:
     import probables  # noqa: F401
     from ._filters import ScalableCuckooFilter, ScalableBloomFilter  # noqa: F401

fakeredis/model/_vectorset.py

@@ -0,0 +1,181 @@
+import json
+import re
+import struct
+from typing import List, Dict, Any, Literal, Optional, Iterator, Self, Union
+
+import numpy as np
+from jsonpath_ng import JSONPath
+from jsonpath_ng.exceptions import JsonPathParserError
+from jsonpath_ng.ext import parse
+
+from fakeredis import _msgs as msgs
+from fakeredis._helpers import SimpleError
+
+QUANTIZATION_TYPE = Literal["noquant", "bin", "int8"]
+
+
+def _update_to_jsonpath_format(path: Union[bytes, str]) -> str:
+    path_str = path.decode() if isinstance(path, bytes) else path
+    path_str = path_str.replace("and", "&").replace("or", "|").replace("not", "!").replace(".", "@.")
+
+    # Replace `v in [x, y, z]` with `(v=~'x|y|z')`
+    def expand_in(m: re.Match) -> str:
+        var = m.group(1)
+        items = [item.strip().replace("'", "") for item in m.group(2).split(",")]
+        return f"({var}=~'{'|'.join(items)}')"
+
+    path_str = re.sub(r"(\S+)\s+in\s+\[([^]]+)]", expand_in, path_str)
+
+    return f"$[?({path_str})]"
+
+
+def _parse_jsonfilter(path: Union[str, bytes]) -> JSONPath:
+    path_str: str = _update_to_jsonpath_format(path)
+    try:
+        return parse(path_str)
+    except JsonPathParserError:
+        raise SimpleError(msgs.JSON_PATH_DOES_NOT_EXIST.format(path_str))
+
+
+def quantize_int8(x):
+    qmin = -(2.0**7) if (x < 0).any() else 0  # Signed or unsigned range
+    qmax = 2.0**7 - 1 if (x < 0).any() else 2.0**8 - 1
+
+    min_val, max_val = x.min(), x.max()
+
+    # Calculate the scale factor
+    scale = (max_val - min_val) / (qmax - qmin)
+
+    # Calculate the initial zero point and clamp it to the valid range
+    initial_zero_point = qmin - min_val / scale
+    zero_point = int(np.clip(initial_zero_point, qmin, qmax))
+
+    # Quantize the values and round them
+    q_x = zero_point + x / scale
+    q_x = np.clip(q_x, qmin, qmax).round().astype(np.int8)  # Use np.int8 for the final data type
+
+    return q_x, scale, zero_point
+
+
+class Vector:
+    def __init__(
+        self, name: bytes, values: List[float], attributes: Optional[bytes], quantization: QUANTIZATION_TYPE, ef: int
+    ) -> None:
+        self.name = name
+        self.values = values
+        self.attributes = attributes
+        self.quantization = quantization
+        self.l2_norm = sum(v * v for v in values) ** 0.5
+        if self.quantization == "bin":
+            self.values = [1 if v > 0 else -1 for v in self.values]
+
+    def __repr__(self):
+        return f"Vector(name={self.name}, values={self.values}, attributes={self.attributes}, quantization={self.quantization})"
+
+    def __hash__(self):
+        return hash(self.name)
+
+    @classmethod
+    def from_vector_values(cls, values: List[float]) -> Self:
+        return cls("", values, b"", "int8", 0)
+
+    def raw(self) -> List[Any]:
+        raw_bytes = struct.pack(f"{len(self.values)}f", *self.values)
+        if self.quantization == "int8":
+            q_values, scale, zero_point = quantize_int8(np.array(self.values))
+            return [self.quantization.encode(), raw_bytes, self.l2_norm, scale, zero_point]
+        if self.quantization == "bin":
+            return [self.quantization.encode(), raw_bytes, self.l2_norm]
+
+        return [b"f32", raw_bytes, self.l2_norm]
+
+    def similarity(self, other: Self) -> float:
+        me = np.array(self.values)
+        o = np.array(other.values)
+        return float(np.dot(me, o)) / (self.l2_norm * other.l2_norm)
+
+    def accept_filter(self, filter_expression: Optional[bytes]) -> bool:
+        if filter_expression is None:
+            return True
+        if self.attributes is None:
+            return False
+        json_obj = json.loads(self.attributes)
+        return len(_parse_jsonfilter(filter_expression).find([json_obj])) > 0
+
+
+class VectorSet:
+    def __init__(self, dimensions: int):
+        self._dimensions = dimensions
+        self._vectors: Dict[bytes, Vector] = dict()
+        self._links: Dict[bytes, int] = dict()
+
+    @property
+    def dimensions(self) -> int:
+        return self._dimensions
+
+    @property
+    def card(self) -> int:
+        return len(self._vectors)
+
+    def vector_names(self) -> List[bytes]:
+        return list(self._vectors.keys())
+
+    def exists(self, name: bytes) -> bool:
+        return name in self._vectors
+
+    def add(self, vector: Vector, numlinks: int) -> None:
+        self._vectors[vector.name] = vector
+        self._links[vector.name] = numlinks  # type: ignore
+
+    def remove(self, name: bytes) -> int:
+        if name not in self._vectors:
+            return 0
+        del self._vectors[name]
+        del self._links[name]
+        return 1
+
+    def info(self) -> Dict[str, Any]:
+        return {
+            "quant-type": "fp32",
+            "vector-dim": self._dimensions,
+            "size": len(self._vectors),
+            "max-level": 0,  # TODO
+            "vset-uid": 1,  # TODO
+            "hnsw-max-node-uid": 0,  # TODO
+        }
+
+    def range(
+        self,
+        min_value: Optional[bytes],
+        include_min: bool,
+        max_value: Optional[bytes],
+        include_max: bool,
+        count: Optional[int],
+    ) -> List[bytes]:
+        if count is not None and count < 0:
+            count = None
+        res: List[bytes] = []
+        for name in self._vectors.keys():
+            if (min_value is None or name > min_value or (include_min and name == min_value)) and (
+                max_value is None or name < max_value or (include_max and name == max_value)
+            ):
+                res.append(name)
+            if count is not None and len(res) >= count:
+                break
+        return res
+
+    def __contains__(self, k: bytes) -> bool:
+        return k in self._vectors
+
+    def __getitem__(self, k: bytes) -> Vector:
+        if k not in self._vectors:
+            raise KeyError(f"Vector with name {k} does not exist.")
+        return self._vectors[k]
+
+    def __iter__(self) -> Iterator[Vector]:
+        return iter(self._vectors.values())
+
+    def get(self, k: bytes) -> Optional[Vector]:
+        if k in self._vectors:
+            return self._vectors[k]
+        return None

fakeredis/stack/__init__.py

@@ -2,10 +2,19 @@
 from ._timeseries_mixin import TimeSeriesCommandsMixin
 from ._topk_mixin import TopkCommandsMixin  # noqa: F401
 
+try:
+    import numpy  # noqa: F401
+    from ._vectorset_mixin import VectorSetCommandsMixin  # noqa: F401
+except ImportError:
+
+    class VectorSetCommandsMixin:
+        pass
+
+
 try:
     from jsonpath_ng.ext import parse  # noqa: F401
     from redis.commands.json.path import Path  # noqa: F401
-    from ._json_mixin import JSONCommandsMixin, JSONObject  # noqa: F401
+    from ._json_mixin import JSONCommandsMixin  # noqa: F401
 except ImportError as e:
     if e.name == "fakeredis.stack._json_mixin":
         raise e
@@ -37,10 +46,10 @@ class CMSCommandsMixin:  # type: ignore # noqa: E303
 __all__ = [
     "TopkCommandsMixin",
     "JSONCommandsMixin",
-    "JSONObject",
     "BFCommandsMixin",
     "CFCommandsMixin",
     "CMSCommandsMixin",
     "TDigestCommandsMixin",
     "TimeSeriesCommandsMixin",
+    "VectorSetCommandsMixin",
 ]

fakeredis/stack/_json_mixin.py

@@ -15,8 +15,7 @@
 from fakeredis._commands import Key, command, delete_keys, CommandItem, Int, Float
 from fakeredis._helpers import SimpleString
 from fakeredis.model import ZSet, ClientInfo
-
-JsonType = Union[str, int, float, bool, None, Dict[str, Any], List[Any]]
+from fakeredis._typing import JsonType
 
 
 def _format_path(path: Union[bytes, str]) -> str:
@@ -75,8 +74,7 @@ def decode(cls, value: bytes) -> Any:
 
     @classmethod
     def encode(cls, value: Any) -> Optional[bytes]:
-        """Serialize the supplied Python object into a valid, JSON-formatted
-        byte-encoded string."""
+        """Serialize the supplied Python object into a valid, JSON-formatted byte-encoded string."""
         return json.dumps(value, default=str).encode() if value is not None else None