numpy-vector-store 0.1.0__py3-none-any.whl

numpy_vector_store/__init__.py

@@ -0,0 +1,13 @@
+"""
+NumPy Vector Store - A simple vector store implementation using NumPy.
+
+This package provides a lightweight vector store implementation for storing
+and searching vector embeddings using NumPy arrays.
+"""
+
+__version__ = "0.1.0"
+__author__ = "Tim VanReenen"
+
+from .vector_store import VectorStore
+
+__all__ = ["VectorStore"]

numpy_vector_store/vector_store.py

@@ -0,0 +1,253 @@
+from pathlib import Path
+from typing import Any, Optional
+
+import numpy as np
+
+
+class VectorStore:
+    """
+    A simple vector store implementation using NumPy.
+
+    This class provides basic functionality for storing and searching
+    vector embeddings using cosine similarity.
+    """
+
+    def __init__(
+        self,
+        dimensions: int,
+        file_path: Optional[str] = None,
+        metadata_schema: Optional[dict] = None,
+    ):
+        """
+        Initialize the vector store.
+
+        Args:
+            dimensions: The number of dimensions for vectors to be stored.
+            file_path: Optional path to save/load vectors from.
+            metadata_schema: Optional schema for structured metadata operations.
+        """
+        self.dimensions = dimensions
+        self.file_path = Path(file_path) if file_path else None
+        # Initialize as empty 2D array with correct dimensions
+        self.vectors: np.ndarray = np.array([]).reshape(0, dimensions)
+        self._loaded = False
+
+        if metadata_schema:
+            # Use structured NumPy array for performance
+            self.metadata = np.array([], dtype=self._create_dtype(metadata_schema))
+            self._use_structured = True
+            self._metadata_schema: Optional[dict] = metadata_schema
+        else:
+            # Use object array for flexibility (still NumPy)
+            self.metadata = np.array([], dtype=object)
+            self._use_structured = False
+            self._metadata_schema = None
+
+    def _create_dtype(self, schema: dict) -> list:
+        """Convert schema dict to NumPy dtype."""
+        dtype = []
+        for field, field_type in schema.items():
+            if isinstance(field_type, str):
+                dtype.append((field, field_type))
+            else:
+                dtype.append((field, "O"))  # Object type for complex data
+        return dtype
+
+    def add_vectors(self, vectors_2d: np.ndarray, metadata_array: np.ndarray) -> None:
+        """
+        Add vectors and metadata directly as NumPy arrays (most efficient).
+
+        Args:
+            vectors_2d: 2D NumPy array of shape (n_vectors, dimensions)
+            metadata_array: 1D NumPy array of metadata objects
+        """
+        if vectors_2d.ndim != 2:
+            raise ValueError("vectors_2d must be a 2D NumPy array")
+
+        if metadata_array.ndim != 1:
+            raise ValueError("metadata_array must be a 1D NumPy array")
+
+        if vectors_2d.shape[1] != self.dimensions:
+            raise ValueError(
+                f"Vector dimensions {vectors_2d.shape[1]} doesn't match store dimensions {self.dimensions}"
+            )
+
+        if len(vectors_2d) != len(metadata_array):
+            raise ValueError("Number of vectors must match number of metadata items")
+
+        # Normalize vectors in batch
+        norms = np.linalg.norm(vectors_2d, axis=1, keepdims=True)
+        if np.any(norms == 0):
+            raise ValueError("Cannot add zero-norm vectors")
+        normalized_vectors = vectors_2d / norms
+
+        # Add vectors directly
+        if len(self.vectors) == 0:
+            self.vectors = normalized_vectors
+        else:
+            self.vectors = np.vstack([self.vectors, normalized_vectors])
+
+        # Add metadata directly
+        self.metadata = np.append(self.metadata, metadata_array)
+
+    def load(self) -> None:
+        """
+        Load vectors from file if file_path is specified.
+        """
+        if self._loaded or not self.file_path:
+            return
+
+        if self.file_path.exists():
+            data = np.load(self.file_path, allow_pickle=True)
+            loaded_vectors = data["vectors"]
+            loaded_metadata = data["metadata"]
+            data.close()
+
+            if loaded_vectors.ndim != 2:
+                raise ValueError("Loaded vectors must be a 2D array")
+
+            if loaded_vectors.shape[1] != self.dimensions:
+                raise ValueError(
+                    f"Loaded vector dimension {loaded_vectors.shape[1]} doesn't match store dimensions {self.dimensions}"
+                )
+
+            if len(loaded_vectors) != len(loaded_metadata):
+                raise ValueError("Loaded vectors and metadata length mismatch")
+
+            if self._use_structured and self._metadata_schema is not None:
+                expected_fields = tuple(self._metadata_schema.keys())
+                loaded_fields = loaded_metadata.dtype.names
+                if loaded_fields != expected_fields:
+                    raise ValueError(
+                        f"Loaded metadata schema {loaded_fields} doesn't match expected schema {expected_fields}"
+                    )
+                loaded_metadata = loaded_metadata.astype(
+                    np.dtype(self._create_dtype(self._metadata_schema)), copy=False
+                )
+
+            self.vectors = loaded_vectors
+            self.metadata = loaded_metadata
+
+        self._loaded = True
+
+    def save(self) -> None:
+        """
+        Save vectors to file if file_path is specified.
+        """
+        if not self.file_path:
+            return
+
+        vectors_array: np.ndarray = self.vectors.astype(np.float32)
+        # Preserve metadata dtype so structured arrays remain structured
+        # after save/load round-trips.
+        metadata_array = np.array(self.metadata, copy=True)
+        np.savez_compressed(self.file_path, vectors=vectors_array, metadata=metadata_array)
+
+    def search(
+        self, query_vector: np.ndarray, top_k: int = 10, score_cutoff: float = 0.0
+    ) -> list[tuple[int, float, dict]]:
+        """
+        Search for the most similar vectors.
+
+        Args:
+            query_vector: The query vector to search with.
+            top_k: Number of top results to return.
+            score_cutoff: Minimum similarity score to include in results.
+
+        Returns:
+            List of tuples containing (index, similarity_score, metadata).
+        """
+        if len(query_vector) != self.dimensions:
+            raise ValueError(
+                f"Query vector dimension {len(query_vector)} doesn't match store dimensions {self.dimensions}"
+            )
+
+        if top_k <= 0:
+            raise ValueError("top_k must be greater than 0")
+
+        if len(self.vectors) == 0:
+            return []
+
+        # Use efficient cosine similarity computation
+        similarities = self._cosine_similarity_numpy(query_vector, self.vectors)
+
+        # Filter by similarity cutoff first
+        valid_indices = np.where(similarities >= score_cutoff)[0]
+        if len(valid_indices) == 0:
+            return []
+
+        # Get top_k results from valid indices
+        valid_similarities = similarities[valid_indices]
+        top_valid_indices = np.argsort(valid_similarities)[::-1][:top_k]
+        top_indices = valid_indices[top_valid_indices]
+
+        results = []
+        for idx in top_indices:
+            results.append((int(idx), float(similarities[idx]), self.metadata[idx]))
+
+        return results
+
+    def _cosine_similarity_numpy(
+        self, query: np.ndarray, vectors: np.ndarray
+    ) -> np.ndarray:
+        """
+        Efficient cosine similarity computation using NumPy.
+
+        Args:
+            query: Query vector.
+            vectors: Array of vectors to compare against.
+
+        Returns:
+            Array of similarity scores.
+        """
+        query_magnitude = np.linalg.norm(query)
+        if query_magnitude == 0:
+            raise ValueError("Cannot search with zero-norm query vector")
+        query_norm = query / query_magnitude
+        vector_norms = np.linalg.norm(vectors, axis=1, keepdims=True)
+        if np.any(vector_norms == 0):
+            raise ValueError("Store contains zero-norm vectors")
+        vectors_norm = vectors / vector_norms
+        return np.dot(vectors_norm, query_norm)  # type: ignore[no-any-return]
+
+    def get(self, index: int) -> Optional[tuple[np.ndarray, dict]]:
+        """
+        Get vector and metadata by index.
+
+        Args:
+            index: The index of the vector to retrieve.
+
+        Returns:
+            Tuple of (vector, metadata) or None if index is out of bounds.
+        """
+        if 0 <= index < len(self.vectors):
+            if self._use_structured and self._metadata_schema is not None:
+                # Convert structured array row to dict
+                metadata_row = self.metadata[index]
+                metadata_dict = {
+                    field: metadata_row[field] for field in self._metadata_schema.keys()
+                }
+                return (self.vectors[index], metadata_dict)
+            else:
+                # Object array - return the dict directly
+                return (self.vectors[index], self.metadata[index])
+        return None
+
+    def clear(self) -> None:
+        """Clear all vectors and metadata from the store."""
+        self.vectors = np.array([]).reshape(0, self.dimensions)
+        if self._use_structured and self._metadata_schema is not None:
+            self.metadata = np.array(
+                [], dtype=self._create_dtype(self._metadata_schema)
+            )
+        else:
+            self.metadata = np.array([], dtype=object)
+
+    def __enter__(self) -> "VectorStore":
+        """Enter the context manager."""
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Exit the context manager, auto-save if file_path is specified."""
+        if self.file_path:
+            self.save()

numpy_vector_store-0.1.0.dist-info/METADATA

@@ -0,0 +1,205 @@
+Metadata-Version: 2.4
+Name: numpy-vector-store
+Version: 0.1.0
+Summary: A fast, lightweight, and zero-setup in-memory vector store powered by NumPy
+Project-URL: Homepage, https://github.com/tvanreenen/numpy-vector-store
+Project-URL: Repository, https://github.com/tvanreenen/numpy-vector-store
+Project-URL: Issues, https://github.com/tvanreenen/numpy-vector-store/issues
+Author: Tim VanReenen
+License: MIT
+License-File: LICENSE
+Keywords: embeddings,numpy,search,store,vector
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.9
+Requires-Dist: numpy>=1.20.0
+Description-Content-Type: text/markdown
+
+# NumPy Vector Store
+
+A fast, lightweight, and zero-setup in-memory vector store powered by NumPy.
+
+- **High-performance** index-free O(n) cosine similarity searches
+- **Vectorized** metadata queries with NumPy operations
+- **Easy persistence** to and from compressed NumPy binary files (.npz)
+- **Zero dependency** on external services or data stores
+
+## Why?
+
+This library is purpose-built for small to medium-scale vector search tasks and offers a simple, lightweight alternative to heavyweight solutions like Pinecone, Qdrant, Weaviate, Postgres + pgvector, or Azure AI Search—no complex setup or infrastructure required.
+
+> Sometimes you don't need a sledgehammer to crack a nut.
+
+## When/Where?
+
+Below are benchmark results for the module's search method to help you assess its suitability for your use case.
+
+| Embedding Type | Dimensions | ~5ms | ~25ms | ~100ms | ~500ms |
+|----------------|------------|------|-------|--------|--------|
+| **Sentence Transformers** | 384 | 1K vectors<br/>1.5MB | 10K vectors<br/>15MB | 100K vectors<br/>147MB | 500K vectors<br/>732MB |
+| **OpenAI Small** | 1536 | 500 vectors<br/>3MB | 5K vectors<br/>29MB | 25K vectors<br/>147MB | 100K vectors<br/>586MB |
+| **OpenAI Large** | 3072 | 200 vectors<br/>2MB | 2.5K vectors<br/>29MB | 5K vectors<br/>59MB | 25K vectors<br/>293MB |
+
+*Benchmarks performed on Apple M2 hardware*
+
+## Installation
+
+**⚠️ Pending submission to PyPI**
+
+```bash
+uv add numpy-vector-store
+```
+
+## Quick Start
+
+```python
+import numpy as np
+from numpy_vector_store import VectorStore
+
+# Load your vector store
+store = VectorStore(dimensions=1536, file_path='vectors.npz')
+store.load()
+
+# Embed your search query
+query = np.array([0.2, 0.3, 0.4, ...])
+
+# Search using cosine similarity
+results = store.search(query, top_k=3)
+
+# Compare the results
+for index, similarity, meta in results:
+    print(f"{meta['title']}: {similarity:.3f}")
+```
+
+## Usage Examples
+
+### Adding Vectors
+
+#### Adding Vectors in Batch
+
+The `add_vectors` method takes a **2D NumPy array** where each row is a vector, and a **1D NumPy array** of metadata objects.
+
+```python
+# 2D np.array of vectors
+embeddings = np.array([
+    [0.1, 0.2, 0.3, ...],  # Text embedding 1
+    [0.4, 0.5, 0.6, ...],  # Text embedding 2
+    [0.7, 0.8, 0.9, ...]   # Text embedding 3
+])
+
+# 1D np.array of metadata objects
+metadata = np.array([
+    {"title": "AI Overview", "word_count": 12},
+    {"title": "Python Guide", "word_count": 10},
+    {"title": "Vector DBs", "word_count": 8}
+])
+
+# Vectors and metadata added using efficient vectorized NumPy operations
+store.add_vectors(embeddings, metadata)
+```
+
+#### Adding Vectors Individually
+
+Batch operations are generally preferable and more efficient. But individual vectors can be added with the same method.
+
+```python
+store.add_vectors(
+    np.array([new_embedding]), # or np.atleast_2d(new_embedding)
+    np.array([{"title": "Neural Networks Paper", "word_count": 15}])
+)
+```
+
+### Save the Vector Store
+
+Save to a gzip compressed NumPy binary file directly
+
+```python
+store.save()
+```
+
+Or use context manager for automatic persistence
+
+```python
+with VectorStore(dimensions=3, file_path="vectors.npz") as store:
+    store.add_vectors(vectors_2d, metadata_array)
+
+# Automatically saves when exiting the context
+```
+
+### Working with Metadata
+
+Work with flexible, unstructured metadata using standard Python operations. No schema required - perfect for getting started quickly:
+
+```python
+# Access individual metadata
+first_metadata = store.metadata[0]
+print(f"First entry: {first_metadata}")
+
+# Iterate through all metadata
+for i, metadata in enumerate(store.metadata):
+    print(f"Entry {i}: {metadata}")
+```
+
+### Advanced: Structured Metadata for Performance
+
+If you define a homogeneous NumPy schema upfront, you get **significant performance improvements** for metadata operations. Instead of Python loops and dictionary lookups, you get **vectorized NumPy operations** that are orders of magnitude faster.
+
+> **Learn more**: [NumPy Structured Arrays Documentation](https://numpy.org/doc/stable/user/basics.rec.html)
+
+```python
+# Define schema for performance
+store = VectorStore(
+    dimensions=512,
+    metadata_schema={'title': 'U200', 'year': 'i4', 'citations': 'i4'}
+)
+store.add_vectors(
+    np.array([vector1, vector2]),
+    np.array([
+        {"title": "Paper 1", "year": 2023, "citations": 100},
+        {"title": "Paper 2", "year": 2022, "citations": 50}
+    ])
+)
+
+# Perform vectorized NumPy operations on metadata
+recent_mask = store.metadata['year'] == 2023
+recent_vectors = np.array(store.vectors)[recent_mask]
+
+# Sort by citations
+sorted_indices = np.argsort(store.metadata['citations'])[::-1]
+top_vectors = np.array(store.vectors)[sorted_indices]
+
+# Complex filtering
+high_impact = store.metadata['citations'] > 100
+recent = store.metadata['year'] > 2020
+combined_mask = high_impact & recent
+filtered_vectors = np.array(store.vectors)[combined_mask]
+```
+
+## Contributing
+
+### Setup Development Environment
+
+```bash
+git clone https://github.com/tvanreenen/numpy-vector-store.git
+cd numpy-vector-store
+uv sync --frozen --group dev
+```
+
+### Before Submitting a Pull Request
+
+Please ensure:
+
+1. **Code Quality**: Run `uv run ruff check` - should show no issues
+2. **Formatting**: Run `uv run ruff format` - should show "files left unchanged"
+3. **Type Checking**: Run `uv run mypy src/` - should show no errors
+4. **Tests**: Run `uv run pytest` - all tests should pass
+
+## License
+
+MIT License - see [LICENSE](LICENSE) file for details.

numpy_vector_store-0.1.0.dist-info/RECORD

@@ -0,0 +1,6 @@
+numpy_vector_store/__init__.py,sha256=0_LXVUazbViKSm_aaZvRX7kEIti9YoIATGV31ycifKo,326
+numpy_vector_store/vector_store.py,sha256=-FFqRg0-tG6xJdz8eg0aaUdIqq5WB9LUf3FTLfdTQRI,9399
+numpy_vector_store-0.1.0.dist-info/METADATA,sha256=kK2_AnERaUVSWtEY3SCT9nWDpVt1NPNZLqg8l1Nc424,6478
+numpy_vector_store-0.1.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+numpy_vector_store-0.1.0.dist-info/licenses/LICENSE,sha256=SUVvvB3m6p_RRGYL6l9PPdlUfg39dxYNqeF0b9-z_bg,1070
+numpy_vector_store-0.1.0.dist-info/RECORD,,

numpy_vector_store-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any

numpy_vector_store-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Tim VanReenen
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.