feat: Integrate Flock Extension

docs/usage-dense-flock.md

@@ -0,0 +1,111 @@
+# QuackIR: Dense Retrieval with Flock + Ollama
+
+This short guide shows how to run a minimal dense-retrieval pipeline with the Flock extension and a local Ollama embedding model.
+It assumes you’ve followed the NFCorpus setup in this [experiments guide](./experiments-nfcorpus.md).
+
+## Install and start Ollama
++ Download and install Ollama from the [download page](https://ollama.com/download).
++ Ensure the service is running locally (default: `127.0.0.1:11434`). Start it with:
+
+```bash
+OLLAMA_NUM_PARALLEL=2 \
+OLLAMA_MAX_QUEUE=2048 \
+ollama serve &
+```
+
+The `ollama serve &` command starts the Ollama server in the background.
+The environment variables control concurrency and memory usage.
+If you encounter `server busy, please try again. maximum pending requests exceeded` error while running the script, consider increasing the `OLLAMA_MAX_QUEUE` value, or reducing the `batch_size` parameter in the `options_json` field when registering the model.
+Tune these settings based on your hardware capacity and workload.
+
++ Pull an embedding model:
+
+```bash
+ollama pull embeddinggemma
+```
+
+## Full Python walkthrough
+Paste the following Python script into a file and execute it.
+
+```python
+from quackir.flock import FlockManager
+from quackir.index import DuckDBIndexer
+from quackir.search import DuckDBSearcher
+from quackir import IndexType
+from pathlib import Path
+import csv
+
+"""
+1) Configuration
+    - Set table names, paths, and embedding dimension for the model you'll use.
+"""
+table_name = "corpus_dense"
+corpus_file = "collections/nfcorpus/quackir_corpus.jsonl"
+embedding_dim = 768
+model_alias = "Embedder"
+queries_file = "collections/nfcorpus/queries.tsv"
+output_path = Path("runs/run.quackir.duckdb.dense.flock.nfcorpus.txt")
+top_k = 10
+
+"""
+2) Initialize Flock + register model alias
+    - Loads the extension and creates the model alias bound to the local Ollama model.
+"""
+flock = FlockManager()
+flock.create_model(
+    alias=model_alias,
+    provider_model="embeddinggemma",
+    provider="ollama",
+    options_json='{"tuple_format":"json", "batch_size":16}',
+)
+
+"""
+3) Create dense table schema
+    - A DuckDB table with (id VARCHAR, embedding DOUBLE[embedding_dim]).
+"""
+indexer = DuckDBIndexer(flock_manager=flock)
+indexer.init_table(table_name, IndexType.DENSE, embedding_dim=embedding_dim)
+
+"""
+4) Populate embeddings directly from the corpus file
+    - with_flock=True invokes Flock to compute an embedding per row and insert into the table.
+"""
+indexer.load_table(
+    table_name,
+    corpus_file,
+    with_flock=True,
+    id_column="id",
+    contents_column="contents",
+    embedding_dim=embedding_dim,
+)
+indexer.close()
+
+"""
+5) Embed queries on-the-fly and write a TREC run file
+    - embedding_search(..., with_flock=True) computes the query embedding via Flock, then scores by cosine similarity.
+"""
+searcher = DuckDBSearcher(flock_manager=flock)
+
+with output_path.open("w") as out, open(queries_file) as f:
+    reader = csv.reader(f, delimiter="\t")
+    for qid, qtext in reader:
+        hits = searcher.embedding_search(
+            query_embedding=qtext,  # raw query text; Flock generates the vector
+            top_n=top_k,
+            table_name=table_name,
+            with_flock=True,
+            embedding_dim=embedding_dim,
+        )
+        for rank, (docid, score) in enumerate(hits, start=1):
+            out.write(f"{qid} Q0 {docid} {rank} {score:.6f} QuackIR\n")
+
+searcher.close()
+```
+
+## Evaluate with trec_eval
+
+```bash
+python -m pyserini.eval.trec_eval \
+    -c -m ndcg_cut.10 collections/nfcorpus/qrels/test.qrels \
+    runs/run.quackir.duckdb.dense.flock.nfcorpus.txt
+```

quackir/_base.py

@@ -33,6 +33,11 @@ class SearchDB(Enum):
     SQLITE = 'sqlite'
     POSTGRES = 'postgres'
 
+class SecretProvider(Enum):
+    OLLAMA = 'OLLAMA'
+    OPENAI = 'OPENAI'
+    AZURE = 'AZURE_LLM'
+
 def count_lines(filename, open_cmd):
     with open_cmd(filename, 'r') as file:
         return sum(1 for _ in file)

quackir/flock.py

@@ -0,0 +1,204 @@
+#
+# QuackIR: Reproducible IR research in RDBMS
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import duckdb
+from ._base import SecretProvider
+
+
+class FlockManager:
+    def __init__(
+        self,
+        conn: duckdb.DuckDBPyConnection | None = None,
+        db_path: str = "duck.db",
+        secret_type: SecretProvider = SecretProvider.OLLAMA,
+        api_url: str = "127.0.0.1:11434",
+        api_key: str | None = None,
+        resource_name: str | None = None,
+        api_version: str | None = None,
+    ):
+        """Create a FlockManager.
+
+        Parameters
+        ----------
+        conn : duckdb.DuckDBPyConnection | None
+            Existing connection (optional). If not provided a new one is opened at db_path.
+        db_path : str
+            Path to DuckDB file when creating a new connection.
+        secret_type : SecretProvider
+            One of SecretProvider.OLLAMA, SecretProvider.OPENAI, SecretProvider.AZURE.
+        api_url : str
+            For OLLAMA only.
+        api_key : str | None
+            For OPENAI and AZURE.
+        resource_name : str | None
+            For AZURE only.
+        api_version : str | None
+            For AZURE only.
+        """
+        self.conn = conn or duckdb.connect(db_path)
+        self.db_path = db_path
+        self.secret_type = secret_type
+        self.secret_name = f"__default_{secret_type.name.lower()}"
+        self.api_url = api_url
+        self.api_key = api_key
+        self.resource_name = resource_name
+        self.api_version = api_version
+        self.model_alias = None
+
+        self.ensure_loaded()
+        self.ensure_secret(
+            secret_type=self.secret_type,
+            api_url=self.api_url,
+            api_key=self.api_key,
+            resource_name=self.resource_name,
+            api_version=self.api_version,
+        )
+
+    def ensure_loaded(self) -> None:
+        self.conn.execute("INSTALL flock FROM community;")
+        self.conn.execute("LOAD flock;")
+
+    def ensure_secret(
+        self,
+        *,
+        secret_type: SecretProvider,
+        api_url: str,
+        api_key: str | None,
+        resource_name: str | None,
+        api_version: str | None,
+    ) -> None:
+        provider = secret_type.value
+        params = {"TYPE": provider}
+
+        if secret_type == SecretProvider.OLLAMA:
+            if not api_url:
+                raise ValueError("api_url required for OLLAMA secret")
+            params["API_URL"] = api_url
+
+        elif secret_type == SecretProvider.OPENAI:
+            if not api_key:
+                raise ValueError("api_key required for OPENAI secret")
+            params["API_KEY"] = api_key
+
+        elif secret_type == SecretProvider.AZURE:
+            if not api_key or not resource_name or not api_version:
+                raise ValueError("Missing required AZURE secret parameters")
+            params.update(
+                {
+                    "API_KEY": api_key,
+                    "RESOURCE_NAME": resource_name,
+                    "API_VERSION": api_version,
+                }
+            )
+        else:
+            raise ValueError(f"Unsupported secret_type: {secret_type}")
+
+        self.conn.execute(f'DROP SECRET IF EXISTS "{self.secret_name}";')
+        fields = ", ".join(f"{k} ?" for k in params)
+        values = list(params.values())
+        self.conn.execute(f'CREATE SECRET "{self.secret_name}" ({fields});', values)
+
+    def create_model(
+        self,
+        alias: str,
+        provider_model: str,
+        provider: str = "ollama",
+        options_json: str = "{}",
+        skip_if_exists: bool = True,
+    ) -> None:
+        self.model_alias = alias
+
+        try:
+            self.conn.execute(
+                f"""
+                CREATE MODEL(
+                  '{alias}',
+                  '{provider_model}',
+                  '{provider}',
+                  {options_json}
+                )
+                """
+            )
+        except Exception as e:
+            msg = str(e)
+            if skip_if_exists and ("Duplicate key" in msg or "already exists" in msg):
+                return
+            raise
+
+    def create_embedding(
+        self,
+        dest_table: str,
+        file_path: str,
+        id_column: str = "id",
+        contents_column: str = "contents",
+        embedding_dim: int = 768,
+    ):
+        model_cfg = (
+            f"{{'model_name': '{self.model_alias}', 'secret': '{self.secret_name}'}}"
+        )
+
+        input_cfg = f"{{'context_columns': [ {{'data': {contents_column}}} ]}}"
+
+        if file_path.endswith(".jsonl"):
+            base_query = f"SELECT {id_column} AS id, {contents_column} FROM read_json_auto('{file_path}')"
+        elif file_path.endswith(".parquet"):
+            base_query = f"SELECT {id_column} AS id, {contents_column} FROM read_parquet('{file_path}')"
+        else:
+            raise ValueError("Unsupported file type (use .jsonl or .parquet)")
+
+        sql_insert = f"""
+            INSERT INTO {dest_table}
+            SELECT
+                id,
+                CAST(
+                    llm_embedding(
+                        {model_cfg},
+                        {input_cfg}
+                    ) AS DOUBLE[{embedding_dim}]
+                ) AS embedding
+            FROM ({base_query})
+            """
+        self.conn.execute(sql_insert)
+
+    def search_embedding(
+        self,
+        query_text: str,
+        table_name,
+        top_n=5,
+        embedding_dim: int = 768,
+    ):
+        model_cfg = (
+            f"{{'model_name': '{self.model_alias}', 'secret': '{self.secret_name}'}}"
+        )
+
+        sql_stmt = f"""
+            WITH q AS (
+                SELECT CAST(
+                    llm_embedding(
+                        {model_cfg},
+                        {{'context_columns':[{{'data': ?}}]}}
+                    ) AS DOUBLE[{embedding_dim}]
+                ) AS embedding
+            )
+            SELECT
+                t.id,
+                array_cosine_similarity(t.embedding, q.embedding) AS score
+            FROM {table_name} AS t, q
+            ORDER BY score DESC
+            LIMIT {top_n}
+            """
+
+        return self.conn.execute(sql_stmt, [query_text]).fetchall()

quackir/index/_base.py

@@ -27,7 +27,7 @@ def get_num_rows(self, table_name: str) -> int:
         """Get the number of rows in the specified table."""
         pass
 
-    def load_table(self, table_name: str, file_path: str, index_type: IndexType = None, pretokenized=False):
+    def load_table(self, table_name: str, file_path: str, index_type: IndexType | None = None, pretokenized=False):
         """Load data into the specified table."""
         if index_type == None:
             index_type = self.get_index_type(table_name)

quackir/index/_duck.py

@@ -17,12 +17,14 @@
 from ._base import Indexer
 from quackir._base import IndexType
 from quackir.analysis import tokenize
+from quackir.flock import FlockManager
 import duckdb
 import json
 
 class DuckDBIndexer(Indexer):
-    def __init__(self, db_path="duck.db"):
+    def __init__(self, db_path="duck.db", flock_manager: FlockManager | None = None):
         self.conn = duckdb.connect(db_path)
+        self.flock_manager = flock_manager
 
     def get_index_type(self, table_name: str) -> IndexType:
         table_description = self.conn.execute(f"DESCRIBE {table_name}").fetchall()
@@ -34,6 +36,33 @@ def get_index_type(self, table_name: str) -> IndexType:
         else:
             raise ValueError(f"Unknown index type for table {table_name}. Ensure it has either an 'embedding' column or a 'contents' column.")
 
+    def load_table(
+        self,
+        table_name: str,
+        file_path: str,
+        index_type: IndexType | None = None,
+        pretokenized=False,
+        with_flock=False,
+        **flock_kwargs,
+    ):
+        """
+        Load a file into a table.
+
+        If with_flock=True, build a dense embedding table directly from the source file
+        using the Flock extension instead of loading raw contents.
+
+        Extra keyword arguments are forwarded to create_embedding, e.g.:
+            id_column="id", contents_column="contents", embedding_dim=768
+        """
+        if with_flock:
+            if self.flock_manager is None:
+                raise ValueError("flock_manager must be provided when with_flock=True.")
+            self.flock_manager.create_embedding(
+                dest_table=table_name, file_path=file_path, **flock_kwargs
+            )
+        else:
+            super().load_table(table_name, file_path, index_type, pretokenized)
+
     def init_table(self, table_name: str, index_type: IndexType, embedding_dim=768):
         self.conn.execute(f"""DROP TABLE IF EXISTS {table_name}""")
         if index_type == IndexType.SPARSE: