Use quantized ONNX defaults and PyArrow bulk import for ~9x speedup

Author: bigbag

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "microrag"
-version = "0.2.1"
+version = "0.2.2"
 description = "A feature-rich, universal RAG library for Python with ONNX-backed embeddings and DuckDB storage"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -21,6 +21,7 @@ dependencies = [
     "duckdb>=0.9.0",
     "rank-bm25>=0.2.2",
     "numpy>=1.20.0",
+    "pyarrow>=23.0.0",
 ]
 
 [dependency-groups]

src/microrag/__init__.py

@@ -1,7 +1,7 @@
 """MicroRAG - A feature-rich, universal RAG library for Python.
 
 MicroRAG provides:
-- ONNX-backed embeddings (CPU-only, no PyTorch at runtime)
+- ONNX-backed embeddings with quantized models for fast CPU inference
 - DuckDB storage with HNSW vector indexes
 - Three-tier hybrid search (semantic + BM25 + FTS) with RRF fusion
 - Query preprocessing with abbreviation expansion
@@ -10,10 +10,8 @@
     ```python
     from microrag import MicroRAG, RAGConfig
 
-    config = RAGConfig(
-        model_path="/path/to/all-MiniLM-L6-v2",
-        db_path="./rag.duckdb",
-    )
+    # Uses sentence-transformers/all-MiniLM-L6-v2 with quantized ONNX by default
+    config = RAGConfig(db_path="./rag.duckdb")
 
     with MicroRAG(config) as rag:
         rag.add_documents(["Document 1", "Document 2"])
@@ -35,7 +33,7 @@
 )
 from microrag.models import Document, SearchResult
 
-__version__ = "0.2.1"
+__version__ = "0.2.2"
 
 __all__ = [
     # Main classes

src/microrag/config.py

@@ -5,15 +5,23 @@
 
 from microrag.stopwords import ENGLISH_STOPWORDS
 
+# Default model for sentence-transformers backend (fast, good quality, 384 dims)
+_DEFAULT_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+# Quantized ONNX model for ~2x faster inference on CPU
+_DEFAULT_MODEL_FILE = "onnx/model_qint8_avx512.onnx"
+
 
 @dataclass(frozen=True)
 class RAGConfig:
     """Configuration for MicroRAG instance.
 
     Attributes:
         model_path: Model path (sentence-transformers) or model name (fastembed).
+            Defaults to "sentence-transformers/all-MiniLM-L6-v2".
         embedding_backend: Embedding backend ("auto", "sentence-transformers", "fastembed").
+            Defaults to "auto" which prefers sentence-transformers for best performance.
         model_file: ONNX model filename within model_path (sentence-transformers only).
+            Defaults to quantized model for ~2x faster CPU inference.
         fastembed_cache_dir: Cache directory for fastembed models.
         db_path: DuckDB database path. Use ":memory:" for in-memory database.
         embedding_dim: Dimension of embedding vectors.
@@ -32,9 +40,9 @@ class RAGConfig:
         batch_size: Batch size for embedding generation.
     """
 
-    model_path: str = ""
+    model_path: str = _DEFAULT_MODEL
     embedding_backend: str = "auto"
-    model_file: str | None = None
+    model_file: str | None = _DEFAULT_MODEL_FILE
     fastembed_cache_dir: str | None = None
     db_path: str = ":memory:"
     embedding_dim: int = 384
@@ -50,7 +58,7 @@ class RAGConfig:
     hnsw_ef_search: int = 100
     hnsw_m: int = 16
     hnsw_enable_persistence: bool = False
-    batch_size: int = 32
+    batch_size: int = 64
 
     def __post_init__(self) -> None:
         valid_backends = ("auto", "sentence-transformers", "fastembed")

src/microrag/storage/duckdb.py

@@ -8,6 +8,7 @@
 
 import duckdb
 import numpy as np
+import pyarrow as pa  # type: ignore[import-untyped]
 from numpy.typing import NDArray
 
 from microrag.exceptions import StorageError
@@ -88,37 +89,51 @@ def _init_schema(self) -> None:
             conn.execute("SET hnsw_enable_experimental_persistence = true")
 
         # Create documents table
-        conn.execute(f"""
+        conn.execute(
+            f"""
             CREATE TABLE IF NOT EXISTS documents (
                 id VARCHAR PRIMARY KEY,
                 content TEXT NOT NULL,
                 metadata JSON,
                 embedding FLOAT[{self._embedding_dim}]
             )
-        """)
+        """
+        )
 
     def add_documents(self, documents: Sequence[Document]) -> None:
-        """Add documents to storage."""
+        """Add documents to storage using PyArrow bulk import for performance."""
         if not documents:
             return
 
-        logger.debug("Storing %d document(s) in DuckDB", len(documents))
+        logger.debug("Storing %d document(s) in DuckDB via PyArrow", len(documents))
         conn = self.conn
         try:
+            # Validate embeddings
             for doc in documents:
                 if doc.embedding is None:
                     raise StorageError(f"Document {doc.id} has no embedding")
 
-                embedding_list = doc.embedding.tolist()
-                metadata_json = json.dumps(doc.metadata)
+            # Create PyArrow table - DuckDB can query it directly by name
+            arrow_table = pa.table(  # noqa: F841
+                {
+                    "id": [doc.id for doc in documents],
+                    "content": [doc.content for doc in documents],
+                    "metadata": [json.dumps(doc.metadata) for doc in documents],
+                    "embedding": pa.array(
+                        [doc.embedding.tolist() for doc in documents],  # type: ignore[union-attr]
+                        type=pa.list_(pa.float32()),
+                    ),
+                }
+            )
 
-                conn.execute(
-                    """
-                    INSERT OR REPLACE INTO documents (id, content, metadata, embedding)
-                    VALUES (?, ?, ?, ?)
-                    """,
-                    [doc.id, doc.content, metadata_json, embedding_list],
-                )
+            # Bulk import directly from PyArrow table (no temp file needed)
+            conn.execute(
+                f"""
+                INSERT OR REPLACE INTO documents
+                SELECT id, content, metadata::JSON, embedding::FLOAT[{self._embedding_dim}]
+                FROM arrow_table
+                """
+            )
 
             # Invalidate indexes after adding documents
             self._vector_index_built = False
@@ -223,7 +238,8 @@ def build_vector_index(
             conn.execute("DROP INDEX IF EXISTS documents_embedding_idx")
 
             # Create HNSW index
-            conn.execute(f"""
+            conn.execute(
+                f"""
                 CREATE INDEX documents_embedding_idx ON documents
                 USING HNSW (embedding)
                 WITH (
@@ -232,7 +248,8 @@ def build_vector_index(
                     ef_search = {ef_search},
                     m = {m}
                 )
-            """)
+            """
+            )
 
             self._vector_index_built = True
             logger.debug("HNSW vector index built")
@@ -246,7 +263,8 @@ def build_fts_index(self) -> None:
             conn = self.conn
 
             # Create FTS index using PRAGMA
-            conn.execute("""
+            conn.execute(
+                """
                 PRAGMA create_fts_index(
                     'documents',
                     'id',
@@ -257,7 +275,8 @@ def build_fts_index(self) -> None:
                     strip_accents = 1,
                     lower = 1
                 )
-            """)
+            """
+            )
 
             self._fts_index_built = True
             logger.debug("FTS index built")

tests/test_config.py

@@ -27,10 +27,12 @@ def test_create_config(self):
         assert config.abbreviations == {"ML": "machine learning"}
 
     def test_default_embedding_backend(self):
-        """Test that embedding_backend defaults to 'auto'."""
+        """Test that embedding_backend defaults to 'auto' with quantized model."""
         config = RAGConfig()
         assert config.embedding_backend == "auto"
-        assert config.model_path == ""
+        assert config.model_path == "sentence-transformers/all-MiniLM-L6-v2"
+        assert config.model_file == "onnx/model_qint8_avx512.onnx"
+        assert config.batch_size == 64
 
     def test_invalid_embedding_backend_raises_error(self):
         """Test that invalid embedding_backend raises ValueError."""