diff --git a/.gitignore b/.gitignore index b0ce04a..30309fd 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,7 @@ __pycache__/ *.ipynb *.duckdb duckdb_tmp/ +*.block evaluation_datasets/ @@ -168,3 +169,23 @@ cython_debug/ # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ +arguana +benchmark_bm25s.py +benchmark_crud.py +climate-fever +fever.tmp/duckdb_temp_block-4611686018432402649.block +fever.tmp/duckdb_temp_block-4611686018432404521.block +fever.tmp/duckdb_temp_block-4611686018432404963.block +fever.tmp/duckdb_temp_storage-4.tmp +metrics.json +metrics_20K.json +metrics_bm25s.json +mmetrics_30K.json +msmarco +nfcorpus +nq +quora +scidocs +scifact +trec-covid +webis-touche2020 diff --git a/Makefile b/Makefile index 30a3dd9..af4ffbc 100644 --- a/Makefile +++ b/Makefile @@ -10,28 +10,22 @@ tests: @echo "Removing test.duckdb if it exists..." rm -rf test.duckdb rm -rf test.duckdb.wal - pytest ducksearch/tables/create.py - pytest ducksearch/tables/insert.py - pytest ducksearch/tables/select.py + pytest ducksearch/tables/create.py --disable-warnings + pytest ducksearch/tables/insert.py --disable-warnings + pytest ducksearch/tables/select.py --disable-warnings rm -rf test.duckdb rm -rf test.duckdb.wal - pytest ducksearch/hf/insert.py + pytest ducksearch/hf/insert.py --disable-warnings rm -rf test.duckdb rm -rf test.duckdb.wal - pytest ducksearch/delete/documents.py + pytest ducksearch/evaluation/evaluation.py --disable-warnings rm -rf test.duckdb rm -rf test.duckdb.wal - pytest ducksearch/evaluation/evaluation.py + pytest ducksearch/search/create.py --disable-warnings + pytest ducksearch/search/select.py --disable-warnings rm -rf test.duckdb rm -rf test.duckdb.wal - pytest ducksearch/upload/upload.py - rm -rf test.duckdb - rm -rf test.duckdb.wal - pytest ducksearch/search/create.py - pytest ducksearch/search/select.py - rm -rf test.duckdb - rm -rf test.duckdb.wal - pytest ducksearch/search/graphs.py + pytest ducksearch/search/graphs.py --disable-warnings rm -rf test.duckdb rm -rf test.duckdb.wal diff --git a/README.md b/README.md index 873d7a8..4f7bf37 100644 --- a/README.md +++ b/README.md @@ -13,30 +13,24 @@

-DuckSearch is a lightweight and easy-to-use library that allows to index and search documents. DuckSearch is built on top of DuckDB, a high-performance analytical database. DuckDB is designed to execute analytical SQL queries fast, and DuckSearch leverages this to provide efficient and scallable search / filtering capabilities. +DuckSearch is a lightweight and easy-to-use library to search documents. DuckSearch is built on top of DuckDB, a high-performance analytical database. DuckDB is designed to execute analytical SQL queries fast, and DuckSearch leverages this to provide efficient search and filtering features. DuckSearch index can be updated with new documents and documents can be deleted as well. DuckSearch also supports HuggingFace datasets, allowing to index datasets directly from the HuggingFace Hub.

## Installation -We can install DuckSearch using pip: +Install DuckSearch using pip: ```bash pip install ducksearch ``` -For evaluation dependencies, we can install DuckSearch with the `eval` extra: - -```bash -pip install "ducksearch[eval]" -``` - ## Documentation The complete documentation is available [here](https://lightonai.github.io/ducksearch/), which includes in-depth guides, examples, and API references. ### Upload -We can upload documents to DuckDB using the `upload.documents` function. The documents are stored in a DuckDB database, and the fields are indexed with BM25. +We can upload documents to DuckDB using the `upload.documents` function. The documents are stored in a DuckDB database, and the `fields` are indexed with BM25. ```python from ducksearch import upload @@ -79,7 +73,7 @@ upload.documents( ## Search -We can search documents using the `search.documents` function. The function returns the documents that match the query, sorted by the BM25 score. The `top_k` parameter controls the number of documents to return. We can also filter the results using SQL syntax which will be evaluated by DuckDB, therefore all DuckDB functions are available. +`search.documents` returns a list of list of documents ordered by relevance. We can control the number of documents to return using the `top_k` parameter. The following example demonstrates how to search for documents with the queries "punk" and "california" while filtering the results to include only documents with a date after 1970 and a popularity score greater than 8. ```python from ducksearch import search @@ -117,7 +111,22 @@ search.documents( ] ``` -List of DuckDB functions such as date functions can be found [here](https://duckdb.org/docs/sql/functions/date). +Filters are SQL expressions that are applied to the search results. We can use every filtering function DuckDB provides such as [date functions](https://duckdb.org/docs/sql/functions/date). + +## Delete and update index + +We can delete documents and update the BM25 weights accordingly using the `delete.documents` function. + +```python +from ducksearch import delete + +delete.documents( + database="ducksearch.duckdb", + ids=[0, 1], +) +``` + +To update the index, we should first delete the documents and then upload the updated documents. ## Extra features @@ -152,7 +161,6 @@ search.documents( database="fineweb.duckdb", queries="earth science", top_k=2, - filters="token_count > 200", ) ``` @@ -180,82 +188,31 @@ search.documents( ] ``` -### Graphs +## Benchmark -The `search.graphs` function can be used to search documents with a graph query. This function is useful if we have paired documents and queries. The search will retrieve the set of documents and queries that match the input query. Then it will build a graph and compute the weight of each document using a graph-based scoring function. -```python -from ducksearch import search, upload +| Dataset | ndcg@10 | hits@1 | hits@10 | mrr@10 | map@10 | r-precision | qps | Indexation Time (s) | Number of Documents and Queries | +|-------------------|-----------|---------|----------|----------|---------|-------------|----------------|---------------------|--------------------------------| +| arguana | 0.3779 | 0.0 | 0.8267 | 0.2491 | 0.2528 | 0.0108 | 117.80 | 1.42 | 1,406 queries, 8.67K documents | +| climate-fever | 0.1184 | 0.1068 | 0.3648 | 0.1644 | 0.0803 | 0.0758 | 5.88 | 302.39 | 1,535 queries, 5.42M documents | +| dbpedia-entity | 0.6046 | 0.7669 | 5.6241 | 0.8311 | 0.0649 | 0.0741 | 113.20 | 181.42 | 400 queries, 4.63M documents | +| fever | 0.3861 | 0.2583 | 0.5826 | 0.3525 | 0.3329 | 0.2497 | 74.40 | 329.70 | 6,666 queries, 5.42M documents | +| fiqa | 0.2445 | 0.2207 | 0.6790 | 0.3002 | 0.1848 | 0.1594 | 545.77 | 6.04 | 648 queries, 57K documents | +| hotpotqa | 0.4487 | 0.5059 | 0.9699 | 0.5846 | 0.3642 | 0.3388 | 48.15 | 163.14 | 7,405 queries, 5.23M documents | +| msmarco | 0.8951 | 1.0 | 8.6279 | 1.0 | 0.0459 | 0.0473 | 35.11 | 202.37 | 6,980 queries, 8.84M documents | +| nfcorpus | 0.3301 | 0.4396 | 2.4087 | 0.5292 | 0.1233 | 0.1383 | 3464.66 | 0.99 | 323 queries, 3.6K documents | +| nq | 0.2451 | 0.1272 | 0.4574 | 0.2099 | 0.1934 | 0.1240 | 150.23 | 71.43 | 3,452 queries, 2.68M documents | +| quora | 0.7705 | 0.6783 | 1.1749 | 0.7606 | 0.7206 | 0.6502 | 741.13 | 3.78 | 10,000 queries, 523K documents | +| scidocs | 0.1025 | 0.1790 | 0.8240 | 0.2754 | 0.0154 | 0.0275 | 879.11 | 4.46 | 1,000 queries, 25K documents | +| scifact | 0.6908 | 0.5533 | 0.9133 | 0.6527 | 0.6416 | 0.5468 | 2153.64 | 1.22 | 300 queries, 5K documents | +| trec-covid | 0.9533 | 1.0 | 9.4800 | 1.0 | 0.0074 | 0.0077 | 112.38 | 22.15 | 50 queries, 171K documents | +| webis-touche2020 | 0.4130 | 0.5510 | 3.7347 | 0.7114 | 0.0564 | 0.0827 | 104.65 | 44.14 | 49 queries, 382K documents | -documents = [ - { - "id": 0, - "title": "Hotel California", - "style": "rock", - "date": "1977-02-22", - "popularity": 9, - }, - { - "id": 1, - "title": "Here Comes the Sun", - "style": "rock", - "date": "1969-06-10", - "popularity": 10, - }, - { - "id": 2, - "title": "Alive", - "style": "electro, punk", - "date": "2007-11-19", - "popularity": 9, - }, -] - -upload.documents( - database="ducksearch.duckdb", - key="id", - fields=["title", "style", "date", "popularity"], - documents=documents, - dtypes={ - "date": "DATE", - "popularity": "INT", - }, -) - -# Mapping between documents ids and queries -documents_queries = { - 0: ["the beatles", "rock band"], - 1: ["rock band", "california"], - 2: ["daft"], -} - -upload.queries( - database="ducksearch.duckdb", - documents_queries=documents_queries, -) - -search.graphs( - database="ducksearch.duckdb", - queries="daft punk", - top_k=10, -) -``` - -```python -[ - { - "id": "2", - "title": "Alive", - "style": "electro, punk", - "date": Timestamp("2007-11-19 00:00:00"), - "popularity": 9, - "score": 2.877532958984375, - } -] -``` +## References -## Lightning fast +- [DuckDB](https://duckdb.org/) +- [DuckDB Full Text Search](https://duckdb.org/docs/extensions/full_text_search.html): Note that DuckSearch rely partially on the DuckDB Full Text Search extension but accelerate the search process via `top_k_token` approximation, pre-computation of scores and multi-threading. ## License diff --git a/benchmark.py b/benchmark.py index be7a304..b4a658a 100644 --- a/benchmark.py +++ b/benchmark.py @@ -1,39 +1,40 @@ -import time +from nltk import download +from nltk.corpus import stopwords from ducksearch import evaluation, search, upload +download("stopwords") + +stopword = list(stopwords.words("english")) + dataset_name = "quora" documents, queries, qrels = evaluation.load_beir( - dataset_name=dataset_name, split="test" + dataset_name=dataset_name, + split="test", ) upload.documents( - database=dataset_name, documents=documents, key="id", fields=["title", "text"] + database=dataset_name, + documents=documents, + key="id", + fields=["title", "text"], + stopwords=stopword, ) -upload.indexes(database=dataset_name) - - -start = time.time() - scores = search.documents( database=dataset_name, queries=queries, top_k=10, - top_k_token=10_000, - batch_size=30, + top_k_token=30_000, + batch_size=32, ) -end = time.time() - -print(f"Search took {end - start:.2f} seconds, QPS: {len(queries) / (end - start):.2f}") - evaluation_scores = evaluation.evaluate( scores=scores, qrels=qrels, queries=queries, - metrics=["ndcg@10", "hits@1", "hits@2", "hits@3", "hits@4", "hits@5", "hits@10"], + metrics=["ndcg@10", "hits@1", "hits@10", "mrr@10", "map@10", "r-precision"], ) print(evaluation_scores) diff --git a/docs/api/evaluation/evaluate.md b/docs/api/evaluation/evaluate.md index 1f57c9d..9391c6a 100644 --- a/docs/api/evaluation/evaluate.md +++ b/docs/api/evaluation/evaluate.md @@ -47,16 +47,5 @@ Evaluate the performance of document retrieval using relevance judgments. ... queries=queries, ... top_k=10, ... ) - ->>> evaluation_scores = evaluation.evaluate( -... scores=scores, -... qrels=qrels, -... queries=queries, -... metrics=["ndcg@10", "hits@1", "hits@2", "hits@3", "hits@4", "hits@5", "hits@10"], -... ) - ->>> assert evaluation_scores["ndcg@10"] > 0.68 ->>> assert evaluation_scores["hits@1"] > 0.54 ->>> assert evaluation_scores["hits@10"] > 0.90 ``` diff --git a/docs/api/search/documents.md b/docs/api/search/documents.md index 31eb637..50fb956 100644 --- a/docs/api/search/documents.md +++ b/docs/api/search/documents.md @@ -14,7 +14,7 @@ Search for documents in the documents table using specified queries. A string or list of query strings to search for. -- **batch_size** (*int*) – defaults to `30` +- **batch_size** (*int*) – defaults to `32` The batch size for query processing. @@ -22,7 +22,7 @@ Search for documents in the documents table using specified queries. The number of top documents to retrieve for each query. -- **top_k_token** (*int*) – defaults to `10000` +- **top_k_token** (*int*) – defaults to `30000` The number of documents to score per token. @@ -38,17 +38,22 @@ Search for documents in the documents table using specified queries. Optional SQL filters to apply during the search. -- **kwargs** - ## Examples ```python >>> from ducksearch import evaluation, upload, search ->>> documents, queries, qrels = evaluation.load_beir("scifact", split="test") ->>> scores = search.documents(database="test.duckdb", queries=queries, top_k_token=1000) ->>> evaluation_scores = evaluation.evaluate(scores=scores, qrels=qrels, queries=queries) ->>> assert evaluation_scores["ndcg@10"] > 0.68 + +>>> documents, queries, qrels = evaluation.load_beir( +... "scifact", +... split="test", +... ) + +>>> scores = search.documents( +... database="test.duckdb", +... queries=queries, +... top_k_token=1000, +... ) ``` diff --git a/docs/api/search/graphs.md b/docs/api/search/graphs.md index 3ed008d..bb2339a 100644 --- a/docs/api/search/graphs.md +++ b/docs/api/search/graphs.md @@ -22,7 +22,7 @@ Search for graphs in DuckDB using the provided queries. The number of top documents to retrieve for each query. -- **top_k_token** (*int*) – defaults to `10000` +- **top_k_token** (*int*) – defaults to `30000` The number of top tokens to retrieve. @@ -65,16 +65,5 @@ Search for graphs in DuckDB using the provided queries. ... queries=queries, ... top_k=10, ... ) - ->>> assert len(scores) > 0 - ->>> evaluation_scores = evaluation.evaluate( -... scores=scores, -... qrels=qrels, -... queries=queries, -... metrics=["ndcg@10", "hits@1", "hits@10"] -... ) - ->>> assert evaluation_scores["ndcg@10"] > 0.74 ``` diff --git a/docs/api/search/queries.md b/docs/api/search/queries.md index e3961b5..bd2c8fe 100644 --- a/docs/api/search/queries.md +++ b/docs/api/search/queries.md @@ -14,7 +14,7 @@ Search for queries in the queries table using specified queries. A string or list of query strings to search for. -- **batch_size** (*int*) – defaults to `30` +- **batch_size** (*int*) – defaults to `32` The batch size for query processing. @@ -22,7 +22,7 @@ Search for queries in the queries table using specified queries. The number of top matching queries to retrieve. -- **top_k_token** (*int*) – defaults to `10000` +- **top_k_token** (*int*) – defaults to `30000` The number of documents to score per token. @@ -38,8 +38,6 @@ Search for queries in the queries table using specified queries. Optional SQL filters to apply during the search. -- **kwargs** - ## Examples diff --git a/docs/api/search/search.md b/docs/api/search/search.md index fd64555..31514f1 100644 --- a/docs/api/search/search.md +++ b/docs/api/search/search.md @@ -26,7 +26,7 @@ Run the search for documents or queries in parallel. A string or list of query strings to search for. -- **batch_size** (*int*) – defaults to `30` +- **batch_size** (*int*) – defaults to `64` The batch size for query processing. @@ -34,7 +34,7 @@ Run the search for documents or queries in parallel. The number of top results to retrieve for each query. -- **top_k_token** (*int*) – defaults to `10000` +- **top_k_token** (*int*) – defaults to `30000` The number of documents to score per token. diff --git a/docs/api/search/update-index-documents.md b/docs/api/search/update-index-documents.md index 0101b64..35800b2 100644 --- a/docs/api/search/update-index-documents.md +++ b/docs/api/search/update-index-documents.md @@ -22,7 +22,7 @@ Update the BM25 search index for documents. The stemming algorithm to use (e.g., 'porter'). -- **stopwords** (*str | list[str]*) – defaults to `english` +- **stopwords** (*str | list[str]*) – defaults to `None` The list of stopwords to exclude from indexing. Can be a list or a string specifying the language (e.g., "english"). diff --git a/docs/api/search/update-index-queries.md b/docs/api/search/update-index-queries.md index a223a2e..c8c8576 100644 --- a/docs/api/search/update-index-queries.md +++ b/docs/api/search/update-index-queries.md @@ -22,7 +22,7 @@ Update the BM25 search index for queries. The stemming algorithm to use (e.g., 'porter'). -- **stopwords** (*str | list[str]*) – defaults to `english` +- **stopwords** (*str | list[str]*) – defaults to `None` The list of stopwords to exclude from indexing. Can be a list or a string specifying the language (e.g., "english"). diff --git a/docs/api/upload/documents.md b/docs/api/upload/documents.md index d30c09a..60691a6 100644 --- a/docs/api/upload/documents.md +++ b/docs/api/upload/documents.md @@ -34,7 +34,7 @@ Upload documents to DuckDB, create necessary schema, and index using BM25. Stemming algorithm to use (e.g., 'porter'). The type of stemmer to be used. One of 'arabic', 'basque', 'catalan', 'danish', 'dutch', 'english', 'finnish', 'french', 'german', 'greek', 'hindi', 'hungarian', 'indonesian', 'irish', 'italian', 'lithuanian', 'nepali', 'norwegian', 'porter', 'portuguese', 'romanian', 'russian', 'serbian', 'spanish', 'swedish', 'tamil', 'turkish', or 'none' if no stemming is to be used. -- **stopwords** (*str | list[str]*) – defaults to `english` +- **stopwords** (*str | list[str]*) – defaults to `None` List of stopwords to exclude from indexing. Can be a custom list or a language string. diff --git a/docs/api/upload/queries.md b/docs/api/upload/queries.md index 87c9223..6b8e0ad 100644 --- a/docs/api/upload/queries.md +++ b/docs/api/upload/queries.md @@ -30,7 +30,7 @@ Upload queries to DuckDB, map documents to queries, and index using BM25. Stemming algorithm to use (e.g., 'porter'). The type of stemmer to be used. One of 'arabic', 'basque', 'catalan', 'danish', 'dutch', 'english', 'finnish', 'french', 'german', 'greek', 'hindi', 'hungarian', 'indonesian', 'irish', 'italian', 'lithuanian', 'nepali', 'norwegian', 'porter', 'portuguese', 'romanian', 'russian', 'serbian', 'spanish', 'swedish', 'tamil', 'turkish', or 'none' if no stemming is to be used. -- **stopwords** (*str | list[str]*) – defaults to `english` +- **stopwords** (*str | list[str]*) – defaults to `None` List of stopwords to exclude from indexing. Default can be a custom list or a language string. diff --git a/docs/benchmarks/benchmarks.md b/docs/benchmarks/benchmarks.md index e983feb..a9e772a 100644 --- a/docs/benchmarks/benchmarks.md +++ b/docs/benchmarks/benchmarks.md @@ -1,2 +1,96 @@ ## Benchmarks +### Ducksearch and BM25s + +While DuckSearch provide advanced filtering features / updates on the index, DuckSearch only score `top_k_token` document per query token. Benchmark might evolve with DuckDB improvements and DuckSearch updates. + +=== "Table" + + | Dataset | Metric | Ducksearch | BM25s | Difference (Ducksearch - BM25s) | + |-------------------|---------------|-------------|-----------|---------------------------------| + | **arguana** | ndcg@10 | 0.3779 | 0.3663 | +0.0116 | + | | hits@1 | 0.0 | 0.0 | 0.0 | + | | mrr@10 | 0.2491 | 0.2443 | +0.0048 | + | | map@10 | 0.2528 | 0.2430 | +0.0098 | + | | qps | 117.80 | 2113.50 | -1995.70 | + | | Index Time(s) | 1.42 | 0.48 | +0.94 | + | **climate-fever** | ndcg@10 | 0.1184 | 0.1313 | -0.0129 | + | | hits@1 | 0.1068 | 0.1186 | -0.0118 | + | | mrr@10 | 0.1644 | 0.1809 | -0.0165 | + | | map@10 | 0.0803 | 0.0907 | -0.0104 | + | | qps | 5.88 | 99.49 | -93.61 | + | | Index Time(s) | 302.39 | 209.97 | +92.42 | + | **dbpedia-entity**| ndcg@10 | 0.6046 | 0.6172 | -0.0126 | + | | hits@1 | 0.7669 | 0.7744 | -0.0075 | + | | mrr@10 | 0.8311 | 0.8382 | -0.0071 | + | | map@10 | 0.0649 | 0.0672 | -0.0023 | + | | qps | 113.20 | 182.79 | -69.59 | + | | Index Time(s) | 181.42 | 119.18 | +62.24 | + | **fever** | ndcg@10 | 0.3861 | 0.4825 | -0.0964 | + | | hits@1 | 0.2583 | 0.3312 | -0.0729 | + | | mrr@10 | 0.3525 | 0.4423 | -0.0898 | + | | map@10 | 0.3329 | 0.4212 | -0.0883 | + | | qps | 74.40 | 104.97 | -30.57 | + | | Index Time(s) | 329.70 | 207.52 | +122.18 | + | **fiqa** | ndcg@10 | 0.2445 | 0.2326 | +0.0119 | + | | hits@1 | 0.2207 | 0.2160 | +0.0047 | + | | mrr@10 | 0.3002 | 0.2875 | +0.0127 | + | | map@10 | 0.1848 | 0.1726 | +0.0122 | + | | qps | 545.77 | 2157.35 | -1611.58 | + | | Index Time(s) | 6.04 | 4.27 | +1.77 | + | **hotpotqa** | ndcg@10 | 0.4487 | 0.5630 | -0.1143 | + | | hits@1 | 0.5059 | 0.6523 | -0.1464 | + | | mrr@10 | 0.5846 | 0.7249 | -0.1403 | + | | map@10 | 0.3642 | 0.4697 | -0.1055 | + | | qps | 48.15 | 104.43 | -56.28 | + | | Index Time(s) | 163.14 | 123.39 | +39.75 | + | **msmarco** | ndcg@10 | 0.8951 | 0.9705 | -0.0754 | + | | hits@1 | 1.0 | 1.0 | 0.0 | + | | mrr@10 | 1.0 | 1.0 | 0.0 | + | | map@10 | 0.0459 | 0.0532 | -0.0073 | + | | qps | 35.11 | 71.26 | -36.15 | + | | Index Time(s) | 202.37 | 229.22 | -26.85 | + | **nfcorpus** | ndcg@10 | 0.3301 | 0.3059 | +0.0242 | + | | hits@1 | 0.4396 | 0.4458 | -0.0062 | + | | mrr@10 | 0.5292 | 0.5205 | +0.0087 | + | | map@10 | 0.1233 | 0.1168 | +0.0065 | + | | qps | 3464.66 | 3933.12 | -468.46 | + | | Index Time(s) | 0.99 | 1.67 | -0.68 | + | **nq** | ndcg@10 | 0.2451 | 0.2735 | -0.0284 | + | | hits@1 | 0.1272 | 0.1460 | -0.0188 | + | | mrr@10 | 0.2099 | 0.2366 | -0.0267 | + | | map@10 | 0.1934 | 0.2177 | -0.0243 | + | | qps | 150.23 | 272.62 | -122.39 | + | | Index Time(s) | 71.43 | 87.98 | -16.55 | + | **quora** | ndcg@10 | 0.7705 | 0.7491 | +0.0214 | + | | hits@1 | 0.6783 | 0.6622 | +0.0161 | + | | mrr@10 | 0.7606 | 0.7433 | +0.0173 | + | | map@10 | 0.7206 | 0.6988 | +0.0218 | + | | qps | 741.13 | 1004.44 | -263.31 | + | | Index Time(s) | 3.78 | 6.57 | -2.79 | + | **scidocs** | ndcg@10 | 0.1025 | 0.0993 | +0.0032 | + | | hits@1 | 0.1790 | 0.1910 | -0.0120 | + | | mrr@10 | 0.2754 | 0.2765 | -0.0011 | + | | map@10 | 0.0154 | 0.0147 | +0.0007 | + | | qps | 879.11 | 3570.06 | -2690.95 | + | | Index Time(s) | 4.46 | 1.64 | +2.82 | + | **scifact** | ndcg@10 | 0.6908 | 0.6617 | +0.0291 | + | | hits@1 | 0.5533 | 0.5433 | +0.0100 | + | | mrr@10 | 0. + 6527 | 0.6312 | +0.0215 | + | | map@10 | 0.6416 | 0.6199 | +0.0217 | + | | qps | 2153.64 | 3708.28 | -1554.64 | + | | Index Time(s) | 1.22 | 0.41 | +0.81 | + | **trec-covid** | ndcg@10 | 0.9533 | 0.8983 | +0.0550 | + | | hits@1 | 1.0 | 0.92 | +0.08 | + | | mrr@10 | 1.0 | 0.96 | +0.04 | + | | map@10 | 0.0074 | 0.0069 | +0.0005 | + | | qps | 112.38 | 1275.41 | -1163.03 | + | | Index Time(s) | 22.15 | 10.15 | +12.00 | + | **webis-touche2020** | ndcg@10 | 0.4130 | 0.4671 | -0.0541 | + | | hits@1 | 0.5510 | 0.6122 | -0.0612 | + | | mrr@10 | 0.7114 | 0.7541 | -0.0427 | + | | map@10 | 0.0564 | 0.0659 | -0.0095 | + | | qps | 104.65 | 961.73 | -857.08 | + | | Index Time(s) | 44.14 | 34.89 | +9.25 | + diff --git a/docs/documentation/graph.md b/docs/documentation/graph.md new file mode 100644 index 0000000..4851d27 --- /dev/null +++ b/docs/documentation/graph.md @@ -0,0 +1,100 @@ +## Graph + +The `search.graphs` function can be used to search documents with a graph-based query. This function is useful if we have paired documents and queries. The search will retrieve the set of documents and queries that match the input query. Then it will build a graph and compute the weight of each document using a graph-based scoring function. + +The `search.graphs` function is much slower than the `search.documents` function, but might provide better results with decent amount of paired documents / queries. + +### Documents queries interactions + +We can upload documents queries interactions in order to call the `search.graphs` function. The following example demonstrates how to upload documents queries interactions: + +```python +from ducksearch import search, upload + +documents = [ + { + "id": 0, + "title": "Hotel California", + "style": "rock", + "date": "1977-02-22", + "popularity": 9, + }, + { + "id": 1, + "title": "Here Comes the Sun", + "style": "rock", + "date": "1969-06-10", + "popularity": 10, + }, + { + "id": 2, + "title": "Alive", + "style": "electro, punk", + "date": "2007-11-19", + "popularity": 9, + }, +] + +upload.documents( + database="ducksearch.duckdb", + key="id", + fields=["title", "style", "date", "popularity"], + documents=documents, + dtypes={ + "date": "DATE", + "popularity": "INT", + }, +) + +# Mapping between documents ids and queries +documents_queries = { + 0: ["the beatles", "rock band"], + 1: ["rock band", "california"], + 2: ["daft"], +} + +upload.queries( + database="ducksearch.duckdb", + documents_queries=documents_queries, +) +``` + +???+ tip + We can write documents queries mapping as a list of dict with the weight between the document and the query. The weight is used to compute the score in the `search.graphs` function: + + ```python + documents_queries = { + 0: {"the beatles": 30, "rock band": 10}, + 1: {"rock band": 10, "california": 1}, + 2: {"daft": 60}, + } + ``` + + When the weight is not specified, the default value is 1. + +### Search Graphs + +The following example demonstrates how to search documents with a graph-based query: + +```python +from ducksearch import search + +search.graphs( + database="ducksearch.duckdb", + queries="daft punk", + top_k=10, +) +``` + +```python +[ + { + "id": "2", + "title": "Alive", + "style": "electro, punk", + "date": Timestamp("2007-11-19 00:00:00"), + "popularity": 9, + "score": 2.877532958984375, + } +] +``` \ No newline at end of file diff --git a/docs/documentation/search.md b/docs/documentation/search.md index 4eea581..dfe1ecd 100644 --- a/docs/documentation/search.md +++ b/docs/documentation/search.md @@ -16,7 +16,7 @@ search.documents( queries=["daft punk", "rock"], top_k=10, top_k_token=10_000, - batch_size=30, + batch_size=32, n_jobs=-1, ) ``` @@ -69,7 +69,7 @@ search.documents( queries=["rock", "california"], top_k=10, top_k_token=10_000, - batch_size=30, + batch_size=32, filters="YEAR(date) <= 1990 AND YEAR(date) >= 1970", n_jobs=-1, ) @@ -103,35 +103,3 @@ search.documents( ???+ info The filters are evaluated by DuckDB, so all DuckDB functions are available for use in the filters. You can find more information about DuckDB functions in the [DuckDB documentation](https://duckdb.org/docs/sql/functions/overview). -### Graphs - -???+ info - To benefit from the `search.graphs` function, we need to upload documents and queries to DuckDB using the `upload.documents` and `upload.queries` functions. - -The `search.graphs` function retrieves the top documents. Then it retrieves the top queries indexed from `upload.queries`. Finally, it computes a graph-based ranking of the documents based on the queries. - -```python -from ducksearch import search - -search.graphs( - database="ducksearch.duckdb", - queries="daft punk", - top_k=10, - top_k_token=10_000, - batch_size=30, - n_jobs=-1, -) -``` - -```python -[ - { - "id": "2", - "title": "Alive", - "style": "electro, punk", - "date": Timestamp("2007-11-19 00:00:00"), - "popularity": 9, - "score": 0.17841622233390808, - } -] -``` \ No newline at end of file diff --git a/docs/documentation/upload.md b/docs/documentation/upload.md index b0072ff..7a540df 100644 --- a/docs/documentation/upload.md +++ b/docs/documentation/upload.md @@ -82,78 +82,3 @@ upload.documents( ???+ info More informations about DuckDB and HuggingFace compatibility can be found [here](https://huggingface.co/docs/hub/en/datasets-duckdb) and [here](https://duckdb.org/2024/05/29/access-150k-plus-datasets-from-hugging-face-with-duckdb.html). - - -### Documents queries interactions - -We can upload documents queries interactions in order to call the `search.graphs` function. The following example demonstrates how to upload documents queries interactions: - -```python -from ducksearch import search, upload - -documents = [ - { - "id": 0, - "title": "Hotel California", - "style": "rock", - "date": "1977-02-22", - "popularity": 9, - }, - { - "id": 1, - "title": "Here Comes the Sun", - "style": "rock", - "date": "1969-06-10", - "popularity": 10, - }, - { - "id": 2, - "title": "Alive", - "style": "electro, punk", - "date": "2007-11-19", - "popularity": 9, - }, -] - -upload.documents( - database="ducksearch.duckdb", - key="id", - fields=["title", "style", "date", "popularity"], - documents=documents, - dtypes={ - "date": "DATE", - "popularity": "INT", - }, -) - -# Mapping between documents ids and queries -documents_queries = { - 0: ["the beatles", "rock band"], - 1: ["rock band", "california"], - 2: ["daft"], -} - -upload.queries( - database="ducksearch.duckdb", - documents_queries=documents_queries, -) - -search.graphs( - database="ducksearch.duckdb", - queries="daft punk", - top_k=10, -) -``` - -???+ tip - We can write documents queries mapping as a list of dict with the weight between the document and the query. The weight is used to compute the score in the `search.graphs` function: - - ```python - documents_queries = { - 0: {"the beatles": 30, "rock band": 10}, - 1: {"rock band": 10, "california": 1}, - 2: {"daft": 60}, - } - ``` - - When the weight is not specified, the default value is 1. diff --git a/docs/index.md b/docs/index.md index bba7b7f..7636083 100644 --- a/docs/index.md +++ b/docs/index.md @@ -13,30 +13,26 @@

-DuckSearch is a lightweight and easy-to-use library that allows to index and search documents. DuckSearch is built on top of DuckDB, a high-performance analytical database. DuckDB is designed to execute analytical SQL queries fast, and DuckSearch leverages this to provide efficient and scallable search / filtering capabilities. +DuckSearch is a lightweight and easy-to-use library to search documents. DuckSearch is built on top of DuckDB, a high-performance analytical database. DuckDB is designed to execute analytical SQL queries fast, and DuckSearch leverages this to provide efficient search and filtering features. DuckSearch index can be updated with new documents and documents can be deleted as well. + +DuckSearch also supports HuggingFace datasets, allowing to index datasets directly from the HuggingFace Hub.

## Installation -We can install DuckSearch using pip: +Install DuckSearch using pip: ```bash pip install ducksearch ``` -For evaluation dependencies, we can install DuckSearch with the `eval` extra: - -```bash -pip install "ducksearch[eval]" -``` - ## Documentation The complete documentation is available [here](https://lightonai.github.io/ducksearch/), which includes in-depth guides, examples, and API references. ### Upload -We can upload documents to DuckDB using the `upload.documents` function. The documents are stored in a DuckDB database, and the fields are indexed with BM25. +We can upload documents to DuckDB using the `upload.documents` function. The documents are stored in a DuckDB database, and the `fields` are indexed with BM25. ```python from ducksearch import upload @@ -79,7 +75,7 @@ upload.documents( ## Search -We can search documents using the `search.documents` function. The function returns the documents that match the query, sorted by the BM25 score. The `top_k` parameter controls the number of documents to return, and the `top_k_token` parameter controls the number of documents to score for each query token. Increasing `top_k_token` can improve the quality of the results but also increase the computation time. +`search.documents` returns a list of list of documents ordered by relevance. We can control the number of documents to return using the `top_k` parameter. The following example demonstrates how to search for documents with the queries "punk" and "california" while filtering the results to include only documents with a date after 1970 and a popularity score greater than 8. ```python from ducksearch import search @@ -88,7 +84,7 @@ search.documents( database="ducksearch.duckdb", queries=["punk", "california"], top_k=10, - top_k_token=10_000, + filters="YEAR(date) >= 1970 AND popularity > 8", ) ``` @@ -117,36 +113,22 @@ search.documents( ] ``` -### Filters +Filters are SQL expressions that are applied to the search results. We can use every filtering function DuckDB provides such as [date functions](https://duckdb.org/docs/sql/functions/date). -We can also filter the results using SQL syntax which will be evaluated by DuckDB, therefore all DuckDB functions are available. +## Delete and update index + +We can delete documents and update the BM25 weights accordingly using the `delete.documents` function. ```python -from ducksearch import search +from ducksearch import delete -search.documents( +delete.documents( database="ducksearch.duckdb", - queries="rock", - top_k=10, - top_k_token=10_000, - filters="YEAR(date) < 1970 AND popularity > 9 AND style LIKE '%rock%'", + ids=[0, 1], ) ``` -```python -[ - { - "score": 0.08725740015506744, - "id": "1", - "title": "Here Comes the Sun", - "style": "rock", - "date": Timestamp("1969-06-10 00:00:00"), - "popularity": 10, - } -] -``` - -List of DuckDB functions such as date functions can be found [here](https://duckdb.org/docs/sql/functions/date). +To update the index, we should first delete the documents and then upload the updated documents. ## Extra features @@ -181,13 +163,10 @@ search.documents( database="fineweb.duckdb", queries="earth science", top_k=2, - top_k_token=10_000, - filters="token_count > 200", ) ``` ```python - [ { "id": "", @@ -211,82 +190,25 @@ search.documents( ] ``` -### Graphs - -The `search.graphs` function can be used to search documents with a graph query. This function is useful if we have paired documents and queries. The search will retrieve the set of documents and queries that match the input query. Then it will build a graph and compute the weight of each document using a graph-based scoring function. - -```python -from ducksearch import search, upload - -documents = [ - { - "id": 0, - "title": "Hotel California", - "style": "rock", - "date": "1977-02-22", - "popularity": 9, - }, - { - "id": 1, - "title": "Here Comes the Sun", - "style": "rock", - "date": "1969-06-10", - "popularity": 10, - }, - { - "id": 2, - "title": "Alive", - "style": "electro, punk", - "date": "2007-11-19", - "popularity": 9, - }, -] - -upload.documents( - database="ducksearch.duckdb", - key="id", - fields=["title", "style", "date", "popularity"], - documents=documents, - dtypes={ - "date": "DATE", - "popularity": "INT", - }, -) - -documents_queries = { - 0: ["the beatles", "rock band"], - 1: ["rock band", "california"], - 2: ["daft"], -} - -upload.queries( - database="ducksearch.duckdb", - documents_queries=documents_queries, -) - -search.graphs( - database="ducksearch.duckdb", - queries="daft punk", - top_k=10, - top_k_token=10_000, -) -``` - -```python -[ - { - "id": "2", - "title": "Alive", - "style": "electro, punk", - "date": Timestamp("2007-11-19 00:00:00"), - "popularity": 9, - "score": 2.877532958984375, - } -] -``` - -## Lightning fast - +## Benchmark + + +| Dataset | ndcg@10 | hits@1 | hits@10 | mrr@10 | map@10 | r-precision | qps | Indexation Time (s) | Number of Documents and Queries | +|-------------------|-----------|---------|----------|----------|---------|-------------|----------------|---------------------|--------------------------------| +| arguana | 0.3779 | 0.0 | 0.8267 | 0.2491 | 0.2528 | 0.0108 | 117.80 | 1.42 | 1,406 queries, 8.67K documents | +| climate-fever | 0.1184 | 0.1068 | 0.3648 | 0.1644 | 0.0803 | 0.0758 | 5.88 | 302.39 | 1,535 queries, 5.42M documents | +| dbpedia-entity | 0.6046 | 0.7669 | 5.6241 | 0.8311 | 0.0649 | 0.0741 | 113.20 | 181.42 | 400 queries, 4.63M documents | +| fever | 0.3861 | 0.2583 | 0.5826 | 0.3525 | 0.3329 | 0.2497 | 74.40 | 329.70 | 6,666 queries, 5.42M documents | +| fiqa | 0.2445 | 0.2207 | 0.6790 | 0.3002 | 0.1848 | 0.1594 | 545.77 | 6.04 | 648 queries, 57K documents | +| hotpotqa | 0.4487 | 0.5059 | 0.9699 | 0.5846 | 0.3642 | 0.3388 | 48.15 | 163.14 | 7,405 queries, 5.23M documents | +| msmarco | 0.8951 | 1.0 | 8.6279 | 1.0 | 0.0459 | 0.0473 | 35.11 | 202.37 | 6,980 queries, 8.84M documents | +| nfcorpus | 0.3301 | 0.4396 | 2.4087 | 0.5292 | 0.1233 | 0.1383 | 3464.66 | 0.99 | 323 queries, 3.6K documents | +| nq | 0.2451 | 0.1272 | 0.4574 | 0.2099 | 0.1934 | 0.1240 | 150.23 | 71.43 | 3,452 queries, 2.68M documents | +| quora | 0.7705 | 0.6783 | 1.1749 | 0.7606 | 0.7206 | 0.6502 | 741.13 | 3.78 | 10,000 queries, 523K documents | +| scidocs | 0.1025 | 0.1790 | 0.8240 | 0.2754 | 0.0154 | 0.0275 | 879.11 | 4.46 | 1,000 queries, 25K documents | +| scifact | 0.6908 | 0.5533 | 0.9133 | 0.6527 | 0.6416 | 0.5468 | 2153.64 | 1.22 | 300 queries, 5K documents | +| trec-covid | 0.9533 | 1.0 | 9.4800 | 1.0 | 0.0074 | 0.0077 | 112.38 | 22.15 | 50 queries, 171K documents | +| webis-touche2020 | 0.4130 | 0.5510 | 3.7347 | 0.7114 | 0.0564 | 0.0827 | 104.65 | 44.14 | 49 queries, 382K documents | ## License diff --git a/ducksearch/delete/documents.py b/ducksearch/delete/documents.py index 8b7cf23..e8fdebc 100644 --- a/ducksearch/delete/documents.py +++ b/ducksearch/delete/documents.py @@ -21,6 +21,13 @@ def _update_score() -> None: """Update the score after deleting documents.""" +@execute_with_duckdb( + relative_path="delete/update/df.sql", +) +def _update_df() -> None: + """Update the token frequency deleting documents.""" + + @execute_with_duckdb( relative_path="delete/delete/scores.sql", ) @@ -131,6 +138,12 @@ def documents( config=config, ) + _update_df( + database=database, + parquet_file="_documents_ids.parquet", + config=config, + ) + _update_terms( database=database, parquet_file="_documents_ids.parquet", diff --git a/ducksearch/delete/update/df.sql b/ducksearch/delete/update/df.sql new file mode 100644 index 0000000..60926d4 --- /dev/null +++ b/ducksearch/delete/update/df.sql @@ -0,0 +1,22 @@ +WITH _docs_to_delete AS ( + SELECT DISTINCT + bm25.docid + FROM parquet_scan('{parquet_file}') p + INNER JOIN bm25_documents.docs bm25 + ON p.id = bm25.name +), + +_tf AS ( + SELECT + termid, + sum(tf) as df + FROM bm25_documents.terms + INNER JOIN _docs_to_delete + ON bm25_documents.terms.docid = _docs_to_delete.docid + GROUP BY 1 +) + +UPDATE bm25_documents.dict _dict +SET df = GREATEST(_dict.df - _tf.df, 0) +FROM _tf +WHERE _dict.termid = _tf.termid; \ No newline at end of file diff --git a/ducksearch/evaluation/evaluation.py b/ducksearch/evaluation/evaluation.py index 887ad51..2c4cc69 100644 --- a/ducksearch/evaluation/evaluation.py +++ b/ducksearch/evaluation/evaluation.py @@ -118,17 +118,6 @@ def evaluate( ... top_k=10, ... ) - >>> evaluation_scores = evaluation.evaluate( - ... scores=scores, - ... qrels=qrels, - ... queries=queries, - ... metrics=["ndcg@10", "hits@1", "hits@2", "hits@3", "hits@4", "hits@5", "hits@10"], - ... ) - - >>> assert evaluation_scores["ndcg@10"] > 0.68 - >>> assert evaluation_scores["hits@1"] > 0.55 - >>> assert evaluation_scores["hits@10"] > 0.88 - """ from ranx import Qrels, Run, evaluate diff --git a/ducksearch/search/create/queries_index.sql b/ducksearch/search/create/queries_index.sql index 140cd23..0dcaa87 100644 --- a/ducksearch/search/create/queries_index.sql +++ b/ducksearch/search/create/queries_index.sql @@ -8,4 +8,4 @@ PRAGMA CREATE_FTS_INDEX( STRIP_ACCENTS={strip_accents}, LOWER={lower}, OVERWRITE=1 -); +); \ No newline at end of file diff --git a/ducksearch/search/graphs.py b/ducksearch/search/graphs.py index c106753..1e34895 100644 --- a/ducksearch/search/graphs.py +++ b/ducksearch/search/graphs.py @@ -141,12 +141,23 @@ def graphs( ... fields=["title", "text"], ... documents=documents, ... ) + | Table | Size | + |----------------|------| + | documents | 5183 | + | bm25_documents | 5183 | >>> upload.queries( ... database="test.duckdb", ... queries=queries, ... documents_queries=qrels, ... ) + | Table | Size | + |-------------------|------| + | documents | 5183 | + | queries | 807 | + | bm25_documents | 5183 | + | bm25_queries | 807 | + | documents_queries | 916 | >>> scores = search.graphs( ... database="test.duckdb", @@ -154,17 +165,6 @@ def graphs( ... top_k=10, ... ) - >>> assert len(scores) > 0 - - >>> evaluation_scores = evaluation.evaluate( - ... scores=scores, - ... qrels=qrels, - ... queries=queries, - ... metrics=["ndcg@10", "hits@1", "hits@10"] - ... ) - - >>> assert evaluation_scores["ndcg@10"] > 0.74 - """ resource.setrlimit( resource.RLIMIT_CORE, (resource.RLIM_INFINITY, resource.RLIM_INFINITY) diff --git a/ducksearch/search/select.py b/ducksearch/search/select.py index cf7a1c4..a0ce629 100644 --- a/ducksearch/search/select.py +++ b/ducksearch/search/select.py @@ -94,14 +94,6 @@ def documents( ... top_k_token=1000, ... ) - >>> evaluation_scores = evaluation.evaluate( - ... scores=scores, - ... qrels=qrels, - ... queries=queries - ... ) - - >>> assert evaluation_scores["ndcg@10"] > 0.68 - """ return search( database=database, @@ -311,8 +303,7 @@ def search( ... top_k=10, ... ) - >>> assert len(documents) == 1 - >>> assert len(documents[0]) == 10 + >>> assert len(documents) == 10 """ is_query_str = False @@ -352,7 +343,9 @@ def search( config=config, ) - _matchs = Parallel(n_jobs=n_jobs, backend="threading")( + matchs = [] + + for match in Parallel(n_jobs=n_jobs, backend="threading")( delayed(_search)( database, schema, @@ -368,10 +361,7 @@ def search( for index, batch_queries in enumerate( batchify(queries, batch_size=batch_size, desc="Searching") ) - ) - - matchs = [] - for match in _matchs: + ): matchs.extend(match) return matchs[0] if is_query_str else matchs diff --git a/ducksearch/search/select/search.sql b/ducksearch/search/select/search.sql index 9f6c4fe..b7f3932 100644 --- a/ducksearch/search/select/search.sql +++ b/ducksearch/search/select/search.sql @@ -39,9 +39,7 @@ _matchs_scores AS ( bm25id, SUM(score) AS score FROM _matchs - WHERE score IS NOT NULL AND score > 0 GROUP BY 1, 2 - ), _partition_scores AS ( @@ -49,8 +47,9 @@ _partition_scores AS ( query, bm25id, score, - ROW_NUMBER() OVER (PARTITION BY query ORDER BY score DESC, RANDOM() ASC) AS rank + RANK() OVER (PARTITION BY query ORDER BY score DESC, RANDOM() ASC) AS rank FROM _matchs_scores + QUALIFY rank <= {top_k} ) SELECT @@ -58,7 +57,6 @@ SELECT ps.score, ps.query AS _query FROM _partition_scores ps -JOIN {source_schema}.{source} s +INNER JOIN {source_schema}.{source} s ON ps.bm25id = s.bm25id -WHERE ps.rank <= {top_k} -ORDER BY ps.rank ASC; +ORDER BY score DESC; diff --git a/ducksearch/search/select/search_filters.sql b/ducksearch/search/select/search_filters.sql index bc9b10c..51bdbcb 100644 --- a/ducksearch/search/select/search_filters.sql +++ b/ducksearch/search/select/search_filters.sql @@ -56,11 +56,11 @@ _partition_scores AS ( _query, _score AS score, * EXCLUDE (_score, _query), - ROW_NUMBER() OVER (PARTITION BY _query ORDER BY _score DESC) AS _row_number + RANK() OVER (PARTITION BY _query ORDER BY _score DESC) AS _row_number FROM _filtered_scores + QUALIFY _row_number <= {top_k} ) SELECT * EXCLUDE (_row_number) -FROM _partition_scores -WHERE _row_number <= {top_k}; +FROM _partition_scores; diff --git a/ducksearch/search/select/termids_to_score.sql b/ducksearch/search/select/termids_to_score.sql index 374a4a6..9db3a34 100644 --- a/ducksearch/search/select/termids_to_score.sql +++ b/ducksearch/search/select/termids_to_score.sql @@ -2,11 +2,11 @@ WITH _terms_to_score AS ( SELECT term FROM fts_{schema}__documents.dict + ) -SELECT +SELECT DISTINCT d.termid FROM _terms_to_score t JOIN {schema}.dict d -ON t.term = d.term -ORDER BY RANDOM(); \ No newline at end of file +ON t.term = d.term; \ No newline at end of file